Add vertex attribute descriptor specialization to materials

[mgi388]

Background

At the time of writing, custom vertex attributes need to be added inside Material.specialize. For example, look at line 82 of the custom vertex attributes example:

bevy/examples/shader/custom_vertex_attribute.rs

Lines 74 to 86 in d17fb16

	fn specialize(
	_pipeline: &MaterialPipeline<Self>,
	descriptor: &mut RenderPipelineDescriptor,
	layout: &MeshVertexBufferLayoutRef,
	_key: MaterialPipelineKey<Self>,
	) -> Result<(), SpecializedMeshPipelineError> {
	let vertex_layout = layout.0.get_layout(&[
	Mesh::ATTRIBUTE_POSITION.at_shader_location(0),
	ATTRIBUTE_BLEND_COLOR.at_shader_location(1),
	])?;
	descriptor.vertex.buffers = vec![vertex_layout];
	Ok(())
	}

To add a new attribute you need to overwrite the incoming descriptor.vertex.buffers with your new layout.

However, this means you also need to provide all of the well-known attributes like position, uv, etc. that you want to use in your custom shader. You can see at line 81 it includes the position attribute because it's used in the custom shader.

For the custom vertex attributes Bevy example, this is fine because it is a tiny example that only uses two attributes in its custom shader (position and the custom one):

bevy/assets/shaders/custom_vertex_attribute.wgsl

Lines 8 to 12 in d17fb16

	struct Vertex {
	@builtin(instance_index) instance_index: u32,
	@location(0) position: vec3<f32>,
	@location(1) blend_color: vec4<f32>,
	};

Experience report

However, I found that when writing my own material and shader and trying to reuse whatever existing PBR code I could, my custom attributes caused issues specifically for the prepass shader. This is because my specialize function was always adding my own version of the well-known attributes, specifically color, but the Bevy builtin forward and prepass shaders use different shader locations for color, and it crashed. I worked around the issue by using a key check in my specialize:

fn specialize() {
    // Leave prepass pipeline alone for now so that the well-known
    // prepass pipeline and shaders can be used.
    if key.mesh_key.contains(MeshPipelineKey::DEPTH_PREPASS) {
        return Ok(());
    }

    let attribute_descriptors = vec![
        Mesh::ATTRIBUTE_POSITION.at_shader_location(0),
        Mesh::ATTRIBUTE_NORMAL.at_shader_location(1),
        Mesh::ATTRIBUTE_UV_0.at_shader_location(2),
        Mesh::ATTRIBUTE_TANGENT.at_shader_location(4),
        Mesh::ATTRIBUTE_COLOR.at_shader_location(5),
        ATTRIBUTE_TEXTURE_INDEX.at_shader_location(10),
    ];

    let vertex_layout = layout.0.get_layout(&attribute_descriptors)?;
    descriptor.vertex.buffers = vec![vertex_layout];
    Ok(())
}

If I wanted to provide my own custom attributes to the prepass branch, then I'd need to inline all of those well-known attributes as well, and ideally also include the condition checks like this:

if layout.0.contains(Mesh::ATTRIBUTE_POSITION) {
    shader_defs.push("VERTEX_POSITIONS".into());
    vertex_attributes.push(Mesh::ATTRIBUTE_POSITION.at_shader_location(0));
}

Which is repeating a lot of boilerplate and behavior I really want to inherit.

Proposal

It would be nice if material implementers could implement a simple specialization function for vertex attribute descriptors only. For argument's sake, call it specialize_vertex_attributes (name and required arguments TBD).

I would expect it would take similar arguments to specialize. It's unclear if it needs descriptor, but it would definitely take the "in-progress" list of created attributes from the base material. Then, material implementers could either choose to push their own custom attribute to the end, or replace the whole thing if they need to.

In the base, it would accumulate the well-known attributes as it already does, then call specialize_vertex_attributes with them, and finally it would call M::specialize as it currently does.

For the linked custom vertex attributes example, I expect it would then look something like this:

impl Material for CustomMaterial {
    // ...

