sync 6fdf6301e2bb

Browse files

Files changed (12) hide show

README.md +95 -0
build/webgpu/bench.json +449 -0
build/webgpu/fused-matmul-subgroup-matrix.wgsl.jinja +177 -0
build/webgpu/manifest.json +503 -0
build/webgpu/matmul-band-vec4.wgsl.jinja +80 -0
build/webgpu/matmul-subgroup-matrix-ext.wgsl.jinja +355 -0
build/webgpu/matmul-tiled-general-reg.wgsl.jinja +210 -0
build/webgpu/matmul-tiled-general.wgsl.jinja +168 -0
build/webgpu/matmul-vector-matrix-vec4.wgsl.jinja +55 -0
build/webgpu/metadata.json +41 -0
build/webgpu/reduce-axis0-splitk-combine.wgsl.jinja +95 -0
build/webgpu/test.json +1973 -0

README.md CHANGED Viewed

@@ -1,3 +1,98 @@
 ---
 license: apache-2.0
 ---

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# com.microsoft.TransposeMatMul
+`com.microsoft`  ·  ONNX Runtime contrib operator  ·  contrib since_version 1
+## Description
+Matrix product of two N-dimensional tensors `A` and `B`, following NumPy-style matrix-multiplication broadcasting, with optional transposition of either operand's last two dimensions and a scalar `alpha` multiplier. This is the strict subset of `FusedMatMul` that omits batch-dimension transposition, and it uses the same kernels. Float32 and float16 are supported; double and bfloat16 are not implemented.
+See the [ONNX Runtime `TransposeMatMul` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.TransposeMatMul) for the reference semantics.
+## Inputs
+| Name | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- |
+| `A` | `T` | — | — | N-dimensional matrix A. | required |
+| `B` | `T` | — | — | N-dimensional matrix B. | required |
+## Outputs
+| Name | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- |
+| `Y` | `T` | derived | derived | Matrix-multiplication result whose shape follows NumPy-style rules after applying the requested batch and matrix transpositions. | required |
+## Attributes
+Default values (overridable per request):
+| Attribute | Default | Description |
+| --- | --- | --- |
+| `alpha` | `1` | Scalar multiplier applied to the product of the input tensors. |
+| `transA` | `0` | When non-zero, transposes `A` on its last two dimensions before multiplication. |
+| `transB` | `0` | When non-zero, transposes `B` on its last two dimensions before multiplication. |
+## Type constraints
+| Variable | Allowed dtypes |
+| --- | --- |
+| `T` | `float32`, `float16` |
+## 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.
+- `broadcast_transb_tiled_reg` — Register-blocked broadcast product with physically transposed B. Reuses the shared batch-addressing tile, keeps f32 accumulation, and preserves scalar K order for f16. Low tile count and excessive padding demote this otherwise correct path.
+- `broadcast_transb_subgroup_matrix_f16` — Broadcast transposed-B product using a supported 8x8x8 subgroup-matrix configuration with f32 accumulation. Logical shapes and physical B strides share the existing matrix engine; insufficient output tiles or excessive padding retain the generic tile.
+- `broadcast_transb_subgroup_matrix_f32` — Broadcast transposed-B product using a supported 8x8x8 subgroup-matrix configuration with f32 accumulation. Logical shapes and physical B strides share the existing matrix engine; insufficient output tiles or excessive padding retain the generic tile.
+- `m1_gemv_vec4` — Vector-by-matrix specialization for a single output row: each workgroup owns 32 consecutive vec4 column groups and partitions the reduction across the workgroup's second dimension. The accumulator stays float32 for both tensor types.
+- `rank2_band_vec4_splitk` — Splits the vec4 band's K axis across up to sixteen workgroups. Each range writes an f32 partial band with alpha applied, and a combine pass sums the partials.
+- `rank2_band_vec4` — Few-row band for a rank-2 product without transposes. Each lane owns one vec4 column group and one accumulator per row, so every B word feeds all 2 to 16 rows; alpha is applied at the store.
+- `rank2_band_vec4_f32_preferred` — Few-row band for a rank-2 product without transposes. Each lane owns one vec4 column group and one accumulator per row, so every B word feeds all 2 to 16 rows; alpha is applied at the store.
+- `subgroup_matrix_splitk` — Partitions the K reduction of small-M rank-2 products across subgroup-matrix workgroups, then combines float32 partials that already include alpha.
+- `plain_rank2_tiled_reg` — Register-blocked rank-2 `Y = alpha * A @ B` specialization for non-transposed inputs on tiers without subgroup-matrix support.
+## Device requirements
+Some implementation variants require `subgroup-matrix` and `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
+## 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
+- [`fused-matmul-subgroup-matrix.wgsl.jinja`](build/webgpu/fused-matmul-subgroup-matrix.wgsl.jinja)
+- [`matmul-band-vec4.wgsl.jinja`](build/webgpu/matmul-band-vec4.wgsl.jinja)
+- [`matmul-subgroup-matrix-ext.wgsl.jinja`](build/webgpu/matmul-subgroup-matrix-ext.wgsl.jinja)
+- [`matmul-tiled-general-reg.wgsl.jinja`](build/webgpu/matmul-tiled-general-reg.wgsl.jinja)
+- [`matmul-tiled-general.wgsl.jinja`](build/webgpu/matmul-tiled-general.wgsl.jinja)
+- [`matmul-vector-matrix-vec4.wgsl.jinja`](build/webgpu/matmul-vector-matrix-vec4.wgsl.jinja)
+- [`reduce-axis0-splitk-combine.wgsl.jinja`](build/webgpu/reduce-axis0-splitk-combine.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.
+The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
+Replace each `*Data` placeholder with a typed array containing the corresponding input data.
+```js
+import { getKernel } from "@huggingface/kernels";
+const kernel = await getKernel("webgpu-kernels/com.microsoft.TransposeMatMul", { version: 1 });
+const { Y } = await kernel({ A: { data: AData, shape: [3] }, B: { data: BData, shape: [3] } });
+```

build/webgpu/bench.json ADDED Viewed

	@@ -0,0 +1,449 @@

+{
+  "cases": [
+    {
+      "name": "transposematmul-f32-plain-512x2048x512",
+      "preset": "smoke",
+      "attrs": { "alpha": 1 },
+      "inputs": {
+        "A": { "shape": [512, 2048], "dtype": "float32", "dist": "normal", "seed": 520, "scale": 0.1 },
+        "B": { "shape": [2048, 512], "dtype": "float32", "dist": "normal", "seed": 521, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "shape": [512, 512], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * 512 * 512 * 2048" }] }
+    },
+    {
+      "name": "transposematmul-f16-plain-512x2048x512",
+      "preset": "smoke",
+      "attrs": { "alpha": 1 },
+      "inputs": {
+        "A": { "shape": [512, 2048], "dtype": "float16", "dist": "normal", "seed": 530, "scale": 0.1 },
+        "B": { "shape": [2048, 512], "dtype": "float16", "dist": "normal", "seed": 531, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "shape": [512, 512], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * 512 * 512 * 2048" }] }
+    },
+    {
+      "name": "transposematmul-f32-unaligned-500x2000x500",
+      "preset": "smoke",
+      "attrs": { "alpha": 1 },
+      "inputs": {
+        "A": { "shape": [500, 2000], "dtype": "float32", "dist": "normal", "seed": 540, "scale": 0.1 },
+        "B": { "shape": [2000, 500], "dtype": "float32", "dist": "normal", "seed": 541, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "shape": [500, 500], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * 500 * 500 * 2000" }] }
+    },
+    {
+      "name": "transposematmul-f16-decode-gemv-m1-1x2048x512",
+      "preset": "smoke",
+      "attrs": { "alpha": 1 },
+      "vars": { "M": 1, "K": 2048, "N": 512 },
+      "inputs": {
+        "A": { "shape": [1, 2048], "dtype": "float16", "dist": "normal", "seed": 606, "scale": 0.1 },
+        "B": { "shape": [2048, 512], "dtype": "float16", "dist": "normal", "seed": 607, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "shape": [1, 512], "dtype": "float16" } },
+      "bench": { "primary": true, "metrics": [{ "type": "gflops", "value": "2 * 1 * 512 * 2048" }] }
+    },
+    {
+      "name": "transposematmul-f16-attn-scores-transb-8x512x64",
+      "preset": "smoke",
+      "attrs": { "alpha": 0.125, "transB": 1 },
+      "vars": { "M": 512, "K": 64, "N": 512 },
+      "inputs": {
+        "A": { "shape": [8, 512, 64], "dtype": "float16", "dist": "normal", "seed": 608, "scale": 0.2 },
+        "B": { "shape": [8, 512, 64], "dtype": "float16", "dist": "normal", "seed": 609, "scale": 0.2 }
+      },
+      "outputs": { "Y": { "shape": [8, 512, 512], "dtype": "float16" } },
+      "bench": { "primary": true, "metrics": [{ "type": "gflops", "value": "2 * 8 * 512 * 512 * 64" }] }
+    },
+    {
+      "name": "transposematmul-f16-broadcast-batch-rank4x3-2x8x512x64",
+      "preset": "model",
+      "attrs": { "alpha": 0.125, "transB": 1 },
+      "vars": { "M": 512, "K": 64, "N": 512 },
+      "inputs": {
+        "A": { "shape": [2, 8, 512, 64], "dtype": "float16", "dist": "normal", "seed": 610, "scale": 0.2 },
+        "B": { "shape": [8, 512, 64], "dtype": "float16", "dist": "normal", "seed": 611, "scale": 0.2 }
+      },
+      "outputs": { "Y": { "shape": [2, 8, 512, 512], "dtype": "float16" } },
+      "bench": { "primary": true, "metrics": [{ "type": "gflops", "value": "2 * 2 * 8 * 512 * 512 * 64" }] }
+    },
+    {
+      "name": "transposematmul-f16-rank4-by-rank2-shared-weight-b2h8-m512-k2048-n512",
+      "preset": "model",
+      "provenance": {
+        "notes": "Measures TransposeMatMul over a batched (2x8) projection sharing one rank-2 [2048,512] weight, with M=512, K=2048, N=512 (float16)."
+      },
+      "attrs": { "alpha": 1 },
+      "vars": { "dtype": "float16", "M": 512, "K": 2048, "N": 512 },
+      "inputs": {
+        "A": { "shape": [2, 8, 512, 2048], "dtype": "float16", "dist": "normal", "seed": 744, "scale": 0.05 },
+        "B": { "shape": [2048, 512], "dtype": "float16", "dist": "normal", "seed": 745, "scale": 0.05 }
+      },
+      "outputs": { "Y": { "shape": [2, 8, 512, 512], "dtype": "float16", "dist": "empty" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * args.K" }] }
+    },
+    {
+      "name": "transposematmul-f16-band-m16-k2560-n4096",
+      "preset": "smoke",
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": { "shape": [16, 2560], "dtype": "float16", "dist": "normal", "seed": 752, "scale": 0.1 },
+        "B": { "shape": [2560, 4096], "dtype": "float16", "dist": "normal", "seed": 753, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "shape": [16, 4096], "dtype": "float16", "dist": "empty" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * 16 * 2560 * 4096" }] }
+    },
+    {
+      "name": "broadcast-transb-tails-float16",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float16", "shape": [2, 1, 129, 65], "dist": "normal", "seed": 7302, "scale": 0.1 },
+        "B": { "dtype": "float16", "shape": [3, 129, 65], "dist": "normal", "seed": 7303, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 3, 129, 129] } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, ranks.A - 1)" }] },
+      "provenance": {
+        "notes": "Measures TransposeMatMul over a rank-4 by rank-3 broadcast (batch dims 2x1 against 3) with transposed B: M=129, K=65, N=129 (float16)."
+      }
+    },
+    {
+      "name": "broadcast-transb-rank3x2-float16",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float16", "shape": [4, 256, 128], "dist": "normal", "seed": 7304, "scale": 0.1 },
+        "B": { "dtype": "float16", "shape": [512, 128], "dist": "normal", "seed": 7305, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [4, 256, 512] } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, ranks.A - 1)" }] },
+      "provenance": {
+        "notes": "Measures TransposeMatMul over a rank-3 by rank-2 broadcast with transposed B: M=256, K=128, N=512, batch=4 (float16)."
+      }
+    },
+    {
+      "name": "broadcast-transb-rank5x3-float16",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float16", "shape": [2, 1, 2, 128, 32], "dist": "normal", "seed": 7306, "scale": 0.1 },
+        "B": { "dtype": "float16", "shape": [2, 128, 32], "dist": "normal", "seed": 7307, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 1, 2, 128, 128] } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, ranks.A - 1)" }] },
+      "provenance": {
+        "notes": "Measures TransposeMatMul over a rank-5 by rank-3 broadcast with transposed B: M=128, K=32, N=128, batch dims 2x1x2 (float16)."
+      }
+    },
+    {
+      "name": "broadcast-transb-low_tiles-float16",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float16", "shape": [1, 64, 32], "dist": "normal", "seed": 7308, "scale": 0.1 },
+        "B": { "dtype": "float16", "shape": [64, 32], "dist": "normal", "seed": 7309, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [1, 64, 64] } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, ranks.A - 1)" }] },
+      "provenance": {
+        "notes": "Measures TransposeMatMul over a rank-3 by rank-2 broadcast with transposed B: M=64, K=32, N=64, batch=1 (float16)."
+      }
+    },
+    {
+      "name": "broadcast-transb-large-float32",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [2, 8, 512, 64], "dist": "normal", "seed": 7300, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [8, 512, 64], "dist": "normal", "seed": 7301, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 8, 512, 512] } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, ranks.A - 1)" }] },
+      "provenance": {
+        "notes": "Measures TransposeMatMul over a rank-4 by rank-3 broadcast with transposed B: M=512, K=64, N=512, batch dims 2x8 (float32)."
+      }
+    },
+    {
+      "name": "broadcast-transb-tails-float32",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [2, 1, 129, 65], "dist": "normal", "seed": 7302, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [3, 129, 65], "dist": "normal", "seed": 7303, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 3, 129, 129] } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, ranks.A - 1)" }] },
+      "provenance": {
+        "notes": "Measures TransposeMatMul over a rank-4 by rank-3 broadcast (batch dims 2x1 against 3) with transposed B: M=129, K=65, N=129 (float32)."
+      }
+    },
+    {
+      "name": "broadcast-transb-rank3x2-float32",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [4, 256, 128], "dist": "normal", "seed": 7304, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [512, 128], "dist": "normal", "seed": 7305, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [4, 256, 512] } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, ranks.A - 1)" }] },
+      "provenance": {
+        "notes": "Measures TransposeMatMul over a rank-3 by rank-2 broadcast with transposed B: M=256, K=128, N=512, batch=4 (float32)."
+      }
+    },
+    {
+      "name": "broadcast-transb-rank5x3-float32",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [2, 1, 2, 128, 32], "dist": "normal", "seed": 7306, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [2, 128, 32], "dist": "normal", "seed": 7307, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 1, 2, 128, 128] } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, ranks.A - 1)" }] },
+      "provenance": {
+        "notes": "Measures TransposeMatMul over a rank-5 by rank-3 broadcast with transposed B: M=128, K=32, N=128, batch dims 2x1x2 (float32)."
+      }
+    },
+    {
+      "name": "broadcast-transb-low_tiles-float32",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [1, 64, 32], "dist": "normal", "seed": 7308, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [64, 32], "dist": "normal", "seed": 7309, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [1, 64, 64] } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, ranks.A - 1)" }] },
+      "provenance": {
+        "notes": "Measures TransposeMatMul over a rank-3 by rank-2 broadcast with transposed B: M=64, K=32, N=64, batch=1 (float32)."
+      }
+    },
+    {
+      "name": "broadcast-grid-b3-m128-k64-n256-float16-a0.125",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": 0.125 },
+      "inputs": {
+        "A": { "dtype": "float16", "shape": [3, 128, 64], "dist": "normal", "seed": 7304, "scale": 0.1 },
+        "B": { "dtype": "float16", "shape": [256, 64], "dist": "normal", "seed": 7305, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [3, 128, 256] } },
+      "provenance": {
+        "notes": "Measures TransposeMatMul over a batch=3 grid with transposed B: M=128, K=64, N=256, alpha=0.125 (float16)."
+      },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, ranks.A - 1)" }] }
+    },
+    {
+      "name": "broadcast-grid-b4-m128-k64-n256-float16-a-0.375",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.375 },
+      "inputs": {
+        "A": { "dtype": "float16", "shape": [4, 128, 64], "dist": "normal", "seed": 7304, "scale": 0.1 },
+        "B": { "dtype": "float16", "shape": [256, 64], "dist": "normal", "seed": 7305, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [4, 128, 256] } },
+      "provenance": {
+        "notes": "Measures TransposeMatMul over a batch=4 grid with transposed B: M=128, K=64, N=256, alpha=-0.375 (float16)."
+      },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, ranks.A - 1)" }] }
+    },
+    {
+      "name": "broadcast-grid-b8-m128-k64-n256-float16-a0.5",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": 0.5 },
+      "inputs": {
+        "A": { "dtype": "float16", "shape": [8, 128, 64], "dist": "normal", "seed": 7304, "scale": 0.1 },
+        "B": { "dtype": "float16", "shape": [256, 64], "dist": "normal", "seed": 7305, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [8, 128, 256] } },
+      "provenance": {
+        "notes": "Measures TransposeMatMul over a batch=8 grid with transposed B: M=128, K=64, N=256, alpha=0.5 (float16)."
+      },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, ranks.A - 1)" }] }
+    },
+    {
+      "name": "broadcast-grid-b8-m129-k65-n257-float16-a0.125",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": 0.125 },
+      "inputs": {
+        "A": { "dtype": "float16", "shape": [8, 129, 65], "dist": "normal", "seed": 7304, "scale": 0.1 },
+        "B": { "dtype": "float16", "shape": [257, 65], "dist": "normal", "seed": 7305, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [8, 129, 257] } },
+      "provenance": {
+        "notes": "Measures TransposeMatMul over a batch=8 grid with transposed B: M=129, K=65, N=257, alpha=0.125 (float16)."
+      },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, ranks.A - 1)" }] }
+    },
+    {
+      "name": "broadcast-grid-b8-m128-k96-n128-float16-a0",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": 0 },
+      "inputs": {
+        "A": { "dtype": "float16", "shape": [8, 128, 96], "dist": "normal", "seed": 7304, "scale": 0.1 },
+        "B": { "dtype": "float16", "shape": [128, 96], "dist": "normal", "seed": 7305, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [8, 128, 128] } },
+      "provenance": {
+        "notes": "Measures TransposeMatMul over a batch=8 grid with transposed B: M=128, K=96, N=128, alpha=0 (float16)."
+      },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, ranks.A - 1)" }] }
+    },
+    {
+      "name": "broadcast-grid-b4-m256-k127-n512-float16-a-0.5",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float16", "shape": [4, 256, 127], "dist": "normal", "seed": 7304, "scale": 0.1 },
+        "B": { "dtype": "float16", "shape": [512, 127], "dist": "normal", "seed": 7305, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [4, 256, 512] } },
+      "provenance": {
+        "notes": "Measures TransposeMatMul over a batch=4 grid with transposed B: M=256, K=127, N=512, alpha=-0.5 (float16)."
+      },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, ranks.A - 1)" }] }
+    },
+    {
+      "name": "broadcast-grid-b3-m128-k64-n256-float32-a0.125",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": 0.125 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [3, 128, 64], "dist": "normal", "seed": 7304, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [256, 64], "dist": "normal", "seed": 7305, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [3, 128, 256] } },
+      "provenance": {
+        "notes": "Measures TransposeMatMul over a batch=3 grid with transposed B: M=128, K=64, N=256, alpha=0.125 (float32)."
+      },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, ranks.A - 1)" }] }
+    },
+    {
+      "name": "broadcast-grid-b4-m128-k64-n256-float32-a-0.375",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.375 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [4, 128, 64], "dist": "normal", "seed": 7304, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [256, 64], "dist": "normal", "seed": 7305, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [4, 128, 256] } },
+      "provenance": {
+        "notes": "Measures TransposeMatMul over a batch=4 grid with transposed B: M=128, K=64, N=256, alpha=-0.375 (float32)."
+      },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, ranks.A - 1)" }] }
+    },
+    {
+      "name": "broadcast-grid-b8-m128-k64-n256-float32-a0.5",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": 0.5 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [8, 128, 64], "dist": "normal", "seed": 7304, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [256, 64], "dist": "normal", "seed": 7305, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [8, 128, 256] } },
+      "provenance": {
+        "notes": "Measures TransposeMatMul over a batch=8 grid with transposed B: M=128, K=64, N=256, alpha=0.5 (float32)."
+      },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, ranks.A - 1)" }] }
+    },
+    {
+      "name": "broadcast-grid-b8-m129-k65-n257-float32-a0.125",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": 0.125 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [8, 129, 65], "dist": "normal", "seed": 7304, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [257, 65], "dist": "normal", "seed": 7305, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [8, 129, 257] } },
+      "provenance": {
+        "notes": "Measures TransposeMatMul over a batch=8 grid with transposed B: M=129, K=65, N=257, alpha=0.125 (float32)."
+      },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, ranks.A - 1)" }] }
+    },
+    {
+      "name": "broadcast-grid-b8-m128-k96-n128-float32-a0",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": 0 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [8, 128, 96], "dist": "normal", "seed": 7304, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [128, 96], "dist": "normal", "seed": 7305, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [8, 128, 128] } },
+      "provenance": {
+        "notes": "Measures TransposeMatMul over a batch=8 grid with transposed B: M=128, K=96, N=128, alpha=0 (float32)."
+      },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, ranks.A - 1)" }] }
+    },
+    {
+      "name": "broadcast-grid-b4-m256-k127-n512-float32-a-0.5",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [4, 256, 127], "dist": "normal", "seed": 7304, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [512, 127], "dist": "normal", "seed": 7305, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [4, 256, 512] } },
+      "provenance": {
+        "notes": "Measures TransposeMatMul over a batch=4 grid with transposed B: M=256, K=127, N=512, alpha=-0.5 (float32)."
+      },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, ranks.A - 1)" }] }
+    },
+    {
+      "name": "broadcast-selected-mn-tail-float16",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": 0.125 },
+      "inputs": {
+        "A": { "dtype": "float16", "shape": [8, 129, 64], "dist": "normal", "seed": 7304, "scale": 0.1 },
+        "B": { "dtype": "float16", "shape": [257, 64], "dist": "normal", "seed": 7305, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [8, 129, 257] } },
+      "provenance": {
+        "notes": "Measures TransposeMatMul over a batch=8 grid with transposed B: M=129, K=64, N=257, alpha=0.125 (float16)."
+      },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, ranks.A - 1)" }] }
+    },
+    {
+      "name": "broadcast-selected-mn-tail-float32",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": 0.125 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [8, 129, 64], "dist": "normal", "seed": 7304, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [257, 64], "dist": "normal", "seed": 7305, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [8, 129, 257] } },
+      "provenance": {
+        "notes": "Measures TransposeMatMul over a batch=8 grid with transposed B: M=129, K=64, N=257, alpha=0.125 (float32)."
+      },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, ranks.A - 1)" }] }
+    },
+    {
+      "name": "broadcast-selected-broadcast-tail-float16",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": 0.125 },
+      "inputs": {
+        "A": { "dtype": "float16", "shape": [4, 1, 129, 64], "dist": "normal", "seed": 7304, "scale": 0.1 },
+        "B": { "dtype": "float16", "shape": [3, 257, 64], "dist": "normal", "seed": 7305, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [4, 3, 129, 257] } },
+      "provenance": {
+        "notes": "Measures TransposeMatMul over a rank-4 by rank-3 broadcast (batch dims 4x1 against 3) with transposed B: M=129, K=64, N=257, alpha=0.125 (float16)."
+      },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, ranks.A - 1)" }] }
+    },
+    {
+      "name": "broadcast-selected-broadcast-tail-float32",
+      "preset": "stress",
+      "attrs": { "transB": 1, "alpha": 0.125 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [4, 1, 129, 64], "dist": "normal", "seed": 7304, "scale": 0.1 },
+        "B": { "dtype": "float32", "shape": [3, 257, 64], "dist": "normal", "seed": 7305, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [4, 3, 129, 257] } },
+      "provenance": {
+        "notes": "Measures TransposeMatMul over a rank-4 by rank-3 broadcast (batch dims 4x1 against 3) with transposed B: M=129, K=64, N=257, alpha=0.125 (float32)."
+      },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * dim(shapes.A, ranks.A - 1)" }] }
+    }
+  ]
+}

build/webgpu/fused-matmul-subgroup-matrix.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,177 @@

+// com.microsoft.FusedMatMul subgroup-matrix specialization: Y = alpha * op(A) @ op(B).
+// transA and transB transpose the corresponding matrix operand on load.
+// Dense batches map through workgroup_id.z, and M-tail rows are guarded by row_limit.
+// The row count M arrives per call in `params.M`, so one pipeline serves every M;
+// K, N and the batch layout compile in.
+// A K % 32 == 0 gate keeps the reduction loop whole; N is free of the 64-wide column
+// tile because nTailSafe clamps the trailing tile's B columns to N - 1 and guards
+// every store on col < N. Both operands stage through workgroup memory, so
+// subgroupMatrixLoad only ever reads the full tile_A/tile_B arrays and never sees a
+// partial 8x8 tile at any M or N. The clamp is not interchangeable with a zero fill:
+// an out-of-bounds subgroupMatrixLoad resets to offset 0 and returns a different
+// valid tile, and a duplicated real column keeps the discarded accumulators finite.
+// The batch is required to match between A and B (no broadcast) because a_base and
+// b_base both index by the same workgroup_id.z.
+enable subgroups;
+{% if pinSubgroupSize32 %}
+enable subgroup_size_control;
+{% endif %}
+enable chromium_experimental_subgroup_matrix;
+diagnostic(off, chromium.subgroup_matrix_uniformity);
+{{ env.wgsl.resourceDeclarations }}
+{% set operandScalar = fScalar %}
+{% set accScalar = "f32" %}
+const K: u32 = {{ K }}u;
+const N: u32 = {{ N }}u;
+{% if not transA %}
+const A_M_STRIDE: u32 = K;
+{% endif %}
+const B_BATCH_STRIDE: u32 = K * N;
+const ALPHA: {{ accScalar }} = {{ accScalar }}({{ alpha }});
+const TILE_COLS: u32 = 64u;
+const TILE_ROWS: u32 = 32u;
+const TILE_K: u32 = 32u;
+const SUB_COLS: u32 = 32u;
+const SUB_ROWS: u32 = 16u;
+var<workgroup> tile_A: array<{{ operandScalar }}, 32 * 32>;
+var<workgroup> tile_B: array<{{ operandScalar }}, 64 * 32>;
+var<workgroup> scratch: array<array<array<{{ accScalar }}, 64>, 4>, 4>;
+fn loadSHMA(a_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
+  let a_global = tile_base + row;
+  let col = c_idx * 8u;
+  for (var col_offset = 0u; col_offset < 8u; col_offset = col_offset + 1u) {
+    let k = k_idx + col + col_offset;
+    if (a_global < params.M) {
+{% if transA %}
+      // op(A) = A^T: A stored [.., K, M], so op(A)[a_global, k] = A[k, a_global].
+      tile_A[row * TILE_K + col + col_offset] = {{ operandScalar }}(a[a_base + k * params.M + a_global]);
+{% else %}
+      tile_A[row * TILE_K + col + col_offset] = {{ operandScalar }}(a[a_base + a_global * A_M_STRIDE + k]);
+{% endif %}
+    } else {
+      tile_A[row * TILE_K + col + col_offset] = {{ operandScalar }}(0.0);
+    }
+  }
+}
+fn loadSHMB(b_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
+  let b_col = tile_base + row;
+  let col = c_idx * 16u;
+  for (var i = 0u; i < 16u; i = i + 1u) {
+    let k = k_idx + col + i;
+{% if transB %}
+    // op(B) = B^T: B stored [.., N, K], so op(B)[k, b_col] = B[b_col, k].
+    tile_B[row * TILE_K + col + i] = {{ operandScalar }}(b[b_base + b_col * K + k]);
+{% else %}
+    tile_B[row * TILE_K + col + i] = {{ operandScalar }}(b[b_base + k * N + b_col]);
+{% endif %}
+  }
+}
+fn storeOutput(offset: u32, row: u32, col: u32, src_slot: u32, row_limit: i32) {
+  if (row_limit > 0 && row < u32(row_limit)) {
+    let col2 = col + 1u;
+    y[offset + row * N + col]       = {{ outScalar }}(ALPHA * scratch[src_slot][0][row * 8u + col]);
+    y[offset + row * N + col + 8u]  = {{ outScalar }}(ALPHA * scratch[src_slot][1][row * 8u + col]);
+    y[offset + row * N + col + 16u] = {{ outScalar }}(ALPHA * scratch[src_slot][2][row * 8u + col]);
+    y[offset + row * N + col + 24u] = {{ outScalar }}(ALPHA * scratch[src_slot][3][row * 8u + col]);
+    y[offset + row * N + col2]       = {{ outScalar }}(ALPHA * scratch[src_slot][0][row * 8u + col2]);
+    y[offset + row * N + col2 + 8u]  = {{ outScalar }}(ALPHA * scratch[src_slot][1][row * 8u + col2]);
+    y[offset + row * N + col2 + 16u] = {{ outScalar }}(ALPHA * scratch[src_slot][2][row * 8u + col2]);
+    y[offset + row * N + col2 + 24u] = {{ outScalar }}(ALPHA * scratch[src_slot][3][row * 8u + col2]);
+  }
+}
+@compute @workgroup_size(128, 1, 1){{ " @subgroup_size(32)" if pinSubgroupSize32 else "" }}
+fn main(
+  @builtin(workgroup_id) workgroup_id: vec3<u32>,
+  @builtin(local_invocation_index) local_idx: u32,
+  @builtin(subgroup_invocation_id) sg_id: u32,
+  @builtin(subgroup_size) sg_size: u32
+) {
+  let M = params.M;
+  let A_BATCH_STRIDE = M * K;
+  let C_BATCH_STRIDE = M * N;
+  let batch = workgroup_id.z;
+  let a_base = batch * A_BATCH_STRIDE;
+  let b_base = batch * B_BATCH_STRIDE;
+  let c_base = batch * C_BATCH_STRIDE;
+  let a_global_base = workgroup_id.y * TILE_ROWS;
+  let b_global_base = workgroup_id.x * TILE_COLS;
+  let subtile_id = local_idx / sg_size;
+  let subtile_idx = subtile_id / 2u;
+  let subtile_idy = subtile_id % 2u;
+  let base_A = subtile_idy * SUB_ROWS;
+  let base_B = subtile_idx * SUB_COLS;
+  var matC00: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
+  var matC01: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
+  var matC02: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
+  var matC03: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
+  var matC10: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
+  var matC11: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
+  var matC12: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
+  var matC13: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
+  for (var kidx = 0u; kidx < K; kidx = kidx + TILE_K) {
+    loadSHMA(a_base, a_global_base, kidx, local_idx / 4u, local_idx % 4u);
+    loadSHMB(b_base, b_global_base, kidx, local_idx / 2u, local_idx % 2u);
+    workgroupBarrier();
+    for (var step = 0u; step < TILE_K; step = step + 8u) {
+{% set directInputs = false %}
+      let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
+{% for r in range(2) %}
+      var matA{{ r }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&{{ "w" if directInputs else "tile_A" }}, matrix_a_offset{% if r > 0 %} + 8u * {{ "K" if directInputs else "TILE_K" }}{% endif %}, {{ "K" if directInputs else "TILE_K" }});
+{% endfor %}
+      let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
+{% for c in range(4) %}
+      var matB{{ c }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&{{ "xm" if directInputs else "tile_B" }}, matrix_b_offset{% if c > 0 %} + {{ c * 8 }}u * TILE_K{% endif %}, {{ "N" if directInputs else "TILE_K" }});
+{% endfor %}
+      matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
+      matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);
+      matC02 = subgroupMatrixMultiplyAccumulate(matA0, matB2, matC02);
+      matC03 = subgroupMatrixMultiplyAccumulate(matA0, matB3, matC03);
+      matC10 = subgroupMatrixMultiplyAccumulate(matA1, matB0, matC10);
+      matC11 = subgroupMatrixMultiplyAccumulate(matA1, matB1, matC11);
+      matC12 = subgroupMatrixMultiplyAccumulate(matA1, matB2, matC12);
+      matC13 = subgroupMatrixMultiplyAccumulate(matA1, matB3, matC13);
+    }
+    workgroupBarrier();
+  }
+  // The four scratch banks are reused across the two row-groups, and each is written
+  // by a collective subgroupMatrixStore then read across lanes by storeOutput. Barriers
+  // give the reads visibility of the store and stop the second row-group's store from
+  // clobbering the first's still-in-flight readback when a partial final M-tile
+  // diverges storeOutput's guard. Without both barriers the last valid row can be corrupted.
+  subgroupMatrixStore<row_major>(&scratch[subtile_id][0], 0u, matC00, 8u);
+  subgroupMatrixStore<row_major>(&scratch[subtile_id][1], 0u, matC01, 8u);
+  subgroupMatrixStore<row_major>(&scratch[subtile_id][2], 0u, matC02, 8u);
+  subgroupMatrixStore<row_major>(&scratch[subtile_id][3], 0u, matC03, 8u);
+  workgroupBarrier();
+  let row = sg_id / 4u;
+  let col = (sg_id % 4u) * 2u;
+  var matrix_c_offset = c_base + (a_global_base + base_A) * N + b_global_base + base_B;
+  var row_limit = i32(M) - i32(a_global_base + base_A);
+  storeOutput(matrix_c_offset, row, col, subtile_id, row_limit);
+  workgroupBarrier();
+  subgroupMatrixStore<row_major>(&scratch[subtile_id][0], 0u, matC10, 8u);
+  subgroupMatrixStore<row_major>(&scratch[subtile_id][1], 0u, matC11, 8u);
+  subgroupMatrixStore<row_major>(&scratch[subtile_id][2], 0u, matC12, 8u);
+  subgroupMatrixStore<row_major>(&scratch[subtile_id][3], 0u, matC13, 8u);
+  workgroupBarrier();
+  matrix_c_offset = matrix_c_offset + 8u * N;
+  row_limit = i32(M) - i32(a_global_base + base_A + 8u);
+  storeOutput(matrix_c_offset, row, col, subtile_id, row_limit);
+}

build/webgpu/manifest.json ADDED Viewed

	@@ -0,0 +1,503 @@

+{
+  "domain": "com.microsoft",
+  "name": "TransposeMatMul",
+  "sinceVersion": 1,
+  "inputs": { "A": { "dtype": "T" }, "B": { "dtype": "T" } },
+  "outputs": {
+    "Y": {
+      "dtype": "T",
+      "rank": "max(ranks.A, ranks.B) - (1 if ranks.A == 1 or ranks.B == 1 else 0)",
+      "shape": "matmulShape(logicalAShape, logicalBShape)"
+    }
+  },
+  "attributes": { "alpha": { "default": 1 }, "transA": { "default": 0 }, "transB": { "default": 0 } },
+  "typeConstraints": { "T": ["float32", "float16"] },
+  "tunables": {
+    "TILED_REG_MIN_WORKGROUPS": { "default": 64 },
+    "PLAIN_RANK2_REG_DEEP_K_TILES": { "default": 128 },
+    "GEMV_TARGET_BLOCKS": { "default": 512 },
+    "SUBGROUP_MATRIX_MIN_M": { "default": 2 },
+    "SUBGROUP_MATRIX_SPLITK_TARGET_WGS": { "default": 512 },
+    "SUBGROUP_MATRIX_SPLITK_MIN_K": { "default": 1024 },
+    "SUBGROUP_MATRIX_SPLITK_MAX_TILES": { "default": 128 },
+    "BAND_VEC4_MAX_ROWS": { "default": 16 },
+    "BAND_SPLIT_TARGET_WORKGROUPS": { "default": 256 },
+    "BAND_SPLIT_MAX_COLUMN_GROUPS": { "default": 24 },
+    "BAND_SPLIT_SLICES": { "default": 8 },
+    "BAND_PREFER_MAX_ROWS": { "default": 8 },
+    "BAND_PREFER_DEEP_K": { "default": 4096 },
+    "BROADCAST_TRANSB_MIN_WORKGROUPS": { "default": 64 },
+    "BROADCAST_TRANSB_MAX_PADDING_RATIO": { "default": 2 }
+  },
+  "derive": {
+    "batchMovedAShape": "shapes.A",
+    "batchMovedBShape": "shapes.B",
+    "logicalAShape": "moveAxis(batchMovedAShape, -1, -2) if attrs.transA != 0 and ranks.A > 1 else batchMovedAShape",
+    "logicalBShape": "moveAxis(batchMovedBShape, -1, -2) if attrs.transB != 0 and ranks.B > 1 else batchMovedBShape",
+    "gemvN": "dim(shapes.B, 1)",
+    "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32",
+    "canPinSubgroupSize32": "device.features.has(\"subgroups\") and device.features.has(\"subgroup-size-control\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize <= 32 and device.adapterInfo.subgroupMaxSize >= 32",
+    "pinSubgroupSize32": "canPinSubgroupSize32 and not wave32Adapter",
+    "wave32Effective": "wave32Adapter or pinSubgroupSize32",
+    "variableSubgroup16To32": "device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 16 and device.adapterInfo.subgroupMaxSize == 32",
+    "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
+    "rank2DeepPortableTier": "variableSubgroup16To32 or (has(device.adapterInfo, \"architecture\") and (device.adapterInfo.architecture == \"pascal\" or (not device.features.has(\"subgroups\") and (device.adapterInfo.architecture == \"apple\" or device.adapterInfo.architecture == \"gen-9\"))))",
+    "broadcastTransbM": "dim(shapes.A, ranks.A - 2)",
+    "broadcastTransbN": "dim(shapes.B, ranks.B - 2)",
+    "broadcastTransbK": "dim(shapes.A, ranks.A - 1)",
+    "broadcastTransbBatches": "numel(shapes.Y) / (dim(shapes.A, ranks.A - 2) * dim(shapes.B, ranks.B - 2))",
+    "gemvLanes": 32,
+    "vec4OutputTile": "4 * gemvLanes",
+    "gemvWorkgroups": "ceilDiv(gemvN, vec4OutputTile)",
+    "gemvSliceCap": "min(32, device.limits.maxComputeWorkgroupSizeY, floor(device.limits.maxComputeInvocationsPerWorkgroup / gemvLanes), floor(device.limits.maxComputeWorkgroupStorageSize / (16 * gemvLanes)))",
+    "gemvSlices": "max(1, min(gemvSliceCap, max(8, pow2ceil(ceilDiv(tunables.GEMV_TARGET_BLOCKS, gemvWorkgroups)))))",
+    "gemvResourcesFit": "gemvLanes <= device.limits.maxComputeWorkgroupSizeX and gemvSlices <= device.limits.maxComputeWorkgroupSizeY and gemvLanes * gemvSlices <= device.limits.maxComputeInvocationsPerWorkgroup and 16 * gemvLanes * gemvSlices <= device.limits.maxComputeWorkgroupStorageSize",
+    "registerTile": 64,
+    "generalTile": 32,
+    "tiledRegResourcesFit": "registerTile / 4 <= device.limits.maxComputeWorkgroupSizeX and registerTile / 4 <= device.limits.maxComputeWorkgroupSizeY and registerTile * registerTile / 16 <= device.limits.maxComputeInvocationsPerWorkgroup and 32 * registerTile * dtypeBytes(dtypes.T) <= device.limits.maxComputeWorkgroupStorageSize",
+    "plainRank2RegDeepPreferredTier": "rank2DeepPortableTier and dim(shapes.A, 1) >= tunables.PLAIN_RANK2_REG_DEEP_K_TILES * 16 and dim(shapes.A, 1) % 16 == 0",
+    "subgroupMatrixResourcesFit": "128 <= deviceWorkgroupCap and ((32 * 32 + 64 * 32) * dtypeBytes(dtypes.T) + 4 * 4 * 64 * 4) <= device.limits.maxComputeWorkgroupStorageSize",
+    "sgmatSplitKDepth": "dim(shapes.A, ranks.A - 1)",
+    "sgmatSplitK32Ok": "sgmatSplitKDepth % 1024 == 0",
+    "sgmatSplitK16Ok": "sgmatSplitKDepth % 512 == 0",
+    "sgmatSplitK8Ok": "sgmatSplitKDepth % 256 == 0",
+    "sgmatSplitK4Ok": "sgmatSplitKDepth % 128 == 0",
+    "sgmatSplitK2Ok": "sgmatSplitKDepth % 64 == 0",
+    "bandSplitWant": "pow2ceil(ceilDiv(tunables.BAND_SPLIT_TARGET_WORKGROUPS, max(1, gemvWorkgroups)))",
+    "bandSplitK": "16 if (bandSplitWant >= 16 and dim(shapes.A, ranks.A - 1) >= 4096) else (8 if (bandSplitWant >= 8 and dim(shapes.A, ranks.A - 1) >= 2048) else (4 if (bandSplitWant >= 4 and dim(shapes.A, ranks.A - 1) >= 1024) else (2 if (bandSplitWant >= 2 and dim(shapes.A, ranks.A - 1) >= 512) else 1)))",
+    "scalar": "dtypes.T",
+    "alpha": "attrs.alpha",
+    "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
+    "fScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"",
+    "aRank": "ranks.A",
+    "bRank": "ranks.B",
+    "fusedSgmatRank2Ok": "ranks.A == 2 and ranks.B == 2 and ranks.Y == 2 and attrs.transA == 0 and attrs.transB == 0 and dim(shapes.A, 1) == dim(shapes.B, 0) and dim(shapes.Y, 0) == dim(shapes.A, 0) and dim(shapes.Y, 1) == dim(shapes.B, 1)",
+    "sgmatOutTiles": "ceilDiv(dim(shapes.A, 0), 32) * ceilDiv(dim(shapes.B, 1), 64) if fusedSgmatRank2Ok else 1",
+    "sgmatSplitKWant": "ceilDiv(tunables.SUBGROUP_MATRIX_SPLITK_TARGET_WGS, sgmatOutTiles)",
+    "sgmatSplitK": "32 if (sgmatSplitKWant > 16 and sgmatSplitK32Ok) else (16 if (sgmatSplitKWant > 8 and sgmatSplitK16Ok) else (8 if (sgmatSplitKWant > 4 and sgmatSplitK8Ok) else (4 if (sgmatSplitKWant > 2 and sgmatSplitK4Ok) else (2 if sgmatSplitK2Ok else 1))))",
+    "bandRank2Ok": "ranks.A == 2 and ranks.B == 2 and ranks.Y == 2 and attrs.transA == 0 and attrs.transB == 0 and dim(shapes.A, 1) == dim(shapes.B, 0) and dim(shapes.Y, 0) == dim(shapes.A, 0) and dim(shapes.Y, 1) == dim(shapes.B, 1)",
+    "broadcastTransbContract": "(f16Ok(dtypes.T)) and (attrs.transA == 0 and attrs.transB != 0) and (ranks.A > ranks.B and ranks.B >= 2) and (ranks.Y == ranks.A) and (sameShape(shapes.Y, matmulShape(logicalAShape, logicalBShape))) and (dim(shapes.A, ranks.A - 1) == dim(shapes.B, ranks.B - 1)) and (dim(shapes.A, ranks.A - 2) >= 64) and (dim(shapes.A, ranks.A - 1) >= 32) and (dim(shapes.B, ranks.B - 2) >= 64) and (numel(shapes.Y) / (dim(shapes.A, ranks.A - 2) * dim(shapes.B, ranks.B - 2)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535))"
+  },
+  "bindings": {
+    "a": { "arg": "A", "elementType": "$scalar" },
+    "b": { "arg": "B", "elementType": "$vectorScalar" },
+    "partials": { "buffer": "read-only-storage", "elementType": "f32" },
+    "y": { "arg": "Y", "elementType": "$scalar" },
+    "params": { "struct": [{ "name": "cols", "type": "u32", "value": "numel(shapes.Y)" }] },
+    "b_scalar": { "arg": "B", "name": "b", "elementType": "$scalar" },
+    "params_rows": { "name": "params", "struct": [{ "name": "M", "type": "u32", "value": "rowCount" }] }
+  },
+  "variants": [
+    {
+      "id": "broadcast_transb_tiled_reg",
+      "priority": 6,
+      "when": ["broadcastTransbContract", "tiledRegResourcesFit", "ceilDiv(dim(shapes.A, ranks.A - 2), registerTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceilDiv(dim(shapes.B, ranks.B - 2), registerTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
+      "derive": {
+        "bShape": "logicalBShape",
+        "bTransposed": true,
+        "rowCount": "outer(shapes.A, ranks.A - 1) if ranks.B == 2 else broadcastTransbM",
+        "N": "broadcastTransbN",
+        "batchCount": "broadcastTransbBatches if ranks.B > 2 else 1",
+        "transBatchA": false,
+        "regSequentialK": "dtypes.T == \"f16\""
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "TransposeMatMul.BroadcastTransBTiledReg",
+          "shader": "matmul-tiled-general-reg.wgsl.jinja",
+          "derive": {
+            "aBatchShape": "prefix(shapes.A, ranks.A - 2) if ranks.B > 2 else []",
+            "aRank": "ranks.A if ranks.B > 2 else 2",
+            "K": "dim(shapes.A, ranks.A - 1)"
+          },
+          "bindings": ["a", "b_scalar", "y", "params_rows"],
+          "dispatch": { "x": "ceilDiv(N, registerTile)", "y": "ceilDiv(rowCount, registerTile)", "z": "batchCount" }
+        }
+      ],
+      "demoteWhen": ["broadcastTransbBatches * ceilDiv(broadcastTransbM, registerTile) * ceilDiv(broadcastTransbN, registerTile) < tunables.BROADCAST_TRANSB_MIN_WORKGROUPS", "ceilDiv(broadcastTransbM, registerTile) * registerTile * ceilDiv(broadcastTransbN, registerTile) * registerTile * ceilDiv(broadcastTransbK,16) * 16 > tunables.BROADCAST_TRANSB_MAX_PADDING_RATIO * broadcastTransbM * broadcastTransbN * broadcastTransbK"]
+    },
+    {
+      "id": "broadcast_transb_subgroup_matrix_f16",
+      "priority": 11,
+      "when": ["broadcastTransbContract", "dtypes.T == \"f16\"", "wave32Effective", "subgroupMatrixResourcesFit", "ceilDiv(dim(shapes.A, ranks.A - 2),32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceilDiv(dim(shapes.B, ranks.B - 2),64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }]
+      },
+      "derive": {
+        "bShape": "logicalBShape",
+        "bTransposed": true,
+        "rowCount": "outer(shapes.A, ranks.A - 1) if ranks.B == 2 else broadcastTransbM",
+        "K": "broadcastTransbK",
+        "N": "broadcastTransbN",
+        "batchCount": "broadcastTransbBatches if ranks.B > 2 else 1",
+        "hasBias": false,
+        "fScalar": "dtypes.T",
+        "outScalar": "dtypes.T",
+        "generalAddressing": true,
+        "tailSafe": "dim(shapes.A, ranks.A - 1) % 32 != 0 or dim(shapes.B, ranks.B - 2) % 64 != 0",
+        "outputBuffer": "\"y\""
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "TransposeMatMul.BroadcastTransBSubgroupMatrix",
+          "shader": "matmul-subgroup-matrix-ext.wgsl.jinja",
+          "derive": {
+            "aBatchShape": "prefix(shapes.A, ranks.A - 2) if ranks.B > 2 else []",
+            "aRank": "ranks.A if ranks.B > 2 else 2"
+          },
+          "bindings": ["a", "b_scalar", "y", "params_rows"],
+          "dispatch": { "x": "ceilDiv(N,64)", "y": "ceilDiv(rowCount,32)", "z": "batchCount" }
+        }
+      ],
+      "demoteWhen": ["broadcastTransbBatches * ceilDiv(broadcastTransbM,32) * ceilDiv(broadcastTransbN,64) < tunables.BROADCAST_TRANSB_MIN_WORKGROUPS", "ceilDiv(broadcastTransbM,32) * 32 * ceilDiv(broadcastTransbN,64) * 64 * ceilDiv(broadcastTransbK,32) * 32 > tunables.BROADCAST_TRANSB_MAX_PADDING_RATIO * broadcastTransbM * broadcastTransbN * broadcastTransbK"]
+    },
+    {
+      "id": "broadcast_transb_subgroup_matrix_f32",
+      "priority": 11,
+      "when": ["broadcastTransbContract", "dtypes.T == \"f32\"", "wave32Effective", "subgroupMatrixResourcesFit", "ceilDiv(dim(shapes.A, ranks.A - 2),32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceilDiv(dim(shapes.B, ranks.B - 2),64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
+      },
+      "derive": {
+        "bShape": "logicalBShape",
+        "bTransposed": true,
+        "rowCount": "outer(shapes.A, ranks.A - 1) if ranks.B == 2 else broadcastTransbM",
+        "K": "broadcastTransbK",
+        "N": "broadcastTransbN",
+        "batchCount": "broadcastTransbBatches if ranks.B > 2 else 1",
+        "hasBias": false,
+        "fScalar": "dtypes.T",
+        "outScalar": "dtypes.T",
+        "generalAddressing": true,
+        "tailSafe": "dim(shapes.A, ranks.A - 1) % 32 != 0 or dim(shapes.B, ranks.B - 2) % 64 != 0",
+        "outputBuffer": "\"y\""
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "TransposeMatMul.BroadcastTransBSubgroupMatrix",
+          "shader": "matmul-subgroup-matrix-ext.wgsl.jinja",
+          "derive": {
+            "aBatchShape": "prefix(shapes.A, ranks.A - 2) if ranks.B > 2 else []",
+            "aRank": "ranks.A if ranks.B > 2 else 2"
+          },
+          "bindings": ["a", "b_scalar", "y", "params_rows"],
+          "dispatch": { "x": "ceilDiv(N,64)", "y": "ceilDiv(rowCount,32)", "z": "batchCount" }
+        }
+      ],
+      "demoteWhen": ["broadcastTransbBatches * ceilDiv(broadcastTransbM,32) * ceilDiv(broadcastTransbN,64) < tunables.BROADCAST_TRANSB_MIN_WORKGROUPS", "ceilDiv(broadcastTransbM,32) * 32 * ceilDiv(broadcastTransbN,64) * 64 * ceilDiv(broadcastTransbK,32) * 32 > tunables.BROADCAST_TRANSB_MAX_PADDING_RATIO * broadcastTransbM * broadcastTransbN * broadcastTransbK"]
+    },
+    {
+      "id": "m1_gemv_vec4",
+      "priority": 30,
+      "when": ["(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "attrs.transA == 0", "attrs.transB == 0", "ranks.A == 2", "ranks.B == 2", "ranks.Y == 2", "dim(shapes.A, 0) == 1", "dim(shapes.Y, 0) == 1", "dim(shapes.A, 1) == dim(shapes.B, 0)", "dim(shapes.Y, 1) == dim(shapes.B, 1)", "dim(shapes.B, 1) > 0", "dim(shapes.B, 1) % 4 == 0", "gemvWorkgroups <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "f16Ok(dtypes.T)"],
+      "derive": {
+        "unrollK2": "dtypes.T == \"f16\"",
+        "gemvScalar": "dtypes.T",
+        "gemvVector": "\"vec4<\" ~ dtypes.T ~ \">\"",
+        "alphaScale": "attrs.alpha"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "TransposeMatMul.M1GemvVec4",
+          "shader": "matmul-vector-matrix-vec4.wgsl.jinja",
+          "bindings": [
+            { "arg": "A", "name": "a", "elementType": "$gemvScalar" },
+            { "arg": "B", "name": "b", "elementType": "$gemvVector" },
+            { "arg": "Y", "name": "c", "elementType": "$gemvVector" },
+            {
+              "name": "params",
+              "struct": [
+                { "name": "K", "type": "u32", "value": "dim(shapes.A, 1)" },
+                { "name": "N4", "type": "u32", "value": "dim(shapes.B, 1) / 4" }
+              ]
+            }
+          ],
+          "dispatch": { "x": "gemvWorkgroups" }
+        }
+      ]
+    },
+    {
+      "id": "rank2_band_vec4_splitk",
+      "priority": 11,
+      "when": ["(dtypes.T == \"f32\" or dtypes.T == \"f16\") and f16Ok(dtypes.T)", "bandRank2Ok", "dim(shapes.A, 0) >= 2", "dim(shapes.A, 0) <= tunables.BAND_VEC4_MAX_ROWS", "dim(shapes.A, 1) > 0", "dim(shapes.B, 1) > 0", "dim(shapes.B, 1) % 4 == 0", "gemvWorkgroups <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "gemvLanes <= device.limits.maxComputeWorkgroupSizeX", "gemvWorkgroups <= tunables.BAND_SPLIT_MAX_COLUMN_GROUPS", "bandSplitK >= 2", "bandSplitK * numel(shapes.Y) * 4 <= device.limits.maxStorageBufferBindingSize", "bandSplitK <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tunables.BAND_SPLIT_SLICES <= device.limits.maxComputeWorkgroupSizeY", "gemvLanes * tunables.BAND_SPLIT_SLICES <= device.limits.maxComputeInvocationsPerWorkgroup"],
+      "derive": {
+        "batched": false,
+        "outputBuffer": "\"y\"",
+        "M": "dim(shapes.A, 0)",
+        "K": "dim(shapes.A, 1)",
+        "N": "dim(shapes.B, 1)",
+        "gemvSlices": "tunables.BAND_SPLIT_SLICES",
+        "kSplits": "bandSplitK",
+        "split": "bandSplitK",
+        "workgroupSize": 256
+      },
+      "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[bandSplitK * numel(shapes.Y)]" }],
+      "passes": [
+        {
+          "id": "partial",
+          "name": "TransposeMatMul.Rank2BandVec4SplitK",
+          "shader": "matmul-band-vec4.wgsl.jinja",
+          "bindings": ["a", "b", { "scratch": "partials", "name": "y", "elementType": "vec4<f32>" }],
+          "dispatch": { "x": "gemvWorkgroups", "y": "bandSplitK" }
+        },
+        {
+          "id": "combine",
+          "name": "TransposeMatMul.Rank2BandVec4SplitKCombine",
+          "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
+          "derive": { "op": "\"sum\"", "outputF16": "dtypes.T == \"f16\"" },
+          "bindings": ["partials", "y", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.Y)), (256)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.Y)), (256)), 65535)",
+            "z": 1
+          }
+        }
+      ]
+    },
+    {
+      "id": "rank2_band_vec4",
+      "priority": 11,
+      "when": ["(dtypes.T == \"f32\" or dtypes.T == \"f16\") and f16Ok(dtypes.T)", "bandRank2Ok", "dim(shapes.A, 0) >= 2", "dim(shapes.A, 0) <= tunables.BAND_VEC4_MAX_ROWS", "dim(shapes.A, 1) > 0", "dim(shapes.B, 1) > 0", "dim(shapes.B, 1) % 4 == 0", "gemvWorkgroups <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "gemvResourcesFit", "not (gemvWorkgroups <= tunables.BAND_SPLIT_MAX_COLUMN_GROUPS and bandSplitK >= 2)"],
+      "derive": {
+        "batched": false,
+        "outputBuffer": "\"y\"",
+        "M": "dim(shapes.A, 0)",
+        "K": "dim(shapes.A, 1)",
+        "N": "dim(shapes.B, 1)"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "TransposeMatMul.Rank2BandVec4",
+          "shader": "matmul-band-vec4.wgsl.jinja",
+          "bindings": ["a", "b", { "arg": "Y", "name": "y", "elementType": "$vectorScalar" }],
+          "dispatch": { "x": "gemvWorkgroups" }
+        }
+      ]
+    },
+    {
+      "id": "rank2_band_vec4_f32_preferred",
+      "priority": 13,
+      "when": ["(dtypes.T == \"f32\" or dtypes.T == \"f16\") and f16Ok(dtypes.T)", "bandRank2Ok", "dim(shapes.A, 0) >= 2", "dim(shapes.A, 0) <= tunables.BAND_VEC4_MAX_ROWS", "dim(shapes.A, 1) > 0", "dim(shapes.B, 1) > 0", "dim(shapes.B, 1) % 4 == 0", "gemvWorkgroups <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "gemvResourcesFit", "not (gemvWorkgroups <= tunables.BAND_SPLIT_MAX_COLUMN_GROUPS and bandSplitK >= 2)"],
+      "demoteWhen": ["dtypes.T != \"f32\" or (dim(shapes.A, 0) > tunables.BAND_PREFER_MAX_ROWS and dim(shapes.A, 1) >= tunables.BAND_PREFER_DEEP_K)"],
+      "derive": {
+        "batched": false,
+        "outputBuffer": "\"y\"",
+        "M": "dim(shapes.A, 0)",
+        "K": "dim(shapes.A, 1)",
+        "N": "dim(shapes.B, 1)"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "TransposeMatMul.Rank2BandVec4",
+          "shader": "matmul-band-vec4.wgsl.jinja",
+          "bindings": ["a", "b", { "arg": "Y", "name": "y", "elementType": "$vectorScalar" }],
+          "dispatch": { "x": "gemvWorkgroups" }
+        }
+      ]
+    },
+    {
+      "id": "subgroup_matrix_splitk",
+      "priority": 12,
+      "when": ["(dtypes.T == \"f16\" or dtypes.T == \"f32\") and f16Ok(dtypes.T)", "fusedSgmatRank2Ok", "dim(shapes.A, 0) >= tunables.SUBGROUP_MATRIX_MIN_M", "dim(shapes.A, 1) >= tunables.SUBGROUP_MATRIX_SPLITK_MIN_K", "dim(shapes.B, 1) % 64 == 0", "sgmatSplitK >= 2", "sgmatOutTiles < tunables.SUBGROUP_MATRIX_SPLITK_MAX_TILES", "sgmatSplitK * numel(shapes.Y) * 4 <= device.limits.maxStorageBufferBindingSize", "sgmatSplitK <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceilDiv(dim(shapes.Y, 1), 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceilDiv(dim(shapes.Y, 0), 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "subgroupMatrixResourcesFit", "wave32Effective"],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [
+          { "componentType": "f16", "M": 8, "N": 8, "K": 8 },
+          { "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }
+        ]
+      },
+      "derive": {
+        "hasBias": false,
+        "generalAddressing": true,
+        "tailSafe": false,
+        "outputBuffer": "\"partials\"",
+        "outScalar": "\"f32\"",
+        "rowCount": "dim(shapes.A, 0)",
+        "K": "dim(shapes.A, 1)",
+        "N": "dim(shapes.B, 1)",
+        "batchCount": 1,
+        "splitK": "sgmatSplitK",
+        "kPerSplit": "dim(shapes.A, 1) / sgmatSplitK",
+        "split": "sgmatSplitK",
+        "workgroupSize": 256
+      },
+      "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[sgmatSplitK * numel(shapes.Y)]" }],
+      "passes": [
+        {
+          "id": "partial",
+          "name": "TransposeMatMul.SubgroupMatrixSplitK",
+          "shader": "matmul-subgroup-matrix-ext.wgsl.jinja",
+          "derive": { "bShape": ["dim(shapes.B, 0)", "dim(shapes.B, 1)"], "aRank": 2, "bRank": 2 },
+          "bindings": ["a", "b_scalar", { "name": "partials", "elementType": "f32" }, "params_rows"],
+          "dispatch": { "x": "ceilDiv(dim(shapes.Y, 1), 64)", "y": "ceilDiv(dim(shapes.Y, 0), 32)", "z": "sgmatSplitK" }
+        },
+        {
+          "id": "combine",
+          "name": "TransposeMatMul.SubgroupMatrixSplitKCombine",
+          "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
+          "derive": { "op": "\"sum\"", "outputF16": "dtypes.T == \"f16\"" },
+          "bindings": ["partials", "y", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.Y)), (256)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.Y)), (256)), 65535)",
+            "z": 1
+          }
+        }
+      ]
+    },
+    {
+      "id": "subgroup_matrix_tail_broadcast",
+      "priority": 11,
+      "when": ["dtypes.T == \"f16\"", "f16Ok(dtypes.T)", "attrs.transA == 0", "attrs.transB == 0", "(((ranks.A == 2 and ranks.B == 2 and ranks.Y == 2) or (ranks.A == 4 and ranks.B == 3 and ranks.Y == 4 and dim(shapes.Y, 0) == dim(shapes.A, 0) and (dim(shapes.A, 1) == dim(shapes.B, 0) or dim(shapes.A, 1) == 1 or dim(shapes.B, 0) == 1) and dim(shapes.Y, 1) == max(dim(shapes.A, 1), dim(shapes.B, 0)))) or (ranks.A == 4 and ranks.B == 2 and ranks.Y == 4 and sameShape(prefix(shapes.Y, 2), prefix(shapes.A, 2))))", "dim(shapes.A, ranks.A - 1) == dim(shapes.B, ranks.B - 2)", "dim(shapes.A, ranks.A - 2) >= tunables.SUBGROUP_MATRIX_MIN_M", "dim(shapes.A, ranks.A - 1) >= 32", "dim(shapes.B, ranks.B - 1) >= 64", "dim(shapes.Y, ranks.Y - 2) == dim(shapes.A, ranks.A - 2)", "dim(shapes.Y, ranks.Y - 1) == dim(shapes.B, ranks.B - 1)", "ceil(dim(shapes.B, ranks.B - 1) / 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceil(dim(shapes.A, ranks.A - 2) / 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "numel(shapes.Y) / (dim(shapes.A, ranks.A - 2) * dim(shapes.B, ranks.B - 1)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "wave32Effective"],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }]
+      },
+      "derive": {
+        "hasBias": false,
+        "fScalar": "\"f16\"",
+        "outScalar": "\"f16\"",
+        "generalAddressing": true,
+        "tailSafe": "dim(shapes.A, ranks.A - 1) % 32 != 0 or dim(shapes.B, ranks.B - 1) % 64 != 0",
+        "outputBuffer": "\"y\"",
+        "rowCount": "outer(shapes.A, ranks.A - 1) if ranks.B == 2 else dim(shapes.A, ranks.A - 2)",
+        "K": "dim(shapes.A, ranks.A - 1)",
+        "N": "dim(shapes.B, ranks.B - 1)",
+        "batchCount": "numel(shapes.Y) / (dim(shapes.A, ranks.A - 2) * dim(shapes.B, ranks.B - 1)) if ranks.B > 2 else 1"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "TransposeMatMul.SubgroupMatrixTailBroadcast",
+          "shader": "matmul-subgroup-matrix-ext.wgsl.jinja",
+          "derive": {
+            "aBatchShape": "prefix(shapes.A, ranks.A - 2) if ranks.B > 2 else []",
+            "aRank": "ranks.A if ranks.B > 2 else 2",
+            "bShape": "shapes.B"
+          },
+          "bindings": ["a", "b_scalar", "y", "params_rows"],
+          "dispatch": { "x": "ceil(N / 64)", "y": "ceil(rowCount / 32)", "z": "batchCount" }
+        }
+      ]
+    },
+    {
+      "id": "subgroup_matrix",
+      "priority": 10,
+      "when": ["f16Ok(dtypes.T)", "ranks.A >= 2", "ranks.B == ranks.A", "ranks.Y == ranks.A", "(dim(shapes.A, ranks.A - 2) if attrs.transA != 0 else dim(shapes.A, ranks.A - 1)) == (dim(shapes.B, ranks.B - 1) if attrs.transB != 0 else dim(shapes.B, ranks.B - 2))", "(dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else dim(shapes.A, ranks.A - 2)) >= tunables.SUBGROUP_MATRIX_MIN_M", "(dim(shapes.A, ranks.A - 2) if attrs.transA != 0 else dim(shapes.A, ranks.A - 1)) % 32 == 0", "(dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1)) % 64 == 0", "(ranks.A == 2 or (ranks.A == 3 and dim(shapes.A, 0) == dim(shapes.B, 0) and dim(shapes.Y, 0) == dim(shapes.B, 0)) or (ranks.A == 4 and dim(shapes.A, 0) == dim(shapes.B, 0) and dim(shapes.A, 1) == dim(shapes.B, 1) and dim(shapes.Y, 0) == dim(shapes.A, 0) and dim(shapes.Y, 1) == dim(shapes.A, 1)))", "dim(shapes.Y, ranks.Y - 2) == (dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else dim(shapes.A, ranks.A - 2))", "dim(shapes.Y, ranks.Y - 1) == (dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1))", "ceil((dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1)) / 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceil((dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else dim(shapes.A, ranks.A - 2)) / 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "numel(shapes.Y) / ((dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else dim(shapes.A, ranks.A - 2)) * (dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1))) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "wave32Effective"],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [
+          { "componentType": "f16", "M": 8, "N": 8, "K": 8 },
+          { "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }
+        ]
+      },
+      "derive": {
+        "outScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"",
+        "transA": "attrs.transA != 0",
+        "transB": "attrs.transB != 0",
+        "transBatchA": "false",
+        "rowCount": "(dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else dim(shapes.A, ranks.A - 2))",
+        "K": "(dim(shapes.A, ranks.A - 2) if attrs.transA != 0 else dim(shapes.A, ranks.A - 1))",
+        "N": "(dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1))",
+        "batchCount": "numel(shapes.Y) / ((dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else dim(shapes.A, ranks.A - 2)) * (dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1)))"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "TransposeMatMul.SubgroupMatrix",
+          "shader": "fused-matmul-subgroup-matrix.wgsl.jinja",
+          "bindings": ["a", "b_scalar", "y", "params_rows"],
+          "dispatch": { "x": "ceil(N / 64)", "y": "ceil(rowCount / 32)", "z": "numel(shapes.Y) / (rowCount * N)" }
+        }
+      ]
+    },
+    {
+      "id": "broadcast_rank4_tiled_reg",
+      "priority": 6,
+      "when": ["f16Ok(dtypes.T)", "attrs.transA == 0", "attrs.transB == 0", "ranks.A == 4", "(ranks.B == 2 or ranks.B == 3)", "ranks.Y == 4", "dim(shapes.Y, 0) == dim(shapes.A, 0)", "(ranks.B == 2 or dim(shapes.A, 1) == dim(shapes.B, 0) or dim(shapes.A, 1) == 1 or dim(shapes.B, 0) == 1)", "dim(shapes.Y, 1) == (dim(shapes.A, 1) if ranks.B == 2 else max(dim(shapes.A, 1), dim(shapes.B, 0)))", "dim(shapes.A, 3) == dim(shapes.B, ranks.B - 2)", "dim(shapes.Y, 2) == dim(shapes.A, 2)", "dim(shapes.Y, 3) == dim(shapes.B, ranks.B - 1)", "dim(shapes.A, 2) >= 64", "dim(shapes.A, 3) >= 32", "dim(shapes.B, ranks.B - 1) >= 64", "ceil(dim(shapes.B, ranks.B - 1) / registerTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceil(dim(shapes.A, 2) / registerTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "numel(shapes.Y) / (dim(shapes.A, 2) * dim(shapes.B, ranks.B - 1)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
+      "passes": [
+        {
+          "id": "main",
+          "name": "TransposeMatMul.BroadcastRank4TiledReg",
+          "shader": "matmul-tiled-general-reg.wgsl.jinja",
+          "derive": {
+            "aBatchShape": "prefix(shapes.A, ranks.A - 2) if ranks.B > 2 else []",
+            "aRank": "ranks.A if ranks.B > 2 else 2",
+            "K": "dim(shapes.A, ranks.A - 1)",
+            "bShape": "shapes.B",
+            "transBatchA": "false"
+          },
+          "bindings": ["a", "b_scalar", "y", "params_rows"],
+          "dispatch": {
+            "x": "ceil(dim(shapes.B, ranks.B - 1) / registerTile)",
+            "y": "ceil(rowCount / registerTile)",
+            "z": "batchCount"
+          }
+        }
+      ],
+      "derive": {
+        "rowCount": "outer(shapes.A, 3) if ranks.B == 2 else dim(shapes.A, 2)",
+        "batchCount": "numel(shapes.Y) / (dim(shapes.A, 2) * dim(shapes.B, ranks.B - 1)) if ranks.B > 2 else 1"
+      }
+    },
+    {
+      "id": "plain_rank2_tiled_reg",
+      "priority": 4,
+      "demoteWhen": ["plainRank2RegDeepPreferredTier"],
+      "when": ["f16Ok(dtypes.T)", "fusedSgmatRank2Ok", "dim(shapes.A, 0) >= 64", "dim(shapes.A, 1) >= 32", "dim(shapes.B, 1) >= 64", "ceil(dim(shapes.A, 0) / registerTile) * ceil(dim(shapes.B, 1) / registerTile) >= tunables.TILED_REG_MIN_WORKGROUPS", "ceil(dim(shapes.B, 1) / registerTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceil(dim(shapes.A, 0) / registerTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
+      "passes": [
+        {
+          "id": "main",
+          "name": "TransposeMatMul.PlainRank2TiledReg",
+          "shader": "matmul-tiled-general-reg.wgsl.jinja",
+          "derive": {
+            "rowCount": "dim(shapes.A, 0)",
+            "K": "dim(shapes.A, 1)",
+            "bShape": "shapes.B",
+            "transBatchA": "false"
+          },
+          "bindings": ["a", "b_scalar", "y", "params_rows"],
+          "dispatch": {
+            "x": "ceil(dim(shapes.B, 1) / registerTile)",
+            "y": "ceil(dim(shapes.A, 0) / registerTile)",
+            "z": 1
+          }
+        }
+      ]
+    },
+    {
+      "id": "tiled",
+      "priority": 0,
+      "when": ["ranks.A >= 1", "ranks.B >= 1", "f16Ok(dtypes.T)", "(dim(shapes.A, 0) if ranks.A == 1 else (dim(shapes.A, ranks.A - 1) if attrs.transA == 0 else dim(shapes.A, ranks.A - 2))) == (dim(shapes.B, 0) if ranks.B == 1 else (dim(shapes.B, ranks.B - 1) if attrs.transB != 0 else dim(shapes.B, ranks.B - 2)))", "ceil((1 if ranks.A == 1 else (dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else dim(shapes.A, ranks.A - 2))) / 16) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceil((1 if ranks.B == 1 else (dim(shapes.B, ranks.B - 1) if attrs.transB == 0 else dim(shapes.B, ranks.B - 2))) / 16) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "numel(shapes.Y) / max(1, (1 if ranks.A == 1 else (dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else dim(shapes.A, ranks.A - 2))) * (1 if ranks.B == 1 else (dim(shapes.B, ranks.B - 1) if attrs.transB == 0 else dim(shapes.B, ranks.B - 2)))) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
+      "passes": [
+        {
+          "id": "main",
+          "name": "TransposeMatMul.Tiled",
+          "shader": "matmul-tiled-general.wgsl.jinja",
+          "derive": {
+            "aShape": "shapes.A",
+            "bShape": "shapes.B",
+            "transA": "attrs.transA != 0",
+            "transB": "attrs.transB != 0",
+            "transBatchA": "false",
+            "transBatchB": "false"
+          },
+          "bindings": ["a", "b_scalar", "y"],
+          "dispatch": {
+            "x": "ceil((1 if ranks.B == 1 else (dim(shapes.B, ranks.B - 1) if attrs.transB == 0 else dim(shapes.B, ranks.B - 2))) / generalTile)",
+            "y": "ceil((1 if ranks.A == 1 else (dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else dim(shapes.A, ranks.A - 2))) / generalTile)",
+            "z": "numel(shapes.Y) / max(1, (1 if ranks.A == 1 else (dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else dim(shapes.A, ranks.A - 2))) * (1 if ranks.B == 1 else (dim(shapes.B, ranks.B - 1) if attrs.transB == 0 else dim(shapes.B, ranks.B - 2))))"
+          }
+        }
+      ]
+    }
+  ]
+}

build/webgpu/matmul-band-vec4.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,80 @@

+// Band GEMM y[M, N] = a[M, K] @ B[K, N]. Each lane owns one vec4 column group
+// and carries one accumulator per row, so a loaded B word feeds all M row
+// accumulators and each A value is reused across four adjacent output columns.
+//
+// A batched consumer runs one band per workgroup row: workgroup_id.y selects
+// the matrix, and every operand is offset by its per-matrix extent.
+{{ env.wgsl.resourceDeclarations }}
+const K: u32 = {{ K }}u;
+const N4: u32 = {{ N }}u / 4u;
+const LANES: u32 = {{ gemvLanes }}u;
+// SLICES partitions the K reduction across the workgroup's second dimension.
+const SLICES: u32 = {{ gemvSlices }}u;
+{% set kSplitsValue = kSplits if kSplits is defined else 1 %}
+{% set alphaValue = alpha if alpha is defined else 1 %}
+{% if kSplitsValue > 1 %}
+const K_SPLITS: u32 = {{ kSplitsValue }}u;
+const K_PER_SPLIT: u32 = (K + K_SPLITS - 1u) / K_SPLITS;
+{% endif %}
+// The rows drain through one 32 x SLICES array in turn, so the workgroup
+// footprint does not grow with the band.
+var<workgroup> partials: array<vec4<f32>, LANES * SLICES>;
+@compute @workgroup_size({{ gemvLanes }}, {{ gemvSlices }}, 1)
+fn main(
+  @builtin(workgroup_id) workgroup_id: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>
+) {
+  let lane = lid.x;
+  let slice = lid.y;
+  let cg = workgroup_id.x * LANES + lane;
+{% if kSplitsValue > 1 %}
+  let a_base = 0u;
+  let b_base = 0u;
+  let y_base = workgroup_id.y * ({{ M }}u * N4);
+  let k_begin = workgroup_id.y * K_PER_SPLIT;
+  let k_end = min(K, k_begin + K_PER_SPLIT);
+{% else %}
+  let a_base = 0u;
+  let b_base = 0u;
+  let y_base = 0u;
+{% endif %}
+{% for r in range(M) %}
+  var acc{{ r }} = vec4<f32>(0.0);
+{% endfor %}
+  if (cg < N4) {
+{% if kSplitsValue > 1 %}
+    for (var k = k_begin + slice; k < k_end; k = k + SLICES) {
+{% else %}
+    for (var k = slice; k < K; k = k + SLICES) {
+{% endif %}
+      let bv = vec4<f32>(b[b_base + k * N4 + cg]);
+{% for r in range(M) %}
+      acc{{ r }} = acc{{ r }} + f32(a[a_base + {{ r }}u * K + k]) * bv;
+{% endfor %}
+    }
+  }
+{% for r in range(M) %}
+  partials[slice * LANES + lane] = acc{{ r }};
+  workgroupBarrier();
+  if (slice == 0u && cg < N4) {
+    var total = partials[lane];
+    for (var s = 1u; s < SLICES; s = s + 1u) {
+      total = total + partials[s * LANES + lane];
+    }
+{% if alphaValue != 1 %}
+    total = total * {{ alphaValue }};
+{% endif %}
+{% if kSplitsValue > 1 %}
+    {{ outputBuffer }}[y_base + {{ r }}u * N4 + cg] = total;
+{% else %}
+    {{ outputBuffer }}[y_base + {{ r }}u * N4 + cg] = vec4<{{ T }}>(total);
+{% endif %}
+  }
+{% if not loop.last %}
+  workgroupBarrier();
+{% endif %}
+{% endfor %}
+}

build/webgpu/matmul-subgroup-matrix-ext.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,355 @@

+// Subgroup-matrix matmul over row-major, batch-outermost operands, with alpha,
+// dense/broadcast batching and guarded K/N tails under `generalAddressing`, and
+// an optional fused bias on the direct dense path that omits it.
+enable subgroups;
+{% if pinSubgroupSize32 %}
+enable subgroup_size_control;
+{% endif %}
+enable chromium_experimental_subgroup_matrix;
+diagnostic(off, chromium.subgroup_matrix_uniformity);
+{{ env.wgsl.resourceDeclarations }}
+{% set operandScalar = fScalar %}
+{% set accScalar = "f32" %}
+{% set GENERAL = true %}
+{% set TAIL = tailSafe is defined and tailSafe %}
+{% set SPLIT_K = splitK if splitK is defined else 1 %}
+{% set OUT = outputBuffer if outputBuffer is defined else "c" %}
+{% set OUT_SCALAR = outScalar if outScalar is defined else T %}
+{% set STATIC_M = M is defined %}
+{% set ROWS = "M" if STATIC_M else "params.M" %}
+{% set ROW_TEST = "a_global < M" if STATIC_M else "row_in" %}
+{% set aDims = (aShape | default([])) if STATIC_M else (aBatchShape | default([])) %}
+{% set kPerSplit = kPerSplit | default(0) %}
+{% set aR = aRank %}
+{% set bR = bRank %}
+{% set aBatchLen = aR - 2 %}
+{% set bBatchLen = bR - 2 %}
+{% set batchRank = aBatchLen %}
+{% set aMStride = K %}
+{% set aKStride = 1 %}
+{% set bKStride = 1 if bTransposed is defined and bTransposed else bShape[bR-1] %}
+{% set bNStride = bShape[bR-2] if bTransposed is defined and bTransposed else 1 %}
+{% if STATIC_M %}
+const M: u32 = {{ M }}u;
+{% endif %}
+const K: u32 = {{ K }}u;
+const N: u32 = {{ N }}u;
+const BATCH_COUNT: u32 = {{ batchCount if batchCount is defined else 1 }}u;
+{% if SPLIT_K > 1 %}
+const SPLIT_K: u32 = {{ SPLIT_K }}u;
+const K_PER_SPLIT: u32 = {{ kPerSplit }}u;
+{% endif %}
+const A_M_STRIDE: u32 = {{ aMStride }}u;
+const A_K_STRIDE: u32 = {{ aKStride }}u;
+const B_K_STRIDE: u32 = {{ bKStride }}u;
+const B_N_STRIDE: u32 = {{ bNStride }}u;
+{% if TAIL %}const K_FULL: u32 = (K / 32u) * 32u;
+{% endif %}
+const ALPHA: f32 = f32({{ alpha }});
+{% if STATIC_M %}
+const C_BATCH_STRIDE: u32 = M * N;
+{% endif %}
+const TILE_COLS: u32 = 64u;
+const TILE_ROWS: u32 = 32u;
+const TILE_K: u32 = 32u;
+const SUB_COLS: u32 = 32u;
+const SUB_ROWS: u32 = 16u;
+var<workgroup> tile_A: array<{{ operandScalar }}, 32 * 32>;
+var<workgroup> tile_B: array<{{ operandScalar }}, 64 * 32>;
+var<workgroup> scratch: array<array<array<{{ accScalar }}, 64>, 4>, 4>;
+fn loadSHMA(a_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
+  let a_global = tile_base + row;
+{% if not STATIC_M %}
+  let row_in = a_global < params.M;
+{% endif %}
+  let col = c_idx * 8u;
+  for (var col_offset = 0u; col_offset < 8u; col_offset = col_offset + 1u) {
+    let k = k_idx + col + col_offset;
+    if ({{ ROW_TEST }}) {
+{% if operandScalar == "f16" %}
+      tile_A[row * TILE_K + col + col_offset] = f16(a[a_base + a_global * A_M_STRIDE + k * A_K_STRIDE]);
+{% else %}
+      tile_A[row * TILE_K + col + col_offset] = f32(a[a_base + a_global * A_M_STRIDE + k * A_K_STRIDE]);
+{% endif %}
+    } else {
+{% if operandScalar == "f16" %}
+      tile_A[row * TILE_K + col + col_offset] = 0.0h;
+{% else %}
+      tile_A[row * TILE_K + col + col_offset] = 0.0;
+{% endif %}
+    }
+  }
+}
+{% if GENERAL and TAIL %}
+fn loadSHMAKTail(a_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
+  let a_global = tile_base + row;
+  let row_in = a_global < {{ ROWS }};
+  let col = c_idx * 8u;
+  for (var col_offset = 0u; col_offset < 8u; col_offset = col_offset + 1u) {
+    let k = k_idx + col + col_offset;
+    if (row_in && k < K) {
+      tile_A[row * TILE_K + col + col_offset] = {{ operandScalar }}(a[a_base + a_global * A_M_STRIDE + k * A_K_STRIDE]);
+    } else {
+      tile_A[row * TILE_K + col + col_offset] = {{ operandScalar }}(0);
+    }
+  }
+}
+{% endif %}
+fn loadSHMB(b_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
+  let b_col = tile_base + row;
+  let col = c_idx * 16u;
+  for (var i = 0u; i < 16u; i = i + 1u) {
+    let k = k_idx + col + i;
+{% if TAIL %}
+    let b_safe = min(b_col, N - 1u);
+    tile_B[row * TILE_K + col + i] = {{ operandScalar }}(b[b_base + k * B_K_STRIDE + b_safe * B_N_STRIDE]);
+{% else %}
+{% if operandScalar == "f16" %}
+    tile_B[row * TILE_K + col + i] = f16(b[b_base + k * B_K_STRIDE + b_col * B_N_STRIDE]);
+{% else %}
+    tile_B[row * TILE_K + col + i] = f32(b[b_base + k * B_K_STRIDE + b_col * B_N_STRIDE]);
+{% endif %}
+{% endif %}
+  }
+}
+{% if GENERAL and TAIL %}
+fn loadSHMBKTail(b_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
+  let b_col = min(tile_base + row, N - 1u);
+  let col = c_idx * 16u;
+  for (var i = 0u; i < 16u; i = i + 1u) {
+    let k = k_idx + col + i;
+    if (k < K) {
+      tile_B[row * TILE_K + col + i] = {{ operandScalar }}(b[b_base + k * B_K_STRIDE + b_col * B_N_STRIDE]);
+    } else {
+      tile_B[row * TILE_K + col + i] = {{ operandScalar }}(0);
+    }
+  }
+}
+{% endif %}
+{% set needsColBase = hasBias or (GENERAL and TAIL) %}
+fn storeOutput(offset: u32{% if needsColBase %}, col_base: u32{% endif %}, row: u32, col: u32, src_slot: u32, row_limit: i32) {
+  if (row_limit > 0 && row < u32(row_limit)) {
+    let col2 = col + 1u;
+{% for block in range(4) %}
+{% if TAIL %}
+    if (col_base + col + {{ block * 8 }}u < N) {
+{% endif %}
+      {{ OUT }}[offset + row * N + col + {{ block * 8 }}u] = {{ OUT_SCALAR }}(
+        ALPHA * scratch[src_slot][{{ block }}][row * 8u + col]
+      );
+{% if TAIL %}
+    }
+    if (col_base + col2 + {{ block * 8 }}u < N) {
+{% endif %}
+      {{ OUT }}[offset + row * N + col2 + {{ block * 8 }}u] = {{ OUT_SCALAR }}(
+        ALPHA * scratch[src_slot][{{ block }}][row * 8u + col2]
+      );
+{% if TAIL %}
+    }
+{% endif %}
+{% endfor %}
+  }
+}
+@compute @workgroup_size(128, 1, 1){{ " @subgroup_size(32)" if pinSubgroupSize32 else "" }}
+fn main(
+  @builtin(workgroup_id) workgroup_id: vec3<u32>,
+  @builtin(num_workgroups) num_wg: vec3<u32>,
+  @builtin(local_invocation_index) local_idx: u32,
+  @builtin(subgroup_invocation_id) sg_id: u32,
+  @builtin(subgroup_size) sg_size: u32
+) {
+{% if not STATIC_M %}
+  // The row count arrives per call; the strides it scales are formed here.
+  let M = params.M;
+  let C_BATCH_STRIDE = M * N;
+{% endif %}
+  let b_global_base = workgroup_id.x * TILE_COLS;
+  let subtile_id = local_idx / sg_size;
+  let subtile_idx = subtile_id / 2u;
+  let subtile_idy = subtile_id % 2u;
+  let base_A = subtile_idy * SUB_ROWS;
+  let base_B = subtile_idx * SUB_COLS;
+  // Grid-stride over both the M-tile (y) and batch (z) axes so the dispatch stays
+  // <= maxComputeWorkgroupsPerDimension per dimension even when ceil(M/TILE_ROWS) or BATCH_COUNT exceed the
+  // limit. workgroup_size.z = 1 so num_wg.z is the batch dispatch stride, and
+  // num_wg.y * TILE_ROWS is the row-tile dispatch stride. The loop bounds (M and
+  // BATCH_COUNT are compile-time / uniform; num_wg and workgroup_id are uniform)
+  // are workgroup-uniform, so the trailing workgroupBarrier()s and the subgroup
+  // matrix operations stay reconverged. When neither axis is clamped,
+  // num_wg.y * TILE_ROWS > M and num_wg.z > BATCH_COUNT, so each loop executes
+  // exactly once at workgroup_id.y/workgroup_id.z.
+  let row_tile_stride = num_wg.y * TILE_ROWS;
+  for (var a_global_base = workgroup_id.y * TILE_ROWS; a_global_base < M; a_global_base += row_tile_stride) {
+{% if SPLIT_K > 1 %}
+  // Split-K maps z to (batch, K segment), increasing the independent workgroup
+  // count for narrow matrices. Every segment owns a disjoint contiguous K range.
+  for (var batch_split = workgroup_id.z; batch_split < BATCH_COUNT * SPLIT_K; batch_split += num_wg.z) {
+    let batch = batch_split / SPLIT_K;
+    let split_id = batch_split - batch * SPLIT_K;
+{% else %}
+  // workgroup_size.z = 1, so num_wg.z is the dispatch stride over the batch axis.
+  for (var batch = workgroup_id.z; batch < BATCH_COUNT; batch += num_wg.z) {
+{% endif %}
+{% set hasBatchCoord = namespace(value=false) %}
+{% for i in range(batchRank) %}
+{% set axis = batchRank - 1 - i %}
+{% set aAxis = axis - (batchRank - aBatchLen) %}
+{% set bAxis = axis - (batchRank - bBatchLen) %}
+{% set aDim = aDims[aAxis] %}
+{% set bDim = bShape[bAxis] if bAxis >= 0 else 1 %}
+{% if aDim > 1 or bDim > 1 %}{% set hasBatchCoord.value = true %}{% endif %}
+{% endfor %}
+    // Right-aligned broadcast offsets, decomposed from the flattened output batch.
+{% if hasBatchCoord.value %}
+    var zTmp = batch;
+{% endif %}
+    var a_base: u32 = 0u;
+    var b_base: u32 = 0u;
+{% for i in range(batchRank) %}
+{% set axis = batchRank - 1 - i %}
+{% set aAxis = axis - (batchRank - aBatchLen) %}
+{% set bAxis = axis - (batchRank - bBatchLen) %}
+{% set aDim = aDims[aAxis] %}
+{% set bDim = bShape[bAxis] if bAxis >= 0 else 1 %}
+{% set outDim = aDim if aDim >= bDim else bDim %}
+{% set aStride = namespace(v=1) %}
+{% if aDim != 1 %}{% for j in range(aAxis + 1, aR - 2) %}{% set aStride.v = aStride.v * aDims[j] %}{% endfor %}{% else %}{% set aStride.v = 0 %}{% endif %}
+{% set bStride = namespace(v=1) %}
+{% if bAxis >= 0 and bDim != 1 %}{% for j in range(bAxis + 1, bR) %}{% set bStride.v = bStride.v * bShape[j] %}{% endfor %}{% else %}{% set bStride.v = 0 %}{% endif %}
+{% if outDim > 1 %}
+    let c{{ axis }} = zTmp % {{ outDim }}u;
+    zTmp = zTmp / {{ outDim }}u;
+{% if aStride.v != 0 %}    a_base = a_base + c{{ axis }} * {% if aStride.v != 1 %}{{ aStride.v }}u * {% endif %}M * K;
+{% endif %}
+{% if bStride.v != 0 %}    b_base = b_base + c{{ axis }} * {{ bStride.v }}u;
+{% endif %}
+{% endif %}
+{% endfor %}
+{% if SPLIT_K > 1 %}
+    let c_base = (batch * SPLIT_K + split_id) * C_BATCH_STRIDE;
+{% else %}
+    let c_base = batch * C_BATCH_STRIDE;
+{% endif %}
+    var matC00: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
+    var matC01: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
+    var matC02: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
+    var matC03: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
+    var matC10: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
+    var matC11: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
+    var matC12: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
+    var matC13: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
+{% if SPLIT_K > 1 %}
+    let k_begin = split_id * K_PER_SPLIT;
+    let k_end = min(k_begin + K_PER_SPLIT, K);
+    for (var kidx = k_begin; kidx < k_end; kidx = kidx + TILE_K) {
+{% else %}
+    for (var kidx = 0u; kidx < {% if GENERAL and TAIL %}K_FULL{% else %}K{% endif %}; kidx = kidx + TILE_K) {
+{% endif %}
+      loadSHMA(a_base, a_global_base, kidx, local_idx / 4u, local_idx % 4u);
+      loadSHMB(b_base, b_global_base, kidx, local_idx / 2u, local_idx % 2u);
+      workgroupBarrier();
+      for (var step = 0u; step < TILE_K; step = step + 8u) {
+{% set directInputs = false %}
+      let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
+{% for r in range(2) %}
+      var matA{{ r }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&{{ "w" if directInputs else "tile_A" }}, matrix_a_offset{% if r > 0 %} + 8u * {{ "K" if directInputs else "TILE_K" }}{% endif %}, {{ "K" if directInputs else "TILE_K" }});
+{% endfor %}
+      let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
+{% for c in range(4) %}
+      var matB{{ c }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&{{ "xm" if directInputs else "tile_B" }}, matrix_b_offset{% if c > 0 %} + {{ c * 8 }}u * TILE_K{% endif %}, {{ "N" if directInputs else "TILE_K" }});
+{% endfor %}
+      matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
+      matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);
+      matC02 = subgroupMatrixMultiplyAccumulate(matA0, matB2, matC02);
+      matC03 = subgroupMatrixMultiplyAccumulate(matA0, matB3, matC03);
+      matC10 = subgroupMatrixMultiplyAccumulate(matA1, matB0, matC10);
+      matC11 = subgroupMatrixMultiplyAccumulate(matA1, matB1, matC11);
+      matC12 = subgroupMatrixMultiplyAccumulate(matA1, matB2, matC12);
+      matC13 = subgroupMatrixMultiplyAccumulate(matA1, matB3, matC13);
+      }
+      workgroupBarrier();
+    }
+{% if GENERAL and TAIL %}
+    if (K_FULL < K) {
+      loadSHMAKTail(a_base, a_global_base, K_FULL, local_idx / 4u, local_idx % 4u);
+      loadSHMBKTail(b_base, b_global_base, K_FULL, local_idx / 2u, local_idx % 2u);
+      workgroupBarrier();
+      for (var step = 0u; step < TILE_K; step = step + 8u) {
+{% set directInputs = false %}
+      let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
+{% for r in range(2) %}
+      var matA{{ r }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&{{ "w" if directInputs else "tile_A" }}, matrix_a_offset{% if r > 0 %} + 8u * {{ "K" if directInputs else "TILE_K" }}{% endif %}, {{ "K" if directInputs else "TILE_K" }});
+{% endfor %}
+      let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
+{% for c in range(4) %}
+      var matB{{ c }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&{{ "xm" if directInputs else "tile_B" }}, matrix_b_offset{% if c > 0 %} + {{ c * 8 }}u * TILE_K{% endif %}, {{ "N" if directInputs else "TILE_K" }});
+{% endfor %}
+      matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
+      matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);
+      matC02 = subgroupMatrixMultiplyAccumulate(matA0, matB2, matC02);
+      matC03 = subgroupMatrixMultiplyAccumulate(matA0, matB3, matC03);
+      matC10 = subgroupMatrixMultiplyAccumulate(matA1, matB0, matC10);
+      matC11 = subgroupMatrixMultiplyAccumulate(matA1, matB1, matC11);
+      matC12 = subgroupMatrixMultiplyAccumulate(matA1, matB2, matC12);
+      matC13 = subgroupMatrixMultiplyAccumulate(matA1, matB3, matC13);
+      }
+      workgroupBarrier();
+    }
+{% endif %}
+    // The four scratch banks are reused across the two row-groups, and each is written
+    // by a collective subgroupMatrixStore then read CROSS-LANE by storeOutput. Barriers
+    // give the reads visibility of the store AND stop the second row-group's store from
+    // clobbering the first's still-in-flight readback when a partial final M-tile
+    // diverges storeOutput's guard. Without both barriers the last valid row can be corrupted.
+    subgroupMatrixStore<row_major>(&scratch[subtile_id][0], 0u, matC00, 8u);
+    subgroupMatrixStore<row_major>(&scratch[subtile_id][1], 0u, matC01, 8u);
+    subgroupMatrixStore<row_major>(&scratch[subtile_id][2], 0u, matC02, 8u);
+    subgroupMatrixStore<row_major>(&scratch[subtile_id][3], 0u, matC03, 8u);
+    workgroupBarrier();
+    let row = sg_id / 4u;
+    let col = (sg_id % 4u) * 2u;
+    let col_base = b_global_base + base_B;
+    var matrix_c_offset = c_base + (a_global_base + base_A) * N + col_base;
+    var row_limit = i32(M) - i32(a_global_base + base_A);
+    storeOutput(matrix_c_offset{% if needsColBase %}, col_base{% endif %}, row, col, subtile_id, row_limit);
+    workgroupBarrier();
+    subgroupMatrixStore<row_major>(&scratch[subtile_id][0], 0u, matC10, 8u);
+    subgroupMatrixStore<row_major>(&scratch[subtile_id][1], 0u, matC11, 8u);
+    subgroupMatrixStore<row_major>(&scratch[subtile_id][2], 0u, matC12, 8u);
+    subgroupMatrixStore<row_major>(&scratch[subtile_id][3], 0u, matC13, 8u);
+    workgroupBarrier();
+    matrix_c_offset = matrix_c_offset + 8u * N;
+    row_limit = i32(M) - i32(a_global_base + base_A + 8u);
+    storeOutput(matrix_c_offset{% if needsColBase %}, col_base{% endif %}, row, col, subtile_id, row_limit);
+    // Re-stage workgroup tiles/scratch before the next batch iteration reuses them.
+    workgroupBarrier();
+  }
+  // Re-stage workgroup tiles/scratch before the next M-tile iteration reuses them.
+  // The loop bound is workgroup-uniform (M is a compile-time const or a uniform
+  // read, and num_wg.y and workgroup_id.y are uniform), so every invocation
+  // reaches this barrier together.
+  workgroupBarrier();
+  }
+}

build/webgpu/matmul-tiled-general-reg.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,210 @@

+{{ env.wgsl.resourceDeclarations }}
+// Register-blocked MatMul for the no-subgroup-matrix
+// tier: Y = alpha * A @ B. It retains the bounds-checked addressing and
+// batch-broadcast of the general kernel, and its transposed-batch-A layout.
+// A retains its matrix-axis order; B may transpose its matrix axes.
+// 1-D operands use other variants.
+// Each thread computes a 4x4 micro-tile within a 64x64 workgroup tile, reusing
+// each staged operand across four accumulators. Both tiles are indexed by their
+// own output axis and group four K values per vector word, so the micro-tile
+// accumulates through dot() and one step reads TM + TN words rather than
+// 4 * (TM + TN) scalars. Logical matrix axes specialize to physical operand strides.
+{% set aR = aRank %}
+{% set bR = bRank %}
+{% set aBatchLen = aR - 2 %}
+{% set bBatchLen = bR - 2 %}
+{% set batchRank = aBatchLen %}
+{% set STATIC_M = aShape is defined %}
+{% set aDims = aShape if aShape is defined else (aBatchShape | default([])) %}
+{% set aTailStride = namespace(v=1) %}
+{% if aShape is defined %}{% for j in range(1, aR) %}{% set aTailStride.v = aTailStride.v * aShape[j] %}{% endfor %}{% endif %}
+{% set bTailStride = namespace(v=1) %}
+{% for j in range(1, bR) %}{% set bTailStride.v = bTailStride.v * bShape[j] %}{% endfor %}
+{% if aShape is defined %}
+{% set M = aShape[aR-2] %}{% set K = aShape[aR-1] %}{% endif %}
+{% set N = bShape[bR-1] %}
+{% set aMStride = K %}{% set aKStride = 1 %}{% set bKStride = 1 if bTransposed is defined and bTransposed else bShape[bR-1] %}
+{% set bNStride = bShape[bR-2] if bTransposed is defined and bTransposed else 1 %}{% set is_int = (scalar == "i32" or scalar == "u32") %}
+{% if is_int %}
+// Integer operands accumulate in their integer type, avoiding f32 rounding of
+// values outside the exact 24-bit significand range.
+{% endif %}
+{% set accT = scalar if is_int else "f32" %}
+{% set outScalar = outScalar if outScalar is defined else scalar %}
+{% set splitK = splitK if splitK is defined else 1 %}
+{% if scalar == "f16" %}
+// f16 operands stay packed in workgroup memory and widen on shared load;
+// accumulation remains f32.
+{% endif %}
+{% set tileT = scalar if scalar == "f16" else accT %}
+{% set kTile = 16 %}
+{% if STATIC_M %}
+const M: u32 = {{ M }}u;
+{% endif %}
+const K: u32 = {{ K }}u;
+const N: u32 = {{ N }}u;
+const A_M_STRIDE: u32 = {{ aMStride }}u;
+const A_K_STRIDE: u32 = {{ aKStride }}u;
+const B_K_STRIDE: u32 = {{ bKStride }}u;
+const B_N_STRIDE: u32 = {{ bNStride }}u;
+{% if is_int %}const ALPHA: {{ accT }} = {{ accT }}(1);{% else %}const ALPHA: f32 = f32({{ alpha }});{% endif %}
+// A 4x4 micro-tile over a 64x64 output tile reuses each staged operand across
+// four accumulators. It increases arithmetic work per load without the large
+// per-thread accumulator footprint of an 8x8 micro-tile.
+{% set microTile = 4 %}
+{% set lanes = 16 %}
+const BK: u32 = {{ kTile }}u;
+const BM: u32 = {{ registerTile }}u;
+const BN: u32 = {{ registerTile }}u;
+const TM: u32 = {{ microTile }}u;   // per-thread micro-tile rows
+const TN: u32 = {{ microTile }}u;   // per-thread micro-tile cols
+{% if splitK > 1 %}
+const SPLIT_K: u32 = {{ splitK }}u;
+const K_PER_SPLIT: u32 = {{ kPerSplit }}u;
+{% endif %}
+const K_VECS: u32 = BK / 4u;
+var<workgroup> tileA: array<array<vec4<{{ tileT }}>, K_VECS>, BM>;   // A[m][k/4]
+var<workgroup> tileB: array<array<vec4<{{ tileT }}>, K_VECS>, BN>;   // B[n][k/4]
+{% set hasBatchCoord = namespace(value=false) %}
+{% for i in range(batchRank) %}
+{% set axis = batchRank - 1 - i %}
+{% set aAxis = axis - (batchRank - aBatchLen) %}
+{% set bAxis = axis - (batchRank - bBatchLen) %}
+{% set aStored = (aAxis + 1) if (transBatchA and aAxis >= 0) else aAxis %}
+{% set bStored = bAxis %}
+{% set aDim = aDims[aStored] %}
+{% set bDim = bShape[bStored] if bStored >= 0 else 1 %}
+{% if aDim > 1 or bDim > 1 %}{% set hasBatchCoord.value = true %}{% endif %}
+{% endfor %}
+@compute @workgroup_size({{ lanes }}, {{ lanes }}, 1)
+fn main(
+  @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>
+) {
+{% if not STATIC_M %}
+  let M = params.M;
+{% endif %}
+  let mBase = wg.y * BM;
+  let nBase = wg.x * BN;
+  let li = lid.y * {{ lanes }}u + lid.x;
+  let zOut = wg.z;
+{% if splitK > 1 %}
+  let splitId = wg.z % SPLIT_K;
+{% else %}
+{% if hasBatchCoord.value %}
+  var zTmp = wg.z;
+{% endif %}
+{% endif %}
+  var aBatchOff: u32 = 0u;
+  var bBatchOff: u32 = 0u;
+{% for i in range(batchRank) %}
+{% set axis = batchRank - 1 - i %}
+{% set aAxis = axis - (batchRank - aBatchLen) %}
+{% set bAxis = axis - (batchRank - bBatchLen) %}
+{% set aStored = (aAxis + 1) if (transBatchA and aAxis >= 0) else aAxis %}
+{% set bStored = bAxis %}
+{% set aDim = aDims[aStored] %}
+{% set bDim = bShape[bStored] if bStored >= 0 else 1 %}
+{% set outDim = aDim if aDim >= bDim else bDim %}
+{% set aStride = namespace(v=1) %}
+{% if aStored >= 0 and aDim != 1 %}{% for j in range(aStored + 1, aR if STATIC_M else aR - 2) %}{% set aStride.v = aStride.v * aDims[j] %}{% endfor %}{% else %}{% set aStride.v = 0 %}{% endif %}
+{% set bStride = namespace(v=1) %}
+{% if bStored >= 0 and bDim != 1 %}{% for j in range(bStored + 1, bR) %}{% set bStride.v = bStride.v * bShape[j] %}{% endfor %}{% else %}{% set bStride.v = 0 %}{% endif %}{% if outDim > 1 %}
+  let c{{ axis }} = zTmp % {{ outDim }}u;
+  zTmp = zTmp / {{ outDim }}u;
+{% if aStride.v != 0 %}  aBatchOff = aBatchOff + c{{ axis }} * {% if STATIC_M %}{{ aStride.v }}u{% else %}{% if aStride.v != 1 %}{{ aStride.v }}u * {% endif %}M * K{% endif %};
+{% endif %}
+{% if bStride.v != 0 %}  bBatchOff = bBatchOff + c{{ axis }} * {{ bStride.v }}u;
+{% endif %}
+{% endif %}
+{% endfor %}
+  var acc: array<{{ accT }}, TM * TN>;   // [ti*TN + tj] for the TMxTN micro-tile
+  for (var i: u32 = 0u; i < TM * TN; i = i + 1u) { acc[i] = {{ accT }}(0); }
+{% if splitK > 1 %}
+  let kStart = splitId * K_PER_SPLIT;
+  let numTiles = K_PER_SPLIT / BK;
+{% else %}
+  let numTiles = (K + BK - 1u) / BK;
+{% endif %}
+  for (var kt: u32 = 0u; kt < numTiles; kt = kt + 1u) {
+{% if splitK > 1 %}
+    let kBase = kStart + kt * BK;
+{% else %}
+    let kBase = kt * BK;
+{% endif %}
+    // Cooperative load: one vector word per lane per pass. A's lanes walk K, which
+    // it stores contiguously; B's walk N, which it stores contiguously.
+    for (var idx: u32 = li; idx < BM * K_VECS; idx = idx + {{ lanes * lanes }}u) {
+      let ar = idx / K_VECS;
+      let ac4 = idx % K_VECS;
+      let am = mBase + ar;
+      let ak = kBase + ac4 * 4u;
+      var aWord = vec4<{{ tileT }}>({{ tileT }}(0));
+      if (am < M) {
+        let aRowOff = aBatchOff + am * A_M_STRIDE;
+{% for component in range(4) %}
+        if (ak + {{ component }}u < K) { aWord[{{ component }}u] = {{ tileT }}(a[aRowOff + (ak + {{ component }}u) * A_K_STRIDE]); }
+{% endfor %}
+      }
+      tileA[ar][ac4] = aWord;
+    }
+    for (var idx: u32 = li; idx < BN * K_VECS; idx = idx + {{ lanes * lanes }}u) {
+      let bc = idx % BN;
+      let br4 = idx / BN;
+      let bn = nBase + bc;
+      let bk = kBase + br4 * 4u;
+      var bWord = vec4<{{ tileT }}>({{ tileT }}(0));
+      if (bn < N) {
+        let bColOff = bBatchOff + bn * B_N_STRIDE;
+{% for component in range(4) %}
+        if (bk + {{ component }}u < K) { bWord[{{ component }}u] = {{ tileT }}(b[bColOff + (bk + {{ component }}u) * B_K_STRIDE]); }
+{% endfor %}
+      }
+      tileB[bc][br4] = bWord;
+    }
+    workgroupBarrier();
+{% set regT = accT %}{% filter indent(4, true) %}
+{% set regAccumulator = "acc" %}
+let aRow = lid.y * TM;
+let bCol = lid.x * TN;
+for (var kv: u32 = 0u; kv < BK / 4u; kv = kv + 1u) {
+  var av: array<vec4<{{ regT }}>, TM>;
+  var bv: array<vec4<{{ regT }}>, TN>;
+  for (var i: u32 = 0u; i < TM; i = i + 1u) { av[i] = vec4<{{ regT }}>(tileA[aRow + i][kv]); }
+  for (var j: u32 = 0u; j < TN; j = j + 1u) { bv[j] = vec4<{{ regT }}>(tileB[bCol + j][kv]); }
+  for (var i: u32 = 0u; i < TM; i = i + 1u) {
+    for (var j: u32 = 0u; j < TN; j = j + 1u) {
+{% if regSequentialK is defined and regSequentialK %}
+{% for component in range(4) %}
+      {{ regAccumulator }}[i * TN + j] = {{ regAccumulator }}[i * TN + j] + av[i][{{ component }}] * bv[j][{{ component }}];
+{% endfor %}
+{% else %}
+      {{ regAccumulator }}[i * TN + j] = {{ regAccumulator }}[i * TN + j] + dot(av[i], bv[j]);
+{% endif %}
+    }
+  }
+}
+{% endfilter %}
+    workgroupBarrier();
+  }
+  let rowBase = zOut * M * N;
+  let m0 = mBase + lid.y * TM;
+  let n0 = nBase + lid.x * TN;
+  for (var ti: u32 = 0u; ti < TM; ti = ti + 1u) {
+    let m = m0 + ti;
+    if (m >= M) { continue; }
+    for (var tj: u32 = 0u; tj < TN; tj = tj + 1u) {
+      let n = n0 + tj;
+      if (n < N) {
+        y[rowBase + m * N + n] = {{ outScalar }}(ALPHA * acc[ti * TN + tj]);
+      }
+    }
+  }
+}

build/webgpu/matmul-tiled-general.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,168 @@

+{{ env.wgsl.resourceDeclarations }}
+// Shared tiled matrix multiplication: Y = alpha * op(A) @ op(B), where op
+// transposes the last two axes when requested. This bounds-checked kernel
+// handles any M/K/N (no alignment requirement), all four transpose combinations,
+// right-aligned batch broadcasting, 1-D operand promotion (M==1 / N==1), and empty
+// K/M/N. transA/transB only change which stored stride the logical (m,k)/(k,n)
+// walk, using compiled stride constants here; the batch broadcast strides are
+// compiled the same way (a 0 literal means "broadcast / absent on that operand").
+{% set aR = aRank %}
+{% set bR = bRank %}
+{% set aVec = (aR == 1) %}
+{% set bVec = (bR == 1) %}
+{% set aBatchLen = (aR - 2) if aR >= 2 else 0 %}
+{% set bBatchLen = (bR - 2) if bR >= 2 else 0 %}
+{% set batchRank = aBatchLen if aBatchLen >= bBatchLen else bBatchLen %}
+/* transBatch maps stored [d0, d1, ..., dR-2, dR-1] to logical
+ * [d1, ..., dR-2, d0, dR-1]. Stored axis zero becomes the M axis, the final K
+ * axis is unchanged, and the remaining axes form the batch. This stride
+ * permutation composes with the ordinary last-two-axis transpose. */
+{% set STATIC_M = true %}
+{% set aDims = aShape if aShape is defined else (aBatchShape | default([])) %}
+{% set aTailStride = namespace(v=1) %}
+{% for j in range(1, aR) %}{% set aTailStride.v = aTailStride.v * aShape[j] %}{% endfor %}{% set bTailStride = namespace(v=1) %}
+{% for j in range(1, bR) %}{% set bTailStride.v = bTailStride.v * bShape[j] %}{% endfor %}
+{% if aVec %}{% set M = 1 %}{% set K = aShape[0] %}
+{% elif transA %}{% set M = aShape[aR-1] %}{% set K = aShape[aR-2] %}
+{% else %}{% set M = aShape[aR-2] %}{% set K = aShape[aR-1] %}{% endif %}
+{% if bVec %}{% set N = 1 %}
+{% elif transB %}{% set N = bShape[bR-2] %}
+{% else %}{% set N = bShape[bR-1] %}{% endif %}
+{% if aVec %}{% set aMStride = 0 %}{% set aKStride = 1 %}
+{% elif transA %}{% set aMStride = 1 %}{% set aKStride = M %}
+{% else %}{% set aMStride = K %}{% set aKStride = 1 %}{% endif %}
+{% if bVec %}{% set bKStride = 1 %}{% set bNStride = 0 %}
+{% elif transB %}{% set bKStride = 1 %}{% set bNStride = bShape[bR-1] %}
+{% else %}{% set bKStride = bShape[bR-1] %}{% set bNStride = 1 %}{% endif %}
+{% set is_int = (scalar == "i32" or scalar == "u32") %}
+{% set accT = scalar if is_int else "f32" %}
+{% set tileT = scalar if scalar == "f16" else accT %}
+const M: u32 = {{ M }}u;
+const K: u32 = {{ K }}u;
+const N: u32 = {{ N }}u;
+const A_M_STRIDE: u32 = {{ aMStride }}u;
+const A_K_STRIDE: u32 = {{ aKStride }}u;
+const B_K_STRIDE: u32 = {{ bKStride }}u;
+const B_N_STRIDE: u32 = {{ bNStride }}u;
+{% if is_int %}/* Integer matrix multiplication accumulates in the integer type. Widening
+ * through f32 would round values above 2^24. Integer MatMul has alpha = 1. */
+const ALPHA: {{ accT }} = {{ accT }}(1);{% else %}const ALPHA: f32 = f32({{ alpha }});{% endif %}
+// 2x2 register-blocked tile: 16x16 threads each compute a 2x2 micro-tile, for a
+// 32x32 output tile per workgroup with K stepped in BK=16 chunks. Each loaded
+// shared-mem element feeds 2 FMAs, favoring register reuse in the inner loop.
+{% set lanes = 16 %}
+const BK: u32 = {{ lanes }}u;
+const BM: u32 = {{ generalTile }}u;
+const BN: u32 = {{ generalTile }}u;
+var<workgroup> tileA: array<array<{{ tileT }}, {{ lanes }}>, {{ generalTile }}>;
+var<workgroup> tileB: array<array<{{ tileT }}, {{ generalTile }}>, {{ lanes }}>;
+{% set hasBatchCoord = namespace(value=false) %}
+{% for i in range(batchRank) %}
+{% set axis = batchRank - 1 - i %}
+{% set aAxis = axis - (batchRank - aBatchLen) %}
+{% set bAxis = axis - (batchRank - bBatchLen) %}
+{% set aStored = (aAxis + 1) if (transBatchA and aAxis >= 0) else aAxis %}
+{% set bStored = (bAxis + 1) if (transBatchB and bAxis >= 0) else bAxis %}
+{% set aDim = aDims[aStored] if aStored >= 0 else 1 %}
+{% set bDim = bShape[bStored] if bStored >= 0 else 1 %}
+{% if aDim > 1 or bDim > 1 %}{% set hasBatchCoord.value = true %}{% endif %}
+{% endfor %}
+@compute @workgroup_size({{ lanes }}, {{ lanes }}, 1)
+fn main(
+  @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>
+) {
+  let mBase = wg.y * BM;
+  let nBase = wg.x * BN;
+  let li = lid.y * {{ lanes }}u + lid.x;
+  // Per-batch base offsets into A and B using right-aligned broadcast strides.
+  // Decompose the flat output-batch index from the innermost axis outward.
+  let zOut = wg.z;
+{% if hasBatchCoord.value %}
+  var zTmp = wg.z;
+{% endif %}
+  var aBatchOff: u32 = 0u;
+  var bBatchOff: u32 = 0u;
+{% for i in range(batchRank) %}
+{% set axis = batchRank - 1 - i %}
+{% set aAxis = axis - (batchRank - aBatchLen) %}
+{% set bAxis = axis - (batchRank - bBatchLen) %}
+{% set aStored = (aAxis + 1) if (transBatchA and aAxis >= 0) else aAxis %}
+{% set bStored = (bAxis + 1) if (transBatchB and bAxis >= 0) else bAxis %}
+{% set aDim = aDims[aStored] if aStored >= 0 else 1 %}
+{% set bDim = bShape[bStored] if bStored >= 0 else 1 %}
+{% set outDim = aDim if aDim >= bDim else bDim %}
+{% set aStride = namespace(v=1) %}
+{% if aStored >= 0 and aDim != 1 %}{% for j in range(aStored + 1, aR if STATIC_M else aR - 2) %}{% set aStride.v = aStride.v * aDims[j] %}{% endfor %}{% else %}{% set aStride.v = 0 %}{% endif %}
+{% set bStride = namespace(v=1) %}
+{% if bStored >= 0 and bDim != 1 %}{% for j in range(bStored + 1, bR) %}{% set bStride.v = bStride.v * bShape[j] %}{% endfor %}{% else %}{% set bStride.v = 0 %}{% endif %}
+{% if outDim > 1 %}
+{% if aStride.v != 0 or bStride.v != 0 %}
+  let c{{ axis }} = zTmp % {{ outDim }}u;
+{% endif %}
+  zTmp = zTmp / {{ outDim }}u;
+{% if aStride.v != 0 %}  aBatchOff = aBatchOff + c{{ axis }} * {{ aStride.v }}u;
+{% endif %}
+{% if bStride.v != 0 %}  bBatchOff = bBatchOff + c{{ axis }} * {{ bStride.v }}u;
+{% endif %}
+{% endif %}
+{% endfor %}
+  var acc00: {{ accT }} = {{ accT }}(0);
+  var acc01: {{ accT }} = {{ accT }}(0);
+  var acc10: {{ accT }} = {{ accT }}(0);
+  var acc11: {{ accT }} = {{ accT }}(0);
+  let numTiles = (K + BK - 1u) / BK;
+  for (var kt: u32 = 0u; kt < numTiles; kt = kt + 1u) {
+    let kBase = kt * BK;
+    // Cooperative load: 32x16 A tile + 16x32 B tile, 256 threads x 2 each.
+    for (var e: u32 = 0u; e < 2u; e = e + 1u) {
+      let idx = li + e * {{ lanes * lanes }}u;
+      let ar = idx / BK;
+      let ac = idx % BK;
+      let am = mBase + ar;
+      let ak = kBase + ac;
+      if (am < M && ak < K) {
+        tileA[ar][ac] = {{ tileT }}(a[aBatchOff + am * A_M_STRIDE + ak * A_K_STRIDE]);
+      } else {
+        tileA[ar][ac] = {{ tileT }}(0);
+      }
+      let br = idx / BN;
+      let bc = idx % BN;
+      let bk = kBase + br;
+      let bn = nBase + bc;
+      if (bk < K && bn < N) {
+        tileB[br][bc] = {{ tileT }}(b[bBatchOff + bk * B_K_STRIDE + bn * B_N_STRIDE]);
+      } else {
+        tileB[br][bc] = {{ tileT }}(0);
+      }
+    }
+    workgroupBarrier();
+    for (var kk: u32 = 0u; kk < BK; kk = kk + 1u) {
+      let a0 = {{ accT }}(tileA[lid.y * 2u][kk]);
+      let a1 = {{ accT }}(tileA[lid.y * 2u + 1u][kk]);
+      let b0 = {{ accT }}(tileB[kk][lid.x * 2u]);
+      let b1 = {{ accT }}(tileB[kk][lid.x * 2u + 1u]);
+      acc00 = acc00 + a0 * b0;
+      acc01 = acc01 + a0 * b1;
+      acc10 = acc10 + a1 * b0;
+      acc11 = acc11 + a1 * b1;
+    }
+    workgroupBarrier();
+  }
+  let m0 = mBase + lid.y * 2u;
+  let m1 = m0 + 1u;
+  let n0 = nBase + lid.x * 2u;
+  let n1 = n0 + 1u;
+  let rowBase = zOut * M * N;
+  if (m0 < M && n0 < N) { y[rowBase + m0 * N + n0] = {{ scalar }}(ALPHA * acc00); }
+  if (m0 < M && n1 < N) { y[rowBase + m0 * N + n1] = {{ scalar }}(ALPHA * acc01); }
+  if (m1 < M && n0 < N) { y[rowBase + m1 * N + n0] = {{ scalar }}(ALPHA * acc10); }
+  if (m1 < M && n1 < N) { y[rowBase + m1 * N + n1] = {{ scalar }}(ALPHA * acc11); }
+}

build/webgpu/matmul-vector-matrix-vec4.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,55 @@

+// GEMV specialization for y[N] = a[K] @ B[K, N]. Each workgroup owns 32
+// consecutive vec4 column groups, or 128 output columns.
+{{ env.wgsl.resourceDeclarations }}
+{% set OUT = outputBuffer if outputBuffer is defined else "c" %}
+const LANES: u32 = {{ gemvLanes }}u;
+// SLICES partitions the K reduction across the workgroup's second dimension.
+// Thread zero of each column group combines the slice partials in index order.
+const SLICES: u32 = {{ gemvSlices }}u;
+var<workgroup> partials: array<vec4<f32>, LANES * SLICES>;
+@compute @workgroup_size({{ gemvLanes }}, {{ gemvSlices }}, 1)
+fn main(
+  @builtin(workgroup_id) workgroup_id: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>
+) {
+  let lane = lid.x;
+  let slice = lid.y;
+  let cg = workgroup_id.x * LANES + lane;
+  var acc = vec4<f32>(0.0);
+  if (cg < params.N4) {
+{% if unrollK2 is defined and unrollK2 %}
+    // Two K positions per iteration amortize loop/address arithmetic on long
+    // decode projections while preserving each slice's exact strided order.
+    var k = slice;
+    for (; k + SLICES < params.K; k = k + 2u * SLICES) {
+      acc = acc + f32(a[k]) * vec4<f32>(b[k * params.N4 + cg]);
+      let k1 = k + SLICES;
+      acc = acc + f32(a[k1]) * vec4<f32>(b[k1 * params.N4 + cg]);
+    }
+    if (k < params.K) {
+      acc = acc + f32(a[k]) * vec4<f32>(b[k * params.N4 + cg]);
+    }
+{% else %}
+    for (var k = slice; k < params.K; k = k + SLICES) {
+      acc = acc + f32(a[k]) * vec4<f32>(b[k * params.N4 + cg]);
+    }
+{% endif %}
+  }
+  partials[slice * LANES + lane] = acc;
+  workgroupBarrier();
+  if (slice == 0u && cg < params.N4) {
+    var total = partials[lane];
+    for (var s = 1u; s < SLICES; s = s + 1u) {
+      total = total + partials[s * LANES + lane];
+    }
+{% if alphaScale is defined and alphaScale != 1 %}
+    // This specialization bakes the output multiplier into the shader and
+    // applies it after combining the K slices.
+    total = total * f32({{ alphaScale }});
+{% endif %}
+    {{ OUT }}[cg] = vec4<{{ T }}>(total);
+  }
+}

build/webgpu/metadata.json ADDED Viewed

	@@ -0,0 +1,41 @@

+{
+  "name": "com.microsoft.TransposeMatMul",
+  "id": "_com_microsoft_transposematmul_webgpu_45fbc93",
+  "version": 1,
+  "license": "Apache-2.0",
+  "backend": { "type": "webgpu" },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "bench.json": "XEFUj4H5WRbzIOZky4TK16K+wPpc4vYBtmqD+f6fhvA=",
+      "fused-matmul-subgroup-matrix.wgsl.jinja": "fqGrZsIrSx70rZzmpWys4oTG+JSOnAbqtBdQ3xwQpSw=",
+      "manifest.json": "NeYkA+AcZ/38wnrBvqcXxIiJM296w8JwhwTrkX+l40w=",
+      "matmul-band-vec4.wgsl.jinja": "LKAs6A++OJF0ZEEM4JITZmKXkr/wo9gd5qrR3zDZy1g=",
+      "matmul-subgroup-matrix-ext.wgsl.jinja": "DJxg5GNgH4xumnum/RzuWA5x+JNzB0Tx7U5RMsnlm6o=",
+      "matmul-tiled-general-reg.wgsl.jinja": "X8kCdoMgwvyMiCnLl/UDpfpdVcYFqTLlj6+o+R3pJlQ=",
+      "matmul-tiled-general.wgsl.jinja": "9uRxra70oG70IJCiKqA0DG76W0zS46RxLdT6fndR48w=",
+      "matmul-vector-matrix-vec4.wgsl.jinja": "Qiv2AO8MWj1BAoPLMCH98oXBasrQYFWFh/IUhdhZp5I=",
+      "reduce-axis0-splitk-combine.wgsl.jinja": "Zf7tz8nrapi2KMZIbv8cT2f4pliJoJHj4iwS4hc/6ms=",
+      "test.json": "h+wSc5e5ergXuijy9nMDICwCmTvIi54cexuSorODaWM="
+    }
+  },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.1",
+    "variants": {
+      "broadcast_transb_tiled_reg": ["matmul-tiled-general-reg.wgsl.jinja"],
+      "broadcast_transb_subgroup_matrix_f16": ["matmul-subgroup-matrix-ext.wgsl.jinja"],
+      "broadcast_transb_subgroup_matrix_f32": ["matmul-subgroup-matrix-ext.wgsl.jinja"],
+      "m1_gemv_vec4": ["matmul-vector-matrix-vec4.wgsl.jinja"],
+      "rank2_band_vec4_splitk": ["matmul-band-vec4.wgsl.jinja", "reduce-axis0-splitk-combine.wgsl.jinja"],
+      "rank2_band_vec4": ["matmul-band-vec4.wgsl.jinja"],
+      "rank2_band_vec4_f32_preferred": ["matmul-band-vec4.wgsl.jinja"],
+      "subgroup_matrix_splitk": ["matmul-subgroup-matrix-ext.wgsl.jinja", "reduce-axis0-splitk-combine.wgsl.jinja"],
+      "subgroup_matrix_tail_broadcast": ["matmul-subgroup-matrix-ext.wgsl.jinja"],
+      "subgroup_matrix": ["fused-matmul-subgroup-matrix.wgsl.jinja"],
+      "broadcast_rank4_tiled_reg": ["matmul-tiled-general-reg.wgsl.jinja"],
+      "plain_rank2_tiled_reg": ["matmul-tiled-general-reg.wgsl.jinja"],
+      "tiled": ["matmul-tiled-general.wgsl.jinja"]
+    }
+  }
+}

build/webgpu/reduce-axis0-splitk-combine.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,95 @@

+// Pass 2 of the split-K column-wise reduction. One invocation per output column
+// folds the segment partials and applies the selected reduction's final step.
+// Segments are folded in ascending order for deterministic results. This order
+// differs from the single-pass reduction but remains within the f32 tolerance.
+{% set yv = "f16(" if outputF16 else "" %}
+{% set vy = ")" if outputF16 else "" %}
+{{ env.wgsl.resourceDeclarations }}
+{% macro wgsl_minmax_identity(name, op, scalar="f32") %}
+/* Exact {{ op }} reduction identity. WGSL rejects infinity during constant
+ * evaluation, so an f32 identity is constructed from its IEEE-754 bits. */
+fn {{ name }}() -> {{ scalar }} {
+  var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
+  return bitcast<f32>(bits);
+}{% endmacro %}
+const WG: u32 = {{ workgroupSize }}u;
+const SPLIT: u32 = {{ split }}u;
+{% if op == "logsumexp" %}
+const F32_MIN: f32 = -3.4028234663852886e38;
+const F32_MAX: f32 = 3.4028234663852886e38;
+fn is_nan_f32(value: f32) -> bool {
+  let bits = bitcast<u32>(value);
+  return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
+}
+{% elif op == "max" or op == "min" %}
+{{ wgsl_minmax_identity("reduction_identity", op) }}
+{% endif %}
+@compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>,
+        @builtin(num_workgroups) nwg: vec3<u32>) {
+  // The start already folds gid.y in, so the stride must span every y row too;
+  // an x-only stride would send y = 0 lanes over columns the y >= 1 rows own.
+  let stride = nwg.x * nwg.y * WG;
+  let start = (gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG) + gid.x;
+  for (var col = start; col < params.cols; col = col + stride) {
+{% if op == "logsumexp" %}
+    // Merge SPLIT (segMax, segSumExp) pairs stably; carry NaN / +Inf markers.
+    var nan_value = 0.0;
+    var has_nan = false;
+    var global_max = F32_MIN;
+    for (var seg = 0u; seg < SPLIT; seg = seg + 1u) {
+      let nv = partials[(2u * SPLIT + seg) * params.cols + col];
+      if (nv != 0.0 || is_nan_f32(nv)) {
+        has_nan = true;
+        nan_value = nv;
+      }
+      global_max = max(global_max, partials[seg * params.cols + col]);
+    }
+    var sum = 0.0;
+    for (var seg = 0u; seg < SPLIT; seg = seg + 1u) {
+      let seg_max = partials[seg * params.cols + col];
+      let seg_sum = partials[(SPLIT + seg) * params.cols + col];
+      sum = sum + seg_sum * exp(seg_max - global_max);
+    }
+    let has_positive_inf = global_max > F32_MAX;
+    let finite_or_inf = select(global_max + log(sum), global_max, has_positive_inf);
+    y[col] = {{ yv }}select(finite_or_inf, nan_value, has_nan){{ vy }};
+{% else %}
+{% if op == "max" %}
+    var total = reduction_identity();
+{% elif op == "min" %}
+    var total = reduction_identity();
+{% elif op == "prod" %}
+    var total = 1.0;
+{% else %}
+    var total = 0.0;
+{% endif %}
+    for (var seg = 0u; seg < SPLIT; seg = seg + 1u) {
+      let p = partials[seg * params.cols + col];
+{% if op == "max" or op == "min" %}
+      total = {{ op }}(total, p);
+{% elif op == "prod" %}
+      total = total * p;
+{% else %}
+      total = total + p;
+{% endif %}
+    }
+{% if op == "l2" %}
+    y[col] = {{ yv }}sqrt(total){{ vy }};
+{% elif op == "logsum" %}
+    y[col] = {{ yv }}log(total){{ vy }};
+{% elif op == "mean" %}
+    y[col] = {{ yv }}total / f32(params.rows){{ vy }};
+{% else %}
+{% if outputF16 %}
+    y[col] = f16(total);
+{% else %}
+    y[col] = total;
+{% endif %}
+{% endif %}
+{% endif %}
+  }
+}

build/webgpu/test.json ADDED Viewed

	@@ -0,0 +1,1973 @@

+{
+  "fixtureArrays": {
+    "ort_float32_broadcast_rank3_by_rank4_output_Y": [1, 3, 5, 33, 43, 53, 5, 23, 41, 85, 111, 137, 9, 43, 77, 137, 179, 221],
+    "ort_float32_rank3_by_rank2_output_Y": [20, 23, 26, 29, 56, 68, 80, 92, 92, 113, 134, 155, 128, 158, 188, 218],
+    "ort_float32_batched_rank4_input_A": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]
+  },
+  "cases": [
+    {
+      "name": "ort_float32_broadcast_rank4_by_rank3",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
+        "test": "FusedMatMulOpTest.FloatTypeNoTranspose",
+        "notes": "Derived from the com.microsoft.FusedMatMul corpus. Upstream registers one kernel for both operators and TransposeMatMul is the strict subset that omits transBatchA/transBatchB, so an expectation taken with those flags at zero describes this operator exactly."
+      },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [3, 1, 1, 2],
+          "data": { "kind": "values", "values": [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [2, 2, 2],
+          "data": { "kind": "values", "values": [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0] }
+        }
+      },
+      "outputs": {
+        "Y": {
+          "dtype": "float32",
+          "shape": [3, 2, 1, 2],
+          "tolerance": 0.000001,
+          "data": { "kind": "values", "values": [2.0, 3.0, 6.0, 7.0, 6.0, 11.0, 26.0, 31.0, 10.0, 19.0, 46.0, 55.0] }
+        }
+      }
+    },
+    {
+      "name": "ort_float32_broadcast_rank3_by_rank4",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
+        "test": "FusedMatMulOpTest.FloatTypeNoTranspose",
+        "notes": "Derived from the com.microsoft.FusedMatMul corpus. Upstream registers one kernel for both operators and TransposeMatMul is the strict subset that omits transBatchA/transBatchB, so an expectation taken with those flags at zero describes this operator exactly."
+      },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 3, 2],
+          "data": { "kind": "values", "values": [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0] }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [3, 2, 2, 1],
+          "data": { "kind": "values", "values": [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0] }
+        }
+      },
+      "outputs": {
+        "Y": {
+          "dtype": "float32",
+          "shape": [3, 2, 3, 1],
+          "tolerance": 0.000001,
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/ort_float32_broadcast_rank3_by_rank4_output_Y" }
+          }
+        }
+      }
+    },
+    {
+      "name": "ort_float32_left_1d_batched_rhs",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
+        "test": "FusedMatMulOpTest.FloatTypeNoTranspose",
+        "notes": "Derived from the com.microsoft.FusedMatMul corpus. Upstream registers one kernel for both operators and TransposeMatMul is the strict subset that omits transBatchA/transBatchB, so an expectation taken with those flags at zero describes this operator exactly."
+      },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [2], "data": { "kind": "values", "values": [0.0, 1.0] } },
+        "B": {
+          "dtype": "float32",
+          "shape": [3, 2, 1],
+          "data": { "kind": "values", "values": [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] }
+        }
+      },
+      "outputs": {
+        "Y": {
+          "dtype": "float32",
+          "shape": [3, 1],
+          "tolerance": 0.000001,
+          "data": { "kind": "values", "values": [1.0, 3.0, 5.0] }
+        }
+      }
+    },
+    {
+      "name": "ort_float32_right_1d_batched_lhs",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
+        "test": "FusedMatMulOpTest.FloatTypeNoTranspose",
+        "notes": "Derived from the com.microsoft.FusedMatMul corpus. Upstream registers one kernel for both operators and TransposeMatMul is the strict subset that omits transBatchA/transBatchB, so an expectation taken with those flags at zero describes this operator exactly."
+      },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [3, 1, 2],
+          "data": { "kind": "values", "values": [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] }
+        },
+        "B": { "dtype": "float32", "shape": [2], "data": { "kind": "values", "values": [0.0, 1.0] } }
+      },
+      "outputs": {
+        "Y": {
+          "dtype": "float32",
+          "shape": [3, 1],
+          "tolerance": 0.000001,
+          "data": { "kind": "values", "values": [1.0, 3.0, 5.0] }
+        }
+      }
+    },
+    {
+      "name": "ort_float32_plain_2d",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
+        "test": "FusedMatMulOpTest.FloatTypeNoTranspose",
+        "notes": "Derived from the com.microsoft.FusedMatMul corpus. Upstream registers one kernel for both operators and TransposeMatMul is the strict subset that omits transBatchA/transBatchB, so an expectation taken with those flags at zero describes this operator exactly."
+      },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "values", "values": [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0] }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [4, 3],
+          "data": { "kind": "values", "values": [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0] }
+        }
+      },
+      "outputs": {
+        "Y": {
+          "dtype": "float32",
+          "shape": [3, 3],
+          "tolerance": 0.000001,
+          "data": { "kind": "values", "values": [42.0, 48.0, 54.0, 114.0, 136.0, 158.0, 186.0, 224.0, 262.0] }
+        }
+      }
+    },
+    {
+      "name": "ort_float32_rank3_by_rank2",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
+        "test": "FusedMatMulOpTest.FloatTypeNoTranspose",
+        "notes": "Derived from the com.microsoft.FusedMatMul corpus. Upstream registers one kernel for both operators and TransposeMatMul is the strict subset that omits transBatchA/transBatchB, so an expectation taken with those flags at zero describes this operator exactly."
+      },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 2, 3],
+          "data": { "kind": "values", "values": [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0] }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "values", "values": [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0] }
+        }
+      },
+      "outputs": {
+        "Y": {
+          "dtype": "float32",
+          "shape": [2, 2, 4],
+          "tolerance": 0.000001,
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/ort_float32_rank3_by_rank2_output_Y" } }
+        }
+      }
+    },
+    {
+      "name": "ort_float32_rank3_by_broadcast_rank3",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
+        "test": "FusedMatMulOpTest.FloatTypeNoTranspose",
+        "notes": "Derived from the com.microsoft.FusedMatMul corpus. Upstream registers one kernel for both operators and TransposeMatMul is the strict subset that omits transBatchA/transBatchB, so an expectation taken with those flags at zero describes this operator exactly."
+      },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 2, 3],
+          "data": { "kind": "values", "values": [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0] }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [1, 3, 4],
+          "data": { "kind": "values", "values": [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0] }
+        }
+      },
+      "outputs": {
+        "Y": {
+          "dtype": "float32",
+          "shape": [2, 2, 4],
+          "tolerance": 0.000001,
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/ort_float32_rank3_by_rank2_output_Y" } }
+        }
+      }
+    },
+    {
+      "name": "ort_float32_singleton_rank3_by_rank3",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
+        "test": "FusedMatMulOpTest.FloatTypeNoTranspose",
+        "notes": "Derived from the com.microsoft.FusedMatMul corpus. Upstream registers one kernel for both operators and TransposeMatMul is the strict subset that omits transBatchA/transBatchB, so an expectation taken with those flags at zero describes this operator exactly."
+      },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [1, 2, 3],
+          "data": { "kind": "values", "values": [0.0, 1.0, 2.0, 3.0, 4.0, 5.0] }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [1, 3, 4],
+          "data": { "kind": "values", "values": [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0] }
+        }
+      },
+      "outputs": {
+        "Y": {
+          "dtype": "float32",
+          "shape": [1, 2, 4],
+          "tolerance": 0.000001,
+          "data": { "kind": "values", "values": [20.0, 23.0, 26.0, 29.0, 56.0, 68.0, 80.0, 92.0] }
+        }
+      }
+    },
+    {
+      "name": "ort_float32_batched_rank4",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
+        "test": "FusedMatMulOpTest.FloatTypeNoTranspose",
+        "notes": "Derived from the com.microsoft.FusedMatMul corpus. Upstream registers one kernel for both operators and TransposeMatMul is the strict subset that omits transBatchA/transBatchB, so an expectation taken with those flags at zero describes this operator exactly."
+      },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 2, 2, 2],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/ort_float32_batched_rank4_input_A" } }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [2, 2, 2, 2],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/ort_float32_batched_rank4_input_A" } }
+        }
+      },
+      "outputs": {
+        "Y": {
+          "dtype": "float32",
+          "shape": [2, 2, 2, 2],
+          "tolerance": 0.000001,
+          "data": {
+            "kind": "values",
+            "values": [2.0, 3.0, 6.0, 11.0, 46.0, 55.0, 66.0, 79.0, 154.0, 171.0, 190.0, 211.0, 326.0, 351.0, 378.0, 407.0]
+          }
+        }
+      }
+    },
+    {
+      "name": "ort_float32_broadcast_rank4_by_rank4",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
+        "test": "FusedMatMulOpTest.FloatTypeNoTranspose",
+        "notes": "Derived from the com.microsoft.FusedMatMul corpus. Upstream registers one kernel for both operators and TransposeMatMul is the strict subset that omits transBatchA/transBatchB, so an expectation taken with those flags at zero describes this operator exactly."
+      },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [1, 2, 3, 2],
+          "data": { "kind": "values", "values": [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0] }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [3, 2, 2, 1],
+          "data": { "kind": "values", "values": [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0] }
+        }
+      },
+      "outputs": {
+        "Y": {
+          "dtype": "float32",
+          "shape": [3, 2, 3, 1],
+          "tolerance": 0.000001,
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/ort_float32_broadcast_rank3_by_rank4_output_Y" }
+          }
+        }
+      }
+    },
+    {
+      "name": "ort_float32_vector_dot_scalar_output",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
+        "test": "FusedMatMulOpTest.FloatTypeNoTranspose",
+        "notes": "Derived from the com.microsoft.FusedMatMul corpus. Upstream registers one kernel for both operators and TransposeMatMul is the strict subset that omits transBatchA/transBatchB, so an expectation taken with those flags at zero describes this operator exactly."
+      },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [3], "data": { "kind": "values", "values": [0.0, 1.0, 2.0] } },
+        "B": { "dtype": "float32", "shape": [3], "data": { "kind": "values", "values": [0.0, 1.0, 2.0] } }
+      },
+      "outputs": {
+        "Y": { "dtype": "float32", "shape": [], "tolerance": 0.000001, "data": { "kind": "values", "values": [5.0] } }
+      }
+    },
+    {
+      "name": "ort_float32_alpha_zero_outputs_zero",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
+        "test": "FusedMatMulOpTest.DoubleTypeAlphaZero",
+        "notes": "Diverges from the upstream test's inputs (inputs.A values [1.0, 2.0, 3.0, 4.0] -> constant 2.0; inputs.B values [5.0, 6.0, 7.0, 8.0] -> constant 3.0); the expected output is recomputed by the CPU reference for the new inputs. A zero alpha scales the whole product away, so no operand value can reach the result and both operands are uniform fills."
+      },
+      "attrs": { "alpha": 0 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "constant", "value": 2.0 } },
+        "B": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "constant", "value": 3.0 } }
+      },
+      "outputs": {
+        "Y": {
+          "dtype": "float32",
+          "shape": [2, 2],
+          "tolerance": 0.000001,
+          "data": { "kind": "values", "values": [0.0, 0.0, 0.0, 0.0] }
+        }
+      }
+    },
+    {
+      "name": "ort_float32_empty_k_dimension_outputs_zero",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
+        "test": "FusedMatMulOpTest.DoubleTypeEmptyKDim",
+        "notes": "Derived from the com.microsoft.FusedMatMul corpus. Upstream registers one kernel for both operators and TransposeMatMul is the strict subset that omits transBatchA/transBatchB, so an expectation taken with those flags at zero describes this operator exactly."
+      },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [2, 0], "data": { "kind": "values", "values": [] } },
+        "B": { "dtype": "float32", "shape": [0, 3], "data": { "kind": "values", "values": [] } }
+      },
+      "outputs": {
+        "Y": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "tolerance": 0.000001,
+          "data": { "kind": "values", "values": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0] }
+        }
+      }
+    },
+    {
+      "name": "ort_float32_transpose_a_scaled",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
+        "test": "FusedMatMulOpTest.DoubleTypeScale",
+        "notes": "Derived from the com.microsoft.FusedMatMul corpus. Upstream registers one kernel for both operators and TransposeMatMul is the strict subset that omits transBatchA/transBatchB, so an expectation taken with those flags at zero describes this operator exactly."
+      },
+      "attrs": { "alpha": 0.5, "transA": 1 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "values", "values": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "values", "values": [7.0, 8.0, 9.0, 10.0, 11.0, 12.0] }
+        }
+      },
+      "outputs": {
+        "Y": {
+          "dtype": "float32",
+          "shape": [3, 3],
+          "tolerance": 0.000001,
+          "data": { "kind": "values", "values": [23.5, 26.0, 28.5, 32.0, 35.5, 39.0, 40.5, 45.0, 49.5] }
+        }
+      }
+    },
+    {
+      "name": "ort_float32_transpose_b",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
+        "test": "FusedMatMulOpTest.FloatTypeTransposeB",
+        "notes": "Derived from the com.microsoft.FusedMatMul corpus. Upstream registers one kernel for both operators and TransposeMatMul is the strict subset that omits transBatchA/transBatchB, so an expectation taken with those flags at zero describes this operator exactly."
+      },
+      "attrs": { "transB": 1 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "values", "values": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [4, 3],
+          "data": { "kind": "values", "values": [7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0, 18.0] }
+        }
+      },
+      "outputs": {
+        "Y": {
+          "dtype": "float32",
+          "shape": [2, 4],
+          "tolerance": 0.000001,
+          "data": { "kind": "values", "values": [50.0, 68.0, 86.0, 104.0, 122.0, 167.0, 212.0, 257.0] }
+        }
+      }
+    },
+    {
+      "name": "ort_float32_transpose_ab_scaled",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
+        "test": "FusedMatMulOpTest.FloatTypeScale",
+        "notes": "Derived from the com.microsoft.FusedMatMul corpus. Upstream registers one kernel for both operators and TransposeMatMul is the strict subset that omits transBatchA/transBatchB, so an expectation taken with those flags at zero describes this operator exactly."
+      },
+      "attrs": { "alpha": 4, "transA": 1, "transB": 1 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [3, 2],
+          "data": { "kind": "values", "values": [1.0, 4.0, 2.0, 5.0, 3.0, 6.0] }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [4, 3],
+          "data": { "kind": "values", "values": [7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0, 18.0] }
+        }
+      },
+      "outputs": {
+        "Y": {
+          "dtype": "float32",
+          "shape": [2, 4],
+          "tolerance": 0.000001,
+          "data": { "kind": "values", "values": [200.0, 272.0, 344.0, 416.0, 488.0, 668.0, 848.0, 1028.0] }
+        }
+      }
+    },
+    {
+      "name": "ort_float32_scaled_no_transpose",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
+        "test": "FusedMatMulOpTest.FloatTypeScale",
+        "notes": "Derived from the com.microsoft.FusedMatMul corpus. Upstream registers one kernel for both operators and TransposeMatMul is the strict subset that omits transBatchA/transBatchB, so an expectation taken with those flags at zero describes this operator exactly."
+      },
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "values", "values": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [3, 2],
+          "data": { "kind": "values", "values": [7.0, 8.0, 9.0, 10.0, 11.0, 12.0] }
+        }
+      },
+      "outputs": {
+        "Y": {
+          "dtype": "float32",
+          "shape": [2, 2],
+          "tolerance": 0.000001,
+          "data": { "kind": "values", "values": [29.0, 32.0, 69.5, 77.0] }
+        }
+      }
+    },
+    {
+      "name": "ort_float32_empty_input_m_zero",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
+        "test": "FusedMatMulOpTest.DoubleTypeEmptyInput",
+        "notes": "Derived from the com.microsoft.FusedMatMul corpus. Upstream registers one kernel for both operators and TransposeMatMul is the strict subset that omits transBatchA/transBatchB, so an expectation taken with those flags at zero describes this operator exactly."
+      },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [0, 3], "data": { "kind": "values", "values": [] } },
+        "B": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "values", "values": [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0] }
+        }
+      },
+      "outputs": {
+        "Y": { "dtype": "float32", "shape": [0, 4], "tolerance": 0.000001, "data": { "kind": "values", "values": [] } }
+      }
+    },
+    {
+      "name": "aligned_plain_64x32x64",
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.027, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [64, 64], "tolerance": 0.0001 } }
+    },
+    {
+      "name": "register_blocked_plain_512x64x512_alpha_scaled",
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [512, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.027, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [64, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [512, 512], "tolerance": 0.0001 } },
+      "attrs": { "alpha": 0.5 },
+      "provenance": {
+        "notes": "Rank-2 M=N=512 and K=64 produce 64 aligned 64x64 workgroup tiles, exercising register-blocked vec4 staging and 4x4 per-thread accumulation. alpha=0.5 verifies scaling in the output epilogue."
+      }
+    },
+    {
+      "name": "aligned_transB_alpha_64x32",
+      "attrs": { "transB": 1, "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.023, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.009, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [64, 64], "tolerance": 0.0001 } }
+    },
+    {
+      "name": "aligned_batched_plain_2x64x32x64",
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.027, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [2, 32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 64, 64], "tolerance": 0.0001 } }
+    },
+    {
+      "name": "aligned_batched_transB_alpha_2x64x32",
+      "attrs": { "transB": 1, "alpha": 0.25 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.023, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [2, 64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.009, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 64, 64], "tolerance": 0.0001 } }
+    },
+    {
+      "name": "aligned_mtail_50x32x128",
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [50, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.015, "cosStep": 0.021, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [32, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.011, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [50, 128], "tolerance": 0.0001 } }
+    },
+    {
+      "name": "aligned_transA_64x32",
+      "attrs": { "transA": 1 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.027, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [64, 64], "tolerance": 0.0001 } }
+    },
+    {
+      "name": "aligned_transA_transB_alpha_64x32",
+      "attrs": { "transA": 1, "transB": 1, "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.023, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.009, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [64, 64], "tolerance": 0.0001 } }
+    },
+    {
+      "name": "f32_subgroup_matrix_subnormal_dot_products_gpu_gap",
+      "skipGpu": {
+        "category": "permanent",
+        "reason": "Portable WGSL floating-point semantics do not guarantee preservation of the subnormal values required by this fixture. Backend evidence: WebGPU/Metal flushes subnormals to zero in f32; the ~3e-39 subnormal dot products collapse to zero (subgroup-matrix path)."
+      },
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
+        "test": "FusedMatMulOpTest.FloatTypeNoTranspose",
+        "notes": "M=32, K=32, N=64 with finite subnormal dot products that must not flush to zero."
+      },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [32, 32], "data": { "kind": "constant", "value": 1e-20 } },
+        "B": { "dtype": "float32", "shape": [32, 64], "data": { "kind": "constant", "value": 1e-20 } }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [32, 64], "tolerance": 1e-43 } }
+    },
+    {
+      "name": "f32_subgroup_matrix_scaled_subnormal_dot_products_gpu_gap",
+      "skipGpu": {
+        "category": "permanent",
+        "reason": "Portable WGSL floating-point semantics do not guarantee preservation of the subnormal values required by this fixture. Backend evidence: WebGPU/Metal flushes subnormals to zero in f32; the ~3e-39 subnormal dot products collapse to zero before alpha scaling (subgroup-matrix path)."
+      },
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
+        "test": "FusedMatMulOpTest.FloatTypeScale",
+        "notes": "Alpha scaling is applied after accumulation, so finite subnormal products remain valid nonzero outputs."
+      },
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [32, 32], "data": { "kind": "constant", "value": 1e-20 } },
+        "B": { "dtype": "float32", "shape": [32, 64], "data": { "kind": "constant", "value": 1e-20 } }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [32, 64], "tolerance": 1e-43 } }
+    },
+    {
+      "name": "f32_subgroup_matrix_transB_subnormal_dot_products_gpu_gap",
+      "skipGpu": {
+        "category": "permanent",
+        "reason": "Portable WGSL floating-point semantics do not guarantee preservation of the subnormal values required by this fixture. Backend evidence: WebGPU/Metal flushes subnormals to zero in f32; the ~3e-39 subnormal dot products collapse to zero (subgroup-matrix transB path)."
+      },
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
+        "test": "FusedMatMulOpTest.FloatTypeTransposeB",
+        "notes": "The transposed-B subgroup-matrix path has the same finite subnormal accumulation requirement."
+      },
+      "attrs": { "transB": 1 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [32, 32], "data": { "kind": "constant", "value": 1e-20 } },
+        "B": { "dtype": "float32", "shape": [64, 32], "data": { "kind": "constant", "value": 1e-20 } }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [32, 64], "tolerance": 1e-43 } }
+    },
+    {
+      "name": "aligned_f16_plain_64x32x64",
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.027, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [64, 64], "tolerance": 0.002 } }
+    },
+    {
+      "name": "aligned_f16_transB_alpha_64x32",
+      "attrs": { "transB": 1, "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.023, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.009, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [64, 64], "tolerance": 0.02 } }
+    },
+    {
+      "name": "f16_unaligned_3x5x7",
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [3, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.015, "cosStep": 0.021, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [5, 7],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.011, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [3, 7], "tolerance": 0.003 } }
+    },
+    {
+      "name": "aligned_f16_transA_64x32",
+      "attrs": { "transA": 1 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.027, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [64, 64], "tolerance": 0.001 } }
+    },
+    {
+      "name": "aligned_f16_batched_plain_2x64x32x64",
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [2, 64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.027, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [2, 32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 64, 64], "tolerance": 0.002 } }
+    },
+    {
+      "name": "f16_rank3_by_broadcast_rank3",
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [2, 2, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.027, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [1, 3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 2, 4], "tolerance": 0.002 } }
+    },
+    {
+      "name": "aligned_f16_transA_transB_alpha_64x32",
+      "attrs": { "transA": 1, "transB": 1, "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.023, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.009, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [64, 64], "tolerance": 0.002 } }
+    },
+    {
+      "name": "subgroup_matrix_m_tail_57_partial_block_f16",
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [57, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.023, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.011, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [57, 64], "tolerance": 0.001 } }
+    },
+    {
+      "name": "subgroup_matrix_m_tail_33_alpha_scaled_f32",
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [33, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.015, "cosStep": 0.021, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.019, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [33, 64], "tolerance": 0.0002 } }
+    },
+    {
+      "name": "empty_n_dimension_zero_width_output",
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [3, 4],
+          "data": { "kind": "values", "values": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0] }
+        },
+        "B": { "dtype": "float32", "shape": [4, 0], "data": { "kind": "values", "values": [] } }
+      },
+      "outputs": {
+        "Y": { "dtype": "float32", "shape": [3, 0], "tolerance": 0.000001, "data": { "kind": "values", "values": [] } }
+      }
+    },
+    {
+      "name": "transA_transB_subgroup_matrix_m_tail_50_f16",
+      "attrs": { "transA": 1, "transB": 1, "alpha": 1 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [32, 50],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.023, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.009, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [50, 64], "tolerance": 0.002 } }
+    },
+    {
+      "name": "f32_decode_gemv_m1_k65_n68_vec4_compact",
+      "provenance": {
+        "notes": "A compact M=1 float32 matrix product with odd K and N=68 checks the reduction and final output-column tail."
+      },
+      "attrs": { "alpha": 1 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [1, 65],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [65, 68],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [1, 68], "tolerance": 0.0002 } }
+    },
+    {
+      "name": "f16_decode_gemv_m1_k65_n68_vec4_compact",
+      "provenance": {
+        "notes": "A float16 M=1 matrix product with odd K and N=68 checks both reduction and output-column tails. The dot product accumulates in float32 and narrows only at the store."
+      },
+      "attrs": { "alpha": 1 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [1, 65],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [65, 68],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [1, 68], "tolerance": 0.004, "relTolerance": 0.004 } }
+    },
+    {
+      "name": "f16_decode_gemv_m1_k64_n128_alpha_half",
+      "provenance": {
+        "notes": "A float16 M=1 matrix product with even K and non-unit alpha checks complete reduction. The dot product accumulates in float32 and narrows only at the store."
+      },
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [64, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [1, 128], "tolerance": 0.0002, "relTolerance": 0.004 } }
+    },
+    {
+      "name": "f32_rank4_by_rank2_shared_weight_compact",
+      "provenance": {
+        "notes": "A rank-2 weight shared across a 2x3 batch multiplies M=5, K=7, N=9 (alpha=0.5); the odd dimensions check batch-offset indexing and non-power-of-two tails."
+      },
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 3, 5, 7],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [7, 9],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 3, 5, 9], "tolerance": 0.0002 } }
+    },
+    {
+      "name": "subgroup_matrix_kn_tail_f16_compact",
+      "provenance": {
+        "notes": "M=33 and K=34 are one and two past a 32-element boundary and N=66 is two past a 64-element boundary (float16, alpha=0.5), checking partial tiles in all three dimensions."
+      },
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [33, 34],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.1 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [34, 66],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.1 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [33, 66], "tolerance": 0.0004 } }
+    },
+    {
+      "name": "subgroup_matrix_broadcast_rank4x3_f16_compact",
+      "provenance": {
+        "notes": "A rank-4 [1,2,33,32] by rank-3 [2,32,64] float16 broadcast multiplies M=33, K=32, N=64 across a batch of 2, checking batch-dimension broadcasting at a compact scale."
+      },
+      "attrs": { "alpha": 1 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [1, 2, 33, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.1 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [2, 32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.1 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [1, 2, 33, 64], "tolerance": 0.0005 } }
+    },
+    {
+      "name": "broadcast_rank4_tiled_reg_f16_compact",
+      "provenance": {
+        "notes": "Compact rank-4 by rank-3 float16 broadcast with odd M/K/N checks output and reduction tails."
+      },
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [1, 2, 65, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.1 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [2, 33, 67],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.1 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [1, 2, 65, 67], "tolerance": 0.0004 } }
+    },
+    {
+      "name": "broadcast_rank4_tiled_reg_shared_f32_compact",
+      "provenance": {
+        "notes": "Float32 shared rank-2 weight counterpart for the register-blocked rank-4 path. Odd M/K/N cover all output and reduction tails."
+      },
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [1, 2, 65, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.1 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [33, 67],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.1 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [1, 2, 65, 67], "tolerance": 0.0003 } }
+    },
+    {
+      "name": "rank5_three_batch_dims",
+      "attrs": { "alpha": 1, "transA": 0, "transB": 0 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 1, 2, 2, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [1, 3, 1, 3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.25 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 3, 2, 2, 4], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "subgroup_matrix_kn_tail_f16_offset_alpha_scale",
+      "provenance": {
+        "notes": "Offset float16 operands keep each output near `alpha * K * aOffset * bOffset` (about 3.4), making the K=34 reduction tail, N=66 column tail, and `alpha = 0.5` epilogue observable on the subgroup-matrix route."
+      },
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [33, 34],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.1, "offset": 0.5 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [34, 66],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.1, "offset": 0.4 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [33, 66], "tolerance": 0.03, "relTolerance": 0.01 } }
+    },
+    {
+      "name": "subgroup_matrix_broadcast_rank4x3_f16_offset_scale",
+      "provenance": {
+        "notes": "Offset operands keep outputs near `K * aOffset * bOffset`, making the per-batch B slice and K=32 contraction observable in a rank-4 by rank-3 broadcast."
+      },
+      "attrs": { "alpha": 1 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [1, 2, 33, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.1, "offset": 0.5 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [2, 32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.1, "offset": 0.4 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [1, 2, 33, 64], "tolerance": 0.03, "relTolerance": 0.01 } }
+    },
+    {
+      "name": "subgroup_matrix_a_batch_broadcast_rank4x3_f16",
+      "provenance": {
+        "notes": "A has batch extent 1 while B has extent 2, so one A slice feeds both output batches. Offset operands keep the expected magnitude nonzero, exposing a swapped or nonzero A batch stride."
+      },
+      "attrs": { "alpha": 1 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [1, 1, 33, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.019, "scale": 0.1, "offset": 0.5 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [2, 32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.007, "cosStep": 0.031, "scale": 0.1, "offset": 0.4 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [1, 2, 33, 64], "tolerance": 0.03, "relTolerance": 0.01 } }
+    },
+    {
+      "name": "broadcast_rank4_tiled_reg_f16_offset_alpha_scale",
+      "provenance": {
+        "notes": "Offset operands keep outputs near `alpha * K * aOffset * bOffset` with K=33, making the one-element reduction tail and `alpha` multiplier observable on the register-blocked rank-4 route."
+      },
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [1, 2, 65, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.1, "offset": 0.5 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [2, 33, 67],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.1, "offset": 0.4 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [1, 2, 65, 67], "tolerance": 0.03, "relTolerance": 0.01 } }
+    },
+    {
+      "name": "aligned_f16_transA_transB_alpha_offset_scale",
+      "provenance": {
+        "notes": "Offset operands keep the doubly transposed output proportional to `alpha * K`, making the K=32 contraction and `alpha = 0.5` epilogue observable on subgroup-matrix and portable tiled routes."
+      },
+      "attrs": { "transA": 1, "transB": 1, "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.023, "scale": 0.2, "offset": 0.5 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.009, "scale": 0.2, "offset": 0.4 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [64, 64], "tolerance": 0.03, "relTolerance": 0.01 } }
+    },
+    {
+      "name": "f16_unaligned_3x5x7_offset_scale",
+      "provenance": {
+        "notes": "Offset operands in a 3x5 by 5x7 multiply keep outputs proportional to K, exposing dropped reduction elements or doubled tails on the unaligned scalar and tiled routes."
+      },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [3, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.015, "cosStep": 0.021, "scale": 0.2, "offset": 0.6 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [5, 7],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.011, "scale": 0.2, "offset": 0.5 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [3, 7], "tolerance": 0.01, "relTolerance": 0.005 } }
+    },
+    {
+      "name": "aligned_f16_plain_64x32x64_offset_scale",
+      "provenance": {
+        "notes": "Offset operands keep each fully aligned M=64, K=32, N=64 output proportional to K, making the subgroup-matrix reduction count and scratch drain observable."
+      },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.027, "scale": 0.2, "offset": 0.5 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.2, "offset": 0.4 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [64, 64], "tolerance": 0.03, "relTolerance": 0.01 } }
+    },
+    {
+      "name": "aligned_f16_batched_plain_2x64x32x64_offset_scale",
+      "provenance": {
+        "notes": "Two batches carry distinct offset operands, keeping outputs proportional to K and making both the batch stride and aligned subgroup-matrix reduction count observable."
+      },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [2, 64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.027, "scale": 0.2, "offset": 0.5 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [2, 32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.2, "offset": 0.4 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 64, 64], "tolerance": 0.03, "relTolerance": 0.01 } }
+    },
+    {
+      "name": "subgroup_matrix_m_tail_57_partial_block_f16_offset_scale",
+      "provenance": {
+        "notes": "M=57 leaves 25 rows after one full 32-row tile. Offset operands require tail rows to match the full rows' expected magnitude, exposing a short reduction or stale scratch value."
+      },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [57, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.023, "scale": 0.2, "offset": 0.5 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.011, "scale": 0.2, "offset": 0.4 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [57, 64], "tolerance": 0.03, "relTolerance": 0.01 } }
+    },
+    {
+      "name": "aligned_f16_transA_64x32_offset_scale",
+      "provenance": {
+        "notes": "With only A transposed, offset operands keep each output proportional to K and expose both a transposed-A stride error and an incorrect reduction count."
+      },
+      "attrs": { "transA": 1 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.027, "scale": 0.2, "offset": 0.5 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.2, "offset": 0.4 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [64, 64], "tolerance": 0.03, "relTolerance": 0.01 } }
+    },
+    {
+      "name": "transA_transB_subgroup_matrix_m_tail_50_f16_offset_scale",
+      "provenance": {
+        "notes": "Both operands are transposed and M=50 leaves an 18-row tail. Offset operands make the guarded tail rows' magnitude and placement independently observable."
+      },
+      "attrs": { "transA": 1, "transB": 1, "alpha": 1 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [32, 50],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.023, "scale": 0.2, "offset": 0.5 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.009, "scale": 0.2, "offset": 0.4 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [50, 64], "tolerance": 0.03, "relTolerance": 0.01 } }
+    },
+    {
+      "name": "f16_rank3_by_broadcast_rank3_offset_scale",
+      "provenance": {
+        "notes": "f16_rank3_by_broadcast_rank3 cancels to 0.076 under a 0.02 absolute tolerance (26% blind). Offsetting both operands makes each element ~K * aOffset * bOffset over K=3, so the shared single-batch B - read by both output batches - is pinned for value as well as for broadcast addressing."
+      },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [2, 2, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.027, "scale": 0.2, "offset": 1.0 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [1, 3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.2, "offset": 0.8 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 2, 4], "tolerance": 0.01, "relTolerance": 0.005 } }
+    },
+    {
+      "name": "rank6_four_batch_dims",
+      "attrs": { "alpha": 1, "transA": 0, "transB": 0 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 2, 1, 2, 2, 3],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [1, 1, 3, 1, 3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.25 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 2, 3, 2, 2, 4], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "subgroup_matrix_band_m8_f16",
+      "attrs": { "alpha": 2 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [8, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [64, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.029 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [8, 64], "tolerance": 0.005 } }
+    },
+    {
+      "name": "subgroup_matrix_splitk_m_tail_f16",
+      "attrs": { "alpha": 2 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [10, 2048],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [2048, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.031 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [10, 256], "tolerance": 0.005 } }
+    },
+    {
+      "name": "subgroup_matrix_splitk_alpha_scaled_f32",
+      "attrs": { "alpha": 1.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [16, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.03, "cosStep": 0.07, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [1024, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.041, "cosStep": 0.089, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [16, 128], "tolerance": 0.0002 } }
+    },
+    {
+      "name": "band_vec4_alpha_scaled_m8_k256_n512",
+      "provenance": {
+        "notes": "Eight rows, 256 reduction elements, and 512 output columns with alpha=0.5 verify that the non-unit scale factor is applied correctly to the matrix product."
+      },
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [8, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.13, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [256, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.05, "cosStep": 0.19, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [8, 512], "tolerance": 0.0001 } }
+    },
+    {
+      "name": "subgroup_matrix_batched_transB_small_m_f16",
+      "attrs": { "transB": 1, "alpha": 0.25 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [2, 4, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.023, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [2, 64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.009, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 4, 64], "tolerance": 0.01 } }
+    },
+    {
+      "name": "subgroup_matrix_transA_small_m_f16",
+      "attrs": { "transA": 1, "alpha": 3 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [64, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [64, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.029 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [8, 64], "tolerance": 0.005 } }
+    },
+    {
+      "name": "f32_decode_gemv_m1_k65_n68_alpha_half",
+      "provenance": {
+        "notes": "A compact M=1 GEMV with alpha=0.5 exercises a non-unit multiplier folded into the store as a baked constant. The expected output verifies that the multiplier is applied exactly once."
+      },
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [1, 65],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [65, 68],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [1, 68], "tolerance": 0.0002 } }
+    },
+    {
+      "name": "f16_rank4_by_rank2_shared_weight_m33_k34_n66",
+      "provenance": {
+        "notes": "A [2, 3, 33, 34] float16 A shares one [34, 66] B across both batch axes, with alpha 0.5; none of M=33, K=34 or N=66 is a multiple of 32."
+      },
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [2, 3, 33, 34],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [34, 66],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 3, 33, 66], "tolerance": 0.001 } }
+    },
+    {
+      "name": "f16_rank4_by_rank2_shared_weight_m8_k32_n64",
+      "provenance": {
+        "notes": "A [2, 3, 8, 32] float16 A shares one [32, 64] B across both batch axes, with alpha 0.5, at an 8-row M."
+      },
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [2, 3, 8, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 3, 8, 64], "tolerance": 0.001 } }
+    },
+    {
+      "name": "f16_rank4_by_rank2_shared_weight_m1_k32_n64",
+      "provenance": {
+        "notes": "A [2, 3, 1, 32] float16 A shares one [32, 64] B across both batch axes, with alpha 0.5, at a single-row M."
+      },
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [2, 3, 1, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 3, 1, 64], "tolerance": 0.001 } }
+    },
+    {
+      "name": "f32_band_preferred_m4_k2048_n4096",
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [4, 2048],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.1, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [2048, 4096],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.041, "scale": 0.1, "offset": 0.03 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [4, 4096], "tolerance": 0.0001, "relTolerance": 0.00001 } },
+      "provenance": {
+        "notes": "M=4, K=2048, N=4096 float32 operands with alpha=0.5 and different per-operand offsets (0.02 vs 0.03) avoid cancellation in the matrix product."
+      }
+    },
+    {
+      "name": "f32_band_preferred_m16_k2560_n4096",
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [16, 2560],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.1, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [2560, 4096],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.041, "scale": 0.1, "offset": 0.03 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [16, 4096], "tolerance": 0.0001, "relTolerance": 0.00001 } },
+      "provenance": {
+        "notes": "M=16, K=2560, N=4096 float32 operands with alpha=0.5 and different per-operand offsets (0.02 vs 0.03) avoid cancellation in the matrix product."
+      }
+    },
+    {
+      "name": "broadcast-transb-tails-float16",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [2, 1, 129, 65],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.075, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [3, 129, 65],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.017, "scale": 0.075, "offset": -0.02 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 3, 129, 129], "tolerance": 0.001, "relTolerance": 0 } },
+      "provenance": { "notes": "Broadcast transposed-B geometry, alpha, padding and output-grid boundary." }
+    },
+    {
+      "name": "broadcast-transb-rank3x2-float16",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [4, 256, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.075, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [512, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.017, "scale": 0.075, "offset": -0.02 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [4, 256, 512], "tolerance": 0.001, "relTolerance": 0 } },
+      "provenance": { "notes": "Broadcast transposed-B geometry, alpha, padding and output-grid boundary." }
+    },
+    {
+      "name": "broadcast-transb-rank5x3-float16",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [2, 1, 2, 128, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.075, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [2, 128, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.017, "scale": 0.075, "offset": -0.02 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [2, 1, 2, 128, 128], "tolerance": 0.001, "relTolerance": 0 } },
+      "provenance": { "notes": "Broadcast transposed-B geometry, alpha, padding and output-grid boundary." }
+    },
+    {
+      "name": "broadcast-transb-low_tiles-float16",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [1, 64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.075, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.017, "scale": 0.075, "offset": -0.02 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [1, 64, 64], "tolerance": 0.001, "relTolerance": 0 } },
+      "provenance": { "notes": "Broadcast transposed-B geometry, alpha, padding and output-grid boundary." }
+    },
+    {
+      "name": "broadcast-transb-large-float32",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 8, 512, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.075, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [8, 512, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.017, "scale": 0.075, "offset": -0.02 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 8, 512, 512], "tolerance": 0.00001, "relTolerance": 0 } },
+      "provenance": { "notes": "Broadcast transposed-B geometry, alpha, padding and output-grid boundary." }
+    },
+    {
+      "name": "broadcast-transb-tails-float32",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 1, 129, 65],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.075, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [3, 129, 65],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.017, "scale": 0.075, "offset": -0.02 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 3, 129, 129], "tolerance": 0.00001, "relTolerance": 0 } },
+      "provenance": { "notes": "Broadcast transposed-B geometry, alpha, padding and output-grid boundary." }
+    },
+    {
+      "name": "broadcast-transb-rank3x2-float32",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [4, 256, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.075, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [512, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.017, "scale": 0.075, "offset": -0.02 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [4, 256, 512], "tolerance": 0.00001, "relTolerance": 0 } },
+      "provenance": { "notes": "Broadcast transposed-B geometry, alpha, padding and output-grid boundary." }
+    },
+    {
+      "name": "broadcast-transb-rank5x3-float32",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 1, 2, 128, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.075, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [2, 128, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.017, "scale": 0.075, "offset": -0.02 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 1, 2, 128, 128], "tolerance": 0.00001, "relTolerance": 0 } },
+      "provenance": { "notes": "Broadcast transposed-B geometry, alpha, padding and output-grid boundary." }
+    },
+    {
+      "name": "broadcast-transb-low_tiles-float32",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [1, 64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.075, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [64, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.017, "scale": 0.075, "offset": -0.02 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [1, 64, 64], "tolerance": 0.00001, "relTolerance": 0 } },
+      "provenance": { "notes": "Broadcast transposed-B geometry, alpha, padding and output-grid boundary." }
+    },
+    {
+      "name": "broadcast-grid-b3-m128-k64-n256-float16-a0.125",
+      "attrs": { "transB": 1, "alpha": 0.125 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [3, 128, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.075, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [256, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.017, "scale": 0.075, "offset": -0.02 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [3, 128, 256], "tolerance": 0.001, "relTolerance": 0 } },
+      "provenance": { "notes": "Broadcast transposed-B geometry, alpha, padding and output-grid boundary." }
+    },
+    {
+      "name": "broadcast-grid-b4-m128-k64-n256-float16-a-0.375",
+      "attrs": { "transB": 1, "alpha": -0.375 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [4, 128, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.075, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [256, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.017, "scale": 0.075, "offset": -0.02 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [4, 128, 256], "tolerance": 0.001, "relTolerance": 0 } },
+      "provenance": { "notes": "Broadcast transposed-B geometry, alpha, padding and output-grid boundary." }
+    },
+    {
+      "name": "broadcast-grid-b8-m128-k64-n256-float16-a0.5",
+      "attrs": { "transB": 1, "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [8, 128, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.075, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [256, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.017, "scale": 0.075, "offset": -0.02 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [8, 128, 256], "tolerance": 0.001, "relTolerance": 0 } },
+      "provenance": { "notes": "Broadcast transposed-B geometry, alpha, padding and output-grid boundary." }
+    },
+    {
+      "name": "broadcast-grid-b8-m129-k65-n257-float16-a0.125",
+      "attrs": { "transB": 1, "alpha": 0.125 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [8, 129, 65],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.075, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [257, 65],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.017, "scale": 0.075, "offset": -0.02 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [8, 129, 257], "tolerance": 0.001, "relTolerance": 0 } },
+      "provenance": { "notes": "Broadcast transposed-B geometry, alpha, padding and output-grid boundary." }
+    },
+    {
+      "name": "broadcast-grid-b8-m128-k96-n128-float16-a0",
+      "attrs": { "transB": 1, "alpha": 0 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [8, 128, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.075, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [128, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.017, "scale": 0.075, "offset": -0.02 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [8, 128, 128], "tolerance": 0.001, "relTolerance": 0 } },
+      "provenance": { "notes": "Broadcast transposed-B geometry, alpha, padding and output-grid boundary." }
+    },
+    {
+      "name": "broadcast-grid-b4-m256-k127-n512-float16-a-0.5",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [4, 256, 127],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.075, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [512, 127],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.017, "scale": 0.075, "offset": -0.02 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [4, 256, 512], "tolerance": 0.001, "relTolerance": 0 } },
+      "provenance": { "notes": "Broadcast transposed-B geometry, alpha, padding and output-grid boundary." }
+    },
+    {
+      "name": "broadcast-grid-b3-m128-k64-n256-float32-a0.125",
+      "attrs": { "transB": 1, "alpha": 0.125 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [3, 128, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.075, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [256, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.017, "scale": 0.075, "offset": -0.02 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [3, 128, 256], "tolerance": 0.00001, "relTolerance": 0 } },
+      "provenance": { "notes": "Broadcast transposed-B geometry, alpha, padding and output-grid boundary." }
+    },
+    {
+      "name": "broadcast-grid-b4-m128-k64-n256-float32-a-0.375",
+      "attrs": { "transB": 1, "alpha": -0.375 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [4, 128, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.075, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [256, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.017, "scale": 0.075, "offset": -0.02 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [4, 128, 256], "tolerance": 0.00001, "relTolerance": 0 } },
+      "provenance": { "notes": "Broadcast transposed-B geometry, alpha, padding and output-grid boundary." }
+    },
+    {
+      "name": "broadcast-grid-b8-m128-k64-n256-float32-a0.5",
+      "attrs": { "transB": 1, "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [8, 128, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.075, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [256, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.017, "scale": 0.075, "offset": -0.02 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [8, 128, 256], "tolerance": 0.00001, "relTolerance": 0 } },
+      "provenance": { "notes": "Broadcast transposed-B geometry, alpha, padding and output-grid boundary." }
+    },
+    {
+      "name": "broadcast-grid-b8-m129-k65-n257-float32-a0.125",
+      "attrs": { "transB": 1, "alpha": 0.125 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [8, 129, 65],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.075, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [257, 65],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.017, "scale": 0.075, "offset": -0.02 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [8, 129, 257], "tolerance": 0.00001, "relTolerance": 0 } },
+      "provenance": { "notes": "Broadcast transposed-B geometry, alpha, padding and output-grid boundary." }
+    },
+    {
+      "name": "broadcast-grid-b8-m128-k96-n128-float32-a0",
+      "attrs": { "transB": 1, "alpha": 0 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [8, 128, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.075, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [128, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.017, "scale": 0.075, "offset": -0.02 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [8, 128, 128], "tolerance": 0.00001, "relTolerance": 0 } },
+      "provenance": { "notes": "Broadcast transposed-B geometry, alpha, padding and output-grid boundary." }
+    },
+    {
+      "name": "broadcast-grid-b4-m256-k127-n512-float32-a-0.5",
+      "attrs": { "transB": 1, "alpha": -0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [4, 256, 127],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.075, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [512, 127],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.017, "scale": 0.075, "offset": -0.02 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [4, 256, 512], "tolerance": 0.00001, "relTolerance": 0 } },
+      "provenance": { "notes": "Broadcast transposed-B geometry, alpha, padding and output-grid boundary." }
+    },
+    {
+      "name": "broadcast-selected-mn-tail-float16",
+      "attrs": { "transB": 1, "alpha": 0.125 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [8, 129, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.075, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [257, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.017, "scale": 0.075, "offset": -0.02 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [8, 129, 257], "tolerance": 0.001, "relTolerance": 0 } },
+      "provenance": { "notes": "Broadcast transposed-B geometry, alpha, padding and output-grid boundary." }
+    },
+    {
+      "name": "broadcast-selected-mn-tail-float32",
+      "attrs": { "transB": 1, "alpha": 0.125 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [8, 129, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.075, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [257, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.017, "scale": 0.075, "offset": -0.02 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [8, 129, 257], "tolerance": 0.00001, "relTolerance": 0 } },
+      "provenance": { "notes": "Broadcast transposed-B geometry, alpha, padding and output-grid boundary." }
+    },
+    {
+      "name": "broadcast-selected-broadcast-tail-float16",
+      "attrs": { "transB": 1, "alpha": 0.125 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [4, 1, 129, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.075, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float16",
+          "shape": [3, 257, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.017, "scale": 0.075, "offset": -0.02 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float16", "shape": [4, 3, 129, 257], "tolerance": 0.001, "relTolerance": 0 } },
+      "provenance": { "notes": "Broadcast transposed-B geometry, alpha, padding and output-grid boundary." }
+    },
+    {
+      "name": "broadcast-selected-broadcast-tail-float32",
+      "attrs": { "transB": 1, "alpha": 0.125 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [4, 1, 129, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.075, "offset": 0.02 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [3, 257, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.017, "scale": 0.075, "offset": -0.02 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [4, 3, 129, 257], "tolerance": 0.00001, "relTolerance": 0 } },
+      "provenance": { "notes": "Broadcast transposed-B geometry, alpha, padding and output-grid boundary." }
+    },
+    {
+      "name": "plain_rank2_tiled_reg_row_tail_m100_k64_n2048",
+      "provenance": {
+        "notes": "M=100 leaves 36 rows past a 64-row boundary (K=64, N=2048), with alpha=0.5 scaling the output; the extra rows must be included correctly in the result."
+      },
+      "attrs": { "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [100, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.037, "cosStep": 0.061, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [64, 2048],
+          "data": { "kind": "fillFloat32", "sinStep": 0.043, "cosStep": 0.079, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [100, 2048], "tolerance": 0.0001 } }
+    }
+  ]
+}