add additional stage to optimize combination of 3 bsdfs that have been mixed with 2 nodes.

Author: ld-kerley

source/MaterialXGenShader/ShaderGraph.cpp

@@ -796,6 +796,22 @@ ShaderNode* ShaderGraph::getNode(const string& name)
     return it != _nodeMap.end() ? it->second.get() : nullptr;
 }
 
+bool ShaderGraph::removeNode(ShaderNode* node)
+{
+    auto mapIt = _nodeMap.find(node->getName());
+    auto vecIt = std::find(_nodeOrder.begin(), _nodeOrder.end(), node);
+    if (mapIt == _nodeMap.end() || vecIt == _nodeOrder.end())
+        return false;
+
+    for (ShaderInput* input : node->getInputs())
+    {
+        input->breakConnection();
+    }
+    _nodeMap.erase(mapIt);
+    _nodeOrder.erase(vecIt);
+    return true;
+}
+
 const ShaderNode* ShaderGraph::getNode(const string& name) const
 {
     return const_cast<ShaderGraph*>(this)->getNode(name);
@@ -976,9 +992,17 @@ void ShaderGraph::optimize(GenContext& context)
         const vector<ShaderNode*> nodeList = getNodes();
         for (ShaderNode* node : nodeList)
         {
+            // first check the node is still in the graph, and hasn't been removed by a
+            // prior optimization
+            if (!getNode(node->getName()))
+                continue;
+
             if (node->hasClassification(ShaderNode::Classification::MIX_BSDF))
             {
-                optimizeMixBsdf(node, context);
+                if (!optimizeMixMixBsdf(node, context))
+                {
+                    optimizeMixBsdf(node, context);
+                }
             }
         }
     }
@@ -1018,7 +1042,7 @@ void ShaderGraph::optimize(GenContext& context)
 //
 // Motivation - the new graph is more efficient for shader backends to optimize
 // away possible expensive BSDF nodes that might not be used at all.
