EMSUSD-1975 - Optimize material libraries

lib/mayaUsd/render/MaterialXGenOgsXml/LobePruner.cpp

@@ -85,10 +85,14 @@ class LobePrunerImpl
     LobePrunerImpl(LobePrunerImpl&&) = delete;
     LobePrunerImpl& operator=(LobePrunerImpl&&) = delete;
 
-    bool          getOptimizedNodeCategory(const mx::Node& node, std::string& nodeCategory);
+    mx::NodeDefPtr getOptimizedNodeDef(const mx::Node& node);
+
     mx::StringVec getOptimizedAttributeNames(const mx::NodeDefPtr& nodeDef) const;
 
     PXR_NS::TfToken getOptimizedNodeId(const PXR_NS::HdMaterialNode2& node);
+    bool            isOptimizedNodeId(const PXR_NS::TfToken& nodeId);
+
+    void optimizeLibrary(const MaterialX::DocumentPtr& library);
 
     static const std::string ND_PREFIX;
     static const std::string DARK_BASE;
@@ -105,8 +109,8 @@ class LobePrunerImpl
         const mx::InputPtr&     input,
         const mx::NodeGraphPtr& ng,
         const mx::NodeDefPtr&   nd);
-    void
-         ensureLibraryHasOptimizedShader(const PXR_NS::TfToken& nodeDefName, const std::string& flags);
+    mx::NodeDefPtr
+         getOrAddOptimizedNodeDef(const PXR_NS::TfToken& nodeDefName, const std::string& flags);
     void optimizeZeroValue(
         mx::NodeGraphPtr&          optimizedNodeGraph,
         const OptimizableValueMap& optimizationMap,
@@ -132,6 +136,7 @@ class LobePrunerImpl
 
     std::unordered_map<PXR_NS::TfToken, NodeDefData, PXR_NS::TfToken::HashFunctor> _prunerData;
     mx::DocumentPtr                                                                _library;
+    PXR_NS::TfToken::HashSet _optimizedNodeIds;
 };
 
 const std::string LobePrunerImpl::ND_PREFIX = "LPOPTIND_";
@@ -216,6 +221,90 @@ bool LobePrunerImpl::isLobeInput(const mx::InputPtr& input, const mx::NodeDefPtr
     return true;
 }
 
