sync 6fdf6301e2bb

Browse files

Files changed (9) hide show

README.md +9 -2
build/webgpu/bench.json +97 -1
build/webgpu/manifest.json +46 -17
build/webgpu/metadata.json +12 -10
build/webgpu/scan-multichunk-apply.wgsl.jinja +7 -2
build/webgpu/scan-multichunk-carries.wgsl.jinja +7 -2
build/webgpu/scan-prefix-sum.wgsl.jinja +7 -3
build/webgpu/scan-strided-row.wgsl.jinja +169 -0
build/webgpu/test.json +238 -3

README.md CHANGED Viewed

@@ -44,22 +44,29 @@ Default values (overridable per request):
 | --- | --- |
 | `T` | `float32` |
 ## Files
 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
-- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 - [`scan-coop-channel.wgsl.jinja`](build/webgpu/scan-coop-channel.wgsl.jinja)
 - [`scan-multichunk-apply.wgsl.jinja`](build/webgpu/scan-multichunk-apply.wgsl.jinja)
 - [`scan-multichunk-carries.wgsl.jinja`](build/webgpu/scan-multichunk-carries.wgsl.jinja)
 - [`scan-multichunk-local.wgsl.jinja`](build/webgpu/scan-multichunk-local.wgsl.jinja)
 - [`scan-prefix-sum.wgsl.jinja`](build/webgpu/scan-prefix-sum.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
-npm install --save-exact @huggingface/kernels@0.0.1-preview.2
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

 | --- | --- |
 | `T` | `float32` |
+## Implementation variants
+One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
+- `vector_channel_prefix_sum` — Scan four adjacent state channels together with float32 accumulation; bound the power-of-two workgroup by the scan length, tuning and device resources, and use subgroups only when it contains complete groups.
 ## Files
 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark cases
 - [`scan-coop-channel.wgsl.jinja`](build/webgpu/scan-coop-channel.wgsl.jinja)
 - [`scan-multichunk-apply.wgsl.jinja`](build/webgpu/scan-multichunk-apply.wgsl.jinja)
 - [`scan-multichunk-carries.wgsl.jinja`](build/webgpu/scan-multichunk-carries.wgsl.jinja)
 - [`scan-multichunk-local.wgsl.jinja`](build/webgpu/scan-multichunk-local.wgsl.jinja)
 - [`scan-prefix-sum.wgsl.jinja`](build/webgpu/scan-prefix-sum.wgsl.jinja)
+- [`scan-strided-row.wgsl.jinja`](build/webgpu/scan-strided-row.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.3
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

build/webgpu/bench.json CHANGED Viewed

@@ -50,7 +50,7 @@
       "bench": { "metrics": [{ "type": "bandwidth", "value": "2 * 4 * numel(shapes.scan_input)" }] }
     },
     {
-      "name": "coop_healthy_dim512_steps4096",
       "preset": "smoke",
       "inputs": {
         "initial_state": { "dtype": "float32", "shape": [512], "dist": "normal", "seed": 753, "scale": 0.1 },
@@ -140,6 +140,102 @@
         "scan_output": { "dtype": "float32", "shape": [131072, 8] }
       },
       "bench": { "primary": true, "metrics": [{ "type": "bandwidth", "value": "2 * 4 * numel(shapes.scan_input)" }] }
     }
   ]
 }

       "bench": { "metrics": [{ "type": "bandwidth", "value": "2 * 4 * numel(shapes.scan_input)" }] }
     },
     {
+      "name": "coop_control_dim512_steps4096",
       "preset": "smoke",
       "inputs": {
         "initial_state": { "dtype": "float32", "shape": [512], "dist": "normal", "seed": 753, "scale": 0.1 },
         "scan_output": { "dtype": "float32", "shape": [131072, 8] }
       },
       "bench": { "primary": true, "metrics": [{ "type": "bandwidth", "value": "2 * 4 * numel(shapes.scan_input)" }] }
+    },
+    {
+      "name": "vector_channel_boundary_t63_d12_reverse0",
+      "attrs": { "reverse": 0 },
+      "inputs": {
+        "initial_state": { "dtype": "float32", "shape": [12], "dist": "normal", "seed": 1301, "scale": 0.1 },
+        "scan_input": { "dtype": "float32", "shape": [63, 12], "dist": "normal", "seed": 1302, "scale": 0.1 }
+      },
+      "outputs": {
+        "final_state": { "dtype": "float32", "shape": [12] },
+        "scan_output": { "dtype": "float32", "shape": [63, 12] }
+      },
+      "preset": "edge",
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "8 * numel(shapes.scan_input) + 8 * dim(shapes.initial_state, 0)" }]
+      }
+    },
+    {
+      "name": "vector_channel_boundary_t64_d12_reverse0",
+      "attrs": { "reverse": 0 },
+      "inputs": {
+        "initial_state": { "dtype": "float32", "shape": [12], "dist": "normal", "seed": 1301, "scale": 0.1 },
+        "scan_input": { "dtype": "float32", "shape": [64, 12], "dist": "normal", "seed": 1302, "scale": 0.1 }
+      },
+      "outputs": {
+        "final_state": { "dtype": "float32", "shape": [12] },
+        "scan_output": { "dtype": "float32", "shape": [64, 12] }
+      },
+      "preset": "edge",
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "8 * numel(shapes.scan_input) + 8 * dim(shapes.initial_state, 0)" }]
+      }
+    },
+    {
+      "name": "vector_channel_boundary_t65_d12_reverse1",
+      "attrs": { "reverse": 1 },
+      "inputs": {
+        "initial_state": { "dtype": "float32", "shape": [12], "dist": "normal", "seed": 1301, "scale": 0.1 },
+        "scan_input": { "dtype": "float32", "shape": [65, 12], "dist": "normal", "seed": 1302, "scale": 0.1 }
+      },
+      "outputs": {
+        "final_state": { "dtype": "float32", "shape": [12] },
+        "scan_output": { "dtype": "float32", "shape": [65, 12] }
+      },
+      "preset": "edge",
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "8 * numel(shapes.scan_input) + 8 * dim(shapes.initial_state, 0)" }]
+      }
+    },
+    {
+      "name": "vector_channel_boundary_t255_d4_reverse0",
+      "attrs": { "reverse": 0 },
+      "inputs": {
+        "initial_state": { "dtype": "float32", "shape": [4], "dist": "normal", "seed": 1301, "scale": 0.1 },
+        "scan_input": { "dtype": "float32", "shape": [255, 4], "dist": "normal", "seed": 1302, "scale": 0.1 }
+      },
+      "outputs": {
+        "final_state": { "dtype": "float32", "shape": [4] },
+        "scan_output": { "dtype": "float32", "shape": [255, 4] }
+      },
+      "preset": "edge",
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "8 * numel(shapes.scan_input) + 8 * dim(shapes.initial_state, 0)" }]
+      }
+    },
+    {
+      "name": "vector_channel_boundary_t257_d68_reverse1",
+      "attrs": { "reverse": 1 },
+      "inputs": {
+        "initial_state": { "dtype": "float32", "shape": [68], "dist": "normal", "seed": 1301, "scale": 0.1 },
+        "scan_input": { "dtype": "float32", "shape": [257, 68], "dist": "normal", "seed": 1302, "scale": 0.1 }
+      },
+      "outputs": {
+        "final_state": { "dtype": "float32", "shape": [68] },
+        "scan_output": { "dtype": "float32", "shape": [257, 68] }
+      },
+      "preset": "edge",
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "8 * numel(shapes.scan_input) + 8 * dim(shapes.initial_state, 0)" }]
+      }
+    },
+    {
+      "name": "vector_channel_boundary_t257_d69_reverse1",
+      "attrs": { "reverse": 1 },
+      "inputs": {
+        "initial_state": { "dtype": "float32", "shape": [69], "dist": "normal", "seed": 1301, "scale": 0.1 },
+        "scan_input": { "dtype": "float32", "shape": [257, 69], "dist": "normal", "seed": 1302, "scale": 0.1 }
+      },
+      "outputs": {
+        "final_state": { "dtype": "float32", "shape": [69] },
+        "scan_output": { "dtype": "float32", "shape": [257, 69] }
+      },
+      "preset": "edge",
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "8 * numel(shapes.scan_input) + 8 * dim(shapes.initial_state, 0)" }]
+      }
     }
   ]
 }

