Composable Pipeline Specialization

[ecoskey]

Currently, our specialization API works through a series of wrapper structs and traits, which make things confusing to follow and difficult to generalize.

This pr takes a different approach, where "specializers" (types that implement Specialize) are composable, but "flat" rather than composed of a series of wrappers. The key design choice is that specializers don't produce pipeline descriptors, but instead modify existing ones:

pub trait Specialize<T: Specializable> {
    type Key: Clone + Hash + Eq;
    
    fn specialize(&self, key: Self::Key, descriptor: &mut T::Descriptor);
}

This lets us use some derive magic to stick multiple specializers together:

pub struct A;
pub struct B;

impl Specialize<RenderPipeline> for A { ... }
impl Specialize<RenderPipeline> for A { ... }

#[derive(Specialize)]
#[specialize(RenderPipeline)]
struct C {
    // specialization is applied in struct field order
    applied_first: A,
    applied_second: B,
}

type C::Key = (A::Key, B::Key);

This approach is much easier to understand, IMO, and also lets us separate concerns better. Specializers can be placed in fully separate crates/modules, and key computation can be shared as well.

The most significant change workflow-wise for this feature is that since specializers only modify descriptors, we need a "base" descriptor to work off of. This can either be manually supplied when constructing a Specializer (the new collection replacing Specialized[Render/Compute]Pipelines), or supplied by implementing HasBaseDescriptor on a specializer. See examples/shader/custom_phase_item.rs for an example implementation.

Testing

• Did some simple manual testing of the derive macro, it seems robust.

---

Showcase

What material specialization might look like after this change:

#[derive(Specialize, HasBaseDescriptor)]
#[specialize(RenderPipeline)]
pub struct SpecializeMeshMaterial<M: Material> {
    mesh: SpecializeMesh, // set mesh bind group layout and shader defs
    view: SpecializeView, // set view bind group layout and shader defs
    #[key(default)] // since type SpecializeMaterial::Key = (), we can hide it from the wrapper's external API
    #[base_descriptor] //defer to the HasBaseDescriptor impl of SpecializeMaterial, since it carries the vertex and fragment handles
    material: SpecializeMaterial<M>, // set material bind group layout
}

//generated implementation by the derive macro:
impl <M: Material> Specialize<RenderPipeline> for SpecializeMeshMaterial<M> {
    type Key = (MeshKey, ViewKey);

    fn specialize(&self, key: Self::Key, descriptor: &mut RenderPipelineDescriptor) {
        self.mesh.specialize(key.0, descriptor);
        self.view.specialize(key.1, descriptor);
        self.material.specialize((), descriptor);
    }
}

impl <M: Material> HasBaseDescriptor<RenderPipeline> for SpecializeMeshMaterial<M> {
    fn base_descriptor(&self) -> RenderPipelineDescriptor {
        self.material.base_descriptor()
    }
}

Migration Guide

For manual implementers

SpecializedRenderPipeline -> Specialize<RenderPipeline>.
SpecializedMeshPipeline -> Specialize<RenderPipeline>
SpecializedComputePipeline -> Specialize<ComputePipeline>.

SpecializedXPipelines<S> -> Specializer<X, S>

For derive macro usage

Derive macro names can't have generic parameters, so to derive Specialize<RenderPipeline> use #[derive(Specialize)] and #[specialize(RenderPipeline)].
To derive Specialize<ComputePipeline> use #[derive(Specialize)] and #[specialize(ComputePipeline)]. #[specialize(..targets)] can take multiple specialize targets, or use #[specialize(all)] to generate a fully generic implementation, at the cost of slightly worse error messages when the trait bounds aren't satisfied.

Questions

• optionally prehash sub-keys?

[Bleachfuel]

u can use bevy_macro_utils::get_struct_fields. Probably also for more things here.

[ecoskey]

I used a custom version originally for better error messages, but I should probably just upstream those changes to get_struct_fields

[tychedelia]

How does this account for SpecializedMeshPipeline which needs a MeshVertexBufferLayoutRef that's derived from the mesh that's currently being specialized? This is the problem that I ran into when trying to make something similarly generic. Maybe we could have some kind of "extra data" that could piggyback along with the key? Or just add an extra parameter slot / associated type that's () for most implementations?