Propose an initial graph refactoring framework for shader generation

source/JsMaterialX/JsMaterialXGenShader/JsGenOptions.cpp

@@ -36,6 +36,9 @@ EMSCRIPTEN_BINDINGS(GenOptions)
         .property("targetDistanceUnit", &mx::GenOptions::targetDistanceUnit)
         .property("addUpstreamDependencies", &mx::GenOptions::addUpstreamDependencies)
         .property("emitColorTransforms", &mx::GenOptions::emitColorTransforms)
+        .property("elideConstantNodes", &mx::GenOptions::elideConstantNodes)
+        .property("premultipliedBsdfAdd", &mx::GenOptions::premultipliedBsdfAdd)
+        .property("distributeLayerOverBsdfMix", &mx::GenOptions::distributeLayerOverBsdfMix)
         .property("hwTransparency", &mx::GenOptions::hwTransparency)
         .property("hwSpecularEnvironmentMethod", &mx::GenOptions::hwSpecularEnvironmentMethod)
         .property("hwDirectionalAlbedoMethod", &mx::GenOptions::hwDirectionalAlbedoMethod)

source/MaterialXGenHw/HwShaderGenerator.cpp

@@ -89,6 +89,12 @@ HwShaderGenerator::HwShaderGenerator(TypeSystemPtr typeSystem, SyntaxPtr syntax)
     _tokenSubstitutions[HW::T_CLOSURE_DATA_CONSTRUCTOR] = HW::CLOSURE_DATA_CONSTRUCTOR;
 }
 
+void HwShaderGenerator::applyDefaultOptions(GenOptions& options) const
+{
+    ShaderGenerator::applyDefaultOptions(options);
+    options.premultipliedBsdfAdd = true;
+}
+
 ShaderPtr HwShaderGenerator::createShader(const string& name, ElementPtr element, GenContext& context) const
 {
     // Create the root shader graph

source/MaterialXGenHw/HwShaderGenerator.h

@@ -43,6 +43,9 @@ class MX_GENHW_API HwShaderGenerator : public ShaderGenerator
     void emitClosureDataArg(const ShaderNode& node, GenContext& context, ShaderStage& stage) const override;
     void emitClosureDataParameter(const ShaderNode& node, GenContext& context, ShaderStage& stage) const override;
 
+    /// Apply the default GenOptions for hardware shader generation.
+    void applyDefaultOptions(GenOptions& options) const override;
+
     /// Logic to indicate whether code to support direct lighting should be emitted.
     /// By default if the graph is classified as a shader, or BSDF node then lighting is assumed to be required.
     /// Derived classes can override this logic.

source/MaterialXGenMdl/MdlShaderGenerator.cpp

@@ -125,6 +125,12 @@ MdlShaderGenerator::MdlShaderGenerator(TypeSystemPtr typeSystem) :
     registerImplementation("IM_image_vector4_" + MdlShaderGenerator::TARGET, ImageNodeMdl::create);
 }
 
+void MdlShaderGenerator::applyDefaultOptions(GenOptions& options) const
+{
+    ShaderGenerator::applyDefaultOptions(options);
+    options.distributeLayerOverBsdfMix = true;
+}
+
 ShaderPtr MdlShaderGenerator::generate(const string& name, ElementPtr element, GenContext& context) const
 {
     // For MDL we cannot cache node implementations between generation calls,

source/MaterialXGenMdl/MdlShaderGenerator.h

@@ -70,6 +70,9 @@ class MX_GENMDL_API MdlShaderGenerator : public ShaderGenerator
     /// Return a unique identifier for the target this generator is for
     const string& getTarget() const override { return TARGET; }
 
+    /// Apply the default GenOptions for MDL shader generation.
+    void applyDefaultOptions(GenOptions& options) const override;
+
     /// Generate a shader starting from the given element, translating
     /// the element and all dependencies upstream into shader code.
     ShaderPtr generate(const string& name, ElementPtr element, GenContext& context) const override;

source/MaterialXGenShader/GenContext.cpp

@@ -20,6 +20,9 @@ GenContext::GenContext(ShaderGeneratorPtr sg) :
         throw ExceptionShaderGenError("GenContext must have a valid shader generator");
     }
 
