What is the right approach to creating an UberShader using function_constants in Metal?

Question

I just learnt about function_constants in "What's new in Metal" video from WWDC 2016 and it mentioned UberShaders quite a few times. I want to create a fragment uber shader that can be used for different kind of passes, like simplePassThrough, defferred etc. Below is how I want to use it.

constant int passType [[function_constant(0)]];
constant bool simplePassThrough = (passType == 0);
constant bool forwardShading = (passType == 1);
constant bool deferredShading = (passType == 2);

fragment FragmentOutStruct UberFragmentShader()
{
FragmentOutputStruct frgOut;
if (simplePassThrough) {
    // Update frgOut
} else if (forwardShading) {
    // Update frgOut
} else if (deferredShading) {
    // Update frgOut
}
return frgOut;
}

Is this the right approach here? Will my final compiled MTLFunction see too many branches if I use this approach?

Answer 1

This is a legitimate use case for function constants, and will have no branching cost at runtime. This is because the compiler will eliminate code it determines can never be executed (e.g., because it is equivalent to if(false) { ... }).

Answer 2

Yes, you're on the right track. (As @warrenm noted already. But to expand on his answer a bit...)

Your example is essentially the same as Apple shows in the WWDC16 session introducing function constants: your "branches" are all directly derived from a function-constant value, which means that the shader compiler can (when you build your app) generate IR variants for each of the possible paths through your code that depend on function constant values.

Here, you're passing an int to the shader, but that doesn't mean it has to compile 2³² shader variants — the compiler can do some static analysis and see that there are four possible code paths based on that value (0, 1, 2, and anything-else, the last of which just elides the if statements entirely and returns frgOut).

At run time, the Metal framework determines which of the four shaders to send to the GPU based on which value you pass for the constant, so there's no branching in the shader / on the GPU. For example, if you pass a value of 1, you're running a shader that essentially looks like this:

fragment FragmentOutStruct UberFragmentShader() {
    FragmentOutputStruct frgOut;
    // Update frgOut per `if (forwardShading)` chunk of original shader source
    return frgOut;
}

And as you can see, there's no branching in that shader.