    fn specialize_vertex_attributes(
        _pipeline: &MaterialPipeline<Self>,
        _descriptor: &mut RenderPipelineDescriptor,
        attrs: &mut Vec<VertexAttributeDescriptor>,
    ) -> Result<(), SpecializedMeshPipelineError> {
        attrs.push(ATTRIBUTE_BLEND_COLOR.at_shader_location(100)); // Note: Changed from 1 to 100
        Ok(())
    }
}

This clearly illustrates to readers of the code that the material just wants to specialize the base material with an extra vertex attribute but do nothing else.

Note: I changed the location to 100 because the well-known attributes already use position 1.

A version of specialize_vertex_attributes that overwrites the entire list, to match the existing example, would look something like this:

impl Material for CustomMaterial {
    fn specialize_vertex_attributes(
        _pipeline: &MaterialPipeline<Self>,
        _descriptor: &mut RenderPipelineDescriptor,
        attrs: &mut Vec<VertexAttributeDescriptor>,
    ) -> Result<(), SpecializedMeshPipelineError> {
        *attrs = vec![
            Mesh::ATTRIBUTE_POSITION.at_shader_location(0),
            ATTRIBUTE_BLEND_COLOR.at_shader_location(1),
        ];
        Ok(())
    }
}

To be clear, the custom shader still requires the author to pull in the vertex attribute bindings they need. This is because the custom shader defines its own Vertex. However, being able to copy and paste straight from Bevy's own Vertex and VertexOutput into your own Vertex and VertexOutput is valuable.

I also wonder if this proposal would help mitigate this kind of concern written by one of the Bevy core developers

bevy/crates/bevy_pbr/src/prepass/mod.rs

Lines 593 to 596 in f61c55f

	// This is a bit risky because it's possible to change something that would
	// break the prepass but be fine in the main pass.
	// Since this api is pretty low-level it doesn't matter that much, but it is a potential issue.
	M::specialize(&self.material_pipeline, &mut descriptor, layout, key)?;

In the custom vertex attributes example, the version using specialize_vertex_attributes no longer needs to implement specialize and so the Bevy core developer's concern is completely eliminated in this case.

Alternatives considered

It's possible that you can already do this using specialize, but I spent some time on it and couldn't work it out. In any case, the proposal in this issue is more of a "decoupled" approach where you can just tell Bevy what attributes you want, rather than modifying the internals of descriptor.vertex.buffers inside a very general (IMO) specialize function.

[mintlu8]

Second this, the current implementation is confusing and difficult to get right in 3d.

[flash-freezing-lava]

To share my current work-around:

Considering the implementation of MeshVertexBufferLayout::get_layout, I think you only have to push VertexAttributes to the first buffer.
At least for me it worked to copy the below lines:

bevy/crates/bevy_render/src/mesh/mesh/mod.rs

Lines 1357 to 1374 in ddf4666

	if let Some(index) = self
	.attribute_ids
	.iter()
	.position(|id| *id == attribute_descriptor.id)
	{
	let layout_attribute = &self.layout.attributes[index];
	attributes.push(VertexAttribute {
	format: layout_attribute.format,
	offset: layout_attribute.offset,
	shader_location: attribute_descriptor.shader_location,
	});
	} else {
	return Err(MissingVertexAttributeError {
	id: attribute_descriptor.id,
	name: attribute_descriptor.name,
	pipeline_type: None,
	});
	}

into:

fn specialize(
    pipeline: &bevy::pbr::MaterialExtensionPipeline,
    descriptor: &mut bevy::render::render_resource::RenderPipelineDescriptor,
    layout: &bevy::render::mesh::MeshVertexBufferLayoutRef,
    key: bevy::pbr::MaterialExtensionKey<Self>,
) -> Result<(), bevy::render::render_resource::SpecializedMeshPipelineError> {
  for (attribute_descriptor, shader_location) in [(ATTRIBUTE_1, 20), (ATTRIBUTE_2, 21), (ATTRIBUTE_3, 22)] {
      if let Some(index) = layout.0.attribute_ids()
          .iter()
          .position(|id| *id == attribute_descriptor.id)
      {
          let layout_attribute = &layout.0.layout().attributes[index];
          descriptor.vertex.buffers[0].attributes.push(VertexAttribute {
              format: layout_attribute.format,
              offset: layout_attribute.offset,
              shader_location,
          });
      }
  }
  Ok(())
}