Rust-GPU / Rust-GPU/rust-gpu

non-literal workgroup size

Open
#299 8 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement
Dominant language
Rust
Stars
3.4k
Forks
125
PR merge metrics
No merged PRs in 30d

Description

Original discussion in https://github.com/Rust-GPU/rust-gpu/pull/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 https://github.com/Rust-GPU/rust-gpu/pull/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,
];
  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.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by tracing the proc macro handling of #[bindless(compute(threads(...)))] and how it currently emits OpExecutionMode LocalSize. Read the original discussion in pull request 298, then compare the existing literal path with the proposed LocalSizeId and constant references. Done means const expressions are accepted for workgroup sizes while preserving literal support.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
compilers
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.