Realtime Compositor: Add needed GPU module changes

This patch implements the necessary changes to the GPU module that are
needed by the realtime compositor.

A new function GPU_material_from_callbacks was added to construct a GPU
material from a number of callbacks. A callback to construct the
material graph by adding and linking the necessary GPU material nodes.
And the existing code generator callback. This essentially allows the
construction of GPU materials independent of node trees and without the
need to do any node tree localization.

A new composite source output to the code generator was added. This
output contains the serialization of nodes that are tagged with
GPU_NODE_TAG_COMPOSITOR, which are the nodes linked to the newly added
composite output links.

Two new GPU uniform setters were added for int2 and matrix3 types.

Shader create info now supports generated compute sources.

Shaders starting with gpu_shader_compositor are now considered part of
the shader library.

Additionally, two fixes were implemented. First, GPU setter node
de-duplication now appropriately increments the reference count of the
references resources. Second, unlinked sockets now get their value from
their associated GPU node stack instead of the socket itself.

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

Reviewed By: Clement

9 files changed, 142 insertions, 33 deletions

diff --git a/source/blender/gpu/GPU_material.h b/source/blender/gpu/GPU_material.h
index 3ca465fa57a..12a7cf49f9d 100644
--- a/source/blender/gpu/GPU_material.h
+++ b/source/blender/gpu/GPU_material.h
@@ -121,6 +121,7 @@ typedef struct GPUCodegenOutput {
   char *surface;
   char *volume;
   char *thickness;
+  char *composite;
   char *material_functions;
 
   GPUShaderCreateInfo *create_info;
@@ -178,6 +179,8 @@ void GPU_material_output_thickness(GPUMaterial *material, GPUNodeLink *link);
 
 void GPU_material_add_output_link_aov(GPUMaterial *material, GPUNodeLink *link, int hash);
 
+void GPU_material_add_output_link_composite(GPUMaterial *material, GPUNodeLink *link);
+
 /**
  * Wrap a part of the material graph into a function. You need then need to call the function by
  * using something like #GPU_differentiate_float_function.
@@ -218,6 +221,7 @@ GPUMaterial *GPU_material_from_nodetree(struct Scene *scene,
                                         void *thunk);
 
 void GPU_material_compile(GPUMaterial *mat);
+void GPU_material_free_single(GPUMaterial *material);
 void GPU_material_free(struct ListBase *gpumaterial);
 
 void GPU_material_acquire(GPUMaterial *mat);
@@ -319,6 +323,16 @@ struct GHash *GPU_uniform_attr_list_hash_new(const char *info);
 void GPU_uniform_attr_list_copy(GPUUniformAttrList *dest, GPUUniformAttrList *src);
 void GPU_uniform_attr_list_free(GPUUniformAttrList *set);
 
+/* A callback passed to GPU_material_from_callbacks to construct the material graph by adding and
+ * linking the necessary GPU material nodes. */
+typedef void (*ConstructGPUMaterialFn)(void *thunk, GPUMaterial *material);
+
+/* Construct a GPU material from a set of callbacks. See the callback types for more information.
+ * The given thunk will be passed as the first parameter of each callback. */
+GPUMaterial *GPU_material_from_callbacks(ConstructGPUMaterialFn construct_function_cb,
+                                         GPUCodegenCallbackFn generate_code_function_cb,
+                                         void *thunk);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/source/blender/gpu/GPU_shader.h b/source/blender/gpu/GPU_shader.h
index 3460d33fe68..c154f1adc8b 100644
--- a/source/blender/gpu/GPU_shader.h
+++ b/source/blender/gpu/GPU_shader.h
@@ -177,7 +177,9 @@ void GPU_shader_uniform_4f(GPUShader *sh, const char *name, float x, float y, fl
 void GPU_shader_uniform_2fv(GPUShader *sh, const char *name, const float data[2]);
 void GPU_shader_uniform_3fv(GPUShader *sh, const char *name, const float data[3]);
 void GPU_shader_uniform_4fv(GPUShader *sh, const char *name, const float data[4]);
+void GPU_shader_uniform_2iv(GPUShader *sh, const char *name, const int data[2]);
 void GPU_shader_uniform_mat4(GPUShader *sh, const char *name, const float data[4][4]);
+void GPU_shader_uniform_mat3_as_mat4(GPUShader *sh, const char *name, const float data[3][3]);
 void GPU_shader_uniform_2fv_array(GPUShader *sh, const char *name, int len, const float (*val)[2]);
 void GPU_shader_uniform_4fv_array(GPUShader *sh, const char *name, int len, const float (*val)[4]);
 
diff --git a/source/blender/gpu/intern/gpu_codegen.cc b/source/blender/gpu/intern/gpu_codegen.cc
index 82441c3c89c..4a45a3e63ed 100644
--- a/source/blender/gpu/intern/gpu_codegen.cc
+++ b/source/blender/gpu/intern/gpu_codegen.cc
@@ -280,6 +280,7 @@ class GPUCodegen {
 
   void node_serialize(std::stringstream &eval_ss, const GPUNode *node);
   char *graph_serialize(eGPUNodeTag tree_tag, GPUNodeLink *output_link);
+  char *graph_serialize(eGPUNodeTag tree_tag);
 
   static char *extract_c_str(std::stringstream &stream)
   {
@@ -500,6 +501,19 @@ char *GPUCodegen::graph_serialize(eGPUNodeTag tree_tag, GPUNodeLink *output_link
   return eval_c_str;
 }
 
+char *GPUCodegen::graph_serialize(eGPUNodeTag tree_tag)
+{
+  std::stringstream eval_ss;
+  LISTBASE_FOREACH (GPUNode *, node, &graph.nodes) {
+    if (node->tag & tree_tag) {
+      node_serialize(eval_ss, node);
+    }
+  }
+  char *eval_c_str = extract_c_str(eval_ss);
+  BLI_hash_mm2a_add(&hm2a_, (uchar *)eval_c_str, eval_ss.str().size());
+  return eval_c_str;
+}
+
 void GPUCodegen::generate_uniform_buffer()
 {
   /* Extract uniform inputs. */
@@ -539,6 +553,9 @@ void GPUCodegen::generate_graphs()
   output.volume = graph_serialize(GPU_NODE_TAG_VOLUME, graph.outlink_volume);
   output.displacement = graph_serialize(GPU_NODE_TAG_DISPLACEMENT, graph.outlink_displacement);
   output.thickness = graph_serialize(GPU_NODE_TAG_THICKNESS, graph.outlink_thickness);
+  if (!BLI_listbase_is_empty(&graph.outlink_compositor)) {
+    output.composite = graph_serialize(GPU_NODE_TAG_COMPOSITOR);
+  }
 
   if (!BLI_listbase_is_empty(&graph.material_functions)) {
     std::stringstream eval_ss;
@@ -569,9 +586,10 @@ GPUPass *GPU_generate_pass(GPUMaterial *material,
                            GPUCodegenCallbackFn finalize_source_cb,
                            void *thunk)
 {
-  /* Prune the unused nodes and extract attributes before compiling so the
-   * generated VBOs are ready to accept the future shader. */
   gpu_node_graph_prune_unused(graph);
+
+  /* Extract attributes before compiling so the generated VBOs are ready to accept the future
+   * shader. */
   gpu_node_graph_finalize_uniform_attrs(graph);
 
   GPUCodegen codegen(material, graph);
diff --git a/source/blender/gpu/intern/gpu_material.c b/source/blender/gpu/intern/gpu_material.c
index 4d3ea3e0c99..a4842ef0e43 100644
--- a/source/blender/gpu/intern/gpu_material.c
+++ b/source/blender/gpu/intern/gpu_material.c
@@ -141,7 +141,7 @@ static void gpu_material_ramp_texture_build(GPUMaterial *mat)
   mat->coba_builder = NULL;
 }
 
-static void gpu_material_free_single(GPUMaterial *material)
+void GPU_material_free_single(GPUMaterial *material)
 {
   bool do_free = atomic_sub_and_fetch_uint32(&material->refcount, 1) == 0;
   if (!do_free) {
@@ -173,7 +173,7 @@ void GPU_material_free(ListBase *gpumaterial)
   LISTBASE_FOREACH (LinkData *, link, gpumaterial) {
     GPUMaterial *material = link->data;
     DRW_deferred_shader_remove(material);
-    gpu_material_free_single(material);
+    GPU_material_free_single(material);
   }
   BLI_freelistN(gpumaterial);
 }
@@ -538,6 +538,13 @@ void GPU_material_add_output_link_aov(GPUMaterial *material, GPUNodeLink *link,
   BLI_addtail(&material->graph.outlink_aovs, aov_link);
 }
 
+void GPU_material_add_output_link_composite(GPUMaterial *material, GPUNodeLink *link)
+{
+  GPUNodeGraphOutputLink *compositor_link = MEM_callocN(sizeof(GPUNodeGraphOutputLink), __func__);
+  compositor_link->outlink = link;
+  BLI_addtail(&material->graph.outlink_compositor, compositor_link);
+}
+
 char *GPU_material_split_sub_function(GPUMaterial *material,
                                       eGPUType return_type,
                                       GPUNodeLink **link)
@@ -721,7 +728,7 @@ void GPU_material_acquire(GPUMaterial *mat)
 
 void GPU_material_release(GPUMaterial *mat)
 {
-  gpu_material_free_single(mat);
+  GPU_material_free_single(mat);
 }
 
 void GPU_material_compile(GPUMaterial *mat)
@@ -772,3 +779,42 @@ void GPU_materials_free(Main *bmain)
   // BKE_world_defaults_free_gpu();
   BKE_material_defaults_free_gpu();
 }
+
+GPUMaterial *GPU_material_from_callbacks(ConstructGPUMaterialFn construct_function_cb,
+                                         GPUCodegenCallbackFn generate_code_function_cb,
+                                         void *thunk)
+{
+  /* Allocate a new material and its material graph, and initialize its reference count. */
+  GPUMaterial *material = MEM_callocN(sizeof(GPUMaterial), "GPUMaterial");
+  material->graph.used_libraries = BLI_gset_new(
+      BLI_ghashutil_ptrhash, BLI_ghashutil_ptrcmp, "GPUNodeGraph.used_libraries");
+  material->refcount = 1;
+
+  /* Construct the material graph by adding and linking the necessary GPU material nodes. */
+  construct_function_cb(thunk, material);
+
+  /* Create and initialize the texture storing color bands used by Ramp and Curve nodes. */
+  gpu_material_ramp_texture_build(material);
+
+  /* Lookup an existing pass in the cache or generate a new one. */
+  material->pass = GPU_generate_pass(material, &material->graph, generate_code_function_cb, thunk);
+
+  /* The pass already exists in the pass cache but its shader already failed to compile. */
+  if (material->pass == NULL) {
+    material->status = GPU_MAT_FAILED;
+    gpu_node_graph_free(&material->graph);
+    return material;
+  }
+
+  /* The pass already exists in the pass cache and its shader is already compiled. */
+  GPUShader *shader = GPU_pass_shader_get(material->pass);
+  if (shader != NULL) {
+    material->status = GPU_MAT_SUCCESS;
+    gpu_node_graph_free_nodes(&material->graph);
+    return material;
+  }
+
+  /* The material was created successfully but still needs to be compiled. */
+  material->status = GPU_MAT_CREATED;
+  return material;
+}
diff --git a/source/blender/gpu/intern/gpu_node_graph.c b/source/blender/gpu/intern/gpu_node_graph.c
index 684070dbdc0..c7b2fde492f 100644
--- a/source/blender/gpu/intern/gpu_node_graph.c
+++ b/source/blender/gpu/intern/gpu_node_graph.c
@@ -75,9 +75,26 @@ static void gpu_node_input_link(GPUNode *node, GPUNodeLink *link, const eGPUType
 
     if (STR_ELEM(name, "set_value", "set_rgb", "set_rgba") && (input->type == type)) {
       input = MEM_dupallocN(outnode->inputs.first);
+
+      switch (input->source) {
+        case GPU_SOURCE_ATTR:
+          input->attr->users++;
+          break;
+        case GPU_SOURCE_UNIFORM_ATTR:
+          input->uniform_attr->users++;
+          break;
+        case GPU_SOURCE_TEX:
+        case GPU_SOURCE_TEX_TILED_MAPPING:
+          input->texture->users++;
+          break;
+        default:
+          break;
+      }
+
       if (input->link) {
         input->link->users++;
       }
+
       BLI_addtail(&node->inputs, input);
       return;
     }
@@ -179,35 +196,21 @@ static GPUNodeLink *gpu_uniformbuffer_link(GPUMaterial *mat,
   BLI_assert(socket != NULL);
   BLI_assert(socket->in_out == in_out);
 
-  if ((socket->flag & SOCK_HIDE_VALUE) == 0) {
-    GPUNodeLink *link;
-    switch (socket->type) {
-      case SOCK_FLOAT: {
-        bNodeSocketValueFloat *socket_data = socket->default_value;
-        link = GPU_uniform(&socket_data->value);
-        break;
-      }
-      case SOCK_VECTOR: {
-        bNodeSocketValueVector *socket_data = socket->default_value;
-        link = GPU_uniform(socket_data->value);
-        break;
-      }
-      case SOCK_RGBA: {
-        bNodeSocketValueRGBA *socket_data = socket->default_value;
-        link = GPU_uniform(socket_data->value);
-        break;
-      }
-      default:
-        return NULL;
-        break;
-    }
+  if (socket->flag & SOCK_HIDE_VALUE) {
+    return NULL;
+  }
 
-    if (in_out == SOCK_IN) {
-      GPU_link(mat, gpu_uniform_set_function_from_type(socket->type), link, &stack->link);
-    }
-    return link;
+  if (!ELEM(socket->type, SOCK_FLOAT, SOCK_VECTOR, SOCK_RGBA)) {
+    return NULL;
+  }
+
+  GPUNodeLink *link = GPU_uniform(stack->vec);
+
+  if (in_out == SOCK_IN) {
+    GPU_link(mat, gpu_uniform_set_function_from_type(socket->type), link, &stack->link);
   }
-  return NULL;
+
+  return link;
 }
 
 static void gpu_node_input_socket(
@@ -803,6 +806,7 @@ void gpu_node_graph_free(GPUNodeGraph *graph)
 {
   BLI_freelistN(&graph->outlink_aovs);
   BLI_freelistN(&graph->material_functions);
+  BLI_freelistN(&graph->outlink_compositor);
   gpu_node_graph_free_nodes(graph);
 
   BLI_freelistN(&graph->textures);
@@ -855,6 +859,9 @@ void gpu_node_graph_prune_unused(GPUNodeGraph *graph)
   LISTBASE_FOREACH (GPUNodeGraphFunctionLink *, funclink, &graph->material_functions) {
     gpu_nodes_tag(funclink->outlink, GPU_NODE_TAG_FUNCTION);
   }
+  LISTBASE_FOREACH (GPUNodeGraphOutputLink *, compositor_link, &graph->outlink_compositor) {
+    gpu_nodes_tag(compositor_link->outlink, GPU_NODE_TAG_COMPOSITOR);
+  }
 
   for (GPUNode *node = graph->nodes.first, *next = NULL; node; node = next) {
     next = node->next;
diff --git a/source/blender/gpu/intern/gpu_node_graph.h b/source/blender/gpu/intern/gpu_node_graph.h
index ae472d5b7aa..08ff8bbef58 100644
--- a/source/blender/gpu/intern/gpu_node_graph.h
+++ b/source/blender/gpu/intern/gpu_node_graph.h
@@ -59,6 +59,7 @@ typedef enum {
   GPU_NODE_TAG_THICKNESS = (1 << 3),
   GPU_NODE_TAG_AOV = (1 << 4),
   GPU_NODE_TAG_FUNCTION = (1 << 5),
+  GPU_NODE_TAG_COMPOSITOR = (1 << 6),
 } eGPUNodeTag;
 
 ENUM_OPERATORS(eGPUNodeTag, GPU_NODE_TAG_FUNCTION)
@@ -158,6 +159,8 @@ typedef struct GPUNodeGraph {
   ListBase outlink_aovs;
   /* List of GPUNodeGraphFunctionLink */
   ListBase material_functions;
+  /* List of GPUNodeGraphOutputLink */
+  ListBase outlink_compositor;
 
   /* Requested attributes and textures. */
   ListBase attributes;
diff --git a/source/blender/gpu/intern/gpu_shader.cc b/source/blender/gpu/intern/gpu_shader.cc
index fe9aacb95f9..8a630d9e3ca 100644
--- a/source/blender/gpu/intern/gpu_shader.cc
+++ b/source/blender/gpu/intern/gpu_shader.cc
@@ -7,6 +7,7 @@
 
 #include "MEM_guardedalloc.h"
 
+#include "BLI_math_matrix.h"
 #include "BLI_string_utils.h"
 
 #include "GPU_capabilities.h"
@@ -382,6 +383,8 @@ GPUShader *GPU_shader_create_from_info(const GPUShaderCreateInfo *_info)
     sources.append(resources.c_str());
     sources.append(layout.c_str());
     sources.extend(code);
+    sources.extend(info.dependencies_generated);
+    sources.append(info.compute_source_generated.c_str());
 
     shader->compute_shader_from_glsl(sources);
   }
@@ -702,12 +705,25 @@ void GPU_shader_uniform_4fv(GPUShader *sh, const char *name, const float data[4]
   GPU_shader_uniform_vector(sh, loc, 4, 1, data);
 }
 
+void GPU_shader_uniform_2iv(GPUShader *sh, const char *name, const int data[2])
+{
+  const int loc = GPU_shader_get_uniform(sh, name);
+  GPU_shader_uniform_vector_int(sh, loc, 2, 1, data);
+}
+
 void GPU_shader_uniform_mat4(GPUShader *sh, const char *name, const float data[4][4])
 {
   const int loc = GPU_shader_get_uniform(sh, name);
   GPU_shader_uniform_vector(sh, loc, 16, 1, (const float *)data);
 }
 
+void GPU_shader_uniform_mat3_as_mat4(GPUShader *sh, const char *name, const float data[3][3])
+{
+  float matrix[4][4];
+  copy_m4_m3(matrix, data);
+  GPU_shader_uniform_mat4(sh, name, matrix);
+}
+
 void GPU_shader_uniform_2fv_array(GPUShader *sh, const char *name, int len, const float (*val)[2])
 {
   const int loc = GPU_shader_get_uniform(sh, name);
diff --git a/source/blender/gpu/intern/gpu_shader_create_info.hh b/source/blender/gpu/intern/gpu_shader_create_info.hh
index 8e05412d0ee..fb8efbb209a 100644
--- a/source/blender/gpu/intern/gpu_shader_create_info.hh
+++ b/source/blender/gpu/intern/gpu_shader_create_info.hh
@@ -298,6 +298,7 @@ struct ShaderCreateInfo {
   /** Manually set generated code. */
   std::string vertex_source_generated = "";
   std::string fragment_source_generated = "";
+  std::string compute_source_generated = "";
   std::string geometry_source_generated = "";
   std::string typedef_source_generated = "";
   /** Manually set generated dependencies. */
@@ -818,6 +819,7 @@ struct ShaderCreateInfo {
     TEST_EQUAL(*this, b, builtins_);
     TEST_EQUAL(*this, b, vertex_source_generated);
     TEST_EQUAL(*this, b, fragment_source_generated);
+    TEST_EQUAL(*this, b, compute_source_generated);
     TEST_EQUAL(*this, b, typedef_source_generated);
     TEST_VECTOR_EQUAL(*this, b, vertex_inputs_);
     TEST_EQUAL(*this, b, geometry_layout_);
diff --git a/source/blender/gpu/intern/gpu_shader_dependency.cc b/source/blender/gpu/intern/gpu_shader_dependency.cc
index d91e15243f3..aff7df9ac33 100644
--- a/source/blender/gpu/intern/gpu_shader_dependency.cc
+++ b/source/blender/gpu/intern/gpu_shader_dependency.cc
@@ -593,7 +593,8 @@ struct GPUSource {
   bool is_from_material_library() const
   {
     return (filename.startswith("gpu_shader_material_") ||
-            filename.startswith("gpu_shader_common_")) &&
+            filename.startswith("gpu_shader_common_") ||
+            filename.startswith("gpu_shader_compositor_")) &&
            filename.endswith(".glsl");
   }
 };