sync 2e7068faf55e

Browse files

Files changed (10) hide show

README.md +77 -0
build/webgpu/bench.json +217 -0
build/webgpu/fused-matmul-subgroup-matrix.wgsl.jinja +177 -0
build/webgpu/manifest.json +286 -0
build/webgpu/matmul-subgroup-matrix-ext.wgsl.jinja +365 -0
build/webgpu/matmul-tiled-general-reg.wgsl.jinja +176 -0
build/webgpu/matmul-tiled-general.wgsl.jinja +176 -0
build/webgpu/matmul-vector-matrix-vec4.wgsl.jinja +42 -0
build/webgpu/metadata.json +22 -0
build/webgpu/test.json +1420 -0

README.md CHANGED Viewed

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

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# com.microsoft.FusedMatMul
+`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. Supports optional transposition of either operand's last two dimensions, optional batch-dimension transposition, and a scalar `alpha` multiplier. Float32 and float16 are supported; double and bfloat16 are not.
+See the [ONNX Runtime `FusedMatMul` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.FusedMatMul) for the reference semantics.
+## Inputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `A` | `A` | `T` | — | — | N-dimensional matrix A. | required |
+| `B` | `B` | `T` | — | — | N-dimensional matrix B. | required |
+## Outputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `Y` | `Y` | `T` | derived | derived; see description | 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. |
+| `transBatchA` | `0` | When non-zero, transposes `A` on its first dimension and batch dimensions (dim-1 to dim-rank-2) before multiplication. |
+| `transBatchB` | `0` | When non-zero, transposes `B` on its first dimension and batch dimensions (dim-1 to dim-rank-2) before multiplication. |
+## Type constraints
+| Variable | Allowed dtypes |
+| --- | --- |
+| `T` | `float32`, `float16` |
+## 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, provenance)
+- [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
+- [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
+- [`fused-matmul-subgroup-matrix.wgsl.jinja`](build/webgpu/fused-matmul-subgroup-matrix.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)
+## Use with `@huggingface/kernels`
+The loader derives every required output's shape and logical dtype from the manifest contract and this call.
+It then allocates the result tensors automatically.
+The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model 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.FusedMatMul", { 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,217 @@

+{
+  "op": "com.microsoft.FusedMatMul",
+  "cases": [
+    {
+      "name": "fusedmatmul-f32-attn-scores-transB-8x512x64",
+      "preset": "model",
+      "attrs": { "alpha": 0.125, "transB": 1 },
+      "inputs": {
+        "A": { "shape": [8, 512, 64], "dtype": "float32", "dist": "normal", "seed": 510, "scale": 0.2 },
+        "B": { "shape": [8, 512, 64], "dtype": "float32", "dist": "normal", "seed": 511, "scale": 0.2 }
+      },
+      "outputs": { "Y": { "shape": [8, 512, 512], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * 8 * 512 * 512 * 64" }] }
+    },
+    {
+      "name": "fusedmatmul-f32-plain-512x2048x512",
+      "preset": "model",
+      "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": "fusedmatmul-f16-plain-512x2048x512",
+      "preset": "model",
+      "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": "fusedmatmul-f32-unaligned-500x2000x500",
+      "preset": "model",
+      "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": "fusedmatmul-f32-unaligned2-384x1536x400",
+      "preset": "model",
+      "attrs": { "alpha": 1 },
+      "inputs": {
+        "A": { "shape": [384, 1536], "dtype": "float32", "dist": "normal", "seed": 560, "scale": 0.1 },
+        "B": { "shape": [1536, 400], "dtype": "float32", "dist": "normal", "seed": 561, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "shape": [384, 400], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * 384 * 400 * 1536" }] }
+    },
+    {
+      "name": "fusedmatmul-f16-aligned-n512-512x2048x512-healthy",
+      "preset": "smoke",
+      "attrs": { "alpha": 1 },
+      "vars": { "M": 512, "K": 2048, "N": 512 },
+      "inputs": {
+        "A": { "shape": [512, 2048], "dtype": "float16", "dist": "normal", "seed": 600, "scale": 0.1 },
+        "B": { "shape": [2048, 512], "dtype": "float16", "dist": "normal", "seed": 601, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "shape": [512, 512], "dtype": "float16" } },
+      "bench": { "primary": true, "metrics": [{ "type": "gflops", "value": "2 * 512 * 512 * 2048" }] }
+    },
+    {
+      "name": "fusedmatmul-f16-unaligned-n500-512x2048x500-cliff",
+      "preset": "smoke",
+      "attrs": { "alpha": 1 },
+      "vars": { "M": 512, "K": 2048, "N": 500 },
+      "inputs": {
+        "A": { "shape": [512, 2048], "dtype": "float16", "dist": "normal", "seed": 602, "scale": 0.1 },
+        "B": { "shape": [2048, 500], "dtype": "float16", "dist": "normal", "seed": 603, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "shape": [512, 500], "dtype": "float16" } },
+      "bench": { "primary": true, "metrics": [{ "type": "gflops", "value": "2 * 512 * 500 * 2048" }] }
+    },
+    {
+      "name": "fusedmatmul-f16-unaligned-k2050-512x2050x512-cliff",
+      "preset": "smoke",
+      "attrs": { "alpha": 1 },
+      "vars": { "M": 512, "K": 2050, "N": 512 },
+      "inputs": {
+        "A": { "shape": [512, 2050], "dtype": "float16", "dist": "normal", "seed": 604, "scale": 0.1 },
+        "B": { "shape": [2050, 512], "dtype": "float16", "dist": "normal", "seed": 605, "scale": 0.1 }
+      },
+      "outputs": { "Y": { "shape": [512, 512], "dtype": "float16" } },
+      "bench": { "primary": true, "metrics": [{ "type": "gflops", "value": "2 * 512 * 512 * 2050" }] }
+    },
+    {
+      "name": "fusedmatmul-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": "fusedmatmul-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": "fusedmatmul-f16-broadcast-batch-rank4x3-2x8x512x64-tiled",
+      "preset": "smoke",
+      "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": "fusedmatmul-f32-decode-gemv-m1-k4096-n4096-pathology",
+      "preset": "stress",
+      "attrs": { "alpha": 1 },
+      "vars": { "dtype": "float32", "M": 1, "K": 4096, "N": 4096 },
+      "inputs": {
+        "A": { "shape": [1, 4096], "dtype": "float32", "dist": "normal", "seed": 740, "scale": 0.1 },
+        "B": { "shape": [4096, 4096], "dtype": "float32", "dist": "normal", "seed": 741, "scale": 0.05 }
+      },
+      "outputs": { "Y": { "shape": [1, 4096], "dtype": "float32", "dist": "empty" } },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "(args.M * args.K + args.K * args.N + args.M * args.N) * dtypeBytes(args.dtype)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fusedmatmul-f32-rank4-by-rank2-shared-weight-b2h8-m128-k64-n128-pathology",
+      "preset": "stress",
+      "attrs": { "alpha": 1 },
+      "vars": { "dtype": "float32", "M": 128, "K": 64, "N": 128 },
+      "inputs": {
+        "A": { "shape": [2, 8, 128, 64], "dtype": "float32", "dist": "normal", "seed": 742, "scale": 0.2 },
+        "B": { "shape": [64, 128], "dtype": "float32", "dist": "normal", "seed": 743, "scale": 0.2 }
+      },
+      "outputs": { "Y": { "shape": [2, 8, 128, 128], "dtype": "float32", "dist": "empty" } },
+      "bench": { "primary": true, "metrics": [{ "type": "gflops", "value": "2 * numel(shapes.Y) * args.K" }] }
+    },
+    {
+      "name": "fusedmatmul-f16-rank4-by-rank2-shared-weight-b2h8-m512-k2048-n512-pathology",
+      "preset": "stress",
+      "provenance": {
+        "notes": "Production-scale batched projection with a shared rank-2 weight. This common numpy.matmul broadcast shape is excluded from the subgroup-matrix rank combinations and locks the register-blocked portable fallback that replaced the lower-throughput generic tiled kernel."
+      },
+      "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": "fusedmatmul-f16-transbatch-a-forces-tiled-8x512x2048x512-stress",
+      "preset": "stress",
+      "attrs": { "alpha": 1, "transBatchA": 1 },
+      "vars": { "batch": 8, "M": 512, "K": 2048, "N": 512 },
+      "inputs": {
+        "A": { "shape": [512, 8, 2048], "dtype": "float16", "dist": "normal", "seed": 720, "scale": 0.05 },
+        "B": { "shape": [8, 2048, 512], "dtype": "float16", "dist": "normal", "seed": 721, "scale": 0.05 }
+      },
+      "outputs": { "Y": { "shape": [8, 512, 512], "dtype": "float16", "dist": "empty" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * 8 * 512 * 512 * 2048" }] }
+    },
+    {
+      "name": "fusedmatmul-f16-transbatch-b-8x512x2048x512-pathology",
+      "preset": "stress",
+      "provenance": {
+        "notes": "The operator's valid transBatchB semantic at a production projection size. No subgroup-matrix or specialized trans-batch-B variant accepts this layout, so it selects the lower-throughput generic tiled kernel."
+      },
+      "attrs": { "alpha": 1, "transBatchB": 1 },
+      "vars": { "batch": 8, "M": 512, "K": 2048, "N": 512 },
+      "inputs": {
+        "A": { "shape": [8, 512, 2048], "dtype": "float16", "dist": "normal", "seed": 722, "scale": 0.05 },
+        "B": { "shape": [2048, 8, 512], "dtype": "float16", "dist": "normal", "seed": 723, "scale": 0.05 }
+      },
+      "outputs": { "Y": { "shape": [8, 512, 512], "dtype": "float16", "dist": "empty" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * 8 * 512 * 512 * 2048" }] }
+    },
+    {
+      "name": "fusedmatmul-f16-broadcast-batch-forces-tiled-1x8x512x2048x512-stress",
+      "preset": "stress",
+      "attrs": { "alpha": 1 },
+      "inputs": {
+        "A": { "shape": [1, 8, 512, 2048], "dtype": "float16", "dist": "normal", "seed": 730, "scale": 0.05 },
+        "B": { "shape": [8, 2048, 512], "dtype": "float16", "dist": "normal", "seed": 731, "scale": 0.05 }
+      },
+      "outputs": { "Y": { "shape": [1, 8, 512, 512], "dtype": "float16", "dist": "empty" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * 8 * 512 * 512 * 2048" }] }
+    }
+  ]
+}

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

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

+// com.microsoft.FusedMatMul subgroup-matrix specialization: Y = alpha * A @ op(B),
+// op(B) = B^T when transB. Dense batches map through workgroup_id.z, and
+// M-tail rows are guarded by row_limit. transA uses the tiled specialization.
+// Alignment gates keep K % 32 == 0 and N % 64 == 0 so subgroupMatrixLoad never
+// sees partial 8x8 tiles. The batch is required to match between A and B
+// (no broadcast) because a_base/b_base both index by the same workgroup_id.z.
+{% if usesF16 %}
+enable f16;
+{% endif %}
+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 M: u32 = {{ M }}u;
+const K: u32 = {{ K }}u;
+const N: u32 = {{ N }}u;
+{% if transBatchA %}
+const BATCH_COUNT: u32 = {{ batchCount }}u;
+const A_BATCH_STRIDE: u32 = K;
+const A_M_STRIDE: u32 = BATCH_COUNT * K;
+{% else %}
+const A_BATCH_STRIDE: u32 = M * K;
+{% if not transA %}
+const A_M_STRIDE: u32 = K;
+{% endif %}
+{% endif %}
+const B_BATCH_STRIDE: u32 = K * N;
+const C_BATCH_STRIDE: u32 = M * N;
+const ALPHA: {{ accScalar }} = {{ accScalar }}({{ source.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 < 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 * 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 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) {
+      let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
+      var matA0: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>>(&tile_A, matrix_a_offset, false, TILE_K);
+      var matA1: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>>(&tile_A, matrix_a_offset + 8u * TILE_K, false, TILE_K);
+      let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
+      var matB0: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset, true, TILE_K);
+      var matB1: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 8u * TILE_K, true, TILE_K);
+      var matB2: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 16u * TILE_K, true, TILE_K);
+      var matB3: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 24u * TILE_K, true, TILE_K);
+      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(&scratch[subtile_id][0], 0u, matC00, false, 8u);
+  subgroupMatrixStore(&scratch[subtile_id][1], 0u, matC01, false, 8u);
+  subgroupMatrixStore(&scratch[subtile_id][2], 0u, matC02, false, 8u);
+  subgroupMatrixStore(&scratch[subtile_id][3], 0u, matC03, false, 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(&scratch[subtile_id][0], 0u, matC10, false, 8u);
+  subgroupMatrixStore(&scratch[subtile_id][1], 0u, matC11, false, 8u);
+  subgroupMatrixStore(&scratch[subtile_id][2], 0u, matC12, false, 8u);
+  subgroupMatrixStore(&scratch[subtile_id][3], 0u, matC13, false, 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,286 @@

+{
+  "domain": "com.microsoft",
+  "name": "FusedMatMul",
+  "sinceVersion": 1,
+  "description": "Matrix product of two N-dimensional tensors `A` and `B`, following NumPy-style matrix-multiplication broadcasting. Supports optional transposition of either operand's last two dimensions, optional batch-dimension transposition, and a scalar `alpha` multiplier. Float32 and float16 are supported; double and bfloat16 are not.",
+  "inputs": [
+    { "role": "A", "dtype": "T", "description": "N-dimensional matrix A." },
+    { "role": "B", "dtype": "T", "description": "N-dimensional matrix B." }
+  ],
+  "outputs": [
+    {
+      "role": "Y",
+      "dtype": "T",
+      "rank": "max(ranks.A, ranks.B) - (1 if ranks.A == 1 or ranks.B == 1 else 0)",
+      "shape": "matmulShape(logicalAShape, logicalBShape)",
+      "description": "Matrix-multiplication result whose shape follows NumPy-style rules after applying the requested batch and matrix transpositions."
+    }
+  ],
+  "attributes": { "alpha": 1, "transA": 0, "transB": 0, "transBatchA": 0, "transBatchB": 0 },
+  "attributeDescriptions": {
+    "alpha": "Scalar multiplier applied to the product of the input tensors.",
+    "transA": "When non-zero, transposes `A` on its last two dimensions before multiplication.",
+    "transB": "When non-zero, transposes `B` on its last two dimensions before multiplication.",
+    "transBatchA": "When non-zero, transposes `A` on its first dimension and batch dimensions (dim-1 to dim-rank-2) before multiplication.",
+    "transBatchB": "When non-zero, transposes `B` on its first dimension and batch dimensions (dim-1 to dim-rank-2) before multiplication."
+  },
+  "typeConstraints": { "T": ["float32", "float16"] },
+  "args": {
+    "A": { "kind": "tensor", "semantic": "A", "role": "input" },
+    "B": { "kind": "tensor", "semantic": "B", "role": "input" },
+    "Y": { "kind": "tensor", "semantic": "Y", "role": "output" }
+  },
+  "tunables": { "TILED_REG_MIN_WORKGROUPS": 64, "GEMV_TARGET_BLOCKS": 512 },
+  "bindingSets": {
+    "matrix": [
+      { "name": "a", "arg": "A", "semantic": "A", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "b", "arg": "B", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "y", "arg": "Y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" }
+    ],
+    "gemvF32": [
+      { "name": "a", "arg": "A", "semantic": "A", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      {
+        "name": "b",
+        "arg": "B",
+        "semantic": "B",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "vec4<f32>"
+      },
+      { "name": "c", "arg": "Y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "vec4<f32>" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "K", "type": "u32", "value": "dim(shapes.A, 1)" },
+            { "name": "N4", "type": "u32", "value": "dim(shapes.B, 1) / 4" }
+          ]
+        }
+      }
+    ]
+  },
+  "derive": {
+    "gemvWorkgroups": "ceilDiv(dim(shapes.B, 1), 128)",
+    "gemvSliceCap": "min(32, device.limits.maxComputeWorkgroupSizeY, floor(device.limits.maxComputeInvocationsPerWorkgroup / 32), floor(device.limits.maxComputeWorkgroupStorageSize / 512))",
+    "gemvSlices": "max(1, min(gemvSliceCap, max(8, pow2ceil(ceilDiv(tunables.GEMV_TARGET_BLOCKS, gemvWorkgroups)))))",
+    "batchMovedAShape": "moveAxis(shapes.A, 0, -2) if attrs.transBatchA != 0 else shapes.A",
+    "batchMovedBShape": "moveAxis(shapes.B, 0, -2) if attrs.transBatchB != 0 else 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",
+    "transBatchContract": "(attrs.transBatchA == 0 and attrs.transBatchB == 0) or (ranks.A == ranks.B and ranks.A >= 3)",
+    "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"
+  },
+  "variants": [
+    {
+      "id": "f32_m1_gemv_vec4",
+      "priority": 30,
+      "when": ["dtypes.T == \"f32\"", "attrs.alpha == 1", "attrs.transA == 0", "attrs.transB == 0", "attrs.transBatchA == 0", "attrs.transBatchB == 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", "ceil(dim(shapes.B, 1) / 128) <= device.limits.maxComputeWorkgroupsPerDimension"],
+      "constants": { "T": "\"f32\"", "usesF16": false, "gemvSlices": "gemvSlices" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "FusedMatMul.F32M1GemvVec4",
+          "shader": "matmul-vector-matrix-vec4.wgsl.jinja",
+          "bindings": "gemvF32",
+          "dispatch": { "x": "ceil(dim(shapes.B, 1) / 128)" }
+        }
+      ]
+    },
+    {
+      "id": "subgroup_matrix_tail_broadcast",
+      "priority": 11,
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }]
+      },
+      "when": ["dtypes.T == \"f16\"", "f16Ok(dtypes.T)", "attrs.transA == 0", "attrs.transB == 0", "attrs.transBatchA == 0", "attrs.transBatchB == 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))))", "dim(shapes.A, ranks.A - 1) == dim(shapes.B, ranks.B - 2)", "dim(shapes.A, ranks.A - 2) >= 32", "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) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil(dim(shapes.A, ranks.A - 2) / 32) <= device.limits.maxComputeWorkgroupsPerDimension", "numel(shapes.Y) / (dim(shapes.A, ranks.A - 2) * dim(shapes.B, ranks.B - 1)) <= device.limits.maxComputeWorkgroupsPerDimension", "wave32Effective"],
+      "constants": {
+        "hasBias": false,
+        "usesF16": true,
+        "fScalar": "\"f16\"",
+        "outScalar": "\"f16\"",
+        "scalar": "dtypes.T",
+        "generalAddressing": true,
+        "tailSafe": "dim(shapes.A, ranks.A - 1) % 32 != 0 or dim(shapes.B, ranks.B - 1) % 64 != 0",
+        "outputBuffer": "\"y\"",
+        "alpha": "attrs.alpha",
+        "M": "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))"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "FusedMatMul.SubgroupMatrixTailBroadcast",
+          "source": {
+            "shader": "matmul-subgroup-matrix-ext.wgsl.jinja",
+            "inputs": { "aShape": "shapes.A", "bShape": "shapes.B", "aRank": "ranks.A", "bRank": "ranks.B" }
+          },
+          "bindings": "matrix",
+          "dispatch": { "x": "ceil(constants.N / 64)", "y": "ceil(constants.M / 32)", "z": "constants.batchCount" }
+        }
+      ]
+    },
+    {
+      "id": "subgroup_matrix",
+      "priority": 10,
+      "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 }
+        ]
+      },
+      "when": ["f16Ok(dtypes.T)", "attrs.transBatchA == 0 or (attrs.transA == 0 and ranks.A == 3)", "attrs.transBatchB == 0", "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, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2))) >= 32", "(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, 1) if attrs.transBatchA != 0 else 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, 0) if attrs.transBatchA != 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) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil((dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2))) / 32) <= device.limits.maxComputeWorkgroupsPerDimension", "numel(shapes.Y) / ((dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 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))) <= device.limits.maxComputeWorkgroupsPerDimension", "wave32Effective"],
+      "constants": {
+        "usesF16": "dtypes.T == \"f16\"",
+        "fScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"",
+        "outScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"",
+        "scalar": "dtypes.T",
+        "transA": "attrs.transA != 0",
+        "transB": "attrs.transB != 0",
+        "transBatchA": "attrs.transBatchA != 0",
+        "M": "(dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 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, 0) if attrs.transBatchA != 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": "FusedMatMul.SubgroupMatrix",
+          "source": { "shader": "fused-matmul-subgroup-matrix.wgsl.jinja", "inputs": { "alpha": "attrs.alpha" } },
+          "bindings": "matrix",
+          "dispatch": {
+            "x": "ceil(constants.N / 64)",
+            "y": "ceil(constants.M / 32)",
+            "z": "numel(shapes.Y) / (constants.M * constants.N)"
+          }
+        }
+      ]
+    },
+    {
+      "id": "broadcast_rank4_tiled_reg",
+      "priority": 6,
+      "when": ["f16Ok(dtypes.T)", "attrs.transA == 0", "attrs.transB == 0", "attrs.transBatchA == 0", "attrs.transBatchB == 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) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil(dim(shapes.A, 2) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "numel(shapes.Y) / (dim(shapes.A, 2) * dim(shapes.B, ranks.B - 1)) <= device.limits.maxComputeWorkgroupsPerDimension"],
+      "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "FusedMatMul.BroadcastRank4TiledReg",
+          "source": {
+            "shader": "matmul-tiled-general-reg.wgsl.jinja",
+            "inputs": {
+              "aShape": "shapes.A",
+              "bShape": "shapes.B",
+              "alpha": "attrs.alpha",
+              "aRank": "ranks.A",
+              "bRank": "ranks.B",
+              "transBatchA": "false"
+            }
+          },
+          "bindings": "matrix",
+          "dispatch": {
+            "x": "ceil(dim(shapes.B, ranks.B - 1) / 64)",
+            "y": "ceil(dim(shapes.A, 2) / 64)",
+            "z": "numel(shapes.Y) / (dim(shapes.A, 2) * dim(shapes.B, ranks.B - 1))"
+          }
+        }
+      ]
+    },
+    {
+      "id": "plain_rank2_tiled_reg",
+      "description": "Register-blocked rank-2 Y = alpha * A @ B for the tier without subgroup matrices. It is the non-transposed rung the two specialised register-blocked variants beside it never covered, so plain rank-2 shapes fell to the scalar-indexed tile.",
+      "priority": 4,
+      "when": ["f16Ok(dtypes.T)", "attrs.transA == 0", "attrs.transB == 0", "attrs.transBatchA == 0", "attrs.transBatchB == 0", "ranks.A == 2", "ranks.B == 2", "ranks.Y == 2", "dim(shapes.A, 1) == dim(shapes.B, 0)", "dim(shapes.Y, 0) == dim(shapes.A, 0)", "dim(shapes.Y, 1) == dim(shapes.B, 1)", "dim(shapes.A, 0) >= 64", "dim(shapes.A, 1) >= 32", "dim(shapes.B, 1) >= 64", "ceil(dim(shapes.A, 0) / 64) * ceil(dim(shapes.B, 1) / 64) >= tunables.TILED_REG_MIN_WORKGROUPS", "ceil(dim(shapes.B, 1) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil(dim(shapes.A, 0) / 64) <= device.limits.maxComputeWorkgroupsPerDimension"],
+      "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "FusedMatMul.PlainRank2TiledReg",
+          "source": {
+            "shader": "matmul-tiled-general-reg.wgsl.jinja",
+            "inputs": {
+              "aShape": "shapes.A",
+              "bShape": "shapes.B",
+              "alpha": "attrs.alpha",
+              "aRank": "ranks.A",
+              "bRank": "ranks.B",
+              "transBatchA": "false"
+            }
+          },
+          "bindings": "matrix",
+          "dispatch": { "x": "ceil(dim(shapes.B, 1) / 64)", "y": "ceil(dim(shapes.A, 0) / 64)", "z": 1 }
+        }
+      ]
+    },
+    {
+      "id": "transbatch_a_tiled_reg",
+      "priority": 5,
+      "when": ["f16Ok(dtypes.T)", "attrs.transBatchA != 0", "attrs.transBatchB == 0", "attrs.transA == 0", "attrs.transB == 0", "ranks.A == 3", "ranks.B == 3", "ranks.Y == 3", "dim(shapes.A, 1) == dim(shapes.B, 0)", "dim(shapes.Y, 0) == dim(shapes.B, 0)", "dim(shapes.A, 2) == dim(shapes.B, 1)", "dim(shapes.Y, 1) == dim(shapes.A, 0)", "dim(shapes.Y, 2) == dim(shapes.B, 2)", "dim(shapes.A, 0) >= 64", "dim(shapes.A, 2) >= 32", "dim(shapes.B, 2) >= 64", "ceil(dim(shapes.B, 2) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil(dim(shapes.A, 0) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "dim(shapes.Y, 0) <= device.limits.maxComputeWorkgroupsPerDimension"],
+      "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "FusedMatMul.TransBatchATiledReg",
+          "source": {
+            "shader": "matmul-tiled-general-reg.wgsl.jinja",
+            "inputs": {
+              "aShape": "shapes.A",
+              "bShape": "shapes.B",
+              "alpha": "attrs.alpha",
+              "aRank": "ranks.A",
+              "bRank": "ranks.B",
+              "transBatchA": "true",
+              "kTile": "4"
+            }
+          },
+          "bindings": "matrix",
+          "dispatch": {
+            "x": "ceil(dim(shapes.B, 2) / 64)",
+            "y": "ceil(dim(shapes.A, 0) / 64)",
+            "z": "dim(shapes.Y, 0)"
+          }
+        }
+      ]
+    },
+    {
+      "id": "tiled",
+      "priority": 0,
+      "when": ["ranks.A >= 1", "ranks.B >= 1", "f16Ok(dtypes.T)", "transBatchContract", "(dim(shapes.A, 0) if ranks.A == 1 else (dim(shapes.A, ranks.A - 1) if attrs.transA == 0 else (dim(shapes.A, 0) if attrs.transBatchA != 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, 0) if attrs.transBatchB != 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, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2)))) / 16) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil((1 if ranks.B == 1 else (dim(shapes.B, ranks.B - 1) if attrs.transB == 0 else (dim(shapes.B, 0) if attrs.transBatchB != 0 else dim(shapes.B, ranks.B - 2)))) / 16) <= device.limits.maxComputeWorkgroupsPerDimension", "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, 0) if attrs.transBatchA != 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, 0) if attrs.transBatchB != 0 else dim(shapes.B, ranks.B - 2))))) <= device.limits.maxComputeWorkgroupsPerDimension"],
+      "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "FusedMatMul.Tiled",
+          "source": {
+            "shader": "matmul-tiled-general.wgsl.jinja",
+            "inputs": {
+              "aShape": "shapes.A",
+              "bShape": "shapes.B",
+              "transA": "attrs.transA != 0",
+              "transB": "attrs.transB != 0",
+              "alpha": "attrs.alpha",
+              "aRank": "ranks.A",
+              "bRank": "ranks.B",
+              "transBatchA": "attrs.transBatchA != 0",
+              "transBatchB": "attrs.transBatchB != 0"
+            }
+          },
+          "bindings": "matrix",
+          "dispatch": {
+            "x": "ceil((1 if ranks.B == 1 else (dim(shapes.B, ranks.B - 1) if attrs.transB == 0 else (dim(shapes.B, 0) if attrs.transBatchB != 0 else dim(shapes.B, ranks.B - 2)))) / 32)",
+            "y": "ceil((1 if ranks.A == 1 else (dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2)))) / 32)",
+            "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, 0) if attrs.transBatchA != 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, 0) if attrs.transBatchB != 0 else dim(shapes.B, ranks.B - 2)))))"
+          }
+        }
+      ]
+    }
+  ]
+}

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

	@@ -0,0 +1,365 @@

+// 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.
+{% if usesF16 %}
+enable f16;
+{% endif %}
+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 = generalAddressing is defined and generalAddressing %}
+{% set TAIL = tailSafe is defined and tailSafe %}
+{% set OUT = outputBuffer if outputBuffer is defined else "c" %}
+{% set OUT_SCALAR = outScalar if outScalar is defined else T %}
+{% if GENERAL %}
+{% set aR = source.aRank %}
+{% set bR = source.bRank %}
+{% set aBatchLen = aR - 2 %}
+{% set bBatchLen = bR - 2 %}
+{% set batchRank = aBatchLen %}
+{% set aMStride = source.aShape[aR-1] %}
+{% set aKStride = 1 %}
+{% set bKStride = source.bShape[bR-1] %}
+{% set bNStride = 1 %}
+{% endif %}
+const M: u32 = {{ M }}u;
+const K: u32 = {{ K }}u;
+const N: u32 = {{ N }}u;
+const BATCH_COUNT: u32 = {{ batchCount if batchCount is defined else 1 }}u;
+{% if GENERAL %}
+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 }});
+{% else %}
+const A_BATCH_STRIDE: u32 = M * K;
+const B_BATCH_STRIDE: u32 = K * N;
+{% endif %}
+const C_BATCH_STRIDE: u32 = M * N;
+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 < M) {
+{% if operandScalar == "f16" %}
+{% if GENERAL %}
+      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] = f16(a[a_base + a_global * K + k]);
+{% endif %}
+{% else %}
+{% if GENERAL %}
+      tile_A[row * TILE_K + col + col_offset] = f32(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 * K + k]);
+{% endif %}
+{% 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 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 < M && 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 GENERAL %}
+{% 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 %}
+{% else %}
+{% set bIndex = "b_col * K + k" if bTransposed is defined and bTransposed else "k * N + b_col" %}
+    tile_B[row * TILE_K + col + i] = {{ operandScalar }}(b[b_base + {{ bIndex }}]);
+{% 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;
+{% if GENERAL %}
+{% 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 %}
+{% else %}
+{% if hasBias %}
+    {{ OUT }}[offset + row * N + col]       = {{ T }}(scratch[src_slot][0][row * 8u + col] + {{ accScalar }}(bias[col_base + col]));
+    {{ OUT }}[offset + row * N + col + 8u]  = {{ T }}(scratch[src_slot][1][row * 8u + col] + {{ accScalar }}(bias[col_base + col + 8u]));
+    {{ OUT }}[offset + row * N + col + 16u] = {{ T }}(scratch[src_slot][2][row * 8u + col] + {{ accScalar }}(bias[col_base + col + 16u]));
+    {{ OUT }}[offset + row * N + col + 24u] = {{ T }}(scratch[src_slot][3][row * 8u + col] + {{ accScalar }}(bias[col_base + col + 24u]));
+    {{ OUT }}[offset + row * N + col2]       = {{ T }}(scratch[src_slot][0][row * 8u + col2] + {{ accScalar }}(bias[col_base + col2]));
+    {{ OUT }}[offset + row * N + col2 + 8u]  = {{ T }}(scratch[src_slot][1][row * 8u + col2] + {{ accScalar }}(bias[col_base + col2 + 8u]));
+    {{ OUT }}[offset + row * N + col2 + 16u] = {{ T }}(scratch[src_slot][2][row * 8u + col2] + {{ accScalar }}(bias[col_base + col2 + 16u]));
+    {{ OUT }}[offset + row * N + col2 + 24u] = {{ T }}(scratch[src_slot][3][row * 8u + col2] + {{ accScalar }}(bias[col_base + col2 + 24u]));
+{% else %}
+    {{ OUT }}[offset + row * N + col]       = {{ T }}(scratch[src_slot][0][row * 8u + col]);
+    {{ OUT }}[offset + row * N + col + 8u]  = {{ T }}(scratch[src_slot][1][row * 8u + col]);
+    {{ OUT }}[offset + row * N + col + 16u] = {{ T }}(scratch[src_slot][2][row * 8u + col]);
+    {{ OUT }}[offset + row * N + col + 24u] = {{ T }}(scratch[src_slot][3][row * 8u + col]);
+    {{ OUT }}[offset + row * N + col2]       = {{ T }}(scratch[src_slot][0][row * 8u + col2]);
+    {{ OUT }}[offset + row * N + col2 + 8u]  = {{ T }}(scratch[src_slot][1][row * 8u + col2]);
+    {{ OUT }}[offset + row * N + col2 + 16u] = {{ T }}(scratch[src_slot][2][row * 8u + col2]);
+    {{ OUT }}[offset + row * N + col2 + 24u] = {{ T }}(scratch[src_slot][3][row * 8u + col2]);
+{% endif %}
+{% endif %}
+  }
+}
+@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
+) {
+  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) {
+  // 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) {
+{% if GENERAL %}
+{% 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 = source.aShape[aAxis] %}
+{% set bDim = source.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 = source.aShape[aAxis] %}
+{% set bDim = source.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) %}{% set aStride.v = aStride.v * source.aShape[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 * source.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 }} * {{ aStride.v }}u;
+{% endif %}
+{% if bStride.v != 0 %}    b_base = b_base + c{{ axis }} * {{ bStride.v }}u;
+{% endif %}
+{% endif %}
+{% endfor %}
+    let c_base = batch * C_BATCH_STRIDE;
+{% else %}
+    let a_base = batch * A_BATCH_STRIDE;
+    let b_base = batch * B_BATCH_STRIDE;
+    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>;
+    for (var kidx = 0u; kidx < {% if GENERAL and TAIL %}K_FULL{% else %}K{% endif %}; 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) {
+        let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
+      var matA0: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>>(&tile_A, matrix_a_offset, false, TILE_K);
+      var matA1: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>>(&tile_A, matrix_a_offset + 8u * TILE_K, false, TILE_K);
+      let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
+      var matB0: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset, true, TILE_K);
+      var matB1: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 8u * TILE_K, true, TILE_K);
+      var matB2: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 16u * TILE_K, true, TILE_K);
+      var matB3: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 24u * TILE_K, true, TILE_K);
+      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) {
+        let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
+      var matA0: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>>(&tile_A, matrix_a_offset, false, TILE_K);
+      var matA1: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>>(&tile_A, matrix_a_offset + 8u * TILE_K, false, TILE_K);
+      let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
+      var matB0: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset, true, TILE_K);
+      var matB1: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 8u * TILE_K, true, TILE_K);
+      var matB2: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 16u * TILE_K, true, TILE_K);
+      var matB3: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 24u * TILE_K, true, TILE_K);
+      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(&scratch[subtile_id][0], 0u, matC00, false, 8u);
+    subgroupMatrixStore(&scratch[subtile_id][1], 0u, matC01, false, 8u);
+    subgroupMatrixStore(&scratch[subtile_id][2], 0u, matC02, false, 8u);
+    subgroupMatrixStore(&scratch[subtile_id][3], 0u, matC03, false, 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(&scratch[subtile_id][0], 0u, matC10, false, 8u);
+    subgroupMatrixStore(&scratch[subtile_id][1], 0u, matC11, false, 8u);
+    subgroupMatrixStore(&scratch[subtile_id][2], 0u, matC12, false, 8u);
+    subgroupMatrixStore(&scratch[subtile_id][3], 0u, matC13, false, 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, 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,176 @@

+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ 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.
+// Transposed and 1-D operands select other variants and are not handled here.
+// 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. A stores K contiguously and B stores N contiguously,
+// so each staging lane walks the axis its operand already has.
+{% set aR = source.aRank %}
+{% set bR = source.bRank %}
+{% set aBatchLen = aR - 2 %}
+{% set bBatchLen = bR - 2 %}
+{% set batchRank = aBatchLen %}
+{% set aTailStride = namespace(v=1) %}
+{% for j in range(1, aR) %}{% set aTailStride.v = aTailStride.v * source.aShape[j] %}{% endfor %}
+{% set bTailStride = namespace(v=1) %}
+{% for j in range(1, bR) %}{% set bTailStride.v = bTailStride.v * source.bShape[j] %}{% endfor %}
+{% if source.transBatchA %}{% set M = source.aShape[0] %}{% set K = source.aShape[aR-1] %}
+{% else %}{% set M = source.aShape[aR-2] %}{% set K = source.aShape[aR-1] %}{% endif %}
+{% set N = source.bShape[bR-1] %}
+{% if source.transBatchA %}{% set aMStride = aTailStride.v %}{% set aKStride = 1 %}
+{% else %}{% set aMStride = source.aShape[aR-1] %}{% set aKStride = 1 %}{% endif %}
+{% set bKStride = source.bShape[bR-1] %}{% set bNStride = 1 %}
+// Integer MatMul accumulates in its integer type because f32's 24-bit mantissa
+// rounds integer values above 2^24.
+{% set is_int = (scalar == "i32" or scalar == "u32") %}
+{% set accT = scalar if is_int else "f32" %}
+{% set outScalar = source.outScalar if source.outScalar is defined else scalar %}
+// f16 operands stay packed in workgroup memory and widen on shared load.
+// Accumulation remains f32 with the same FMA order, while shared traffic halves.
+{% set tileT = scalar if scalar == "f16" else accT %}
+{% set kTile = source.kTile if source.kTile is defined else 16 %}
+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 %}const ALPHA: {{ accT }} = {{ accT }}(1);{% else %}const ALPHA: f32 = f32({{ source.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.
+const BK: u32 = {{ kTile }}u;
+const BM: u32 = 64u;
+const BN: u32 = 64u;
+const TM: u32 = 4u;   // per-thread micro-tile rows
+const TN: u32 = 4u;   // per-thread micro-tile cols
+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 (source.transBatchA and aAxis >= 0) else aAxis %}
+{% set bStored = bAxis %}
+{% set aDim = source.aShape[aStored] %}
+{% set bDim = source.bShape[bStored] if bStored >= 0 else 1 %}
+{% if aDim > 1 or bDim > 1 %}{% set hasBatchCoord.value = true %}{% endif %}
+{% endfor %}
+@compute @workgroup_size(16, 16, 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 * 16u + lid.x;
+  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 (source.transBatchA and aAxis >= 0) else aAxis %}
+{% set bStored = bAxis %}
+{% set aDim = source.aShape[aStored] %}
+{% set bDim = source.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) %}{% set aStride.v = aStride.v * source.aShape[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 * source.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 }} * {{ aStride.v }}u;
+{% 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); }
+  let numTiles = (K + BK - 1u) / BK;
+  for (var kt: u32 = 0u; kt < numTiles; kt = kt + 1u) {
+    let kBase = kt * BK;
+    // 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 + 256u) {
+      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 + 256u) {
+      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 regIndent = "    " %}{% set regT = accT %}{{ regIndent }}let aRow = lid.y * TM;
+{{ regIndent }}let bCol = lid.x * TN;
+{{ regIndent }}for (var kv: u32 = 0u; kv < BK / 4u; kv = kv + 1u) {
+{{ regIndent }}  var av: array<vec4<{{ regT }}>, TM>;
+{{ regIndent }}  var bv: array<vec4<{{ regT }}>, TN>;
+{{ regIndent }}  for (var i: u32 = 0u; i < TM; i = i + 1u) { av[i] = vec4<{{ regT }}>(tileA[aRow + i][kv]); }
+{{ regIndent }}  for (var j: u32 = 0u; j < TN; j = j + 1u) { bv[j] = vec4<{{ regT }}>(tileB[bCol + j][kv]); }
+{{ regIndent }}  for (var i: u32 = 0u; i < TM; i = i + 1u) {
+{{ regIndent }}    for (var j: u32 = 0u; j < TN; j = j + 1u) {
+{{ regIndent }}      acc[i * TN + j] = acc[i * TN + j] + dot(av[i], bv[j]);
+{{ regIndent }}    }
+{{ regIndent }}  }
+{{ regIndent }}}
+    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,176 @@

+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ 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 = source.aRank %}
+{% set bR = source.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 aTailStride = namespace(v=1) %}
+{% for j in range(1, aR) %}{% set aTailStride.v = aTailStride.v * source.aShape[j] %}{% endfor %}
+{% set bTailStride = namespace(v=1) %}
+{% for j in range(1, bR) %}{% set bTailStride.v = bTailStride.v * source.bShape[j] %}{% endfor %}
+{% if aVec %}{% set M = 1 %}{% set K = source.aShape[0] %}
+{% elif source.transBatchA and source.transA %}{% set M = source.aShape[aR-1] %}{% set K = source.aShape[0] %}
+{% elif source.transBatchA %}{% set M = source.aShape[0] %}{% set K = source.aShape[aR-1] %}
+{% elif source.transA %}{% set M = source.aShape[aR-1] %}{% set K = source.aShape[aR-2] %}
+{% else %}{% set M = source.aShape[aR-2] %}{% set K = source.aShape[aR-1] %}{% endif %}
+{% if bVec %}{% set N = 1 %}
+{% elif source.transBatchB and source.transB %}{% set N = source.bShape[0] %}
+{% elif source.transB %}{% set N = source.bShape[bR-2] %}
+{% else %}{% set N = source.bShape[bR-1] %}{% endif %}
+{% if aVec %}{% set aMStride = 0 %}{% set aKStride = 1 %}
+{% elif source.transBatchA and source.transA %}{% set aMStride = 1 %}{% set aKStride = aTailStride.v %}
+{% elif source.transBatchA %}{% set aMStride = aTailStride.v %}{% set aKStride = 1 %}
+{% elif source.transA %}{% set aMStride = 1 %}{% set aKStride = source.aShape[aR-1] %}
+{% else %}{% set aMStride = source.aShape[aR-1] %}{% set aKStride = 1 %}{% endif %}
+{% if bVec %}{% set bKStride = 1 %}{% set bNStride = 0 %}
+{% elif source.transBatchB and source.transB %}{% set bKStride = 1 %}{% set bNStride = bTailStride.v %}
+{% elif source.transBatchB %}{% set bKStride = bTailStride.v %}{% set bNStride = 1 %}
+{% elif source.transB %}{% set bKStride = 1 %}{% set bNStride = source.bShape[bR-1] %}
+{% else %}{% set bKStride = source.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({{ source.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.
+const BK: u32 = 16u;
+const BM: u32 = 32u;
+const BN: u32 = 32u;
+var<workgroup> tileA: array<array<{{ tileT }}, 16>, 32>;
+var<workgroup> tileB: array<array<{{ tileT }}, 32>, 16>;
+{% 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 (source.transBatchA and aAxis >= 0) else aAxis %}
+{% set bStored = (bAxis + 1) if (source.transBatchB and bAxis >= 0) else bAxis %}
+{% set aDim = source.aShape[aStored] if aStored >= 0 else 1 %}
+{% set bDim = source.bShape[bStored] if bStored >= 0 else 1 %}
+{% if aDim > 1 or bDim > 1 %}{% set hasBatchCoord.value = true %}{% endif %}
+{% endfor %}
+@compute @workgroup_size(16, 16, 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 * 16u + 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 (source.transBatchA and aAxis >= 0) else aAxis %}
+{% set bStored = (bAxis + 1) if (source.transBatchB and bAxis >= 0) else bAxis %}
+{% set aDim = source.aShape[aStored] if aStored >= 0 else 1 %}
+{% set bDim = source.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) %}{% set aStride.v = aStride.v * source.aShape[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 * source.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 * 256u;
+      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,42 @@

+// GEMV specialization for y[N] = a[K] @ B[K, N]: each workgroup owns 32
+// consecutive vec4 column groups (128 output columns). Shared by MatMul and
+// matrix-shaped Einsum decode projections.
+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+{% set OUT = outputBuffer if outputBuffer is defined else "c" %}
+const LANES: u32 = 32u;
+// SLICES is chosen from the workgroup COUNT, not from a fixed default: this
+// kernel launches ceilDiv(N, 128) workgroups, so a narrow-N projection starves
+// the device unless each workgroup carries more of the K reduction.
+const SLICES: u32 = {{ gemvSlices }}u;
+var<workgroup> partials: array<vec4<f32>, LANES * SLICES>;
+@compute @workgroup_size(32, {{ 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) {
+    for (var k = slice; k < params.K; k = k + SLICES) {
+      acc = acc + f32(a[k]) * vec4<f32>(b[k * params.N4 + cg]);
+    }
+  }
+  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];
+    }
+    {{ OUT }}[cg] = vec4<{{ T }}>(total);
+  }
+}

build/webgpu/metadata.json ADDED Viewed

	@@ -0,0 +1,22 @@

+{
+  "name": "com.microsoft.FusedMatMul",
+  "id": "_com_microsoft_fusedmatmul_webgpu_80c34de",
+  "version": 1,
+  "license": "Apache-2.0",
+  "backend": { "type": "webgpu" },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "bench.json": "uw58QO9KL2kJYYMWMZJSUGIwBmsy6bLImlj0E1NNCyw=",
+      "fused-matmul-subgroup-matrix.wgsl.jinja": "d8LaODNB5B+qCRM5SSDSbasoJTBoTkjR8BtwuO8AW6U=",
+      "manifest.json": "F7nDLMjrVNWkpIGc9rUB5ZrhRoW06xk+PPX+M6X7XCA=",
+      "matmul-subgroup-matrix-ext.wgsl.jinja": "Q/5RQ4YPRDz/EiGmw1HigsSCohWzsx6muR4ghzigfWs=",
+      "matmul-tiled-general-reg.wgsl.jinja": "8E6pn6tg6t7LuoU6ktkjL5AZOgrd0HHkFRe8VJC7jwk=",
+      "matmul-tiled-general.wgsl.jinja": "pqYgOQCvHqeTthCEShE8dbN1GJjpFz3LD1K3sBYARMc=",
+      "matmul-vector-matrix-vec4.wgsl.jinja": "9syYZdAgoDnCn5rurCHP9tHtlwm1IXwKvPoj48hH3+E=",
+      "test.json": "Au5pkdVg7fcbm2EDBuhQvc4CPiri5lrQs394cLC5iL0="
+    }
+  },
+  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
+  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.FusedMatMul" }
+}

build/webgpu/test.json ADDED Viewed

	@@ -0,0 +1,1420 @@

+{
+  "op": "com.microsoft.FusedMatMul",
+  "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],
+    "ort_float32_trans_batch_a_input_A": [0, 1, 2, 3, 12, 13, 14, 15, 4, 5, 6, 7, 16, 17, 18, 19, 8, 9, 10, 11, 20, 21, 22, 23],
+    "ort_float32_trans_batch_b_input_B": [1, 0, 1, 2, 0, 1, -1, 0, 1, 1, 0, 1, 2, -1, 1, 1]
+  },
+  "cases": [
+    {
+      "name": "ort_float32_broadcast_rank4_by_rank3",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
+        "test": "FusedMatMulOpTest.FloatTypeNoTranspose"
+      },
+      "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"
+      },
+      "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"
+      },
+      "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"
+      },
+      "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"
+      },
+      "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"
+      },
+      "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"
+      },
+      "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"
+      },
+      "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"
+      },
+      "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"
+      },
+      "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"
+      },
+      "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"
+      },
+      "attrs": { "alpha": 0 },
+      "inputs": {
+        "A": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "values", "values": [1.0, 2.0, 3.0, 4.0] } },
+        "B": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "values", "values": [5.0, 6.0, 7.0, 8.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"
+      },
+      "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"
+      },
+      "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"
+      },
+      "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"
+      },
+      "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_trans_batch_a",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
+        "test": "FusedMatMulOpTest.FloatTypeTransposeBatch",
+        "notes": "Exercises transBatchA, where logical [batch,N,M] input A is stored as [N,batch,M]."
+      },
+      "attrs": { "transBatchA": 1 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [3, 2, 4],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/ort_float32_trans_batch_a_input_A" } }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [2, 4, 2],
+          "data": {
+            "kind": "values",
+            "values": [1.0, 0.0, 0.0, 1.0, 1.0, 1.0, 2.0, -1.0, 1.0, 2.0, -1.0, 0.0, 0.0, 1.0, 1.0, 1.0]
+          }
+        }
+      },
+      "outputs": {
+        "Y": {
+          "dtype": "float32",
+          "shape": [2, 3, 2],
+          "tolerance": 0.000001,
+          "data": { "kind": "values", "values": [8.0, 0.0, 24.0, 4.0, 40.0, 8.0, 14.0, 53.0, 18.0, 69.0, 22.0, 85.0] }
+        }
+      }
+    },
+    {
+      "name": "ort_float32_trans_batch_b",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
+        "test": "FusedMatMulOpTest.FloatTypeTransposeBatch",
+        "notes": "Exercises transBatchB, where logical [batch,K,N] input B is stored as [K,batch,N]."
+      },
+      "attrs": { "transBatchB": 1 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 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, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0, 18.0, 19.0, 20.0, 21.0, 22.0, 23.0]
+          }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [4, 2, 2],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/ort_float32_trans_batch_b_input_B" } }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 3, 2], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "ort_float32_trans_batch_ab_scaled",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
+        "test": "FusedMatMulOpTest.FloatTypeTransposeBatch",
+        "notes": "Exercises combined transBatchA and transBatchB with alpha scaling."
+      },
+      "attrs": { "transBatchA": 1, "transBatchB": 1, "alpha": 0.5 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [3, 2, 4],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/ort_float32_trans_batch_a_input_A" } }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [4, 2, 2],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/ort_float32_trans_batch_b_input_B" } }
+        }
+      },
+      "outputs": {
+        "Y": {
+          "dtype": "float32",
+          "shape": [2, 3, 2],
+          "tolerance": 0.000001,
+          "data": { "kind": "values", "values": [4.0, 0.0, 12.0, 2.0, 20.0, 4.0, 7.0, 26.5, 9.0, 34.5, 11.0, 42.5] }
+        }
+      }
+    },
+    {
+      "name": "ort_float32_scaled_no_transpose",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/fused_matmul_op_test.cc",
+        "test": "FusedMatMulOpTest.FloatTypeScale"
+      },
+      "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"
+      },
+      "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 selects the subgroup-matrix path; finite subnormal dot products 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.02 } }
+    },
+    {
+      "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.02 } }
+    },
+    {
+      "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.03 } }
+    },
+    {
+      "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.03 } }
+    },
+    {
+      "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.02 } }
+    },
+    {
+      "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.03 } }
+    },
+    {
+      "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.03 } }
+    },
+    {
+      "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": "trans_batch_a_broadcast_rank4_f32",
+      "attrs": { "transBatchA": 1 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [2, 3, 2, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.027, "scale": 0.2 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [3, 1, 4, 5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.007, "scale": 0.2 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [3, 2, 2, 5], "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.05 } }
+    },
+    {
+      "name": "f32_decode_gemv_m1_k65_n68_vec4_compact",
+      "provenance": {
+        "notes": "Compact M=1 float32 GEMV correctness lock for the model-shaped K=4096,N=4096 bandwidth-bound benchmark. Odd K preserves the sliced reduction while N=68 exercises the final partial 128-column workgroup."
+      },
+      "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": "f32_rank4_by_rank2_shared_weight_compact",
+      "provenance": {
+        "notes": "Compact rank-4 by rank-2 shared-weight broadcast lock for the attention-shaped benchmark. Odd M/K/N exercise batch offset and tile-tail handling."
+      },
+      "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": "Compact f16 subgroup-matrix lock with both a partial K=34 tile and N=66 output tail." },
+      "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.05 } }
+    },
+    {
+      "name": "subgroup_matrix_broadcast_rank4x3_f16_compact",
+      "provenance": {
+        "notes": "Compact rank-4 by rank-3 broadcast lock for the model-shaped [1,8,M,K] x [8,K,N] stress case."
+      },
+      "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.05 } }
+    },
+    {
+      "name": "broadcast_rank4_tiled_reg_f16_compact",
+      "provenance": {
+        "notes": "Compact rank-4 by rank-3 broadcast lock for the register-blocked non-subgroup-matrix path. Odd M/K/N exercise every output and reduction tail."
+      },
+      "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.05 } }
+    },
+    {
+      "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": "transbatch_a_dense_m_tail_f16_compact",
+      "provenance": {
+        "notes": "Compact lock for stored [M,batch,K] transBatchA addressing. M=65 exercises the subgroup-matrix row tail; all/no-mma/no-subgroups select the MMA/register-blocked portable paths used by the model-shaped stress case."
+      },
+      "attrs": { "alpha": 0.5, "transBatchA": 1 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [65, 2, 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": [2, 65, 64], "tolerance": 0.05 } }
+    },
+    {
+      "name": "transbatch_a_dense_m_tail_f32_compact",
+      "provenance": {
+        "notes": "Float32 counterpart proving the same transBatchA MMA/register-blocked split remains correct when shader-f16 is unavailable."
+      },
+      "attrs": { "alpha": 0.5, "transBatchA": 1 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [65, 2, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.1 }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [2, 32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.009, "cosStep": 0.029, "scale": 0.1 }
+        }
+      },
+      "outputs": { "Y": { "dtype": "float32", "shape": [2, 65, 64], "tolerance": 0.0003 } }
+    },
+    {
+      "name": "rank5_three_batch_dims",
+      "attrs": { "alpha": 1, "transA": 0, "transB": 0, "transBatchA": 0, "transBatchB": 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_lock",
+      "provenance": {
+        "notes": "subgroup_matrix_kn_tail_f16_compact draws both operands from a zero-mean sinusoid, so the K=34 dot random-walks to 0.013 and the 0.05 absolute tolerance is 3.7x the whole output: alpha could be ignored, a K tile dropped, or Y halved and it still passes. Offsetting both operands makes every element approach alpha * K * aOffset * bOffset (~3.4), a value proportional to the contraction length, so the K=34 partial tile, the N=66 column tail and the alpha=0.5 epilogue are all under test on the tail-broadcast 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_lock",
+      "provenance": {
+        "notes": "Well-conditioned sibling of subgroup_matrix_broadcast_rank4x3_f16_compact, whose zero-mean operands cancel to 0.021 under a 0.05 absolute tolerance (2.4x blind). Offsetting both operands drives each output to ~K * aOffset * bOffset, so the rank-4 by rank-3 batch broadcast now pins the per-batch B slice and the K=32 count instead of only the addressing pattern."
+      },
+      "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-side batch broadcast: A carries dim 1 on the batch axis while B carries 2, so one A slice feeds both output batches. Every other rank-4 by rank-3 fixture broadcasts B (or matches both extents), leaving the A-broadcast stride - a zero A batch stride against a non-zero B one - unrendered. Operands are offset so each output lands near alpha * K * aOffset * bOffset and a swapped or dropped batch stride moves the result rather than cancelling."
+      },
+      "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_lock",
+      "provenance": {
+        "notes": "broadcast_rank4_tiled_reg_f16_compact cancels to 0.015 against a 0.05 absolute tolerance, so the register-blocked rank-4 route was 3.4x blind to any uniform scale error - including a silently dropped alpha. Offsetting both operands makes each output ~alpha * K * aOffset * bOffset with K=33 (one full 32-wide k tile plus a 1-element tail), so a mis-counted k tile, a doubled accumulator or a missing alpha all move it far outside tolerance."
+      },
+      "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": "transbatch_a_dense_m_tail_f16_offset_alpha_scale_lock",
+      "provenance": {
+        "notes": "transbatch_a_dense_m_tail_f16_compact was the most blind fixture of this op: a 0.05 absolute tolerance over a 0.010 output, so a 4.8x uniform scale error passes. Offsetting both operands makes every element ~alpha * K * aOffset * bOffset, so the [M,batch,K] transBatchA stride, the K=32 contraction and the alpha=0.5 scale are pinned rather than just the row/column addressing."
+      },
+      "attrs": { "alpha": 0.5, "transBatchA": 1 },
+      "inputs": {
+        "A": {
+          "dtype": "float16",
+          "shape": [65, 2, 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": [2, 65, 64], "tolerance": 0.03, "relTolerance": 0.01 } }
+    },
+    {
+      "name": "aligned_f16_transA_transB_alpha_offset_scale_lock",
+      "provenance": {
+        "notes": "aligned_f16_transA_transB_alpha_64x32 cancels to 0.11 under a 0.03 absolute tolerance, so a 27% scale error - alpha 0.5 applied twice, say - still passes. Offsetting both operands makes each output ~alpha * K * aOffset * bOffset, which keeps the doubly-transposed addressing under test while making the alpha epilogue and the K=32 count observable on both the 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_lock",
+      "provenance": {
+        "notes": "f16_unaligned_3x5x7 cancels to 0.15 under a 0.02 absolute tolerance (13% blind). The same 3x5 by 5x7 shape with offset operands lands at ~K * aOffset * bOffset, so the odd-K/odd-N scalar and tiled fallbacks - the routes small unaligned shapes actually take - now catch a dropped K element or a doubled tail."
+      },
+      "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_lock",
+      "provenance": {
+        "notes": "aligned_f16_plain_64x32x64 cancels to 0.124 under a 0.02 absolute tolerance, so a 16% uniform scale error passes on the fully aligned f16 subgroup-matrix path (M=64, K=32, N=64 - no tail at all). Offsetting both operands makes each element ~K * aOffset * bOffset, so the aligned 32x64 tile now pins the k-block count and the scratch drain, not just the addressing."
+      },
+      "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_lock",
+      "provenance": {
+        "notes": "Batched twin of aligned_f16_plain_64x32x64_offset_scale_lock. aligned_f16_batched_plain_2x64x32x64 cancels to 0.124 under a 0.03 absolute tolerance (24% blind); with offset operands each element is ~K * aOffset * bOffset and the two batches carry different data, so the batch stride is pinned alongside the magnitude."
+      },
+      "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_lock",
+      "provenance": {
+        "notes": "subgroup_matrix_m_tail_57_partial_block_f16 cancels to 0.062 under a 0.03 absolute tolerance, so the partial 8-row MMA block at M=57 (one full 32-row tile plus a 25-row remainder) was 49% blind to a scale error. With offset operands each element is ~K * aOffset * bOffset, so the guarded rows have to carry the same magnitude as the full ones - a tail block that accumulates one k step short, or drains a stale scratch slot, no longer looks like noise."
+      },
+      "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_lock",
+      "provenance": {
+        "notes": "aligned_f16_transA_64x32 cancels to 0.053 under a 0.03 absolute tolerance - 57% blind, the loosest of the transpose fixtures. transA alone (transB = 0) is its own render of the operand loader, and with offset operands each element is ~K * aOffset * bOffset, so the transposed A stride is pinned against a k-count error rather than only against a wrong-element read."
+      },
+      "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_lock",
+      "provenance": {
+        "notes": "Combines both transposes with an M tail (50 = 32 + 18). Offset operands make each element approximately K * aOffset * bOffset, checking the guarded rows of the doubly-transposed load for magnitude as well as placement; the zero-mean companion cancels to 0.099 under a 0.05 absolute tolerance and cannot provide that scale lock."
+      },
+      "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_lock",
+      "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, "transBatchA": 0, "transBatchB": 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": "trans_batch_a_and_trans_a_with_trans_batch_b_and_trans_b",
+      "provenance": {
+        "notes": "Combines transBatchA with transA and transBatchB with transB, checking both logical extents and both stride pairs. Expected values follow the documented storage orders A[k][batch][m] and B[n][batch][k]."
+      },
+      "attrs": { "transA": 1, "transB": 1, "transBatchA": 1, "transBatchB": 1 },
+      "inputs": {
+        "A": {
+          "dtype": "float32",
+          "shape": [4, 2, 3],
+          "data": {
+            "kind": "values",
+            "values": [0.0, 1.0, 2.0, 4.0, 5.0, 6.0, 10.0, 11.0, 12.0, 14.0, 15.0, 16.0, 20.0, 21.0, 22.0, 24.0, 25.0, 26.0, 30.0, 31.0, 32.0, 34.0, 35.0, 36.0]
+          }
+        },
+        "B": {
+          "dtype": "float32",
+          "shape": [5, 2, 4],
+          "data": {
+            "kind": "values",
+            "values": [0.0, 1.0, 2.0, 3.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0, 17.0, 18.0, 19.0, 20.0, 21.0, 22.0, 23.0, 24.0, 24.0, 25.0, 26.0, 27.0, 28.0, 29.0, 30.0, 31.0, 31.0, 32.0, 33.0, 34.0]
+          }
+        }
+      },
+      "outputs": {
+        "Y": {
+          "dtype": "float32",
+          "shape": [2, 3, 5],
+          "tolerance": 0.000001,
+          "data": {
+            "kind": "values",
+            "values": [140.0, 560.0, 980.0, 1400.0, 1820.0, 146.0, 594.0, 1042.0, 1490.0, 1938.0, 152.0, 628.0, 1104.0, 1580.0, 2056.0, 392.0, 924.0, 1456.0, 1988.0, 2520.0, 410.0, 970.0, 1530.0, 2090.0, 2650.0, 428.0, 1016.0, 1604.0, 2192.0, 2780.0]
+          }
+        }
+      }
+    }
+  ]
+}