vello_shaders crate for AOT compilation and lightweight integration with external renderer projects

[dfrg]

This is a general proposal for a new crate to handle shader management.

The motivating desire was to support build time naga generation for native backends, but this would also be useful to clean up our current hand written shader code and act as a base for including more complex permutations.

This is a bit gnarly as a library + build script that depends on code in the same library. The modules are setup to support this.

The basic idea here is that we read the shader directory (along with the new permutations file in that dir) and preprocess all shaders, parse them and extract the resource bindings. This information is then used to generate constant data for all shaders inline in the crate in WGSL and/or MSL depending on enabled features.

There is additional infrastructure (see the PipelineHost trait) to abstract building pipelines from those shaders.

An example of the generated code can be found in this gist.

[dfrg]

To make things a bit more clear, integration into the current wgpu code would look something like:

// New type for our prepared pipelines
pub type FullShaders = vello_shaders::Pipelines<ShaderId>;

// This is just a trait to abstract pipeline creation for different backends
impl vello_shaders::PipelineHost for Engine {
    type Device = wgpu::Device;
    type ComputePipeline = ShaderId;
    type Error = Error;

   fn new_compute_pipeline(
        &mut self,
        device: &Self::Device,
        shader: &ComputeShader,
    ) -> Result<Self::ComputePipeline, Self::Error> {
        // basically same as current add_shader method
    }
}

// vello_shaders::wgsl::SHADERS contains the embedded preprocessed source permutations. This
// simply calls new_compute_pipeline() on engine for each shader and returns the result in a new struct.
let shaders = FullShaders::from_shaders(&vello_shaders::wgsl::SHADERS, &device, &mut engine).unwrap();

[raphlinus]

This needs to pin a git hash, as it's currently failing (missing msl::Options::zero_initialize_workgroup_memory)

[armansito]

gfx-rs/naga#2237 has been merged into mainline naga, which replaces PerStageResources with EntryPointResources and indexes them by entry point name. You can use that here if you pin naga against HEAD.

You should probably cache the entry-point name that the ShaderInfo is initialized with and use that here as the key for EntryPointResources.

Alternately you could process all reachable bindings for every entry-point in the module, set the map once and translate it to MSL in one go with the same EntryPointResourceMap. That might also make sense if we wanted a single ShaderInfo to support multiple entry-points.

[dfrg]

Thanks for the update. I've been waiting on that PR to continue iterating on this.

I'd like to experiment with merging some of our shaders but am deferring that until multisampled path rendering lands to avoid creating unnecessary conflicts.

In the meantime, I'll go ahead with the changes you suggest here and also add a path to generate a C++ header as we discussed.

[dfrg]

This is technically my PR so I won't click approve but additional changes by @armansito LGTM.

[armansito]

As we discussed we now have this crate hosted under crates/shaders. I have successfully integrated this into a native C++ Metal renderer, so I think we can move forward with this as a functional prototype.

There are some additional things that I would like to see but I think we should do these in follow-up PRs (possibly pending the encoding crate):

1. Actually interfacing with the WGSL metadata exported from vello_shaders and removing the shader pre-processing code out of src/ entirely (and also figure out how hot-reload should be handled by the shaders crate).
2. Adding a more detailed README that documents the API for Rust clients.
3. More documentation around the permutations file and extending support to permutations other than full (which we can handle as we resurrect and clean up the coverage mask rendering stages)
4. Documentation on the plans and timelines around other native shading language code that this crate will generate (e.g. SPIR-V).

I think much of this can wait until after I pull the encoding logic out to shaders/encoding. I think it would have been nice for one of us to document the motivations and the specifics of what we're building in an issue (to benefit the community), that said all of this work already falls under the Semi-stable encoding format and Native renderers items detailed on the 2023 roadmap.

[armansito]

I tested and built upon the changes that @dfrg originally proposed all of which lgtm. @dfrg has approved the changes that I pushed to this branch.

[dfrg]

Plans for future work seem on point. Thanks for pushing this forward!