Realtime Compositor: Implement variable size bokeh blur

This patch implements the variable size blur option in the Bokeh Blur
node. The implementation is different from the CPU one in that it also
takes the Bounding Box input into account, which is ignored for some
reason for the CPU. Additionally, this implementation does not do the
optimization where the search radius is limited relative to the maximum
value in the size texture. That's because the cost of computing the
maximum is not worth it for most use cases.

The reference implementation does three unexpected things that are
replicated here nonetheless. First, the center bokeh weight is always
ignored and assumed to be 1. Second the size of the center pixel is
taken into account. Third, a unidimensional distance is used instead of
a 2D euclidean one. Those need to be considered independently.

Differential Revision: https://developer.blender.org/D16185

Reviewed By: Clement Foucault

3 files changed, 77 insertions, 0 deletions

diff --git a/source/blender/gpu/CMakeLists.txt b/source/blender/gpu/CMakeLists.txt
index 34d53eaa230..f387a4588b6 100644
--- a/source/blender/gpu/CMakeLists.txt
+++ b/source/blender/gpu/CMakeLists.txt
@@ -332,6 +332,7 @@ set(GLSL_SRC
   shaders/compositor/compositor_alpha_crop.glsl
   shaders/compositor/compositor_bilateral_blur.glsl
   shaders/compositor/compositor_blur.glsl
+  shaders/compositor/compositor_blur_variable_size.glsl
   shaders/compositor/compositor_bokeh_image.glsl
   shaders/compositor/compositor_box_mask.glsl
   shaders/compositor/compositor_convert.glsl
@@ -613,6 +614,7 @@ set(SRC_SHADER_CREATE_INFOS
   shaders/compositor/infos/compositor_alpha_crop_info.hh
   shaders/compositor/infos/compositor_bilateral_blur_info.hh
   shaders/compositor/infos/compositor_blur_info.hh
+  shaders/compositor/infos/compositor_blur_variable_size_info.hh
   shaders/compositor/infos/compositor_bokeh_image_info.hh
   shaders/compositor/infos/compositor_box_mask_info.hh
   shaders/compositor/infos/compositor_convert_info.hh
diff --git a/source/blender/gpu/shaders/compositor/compositor_blur_variable_size.glsl b/source/blender/gpu/shaders/compositor/compositor_blur_variable_size.glsl
new file mode 100644
index 00000000000..e7e5aac12a5
--- /dev/null
+++ b/source/blender/gpu/shaders/compositor/compositor_blur_variable_size.glsl
@@ -0,0 +1,60 @@
+#pragma BLENDER_REQUIRE(gpu_shader_common_math_utils.glsl)
+#pragma BLENDER_REQUIRE(gpu_shader_compositor_texture_utilities.glsl)
+
+/* Given the texel in the range [-radius, radius] in both axis, load the appropriate weight from
+ * the weights texture, where the texel (0, 0) is considered the center of weights texture. */
+vec4 load_weight(ivec2 texel, float radius)
+{
+  /* The center zero texel is always assigned a unit weight regardless of the corresponding weight
+   * in the weights texture. That's to guarantee that at last the center pixel will be accumulated
+   * even if the weights texture is zero at its center. */
+  if (texel == ivec2(0)) {
+    return vec4(1.0);
+  }
+
+  /* Add the radius to transform the texel into the range [0, radius * 2], then divide by the upper
+   * bound plus one to transform the texel into the normalized range [0, 1] needed to sample the
+   * weights sampler. Finally, also add 0.5 to sample at the center of the pixels. */
+  return texture(weights_tx, (texel + vec2(radius + 0.5)) / (radius * 2 + 1));
+}
+
+void main()
+{
+  ivec2 texel = ivec2(gl_GlobalInvocationID.xy);
+
+  /* The mask input is treated as a boolean. If it is zero, then no blurring happens for this
+   * pixel. Otherwise, the pixel is blurred normally and the mask value is irrelevant. */
+  float mask = texture_load(mask_tx, texel).x;
+  if (mask == 0.0) {
+    imageStore(output_img, texel, texture_load(input_tx, texel));
+    return;
+  }
+
+  float center_size = texture_load(size_tx, texel).x * base_size;
+
+  /* Go over the window of the given search radius and accumulate the colors multiplied by their
+   * respective weights as well as the weights themselves, but only if both the size of the center
+   * pixel and the size of the candidate pixel are less than both the x and y distances of the
+   * candidate pixel. */
+  vec4 accumulated_color = vec4(0.0);
+  vec4 accumulated_weight = vec4(0.0);
+  for (int y = -search_radius; y <= search_radius; y++) {
+    for (int x = -search_radius; x <= search_radius; x++) {
+      float candidate_size = texture_load(size_tx, texel + ivec2(x, y)).x * base_size;
+
+      /* Skip accumulation if either the x or y distances of the candidate pixel are larger than
+       * either the center or candidate pixel size. Note that the max and min functions here denote
+       * "either" in the aforementioned description. */
+      float size = min(center_size, candidate_size);
+      if (max(abs(x), abs(y)) > size) {
+        continue;
+      }
+
+      vec4 weight = load_weight(ivec2(x, y), size);
+      accumulated_color += texture_load(input_tx, texel + ivec2(x, y)) * weight;
+      accumulated_weight += weight;
+    }
+  }
+
+  imageStore(output_img, texel, safe_divide(accumulated_color, accumulated_weight));
+}
diff --git a/source/blender/gpu/shaders/compositor/infos/compositor_blur_variable_size_info.hh b/source/blender/gpu/shaders/compositor/infos/compositor_blur_variable_size_info.hh
new file mode 100644
index 00000000000..05b6385fd1e
--- /dev/null
+++ b/source/blender/gpu/shaders/compositor/infos/compositor_blur_variable_size_info.hh
@@ -0,0 +1,15 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later */
+
+#include "gpu_shader_create_info.hh"
+
+GPU_SHADER_CREATE_INFO(compositor_blur_variable_size)
+    .local_group_size(16, 16)
+    .push_constant(Type::FLOAT, "base_size")
+    .push_constant(Type::INT, "search_radius")
+    .sampler(0, ImageType::FLOAT_2D, "input_tx")
+    .sampler(1, ImageType::FLOAT_2D, "weights_tx")
+    .sampler(2, ImageType::FLOAT_2D, "size_tx")
+    .sampler(3, ImageType::FLOAT_2D, "mask_tx")
+    .image(0, GPU_RGBA16F, Qualifier::WRITE, ImageType::FLOAT_2D, "output_img")
+    .compute_source("compositor_blur_variable_size.glsl")
+    .do_static_compilation(true);