+void LobePrunerImpl::optimizeLibrary(const MaterialX::DocumentPtr& library)
+{
+    if (!_library || _prunerData.empty()) {
+        return;
+    }
+
+    std::set<std::string> allDefinedNodeGraphs;
+    // Go thru all NodeGraphs found in the library that have an associated NodeDef:
+    for (const auto& ng : library->getNodeGraphs()) {
+        if (ng->hasNodeDefString()) {
+            allDefinedNodeGraphs.insert(ng->getName());
+        }
+    }
+    for (const auto& impl : library->getImplementations()) {
+        if (impl->hasNodeGraph()) {
+            allDefinedNodeGraphs.insert(impl->getNodeGraph());
+        }
+    }
+
+    for (const auto& ngName : allDefinedNodeGraphs) {
+        const auto ng = library->getNodeGraph(ngName);
+        // Go thru all the nodes of that NodeGraph
+        for (const auto& node : ng->getNodes()) {
+            // Can this node be optimized?
+            const auto& nd = node->getNodeDef();
+            if (!nd) {
+                continue;
+            }
+
+            const auto ndName = PXR_NS::TfToken(nd->getName());
+            const auto ndIt = _prunerData.find(ndName);
+            if (ndIt == _prunerData.end()) {
+                continue;
+            }
+
+            // This NodeGraph contains an optimizable embedded surface shader node.
+            std::string flags(ndIt->second._attributeData.size(), 'x');
+
+            bool canOptimize = false;
+
+            auto attrIt = ndIt->second._attributeData.cbegin();
+            for (size_t i = 0; attrIt != ndIt->second._attributeData.cend(); ++attrIt, ++i) {
+                const auto nodeinput = node->getActiveInput(attrIt->first);
+                float      inputValue = 0.5F;
+                if (nodeinput) {
+                    // Can not optimize if connected in any way.
+                    if (nodeinput->hasNodeName() || nodeinput->hasOutputString()
+                        || nodeinput->hasInterfaceName()) {
+                        continue;
+                    }
+                    inputValue = nodeinput->getValue()->asA<float>();
+                } else {
+                    const auto defInput = nd->getActiveInput(attrIt->first);
+                    inputValue = defInput->getValue()->asA<float>();
+                }
+
+                for (const auto& optimizableValue : attrIt->second) {
+                    if (optimizableValue.first == inputValue) {
+                        if (inputValue == 0.0F) {
+                            flags[i] = '0';
+                        } else {
+                            flags[i] = '1';
+                        }
+                        canOptimize = true;
+                    }
+                }
+            }
+
+            if (canOptimize) {
+                const auto optimizedNodeDef = getOrAddOptimizedNodeDef(ndName, flags);
+                // Replace the node with an optimized one:
+                const auto nsPrefix = optimizedNodeDef->hasNamespace()
+                    ? optimizedNodeDef->getNamespace() + ":"
+                    : std::string {};
+
+                node->setCategory(nsPrefix + optimizedNodeDef->getNodeString());
+                if (node->hasNodeDefString()) {
+                    node->setNodeDefString(optimizedNodeDef->getName());
+                }
+            }
+        }
+    }
+}
+
 void LobePrunerImpl::addOptimizableValue(
     float                   value,
     const mx::InputPtr&     input,
@@ -242,17 +331,17 @@ void LobePrunerImpl::addOptimizableValue(
     valueMap.find(value)->second.push_back(PXR_NS::TfToken(input->getParent()->getName()));
 }
 
-bool LobePrunerImpl::getOptimizedNodeCategory(const mx::Node& node, std::string& nodeCategory)
+mx::NodeDefPtr LobePrunerImpl::getOptimizedNodeDef(const mx::Node& node)
 {
     const auto& nd = node.getNodeDef();
     if (!nd) {
-        return false;
+        return {};
     }
 
     const auto ndName = PXR_NS::TfToken(nd->getName());
     const auto ndIt = _prunerData.find(ndName);
     if (ndIt == _prunerData.end()) {
-        return false;
+        return {};
     }
 
     std::string flags(ndIt->second._attributeData.size(), 'x');
@@ -288,12 +377,10 @@ bool LobePrunerImpl::getOptimizedNodeCategory(const mx::Node& node, std::string&
     }
 
     if (canOptimize) {
-        ensureLibraryHasOptimizedShader(ndName, flags);
-        nodeCategory = node.getCategory() + "_" + flags;
-        return true;
+        return getOrAddOptimizedNodeDef(ndName, flags);
     }
 
-    return false;
+    return {};
 }
 
 mx::StringVec LobePrunerImpl::getOptimizedAttributeNames(const mx::NodeDefPtr& nodeDef) const
@@ -356,39 +443,53 @@ PXR_NS::TfToken LobePrunerImpl::getOptimizedNodeId(const PXR_NS::HdMaterialNode2
     }
 
     if (canOptimize) {
-        ensureLibraryHasOptimizedShader(node.nodeTypeId, flags);
-        return PXR_NS::TfToken(
-            ND_PREFIX + nodeDef->GetFamily().GetString() + "_" + flags + "_surfaceshader");
+        return PXR_NS::TfToken(getOrAddOptimizedNodeDef(node.nodeTypeId, flags)->getName());
     }
 
     return retVal;
 }
 
