EEVEE: Fix AOV implementation

13 files changed, 66 insertions, 35 deletions

diff --git a/source/blender/draw/engines/eevee/eevee_private.h b/source/blender/draw/engines/eevee/eevee_private.h
index 3b9d6630933..30768c98016 100644
--- a/source/blender/draw/engines/eevee/eevee_private.h
+++ b/source/blender/draw/engines/eevee/eevee_private.h
@@ -462,7 +462,7 @@ typedef struct EEVEE_RenderPassData {
   int renderPassSSSColor;
   int renderPassEnvironment;
   int renderPassAOV;
-  int renderPassAOVActive;
+  uint renderPassAOVActive;
   int _pad[3];
 } EEVEE_RenderPassData;
 
@@ -1050,7 +1050,7 @@ typedef struct EEVEE_PrivateData {
   /* Renderpasses */
   /* Bitmask containing the active render_passes */
   eViewLayerEEVEEPassType render_passes;
-  int aov_hash;
+  uint aov_hash;
   int num_aovs_used;
   struct CryptomatteSession *cryptomatte_session;
   bool cryptomatte_accurate_mode;
@@ -1518,7 +1518,7 @@ bool EEVEE_renderpasses_only_first_sample_pass_active(EEVEE_Data *vedata);
  * Calculate the hash for an AOV. The least significant bit is used to store the AOV
  * type the rest of the bits are used for the name hash.
  */
-int EEVEE_renderpasses_aov_hash(const ViewLayerAOV *aov);
+uint EEVEE_renderpasses_aov_hash(const ViewLayerAOV *aov);
 
 /* eevee_temporal_sampling.c */
 
diff --git a/source/blender/draw/engines/eevee/eevee_renderpasses.c b/source/blender/draw/engines/eevee/eevee_renderpasses.c
index 3814cf44fd6..237c830d547 100644
--- a/source/blender/draw/engines/eevee/eevee_renderpasses.c
+++ b/source/blender/draw/engines/eevee/eevee_renderpasses.c
@@ -60,9 +60,9 @@ bool EEVEE_renderpasses_only_first_sample_pass_active(EEVEE_Data *vedata)
   return (g_data->render_passes & ~EEVEE_RENDERPASSES_POST_PROCESS_ON_FIRST_SAMPLE) == 0;
 }
 
-int EEVEE_renderpasses_aov_hash(const ViewLayerAOV *aov)
+uint EEVEE_renderpasses_aov_hash(const ViewLayerAOV *aov)
 {
-  int hash = BLI_hash_string(aov->name) << 1;
+  uint hash = BLI_hash_string(aov->name) << 1u;
   SET_FLAG_FROM_TEST(hash, aov->type == AOV_TYPE_COLOR, EEVEE_AOV_HASH_COLOR_TYPE_MASK);
   return hash;
 }
diff --git a/source/blender/draw/engines/eevee/shaders/closure_eval_surface_lib.glsl b/source/blender/draw/engines/eevee/shaders/closure_eval_surface_lib.glsl
index 4da70eb3f87..784a0adc1ed 100644
--- a/source/blender/draw/engines/eevee/shaders/closure_eval_surface_lib.glsl
+++ b/source/blender/draw/engines/eevee/shaders/closure_eval_surface_lib.glsl
@@ -21,6 +21,9 @@ float out_ssr_roughness;
 vec3 out_ssr_color;
 vec3 out_ssr_N;
 
