Move ShaderCache shader defs into PipelineCache

[Shfty]

Objective

• NO_STORAGE_BUFFERS_SUPPORT, NO_ARRAY_TEXTURES_SUPPORT and SIXTEEN_BYTE_ALIGNMENT are injected by ShaderCache::get, which occurs after pipeline specialization, and therefore after Material gets the chance to inspect or mutate them in Material::specialize.

• Fixes Material::specialize is provided an incomplete set of engine-sourced shader definitions #8190

• Related: Allow SPIR-V shaders to process when shader defs are present #7772 lifted an incidental restriction preventing the use of SPIR-V in Material, but full access to engine-sourced shader definitions in Material::specialize is still an outstanding issue.

Solution

• Allow Material implementors to access the aforementioned defs by moving them into PipelineCache, and extended the existing pipeline specialization machinery to pass them down the call chain from Specialized*Pipelines::specialize.

---

Changelog

• PipelineCache::base_shader_defs added (populated from the render device in PipelineCache::new), and exposed via a new PipelineCache::base_shader_defs() by-clone accessor
• shader_defs: Vec<ShaderDefVal> param added to SpecializedRenderPipeline::specialize, SpecializedComputePipeline::specialize, and SpecializedMeshPipeline::specialize
• Defs are cloned from PipelineCache in Specialized*Pipelines::specialize and passed into the respective trait methods
• Existing implementors have been modified to use the new trait param instead of constructing their own Vec<ShaderDefVal>

Migration Guide

This is a breaking change.

Specialized*Pipeline implementors should use the newly-provided shader_defs vector instead of constructing their own, extending it with new definitions as necessary, and ensuring that it gets passed into the appropriate pipeline descriptors for access in downstream code.

[superdump]

One problem I see with this approach is that it could be easy to forget to initialise one’s shader def Vec with these shader defs. It feels perhaps a little easy to do the wrong thing at the moment.

In that sense it might be easier to do the right thing if the shader defs were passed into the specialisation function. What do you think @Shfty @robtfm ?

[Shfty]

@superdump Yeah, I agree - passing shader defs in from outside would remove the footgun of soft-requiring the user to magic them up from inside their own impl. If API change is acceptable, I'm happy to go through and rework this appropriately.

Also, now #7771 has been merged, should we consider restoring the ShaderProcessor def count check as part of this PR? I'm not sure of the concensus on how valuable the previous "must have no defs regardless of format" semantic was, but this would be a good opportunity to bring it back if it's desirable.

[Shfty]

@superdump @robtfm I've reworked this PR to make PipelineCache the single source of truth for these platform-level shader defs, propagating them down the specialization chain via newly-added trait parameters.

Overall, I'm much happier with this approach; it slots comfortably into the existing data structure, centralizes the defs in question, avoids adding unnecessary fields to Specialized*Pipeline implementors, and avoids the footgun of requiring consuming code to call an arbitrary method in its construction routines.

As such, this is now an explicitly breaking change, which will better signal the need for migration to existing implementors.

Testing-wise, I've gone through the rendering examples and noticed nothing obviously amiss, and bevy-rust-gpu functions correctly with this patched into its development fork.

I've updated the top post with the appropriate outline, changelog and migration guide.

[robtfm]