build/webgpu/manifest.json CHANGED Viewed

@@ -17,22 +17,36 @@
     "subgroupsWave32": "device.features.has(\"subgroups\") and wave32Adapter",
     "variableNarrowSubgroups": "device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize < device.adapterInfo.subgroupMaxSize and device.adapterInfo.subgroupMaxSize <= 16",
     "parallelSubgroupScan": "subgroupsWave32 or variableNarrowSubgroups",
-    "scanUseSubgroups": "parallelSubgroupScan"
   },
   "when": ["ranks.initial_state == 1", "ranks.scan_input == 2", "ranks.final_state == 1", "ranks.scan_output == 2", "dim(shapes.initial_state, 0) == dim(shapes.scan_input, 1)", "dim(shapes.final_state, 0) == dim(shapes.initial_state, 0)", "dim(shapes.scan_output, 0) == dim(shapes.scan_input, 0)", "dim(shapes.scan_output, 1) == dim(shapes.scan_input, 1)"],
   "bindings": {
-    "scan_input": { "buffer": "read-only-storage", "elementType": "f32" },
-    "scan_output": { "buffer": "storage", "elementType": "f32" },
     "params": {
-      "buffer": "uniform",
       "struct": [
         { "name": "steps", "type": "u32", "value": "dim(shapes.scan_input, 0)" },
         { "name": "dim", "type": "u32", "value": "dim(shapes.scan_input, 1)" },
         { "name": "reverse", "type": "u32", "value": "attrs.reverse" }
       ]
     },
-    "initial_state": { "buffer": "read-only-storage", "elementType": "f32" },
-    "final_state": { "buffer": "storage", "elementType": "f32" }
   },
   "variants": [
     {
@@ -53,14 +67,8 @@
           "name": "Scan.MultichunkLocal",
           "shader": "scan-multichunk-local.wgsl.jinja",
           "derive": { "useSubgroups": "scanUseSubgroups" },
-          "bindings": [
-            "scan_input",
-            "scan_output",
-            { "name": "chunkTotals", "buffer": "storage", "elementType": "f32" },
-            "params"
-          ],
-          "dispatch": { "x": "min(totalChunks, 65535)", "y": "ceilDiv(totalChunks, 65535)", "z": 1 },
-          "subgroupCollectivesWidth": "portable"
         },
         {
           "id": "carries",
@@ -69,7 +77,7 @@
           "bindings": [
             "initial_state",
             { "name": "chunkTotals", "buffer": "read-only-storage", "elementType": "f32" },
-            { "name": "chunkCarries", "buffer": "storage", "elementType": "f32" },
             "final_state",
             { "name": "params", "struct": [{ "name": "dim", "type": "u32", "value": "dim(shapes.scan_input, 1)" }] }
           ],
@@ -119,8 +127,7 @@
             "x": "min(dim(shapes.scan_input, 1), 65535)",
             "y": "ceilDiv(dim(shapes.scan_input, 1), 65535)",
             "z": 1
-          },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -139,6 +146,28 @@
           }
         }
       ]
     }
   ]
 }

     "subgroupsWave32": "device.features.has(\"subgroups\") and wave32Adapter",
     "variableNarrowSubgroups": "device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize < device.adapterInfo.subgroupMaxSize and device.adapterInfo.subgroupMaxSize <= 16",
     "parallelSubgroupScan": "subgroupsWave32 or variableNarrowSubgroups",