+    // Apply the generator's default options for its target.
+    _sg->applyDefaultOptions(_options);
+
     // Collect and cache reserved words from the shader generator
     StringSet reservedWords;
 

source/MaterialXGenShader/GenOptions.h

@@ -82,6 +82,8 @@ class MX_GENSHADER_API GenOptions
         libraryPrefix("libraries"),
         emitColorTransforms(true),
         elideConstantNodes(true),
+        premultipliedBsdfAdd(false),
+        distributeLayerOverBsdfMix(false),
         hwTransparency(false),
         hwSpecularEnvironmentMethod(SPECULAR_ENVIRONMENT_FIS),
         hwDirectionalAlbedoMethod(DIRECTIONAL_ALBEDO_ANALYTIC),
@@ -141,6 +143,17 @@ class MX_GENSHADER_API GenOptions
     /// Enable eliding constant nodes. Defaults to true.
     bool elideConstantNodes;
 
+    /// Enable replacing BSDF mix nodes with premultiplied add nodes.
+    /// This folds the mix weight into each BSDF's weight input, enabling
+    /// hardware shading languages to skip BSDF evaluation via dynamic
+    /// branching when the weight is zero. Defaults to false.
+    bool premultipliedBsdfAdd;
+
+    /// Enable distributing layer operations over mix nodes.
+    /// Transforms layer(mix(A, B, w), C) into mix(layer(A, C), layer(B, C), w)
+    /// for backends with limited layering capabilities. Defaults to false.
+    bool distributeLayerOverBsdfMix;
+
     /// Sets if transparency is needed or not for HW shaders.
     /// If a surface shader has potential of being transparent
     /// this must be set to true, otherwise no transparency

source/MaterialXGenShader/ShaderGenerator.cpp

@@ -32,6 +32,16 @@ ShaderGenerator::ShaderGenerator(TypeSystemPtr typeSystem, SyntaxPtr syntax) :
     _typeSystem(typeSystem),
     _syntax(syntax)
 {
+    // Register all graph refactoring passes.
+    registerRefactor(std::make_shared<NodeElisionRefactor>());
+    registerRefactor(std::make_shared<PremultipliedBsdfAddRefactor>());
+    registerRefactor(std::make_shared<DistributeLayerOverMixRefactor>());
+}
+
+void ShaderGenerator::applyDefaultOptions(GenOptions& /*options*/) const
+{
+    // Base implementation sets no additional defaults.
+    // Derived generators override to set target-specific defaults.
 }
 
 void ShaderGenerator::emitScopeBegin(ShaderStage& stage, Syntax::Punctuation punc) const

source/MaterialXGenShader/ShaderGenerator.h

@@ -14,6 +14,7 @@
 #include <MaterialXGenShader/ColorManagementSystem.h>
 #include <MaterialXGenShader/Exception.h>
 #include <MaterialXGenShader/Factory.h>
+#include <MaterialXGenShader/ShaderGraphRefactor.h>
 #include <MaterialXGenShader/ShaderStage.h>
 #include <MaterialXGenShader/Syntax.h>
 
@@ -214,6 +215,21 @@ class MX_GENSHADER_API ShaderGenerator
         return _tokenSubstitutions;
     }
 
+    /// Register a shader graph refactoring pass.
+    void registerRefactor(ShaderGraphRefactorPtr refactor)
+    {
+        _refactors.push_back(refactor);
+    }
+
+    /// Return the registered graph refactoring passes.
+    const vector<ShaderGraphRefactorPtr>& getRefactors() const
+    {
+        return _refactors;
+    }
+
+    /// Apply the default GenOptions for this generator's target.
+    virtual void applyDefaultOptions(GenOptions& options) const;
+
     /// Register type definitions from the document.
     virtual void registerTypeDefs(const DocumentPtr& doc);
 
@@ -261,6 +277,7 @@ class MX_GENSHADER_API ShaderGenerator
     ColorManagementSystemPtr _colorManagementSystem;
     UnitSystemPtr _unitSystem;
     mutable StringMap _tokenSubstitutions;
+    vector<ShaderGraphRefactorPtr> _refactors;
 
     friend ShaderGraph;
 };

source/MaterialXGenShader/ShaderGraph.cpp