+bool aov_is_valid = false;
+vec3 out_aov;
+
 bool output_sss(ClosureDiffuse diffuse, ClosureOutputDiffuse diffuse_out)
 {
   if (diffuse.sss_id == 0u || !do_sss || !sssToggle || outputSssId == 0) {
@@ -48,6 +51,22 @@ bool output_ssr(ClosureReflection reflection)
   return true;
 }
 
+void output_aov(vec4 color, float value, uint hash)
+{
+  /* Keep in sync with `render_pass_aov_hash` and `EEVEE_renderpasses_aov_hash`. */
+  hash <<= 1u;
+
+  if (renderPassAOV && !aov_is_valid && hash == render_pass_aov_hash()) {
+    aov_is_valid = true;
+    if (render_pass_aov_is_color()) {
+      out_aov = color.rgb;
+    }
+    else {
+      out_aov = vec3(value);
+    }
+  }
+}
+
 /* Single BSDFs. */
 CLOSURE_EVAL_FUNCTION_DECLARE_1(DiffuseBSDF, Diffuse)
 Closure closure_eval(ClosureDiffuse diffuse)
diff --git a/source/blender/draw/engines/eevee/shaders/closure_eval_volume_lib.glsl b/source/blender/draw/engines/eevee/shaders/closure_eval_volume_lib.glsl
index 9e11b7410be..b726c572b3c 100644
--- a/source/blender/draw/engines/eevee/shaders/closure_eval_volume_lib.glsl
+++ b/source/blender/draw/engines/eevee/shaders/closure_eval_volume_lib.glsl
@@ -1,14 +1,8 @@
 
-bool do_sss = true;
-bool do_ssr = true;
-
-vec3 out_sss_radiance;
-vec3 out_sss_color;
-float out_sss_radius;
-
-float out_ssr_roughness;
-vec3 out_ssr_color;
-vec3 out_ssr_N;
+void output_aov(vec4 color, float value, uint hash)
+{
+  /* Unsupported. */
+}
 
 /* Surface BSDFs. */
 Closure closure_eval(ClosureDiffuse diffuse)
diff --git a/source/blender/draw/engines/eevee/shaders/closure_type_lib.glsl b/source/blender/draw/engines/eevee/shaders/closure_type_lib.glsl
index c50d4234ec9..0f35acfcc86 100644
--- a/source/blender/draw/engines/eevee/shaders/closure_type_lib.glsl
+++ b/source/blender/draw/engines/eevee/shaders/closure_type_lib.glsl
@@ -29,6 +29,7 @@ struct Closure {
 /* Prototype */
 Closure nodetree_exec();
 vec4 closure_to_rgba(Closure);
+void output_aov(vec4 color, float value, uint hash);
 vec3 coordinate_camera(vec3 P);
 vec3 coordinate_screen(vec3 P);
 vec3 coordinate_reflect(vec3 P, vec3 N);
diff --git a/source/blender/draw/engines/eevee/shaders/renderpass_lib.glsl b/source/blender/draw/engines/eevee/shaders/renderpass_lib.glsl
index c7728cdf76f..f276e4f26ca 100644
--- a/source/blender/draw/engines/eevee/shaders/renderpass_lib.glsl
+++ b/source/blender/draw/engines/eevee/shaders/renderpass_lib.glsl
@@ -1,4 +1,4 @@
-#define EEVEE_AOV_HASH_COLOR_TYPE_MASK 1
+#define EEVEE_AOV_HASH_COLOR_TYPE_MASK 1u
 
 /* ---------------------------------------------------------------------- */
 /** \name Resources
@@ -14,7 +14,7 @@ layout(std140) uniform renderpass_block
   bool renderPassSSSColor;
   bool renderPassEnvironment;
   bool renderPassAOV;
-  int renderPassAOVActive;
+  uint renderPassAOVActive;
 };
 
 /** \} */
@@ -40,10 +40,10 @@ vec3 render_pass_emission_mask(vec3 emission_light)
 
 bool render_pass_aov_is_color()
 {
-  return (renderPassAOVActive & EEVEE_AOV_HASH_COLOR_TYPE_MASK) != 0;
+  return (renderPassAOVActive & EEVEE_AOV_HASH_COLOR_TYPE_MASK) != 0u;
 }
 
-int render_pass_aov_hash()
+uint render_pass_aov_hash()
 {
   return renderPassAOVActive & ~EEVEE_AOV_HASH_COLOR_TYPE_MASK;
 }
diff --git a/source/blender/draw/engines/eevee/shaders/surface_frag.glsl b/source/blender/draw/engines/eevee/shaders/surface_frag.glsl
index b2b1286df7b..e31302b1b2e 100644
--- a/source/blender/draw/engines/eevee/shaders/surface_frag.glsl
+++ b/source/blender/draw/engines/eevee/shaders/surface_frag.glsl
@@ -122,6 +122,15 @@ void main()
   outRadiance.rgb *= alpha_div;
   ssrData.rgb *= alpha_div;
   sssAlbedo.rgb *= alpha_div;
+
+  if (renderPassAOV) {
+    if (aov_is_valid) {
+      outRadiance = vec4(out_aov, 1.0);
+    }
+    else {
+      outRadiance = vec4(0.0);
+    }
+  }
 #endif
 
 #ifdef LOOKDEV
diff --git a/source/blender/gpu/GPU_material.h b/source/blender/gpu/GPU_material.h
index 7301b003e0e..f38b9681ad7 100644
--- a/source/blender/gpu/GPU_material.h
+++ b/source/blender/gpu/GPU_material.h
@@ -79,6 +79,7 @@ typedef enum eGPUMaterialFlag {
   GPU_MATFLAG_AO = (1 << 8),
 
   GPU_MATFLAG_OBJECT_INFO = (1 << 10),
+  GPU_MATFLAG_AOV = (1 << 11),
 
   GPU_MATFLAG_BARYCENTRIC = (1 << 20),
 
diff --git a/source/blender/gpu/intern/gpu_codegen.cc b/source/blender/gpu/intern/gpu_codegen.cc
index a53843842c3..8963fa45c96 100644
--- a/source/blender/gpu/intern/gpu_codegen.cc
+++ b/source/blender/gpu/intern/gpu_codegen.cc
@@ -565,7 +565,7 @@ void GPUCodegen::generate_graphs()
 {
   set_unique_ids();
 
-  output.surface = graph_serialize(GPU_NODE_TAG_SURFACE, graph.outlink_surface);
+  output.surface = graph_serialize(GPU_NODE_TAG_SURFACE | GPU_NODE_TAG_AOV, graph.outlink_surface);
   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);
diff --git a/source/blender/gpu/intern/gpu_node_graph.h b/source/blender/gpu/intern/gpu_node_graph.h
index 4bebb4f77cb..024119e1c24 100644
--- a/source/blender/gpu/intern/gpu_node_graph.h
+++ b/source/blender/gpu/intern/gpu_node_graph.h
@@ -65,6 +65,8 @@ typedef enum {
   GPU_NODE_TAG_FUNCTION = (1 << 5),
 } eGPUNodeTag;
 
+ENUM_OPERATORS(eGPUNodeTag, GPU_NODE_TAG_FUNCTION)
+
 struct GPUNode {
   struct GPUNode *next, *prev;
 
diff --git a/source/blender/gpu/shaders/material/gpu_shader_material_output_aov.glsl b/source/blender/gpu/shaders/material/gpu_shader_material_output_aov.glsl
index 648994739bf..b166c3e4e9f 100644
--- a/source/blender/gpu/shaders/material/gpu_shader_material_output_aov.glsl
+++ b/source/blender/gpu/shaders/material/gpu_shader_material_output_aov.glsl
@@ -1,13 +1,5 @@
 
-void node_output_aov(vec4 color, float value, out Closure result)
+void node_output_aov(vec4 color, float value, float hash, out Closure dummy)
 {
-  result = CLOSURE_DEFAULT;
-#ifndef VOLUMETRICS
-  if (render_pass_aov_is_color()) {
-    result.radiance = color.rgb;
-  }
-  else {
-    result.radiance = vec3(value);
-  }
-#endif
+  output_aov(color, value, floatBitsToUint(hash));
 }
diff --git a/source/blender/nodes/shader/node_shader_tree.cc b/source/blender/nodes/shader/node_shader_tree.cc
index a8ff7950d82..42a31f39041 100644
--- a/source/blender/nodes/shader/node_shader_tree.cc
+++ b/source/blender/nodes/shader/node_shader_tree.cc
@@ -1007,6 +1007,12 @@ static void ntree_shader_pruned_unused(bNodeTree *ntree, bNode *output_node)
 
   nodeChainIterBackwards(ntree, output_node, ntree_branch_node_tag, nullptr, 0);
 
+  LISTBASE_FOREACH (bNode *, node, &ntree->nodes) {
+    if (node->type == SH_NODE_OUTPUT_AOV) {
+      nodeChainIterBackwards(ntree, node, ntree_branch_node_tag, nullptr, 0);
+    }
+  }
+
   LISTBASE_FOREACH_MUTABLE (bNode *, node, &ntree->nodes) {
     if (node->tmp_flag == 0) {
       ntreeFreeLocalNode(ntree, node);
diff --git a/source/blender/nodes/shader/nodes/node_shader_output_aov.cc b/source/blender/nodes/shader/nodes/node_shader_output_aov.cc
index c88ec58d3dd..78dabc5c1c2 100644
--- a/source/blender/nodes/shader/nodes/node_shader_output_aov.cc
+++ b/source/blender/nodes/shader/nodes/node_shader_output_aov.cc
@@ -31,13 +31,20 @@ static int node_shader_gpu_output_aov(GPUMaterial *mat,
                                       bNode *node,
                                       bNodeExecData *UNUSED(execdata),
                                       GPUNodeStack *in,
-                                      GPUNodeStack *UNUSED(out))
+                                      GPUNodeStack *out)
 {
+  GPUNodeLink *outlink;
   NodeShaderOutputAOV *aov = (NodeShaderOutputAOV *)node->storage;
-  /* Keep in sync with `renderpass_lib.glsl#render_pass_aov_hash` and
-   * `EEVEE_renderpasses_aov_hash`. */
-  unsigned int hash = BLI_hash_string(aov->name) << 1;
-  GPU_material_add_output_link_aov(mat, in[0].link, hash);
+  uint hash = BLI_hash_string(aov->name);
+  /* WORKAROUND: We don't support int/uint constants for now. So make sure the aliasing works.
+   * We cast back to uint in GLSL. */
+  BLI_STATIC_ASSERT(sizeof(float) == sizeof(uint),
+                    "GPUCodegen: AOV hash needs float and uint to be the same size.");
+  GPUNodeLink *hash_link = GPU_constant((float *)&hash);
+
+  GPU_material_flag_set(mat, GPU_MATFLAG_AOV);
+  GPU_stack_link(mat, node, "node_output_aov", in, out, hash_link, &outlink);
+  GPU_material_add_output_link_aov(mat, outlink, hash);
   return true;
 }