+    "scanUseSubgroups": "parallelSubgroupScan",
+    "vectorScanChannels": "dim(shapes.scan_input, 1) / 4",
+    "vectorScanWg": "pow2ceil(min(tunables.WORKGROUP_SIZE, pow2ceil(dim(shapes.scan_input, 0)), device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX, floor((device.limits.maxComputeWorkgroupStorageSize - 4) / 16)) + 1) / 2",
+    "vectorScanUseSubgroups": "scanUseSubgroups and vectorScanWg >= device.adapterInfo.subgroupMaxSize"
   },
   "when": ["ranks.initial_state == 1", "ranks.scan_input == 2", "ranks.final_state == 1", "ranks.scan_output == 2", "dim(shapes.initial_state, 0) == dim(shapes.scan_input, 1)", "dim(shapes.final_state, 0) == dim(shapes.initial_state, 0)", "dim(shapes.scan_output, 0) == dim(shapes.scan_input, 0)", "dim(shapes.scan_output, 1) == dim(shapes.scan_input, 1)"],
   "bindings": {
+    "scan_input": { "elementType": "f32" },
+    "scan_output": { "elementType": "f32" },
     "params": {
       "struct": [
         { "name": "steps", "type": "u32", "value": "dim(shapes.scan_input, 0)" },
         { "name": "dim", "type": "u32", "value": "dim(shapes.scan_input, 1)" },
         { "name": "reverse", "type": "u32", "value": "attrs.reverse" }
       ]
     },
+    "initial_state": { "elementType": "f32" },
+    "final_state": { "elementType": "f32" },
+    "vectorInitial": { "arg": "initial_state", "name": "initial_state", "elementType": "vec4<f32>" },
+    "vectorInput": { "arg": "scan_input", "name": "x", "elementType": "vec4<f32>" },
+    "vectorOutput": { "arg": "scan_output", "name": "y", "elementType": "vec4<f32>" },
+    "vectorFinal": { "arg": "final_state", "name": "final_state", "elementType": "vec4<f32>" },
+    "vectorParams": {
+      "name": "params",
+      "struct": [
+        { "name": "rows", "type": "u32", "value": "vectorScanChannels" },
+        { "name": "dim", "type": "u32", "value": "dim(shapes.scan_input, 0)" },
+        { "name": "inner", "type": "u32", "value": "vectorScanChannels" }
+      ]
+    }
   },
   "variants": [
     {
           "name": "Scan.MultichunkLocal",
           "shader": "scan-multichunk-local.wgsl.jinja",
           "derive": { "useSubgroups": "scanUseSubgroups" },
+          "bindings": ["scan_input", "scan_output", { "name": "chunkTotals", "elementType": "f32" }, "params"],
+          "dispatch": { "x": "min(totalChunks, 65535)", "y": "ceilDiv(totalChunks, 65535)", "z": 1 }
         },
         {
           "id": "carries",
           "bindings": [
             "initial_state",
             { "name": "chunkTotals", "buffer": "read-only-storage", "elementType": "f32" },
+            { "name": "chunkCarries", "elementType": "f32" },
             "final_state",
             { "name": "params", "struct": [{ "name": "dim", "type": "u32", "value": "dim(shapes.scan_input, 1)" }] }
           ],
             "x": "min(dim(shapes.scan_input, 1), 65535)",
             "y": "ceilDiv(dim(shapes.scan_input, 1), 65535)",
             "z": 1
+          }
         }
       ]
     },
           }
         }
       ]