@@ -7,6 +7,7 @@
 
 #include <MaterialXGenShader/Exception.h>
 #include <MaterialXGenShader/GenContext.h>
+#include <MaterialXGenShader/ShaderGraphRefactor.h>
 #include <MaterialXGenShader/Util.h>
 
 #include <iostream>
@@ -910,8 +911,18 @@ void ShaderGraph::finalize(GenContext& context)
     _inputUnitTransformMap.clear();
     _outputUnitTransformMap.clear();
 
-    // Optimize the graph, removing redundant paths.
-    optimize(context);
+    // Run registered graph refactoring passes.
+    size_t totalEdits = 0;
+    for (auto& refactor : context.getShaderGenerator().getRefactors())
+    {
+        totalEdits += refactor->execute(*this, context);
+    }
+
+    // Remove unused nodes if any refactoring pass made edits.
+    if (totalEdits > 0)
+    {
+        removeUnusedNodes();
+    }
 
     // Sort the nodes in topological order.
     topologicalSort();
@@ -971,97 +982,57 @@ void ShaderGraph::disconnect(ShaderNode* node) const
     }
 }
 
-void ShaderGraph::optimize(GenContext& context)
+void ShaderGraph::removeUnusedNodes()
 {
-    size_t numEdits = 0;
-    for (ShaderNode* node : getNodes())
+    std::set<ShaderNode*> usedNodesSet;
+    std::vector<ShaderNode*> usedNodesVec;
+
+    // Traverse the graph to find nodes still in use.
+    for (ShaderGraphOutputSocket* outputSocket : getOutputSockets())
     {
-        if (node->hasClassification(ShaderNode::Classification::CONSTANT))
+        // Make sure to not include connections to the graph itself.
+        ShaderOutput* upstreamPort = outputSocket->getConnection();
+        if (upstreamPort && upstreamPort->getNode() != this)
         {
-            if (node->numInputs() != 1 || node->numOutputs() != 1)
-            {
-                // Constant node doesn't follow expected interface, cannot elide.
-                continue;
-            }
-            // Constant nodes can be elided by moving their value downstream.
-            bool canElide = context.getOptions().elideConstantNodes;
-            if (!canElide)
+            for (ShaderGraphEdge edge : traverseUpstream(upstreamPort))
             {
-                // We always elide filename constant nodes regardless of the
-                // option. See DOT below.
-                ShaderInput* in = node->getInput("value");
-                if (in && in->getType() == Type::FILENAME)
+                ShaderNode* node = edge.upstream->getNode();
+                if (usedNodesSet.count(node) == 0)
                 {
-                    canElide = true;
+                    usedNodesSet.insert(node);
+                    usedNodesVec.push_back(node);
                 }
             }
-            if (canElide)
-            {
-                bypass(node, 0);
-                ++numEdits;
-            }
-        }
-        else if (node->hasClassification(ShaderNode::Classification::DOT))
-        {
-            if (node->numOutputs() != 1)
-            {
-                // Dot node dosen't follow expected interface, cannot elide.
-                continue;
-            }
-            // Filename dot nodes must be elided so they do not create extra samplers.
-            ShaderInput* in = node->getInput("in");
-            if (in && in->getType() == Type::FILENAME)
-            {
-                bypass(node, 0);
-                ++numEdits;
-            }
         }
-        // Adding more nodes here requires them to have an input that is tagged
-        // "uniform" in the NodeDef or to handle very specific cases, like FILENAME.
     }
 
-    if (numEdits > 0)
+    // Remove any unused nodes.
+    for (auto it = _nodeMap.begin(); it != _nodeMap.end();)
     {
-        std::set<ShaderNode*> usedNodesSet;
-        std::vector<ShaderNode*> usedNodesVec;
-
-        // Traverse the graph to find nodes still in use
-        for (ShaderGraphOutputSocket* outputSocket : getOutputSockets())
+        if (usedNodesSet.count(it->second.get()) == 0)
         {
-            // Make sure to not include connections to the graph itself.
-            ShaderOutput* upstreamPort = outputSocket->getConnection();
-            if (upstreamPort && upstreamPort->getNode() != this)
-            {
-                for (ShaderGraphEdge edge : traverseUpstream(upstreamPort))
-                {
-                    ShaderNode* node = edge.upstream->getNode();
-                    if (usedNodesSet.count(node) == 0)
-                    {
-                        usedNodesSet.insert(node);
-                        usedNodesVec.push_back(node);
-                    }
-                }
-            }
-        }
+            // Break all connections.
+            disconnect(it->second.get());
 
-        // Remove any unused nodes
-        for (auto it = _nodeMap.begin(); it != _nodeMap.end();)
+            // Erase from storage.
+            it = _nodeMap.erase(it);
+        }
+        else
         {
-            if (usedNodesSet.count(it->second.get()) == 0)
-            {
-                // Break all connections
-                disconnect(it->second.get());
-
-                // Erase from storage
-                it = _nodeMap.erase(it);
-            }
-            else
-            {
-                ++it;
-            }
+            ++it;
         }
+    }
+
+    _nodeOrder = usedNodesVec;
+}
 