i haven't been through in detail but i wonder if you have a feeling for how this interacts with what i proposed in #5703 -
in that pr, we can add shader defs to importable (wgsl or glsl) shader modules, so that when you import a module into your shader you get its defs as well. then (i didn't do this in that pr yet but) things like NO_STORAGE_ARRAYS would be associated with a platform.wgsl module that any shader that needs it could import.

i suspect this would solve your problem as well since those platform defs then woudn't be present in any spirv shaders since there's no mechanism to use imports with spirv anyway.

[Shfty]

@robtfm I believe it's orthogonal to the intent of this PR:

The motivation behind #7772 was to remove the blocker caused by ShaderProcessor rejecting SPIR-V shaders in the presence of shader defs. Now that's merged, SPIR-V users don't have to worry about manually clearing defs before they reach ShaderProcessor, but still need a way to access them during specialization.

With this PR, SPIR-V users can receive the previously-unavailable defs in the VertexState and FragmentState provided to Material::specialize, and use them to apply appropriate pre-compiled / pre-permutated entry points (i.e. my_fragment_shader__buffers_storage__textures_array__tonemap_in_shader vs my_fragment_shader__buffers_uniform__textures_single__tonemap_external).

(For a more complex practical example, bevy-rust-gpu parses those defs into a JSON file, which gets watched by an external SPIR-V compilation daemon, and read by a #[cfg()]-like proc-macro in the compiled shader crates to generate entrypoint variants. These are then output to a SPIR-V + metadata file, hot-reloaded as a bevy asset, and applied to a material.)

If the primary source of truth for those shader defs were to move into a WGSL module exclusively accessible via #import by textual shader languages, it would prevent SPIR-V users from specializing their shaders in a manner analogous to the preprocessor's #ifdef / const substitution features.

[Shfty]

(Also, I'm realising that much of this context is absent from the top post, having mainly been discussed in Discord. Will edit to amend.)

[robtfm]

thanks, that context makes it clearer for me. since i'm late i'm going to summarize to make sure we are on the same page

• shader defs consist of platform/device defs (NO_STORAGE_ARRAYS) and arbitrary user defs (TONEMAP_METHOD_XYZ)
• all these defs are used in the preprocessor within the shader cache to customise non-spirv shaders (to handle availability of features like storage arrays, and to statically select other custom functionality)
• also in non-spirv context, shader defs can be defined in shaders via the '#define' directive, so the full set of shader defs is not knowable until the shader is built by the shader cache (since #imports can be conditional on defs as well), so they can't really be used outside of the shader scope. in the current rust code the only read of the defs is in the preprocessor itself, and (now removed?) for spirv checking that there are no defs

if i understood correctly, what you're suggesting is

• move the platform/device defs into the pipeline cache so that they can be passed to Pipeline::specialize
• access the defs in Pipeline::specialize in order to select a spirv entrypoint based on them
• continue to handle non-spirv defs in the shader cache's preprocessor, just sourcing them from the pipeline cache instead

so in the non-spirv case, we would still have a variant being "selected"/built by the shader cache + preprocessor based on what defs are passed in from pipeline cache plus specialize, and in the spirv case we have the variant being selected by the pipeline itself, based on the (platform-subset-of) defs generated by the pipeline cache, and the defs that are output from specialize are ignored.

it feels pretty convoluted to me to have the usage location of the defs vary like this depending on the shader type. in particular i think it would make on-the-fly compiling / caching of shaders (from naga to spirv) harder in the future, since pipelines would need to know ahead of time what "style" of shader they are using, where currently they just do their specializing job regardless of the shader format. it would also maybe mean you can't share spirv shaders so easily between pipelines by just varying the defs that the pipeline outputs.

i think the better answer is to use the defs in the shader cache itself to select the variant somehow for spirv, similarly to what we do for non-spirv. i don't know exactly how to do that - perhaps the Shader::Spirv variant could have a lookup based on entry point plus shader defs that the shader cache uses to select a final entry point, or something like that? or maybe there's a better mechanism to operate on spirv itself like there is with naga, to actually modify the shader code somehow. i'm not sure exactly how it should look but i would need some strong reason why Pipeline::specialize is the right place to do this for spirv, rather than putting that functionality alongside the equivalent non-spirv functionality in the shader cache itself.

[Shfty]

thanks, that context makes it clearer for me. since i'm late i'm going to summarize to make sure we are on the same page

• shader defs consist of platform/device defs (NO_STORAGE_ARRAYS) and arbitrary user defs (TONEMAP_METHOD_XYZ)
• all these defs are used in the preprocessor within the shader cache to customise non-spirv shaders (to handle availability of features like storage arrays, and to statically select other custom functionality)
• also in non-spirv context, shader defs can be defined in shaders via the '#define' directive, so the full set of shader defs is not knowable until the shader is built by the shader cache (since #imports can be conditional on defs as well), so they can't really be used outside of the shader scope. in the current rust code the only read of the defs is in the preprocessor itself, and (now removed?) for spirv checking that there are no defs

Indeed - though I consider definitions made in shader code via #define to be outside the scope of this change, since you could do the equivalent in a Rust shader using #[cfg()] or such and be similarly unable to share it across the binary / textual boundary.

My "all shader defs" bulletpoint is probably a little broad in that sense!

if i understood correctly, what you're suggesting is

• move the platform/device defs into the pipeline cache so that they can be passed to Pipeline::specialize
• access the defs in Pipeline::specialize in order to select a spirv entrypoint based on them
• continue to handle non-spirv defs in the shader cache's preprocessor, just sourcing them from the pipeline cache instead

Somewhat. The defs are accessed in Material::specialize in my case, which is downstream of Specialized*Pipeline::specialize, but the broader point is to make them generally available (as per pipeline-sourced defs) rather than exclusive to code downstream of ShaderCache (i.e. WGSL / GLSL).

I'm not proposing to codify bevy SPIR-V handling as part of this PR. That's beyond scope; the intent here is to open the door for user code to codify its own SPIR-V handling in the absence of official support for something like rust-gpu (which is out-of-scope in the near term as per #5634.)

so in the non-spirv case, we would still have a variant being "selected"/built by the shader cache + preprocessor based on what defs are passed in from pipeline cache plus specialize, and in the spirv case we have the variant being selected by the pipeline itself, based on the (platform-subset-of) defs generated by the pipeline cache, and the defs that are output from specialize are ignored.

As-is, SPIR-V has no intrinsic (or bevy-provided) way to specialize. naga translation notwithstanding, what you put in is what you get, with an error + non-rendering material if whatever's compiled in doesn't align with the configured pipeline.

What you describe is effectively the case for SPIR-V in the existing codebase outside of ShaderCache late-injecting its definitions; pipeline-sourced defs propagate through the specialization hierarchy, are augmented by ShaderCache, and end up getting ignored by ShaderProcessor because it can't operate on a binary.

it feels pretty convoluted to me to have the usage location of the defs vary like this depending on the shader type. in particular i think it would make on-the-fly compiling / caching of shaders (from naga to spirv) harder in the future, since pipelines would need to know ahead of time what "style" of shader they are using, where currently they just do their specializing job regardless of the shader format. it would also maybe mean you can't share spirv shaders so easily between pipelines by just varying the defs that the pipeline outputs.

Material::specialize is the only feasible place for user code to step in and assume responsibility for specializing binary shaders as-is, so I consider it reasonable to provide it access to the full range of engine-side defs.

In terms of internal code, it makes sense that a lower-level solution would be preferable. That implies either exposing a lower level trait that can be implemented to custom-specialize at the ShaderCache level, or codifying the way bevy-handles SPIR-V.

i think the better answer is to use the defs in the shader cache itself to select the variant somehow for spirv, similarly to what we do for non-spirv. i don't know exactly how to do that - perhaps the Shader::Spirv variant could have a lookup based on entry point plus shader defs that the shader cache uses to select a final entry point, or something like that? or maybe there's a better mechanism to operate on spirv itself like there is with naga, to actually modify the shader code somehow. i'm not sure exactly how it should look but i would need some strong reason why Pipeline::specialize is the right place to do this for spirv, rather than putting that functionality alongside the equivalent non-spirv functionality in the shader cache itself.

I put together a shader def -> entrypoint permutation mapping in the course of building the trait machinery to specify entrypoints in bevy-rust-gpu. It works well enough to cover the parameters for a Rust reimplementation of bevy_pbr's WGSL, handling enum-like flags, bool / int costs, and potentially generic types in the future.

However, it's arbitrary, as would be any solution that operates on an equally-arbitrary SPIR-V input. My plugin has the luxury of being a third-party codegen-backend-specific integration that can introduce such a restriction with the understanding that it will be handled automatically by the backend, provided that the user follows instructions.

I don't feel the same is true for bevy, since SPIR-V is intrinsically a black box, and the expectation for an engine is to accept that black box without imposing any additional restrictions. If ShaderCache-authority over binary specialization is necessary for proper performance, then it should be exposed in such a way that doesn't restrict usage.

On the topic of a naga_oil-style / transpilation / IR approach, I don't think that's viable as things stand; if nothing else, naga is currently not comprehensive enough to process arbitrary valid SPIR-V (see #8131, spirt/#9, naga/#1977), which would restrict such a solution to its supported subset.

(I'll post some context on how I arrived at the current approach in a separate message re. justifying the Pipeline::specialize approach in structural terms, as this response is running on 😅)

[Shfty]

Some additional context on how I arrived at the PipelineCache approach:

This PR originally avoided API changes by introducing a free-floating function to retrieve the defs in question, which would be called at construction time by any Specialized*Pipeline implementors. As that would have introduced an easy footgun of failing to call the function, I retooled it to move construction to the topmost point in the specialization hierarchy and pass the result down from there.

During the initial rework, I tried using a BaseShaderDefs(Vec<ShaderDefVal>) resource that would be added to the render world at the same time as PipelineCache, construct the defs in FromWorld, and be available to the various render world systems that call the top-level specialize functions, which then waterfall down through Specialized*Pipeline and Material, ultimately arriving at ShaderProcessor. (i.e. the existing implementation, but with top-level BaseShaderDefs replacing an extra ShaderCache stage between Material and ShaderProcessor.)

Unfortunately, this didn't work, as one of the sprite pipeline systems started throwing a "does not implement " error at its add_system call; I'm not certain, but suspect this may have been down to hitting some internal system param count limit, as it was localized to that specific system (which already had a high param count), and no extra type constraints were present.

Looking for alternatives, I noted that PipelineCache was present in every system that would consume BaseShaderDefs. Given that ShaderCache is a member of PipelineCache, it seemed reasonable to move def construction up one level in the type hierarchy in order to solve the ordering issue without necessitating the addition of a new resource.

[robtfm]

i'm basically going to repeat what i said before: I don't think Pipeline::specialize is the right place to do shader specialization. the pipeline should not care if the shader is spirv or wgsl or naga IR or anything else. this is for clean separation of responsibilities and is to support compiled shader-caching which is something we definitely want in future.

if you can select (either the whole spirv module or the entry point) based on defs in the shader cache in a way that works for you, then maybe it'll work for everyone. if not, we can iterate on it in future.

SPIR-V is intrinsically a black box

sure, but bevy's Shader::Spirv doesn't have to be

internal system param count limit

if so you can work around this using a tuple argument like (q_1, q_2): (Query<&One>, Query<&Two>). you could also use the platform resource in a FromWorld impl to avoid having to pass extra data to Pipeline::specialize. i don't think that's the right approach for shader code branching, but it could be useful for texture format selection or similar.

[Shfty]

i'm basically going to repeat what i said before: I don't think Pipeline::specialize is the right place to do shader specialization. the pipeline should not care if the shader is spirv or wgsl or naga IR or anything else. this is for clean separation of responsibilities and is to support compiled shader-caching which is something we definitely want in future.

That's clear, but separable from both the issue of the Material API's access to bevy-side shader defs being incomplete, and my by-necessity implementation of it doing something suboptimal in userland because the engine only affords me so much control.

The former remains a demonstrable bevy-specific problem that should be fixed, or documented on the Material trait to ensure that users understand they won't be receiving the full set of definitions relating to the mesh_view bindings available to their shaders. To that end, #8190 has been opened to track the issue in proper context - I'm not married to this specific implementation, and am willing to look into whatever solution will check the required boxes. Suggestions would be welcome.

And the latter can be removed from the equation entirely. Specialization best practices would not have entered the discussion if my touchstone example had been anything other than a skunkworks rust-gpu plugin project that happens to use Material::specialize as its entrypoint to the bevy rendering ecosystem.

Ergo, given that I am not asking bevy to compromise its own approach to specialization, it's irrelevant unless upcoming changes to the associated machinery are liable to reduce or remove access to shader definitions in Material. In which case, that would be a regression in the existing API, and counter to the discussion with @superdump that motivated this string of issues and pull requests in the first place.

[robtfm]

unless upcoming changes to the associated machinery are liable to reduce or remove access to shader definitions in Material

this is exactly the case with the platform.wgsl i mentioned above (i'm not just being difficult :))

[superdump]

I haven’t had time to catch up on this discussion nor do I right this moment. I just want to say that I haven’t either had the time to look at @robtfm ’s proposed changes and so I hadn’t considered that my proposal might impact them. I’ve chatted with Rob on Discord a bit to align our thinking and it sounds like this platform definitions approach is going to be generally useful. That said, I haven’t quite understood yet what breaks what, though I trust Rob’s judgement to find a solution that fits the use cases and I hope it won’t be too much bother to adjust things to address them all. I’m not the only authority around rendering parts anymore. :) I’ll catch up when I can.

My line of thought was founded in that shader defs impact shader code which impact the pipeline that is constructed. I was thinking then that generally, specialisation should control how pipelines are configured, so it made sense that all shader defs one way or another go through the specialisation function on Material. However, if there are good reasons that globally-constant shader defs should not go through pipeline specialisation, but should be available to shaders to use, then I’m open to understanding what the tradeoffs are and why we want to make an exception.

[Shfty]

unless upcoming changes to the associated machinery are liable to reduce or remove access to shader definitions in Material

this is exactly the case with the platform.wgsl i mentioned above (i'm not just being difficult :))

In which case, it would be productive to detail those changes in the appropriate issue / PR / RFC so potentially-controversial regressions - and their accompanying resolutions - can be discussed, before using them as justification against affected PRs.

#5703 only references shader defs in a limited capacity that doesn't suggest a level of impact comparable to what you've implied here.

It's very well to be told that my existing entrypoints may be taken away because the approach they support doesn't align with bevy's development trajectory - that's an implicit occupational hazard - but I can only contribute effectively and follow said trajectory if the rules of engagement exist on paper.

SPIR-V is intrinsically a black box

sure, but bevy's Shader::Spirv doesn't have to be

Is there an RFC or issue open for discussing potential stratification of bevy's SPIR-V handling? This seems like a controversial change from the perspective of using the format for the sake of total control, so it would be good to contribute to discussion on that front.

I haven’t had time to catch up on this discussion nor do I right this moment. I just want to say that I haven’t either had the time to look at @robtfm ’s proposed changes and so I hadn’t considered that my proposal might impact them. I’ve chatted with Rob on Discord a bit to align our thinking and it sounds like this platform definitions approach is going to be generally useful. That said, I haven’t quite understood yet what breaks what, though I trust Rob’s judgement to find a solution that fits the use cases and I hope it won’t be too much bother to adjust things to address them all. I’m not the only authority around rendering parts anymore. :) I’ll catch up when I can.

Of course, I understand that there's much to do in many spheres outside my localized rust-gpu orbit 🙂

And it's no bother - I just want to make sure the correct solution has been found, or at least progressed toward, when the time comes to give it full consideration.

[JMS55]

Another use case I've seen asked for: Having a custom material that all it does is use the standard material shaders, but removes the ENVIRONMENT_MAP_LIGHT shader def to disable environment map lighting for a specific mesh.