Early lobe pruning during ShaderGraph construction

[ppenenko]

Summary

This PR partially implements #2680 - an optional early pruning optimization to MaterialX shader generation. The implementation is heavily inspired by the OpenUSD PR PixarAnimationStudios/OpenUSD#3525

When GenOptions::enableLobePruning is enabled, the system analyzes NodeGraph topology to identify topological ("permutation-defining") inputs (mix weights, multiply factors) that, when hardcoded to 0 or 1 at a call site, allow entire dead branches to be pruned. In contrast to #2499 and #2832 which rely on rewriting the already-constructed ShaderGraph in the optimize call, this PR applies pruning at an earlier stage - during ShaderGraph construction, during the NodeGraph traversal. If a sub-node is found to be prunable, the traversal skips it entirely. This allows us to save time even at the codegen stage.

Key components

• NodeGraphTopology (new): Statically analyzes a NodeGraph to identify topological inputs and which internal nodes each can eliminate. Results are cached per NodeGraph in GenContext via NodeGraphTopologyCache.
• NodeGraphPermutation (new): Lightweight object representing a specific pruning configuration for a call site, produced by NodeGraphTopology::createPermutation().
• CompoundNode / ShaderGenerator changes: Thread the permutation through createShaderNodeImplForNodeGraph() → ShaderGraph::create() → createConnectedNodes(), where shouldPrune() skips dead nodes before they are instantiated.
• GenOptions::enableLobePruning: Runtime flag (default off) to enable the optimization. Wired through test suite options.

Measurement infrastructure

The optimization was measured with a set of companion PRs that add profiling and comparison tooling to the MaterialX test suite:

• Extended tracing instrumentation for shader codegen and rendering #2820 — Extended tracing instrumentation for shader codegen and rendering
• Add envSampleCount and framesPerMaterial test suite options #2821 — envSampleCount and framesPerMaterial test suite options
• Python comparison scripts for test suite output (diff_test_runs) #2822 — Python comparison scripts for test suite output (diff_test_runs)
• GPU async timing instrumentation via MX_TRACE_ASYNC #2824 — GPU async timing instrumentation via MX_TRACE_ASYNC

Results

Measured on the MaterialX test suite (baseline vs. enableLobePruning=true).

GPU frame time has not improved noticeably.

Shader compilation time shows clear improvement for materials with prunable lobes:

The lines of code metric also shows a clear reduction, reflecting the decreased shader complexity:

Further work

This PR performs what is effectively early dead code elimination (DCE) based on input values that are known to be hardcoded for all possible instances of the material - for example, unconnected subnode inputs assuming default values. This is the same optimization that the downstream shader compiler (e.g. glslang + spirv-opt or the video driver compiler) would eventually perform, resulting in an equally optimized final shader binary, which explains the lack of improvement in GPU frame time. However, DCE performed earlier, during MaterialX codegen is cheaper and saves the more expensive time overhead of DCE in the compiler.

A natural follow-up is to proactively generate more optimal shader permutations by specializing on the topological input values encountered at run time, similar to Jerry's approach in PixarAnimationStudios/OpenUSD#3525.

[jstone-lucasfilm]

Thanks for putting together this proposal, @ppenenko! The shader source size and compile-time reductions are great to see, and the benefits of this change are clear.

One choice that I'd like to discuss before we proceed further along this path: What are your thoughts on organizing this as a new ShaderGraphRefactor pass rather than as construction-time pruning? I could imagine a LobePruningRefactor pass composing well with the NodeElisionRefactor, PremultipliedBsdfAddRefactor, and DistributeLayerOverMixRefactor passes that we already have.

Here are a few potential benefits of this approach:

• Shared analysis surface: The framework runs passes inside ShaderGraph::finalize against a constructed graph where node classifications (BSDF | CLOSURE | MIX | LAYER) and resolved connections are already available. The proposed NodeGraphTopology analyzer reimplements reference counting at the document layer with hardcoded category strings ("mix", "multiply", the kWeightedPbrNodes allowlist), most of which the ShaderGraph layer already gives us for free.
• Natural overlap with PremultipliedBsdfAddRefactor: Both passes special-case BSDF mix nodes with weight inputs, and living side-by-side as refactor passes makes that sharing straightforward.
• Composability: Passes run in sequence with removeUnusedNodes between them, so e.g. constant elision can clean up inputs that became unconnected after pruning, with no bespoke plumbing.
• API consistency: Adding a ShaderGraphRefactor pass would require no changes to createShaderNodeImplForNodeGraph, making this PR transparent to custom shader generators.

