do not over-dispatch

Author: reczkok

apps/typegpu-docs/src/examples/algorithms/mnist-inference/index.ts

@@ -4,6 +4,7 @@ import { downloadLayers } from './helpers.ts';
 import { defineControls } from '../../common/defineControls.ts';
 
 const SIZE = 28;
+const WORKGROUP_SIZE = 64;
 
 const root = await tgpu.init({
   device: { optionalFeatures: ['timestamp-query', 'subgroups', 'shader-f16'] },
@@ -39,14 +40,15 @@ function relu(x: number): number {
 
 const defaultCompute = tgpu.computeFn({
   in: { gid: d.builtin.globalInvocationId },
-  workgroupSize: [64],
+  workgroupSize: [WORKGROUP_SIZE],
 })(({ gid }) => {
   const i = gid.x;
-  const inputSize = ioLayout.$.input.length;
-  if (i >= inputSize) {
+  const outLen = ioLayout.$.output.length;
+  if (i >= outLen) {
     return;
   }
 
+  const inputSize = ioLayout.$.input.length;
   const weightsOffset = i * inputSize;
   let sum = float();
 
@@ -65,7 +67,7 @@ const subgroupCompute = tgpu.computeFn({
     sgid: d.builtin.subgroupId,
     nsg: d.builtin.numSubgroups,
   },
-  workgroupSize: [64],
+  workgroupSize: [WORKGROUP_SIZE],
 })(({ wid, sid, sgid, nsg }) => {
   const outLen = ioLayout.$.output.length;
   const inputSize = ioLayout.$.input.length;
@@ -155,7 +157,8 @@ function createNetwork(layers: [LayerData, LayerData][]): Network {
     }
     input.write(data);
 
-    const pipeline = useSubgroups && pipelines.subgroup ? pipelines.subgroup : pipelines.default;
+    const subgroupPipeline = useSubgroups ? pipelines.subgroup : null;
+    const pipeline = subgroupPipeline ?? pipelines.default;
 
     // Run the network
     for (let i = 0; i < buffers.length; i++) {
@@ -173,7 +176,10 @@ function createNetwork(layers: [LayerData, LayerData][]): Network {
         boundPipeline = boundPipeline.withTimestampWrites(descriptor);
       }
 
-      boundPipeline.dispatchWorkgroups(buffers[i].biases.dataType.elementCount);
+      const outputCount = buffers[i].biases.dataType.elementCount;
+      boundPipeline.dispatchWorkgroups(
+        subgroupPipeline ? outputCount : Math.ceil(outputCount / WORKGROUP_SIZE),
+      );
     }
 
     if (querySet?.available) {

apps/typegpu-docs/tests/individual-example-tests/mnist-inference.test.ts

@@ -26,24 +26,25 @@ describe('mnist inference example', () => {
       "enable subgroups;
       enable f16;
 
+      @group(0) @binding(1) var<storage, read_write> output: array<f32>;
+
       @group(0) @binding(0) var<storage, read> input: array<f32>;
 
       @group(1) @binding(0) var<storage, read> weights: array<f32>;
 
       @group(1) @binding(1) var<storage, read> biases: array<f32>;
 
-      @group(0) @binding(1) var<storage, read_write> output: array<f32>;
-
       fn relu(x: f32) -> f32 {
         return max(0f, x);
       }
 
       @compute @workgroup_size(64) fn defaultCompute(@builtin(global_invocation_id) gid: vec3u) {
         let i = gid.x;
-        let inputSize = arrayLength(&input);
-        if ((i >= inputSize)) {
+        let outLen = arrayLength(&output);
+        if ((i >= outLen)) {
           return;
         }
+        let inputSize = arrayLength(&input);
         let weightsOffset = (i * inputSize);
         var sum = 0f;
         for (var j = 0u; (j < inputSize); j++) {