+    },
+    {
+      "id": "vector_channel_prefix_sum",
+      "priority": 20,
+      "when": ["dim(shapes.scan_input, 1) >= 4", "dim(shapes.scan_input, 1) % 4 == 0", "dim(shapes.scan_input, 0) >= 64", "vectorScanWg >= 1", "ceilDiv(vectorScanChannels, DISPATCH_FOLD_WIDTH) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
+      "derive": {
+        "scalar": "\"vec4<f32>\"",
+        "accType": "\"vec4<f32>\"",
+        "workgroupSize": "vectorScanWg",
+        "reverse": "attrs.reverse != 0",
+        "useSubgroups": "vectorScanUseSubgroups"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "Scan.VectorChannelPrefix",
+          "shader": "scan-strided-row.wgsl.jinja",
+          "bindings": ["vectorInitial", "vectorInput", "vectorOutput", "vectorFinal", "vectorParams"],
+          "dispatch": { "x": "min(vectorScanChannels, 65535)", "y": "ceilDiv(vectorScanChannels, 65535)", "z": 1 }
+        }
+      ],
+      "demoteWhen": ["dim(shapes.scan_input, 1) < vectorScanWg"]
     }
   ]
 }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,29 +1,31 @@
 {
   "name": "ai.onnx.Scan",
-  "id": "_ai_onnx_scan_webgpu_0844f69",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "tElCZSCAVE3/XQ0WO3gJ4ysjWou60Gx57WXs/iIiV70=",
-      "manifest.json": "3nSOTULQv+1MDyOHGRakr7ahH/iELdpqI4TMW0Y/4Ac=",
       "scan-coop-channel.wgsl.jinja": "IoZdnMesPo0LyCYHIQ1NahvE0+BUC9hGFOB8kbKAAE0=",
-      "scan-multichunk-apply.wgsl.jinja": "0xcAy1mf28ZOc4VdaIx9ZBFRLyWwUk+ukwd9hkGoahw=",
-      "scan-multichunk-carries.wgsl.jinja": "Cd9J4zzmy935VHhL6BOPGIVrK0daOx2gQgfkMDlLZQg=",
       "scan-multichunk-local.wgsl.jinja": "E7BNf1L/PxxwcNFCdEfoWf1Zd7BF9CDD0GnPuQ+VinY=",
-      "scan-prefix-sum.wgsl.jinja": "Z5kIBUtpZ99FlX3aMLLkCWZ1bc2JsLjcfTnN30kXId0=",
-      "test.json": "jL0GWqv/vJge/cj98AfdAK4Bh4AATOcIykdfuVZsbnk="
     }
   },
-  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
   "webgpu": {
-    "manifestSpec": "2.0",
     "variants": {
       "multichunk_small_state": ["scan-multichunk-apply.wgsl.jinja", "scan-multichunk-carries.wgsl.jinja", "scan-multichunk-local.wgsl.jinja"],
       "coop_channel_prefix_sum": ["scan-coop-channel.wgsl.jinja"],
-      "lowered_prefix_sum": ["scan-prefix-sum.wgsl.jinja"]
     }
   }
 }

 {
   "name": "ai.onnx.Scan",
+  "id": "_ai_onnx_scan_webgpu_9e49b5b",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "JCxeUmlGL+dliqxE/n7GRey9/Yf0V0QkCzOarWmaurg=",
+      "manifest.json": "gmnmQ85vpC5piLIt5SAtrHXLOtCL7uN1py6wJ4LCeyI=",
       "scan-coop-channel.wgsl.jinja": "IoZdnMesPo0LyCYHIQ1NahvE0+BUC9hGFOB8kbKAAE0=",
+      "scan-multichunk-apply.wgsl.jinja": "yyZ1xgBfo8oVd2kBGbgmEeKtKXwmPRVuYdAGie6AAZs=",
+      "scan-multichunk-carries.wgsl.jinja": "1iwwgIXERfLk8H/ZbAx5G7MA7hfPCLDpAZLqzbgHmLU=",
       "scan-multichunk-local.wgsl.jinja": "E7BNf1L/PxxwcNFCdEfoWf1Zd7BF9CDD0GnPuQ+VinY=",
+      "scan-prefix-sum.wgsl.jinja": "OvW6G/HRG2pyeNrDHFfsvrzN3U7MAUavQnaGXOMU7vw=",
+      "scan-strided-row.wgsl.jinja": "ywymN5rjP420WNzHOIgBF4dDH6a2lxC0SEV9DpWtq3k=",
+      "test.json": "8vWkrtFOsRfBxwiXhIPyZAxC18aH1kt6lcXBCRPxmH8="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "variants": {
       "multichunk_small_state": ["scan-multichunk-apply.wgsl.jinja", "scan-multichunk-carries.wgsl.jinja", "scan-multichunk-local.wgsl.jinja"],
       "coop_channel_prefix_sum": ["scan-coop-channel.wgsl.jinja"],
+      "lowered_prefix_sum": ["scan-prefix-sum.wgsl.jinja"],
+      "vector_channel_prefix_sum": ["scan-strided-row.wgsl.jinja"]
     }
   }
 }