-void LobePrunerImpl::ensureLibraryHasOptimizedShader(
+bool LobePrunerImpl::isOptimizedNodeId(const PXR_NS::TfToken& nodeId)
+{
+    return _optimizedNodeIds.count(nodeId) != 0;
+}
+
+mx::NodeDefPtr LobePrunerImpl::getOrAddOptimizedNodeDef(
     const PXR_NS::TfToken& nodeDefName,
     const std::string&     flags)
 {
     const auto ndIt = _prunerData.find(nodeDefName);
     if (ndIt == _prunerData.end()) {
-        return;
+        return {};
     }
 
-    const auto        originalNodeDef = _library->getNodeDef(nodeDefName.GetString());
-    const auto        originalNodeGraph = _library->getNodeGraph(ndIt->second._nodeGraphName);
-    const std::string optimizedNodeName = originalNodeDef->getNodeString() + "_" + flags;
-    const std::string optimizedNodeDefName = ND_PREFIX + optimizedNodeName + "_surfaceshader";
-    if (_library->getNodeDef(optimizedNodeDefName)) {
+    const auto originalNodeDef = _library->getNodeDef(nodeDefName.GetString());
+    const auto originalNodeGraph = _library->getNodeGraph(ndIt->second._nodeGraphName);
+    const auto nsPrefix = originalNodeDef->hasNamespace()
+        ? originalNodeDef->getNamespace() + mx::NAME_PREFIX_SEPARATOR
+        : std::string {};
+    auto optimizedNodeName = originalNodeDef->getNodeString() + "_" + flags;
+    if (!nsPrefix.empty() && optimizedNodeName.rfind(nsPrefix, 0) == 0) {
+        optimizedNodeName = optimizedNodeName.substr(nsPrefix.size());
+    }
+    const auto        optimizedNodeNameWithNS = nsPrefix + optimizedNodeName;
+    const std::string optimizedNodeDefName
+        = nsPrefix + ND_PREFIX + optimizedNodeName + "_surfaceshader";
+    if (const auto existingNd = _library->getNodeDef(optimizedNodeDefName)) {
         // Already there
-        return;
+        return existingNd;
     }
 
+    _optimizedNodeIds.insert(PXR_NS::TfToken(optimizedNodeDefName));
+
     auto optimizedNodeDef
         = _library->addNodeDef(optimizedNodeDefName, "surfaceshader", optimizedNodeName);
     optimizedNodeDef->copyContentFrom(originalNodeDef);
     optimizedNodeDef->setSourceUri("");
     optimizedNodeDef->setNodeString(optimizedNodeName);
 
-    auto optimizedNodeGraph = _library->addNodeGraph("NG_" + optimizedNodeName + "_surfaceshader");
+    auto optimizedNodeGraph
+        = _library->addNodeGraph(nsPrefix + "LPOPTING_" + optimizedNodeName + "_surfaceshader");
     optimizedNodeGraph->copyContentFrom(originalNodeGraph);
     optimizedNodeGraph->setSourceUri("");
     optimizedNodeGraph->setNodeDefString(optimizedNodeDefName);
@@ -414,6 +515,8 @@ void LobePrunerImpl::ensureLibraryHasOptimizedShader(
         default: continue;
         }
     }
+
+    return optimizedNodeDef;
 }
 
 void LobePrunerImpl::optimizeZeroValue(
@@ -560,9 +663,16 @@ LobePruner::Ptr LobePruner::create() { return std::make_shared<LobePruner>(); }
 LobePruner::~LobePruner() = default;
 LobePruner::LobePruner() = default;
 
-bool LobePruner::getOptimizedNodeCategory(const mx::Node& node, std::string& nodeCategory)
+void LobePruner::optimizeLibrary(const MaterialX::DocumentPtr& library)
+{
+    if (_impl) {
+        _impl->optimizeLibrary(library);
+    }
+}
+
+mx::NodeDefPtr LobePruner::getOptimizedNodeDef(const mx::Node& node)
 {
-    return _impl ? _impl->getOptimizedNodeCategory(node, nodeCategory) : false;
+    return _impl ? _impl->getOptimizedNodeDef(node) : mx::NodeDefPtr {};
 }
 
 mx::StringVec LobePruner::getOptimizedAttributeNames(const mx::NodeDefPtr& nodeDef) const
@@ -582,7 +692,7 @@ void LobePruner::setLibrary(const mx::DocumentPtr& library)
 
 bool LobePruner::isOptimizedNodeId(const PXR_NS::TfToken& nodeId)
 {
-    return nodeId.GetString().rfind(LobePrunerImpl::ND_PREFIX, 0) == 0;
+    return _impl ? _impl->isOptimizedNodeId(nodeId) : false;
 }
 
 const std::string& LobePruner::getOptimizedNodeDefPrefix() { return LobePrunerImpl::ND_PREFIX; }

lib/mayaUsd/render/MaterialXGenOgsXml/LobePruner.h

@@ -54,19 +54,18 @@ class MAYAUSD_CORE_PUBLIC LobePruner
      */
     void setLibrary(const mx::DocumentPtr& library);
 
+    /*! Traverses and optimizes in place all NodeGraphs found in the library. Useful if a library
+     * contains NodeGraphs that embed optimizable nodes.
+     * @param[in] library is the library used to generate shaders.
+     */
+    void optimizeLibrary(const MaterialX::DocumentPtr& library);
+
     /*! Checks if a node is optimizable and if this is the case, create the optimized NodeDef and
-     * NodeGraph in the library and return the optimized node category.  An optimized node category
-     * will consist of the name of the original category followed by a series of characters
-     * describing which attibutes were optimized:
-     *   - 'x' that attribute was not optimized (intermediate value or connected)
-     *   - '0' a zero value was optimized
-     *   - '1' a one value was optimized
-     * The ordered list of attributes names can be found by calling getOptimizedAttributeNames().
+     * NodeGraph in the library and return the optimized NodeDef.
      * @param[in] node is a node we want to optimize. All nodes are welcome.
-     * @param[out] nodeCategory is the node category of the optimized node.
-     *  \return true if an optimization was found
+     *  \return the optimized NodeDef if one was found, nullptr otherwise
      */
-    bool getOptimizedNodeCategory(const mx::Node& node, std::string& nodeCategory);
+    mx::NodeDefPtr getOptimizedNodeDef(const mx::Node& node);
 
     /*! Get the list of attribute names that are optimization targets for a specific NodeDef.
      * @param[in] nodeDef is preferably node definition that has previously been optimized, but all
@@ -95,7 +94,7 @@ class MAYAUSD_CORE_PUBLIC LobePruner
      * @param[in] nodeId is a node:id we want to check.
      *  \return true if that node was generated by a LobePruner.
      */
-    static bool isOptimizedNodeId(const PXR_NS::TfToken& nodeId);
+    bool isOptimizedNodeId(const PXR_NS::TfToken& nodeId);
 
     /*! Returns the NodeDef prefix common to all LobePruner optimized definitions.
      *  \return the LobePruner NodeDef prefix

lib/mayaUsd/render/MaterialXGenOgsXml/ShaderGenUtil.cpp

@@ -233,11 +233,14 @@ mx::NodePtr TopoNeutralGraph::cloneNode(const mx::Node& node, mx::GraphElement&
     if (!nodeDef) {
         throw mx::Exception("Ambiguous node is not fully resolvable");
     }
-    std::string optimizedNodeCategory;
-    if (_lobePruner && _lobePruner->getOptimizedNodeCategory(node, optimizedNodeCategory)) {
-        destNode->setCategory(optimizedNodeCategory);
-        destNode->setNodeDefString(
-            LobePruner::getOptimizedNodeDefPrefix() + optimizedNodeCategory + "_surfaceshader");
+    auto optimizedNodeDef
+        = _lobePruner ? _lobePruner->getOptimizedNodeDef(node) : mx::NodeDefPtr {};
+    if (optimizedNodeDef) {
+        const auto nsPrefix = optimizedNodeDef->hasNamespace()
+            ? optimizedNodeDef->getNamespace() + ":"
+            : std::string {};
+        destNode->setCategory(nsPrefix + optimizedNodeDef->getNodeString());
+        destNode->setNodeDefString(optimizedNodeDef->getName());
         for (const auto& attrName : _lobePruner->getOptimizedAttributeNames(nodeDef)) {
             _optimizedAttributes.push_back(node.getNamePath() + "." + attrName);
         }

lib/mayaUsd/render/vp2RenderDelegate/material.cpp

@@ -359,6 +359,7 @@ struct _MaterialXData
 #if MX_COMBINED_VERSION >= 13808
             _lobePruner = MaterialXMaya::ShaderGenUtil::LobePruner::create();
             _lobePruner->setLibrary(_mtlxLibrary);
+            _lobePruner->optimizeLibrary(_mtlxLibrary);
 
             // TODO: Optimize published shaders.
             // SCENARIO: User publishes a shader with a NodeGraph implementation that encapsulates a
@@ -3038,8 +3039,7 @@ MHWRender::MShaderInstance* HdVP2Material::CompiledNetwork::_CreateMaterialXShad
         const mx::FileSearchPath& crLibrarySearchPath(_GetMaterialXData()._mtlxSearchPath);
 #if MX_COMBINED_VERSION >= 13808
         if (mtlxSdrNode
-            || MaterialXMaya::ShaderGenUtil::LobePruner::isOptimizedNodeId(
-                surfTerminal->nodeTypeId)) {
+            || _GetMaterialXData()._lobePruner->isOptimizedNodeId(surfTerminal->nodeTypeId)) {
 #else
         if (mtlxSdrNode) {
 #endif

lib/mayaUsdAPI/render.cpp

@@ -157,6 +157,13 @@ void LobePruner::setLibrary(const MaterialX::DocumentPtr& library)
     }
 }
 
+void LobePruner::optimizeLibrary(const MaterialX::DocumentPtr& library)
+{
+    if (_imp && _imp->_lobePruner) {
+        _imp->_lobePruner->optimizeLibrary(library);
+    }
+}
+
 struct TopoNeutralGraphImpl
 {
     TopoNeutralGraphImpl(const MaterialX::ElementPtr& material)

lib/mayaUsdAPI/render.h

@@ -135,6 +135,12 @@ class MAYAUSD_API_PUBLIC LobePruner
      */
     void setLibrary(const MaterialX::DocumentPtr& library);
 
+    /*! Traverses and optimizes in place all NodeGraphs found in the library. Useful if a library
+     * contains NodeGraphs that embed optimizable nodes.
+     * @param[in] library is the library used to generate shaders.
+     */
+    void optimizeLibrary(const MaterialX::DocumentPtr& library);
+
 private:
     LobePruner();