schell / schell/wgsl-rs

Configurable workgroup sizes via WGSL override constants

Open
#94 0 comments 0 reactions 0 assignees View on GitHub
enhancement P3
Dominant language
Rust
Stars
61
Forks
4
Avg merge
1d 21h
Merged PRs (30d)
9

Description

## Summary

The `#[workgroup_size(...)]` annotation currently bakes workgroup dimensions as hard-coded literal values in the generated WGSL. This means the workgroup size is fixed at compile time and cannot be tuned at runtime without regenerating the shader source.

This issue proposes making workgroup sizes configurable at pipeline creation time using WGSL `override` declarations and wgpu's `PipelineCompilationOptions::constants`.

## Motivation

Different GPUs have different optimal workgroup sizes. A workgroup size of 64 may be ideal on one GPU but suboptimal on another. Being able to configure the workgroup size at pipeline creation time allows applications to:
- Tune performance per-device
- Experiment with workgroup sizes without recompilation
- Use the annotation value as a sensible default while allowing overrides

## Design

### WGSL output changes

Currently, for `#[compute] #[workgroup_size(64)]`, the generated WGSL is:

```wgsl
@compute @workgroup_size(64) fn main(...) { ... }
```

The new output would emit `override` declarations and reference them in `@workgroup_size`:

```wgsl
override WORKGROUP_SIZE_X: u32 = 64;
override WORKGROUP_SIZE_Y: u32 = 1;
override WORKGROUP_SIZE_Z: u32 = 1;

@compute @workgroup_size(WORKGROUP_SIZE_X, WORKGROUP_SIZE_Y, WORKGROUP_SIZE_Z)
fn main(...) { ... }
```

This is valid per the WGSL spec -- `@workgroup_size` parameters may be override-expressions ([WGSL spec 12.14](https://www.w3.org/TR/WGSL/#workgroup-size-attr)).

For modules with multiple compute entry points, the override names should be scoped per entry point, e.g. `MAIN_WORKGROUP_SIZE_X`, `COMPUTE_FOO_WORKGROUP_SIZE_X`, etc.

### Linkage changes (in `linkage.rs`)

The generated `linkage` module should provide two pipeline creation paths:

**1. Default workgroup size (existing behavior, ergonomic):**

```rust
// Uses the default workgroup size from the annotation
pub fn compute_pipeline(
device: &wgpu::Device,
layout: Option<&wgpu::PipelineLayout>,
module: &wgpu::ShaderModule,
) -> wgpu::ComputePipeline { ... }
```

**2. Custom workgroup size (new):**

```rust
/// Creates a compute pipeline with a custom workgroup size, overriding the default.
pub fn compute_pipeline_with_workgroup_size(
device: &wgpu::Device,
layout: Option<&wgpu::PipelineLayout>,
module: &wgpu::ShaderModule,
workgroup_size: (u32, u32, u32),
) -> wgpu::ComputePipeline { ... }
```

The custom variant would populate `PipelineCompilationOptions::constants` with the override constant names mapped to the provided values:

```rust
let constants = &[
("ENTRY_WORKGROUP_SIZE_X", workgroup_size.0 as f64),
("ENTRY_WORKGROUP_SIZE_Y", workgroup_size.1 as f64),
("ENTRY_WORKGROUP_SIZE_Z", workgroup_size.2 as f64),
];
```

Similarly, `compute_pipeline_descriptor` should get a `_with_workgroup_size` variant.

The existing `WORKGROUP_SIZE` constant should be kept as `DEFAULT_WORKGROUP_SIZE` to reflect that it represents the default from the annotation, not necessarily the runtime value.

### CPU-side runtime changes

The `dispatch_workgroups` function in `runtime.rs` already accepts `workgroup_size: (u32, u32, u32)` as a parameter, so no changes are needed for the CPU-side compute dispatch.

## Affected files

| File | Change |
|------|--------|
| `crates/wgsl-rs-macros/src/code_gen/formatter.rs` (lines 1533-1574) | Emit `override` declarations; reference them in `@workgroup_size(...)` |
| `crates/wgsl-rs-macros/src/linkage.rs` (lines 63-67) | Update `ComputeEntry` if override naming needs to be tracked |
| `crates/wgsl-rs-macros/src/linkage.rs` (lines 523-591) | Generate `_with_workgroup_size` variants; rename `WORKGROUP_SIZE` to `DEFAULT_WORKGROUP_SIZE` |
| `crates/wgsl-rs-macros/src/parse.rs` (lines 3170-3318) | No changes expected (parsing stays the same) |
| `crates/wgsl-rs/src/lib.rs` (lines 22-31) | No changes expected (`Module` struct is unaffected) |

## Open questions

- **Override naming for multi-entry modules:** Should override constants be prefixed with the entry point name (e.g. `MAIN_WORKGROUP_SIZE_X`) or use `@id` numeric identifiers? Name-based is more readable; `@id`-based avoids naming conflicts but is less ergonomic.
- **Validation:** The naga validation test runs at compile time against the generated WGSL. Need to verify naga properly handles `override` in `@workgroup_size` positions.

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.