build/webgpu/scan-multichunk-apply.wgsl.jinja CHANGED Viewed

@@ -1,3 +1,9 @@
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
@@ -5,8 +11,7 @@ const CHUNKS: u32 = {{ chunks }}u;
 @compute @workgroup_size(WG)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-  let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
-  if (i >= params.count) { return; }
   let t = i / params.dim;
   let channel = i % params.dim;
   let iter = select(t, params.steps - 1u - t, params.reverse != 0u);

+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
+  if ({{ name }} >= {{ bound }}) { return; }{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
 @compute @workgroup_size(WG)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d("WG", guardInline=true) }}
   let t = i / params.dim;
   let channel = i % params.dim;
   let iter = select(t, params.steps - 1u - t, params.reverse != 0u);

build/webgpu/scan-multichunk-carries.wgsl.jinja CHANGED Viewed

@@ -1,11 +1,16 @@
 {{ env.wgsl.resourceDeclarations }}
 const CHUNKS: u32 = {{ chunks }}u;
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-  let channel = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
-  if (channel >= params.dim) { return; }
   var carry = initial_state[channel];
   let base = channel * CHUNKS;
   for (var chunk = 0u; chunk < CHUNKS; chunk += 1u) {

+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
+  if ({{ name }} >= {{ bound }}) { return; }{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 const CHUNKS: u32 = {{ chunks }}u;
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d(tunables.WORKGROUP_SIZE, "channel", "params.dim", guardInline=true) }}
   var carry = initial_state[channel];
   let base = channel * CHUNKS;
   for (var chunk = 0u; chunk < CHUNKS; chunk += 1u) {

build/webgpu/scan-prefix-sum.wgsl.jinja CHANGED Viewed

@@ -1,3 +1,9 @@
 {{ env.wgsl.resourceDeclarations }}
 // One invocation owns each channel and runs its prefix sum serially. Channels
@@ -5,9 +11,7 @@
 // order remain sequential.
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-  // gid.y carries channel-index bits past the per-dimension dispatch limit.
-  let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
-  if (i >= params.dim) { return; }
   var state = initial_state[i];
   for (var iter = 0u; iter < params.steps; iter = iter + 1u) {
     let t = select(iter, params.steps - 1u - iter, params.reverse != 0u);

+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
+  if ({{ name }} >= {{ bound }}) { return; }{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 // One invocation owns each channel and runs its prefix sum serially. Channels
 // order remain sequential.
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d(tunables.WORKGROUP_SIZE, bound="params.dim", guardInline=true) }}
   var state = initial_state[i];
   for (var iter = 0u; iter < params.steps; iter = iter + 1u) {
     let t = select(iter, params.steps - 1u - iter, params.reverse != 0u);

build/webgpu/scan-strided-row.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,169 @@

+// Workgroup-per-row chunked prefix scan along a non-innermost axis
+// (inner > 1). One workgroup owns one (outer, inner) row and scans the dim-length
+// axis cooperatively in WORKGROUP_SIZE chunks. Splitting each row across the
+// workgroup exposes parallelism when the tensor has few rows and a long axis.
+//
+// The row's element d is `(outer_i * dim + d) * inner + inner_i`, so successive
+// scan elements are `inner` apart. The vec4 mode loads four adjacent columns
+// per invocation. Distinct workgroups do not share a SIMD memory transaction.
+// Short axes use a scan-length-sized workgroup to avoid idle padded lanes.
+//
+// Floats accumulate in f32 even for f16 storage, and identity padding leaves
+// partial chunks unchanged.
+// Tree reduction reassociates floating-point sums relative to a serial loop;
+// wrapping integer addition remains associative.
+{% if useSubgroups %}
+enable subgroups;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+const WG: u32 = {{ workgroupSize }}u;
+const IDENTITY: {{ accType }} = {{ accType }}(0.0);
+fn combine(a: {{ accType }}, b: {{ accType }}) -> {{ accType }} {
+  return a + b;
+}
+{% if useSubgroups %}
+var<workgroup> sgLaneClaims: atomic<u32>;
+struct SubgroupLogicalLanes {
+  ltid: u32,
+  ord: u32,
+  rank: u32,
+  count: u32,
+}
+fn subgroup_logical_lanes(rank: u32, count: u32) -> SubgroupLogicalLanes {
+  // Every lane performs the atomic so no collective follows a lane guard: a
+  // subgroup op after a closed `if (rank == 0u)` block is a reconvergence
+  // hazard, since every subgroup operation must be reached in uniform control
+  // flow. Only the lane-0 lane adds its subgroup's claim, every other lane
+  // adds 0 and discards its snapshot. Lane 0 is the lowest lane, which is the
+  // lane `subgroupBroadcastFirst` reads.
+  let ticket = atomicAdd(&sgLaneClaims, select(0u, count | (1u << 16u), rank == 0u));
+  let claim = subgroupBroadcastFirst(ticket);
+  return SubgroupLogicalLanes((claim & 0xffffu) + rank, claim >> 16u, rank, count);
+}
+fn subgroup_logical_count() -> u32 {
+  return atomicLoad(&sgLaneClaims) >> 16u;
+}
+// One partial per subgroup, addressed by its logical ordinal. A workgroup of WG
+// invocations holds at most WG subgroups at any adapter width, so WG is a sound
+// bound with no width guard.
+var<workgroup> sgTotals: array<{{ accType }}, WG>;
+fn subgroupAddRows(v: vec4<f32>) -> vec4<f32> {
+  return vec4<f32>(subgroupAdd(v.x), subgroupAdd(v.y), subgroupAdd(v.z), subgroupAdd(v.w));
+}
+fn subgroupScanRows(v: vec4<f32>) -> vec4<f32> {
+  return vec4<f32>(subgroupInclusiveAdd(v.x), subgroupInclusiveAdd(v.y), subgroupInclusiveAdd(v.z), subgroupInclusiveAdd(v.w));
+}
+@compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(workgroup_id) wg: vec3<u32>,
+        @builtin(subgroup_invocation_id) sgLane: u32,
+        @builtin(subgroup_size) sgWidth: u32) {
+  // wg.y carries row-index bits past the per-dimension dispatch limit.
+  let row = wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u;
+  if (row >= params.rows) {
+    return;
+  }
+  let outer_i = row / params.inner;
+  let inner_i = row % params.inner;
+  let base = outer_i * params.dim * params.inner + inner_i;
+  // `ltid` is the axis position this invocation owns and `ord` its subgroup's
+  // carry ordinal; both are dense and independent of the device's partition.
+  let L = subgroup_logical_lanes(sgLane, sgWidth);
+  let tid = L.ltid;
+  workgroupBarrier();
+  let numSub = subgroup_logical_count();
+  var carry: {{ accType }} = initial_state[inner_i];
+  let chunks = (params.dim + WG - 1u) / WG;
+  for (var c = 0u; c < chunks; c = c + 1u) {
+    let j = c * WG + tid;
+    let valid = j < params.dim;
+{% if reverse %}
+    let d = params.dim - 1u - j;
+{% else %}
+    let d = j;
+{% endif %}
+    let src = base + d * params.inner;
+    var v: {{ accType }} = IDENTITY;
+    if (valid) {
+      v = {{ accType }}(x[src]);
+    }
+    let scannedSg = subgroupScanRows(v);
+    let totSg = subgroupAddRows(v);
+    if (L.rank == 0u) {
+      sgTotals[L.ord] = totSg;
+    }
+    workgroupBarrier();
+    var sgOffset: {{ accType }} = IDENTITY;
+    for (var s = 0u; s < L.ord; s = s + 1u) {
+      sgOffset = combine(sgOffset, sgTotals[s]);
+    }
+    var chunkTotal: {{ accType }} = IDENTITY;
+    for (var s = 0u; s < numSub; s = s + 1u) {
+      chunkTotal = combine(chunkTotal, sgTotals[s]);
+    }
+    let value = combine(carry, combine(sgOffset, scannedSg));
+    if (valid) {
+      y[src] = {{ scalar }}(value);
+    }
+    carry = combine(carry, chunkTotal);
+    workgroupBarrier();
+  }
+  if (tid == 0u) { final_state[inner_i] = carry; }
+}
+{% else %}
+var<workgroup> wgScan: array<{{ accType }}, WG>;
+@compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(workgroup_id) wg: vec3<u32>,
+        @builtin(local_invocation_id) lid: vec3<u32>) {
+  let row = wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u;
+  let tid = lid.x;
+  if (row >= params.rows) {
+    return;
+  }
+  let outer_i = row / params.inner;
+  let inner_i = row % params.inner;
+  let base = outer_i * params.dim * params.inner + inner_i;
+  var carry: {{ accType }} = initial_state[inner_i];
+  let chunks = (params.dim + WG - 1u) / WG;
+  for (var c = 0u; c < chunks; c = c + 1u) {
+    let j = c * WG + tid;
+    let valid = j < params.dim;
+{% if reverse %}
+    let d = params.dim - 1u - j;
+{% else %}
+    let d = j;
+{% endif %}
+    let src = base + d * params.inner;
+    var v: {{ accType }} = IDENTITY;
+    if (valid) {
+      v = {{ accType }}(x[src]);
+    }
+    workgroupBarrier();
+    wgScan[tid] = v;
+    for (var step = 1u; step < WG; step = step << 1u) {
+      workgroupBarrier();
+      var prev: {{ accType }} = IDENTITY;
+      if (tid >= step) {
+        prev = wgScan[tid - step];
+      }
+      workgroupBarrier();
+      wgScan[tid] = combine(prev, wgScan[tid]);
+    }
+    workgroupBarrier();
+    let value = combine(carry, wgScan[tid]);
+    if (valid) {
+      y[src] = {{ scalar }}(value);
+    }
+    carry = combine(carry, wgScan[WG - 1u]);
+  }
+  if (tid == 0u) { final_state[inner_i] = carry; }
+}
+{% endif %}

