Repository navigation
non-literal workgroup size #299
Description
Original discussion in #298
Currently, workgroup size must be a number literal and does not accept const expr. So to have a const if your workgroup size, you must copy-paste the literal in two / three locations:
pub const LIGHTING_WG_SIZE: u32 = 64;
const_assert_eq!(LIGHTING_WG_SIZE, 64);
#[bindless(compute(threads(64)))]
pub fn lighting_cs(...) { ... }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(...) { ... }Alternative: WorkgroupSize builtin
The WorkgroupSize builtin as in #298 is not sufficient for these use-cases:
- Accessing the workgroup size from the CPU code to compute the workgroup dimensions:
let groups = [
(image_size.x + LIGHTING_WG_SIZE - 1) / LIGHTING_WG_SIZE,
image_size.y,
1,
];- 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],
) { ... }Potential Implementation Path
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.
Activity
nazar-pc commentedon Jun 27, 2025
I don't remember if const generics support attributes, but if they do, that might be a solution for CPU code with
WorkgroupSize. With current Rust nightly it'll look something like this:This requires experimental nightly features, but it is really powerful even with current limitations.
LegNeato commentedon Jun 27, 2025
I don't believe they do and don't immediately see a nightly feature to flip on.
LegNeato commentedon Jun 27, 2025
What about not requiring it to be
UVec3? We are in a const context and therefore have access to the values at compile time and can support other types as long as they are in the proper range (we'd still support primitive parts of course, but those too don't need to be u32)nazar-pc commentedon Jun 27, 2025
#[cfg(feature = "something")] const WORKGROUP_SIZE: u32works on nightly, meaning it is a valid syntax that macro could ingest and strip for non-SPIR-V target and replace with something else for SPIR-V target.LegNeato commentedon Jun 27, 2025
@Firestar99 we could only do
LocalSizeIdfor vuklan1.2+, right? We'd need anotehr strategy for lower targets. OpenCL apparently needs WorkGroupSize built-in as well, not sure how well those targets are supported though.LegNeato commentedon Jun 27, 2025
How would spec constants look UX-wise?
Firestar99 commentedon Jun 27, 2025
Yes
LocalSizeIdis vulkan1.2+, but that should be 99% of our user base. So I actually don't mind the idea of requiring vulkan1.2 to be able to use constants instead of literals, and if only literals are in use, defer to the existing system.I would not actually worry about spec constants for workgroup size until someone actually needs it. You could probably get away with a bit of
macro_rules!copy-paste for the few use-cases i can imagine.LegNeato commentedon Jun 29, 2025
Found and fixed an issue in
rspirvwithLocalSizeId: gfx-rs/rspirv#263WIP for Rust-GPU#299 - const workgroup size.