Injection of constant values in shaders

[k2d222]

in the MVP task list #54 we mentionned code injection and came to the conclusion that overridable constants were sufficient, at least for M1. I'm just starting the discussion for future versions.

What is value injection

Just being extra clear, we are talking of this:

override some_value: u32;

fn very_hot_path() {
    if (some_value > 1000) { impl_a(); }
    else { impl_b(); }
}

we made some_value overridable because

• it depends on knowlege available at runtime, e.g. the hardware specs; or
• the code is used in several shader variants with different purposes and constraints.

In principle and in practice, the GPU driver will eliminate the branch and one of the impl_[a/b] functions. It generates a shorter and more efficient program to upload to the GPU.

TL;DR value injection is when the linker controls a value that can change the program behavior.

What issues with the override solution

1. Code size
The WGSL output will contain all code reached by branches dependent on overridable constants. Essentially overridables prevent Dead Code Elimination. It can result in larger source code which is not optimal for transfer time on the web and driver parsing/compilation/optimization time at runtime during createShaderModule and create[Compute/Render]Pipeline.

2. Overrides are Globals (not scoped)
See #43. We don't know yet how to deal with overrides defined in packages or submodules. There is currently no way to set an override value from shader code.

3. Overrides don't cover all constant use-cases
See #54 (comment). Overrides can't be used in bindings (@group/@binding); as the number of elements of nested arrays; and a couple more edge-cases.

4. Preprocessing shader variants
A common technique in games is to cache shader variants so they don't have to be computed each time the game is run. There is currently no way to force an override to take a fixed value (turning it to a const) which would allow WESL Dead Code Elimination to run.

5. Conditional translation based on values
Something not covered by overrides, we may want to use values in conditional translation in cases where different values cause a different host<->shader interface. There might be other use-cases for that too. Example:

@if(instance_count < 100) // it fits in a uniform buffer which is typically faster than storage.
@group(0) @binding(0) var<uniform> my_data: array<Data, instance_count>;
@if(instance_count >= 100)
@group(0) @binding(0) var<storage> my_data: array<Data, instance_count>;

What possible solutions

The current workaround is to generate a "virtual" shader module at runtime and import the constants in other modules. Solve issues 3, 4, 5.

A cleaner mechanism is to allow the linker to override constant declarations (and perhaps demote override declarations to constants). We may want to allow the "uninitialized constants" extension to force the linker to provide a value. It's basically reimplementing the concept of override with constants. Solve issues 1, 2, 3, 4.

We could provide a mechanism for shaders to fix to a constant (or forward the override) for overrides defined in dependent modules. In this case we'd want to think about what happens when the same module is used twice with different values. Solve issues 1, 2.

[mighdoll]

excellent summary!

I don't recall the discussion, but this #54 (comment) makes me think there was a reason to want const injection for M1 if possible.

And wesl-rs has the virtual module approach now, right? I'm not pushing for the feature, but we could add something like that to wesl-js pretty easily if it's clear what to do.

[mighdoll]

So the recommendation is to do a virtual shader module for people to play the capability of const injection in M1, since it seems like a few people need it, but make clear in the docs that we're going to do look at doing some cleaner later.

Sound right?

Seems straightforward implement, knock wood. What name should the virtual module have?

[k2d222]

sounds right! I guess the virtual module should be named when it is provided to the linker. This is roughly how I do it:
linker.add_virtual_module("foo::bar", "some wesl string...")

[mighdoll]

@k2d222 suggests we unify on constants as the virtual package name for the interim virtual package solution.

[mighdoll]

added constants option to link() in wesl-js:

  /** Host (ts/js) provided wgsl constants.
   * Users can import the values from wesl code via the `constants' virtual library:
   *  `import constants::num_lights;` */
  constants?: Record<string, string | number>;