Cycles: support for custom shader AOVs

Custom render passes are added in the Shader AOVs panel in the view layer
settings, with a name and data type. In shader nodes, an AOV Output node
is then used to output either a value or color to the pass.

Arbitrary names can be used for these passes, as long as they don't conflict
with built-in passes that are enabled. The AOV Output node can be used in both
material and world shader nodes.

Implemented by Lukas, with tweaks by Brecht.

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

11 files changed, 214 insertions, 35 deletions

diff --git a/intern/cycles/render/buffers.cpp b/intern/cycles/render/buffers.cpp
index 49e15d9eaf1..fe8606e1939 100644
--- a/intern/cycles/render/buffers.cpp
+++ b/intern/cycles/render/buffers.cpp
@@ -234,7 +234,7 @@ bool RenderBuffers::get_denoising_pass_rect(
 }
 
 bool RenderBuffers::get_pass_rect(
-    PassType type, float exposure, int sample, int components, float *pixels, const string &name)
+    const string &name, float exposure, int sample, int components, float *pixels)
 {
   if (buffer.data() == NULL) {
     return false;
@@ -245,18 +245,14 @@ bool RenderBuffers::get_pass_rect(
   for (size_t j = 0; j < params.passes.size(); j++) {
     Pass &pass = params.passes[j];
 
-    if (pass.type != type) {
+    /* Pass is identified by both type and name, multiple of the same type
+     * may exist with a different name. */
+    if (pass.name != name) {
       pass_offset += pass.components;
       continue;
     }
 
-    /* Tell Cryptomatte passes apart by their name. */
-    if (pass.type == PASS_CRYPTOMATTE) {
-      if (pass.name != name) {
-        pass_offset += pass.components;
-        continue;
-      }
-    }
+    PassType type = pass.type;
 
     float *in = buffer.data() + pass_offset;
     int pass_stride = params.get_passes_size();
diff --git a/intern/cycles/render/buffers.h b/intern/cycles/render/buffers.h
index 74bb656c5a4..1042b42810f 100644
--- a/intern/cycles/render/buffers.h
+++ b/intern/cycles/render/buffers.h
@@ -88,12 +88,8 @@ class RenderBuffers {
   void zero();
 
   bool copy_from_device();
-  bool get_pass_rect(PassType type,
-                     float exposure,
-                     int sample,
-                     int components,
-                     float *pixels,
-                     const string &name);
+  bool get_pass_rect(
+      const string &name, float exposure, int sample, int components, float *pixels);
   bool get_denoising_pass_rect(
       int offset, float exposure, int sample, int components, float *pixels);
 };
diff --git a/intern/cycles/render/film.cpp b/intern/cycles/render/film.cpp
index 7f5bec2a66e..379b0e6c214 100644
--- a/intern/cycles/render/film.cpp
+++ b/intern/cycles/render/film.cpp
@@ -163,6 +163,12 @@ void Pass::add(PassType type, vector<Pass> &passes, const char *name)
     case PASS_CRYPTOMATTE:
       pass.components = 4;
       break;
+    case PASS_AOV_COLOR:
+      pass.components = 4;
+      break;
+    case PASS_AOV_VALUE:
+      pass.components = 1;
+      break;
     default:
       assert(false);
       break;
@@ -327,7 +333,7 @@ void Film::device_update(Device *device, DeviceScene *dscene, Scene *scene)
   kfilm->pass_stride = 0;
   kfilm->use_light_pass = use_light_visibility || use_sample_clamp;
 
-  bool have_cryptomatte = false;
+  bool have_cryptomatte = false, have_aov_color = false, have_aov_value = false;
 
   for (size_t i = 0; i < passes.size(); i++) {
     Pass &pass = passes[i];
@@ -464,6 +470,18 @@ void Film::device_update(Device *device, DeviceScene *dscene, Scene *scene)
                                       kfilm->pass_stride;
         have_cryptomatte = true;
         break;
+      case PASS_AOV_COLOR:
+        if (!have_aov_color) {
+          kfilm->pass_aov_color = kfilm->pass_stride;
+          have_aov_color = true;
+        }
+        break;
+      case PASS_AOV_VALUE:
+        if (!have_aov_value) {
+          kfilm->pass_aov_value = kfilm->pass_stride;
+          have_aov_value = true;
+        }
+        break;
       default:
         assert(false);
         break;
@@ -569,4 +587,27 @@ void Film::tag_update(Scene * /*scene*/)
   need_update = true;
 }
 
+int Film::get_aov_offset(string name, bool &is_color)
+{
+  int num_color = 0, num_value = 0;
+  foreach (const Pass &pass, passes) {
+    if (pass.type == PASS_AOV_COLOR) {
+      num_color++;
+    }
+    else if (pass.type == PASS_AOV_VALUE) {
+      num_value++;
+    }
+    else {
+      continue;
+    }
+
+    if (pass.name == name) {
+      is_color = (pass.type == PASS_AOV_COLOR);
+      return (is_color ? num_color : num_value) - 1;
+    }
+  }
+
+  return -1;
+}
+
 CCL_NAMESPACE_END
diff --git a/intern/cycles/render/film.h b/intern/cycles/render/film.h
index 2c4e07d60ae..48673af0114 100644
--- a/intern/cycles/render/film.h
+++ b/intern/cycles/render/film.h
@@ -93,6 +93,8 @@ class Film : public Node {
   bool modified(const Film &film);
   void tag_passes_update(Scene *scene, const vector<Pass> &passes_, bool update_passes = true);
   void tag_update(Scene *scene);
+
+  int get_aov_offset(string name, bool &is_color);
 };
 
 CCL_NAMESPACE_END
diff --git a/intern/cycles/render/graph.cpp b/intern/cycles/render/graph.cpp
index 501d1e33a9b..0e520c700a7 100644
--- a/intern/cycles/render/graph.cpp
+++ b/intern/cycles/render/graph.cpp
@@ -55,6 +55,25 @@ bool check_node_inputs_traversed(const ShaderNode *node, const ShaderNodeSet &do
 
 } /* namespace */
 
+/* Sockets */
+
+void ShaderInput::disconnect()
+{
+  if (link) {
+    link->links.erase(remove(link->links.begin(), link->links.end(), this), link->links.end());
+  }
+  link = NULL;
+}
+
+void ShaderOutput::disconnect()
+{
+  foreach (ShaderInput *sock, links) {
+    sock->link = NULL;
+  }
+
+  links.clear();
+}
+
 /* Node */
 
 ShaderNode::ShaderNode(const NodeType *type) : Node(type)
@@ -285,11 +304,7 @@ void ShaderGraph::disconnect(ShaderOutput *from)
   assert(!finalized);
   simplified = false;
 
-  foreach (ShaderInput *sock, from->links) {
-    sock->link = NULL;
-  }
-
-  from->links.clear();
+  from->disconnect();
 }
 
 void ShaderGraph::disconnect(ShaderInput *to)
@@ -298,10 +313,7 @@ void ShaderGraph::disconnect(ShaderInput *to)
   assert(to->link);
   simplified = false;
 
-  ShaderOutput *from = to->link;
-
-  to->link = NULL;
-  from->links.erase(remove(from->links.begin(), from->links.end(), to), from->links.end());
+  to->disconnect();
 }
 
 void ShaderGraph::relink(ShaderInput *from, ShaderInput *to)
@@ -782,6 +794,11 @@ void ShaderGraph::clean(Scene *scene)
 
   /* break cycles */
   break_cycles(output(), visited, on_stack);
+  foreach (ShaderNode *node, nodes) {
+    if (node->special_type == SHADER_SPECIAL_TYPE_OUTPUT_AOV) {
+      break_cycles(node, visited, on_stack);
+    }
+  }
 
   /* disconnect unused nodes */
   foreach (ShaderNode *node, nodes) {
diff --git a/intern/cycles/render/graph.h b/intern/cycles/render/graph.h
index cade04de374..0ea7935f714 100644
--- a/intern/cycles/render/graph.h
+++ b/intern/cycles/render/graph.h
@@ -67,6 +67,7 @@ enum ShaderNodeSpecialType {
   SHADER_SPECIAL_TYPE_COMBINE_CLOSURE,
   SHADER_SPECIAL_TYPE_OUTPUT,
   SHADER_SPECIAL_TYPE_BUMP,
+  SHADER_SPECIAL_TYPE_OUTPUT_AOV,
 };
 
 /* Input
@@ -104,6 +105,8 @@ class ShaderInput {
     ((Node *)parent)->set(socket_type, f);
   }
 
+  void disconnect();
+
   const SocketType &socket_type;
   ShaderNode *parent;
   ShaderOutput *link;
@@ -130,6 +133,8 @@ class ShaderOutput {
     return socket_type.type;
   }
 
+  void disconnect();
+
   const SocketType &socket_type;
   ShaderNode *parent;
   vector<ShaderInput *> links;
diff --git a/intern/cycles/render/nodes.cpp b/intern/cycles/render/nodes.cpp
index 26f16d5ee80..5e12d79bc6b 100644
--- a/intern/cycles/render/nodes.cpp
+++ b/intern/cycles/render/nodes.cpp
@@ -5709,6 +5709,58 @@ void ClampNode::compile(OSLCompiler &compiler)
   compiler.add(this, "node_clamp");
 }
 
+/* AOV Output */
+
+NODE_DEFINE(OutputAOVNode)
+{
+  NodeType *type = NodeType::add("aov_output", create, NodeType::SHADER);
+
+  SOCKET_IN_COLOR(color, "Color", make_float3(0.0f, 0.0f, 0.0f));
+  SOCKET_IN_FLOAT(value, "Value", 0.0f);
+
+  SOCKET_STRING(name, "AOV Name", ustring(""));
+
+  return type;
+}
+
+OutputAOVNode::OutputAOVNode() : ShaderNode(node_type)
+{
+  special_type = SHADER_SPECIAL_TYPE_OUTPUT_AOV;
+  slot = -1;
+}
+
+void OutputAOVNode::simplify_settings(Scene *scene)
+{
+  slot = scene->film->get_aov_offset(name.string(), is_color);
+  if (slot == -1) {
+    slot = scene->film->get_aov_offset(name.string(), is_color);
+  }
+
+  if (slot == -1 || is_color) {
+    input("Value")->disconnect();
+  }
+  if (slot == -1 || !is_color) {
+    input("Color")->disconnect();
+  }
+}
+
+void OutputAOVNode::compile(SVMCompiler &compiler)
+{
+  assert(slot >= 0);
+
+  if (is_color) {
+    compiler.add_node(NODE_AOV_COLOR, compiler.stack_assign(input("Color")), slot);
+  }
+  else {
+    compiler.add_node(NODE_AOV_VALUE, compiler.stack_assign(input("Value")), slot);
+  }
+}
+
+void OutputAOVNode::compile(OSLCompiler & /*compiler*/)
+{
+  /* TODO */
+}
+
 /* Math */
 
 NODE_DEFINE(MathNode)
diff --git a/intern/cycles/render/nodes.h b/intern/cycles/render/nodes.h
index ae762f3db1c..54124cd2175 100644
--- a/intern/cycles/render/nodes.h
+++ b/intern/cycles/render/nodes.h
@@ -189,6 +189,26 @@ class OutputNode : public ShaderNode {
   }
 };
 
+class OutputAOVNode : public ShaderNode {
+ public:
+  SHADER_NODE_CLASS(OutputAOVNode)
+  virtual void simplify_settings(Scene *scene);
+
+  float value;
+  float3 color;
+
+  ustring name;
+
+  /* Don't allow output node de-duplication. */
+  virtual bool equals(const ShaderNode & /*other*/)
+  {
+    return false;
+  }
+
+  int slot;
+  bool is_color;
+};
+
 class GradientTextureNode : public TextureNode {
  public:
   SHADER_NODE_CLASS(GradientTextureNode)
diff --git a/intern/cycles/render/shader.cpp b/intern/cycles/render/shader.cpp
index 00d9cd5e672..3b73fa4139f 100644
--- a/intern/cycles/render/shader.cpp
+++ b/intern/cycles/render/shader.cpp
@@ -225,6 +225,13 @@ Shader::~Shader()
 
 bool Shader::is_constant_emission(float3 *emission)
 {
+  /* If the shader has AOVs, they need to be evaluated, so we can't skip the shader. */
+  foreach (ShaderNode *node, graph->nodes) {
+    if (node->special_type == SHADER_SPECIAL_TYPE_OUTPUT_AOV) {
+      return false;
+    }
+  }
+
   ShaderInput *surf = graph->output()->input("Surface");
 
   if (surf->link == NULL) {
diff --git a/intern/cycles/render/svm.cpp b/intern/cycles/render/svm.cpp
index f69c89b1fd4..8466742ae38 100644
--- a/intern/cycles/render/svm.cpp
+++ b/intern/cycles/render/svm.cpp
@@ -554,6 +554,24 @@ void SVMCompiler::generated_shared_closure_nodes(ShaderNode *root_node,
   }
 }
 
+void SVMCompiler::generate_aov_node(ShaderNode *node, CompilerState *state)
+{
+  /* execute dependencies for node */
+  foreach (ShaderInput *in, node->inputs) {
+    if (in->link != NULL) {
+      ShaderNodeSet dependencies;
+      find_dependencies(dependencies, state->nodes_done, in);
+      generate_svm_nodes(dependencies, state);
+    }
+  }
+
+  /* compile node itself */
+  generate_node(node, state->nodes_done);
+
+  state->nodes_done.insert(node);
+  state->nodes_done_flag[node->id] = true;
+}
+
 void SVMCompiler::generate_multi_closure(ShaderNode *root_node,
                                          ShaderNode *node,
                                          CompilerState *state)
@@ -703,21 +721,21 @@ void SVMCompiler::compile_type(Shader *shader, ShaderGraph *graph, ShaderType ty
   current_graph = graph;
 
   /* get input in output node */
-  ShaderNode *node = graph->output();
+  ShaderNode *output = graph->output();
   ShaderInput *clin = NULL;
 
   switch (type) {
     case SHADER_TYPE_SURFACE:
-      clin = node->input("Surface");
+      clin = output->input("Surface");
       break;
     case SHADER_TYPE_VOLUME:
-      clin = node->input("Volume");
+      clin = output->input("Volume");
       break;
     case SHADER_TYPE_DISPLACEMENT:
-      clin = node->input("Displacement");
+      clin = output->input("Displacement");
       break;
     case SHADER_TYPE_BUMP:
-      clin = node->input("Normal");
+      clin = output->input("Normal");
       break;
     default:
       assert(0);
@@ -728,10 +746,10 @@ void SVMCompiler::compile_type(Shader *shader, ShaderGraph *graph, ShaderType ty
   memset((void *)&active_stack, 0, sizeof(active_stack));
   current_svm_nodes.clear();
 
-  foreach (ShaderNode *node_iter, graph->nodes) {
-    foreach (ShaderInput *input, node_iter->inputs)
+  foreach (ShaderNode *node, graph->nodes) {
+    foreach (ShaderInput *input, node->inputs)
       input->stack_offset = SVM_STACK_INVALID;
-    foreach (ShaderOutput *output, node_iter->outputs)
+    foreach (ShaderOutput *output, node->outputs)
       output->stack_offset = SVM_STACK_INVALID;
   }
 
@@ -745,6 +763,7 @@ void SVMCompiler::compile_type(Shader *shader, ShaderGraph *graph, ShaderType ty
   }
 
   if (shader->used) {
+    CompilerState state(graph);
     if (clin->link) {
       bool generate = false;
 
@@ -769,13 +788,36 @@ void SVMCompiler::compile_type(Shader *shader, ShaderGraph *graph, ShaderType ty
       }
 
       if (generate) {
-        CompilerState state(graph);
         generate_multi_closure(clin->link->parent, clin->link->parent, &state);
       }
     }
 
     /* compile output node */
-    node->compile(*this);
+    output->compile(*this);
+
+    if (type == SHADER_TYPE_SURFACE) {
+      vector<OutputAOVNode *> aov_outputs;
+      foreach (ShaderNode *node, graph->nodes) {
+        if (node->special_type == SHADER_SPECIAL_TYPE_OUTPUT_AOV) {
+          OutputAOVNode *aov_node = static_cast<OutputAOVNode *>(node);
+          if (aov_node->slot >= 0) {
+            aov_outputs.push_back(aov_node);
+          }
+        }
+      }
+      if (aov_outputs.size() > 0) {
+        /* AOV passes are only written if the object is directly visible, so
+         * there is no point in evaluating all the nodes generated only for the
+         * AOV outputs if that's not the case. Therefore, we insert
+         * NODE_AOV_START into the shader before the AOV-only nodes are
+         * generated which tells the kernel that it can stop evaluation
+         * early if AOVs will not be written. */
+        add_node(NODE_AOV_START, 0, 0, 0);
+        foreach (OutputAOVNode *node, aov_outputs) {
+          generate_aov_node(node, &state);
+        }
+      }
+    }
   }
 
   /* add node to restore state after bump shader has finished */
diff --git a/intern/cycles/render/svm.h b/intern/cycles/render/svm.h
index 7eac3e19bb3..d1534567bea 100644
--- a/intern/cycles/render/svm.h
+++ b/intern/cycles/render/svm.h
@@ -204,6 +204,7 @@ class SVMCompiler {
                          ShaderInput *input,
                          ShaderNode *skip_node = NULL);
   void generate_node(ShaderNode *node, ShaderNodeSet &done);
+  void generate_aov_node(ShaderNode *node, CompilerState *state);
   void generate_closure_node(ShaderNode *node, CompilerState *state);
   void generated_shared_closure_nodes(ShaderNode *root_node,
                                       ShaderNode *node,