SPIR-V runtime Patching

Date: 11.10.2023

So, did you ever want to change parts of you shader code at runtime? Lets say change how a BRDF behaves, or what a GPU-Simulation does based on some user input? Then this article (and the resulting program) are for you!

Use case

So called modern graphics APIs currently have two compile times when it comes to shaders. One is offline, this compiles your source code to something like SPIR-V or DXIL, and one is at runtime, usually when you create a pipeline and supply that Intermediate Representation (IR) to your graphics API and, subsequently to your graphics driver.

A cool thing about this is, that we can transform the IR before actually sending it of to the GPU. Which, in theory would allow for stuff like linking, which then again makes problems like shader permutation easier to handle.

Another use case is shader specialization, which is currently done on the API level and heavily restricts *what** you can specialise (usually only values). So in case you want to specialise a function at runtime, no luck here, you’d have to anticipate all possible outcomes.

This is where the spv-patcher comes in, it allows you to actually change your IRs code, by loading it, analysing it and then allowing you to patch certain locations.

As a side note, this is directly related to my algae experiment, where I implemented a raw, way less checked version of this.

Why not IR-Format XY?

While SPIR-V has its shortcomings, it is uniquely positioned in the opensource IR landscape as below language dependend but above driver level. I think this image shows that quiet nicely:

SPIR-V's position in the OOS IR landscape.

As you can see there are also a lot of translators to other shader-IR formats like DXIL or MSL, so in theory this SPIR-V patching also enables you to change your DXIL in some way.

For more serious transformation however I also include SPIR-T which is a research IR that is closely related to SPIR-V. Currently it is used in the static-patching path to do in-memory module linking. I used spirv-tool’s linking, but that was too slow in practice.

Implementation

The implementation is somewhat structured like a compiler. We load a source file (SPIR-V or SPIR-T) and apply all modification passes. Afterwards we store that back in a file, or send that in-memory SPIR-V representation directly to the API. In theory the patcher can implement any SPIR-V  − > SPIR-V transformation. For instance there is also one that finds and patches non-uniform data access. This is a kind of bug that is really hard to spot in the source code but easy-ish on the IR level.

At runtime we can load a file as a template and then either patch some other SPIR-V code (which we call static-patching, its similar to static linking), or by executing a user defined function on that portion of the code, which is called dynamic patching. Dynamic patching (as opposed to static-patching) could analyze the current code, or emit different code depending on some runtime state.

Usage

For an academic example lets take some compute shader. In its unmodified for it contains a function calculate in the shader:

pub fn calculation(a: u32, b: u32) -> u32 {
    a + b
}

We can now use the Patcher to transform a template shader module like this.


let template_module = spv_patcher::Module::new(spirv_binary)?;
let modified_module = {
    let patcher = template_module.patch()
        .patch(patch_function::DynamicReplace::new(
            patch_function::FuncIdent::Name("calculation".to_owned()),
            |builder, sig| {
                //we use IMUL to multiply the two arguments and store that into a new id that is returned
                let a = &sig.parameter[0];
                let b = &sig.parameter[1];
                assert!(
                    (a.1 == b.1) && (a.1 == sig.return_type),
                    "Types do not match!"
                );
                //add imul
                let res_id = builder.i_mul(sig.return_type, None, a.0, b.0).unwrap();
                
                //now assign result to return instruction
                let _ = builder.ret_value(res_id).unwrap();
                Ok(())
                },
            ));
            
    patcher.assemble()?
};

This will search for a calculation function and then call the user defined modification function with the supplied signature information. The patch (in this case DynamicReplace) should (and in this case does) check, that the patched code is valid. However, it is also trivial to run spirv-val on the generated code.

Validation

Right now we can use SPIR-V’s spirv-val to validate the outcome. However, in practice this does not catch all errors, so you’d be good to check your vulkan validation layers before actually trying to run the new pipeline. In a perfect world you’d use some kind of compiler to generate the new SPIR-V code which would make sure that everything is legal.

Conclusion

Now we are able to modify our shader code on a SPIR-V level. Essentially this gives us the opportunity to transform user generated content into actual shader code. This will come in handy later for me, when I continue working on algae and my signed-distance-function language. I will be able to transform a SDF into actual SPIR-V code, that is then runtime patched into existing shaders. So you’d get the performance as if you’d write a specialised shader, just for that function like on shadertoy.