naga - compile with context

[stefnotch]

Is your feature request related to a problem? Please describe.

Background - naga_oil
I have been working on naga_oil, which is a WGSL preprocessor, for a while. One important feature that naga_oil implements is shader imports.

For shader imports, we independently compile shaders using naga, and then correctly link them together. This works nicely, but comes with the limitation that we need to generate "header files" for the naga compiler.
For example, when we want to compile

// cat.wgsl
#import foo::Bar;

struct Cat {
  a: Bar;
}

we first look at the foo file, and generate code along the lines of

// cat.wgsl
// foo.wgsl header
struct Something { hi: f32 }
struct Bar { baz: f32, henlo: Something }

struct Cat {
  a: Bar;
}

And this scales really badly. As the import tree gets larger, we can end up with massive shaders that end up including almost everything else, just in case the compilation needs it.

So we're doing

1. Compile foo.wgsl to a naga module
2. Extract a header, and generate WGSL code
3. Add the header to cat.wgsl, and compile that to a naga module
   and repeat that, until we've compiled all shaders

Describe the solution you'd like

It would be much nicer if we could do

1. Compile foo.wgsl to a naga module
2. Use the existing naga module to compile cat.wgsl

This request is asking for a way of pre-populating a naga module with some entries. I think this would be a useful building block for other tools, such as naga_oil.

Describe alternatives you've considered
Writing our own WGSL parser for performance, or accepting O(n^2) compile times hit for Bevy's shaders.

Additional context
There might also additional use cases for this, such as incremental compilation. But that's a different story.

[stefnotch]

I would, of course, be willing to dive into naga and provide a PR for this. Pointers on roughly where to look are always appreciated.

[teoxoy]

This is the general flow when you call parse_str():

wgpu/naga/src/front/wgsl/mod.rs

Lines 39 to 41 in 4902e47

	let tu = self.parser.parse(source)?;
	let index = index::Index::generate(&tu)?;
	let module = Lowerer::new(&index).lower(&tu)?;

The parser should be able to handle missing global items just fine and the resulting TranslationUnit has information about all declarations which includes the missing ones. The Lowerer already does type checking and has additional state to do the lowering so providing an initial Module won't work unless the state of the Lowerer is also restored - I haven't looked into how complicated that would be.

For a first attempt at this, I would:

• call parser.parse to get the TranslationUnit
• look for missing items
• patch-up the TranslationUnit with the missing items
• lower to get the Module

This would mean that you have to hold onto all original TranslationUnits but should be better than the current approach and not be as complicated as modifying the Lowerer (which can still be done later).

[stefnotch]

Thank you! I've now had a chance to look at this, and there are a two different options regarding how to implement it.

• The approach where one can keep a TranslationUnit around, and call a function to add globals to it. Which internally patches up the TranslationUnit.
• Using the lowerer's globals. When lower is called, it starts out by creating a GlobalContext, which has a globals: &'temp mut FastHashMap<&'source str, LoweredGlobalDecl>,. From what I can tell, whenever it encounters something unknown, it consults this list. For example, for a function call, it consults the GlobalContext

  wgpu/naga/src/front/wgsl/lower/mod.rs

  Line 1945 in 9f74b15

  match ctx.globals.get(function.name) {

I think the second option could be even simpler to implement and maintain. Should I take a stab at it, or could there be any major issues that I'm overlooking?

[teoxoy]

You can try the 2nd approach but I think the first will be more straight-forward (it's not clear to me what needs to happen for the 2nd approach to work without actually doing it).

[stefnotch]

I found an interesting issue. In WGSL, one can redefine built-ins at any point. For example:

Like compute.toys does, I want to import/inject a struct with the mouse position

struct Mouse { pos: vec2f }
@group(0) @binding(0) var<uniform> mouse: Mouse;

into a shader

fn main() -> vec3 {
  return vec3(mouse.pos.xy);
}

That would work. But the moment the shader author adds an alias

fn main() -> vec3 {
  return vec3(mouse.pos.xy, 1.0);
}

alias vec2f = vec2<i32>

the provided mouse struct would break. Despite it being part of the context, and not part of the shader source code.
With a similar approach, one could alias the cos function and break the "compile with context" in other ways. Or do that any other predefined function.

I think, as with Rust's macros, that this should be prevented. As in, "compile with context" should be hygienic. The context can only add definitions for subsequent code, it cannot be changed by subsequent code. This should be relatively do-able, one just has to always refer to the resolved built-in types in the provided context.

[teoxoy]

I think, as with Rust's macros, that this should be prevented. As in, "compile with context" should be hygienic. The context can only add definitions for subsequent code, it cannot be changed by subsequent code. This should be relatively do-able, one just has to always refer to the resolved built-in types in the provided context.

Why should imports be compared to rust macros instead of rust imports? The examples above would be a non-issue if you had to specify what you wanted to import. Even in rust you can do type f32 = i32;.

There are also plans to add a wgsl namespace to the WGSL spec which would contain all built-ins.

[teoxoy]

@stefnotch I talked with @jimblandy about this and we think that the approach in #5713 (comment) is the way to go.

[stefnotch]

@teoxoy Thank you! Then trying out the approach from said comment will be one of the next things to try out for me :)

Currently, I am also exploring the alternative variant of writing a WGSL parser from scratch. Essentially, that gives me a chance to try building a minimal implementation of "WGSL imports". Afterwards, I can take those ideas and figure out how to combine that with naga.