non-literal workgroup size

[Firestar99]

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:

1. 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,
];

2. 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.

[nazar-pc]

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:

const fn shared_size(workgroup_size: u32) -> usize {
    workgroup_size as usize * 2
}

pub fn directional_shadows<#[spirv(workgroup_size)] const WORKGROUP_SIZE: u32>(
    #[spirv(workgroup)] shared: &[f32; shared_size(WORKGROUP_SIZE)],
)
where
    [(); shared_size(WORKGROUP_SIZE)]:
{ ... }

This requires experimental nightly features, but it is really powerful even with current limitations.

[LegNeato]

I don't believe they do and don't immediately see a nightly feature to flip on.

[LegNeato]

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]

#[cfg(feature = "something")] const WORKGROUP_SIZE: u32 works 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]

@Firestar99 we could only do LocalSizeId for 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.

Issue #756: Deprecated the use of BuiltIn to decorate a constant to set its value and removed the deprecation of the WorkgroupSize built-in. That is, WorkgroupSize is kept but no longer marked as deprecated (it is still required by OpenCL). The use of BuiltIn to decorate a constant to set its value was only for WorkgroupSize, which has been superseded by the LocalSizeId execution mode.

[LegNeato]

How would spec constants look UX-wise?

#[spirv(compute(threads(x_spec, y_spec, 1)]
fn main(
    #[spirv(spec_constant(id = 1))] x_spec: u32,
    #[spirv(spec_constant(id = 9000, default = 4))] y_spec: u32,
) {}

[Firestar99]

Yes LocalSizeId is 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.