Skip to content

Commit d78ee4b

Browse files
mvaligurskyMartin Valigursky
mvaligursky
and
Martin Valigursky
authored
feat(webgpu): automatic bind-group reflection for compute shaders (#8821)
* feat(webgpu): automatic bind-group reflection for compute shaders WebGPU compute shaders can now use the same simplified WGSL syntax as vertex/fragment shaders. The engine reflects resources from the shader source and builds the bind group automatically, so a hand-written computeBindGroupFormat is no longer required. - Reflect loose uniforms (into a generated uniform buffer), textures + samplers, storage buffers, and storage textures declared with simplified syntax. - computeBindGroupFormat / computeUniformBufferFormats are now optional. Reflected resources go in their own bind group (group 1 when a caller format is supplied, otherwise group 0); explicitly-bound resources are left untouched. - Generalize compute pipeline / bind-group execution to support multiple bind groups. - Fix storage-buffer reflection regex to accept read_write. - Convert edge-detect and particles examples to the simplified syntax. - Add unit tests for the compute reflection paths. Fixes #7689 * fix(webgpu): avoid GPUComputePipeline type in generated d.ts The @returns {GPUComputePipeline} JSDoc emitted an explicit return type into the .d.ts, which fails the standalone type check (no WebGPU lib types in scope). * refactor(webgpu): address compute reflection review comments - gpuFormatToPixelFormat: use last-wins so the canonical PIXELFORMAT (e.g. RGBA8 for 'rgba8unorm') is chosen, avoiding a needless bind group layout / pipeline cache split (the key uses the numeric PIXELFORMAT). - Assert computeUniformBufferFormats requires a computeBindGroupFormat, instead of silently dropping the uniform buffers. - STORAGE_BUFFER_REGEX: drop the invalid 'write' token (storage buffers only support read / read_write in WGSL). - Sync runCompute and buildResourceFormats JSDoc to mention storage textures. --------- Co-authored-by: Martin Valigursky <mvaligursky@snapchat.com>
1 parent 75e409b commit d78ee4b

9 files changed

Lines changed: 594 additions & 96 deletions

File tree

‎examples/src/examples/compute/edge-detect.compute-shader.wgsl‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
// Include half-precision type aliases (resolves to f16 when supported, f32 otherwise)
22
#include "halfTypesCS"
33

4-
@group(0) @binding(0) var inputTexture: texture_2d<f32>;
5-
@group(0) @binding(1) var inputTexture_sampler: sampler;
6-
@group(0) @binding(2) var outputTexture: texture_storage_2d<rgba8unorm, write>;
4+
// Simplified-syntax declarations (no @group/@binding) - the engine reflects these into a bind
5+
// group automatically, so the example does not provide a computeBindGroupFormat.
6+
var inputTexture: texture_2d<f32>;
7+
var inputTexture_sampler: sampler;
8+
var outputTexture: texture_storage_2d<rgba8unorm, write>;
79

810
@compute @workgroup_size(8, 8, 1)
911
fn main(@builtin(global_invocation_id) global_id : vec3u) {

‎examples/src/examples/compute/edge-detect.example.mjs‎

Lines changed: 4 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -173,18 +173,13 @@ assetListLoader.load(() => {
173173
const createComputeShader = () => {
174174
if (!device.supportsCompute) return null;
175175

176+
// No computeBindGroupFormat is provided - the input texture (+ sampler) and the output
177+
// storage texture use the simplified WGSL syntax and are reflected automatically by the
178+
// engine from the shader source.
176179
return new pc.Shader(device, {
177180
name: 'EdgeDetect-Shader',
178181
shaderLanguage: pc.SHADERLANGUAGE_WGSL,
179-
cshader: computeShaderWgsl,
180-
181-
// Format of a bind group for the compute shader
182-
computeBindGroupFormat: new pc.BindGroupFormat(device, [
183-
// Input texture with sampler (sampler takes binding slot+1 automatically)
184-
new pc.BindTextureFormat('inputTexture', pc.SHADERSTAGE_COMPUTE, undefined, undefined, true),
185-
// Output storage texture
186-
new pc.BindStorageTextureFormat('outputTexture', pc.PIXELFORMAT_RGBA8, pc.TEXTUREDIMENSION_2D)
187-
])
182+
cshader: computeShaderWgsl
188183
});
189184
};
190185

‎examples/src/examples/compute/particles.example.mjs‎

Lines changed: 5 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -84,31 +84,15 @@ assetListLoader.load(() => {
8484

8585
const numParticles = 1024 * 1024;
8686

87-
// a compute shader that will simulate the particles stored in a storage buffer
87+
// a compute shader that will simulate the particles stored in a storage buffer. No bind group
88+
// or uniform buffer formats are provided - the loose uniforms (count, dt, sphereCount) and the
89+
// storage buffers (particles, spheres) use the simplified WGSL syntax and are reflected
90+
// automatically by the engine from the shader source.
8891
const shader = device.supportsCompute ?
8992
new pc.Shader(device, {
9093
name: 'SimulationShader',
9194
shaderLanguage: pc.SHADERLANGUAGE_WGSL,
92-
cshader: shaderSharedWgsl + shaderSimulationWgsl,
93-
94-
// format of a uniform buffer used by the compute shader
95-
computeUniformBufferFormats: {
96-
ub: new pc.UniformBufferFormat(device, [
97-
new pc.UniformFormat('count', pc.UNIFORMTYPE_UINT),
98-
new pc.UniformFormat('dt', pc.UNIFORMTYPE_FLOAT),
99-
new pc.UniformFormat('sphereCount', pc.UNIFORMTYPE_UINT)
100-
])
101-
},
102-
103-
// format of a bind group, providing resources for the compute shader
104-
computeBindGroupFormat: new pc.BindGroupFormat(device, [
105-
// a uniform buffer we provided the format for
106-
new pc.BindUniformBufferFormat('ub', pc.SHADERSTAGE_COMPUTE),
107-
// particle storage buffer
108-
new pc.BindStorageBufferFormat('particles', pc.SHADERSTAGE_COMPUTE),
109-
// rad only collision spheres
110-
new pc.BindStorageBufferFormat('spheres', pc.SHADERSTAGE_COMPUTE, true)
111-
])
95+
cshader: shaderSharedWgsl + shaderSimulationWgsl
11296
}) :
11397
null;
11498

‎examples/src/examples/compute/particles.shader-simulation.wgsl‎

Lines changed: 12 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,30 +1,29 @@
1-
// uniform buffer for the compute shader
2-
struct ub_compute {
3-
count: u32, // number of particles
4-
dt: f32, // delta time
5-
sphereCount: u32 // number of spheres
6-
}
7-
81
// sphere struct used for the colliders
92
struct Sphere {
103
center: vec3<f32>,
114
radius: f32
125
}
136

14-
@group(0) @binding(0) var<uniform> ubCompute : ub_compute;
15-
@group(0) @binding(1) var<storage, read_write> particles: array<Particle>;
16-
@group(0) @binding(2) var<storage, read> spheres: array<Sphere>;
7+
// Simplified-syntax resources (no @group/@binding) - the engine reflects these into a bind group
8+
// automatically, so the example does not provide computeBindGroupFormat / computeUniformBufferFormats.
9+
// The loose uniforms are collapsed into a single generated uniform buffer.
10+
uniform count: u32; // number of particles
11+
uniform dt: f32; // delta time
12+
uniform sphereCount: u32; // number of spheres
13+
14+
var<storage, read_write> particles: array<Particle>;
15+
var<storage, read> spheres: array<Sphere>;
1716

1817
@compute @workgroup_size(64)
1918
fn main(@builtin(global_invocation_id) global_invocation_id: vec3u) {
2019

2120
// particle index - ignore if out of bounds (as they get batched into groups of 64)
2221
let index = global_invocation_id.x * 1024 + global_invocation_id.y;
23-
if (index >= ubCompute.count) { return; }
22+
if (index >= uniform.count) { return; }
2423

2524
// update times
2625
var particle = particles[index];
27-
particle.collisionTime += ubCompute.dt;
26+
particle.collisionTime += uniform.dt;
2827

2928
// if particle gets too far, reset it to its original position / velocity
3029
var distance = length(particle.position);
@@ -41,7 +40,7 @@ fn main(@builtin(global_invocation_id) global_invocation_id: vec3u) {
4140
var next = particle.position + delta;
4241

4342
// handle collisions with spheres
44-
for (var i = 0u; i < ubCompute.sphereCount; i++) {
43+
for (var i = 0u; i < uniform.sphereCount; i++) {
4544
var center = spheres[i].center;
4645
var radius = spheres[i].radius;
4746

‎src/platform/graphics/webgpu/webgpu-compute-pipeline.js‎

Lines changed: 20 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,8 @@ class CacheEntry {
2929
}
3030

3131
class WebgpuComputePipeline extends WebgpuPipeline {
32-
lookupHashes = new Uint32Array(2);
32+
// shader compute key + up to 2 bind group format keys (caller group 0 + reflected group)
33+
lookupHashes = new Uint32Array(3);
3334

3435
/**
3536
* The cache of compute pipelines
@@ -38,12 +39,22 @@ class WebgpuComputePipeline extends WebgpuPipeline {
3839
*/
3940
cache = new Map();
4041

41-
get(shader, bindGroupFormat) {
42+
/**
43+
* @param {import('../shader.js').Shader} shader - The compute shader.
44+
* @param {import('../bind-group-format.js').BindGroupFormat[]} bindGroupFormats - The bind group
45+
* formats, in bind group index order (dense, no gaps).
46+
* @returns {object} - The compute pipeline (GPUComputePipeline).
47+
*/
48+
get(shader, bindGroupFormats) {
49+
50+
Debug.assert(bindGroupFormats.length <= 2);
4251

43-
// unique hash for the pipeline
52+
// unique hash for the pipeline - shader key followed by each bind group format key (0 for
53+
// an absent group). All slots are written, so no need to clear stale values from reuse.
4454
const lookupHashes = this.lookupHashes;
4555
lookupHashes[0] = shader.impl.computeKey;
46-
lookupHashes[1] = bindGroupFormat.impl.key;
56+
lookupHashes[1] = bindGroupFormats[0] ? bindGroupFormats[0].impl.key : 0;
57+
lookupHashes[2] = bindGroupFormats[1] ? bindGroupFormats[1].impl.key : 0;
4758
const hash = hash32Fnv1a(lookupHashes);
4859

4960
// Check cache
@@ -58,8 +69,11 @@ class WebgpuComputePipeline extends WebgpuPipeline {
5869
}
5970
}
6071

61-
// Cache miss - create new pipeline
62-
const pipelineLayout = this.getPipelineLayout([bindGroupFormat.impl]);
72+
// Cache miss - create new pipeline. Build the impl array explicitly (at most 2 groups).
73+
const impls = [];
74+
if (bindGroupFormats[0]) impls.push(bindGroupFormats[0].impl);
75+
if (bindGroupFormats[1]) impls.push(bindGroupFormats[1].impl);
76+
const pipelineLayout = this.getPipelineLayout(impls);
6377
const cacheEntry = new CacheEntry();
6478
cacheEntry.hashes = new Uint32Array(lookupHashes);
6579
cacheEntry.pipeline = this.create(shader, pipelineLayout);

‎src/platform/graphics/webgpu/webgpu-compute.js‎

Lines changed: 65 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,14 @@ class WebgpuCompute {
1515
/** @type {UniformBuffer[]} */
1616
uniformBuffers = [];
1717

18-
/** @type {BindGroup} */
19-
bindGroup = null;
18+
/**
19+
* Bind groups, indexed by bind group index. A caller-provided format occupies group 0;
20+
* auto-reflected resources occupy their own group (0 when no caller format, otherwise 1).
21+
* The array is dense (no gaps), as required by WebGPU pipeline layouts.
22+
*
23+
* @type {BindGroup[]}
24+
*/
25+
bindGroups = [];
2026

2127
constructor(compute) {
2228
this.compute = compute;
@@ -25,27 +31,57 @@ class WebgpuCompute {
2531

2632
DebugGraphics.pushGpuMarker(device, `Compute:${compute.name}`);
2733

28-
// create bind group
29-
const { computeBindGroupFormat, computeUniformBufferFormats } = shader.impl;
30-
Debug.assert(computeBindGroupFormat, 'Compute shader does not have computeBindGroupFormat specified', shader);
31-
32-
// this.bindGroup = new BindGroup(device, computeBindGroupFormat, this.uniformBuffer);
33-
this.bindGroup = new BindGroup(device, computeBindGroupFormat);
34-
DebugHelper.setName(this.bindGroup, `Compute-BindGroup_${this.bindGroup.id}`);
35-
36-
if (computeUniformBufferFormats) {
37-
for (const name in computeUniformBufferFormats) {
38-
if (computeUniformBufferFormats.hasOwnProperty(name)) {
39-
// TODO: investigate implications of using a non-persistent uniform buffer
40-
const ub = new UniformBuffer(device, computeUniformBufferFormats[name], true);
41-
this.uniformBuffers.push(ub);
42-
this.bindGroup.setUniformBuffer(name, ub);
34+
const {
35+
computeBindGroupFormat, computeUniformBufferFormats,
36+
computeReflectedBindGroupFormat, computeReflectedUniformBufferFormat,
37+
computeReflectedGroupIndex
38+
} = shader.impl;
39+
40+
// caller uniform buffers are bound into the caller bind group, so the format is required
41+
Debug.assert(!computeUniformBufferFormats || computeBindGroupFormat,
42+
'Compute shader specifies computeUniformBufferFormats but no computeBindGroupFormat to bind them into', shader);
43+
44+
// ordered, gapless array of bind group formats (array index === bind group index)
45+
const formats = [];
46+
47+
// group 0: caller-provided resources (if any)
48+
if (computeBindGroupFormat) {
49+
const bindGroup = new BindGroup(device, computeBindGroupFormat);
50+
DebugHelper.setName(bindGroup, `Compute-BindGroup_${bindGroup.id}`);
51+
52+
if (computeUniformBufferFormats) {
53+
for (const name in computeUniformBufferFormats) {
54+
if (computeUniformBufferFormats.hasOwnProperty(name)) {
55+
// TODO: investigate implications of using a non-persistent uniform buffer
56+
const ub = new UniformBuffer(device, computeUniformBufferFormats[name], true);
57+
this.uniformBuffers.push(ub);
58+
bindGroup.setUniformBuffer(name, ub);
59+
}
4360
}
4461
}
62+
63+
formats[0] = computeBindGroupFormat;
64+
this.bindGroups[0] = bindGroup;
65+
}
66+
67+
// auto-reflected resources, at their own bind group (0 when no caller format, otherwise 1)
68+
if (computeReflectedBindGroupFormat) {
69+
const reflectedBindGroup = new BindGroup(device, computeReflectedBindGroupFormat);
70+
DebugHelper.setName(reflectedBindGroup, `Compute-ReflectedBindGroup_${reflectedBindGroup.id}`);
71+
72+
if (computeReflectedUniformBufferFormat) {
73+
// matches the generated 'ub_compute' uniform buffer (see WebgpuShaderProcessorWGSL.runCompute)
74+
const ub = new UniformBuffer(device, computeReflectedUniformBufferFormat, true);
75+
this.uniformBuffers.push(ub);
76+
reflectedBindGroup.setUniformBuffer('ub_compute', ub);
77+
}
78+
79+
formats[computeReflectedGroupIndex] = computeReflectedBindGroupFormat;
80+
this.bindGroups[computeReflectedGroupIndex] = reflectedBindGroup;
4581
}
4682

4783
// pipeline
48-
this.pipeline = device.computePipeline.get(shader, computeBindGroupFormat);
84+
this.pipeline = device.computePipeline.get(shader, formats);
4985

5086
DebugGraphics.popGpuMarker(device);
5187
}
@@ -55,23 +91,27 @@ class WebgpuCompute {
5591
this.uniformBuffers.forEach(ub => ub.destroy());
5692
this.uniformBuffers.length = 0;
5793

58-
this.bindGroup.destroy();
59-
this.bindGroup = null;
94+
this.bindGroups.forEach(bindGroup => bindGroup.destroy());
95+
this.bindGroups.length = 0;
6096
}
6197

6298
updateBindGroup() {
6399

64100
// bind group data
65-
const { bindGroup } = this;
66-
bindGroup.updateUniformBuffers();
67-
bindGroup.update();
101+
for (let i = 0; i < this.bindGroups.length; i++) {
102+
const bindGroup = this.bindGroups[i];
103+
bindGroup.updateUniformBuffers();
104+
bindGroup.update();
105+
}
68106
}
69107

70108
dispatch(x, y, z) {
71109

72-
// bind group
110+
// bind groups
73111
const device = this.compute.device;
74-
device.setBindGroup(0, this.bindGroup);
112+
for (let i = 0; i < this.bindGroups.length; i++) {
113+
device.setBindGroup(i, this.bindGroups[i]);
114+
}
75115

76116
// compute pipeline
77117
const passEncoder = device.passEncoder;

0 commit comments

Comments
 (0)