Add WorkgroupSize builtin support

[LegNeato]

Implement #[spirv(workgroup_size)] builtin that provides workgroup dimensions as a constant to compute shaders. The constant is created from the LocalSize execution mode values.

See https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#BuiltIn and https://github.com/KhronosGroup/SPIRV-Guide/blob/main/chapters/local_size_and_workgroup_size.md

[nazar-pc]

Specialization constants seem to be supported as well, at least according to specs. Is this a temporary restriction? I'd add a TODO if it is.

[Firestar99]

see #298 (comment) for how they work :D

[Firestar99]

I think it's fine the way it is, but I do have some design feedback for potential future iterations:

I'd much rather want the workgroup size be a constant, similar to what I'm doing rn:

pub const LIGHTING_WG_SIZE: u32 = 64;

const_assert_eq!(LIGHTING_WG_SIZE, 64);
#[bindless(compute(threads(64)))]
pub fn lighting_cs(...) { ... }

This allows me to access the workgroup size from the CPU code as well to compute the workgroup dimensions:

let groups = [
	(image_size.x + LIGHTING_WG_SIZE - 1) / LIGHTING_WG_SIZE,
	image_size.y,
	1,
];

And allows you to declare shared memory as a multiple of the workgroup size:

pub const DIRECTIONAL_SHADOWS_WG_SIZE: u32 = 64;
const SHARED_SIZE: usize = DIRECTIONAL_SHADOWS_WG_SIZE as usize * 2;

const_assert_eq!(DIRECTIONAL_SHADOWS_WG_SIZE, 64);
#[bindless(compute(threads(64)))]
pub fn directional_shadows(
    #[spirv(workgroup)] shared: &[f32; SHARED_SIZE],
) { ... }

Instead of this PR, I'd much rather have the proc macro accept a const expr instead of a literal, so we can do this:

pub const LIGHTING_WG_SIZE: u32 = 64;

#[bindless(compute(threads(LIGHTING_WG_SIZE)))]
pub fn lighting_cs(...) { ... }

Not sure how hard this is to implement.

[LegNeato]

I tried to do the rust const route for the past week, the problem is that rustc doesn't support constants in attrs currently so I have to do a bunch of hacky stuff with resolution to make it work. I have it working for most cases, but not all. Not sure if it is better to support it with sharp edge cases or not support it at all.

[LegNeato]

I'll clean up what I have and we can compare and contrast. At the very least we'll need to document as people coming from glsl will expect something like workgroupsize attr to exist.

[Firestar99]

@LegNeato I wonder if we actually need to parse out the actual value within that macro....

Currently we're emitting this, for which we clearly need to parse the literals:

     OpEntryPoint GLCompute %3 "lighting_cs" %bla %bla2 %bla3
     OpExecutionMode %3 LocalSize 64 1 1
%3 = OpFunction %void None %443

But with Vulkan1.2 we get the new fancy LocalSizeId instead of LocalSize, so we can do this:

%4 = OpTypeInt 32 0
     OpDecorate %5 SpecId 123
%5 = OpSpecConstant %4 64
%6 = OpConstant %4 1
     OpEntryPoint GLCompute %3 "lighting_cs" %bla %bla2 %bla3
     OpExecutionMode %3 LocalSizeId %5 %6 %6
%3 = OpFunction %void None %443

Note how at the point we're writing the OpExecutionMode we don't actually write any literaly, just references to constants. Why couldn't they be references to actual constants in rust code? And if we have number literals, we can still parse and emit them as constants. It would also open up a path towards workgroup size being a spec constant, see %5.

We probably should move this design discussion to an issue, I just want to talk about it now in case this has any ramifications for this PR as well.

[LegNeato]

rustc just changed some attr parsing infra, need to look and see if it better enables this

[nazar-pc]

I just read this: https://github.com/KhronosGroup/SPIRV-Guide/blob/1482521d9826fdbb48854dfff8c46cc4a8b3607e/chapters/local_size_and_workgroup_size.md

And while I don't really understand much, it seems that:

• WorkgroupSize is actually deprecated in SPIR-V 1.6, so it is questionably whether it is worth adding it at this point
• the same page claims that LocalSizeId is not available by default unless VK_KHR_maintenance4 extension is used, so it doesn't seem to be a good default

While I think LocalSizeId might be useful in general, I'll most likely stick with fixed workgroup size, for which the most beneficial change would be to support constants in #[spirv(compute(threads(WORKGROUP_SIZE)))] like @Firestar99 suggested above.

A simpler (but less intuitive) way to achieve the same might be to generate a submodule in proc macro, such that this works:

#[spirv(compute(threads(256)))]
pub fn entrypoint() {}

// This is the same as `const_assert_eq!()`, but without macros and external crates
const _: () = {
    // `entrypoint` is a macro-generated module and has a constant defined in it
    assert!(entrypoint::WORKGROUP_SIZE == 256);
};