I think the strongest argument for your construction-time approach is that lobe pruning is call-site-specialized in a way the existing refactor passes aren't. The same NodeGraph needs to yield different ShaderGraphs depending on the parent material's inputs. That's not trivial, but I don't think it's blocking; a refactor pass running inside CompoundNode::initialize already has access to GenContext::getParentNodes, so it can specialize the same way without leaving the framework.

The remaining benefit of your proposed approach is avoiding ShaderNode::create and initialize for pruned nodes. Do you have a sense of how much of the measured code generation time win comes from that specifically, vs. just from emitting fewer nodes downstream? If most of the win is downstream, then a refactor-pass form likely captures most of it while staying consistent with the framework that just landed.

Let me know your thoughts!

[ppenenko]

Hi @jstone-lucasfilm, thanks for the feedback. Let me elaborate what lead me to the current design applying the pruning during shader graph construction.

First of all, let me assure you that I had initially invested a lot of time into a prototype applying the optimization at the shader graph rewriting phase, as we initially discussed. However I didn't find that all the necessary metadata was readily available at that late stage - instead it had to be passed down there from the construction phase, complicating the design.

• Shared analysis surface: The framework runs passes inside ShaderGraph::finalize against a constructed graph where node classifications (BSDF | CLOSURE | MIX | LAYER) and resolved connections are already available. The proposed NodeGraphTopology analyzer reimplements reference counting at the document layer with hardcoded category strings ("mix", "multiply", the kWeightedPbrNodes allowlist), most of which the ShaderGraph layer already gives us for free.

Yes, relying on the magic strings is unfortunate. These strings are inherited from Jerry's original OpenUSD implementation PixarAnimationStudios/OpenUSD#3525 Some, but not all of them do overlap the existing ShaderNode::Classification constants. To be fair, those constants are populated by parsing magic strings as well, e.g.

MaterialX/source/MaterialXGenShader/ShaderNode.cpp

Lines 288 to 297 in e8c12a2

	if (nodeDefName == "ND_layer_bsdf" || nodeDefName == "ND_layer_vdf")
	{
	newNode->_classification |= Classification::LAYER;
	}
	// Check specifically for the mix node
	if (nodeDefName == "ND_mix_bsdf")
	{
	newNode->_classification |= Classification::MIX;
	}

It would definitely be nice to centralize and share the code dealing with such strings, but it could be done by calling the shared parsing code from the two code paths, instead of moving both code paths to the same function.

Actually, I would argue that this is where the token (interned string) pattern would be a natural fit - this is what a compiler would do. Tokens would also help us avoid string comparison overhead and the risk of typos in other areas of the code. For example, the same string literals are compared against in ShaderNode.cpp and in ShaderGraphRefactor.cpp.

Speaking of classifications, wouldn't the node XML attribute value returned by the node's getCategory() API (e.g. mix) method be a cleaner way to derive ShaderNode::Classification than checking all the concrete nodedef names (ND_mix_bsdf etc.)? Also, some of the classifications seem not to be covered by the constants, e.g. MULTIPLY.

Other magic strings used by this optimization are uimin/uimax and the "optimizeable PBR nodes list". These are involved in determining which inputs on which nodes are "topological", i.e. create specializations/permutations. I think that ideally, these should be explicitly defined as metadata in MaterialX documents (or other sources exposed via the future Visitor Pattern), e.g. a topological="true" attribute. However, the current implementation is safe/conservative and I don't see how existing Shader Graph abstractions can help us here.

Regarding reference counting in NodeGraphTopology. First, I think it's not fair to say that it's implemented at the document layer. It queries the document for topology information, and could just as well query the future Visitor Pattern API for the same. I would rather call NodeGraphTopology a lightweight data structure that lives between the shader asset (such as a document) and the Shader Graph and makes it possible to avoid the creation of Shader Graphs as an optimization.

Second, NodeGraphTopology implements reference counting as an optimization, unlike removeUnusedNodes(), which which repeatedly traverses the graph from its outputs. The reference counts are populated once and then the topology is cached for the lifetime of the process. This optimization is easier to pull off at this early stage because the document is immutable from the perspective of codegen, unlike the Shader Graph.

• Natural overlap with PremultipliedBsdfAddRefactor: Both passes special-case BSDF mix nodes with weight inputs, and living side-by-side as refactor passes makes that sharing straightforward.

In my mind, this is another argument for shared utilities, but not necessarily for collocating the code.

