From 3f8d212c0ebf35eaa05858444ddddab73f5bfb4c Mon Sep 17 00:00:00 2001 From: reczkok Date: Tue, 4 Mar 2025 17:08:21 +0100 Subject: [PATCH 1/4] add subgroup mnist --- .../algorithms/mnist-inference/index.ts | 66 +++++++++++++++++-- 1 file changed, 61 insertions(+), 5 deletions(-) diff --git a/apps/typegpu-docs/src/content/examples/algorithms/mnist-inference/index.ts b/apps/typegpu-docs/src/content/examples/algorithms/mnist-inference/index.ts index 6304c4825..08af03ed8 100644 --- a/apps/typegpu-docs/src/content/examples/algorithms/mnist-inference/index.ts +++ b/apps/typegpu-docs/src/content/examples/algorithms/mnist-inference/index.ts @@ -5,16 +5,18 @@ const SIZE = 28; const root = await tgpu.init({ device: { - optionalFeatures: ['timestamp-query'], + optionalFeatures: ['timestamp-query', 'subgroups' as GPUFeatureName], }, }); const hasTimestampQuery = root.enabledFeatures.has('timestamp-query'); +const hasSubgroups = root.enabledFeatures.has('subgroups' as GPUFeatureName); const device = root.device; + const canvasData = new Array(SIZE ** 2).fill(0); // Shader code -const layerShader = ` +const fallbackShader = ` @binding(0) @group(0) var input: array; @binding(1) @group(0) var output: array; @@ -38,8 +40,62 @@ const layerShader = ` sum = sum + input[j] * weights[weightsOffset + j]; } - sum = sum + biases[i]; - output[i] = relu(sum); + let total = sum + biases[i]; + output[i] = relu(total); + } +`; + +const subgroupShader = ` + enable subgroups; + + @binding(0) @group(0) var input: array; + @binding(1) @group(0) var output: array; + + @binding(0) @group(1) var weights: array; + @binding(1) @group(1) var biases: array; + + fn relu(x: f32) -> f32 { + return max(0.0, x); + } + + // WebGPU guarantees a subgroup size of at least 4 + var subgroupSums: array; + + @compute @workgroup_size(64) + fn main( + @builtin(local_invocation_id) lid: vec3u, + @builtin(workgroup_id) wid: vec3u, + @builtin(subgroup_invocation_id) sid: u32, + @builtin(subgroup_size) ssize: u32 + ) { + let neuronIndex = wid.x; + let inputSize = arrayLength(&input); + let weightsOffset = neuronIndex * inputSize; + + var partial: f32 = 0.0; + for (var j = lid.x; j < inputSize; j = j + 64) { + partial = partial + input[j] * weights[weightsOffset + j]; + } + + let subgroupSum = subgroupAdd(partial); + let subgroupId = lid.x / ssize; + + let numSubgroups = 64 / ssize; + + if (sid == 0u) { + subgroupSums[subgroupId] = subgroupSum; + } + + workgroupBarrier(); + + var total: f32 = 0.0; + if (lid.x == 0u) { + for (var i = 0u; i < numSubgroups; i = i + 1u) { + total = total + subgroupSums[i]; + } + total = total + biases[neuronIndex]; + output[neuronIndex] = relu(total); + } } `; @@ -69,7 +125,7 @@ const pipeline = device.createComputePipeline({ }), compute: { module: device.createShaderModule({ - code: layerShader, + code: hasSubgroups ? subgroupShader : fallbackShader, }), }, }); From 1c2a623f0f0ceabec2847a32965c0854ef7e8342 Mon Sep 17 00:00:00 2001 From: reczkok Date: Tue, 17 Jun 2025 18:27:37 +0200 Subject: [PATCH 2/4] Add UI for subgroups status and inference time in MNIST example --- .../algorithms/mnist-inference/index.html | 39 +++++ .../algorithms/mnist-inference/index.ts | 135 ++++++++++++------ 2 files changed, 131 insertions(+), 43 deletions(-) diff --git a/apps/typegpu-docs/src/content/examples/algorithms/mnist-inference/index.html b/apps/typegpu-docs/src/content/examples/algorithms/mnist-inference/index.html index bb181ae39..5c179ebab 100644 --- a/apps/typegpu-docs/src/content/examples/algorithms/mnist-inference/index.html +++ b/apps/typegpu-docs/src/content/examples/algorithms/mnist-inference/index.html @@ -16,6 +16,11 @@
8
9
+ +
+
Subgroups: -
+
Inference: -
+
@@ -66,6 +71,40 @@ } } + .info { + width: 100%; + padding: 0.75rem; + background: #f8f9fa; + border-radius: 0.5rem; + font-size: 0.875rem; + display: flex; + flex-direction: column; + gap: 0.25rem; + @media (max-width: 1024px) { + font-size: 0.75rem; + } + } + + .info div { + display: flex; + justify-content: space-between; + color: #64748b; + } + + .info span { + font-family: 'Monaco', monospace; + font-weight: 600; + color: #1e293b; + } + + .info .enabled { + color: #16a34a; + } + + .info .disabled { + color: #dc2626; + } + .bar { position: relative; height: 1.25rem; diff --git a/apps/typegpu-docs/src/content/examples/algorithms/mnist-inference/index.ts b/apps/typegpu-docs/src/content/examples/algorithms/mnist-inference/index.ts index 08af03ed8..719f02775 100644 --- a/apps/typegpu-docs/src/content/examples/algorithms/mnist-inference/index.ts +++ b/apps/typegpu-docs/src/content/examples/algorithms/mnist-inference/index.ts @@ -10,19 +10,35 @@ const root = await tgpu.init({ }); const hasTimestampQuery = root.enabledFeatures.has('timestamp-query'); const hasSubgroups = root.enabledFeatures.has('subgroups' as GPUFeatureName); +let useSubgroups = hasSubgroups; const device = root.device; const canvasData = new Array(SIZE ** 2).fill(0); -// Shader code +const ReadonlyFloats = { + storage: (n: number) => d.arrayOf(d.f32, n), + access: 'readonly', +} as const; + +const MutableFloats = { + storage: (n: number) => d.arrayOf(d.f32, n), + access: 'mutable', +} as const; + +const ioLayout = tgpu.bindGroupLayout({ + input: ReadonlyFloats, + output: MutableFloats, +}).$idx(0); -const fallbackShader = ` - @binding(0) @group(0) var input: array; - @binding(1) @group(0) var output: array; +const weightsBiasesLayout = tgpu.bindGroupLayout({ + weights: ReadonlyFloats, + biases: ReadonlyFloats, +}).$idx(1); - @binding(0) @group(1) var weights: array; - @binding(1) @group(1) var biases: array; +// Shader code +const fallbackShader = tgpu.resolve({ + template: ` fn relu(x: f32) -> f32 { return max(0.0, x); } @@ -43,17 +59,17 @@ const fallbackShader = ` let total = sum + biases[i]; output[i] = relu(total); } -`; - -const subgroupShader = ` - enable subgroups; - - @binding(0) @group(0) var input: array; - @binding(1) @group(0) var output: array; - - @binding(0) @group(1) var weights: array; - @binding(1) @group(1) var biases: array; +`, + externals: { + ...weightsBiasesLayout.bound, + ...ioLayout.bound, + }, +}); +const subgroupShader = `enable subgroups; + ${ + tgpu.resolve({ + template: ` fn relu(x: f32) -> f32 { return max(0.0, x); } @@ -97,35 +113,21 @@ const subgroupShader = ` output[neuronIndex] = relu(total); } } -`; - -const ReadonlyFloats = { - storage: (n: number) => d.arrayOf(d.f32, n), - access: 'readonly', -} as const; - -const MutableFloats = { - storage: (n: number) => d.arrayOf(d.f32, n), - access: 'mutable', -} as const; - -const ioLayout = tgpu.bindGroupLayout({ - input: ReadonlyFloats, - output: MutableFloats, -}); - -const weightsBiasesLayout = tgpu.bindGroupLayout({ - weights: ReadonlyFloats, - biases: ReadonlyFloats, -}); - -const pipeline = device.createComputePipeline({ +`, + externals: { + ...weightsBiasesLayout.bound, + ...ioLayout.bound, + }, + }) +}`; + +let pipeline = device.createComputePipeline({ layout: device.createPipelineLayout({ bindGroupLayouts: [root.unwrap(ioLayout), root.unwrap(weightsBiasesLayout)], }), compute: { module: device.createShaderModule({ - code: hasSubgroups ? subgroupShader : fallbackShader, + code: useSubgroups ? subgroupShader : fallbackShader, }), }, }); @@ -237,9 +239,12 @@ function createNetwork(layers: [LayerData, LayerData][]): Network { if (querySet?.available) { querySet.resolve(); const results = await querySet.read(); - console.log( - `Inference took ${Number(results[1] - results[0]) / 1_000_000} ms`, - ); + const inferenceTimeMs = Number(results[1] - results[0]) / 1_000_000; + console.log(`Inference took ${inferenceTimeMs} ms`); + + inferenceTimeEl.textContent = `${inferenceTimeMs.toFixed(2)} ms`; + } else { + inferenceTimeEl.textContent = 'N/A'; } // Read the output @@ -254,6 +259,22 @@ function createNetwork(layers: [LayerData, LayerData][]): Network { }; } +const recreatePipeline = () => { + pipeline = device.createComputePipeline({ + layout: device.createPipelineLayout({ + bindGroupLayouts: [ + root.unwrap(ioLayout), + root.unwrap(weightsBiasesLayout), + ], + }), + compute: { + module: device.createShaderModule({ + code: useSubgroups ? subgroupShader : fallbackShader, + }), + }, + }); +}; + const network = createNetwork(await downloadLayers()); // #region Downloading weights & biases @@ -328,6 +349,12 @@ const canvas = document.querySelector('canvas') as HTMLCanvasElement; const context = canvas.getContext('2d') as CanvasRenderingContext2D; const bars = Array.from(document.querySelectorAll('.bar')) as HTMLDivElement[]; +const subgroupsEl = document.getElementById( + 'subgroups-status', +) as HTMLSpanElement; +const inferenceTimeEl = document.getElementById( + 'inference-time', +) as HTMLSpanElement; const uiState = { isDrawing: false, @@ -385,6 +412,20 @@ function run() { } document.querySelector('.loading')?.classList.add('loaded'); + +function updateSubgroupsStatus() { + const text = !hasSubgroups + ? 'Not Supported' + : useSubgroups + ? 'Enabled' + : 'Disabled'; + const cls = !hasSubgroups || !useSubgroups ? 'disabled' : 'enabled'; + subgroupsEl.textContent = text; + subgroupsEl.className = cls; +} + +updateSubgroupsStatus(); + run(); canvas.addEventListener('mousedown', () => { @@ -511,6 +552,14 @@ export const controls = { Reset: { onButtonClick: resetDrawing, }, + 'Use Subgroups': { + initial: hasSubgroups, + onToggleChange: (value: boolean) => { + useSubgroups = value; + recreatePipeline(); + updateSubgroupsStatus(); + }, + }, }; // #endregion From d7fc5e96d200fc7d1d6bd006b9063be9c85af8ab Mon Sep 17 00:00:00 2001 From: reczkok Date: Mon, 23 Jun 2025 10:56:31 +0200 Subject: [PATCH 3/4] Add "subgroups" tag --- .../src/content/examples/algorithms/mnist-inference/meta.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apps/typegpu-docs/src/content/examples/algorithms/mnist-inference/meta.json b/apps/typegpu-docs/src/content/examples/algorithms/mnist-inference/meta.json index 62b0ec151..da25cb794 100644 --- a/apps/typegpu-docs/src/content/examples/algorithms/mnist-inference/meta.json +++ b/apps/typegpu-docs/src/content/examples/algorithms/mnist-inference/meta.json @@ -1,5 +1,5 @@ { "title": "MNIST Inference", "category": "algorithms", - "tags": ["ai", "compute", "inference", "timestamp query"] + "tags": ["ai", "compute", "inference", "timestamp query", "subgroups"] } From b7d4c5809428f9c01d781055e3110e00ca480a54 Mon Sep 17 00:00:00 2001 From: reczkok Date: Mon, 23 Jun 2025 11:02:12 +0200 Subject: [PATCH 4/4] Refactor responsive styles for MNIST inference example --- .../algorithms/mnist-inference/index.html | 42 ++++++++++++++----- 1 file changed, 32 insertions(+), 10 deletions(-) diff --git a/apps/typegpu-docs/src/content/examples/algorithms/mnist-inference/index.html b/apps/typegpu-docs/src/content/examples/algorithms/mnist-inference/index.html index 5c179ebab..a939bc314 100644 --- a/apps/typegpu-docs/src/content/examples/algorithms/mnist-inference/index.html +++ b/apps/typegpu-docs/src/content/examples/algorithms/mnist-inference/index.html @@ -54,9 +54,6 @@ .predictions-label { margin-bottom: 0.5rem; font-size: 1.25rem; - @media (max-width: 1024px) { - font-size: 1rem; - } } .bars-container { @@ -65,10 +62,6 @@ width: 100%; justify-content: flex-start; row-gap: 0.5rem; - @media (max-width: 1024px) { - width: calc(100% + 8rem); - row-gap: 0.1rem; - } } .info { @@ -80,8 +73,34 @@ display: flex; flex-direction: column; gap: 0.25rem; - @media (max-width: 1024px) { - font-size: 0.75rem; + } + + @media (max-width: 1024px) { + .predictions-container { + width: 100%; + gap: 0.5rem; + } + + .predictions-label { + font-size: 1rem; + } + + .predictions-container .bars-container, + .predictions-container .info { + display: inline-block; + vertical-align: top; + } + + .bars-container { + width: calc(100% - 9rem); + row-gap: 0.1rem; + } + + .info { + width: 8rem; + padding: 0.5rem; + font-size: 0.65rem; + margin-left: 1rem; } } @@ -113,7 +132,10 @@ font-size: 1rem; background: linear-gradient(to right, transparent, #e6e6f2); border-radius: 9999px; - @media (max-width: 1024px) { + } + + @media (max-width: 1024px) { + .bar { height: 1rem; font-size: 0.75rem; }