build/webgpu/test.json CHANGED Viewed

@@ -200,9 +200,9 @@
       }
     },
     {
-      "name": "coop_dim1025_above_old_fixed_gate",
       "provenance": {
-        "notes": "A 256-step prefix sum across 1,025 state lanes exercises cooperative dispatch beyond 1,024 lanes when the adapter limits admit it."
       },
       "inputs": {
         "initial_state": { "dtype": "float32", "shape": [1025] },
@@ -340,7 +340,7 @@
     {
       "name": "multichunk_small_state_dim2_steps4096",
       "provenance": {
-        "notes": "Compact correctness lock for a realistic long-stream recurrence with only two independent state lanes. The optimized path partitions the sequential axis into workgroup-sized chunks, scans chunk totals, then applies carries."
       },
       "inputs": {
         "initial_state": { "dtype": "float32", "shape": [2], "data": { "kind": "values", "values": [1.5, -2.0] } },
@@ -354,6 +354,241 @@
         "final_state": { "dtype": "float32", "shape": [2], "tolerance": 0.001, "relTolerance": 0.0001 },
         "scan_output": { "dtype": "float32", "shape": [4096, 2], "tolerance": 0.001, "relTolerance": 0.0001 }
       }
     }
   ]
 }

       }
     },
     {
+      "name": "coop_dim1025_above_fixed_gate",
       "provenance": {
+        "notes": "A 1025-element state, one past 1024, runs a 256-step prefix sum; checks correctness when the per-step state exceeds a 1024-lane count on adapters whose limits allow it."
       },
       "inputs": {
         "initial_state": { "dtype": "float32", "shape": [1025] },
     {
       "name": "multichunk_small_state_dim2_steps4096",
       "provenance": {
+        "notes": "A long sequence with two independent state lanes checks recurrent state across chunk boundaries."
       },
       "inputs": {
         "initial_state": { "dtype": "float32", "shape": [2], "data": { "kind": "values", "values": [1.5, -2.0] } },
         "final_state": { "dtype": "float32", "shape": [2], "tolerance": 0.001, "relTolerance": 0.0001 },
         "scan_output": { "dtype": "float32", "shape": [4096, 2], "tolerance": 0.001, "relTolerance": 0.0001 }
       }
+    },
+    {
+      "name": "vector_channel_full_lowered_prefix_sum_256x512",
+      "inputs": {
+        "initial_state": {
+          "dtype": "float32",
+          "shape": [512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.1 }
+        },
+        "scan_input": {
+          "dtype": "float32",
+          "shape": [256, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.1 }
+        }
+      },
+      "outputs": {
+        "final_state": { "dtype": "float32", "shape": [512], "tolerance": 0.0001 },
+        "scan_output": { "dtype": "float32", "shape": [256, 512], "tolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "vector_channel_full_lowered_prefix_sum_2048x512",
+      "inputs": {
+        "initial_state": {
+          "dtype": "float32",
+          "shape": [512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.1 }
+        },
+        "scan_input": {
+          "dtype": "float32",
+          "shape": [2048, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.1 }
+        }
+      },
+      "outputs": {
+        "final_state": { "dtype": "float32", "shape": [512], "tolerance": 0.0001 },
+        "scan_output": { "dtype": "float32", "shape": [2048, 512], "tolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "vector_channel_full_lowered_prefix_sum_256x128",
+      "inputs": {
+        "initial_state": {
+          "dtype": "float32",
+          "shape": [128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.1 }
+        },
+        "scan_input": {
+          "dtype": "float32",
+          "shape": [256, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.1 }
+        }
+      },
+      "outputs": {
+        "final_state": { "dtype": "float32", "shape": [128], "tolerance": 0.0001 },
+        "scan_output": { "dtype": "float32", "shape": [256, 128], "tolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "vector_channel_full_coop_control_dim512_steps4096",
+      "inputs": {
+        "initial_state": {
+          "dtype": "float32",
+          "shape": [512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.1 }
+        },
+        "scan_input": {
+          "dtype": "float32",
+          "shape": [4096, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.1 }
+        }
+      },
+      "outputs": {
+        "final_state": { "dtype": "float32", "shape": [512], "tolerance": 0.0001 },
+        "scan_output": { "dtype": "float32", "shape": [4096, 512], "tolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "vector_channel_full_lowered_largedim_fallback_dim4096_steps1024",
+      "inputs": {
+        "initial_state": {
+          "dtype": "float32",
+          "shape": [4096],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.1 }
+        },
+        "scan_input": {
+          "dtype": "float32",
+          "shape": [1024, 4096],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.1 }
+        }
+      },
+      "outputs": {
+        "final_state": { "dtype": "float32", "shape": [4096], "tolerance": 0.0001 },
+        "scan_output": { "dtype": "float32", "shape": [1024, 4096], "tolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "vector_channel_full_coop_reverse_partialchunk_dim512_steps300",
+      "attrs": { "reverse": 1 },
+      "inputs": {
+        "initial_state": {
+          "dtype": "float32",
+          "shape": [512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.1 }
+        },
+        "scan_input": {
+          "dtype": "float32",
+          "shape": [300, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.1 }
+        }
+      },
+      "outputs": {
+        "final_state": { "dtype": "float32", "shape": [512], "tolerance": 0.0001 },
+        "scan_output": { "dtype": "float32", "shape": [300, 512], "tolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "vector_channel_boundary_t63_d12_reverse0",
+      "attrs": { "reverse": 0 },
+      "inputs": {
+        "initial_state": {
+          "dtype": "float32",
+          "shape": [12],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.1 }
+        },
+        "scan_input": {
+          "dtype": "float32",
+          "shape": [63, 12],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.1 }
+        }
+      },
+      "outputs": {
+        "final_state": { "dtype": "float32", "shape": [12], "tolerance": 0.0001 },
+        "scan_output": { "dtype": "float32", "shape": [63, 12], "tolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "vector_channel_boundary_t64_d12_reverse0",
+      "attrs": { "reverse": 0 },
+      "inputs": {
+        "initial_state": {
+          "dtype": "float32",
+          "shape": [12],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.1 }
+        },
+        "scan_input": {
+          "dtype": "float32",
+          "shape": [64, 12],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.1 }
+        }
+      },
+      "outputs": {
+        "final_state": { "dtype": "float32", "shape": [12], "tolerance": 0.0001 },
+        "scan_output": { "dtype": "float32", "shape": [64, 12], "tolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "vector_channel_boundary_t65_d12_reverse1",
+      "attrs": { "reverse": 1 },
+      "inputs": {
+        "initial_state": {
+          "dtype": "float32",
+          "shape": [12],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.1 }
+        },
+        "scan_input": {
+          "dtype": "float32",
+          "shape": [65, 12],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.1 }
+        }
+      },
+      "outputs": {
+        "final_state": { "dtype": "float32", "shape": [12], "tolerance": 0.0001 },
+        "scan_output": { "dtype": "float32", "shape": [65, 12], "tolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "vector_channel_boundary_t255_d4_reverse0",
+      "attrs": { "reverse": 0 },
+      "inputs": {
+        "initial_state": {
+          "dtype": "float32",
+          "shape": [4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.1 }
+        },
+        "scan_input": {
+          "dtype": "float32",
+          "shape": [255, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.1 }
+        }
+      },
+      "outputs": {
+        "final_state": { "dtype": "float32", "shape": [4], "tolerance": 0.0001 },
+        "scan_output": { "dtype": "float32", "shape": [255, 4], "tolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "vector_channel_boundary_t257_d68_reverse1",
+      "attrs": { "reverse": 1 },
+      "inputs": {
+        "initial_state": {
+          "dtype": "float32",
+          "shape": [68],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.1 }
+        },
+        "scan_input": {
+          "dtype": "float32",
+          "shape": [257, 68],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.1 }
+        }
+      },
+      "outputs": {
+        "final_state": { "dtype": "float32", "shape": [68], "tolerance": 0.0001 },
+        "scan_output": { "dtype": "float32", "shape": [257, 68], "tolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "vector_channel_boundary_t257_d69_reverse1",
+      "attrs": { "reverse": 1 },
+      "inputs": {
+        "initial_state": {
+          "dtype": "float32",
+          "shape": [69],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.1 }
+        },
+        "scan_input": {
+          "dtype": "float32",
+          "shape": [257, 69],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.029, "scale": 0.1 }
+        }
+      },
+      "outputs": {
+        "final_state": { "dtype": "float32", "shape": [69], "tolerance": 0.0001 },
+        "scan_output": { "dtype": "float32", "shape": [257, 69], "tolerance": 0.0001 }
+      }
     }
   ]
 }