• Composability: Passes run in sequence with removeUnusedNodes between them, so e.g. constant elision can clean up inputs that became unconnected after pruning, with no bespoke plumbing.

I would actually argue that the early pruning is even more trivially composable with the refactoring passes than if it were a refactoring pass itself. Early pruning simplifies the graph for further downstream optimization, saving work for the refactoring passes. If whole subgraphs are pruned early, refactoring wouldn't even have to run on them.

• API consistency: Adding a ShaderGraphRefactor pass would require no changes to createShaderNodeImplForNodeGraph, making this PR transparent to custom shader generators.

The permutation being passed to createShaderNodeImplForNodeGraph is a purely additive and backward-compatible change: NULL is the default, and passing NULL simply disables the optimization, which is already the case when the shader gen option is off.

The remaining benefit of your proposed approach is avoiding ShaderNode::create and initialize for pruned nodes. Do you have a sense of how much of the measured code generation time win comes from that specifically, vs. just from emitting fewer nodes downstream? If most of the win is downstream, then a refactor-pass form likely captures most of it while staying consistent with the framework that just landed.

In my mind, this is the main argument for performing the optimization as early as possible: it saves performance downstream. Creating Shader Graphs is a non-negligible cost - see, for example, the performance capture screenshot in #2820 Construction accounts for ~80% of total codegen time. So even a small percentage reduction in construction is significant in absolute terms.

On average, the savings are roughly split 50/50 between the shader graph construction and code emission stages, a few milliseconds each. I believe it's safe to say that it's always going to be strictly more performant than an equivalent graph refactoring pass. Some measurements on my machine:

Material	Total savings	Construction	Emission	Construction share
Standard_Surface1	-5.4ms	-1.9ms	-2.6ms	35%
ThinFilm	-5.3ms	-2.8ms	-2.3ms	54%
Default	-5.0ms	-2.7ms	-1.9ms	54%
Greysphere	-3.5ms	-2.4ms	-1.3ms	70%
Velvet	-2.8ms	-2.2ms	-0.6ms	79%
Car_Paint	-2.0ms	-0.2ms	-1.7ms	11%
Copper	-1.8ms	-0.1ms	-1.6ms	4%

Finally, I wanted to point out the optimization potential of early topology analysis beyond the current results:

1. The next step would be to proactively generate shader permutations for specific values of topological parameters, just like in @JGamache-autodesk's OpenUSD changes. This would not only generate shaders more optimized in terms of GPU rendering performance, but also show more dramatic optimizations of both codegen and shader compilation times, thanks to larger subgraphs being pruned.
2. Imagine a compound node's final codegen results being cached. If we have a mechanism of identifying unique node permutations just by the node's topological parameter values and codegen settings, we can completely skip codegen for a permutation that we've encountered before. This would be extremely beneficial for material editing workflows - we won't have to regenerate the shader for the whole graph just because a non-topological parameter value has changed. This cache could initially be implemented in memory and serialization could be added later.
3. This design could potentially be generalized beyond compound nodes, to arbitrary node implementations. A custom C++ node impl could then manage and cache its own permutations, driven by a well-encapsulated topological parameter interface.

[jstone-lucasfilm]

Thanks for this detailed reply, @ppenenko, and I'm convinced the construction-time approach you've proposed here is worth keeping.

I think the path forward is to decouple the two questions we've been bundling together. Where the pruning happens is a performance question, and your data settles it. How the analysis identifies topological inputs is a maintainability question we can resolve independently, and I like your proposal above: rather than collocating the code, let's factor the category/uimin/uimax string handling into shared parsing utilities that both ShaderNode classification and NodeGraphTopology call into, so we don't carry a second kWeightedPbrNodes-style allowlist that silently misses new BSDFs. Longer term, a topological="true" mechanism along the lines of your suggestion might allow us to drop the heuristics entirely.

A couple of separate coding issues that I'd also recommend addressing in this PR:

• Potential null dereference in getImplementation: createPermutation returns nullptr when no pruning can be applied at the call site -- e.g. when the graph has no topological inputs or when its topological inputs aren't pinned to 0 or 1 for this instance -- but permutation->getKey() is called before the if (permutation ...) null check.
• SHADER_INTERFACE_COMPLETE correctness: pruning bakes topology from the default/call-site weight value, but in the complete interface those weights are published as editable uniforms. Pruning a lobe the user intends to drive at runtime (e.g. raising a coat weight from 0) would then be incorrect. Could we gate the optimization to the reduced interface, or otherwise exclude inputs that get promoted to the editable interface?