improve shader import model

[robtfm]

Objective

operate on naga IR directly to improve handling of shader modules.

• give codespan reporting into imported modules
• allow glsl to be used from wgsl and vice-versa

the ultimate objective is to make it possible to

• provide user hooks for core shader functions (to modify light behaviour within the standard pbr pipeline, for example)
• make automatic binding slot allocation possible

but ... since this is already big, adds some value and (i think) is at feature parity with the existing code, i wanted to push this now.

Solution

i made a crate called naga_oil (https://github.com/robtfm/naga_oil - unpublished for now, could be part of bevy) which manages modules by

• building each module independantly to naga IR
• creating "header" files for each supported language, which are used to build dependent modules/shaders
• make final shaders by combining the shader IR with the IR for imported modules

then integrated this into bevy, replacing some of the existing shader processing stuff. also reworked examples to reflect this.

Migration Guide

shaders that don't use #import directives should work without changes.

the most notable user-facing difference is that imported functions/variables/etc need to be qualified at point of use, and there's no "leakage" of visible stuff into your shader scope from the imports of your imports, so if you used things imported by your imports, you now need to import them directly and qualify them.

the current strategy of including/'spreading' mesh_vertex_output directly into a struct doesn't work any more, so these need to be modified as per the examples (e.g. color_material.wgsl, or many others). mesh data is assumed to be in bindgroup 2 by default, if mesh data is bound into bindgroup 1 instead then the shader def MESH_BINDGROUP_1 needs to be added to the pipeline shader_defs.

[MavethGH]

I would be happy to help develop and maintain the naga_oil crate.

[robtfm]

I would be happy to help develop and maintain the naga_oil crate.

thanks, i appreciate it - worth noting though that even if it's part of bevy it can still be imported and used as a standalone crate if you find it useful for you outside of bevy. i am mostly concerned with bevy since that's where i want to use it, so if @cart wants to own it that's totally fine with me.

[MavethGH]

I actually have a use for it outside of Bevy, but I think that integrating it with Bevy would be most useful to the most people.

[robtfm]

added overrides, so functions can be replaced dynamically, and added an example shader_material_override to demonstrate.

todo still: provide a mechanism for general overrides (so you can override the default lighting functions - this only allows for overriding on a per-material basis currently).

[robtfm]

and added global/default overrides, with pbr_global_override example

[cart]

Gave this a very quick once over (this will take a lot of time to review properly), but I now have a decent understanding of the approach taken. Some quick thoughts:

1. This is awesome and does go generally in the direction I've wanted for Bevy's shaders. Very nice to know that we can rely on Naga here (although I'll admit the naga side ended up being more complex than I expected).
2. naga_oil covers cases that are "core" enough to the bevy shader experience that I would want that to be under the bevyengine org umbrella and go through our normal governance process (even though that is slower). Controversial changes there are definitely controversial changes within the context of bevy. Although I think theres a reasonable argument to be made to keep it separate from the rest of the bevy codebase (ex: exist in its own repo in the bevy org). We can discuss this.
3. Forcing shader import consumers to define their own PascalCase aliases for modules feels wrong. Module authors should be defining the "canonical names" used, not module consumers (at least by default ... if there are naming collisions the as syntax provides a nice way to resolve that). And PascalCase makes it seem like modules are struct types. Given that we use Rust-like snake_case::module::syntax elsewhere, this feels confusing. At the very least we should be consistent, and I personally think snake_case::module::syntax is what we should use.
4. We probably want to support something like this syntax as well: #import bevy_pbr::skinning::{skin_model, skin_normals} (rather than "whole module" imports). Is that compatible with this implementation? How hard would that be? Not saying we need to support it out of the gate, but it would be good to know how possible it is, if this design can accommodate it, and whether we should make that a part of our plan.
5. Allowing arbitrary function overloads feels a bit dangerous / allows people to take dependencies on internal implementation details. I do like that it lets people insert themselves into pre-built pipelines, but we might want to be a bit more principled about it. We should consider only allowing overriding of functions that have been labeled as "overridable". We should also look in to how other engines expose custom lighting logic in shaders. Its possible that a simple function overloading model isn't the best fit. We might need a bit more structure and/or magic. I don't have strong opinions here yet, but this deserves careful design thought.
6. Given that this removes the ability to import "partial struct fields" / spread them into functions or structs, I think we should discuss our plans for "hiding" the complexity / boilerplate of specifying the default VertexOutputs from shader authors. The current approach is a bit weird / arcane, but defining this manually (like we do in this pr) is worse. If we choose to take the path in this pr, we should make sure it can accommodate this scenario in a reasonable way.

[robtfm]

2. naga_oil covers cases that are "core" enough to the bevy shader experience that I would want that to be under the bevyengine org umbrella and go through our normal governance process (even though that is slower). Controversial changes there are _definitely_ controversial changes within the context of bevy. Although I think theres a reasonable argument to be made to keep it separate from the rest of the bevy codebase (ex: exist in its own repo in the bevy org). We can discuss this.

thanks for looking it over. i'm very happy for it to be under bevy org or part of bevy core as you prefer, and happy it would benefit from bevy's review process as well.

3. Forcing shader import consumers to define their own `PascalCase` aliases for modules feels wrong. Module authors should be defining the "canonical names" used, not module consumers (at least by default ... if there are naming collisions the `as` syntax provides a nice way to resolve that). And  `PascalCase` makes it seem like modules are struct types. Given that we use Rust-like `snake_case::module::syntax` elsewhere, this feels confusing. At the very least we should be consistent, and I personally think `snake_case::module::syntax` is what we should use.

the as syntax is optional, if you don't specify it then the original/canonical name stands. i chose to capitalize the modules just to get different colours in the editor but i agree it would be better to be consistent with rust and i'll change that in the bevy wgsl files. (anyway if this goes through i expect we could tailor the syntax hilighting for it in the future.)

4. We probably want to support something like this syntax as well: `#import bevy_pbr::skinning::{skin_model, skin_normals}` (rather than "whole module" imports). Is that compatible with this implementation? How hard would that be? Not saying we need to support it out of the gate, but it would be good to know how possible it is, if this design can accommodate it, and whether we _should_ make that a part of our plan.

this could work easily if they are still qualified at point of use, but i guess that's not what you're intending.

in the current impl there's no lexing done outside of naga, it's all regex-based - we prefix all the members of a module with a sentinel string and the (base32 encoded) module name to make a "namespace", and then search/replace module_name:: in the calling scope with the encoded prefix (and otherwise deny use of the sentinel string to ensure there's no accidental use of module items). this works for now because :: can't otherwise appear in wgsl or glsl.

if we want to import individual names and use them without qualification it may need more care to make sure we only modify exactly what we should. it's probably fine to just use a regex with initial whitespace but i haven't fully thought it through. at worst, ::skin_normals would work.

5. Allowing arbitrary function overloads feels a bit dangerous / allows people to take dependencies on internal implementation details. I do like that it lets people insert themselves into pre-built pipelines, but we might want to be a bit more principled about it. We should consider only allowing overriding of functions that have been labeled as "overridable".

marking functions as overridable is straight-forward. personally i would prefer not to have artificial restrictions on the injection points i can choose, but i'm more familiar with the pbr internals than many users and maybe the support burden would be an issue. but perhaps just documenting "safe" entry points for user overrides would be enough?

5. ... We should also look in to how other engines expose custom lighting logic in shaders. Its possible that a simple function overloading model isn't the best fit. We might need a bit more structure and/or magic. I don't have strong opinions here yet, but this deserves careful design thought.

one thing to bear in mind: the "API" that we would be fixing through overridable function signatures is not quite as restrictive as it seems initially - we can also take the vertex/fragment input and store it to a global variable, so that overriding functions have access to whatever data they need. at the extreme we could also do this with other derived data and reduce the overridable function signatures down to practically nothing. this might be inefficient (i'm not sure how good the driver dead-code analysis is wrt calculations that get stored to globals but never read) but the purge module could fix it if it's a problem.

definitely interested to hear other people's thoughts here though.

6. Given that this removes the ability to import "partial struct fields" / spread them into functions or structs, I think we should discuss our plans for "hiding" the complexity / boilerplate of specifying the default VertexOutputs from shader authors. The current approach is a bit weird / arcane, but defining this manually (like we do in this pr) is worse. If we choose to take the path in this pr, we should make sure it can accommodate this scenario in a reasonable way.

note it is only manually defined for vertex output (there's a shared struct for frag input) ... it is grim though. the current restriction seems to be that structs can be used in entry-point IO, but only one level deep. i'm not entirely sure where the restriction comes from. ideally in future wgsl (or naga) will allow us to have entry-point IO where only the leaves of a heirarchy need to be located.

[robtfm]

structs can be used in entry-point IO, but only one level deep

i've moved @bulitin(position) into the MeshVertexOutput structs so we can use them directly as the vertex stage output and fragment stage input, without needing to embed them within another struct layer.

this seems good to me ... i'm not sure if there's a reason why we are not doing this already.

edit: but of course it still can't be so easily extended by authors of custom vertex shaders ... oops

[robtfm]

further to the above, the struct depth limitation for entry point IO comes from the WGSL spec. i've tried briefly to extend the naga internals to accept nested structs but have hit a point where my lack of spirv knowledge is blocking me and i am not really prepared to learn it well enough to debug. i think it would be possible though since spirv seems to use globals for stage IO.

but i think it's not so bad as it is. depending on if you use custom vertex shaders for

• generating extra outputs for use in custom fragment shaders: if you don't use the standard frag shader then you don't care what the standard frag shader uses, and so you should define your own IO struct with exactly what you need.
• using funky vertex input (e.g. vertex pulling) to pipe into the standard frag shaders: then your vertex stage must produce exactly the standard output (and can use the standard IO struct), else it won't work anyway.
• producing different output, and then calling into the standard fragment shaders: then you will need to construct a MeshVertexOutput in your top level frag shader to call into the standard frag shader, or a PbrInput to call into pbr_functions::pbr. you'll need to ensure your frag shader can derive everything for the pbr entry point you use, but it won't necessarily need all the standard outputs from the vertex shage anyway.

[cart]

The custom_gltf_vertex_attribute and texture_binding_array examples are also broken (unresolved imports).

[cart]

This is looking great! A huge step forward in the usability and modularity of our shader code. Now that we have these capabilities, I think our next steps should be to consolidate imports and improve the UX of writing shaders / consuming built-in Bevy functionality. I think we have too many different modules right now. Some thoughts:

• Combine things like mesh_bindings, mesh_functions, mesh_vertex_output, etc into mesh. Now that we have granular imports we shouldn't need the granular modules in most cases.
• Consider creating preludes to make it super simple to do common things.