RFC: entry point redesign ideation

[Firestar99]

Roughly ordered from easiest to hardest to design / implement:

• builtins as functions in spirv-std, that any code can just call instead of relying on the callees to patch through the correct variables. Needs validation during linking, since buildins are often valid only in some shader types.
• wgpu "bind group layout" / vulkan "descriptor set layout" should have some kind of direct representation within rust-gpu directly. I'm thinking of a struct declaring all the bindings of one descriptor set, soft of like wgsl does it:

#[spirv(descriptor_set)]
struct BindGroupForSystemX {
  #[binding(0)] data: TypedBuffer<MyData>,
  #[binding(1)] base_color: Image2D,
  #[binding(2)] sampler: Sampler,
}

It would be amazing if you could pass them around within this struct (you currently can't), since you could do something like this:

#[spirv(fragment)]
fn my_entry(
    #[spirv(descriptor_set = 0)] system_x: &BindGroupForSystemX,
) {
  systemx::process(system_x)
}

Note how the descriptor_set / bind group id is declared at the entry point, so you could move it to whatever id you'd want and thus avoid conflicts if two systems use the same id? Also I pass it in as a reference, so you can't change eg. the buffer that's in systemx.data to another. We may want that as a feature, but assigning it non-uniformly makes it potentially problematic.

• small extra: Have a non-generic texture that can be unsafely cast to all the generic textures we have right now. For spirv, we need to automatically add a new texture binding of that texture type and the same binding id when casts happen. Super useful for bindless! See rust-gpu zulip, research, slang __DynamicResource
• "pipeline layout" should probably also some some representation as well. It declares the ids of the descriptor sets and the push constants. I'd move it to a struct similarly to above. This also makes it trivial to ensure a vertex and fragment shader have the same pipeline layout.

#[spirv(pipeline_layout)]
struct MyPipelineLayout {
  #[descriptor_set(0)] system_x: BindGroupForSystemX,
  #[descriptor_set(1)] system_y: BindGroupForSystemY,
  #[push_constant] push: MyPushConstant,
}

#[spirv(vertex)]
fn my_fragment(pipeline: MyPipelineLayout) {...}
#[spirv(fragment)]
fn my_fragment(pipeline: MyPipelineLayout) {...}

• Offer derive macros to create bind group layouts / descriptor set layouts / pipeline layouts directly from the structs above for wgpu / vulkan. We'd likely want our own struct representation of descriptor sets and pipeline layouts to derive to, which may be used by the application as metadata. We'd likely want a generic & const-generic variant for compile time checked code like below (from my project):

pub fn create_graphics_pipeline<T: BufferStruct>(
    vertex_shader: &impl BindlessShader<ShaderType = VertexShader, ParamConstant = T>,
    fragment_shader: &impl BindlessShader<ShaderType = FragmentShader, ParamConstant = T>,
    ...
);

We'd also want to offer a non-generic variant for dynamic use-cases, like within this function. And from there offer conversion functions from there to the various wgpu & vulkan objects, hidden behind features. It's just so annoying having to keep bindings and descriptor set IDs in sync!

• The spirv spec requires that an entry point specifies all the in out and built-ins that are used within the shader. Currently, you specify those explicitly in the entry point. But with all of the above, we'd need a fixup pass during linking to assemble the entry points properly. When we already have that, why not add link-time capabilities as well?

Feel free to add suggestions!

[LegNeato]

I have lots of ideas here and a 1/2 written blog post and spec. I think modeling it as a fn is wrong in general.

[nazar-pc]

I personally really like explicit arguments that are "normal Rust" rather than magical intrinsics-like functions. So current situation with arguments is fine, deriving something on a struct is fine too as long as that struct is then an explicit argument of the shader entrypoint.

I did raise a question about const-like things in Discord. It'd be nice to be able to use things like const SUBGROUP_SIZE: u32 in const generics to size various data structures like arrays. It is workable today as is through DCE, but I'd be interested in maybe seeing a proper integration of those things into Rust type system.

[marstaik]

I think in regards to descriptor sets/bindings, the issue is more on being able to consistently use the same sets over and over again.

In slang, you can declare binds globally, and then by importing the module you have access to it. That has some plus and minuses.

My current approach that has me very happy, but involves my own custom macro, is as follows:

// Hardcoded descriptor mappings: (marker, qualifier, descriptor_set, binding)
// Add or modify entries here to expose new parameter marker attributes.
const DESCRIPTORS: &[(&str, &str, u32, u32)] = &[
	("framebuffer_indo", "storage_buffer", 0, 0),
	("camera", "storage_buffer", 0, 1),
	("mesh_info", "storage_buffer", 0, 2),
	("instance_data", "storage_buffer", 0, 3),
	("material", "storage_buffer", 1, 0),
	("instance_material", "storage_buffer", 1, 1),
];

use proc_macro::TokenStream;
use proc_macro2::{Delimiter, Group, TokenTree};
use quote::{TokenStreamExt, format_ident, quote};
use spirv_std_types::spirv_attr_version::spirv_attr_with_version;

#[proc_macro_attribute]
pub fn scene_descriptors(_attr: TokenStream, item: TokenStream) -> TokenStream {
	let spirv_ident = format_ident!("{}", &spirv_attr_with_version());

	// prepend with #[rust_gpu::spirv(..)]
	// let attr: proc_macro2::TokenStream = attr.into();
	let mut tokens: proc_macro2::TokenStream = proc_macro2::TokenStream::new();

	let item: proc_macro2::TokenStream = item.into();
	for tt in item {
		match tt {
			TokenTree::Group(group) if group.delimiter() == Delimiter::Parenthesis => {
				let mut group_tokens = proc_macro2::TokenStream::new();
				let mut last_token_hashtag = false;
				for tt in group.stream() {
					let is_token_hashtag =
						matches!(&tt, TokenTree::Punct(punct) if punct.as_char() == '#');

					// If this is an attribute group (#[...]) whose marker matches one of our
					// hardcoded descriptors, replace it with the full spirv descriptor attribute.
					if let TokenTree::Group(bracket_group) = &tt {
						if bracket_group.delimiter() == Delimiter::Bracket && last_token_hashtag {
							if let Some(TokenTree::Ident(marker_ident)) =
								bracket_group.stream().into_iter().next()
							{
								let marker = marker_ident.to_string();
								if let Some((_, qualifier, set, binding)) =
									DESCRIPTORS.iter().find(|(name, ..)| *name == marker)
								{
									let qualifier_ident = format_ident!("{}", qualifier);
									group_tokens.extend(quote! {
										[cfg_attr(target_arch="spirv", rust_gpu::#spirv_ident (#qualifier_ident, descriptor_set = #set, binding = #binding))]
									});
									last_token_hashtag = is_token_hashtag;
									continue;
								}
							}
						}
					}

					group_tokens.append(tt);
					last_token_hashtag = is_token_hashtag;
				}
				let mut out = Group::new(Delimiter::Parenthesis, group_tokens);
				out.set_span(group.span());
				tokens.append(out);
			}
			_ => tokens.append(tt),
		}
	}
	tokens.into()
}

Then, when writing shader's, it turns out quite clean IMO:

#[scene_descriptors]
#[spirv(vertex)]
pub fn vertex_color_material_vert(
	#[camera] cameras: &[Camera],
	#[instance_data] instance_data: &[InstanceData],
	#[instance_material] instance_materials: &[VertexColorMaterialData],
	#[spirv(instance_index)] instance_index: u32,
	vertex: PositionColor,
	out_data: &mut VertexOutput,
	#[spirv(position)] out_pos: &mut Vec4,
) {
	let my_instance_data = &instance_data[instance_index as usize];
	let my_material_data = &instance_materials[instance_index as usize];

	let world_pos = my_instance_data.world_matrix * vertex.position;
	let clip_pos = cameras[0].projection_view * world_pos;

	*out_pos = clip_pos;
	out_data.color = my_material_data.color * vertex.color;
}

Maybe we can provide helper macros to setup common descriptor sets like this - and even enforce the actual type they check against.

Mine is just a rough implementation cause I'm the only consumer.

[Firestar99]

I made a first prototype for querying read-only builtins as functions within spirv_std: #459