-void ShaderGraph::optimizeMixBsdf(ShaderNode* mixNode, GenContext& context)
+bool ShaderGraph::optimizeMixBsdf(ShaderNode* mixNode, GenContext& context)
 {
     // criteria for optimization...
     // * upstream nodes for MixBsdf node both have `weight` input ports
@@ -1033,29 +1057,29 @@ void ShaderGraph::optimizeMixBsdf(ShaderNode* mixNode, GenContext& context)
     // the standard data library ND_mix_bsdf should always have "fg" and "bg" inputs
     // but we check here anyway to ensure we're not using a custom data library that doesn't follow that convention
     if (!mixFgInput || !mixBgInput)
-        return;
+        return false;
 
     auto fgNode = mixFgInput->getConnectedSibling();
     auto bgNode = mixBgInput->getConnectedSibling();
 
     // We check to see we have two upstream nodes - there are almost certainly other optimizations possible
     // if this isn't true, but we will leave those for a later PR.
     if (!fgNode && !bgNode)
-        return;
+        return false;
 
     auto fgNodeWeightInput = fgNode->getInput("weight");
     auto bgNodeWeightInput = bgNode->getInput("weight");
 
     // We require both upstream nodes to have a "weight" input for this optimization to work.
     if (!fgNodeWeightInput || !bgNodeWeightInput)
-        return;
+        return false;
 
     // we also require the following list of node definitions
     auto addBsdfNodeDef = _document->getNodeDef("ND_add_bsdf");
     auto floatInvertNodeDef = _document->getNodeDef("ND_invert_float");
     auto floatMultNodeDef = _document->getNodeDef("ND_multiply_float");
     if (!addBsdfNodeDef || !floatInvertNodeDef || !floatMultNodeDef)
-        return;
+        return false;
 
     // We meet the requirements for the optimization.
     // We can now create the new nodes and connect them up.
@@ -1127,6 +1151,213 @@ void ShaderGraph::optimizeMixBsdf(ShaderNode* mixNode, GenContext& context)
     {
         addNode->getOutput("out")->makeConnection(conn);
     }
+
+    removeNode(mixNode);
+
+    return true;
+}
+
+
+// Optimize the combination of two MixBsdf nodes by replacing it with a new node graph.
+//
+// The current nodegraph
+// ┌────────┐┌────────┐┌────────┐
+// │A.weight││B.weight││Mix1.mix│
+// └┬───────┘└┬───────┘└┬───────┘
+// ┌▽───────┐┌▽───────┐ │
+// │A (BSDF)││B (BSDF)│ │
+// └┬───────┘└┬───────┘ │
+// ┌▽─────────▽─────────▽┐┌────────┐┌────────┐
+// │Mix_1                ││C.weight││Mix2.mix│
+// └┬────────────────────┘└┬───────┘└┬───────┘
+//  |                     ┌▽───────┐ |
+//  |                     │C (BSDF)│ |
+//  |                     └┬───────┘ |
+// ┌▽──────────────────────▽─────────▽───────┐
+// |Mix_2                                    |
+// └─────────────────────────────────────────┘
+//
+// New nodegraph //TODO
+// ┌────────┐
+// │Mix.mix │
+// └┬──────┬┘
+// ┌▽─────┐│┌────────┐┌────────┐
+// │Invert│││A.weight││B.weight│
+// └─────┬┘│└┬───────┘└──┬─────┘
+//       └─│─│──┐        │
+// ┌───────▽─▽┐┌▽────────▽┐
+// │Multiply A││Multiply B│
+// └┬─────────┘└┬─────────┘
+// ┌▽───────┐┌──▽─────┐
+// │A (BSDF)││B (BSDF)│
+// └┬───────┘└┬───────┘
+// ┌▽─────────▽┐
+// │Add        │
+// └───────────┘
+//
+// Motivation - the new graph is more efficient for shader backends to optimize
+// away possible expensive BSDF nodes that might not be used at all.
+bool ShaderGraph::optimizeMixMixBsdf(ShaderNode* mixNode_x, GenContext& context)
+{
+    // criteria for optimization...
+    // * upstream nodes for MixBsdf node both have `weight` input ports
+    // * We have the following node definitions available in the library
+    //   * ND_add_bsdf
+    //   * ND_invert_float
+    //   * ND_multiply_float
+
+    ShaderNode* mix2Node = mixNode_x;
+
+    auto mix2FgInput = mix2Node->getInput("fg");
+    auto mix2BgInput = mix2Node->getInput("bg");
+
+    // the standard data library ND_mix_bsdf should always have "fg" and "bg" inputs
+    // but we check here anyway to ensure we're not using a custom data library that doesn't follow that convention
+    if (!mix2FgInput || !mix2BgInput)
+        return false;
+
+    auto mix2FgNode = mix2FgInput->getConnectedSibling();
+    auto mix2BgNode = mix2BgInput->getConnectedSibling();
+
+    // We check to see we have two upstream nodes - there are almost certainly other optimizations possible
+    // if this isn't true, but we will leave those for a later PR.
+    if (!mix2FgNode && !mix2BgNode)
+        return false;
+
+    // we require the node connected to the "fg" input to also be a mixBsdf node
+    if (!mix2FgNode->hasClassification(ShaderNode::Classification::MIX_BSDF))
+        return false;
+
+    ShaderNode* mix1Node = mix2FgNode;
+
+
+    auto mix1FgInput = mix1Node->getInput("fg");
+    auto mix1BgInput = mix1Node->getInput("bg");
+
+    // the standard data library ND_mix_bsdf should always have "fg" and "bg" inputs
+    // but we check here anyway to ensure we're not using a custom data library that doesn't follow that convention
+    if (!mix1FgInput || !mix1BgInput)
+        return false;
+
+    auto mix1FgNode = mix1FgInput->getConnectedSibling();
+    auto mix1BgNode = mix1BgInput->getConnectedSibling();
+
+    // We check to see we have two upstream nodes - there are almost certainly other optimizations possible
+    // if this isn't true, but we will leave those for a later PR.
+    if (!mix1FgNode && !mix1BgNode)
+        return false;
+
+    auto mix1FgNodeWeightInput = mix1FgNode->getInput("weight");
+    auto mix1BgNodeWeightInput = mix1BgNode->getInput("weight");
+    auto mix2BgNodeWeightInput = mix2BgNode->getInput("weight");
+
+    // We require both upstream nodes to have a "weight" input for this optimization to work.
+    if (!mix1FgNodeWeightInput || !mix1BgNodeWeightInput || !mix2BgNodeWeightInput)
+        return false;
+
+
+    // we also require the following list of node definitions
+    auto addBsdfNodeDef = _document->getNodeDef("ND_add_bsdf");
+    auto floatInvertNodeDef = _document->getNodeDef("ND_invert_float");
+    auto floatMultNodeDef = _document->getNodeDef("ND_multiply_float");
+    if (!addBsdfNodeDef || !floatInvertNodeDef || !floatMultNodeDef)
+        return false;
+
+    // We meet the requirements for the optimization.
+    // We can now create the new nodes and connect them up.
+
+    // Helper function to redirect the incoming connection to from one input port
+    // to another.
+    // We intentionally skip error checking here, as we're doing it below.
+    // If this proves useful we should make it a method somewhere, and add
+    // more robust error checking.
+    auto redirectInput = [](ShaderInput* fromPort, ShaderInput* toPort) -> void
+    {
+        auto connection = fromPort->getConnection();
+        if (connection)
+        {
+            // we have a connection - so transfer it
+            toPort->makeConnection(connection);
+        }
+        else
+        {
+            // we just remap the value.
+            toPort->setValue(fromPort->getValue());
+        }
+    };
+
+    // Helper function to connect two nodes together, consolidating the valiation of the ports existance.
+    // If this proves useful we should make it a method somewhere.
+    auto connectNodes = [](ShaderNode* fromNode, const string& fromPortName, ShaderNode* toNode, const string& toPortName) -> void
+    {
+        auto fromPort = fromNode->getOutput(fromPortName);
+        auto toPort = toNode->getInput(toPortName);
+        if (!fromPort || !toPort)
+            return;
+
+        fromPort->makeConnection(toPort);
+    };
+
+    auto mix1WeightInput = mix1Node->getInput("mix");
+    auto mix2WeightInput = mix2Node->getInput("mix");
+
+    // create nodes that represents the inverted mix values, ie. 1.0-mix
+    // to be used for the "bg" side of the mix
+    auto invertMix1Node = this->createNode(mix1Node->getName()+"_INV", floatInvertNodeDef, context);
+    redirectInput(mix1WeightInput, invertMix1Node->getInput("in"));
+    auto invertMix2Node = this->createNode(mix2Node->getName()+"_INV", floatInvertNodeDef, context);
+    redirectInput(mix2WeightInput, invertMix2Node->getInput("in"));
+
+
+    // create a multiply node to calculate the new weight value, weighted by the mix value.
+    auto multFg1WeightNode_intermediate = this->createNode(mix1Node->getName()+"_MULT_FG1_INTERMEDIATE", floatMultNodeDef, context);
+    redirectInput(mix1FgNodeWeightInput, multFg1WeightNode_intermediate->getInput("in1"));
+    redirectInput(mix1WeightInput, multFg1WeightNode_intermediate->getInput("in2"));
+    auto multFg1WeightNode = this->createNode(mix1Node->getName()+"_MULT_FG1", floatMultNodeDef, context);
+    redirectInput(mix2WeightInput, multFg1WeightNode->getInput("in1"));
+    connectNodes(multFg1WeightNode_intermediate, "out", multFg1WeightNode, "in2");
+
+    auto multBg1WeightNode_intermediate = this->createNode(mix1Node->getName()+"_MULT_BG1_INTERMEDIATE", floatMultNodeDef, context);
+    redirectInput(mix1BgNodeWeightInput, multBg1WeightNode_intermediate->getInput("in1"));
+    connectNodes(invertMix1Node, "out", multBg1WeightNode_intermediate, "in2");
+    auto multBg1WeightNode = this->createNode(mix1Node->getName()+"_MULT_BG1", floatMultNodeDef, context);
+    redirectInput(mix2WeightInput, multBg1WeightNode->getInput("in1"));
+    connectNodes(multBg1WeightNode_intermediate, "out", multBg1WeightNode, "in2");
+
+    auto multBg2WeightNode = this->createNode(mix2Node->getName()+"_MULT_BG2", floatMultNodeDef, context);
+    redirectInput(mix2BgNodeWeightInput, multBg2WeightNode->getInput("in1"));
+    connectNodes(invertMix2Node, "out", multBg2WeightNode, "in2");
+
+
+    // connect the two newly created weights to the fg and bg BSDF nodes.
+    connectNodes(multFg1WeightNode, "out", mix1FgNode, "weight");
+    connectNodes(multBg1WeightNode, "out", mix1BgNode, "weight");
+    connectNodes(multBg2WeightNode, "out", mix2BgNode, "weight");
+
+
+    // Create the ND_add_bsdf nodes that will add the three BSDF nodes with the modified weights
+    // this replaces the original mix nodes.
+    auto addNode_intermediate = this->createNode(mix2Node->getName()+"_ADD_INTERMEDIATE", addBsdfNodeDef, context);
+    connectNodes(mix1BgNode, "out", addNode_intermediate, "in1");
+    connectNodes(mix1FgNode, "out", addNode_intermediate, "in2");
+
+    auto addNode = this->createNode(mix2Node->getName()+"_ADD", addBsdfNodeDef, context);
+    connectNodes(addNode_intermediate, "out", addNode, "in1");
+    connectNodes(mix2BgNode, "out", addNode, "in2");
+
+    // Finally for all the previous outgoing connections from the original mix node
+    // replace those with the outgoing connection from the new add node.
+    auto mixNodeOutput = mix2Node->getOutput("out");
+    auto mixNodeOutputConns = mixNodeOutput->getConnections();
+    for (auto conn : mixNodeOutputConns)
+    {
+        addNode->getOutput("out")->makeConnection(conn);
+    }
+
+    removeNode(mix1Node);
+    removeNode(mix2Node);
+
+    return true;
 }
 
 

source/MaterialXGenShader/ShaderGraph.h

@@ -137,6 +137,8 @@ class MX_GENSHADER_API ShaderGraph : public ShaderNode
     /// Add a node to the graph
     void addNode(ShaderNodePtr node);
 
+    bool removeNode(ShaderNode* node);
+
     /// Add input sockets from an interface element (nodedef, nodegraph or node)
     void addInputSockets(const InterfaceElement& elem, GenContext& context);
 
@@ -166,7 +168,8 @@ class MX_GENSHADER_API ShaderGraph : public ShaderNode
     /// Optimize the graph, removing redundant paths.
     void optimize(GenContext& context);
 
-    void optimizeMixBsdf(ShaderNode* node, GenContext& context);
+    bool optimizeMixBsdf(ShaderNode* node, GenContext& context);
+    bool optimizeMixMixBsdf(ShaderNode* node, GenContext& context);
 
     /// Bypass a node for a particular input and output,
     /// effectively connecting the input's upstream connection