Shader Generator ND_mix_bsdf optimization

[ld-kerley]

Taking inspiration from #2459, which is a significant win in terms of HW shader performance. Instead of optimizing the node graph directly, here we apply the same idea, but programmatically at shader generation time.

If we meet the following condition.

"A mix node that has BSDF type inputs, and if those two nodes have weight inputs themselves."

Then we can replace the mix node with an add and modify the weight inputs, by multiplying by the mix input, and the inverse of the mix input.

Old nodegraph

---
config:
      theme: redux
---
flowchart TD
    weightA{{"weight_A"}}
    weightB{{"weight_B"}}
    mix{{"mix"}}
    BsdfA["BSDF_A"]
    BsdfB["BSDF_B"]
    Mix["Mix"]
    weightA --> BsdfA
    weightB --> BsdfB
    BsdfA --> Mix
    BsdfB --> Mix
    mix --> Mix

Loading

New nodegraph

flowchart TD
    weightA{{"weight_A"}}
    weightB{{"weight_B"}}
    mix{{"mix"}}
    Invert["Invert"]
    MultA["Multiply"]
    MultB["Multiply"]
    BsdfA["BSDF_A"]
    BsdfB["BSDF_B"]
    Add["Add"]
    mix --> Invert
    weightA --> MultA
    mix --> MultA
    weightB --> MultB
    Invert --> MultB

    MultA --> BsdfA
    MultB --> BsdfB

    BsdfA --> Add
    BsdfB --> Add
    

Loading

This allows the BSDF nodes to early out if their contribution is not necessary in the mix.

Taking this approach has a couple of advantages:

• We programmatically match all instances of this possible optimization, not just the few that are hand crafted.
• We can leave the uber shader nodegraphs, in their original state. The use of the mix node better conveys the idea behind the nodegraph.
• Its possible that some backends might be able to optimize the original graph more effectively than the current hand-crafted representation.

[jstone-lucasfilm]

I'm very excited about the future of this work, @ld-kerley, and my main recommendation would be that we consider it a new investment of research and engineering effort for MaterialX 1.39.5, so that we have the time to validate this approach for all MaterialX shading models, from both a performance and visual parity perspective.

Once we're confident that this gives us the same or better performance as #2483, #2459, and #2467, we can replace all of those manual optimizations with this automated approach, with the expectation that other shading models (even custom studio models) will benefit in the same way.

[kwokcb]

Hi @ld-kerley ,
What are your thoughts on integrations that don't use shader generation and hence won't ever see / use this optimization.
It feels like this can be a pre-processing step or even a stand-alone optimization utility.
It can still be integrated here to automatically perform these (or other optimizations) but is not directly tied to code generation.

[fpsunflower]

I'm not super familiar with the details of codegen here, but can you explain what is different between the two implementations?

What can the add/mult nodes do that the mix node cannot ?

[ld-kerley]

Related to #2480

[ld-kerley]

@fpsunflower Its not so much a feature of codegen, as much as function of the HW shader language backends here. The implementation of the BSDF nodes that accept a weight input are written to early out if the weight is zero. In the mix case the weight is often not zero, but unfolding the mix math and moving the weight to be combined with the weight input to the BSDF node does often end up with zero weights in the BSDF nodes. In a language like OSL I would expect the OSL optimizer to be able to fold away the unused branches if the mix value is set to 0 or 1.

@bernardkwok today anyone not using codegen wouldn't have access to the optimization in this form. We could pre-process nodegraphs during build with this logic if we wanted. Personally I'm not a fan of optimizing the nodegraphs, when we've only validated it a performance improvement in subset of the backend languages.

I actually wonder if we should put this optimization behind a code gen option switch to make it easier for downstream consumers to profile either way and make their own decision if they want it.

If we grow the number of these sorts of optimizations, which I hope we do, there are a number of other really valuable ones we used at Imageworks, then I can imagine a world where perhaps downstream consumers could run "code gen" to generate an optimized version of the graph - but not go all the way to shader code. I think this might actually be an important part of any future overhaul of the code gen system, to centralize the graph based construction and optimization, and allow an "exit-point" here. But this is all me just shooting from the hip.

[fpsunflower]

Ah I see now -- the issue is the mix happens after the BSDFs are already evaluated, so there's no chance to short-circuit them. In the new layout they can be. This makes sense.

[ld-kerley]

Just for some more context - Slack conversations with Kai Rohmer here, suggest that at least in MDL this optimization might not be ideal. So we should refactor things here to be constrained to the HW shading languages for now.

[ld-kerley]

I've added a GenOption based option to toggle this behavior on and off. It's also enabled for the GLSL/MSL code generators used in MaterialXView and MaterialXGraphEditor.

Before this is merged the specialized versions of the graph should be removed from (libraries/bxdf/genglsl) - but I'm leaving them here for now to allow easier A/B testing between the modified nodegraph approach and the programmatic approach in this PR.

[jstone-lucasfilm]

@ld-kerley Following up on our last MaterialX TSC meeting, I'd love to bring this PR back into focus again, so that we can make it available as a platform for #2680 and further generator optimizations in the future.

When you have a moment, would you mind resolving the handful of merge conflicts in MaterialXGenShader?

I can then take the next step and compare the performance of these latest generator optimizations against the manually optimized shading model graphs, and we can work through any minor differences that may remain.

[ld-kerley]

I rebased and pushed - and re-ran the test suite locally - but didn't have time to validate the optimizations again - I think they should still be valid, but it would be really interesting to find a way to add a unit test to verify they are applied correctly.