-        _nodeOrder = usedNodesVec;
+void ShaderGraph::replaceOutput(ShaderOutput* oldOutput, ShaderOutput* newOutput)
+{
+    ShaderInputVec downstreamConnections = oldOutput->getConnections();
+    for (ShaderInput* downstream : downstreamConnections)
+    {
+        oldOutput->breakConnection(downstream);
+        downstream->makeConnection(newOutput);
     }
 }
 

source/MaterialXGenShader/ShaderGraph.h

@@ -26,6 +26,7 @@ class Syntax;
 class ShaderGraphEdge;
 class ShaderGraphEdgeIterator;
 class GenOptions;
+class ShaderGraphRefactor;
 
 /// An internal input socket in a shader graph,
 /// used for connecting internal nodes to the outside
@@ -126,6 +127,23 @@ class MX_GENSHADER_API ShaderGraph : public ShaderNode
     /// Return the map of unique identifiers used in the scope of this graph.
     IdentifierMap& getIdentifierMap() { return _identifiers; }
 
+    /// Return the document associated with this graph.
+    ConstDocumentPtr getDocument() const { return _document; }
+
+    /// Create a new node in the graph from a node definition.
+    ShaderNode* createNode(const string& name, const string& uniqueId, ConstNodeDefPtr nodeDef, GenContext& context);
+
+    /// Bypass a node for a particular input and output,
+    /// effectively connecting the input's upstream connection
+    /// with the output's downstream connections.
+    void bypass(ShaderNode* node, size_t inputIndex, size_t outputIndex = 0);
+
+    /// Remove nodes that are no longer connected to any output.
+    void removeUnusedNodes();
+
+    /// Rewire all downstream connections from one output to another.
+    void replaceOutput(ShaderOutput* oldOutput, ShaderOutput* newOutput);
+
   protected:
     /// Create node connections corresponding to the connection between a pair of elements.
     /// @param downstreamElement Element representing the node to connect to.
@@ -137,12 +155,6 @@ class MX_GENSHADER_API ShaderGraph : public ShaderNode
                               ElementPtr connectingElement,
                               GenContext& context);
 
-    /// Create a new node in a graph from a node definition.
-    /// The uniqueId argument is used as the node's key in the graph's node map.
-    /// Note - this does not initialize the node instance with any concrete values, but
-    /// instead creates an empty instance of the provided node definition
-    ShaderNode* createNode(const string& name, const string& uniqueId, ConstNodeDefPtr nodeDef, GenContext& context);
-
     /// Add a node to the graph, keyed by the node's unique identifier.
     void addNode(ShaderNodePtr node);
 
@@ -172,14 +184,6 @@ class MX_GENSHADER_API ShaderGraph : public ShaderNode
     /// Perform all post-build operations on the graph.
     void finalize(GenContext& context);
 
-    /// Optimize the graph, removing redundant paths.
-    void optimize(GenContext& context);
-
-    /// Bypass a node for a particular input and output,
-    /// effectively connecting the input's upstream connection
-    /// with the output's downstream connections.
-    void bypass(ShaderNode* node, size_t inputIndex, size_t outputIndex = 0);
-
     /// For inputs and outputs in the graph set the variable names to be used
     /// in generated code. Making sure variable names are valid and unique
     /// to avoid name conflicts during shader generation.