sync 2e7068faf55e

Browse files

Files changed (13) hide show

README.md +88 -0
build/webgpu/bench.json +922 -0
build/webgpu/conv-1x1-gemm-tiled-reg.wgsl.jinja +419 -0
build/webgpu/conv-1x1-gemm-tiled.wgsl.jinja +138 -0
build/webgpu/conv-1x1-subgroup-matrix.wgsl.jinja +353 -0
build/webgpu/conv-direct-nd.wgsl.jinja +86 -0
build/webgpu/conv-direct-unrolled.wgsl.jinja +95 -0
build/webgpu/conv-im2col-nchw.wgsl.jinja +73 -0
build/webgpu/conv1d-tiled-reg.wgsl.jinja +166 -0
build/webgpu/conv2d-grouped-large-w4.wgsl.jinja +176 -0
build/webgpu/manifest.json +0 -0
build/webgpu/metadata.json +25 -0
build/webgpu/test.json +0 -0

README.md CHANGED Viewed

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

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# com.microsoft.FusedConv
+`com.microsoft`  ·  ONNX Runtime contrib operator  ·  contrib since_version 1
+## Description
+Applies an N-dimensional convolution with optional bias `B` and residual `Z`, followed by an optional fused activation. Omitting `activation` leaves the convolution result unchanged. Supported activations are `Relu`, `LeakyRelu`, `Sigmoid`, `Tanh`, `HardSigmoid`, `HardSwish`, and `Clip`; other schema-permitted activation strings are not implemented. The implementation supports one to three spatial dimensions and float16 or float32; higher spatial ranks and float64 are not implemented.
+See the [ONNX Runtime `FusedConv` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.FusedConv) for the reference semantics.
+## Inputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `X` | `x` | `T` | — | — | Input data tensor of shape `(N, C, D1, ..., Dn)` for one to three spatial dimensions. | required |
+| `W` | `w` | `T` | — | — | Convolution filter tensor of shape `(M, C/group, k1, ..., kn)`, with the same spatial rank as `X`. | required |
+| `B` | `bias` | `T` | `1` | — | Optional 1-D bias tensor of length `out_channels`, broadcast-added to each output channel. | optional |
+| `Z` | `zResidual` | `T` | same as `X` | — | Optional residual tensor with the same shape as the output `Y`, added before the activation. | optional |
+## Outputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `Y` | `y` | `T` | same as `X` | derived; see description | Output feature map tensor after convolution, optional bias/residual addition, and the fused activation. | required |
+## Attributes
+Attributes and default values (overridable per request):
+| Attribute | Default | Description |
+| --- | --- | --- |
+| `auto_pad` | `"NOTSET"` | Automatic padding mode. `NOTSET` uses `pads`; `SAME_UPPER` and `SAME_LOWER` choose padding so each output spatial size is `ceil(input / stride)`; `VALID` uses no padding. |
+| `group` | `1` | Number of groups that input and output channels are split into; defaults to 1. |
+| `activation` | — | Optional fused activation name: `Relu`, `LeakyRelu`, `Sigmoid`, `Tanh`, `HardSigmoid`, `HardSwish`, or `Clip`. Omission applies no activation. |
+| `activation_params` | — | Positional parameters for the fused activation: exactly `[alpha]` is required for `LeakyRelu`, and exactly `[alpha, beta]` or `[min, max]` is required for `HardSigmoid` or `Clip`, respectively. Parameter-free activations ignore this attribute. |
+| `dilations` | — | Optional dilation factors, one positive integer per spatial axis. Omission means all ones. |
+| `kernel_shape` | — | Optional kernel shape, one positive integer per spatial axis. When present, it must match the spatial dimensions of the weight tensor; omission infers the shape from the weights. |
+| `pads` | — | Optional explicit padding in ONNX order `[begin_axis_0, ..., begin_axis_n, end_axis_0, ..., end_axis_n]`. Omission means all zeros; it cannot be combined with an automatic padding mode. |
+| `strides` | — | Optional stride factors, one positive integer per spatial axis. Omission means all ones. |
+## 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
+- [`conv-1x1-gemm-tiled-reg.wgsl.jinja`](build/webgpu/conv-1x1-gemm-tiled-reg.wgsl.jinja)
+- [`conv-1x1-gemm-tiled.wgsl.jinja`](build/webgpu/conv-1x1-gemm-tiled.wgsl.jinja)
+- [`conv-1x1-subgroup-matrix.wgsl.jinja`](build/webgpu/conv-1x1-subgroup-matrix.wgsl.jinja)
+- [`conv-direct-nd.wgsl.jinja`](build/webgpu/conv-direct-nd.wgsl.jinja)
+- [`conv-direct-unrolled.wgsl.jinja`](build/webgpu/conv-direct-unrolled.wgsl.jinja)
+- [`conv-im2col-nchw.wgsl.jinja`](build/webgpu/conv-im2col-nchw.wgsl.jinja)
+- [`conv1d-tiled-reg.wgsl.jinja`](build/webgpu/conv1d-tiled-reg.wgsl.jinja)
+- [`conv2d-grouped-large-w4.wgsl.jinja`](build/webgpu/conv2d-grouped-large-w4.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.FusedConv", { version: 1 });
+const { y } = await kernel({
+  x: { data: xData, shape: [1, 32, 8, 8] },
+  w: { data: wData, shape: [32, 32, 1, 1] },
+});
+```

build/webgpu/bench.json ADDED Viewed

	@@ -0,0 +1,922 @@

+{
+  "op": "com.microsoft.FusedConv",
+  "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
+  "cases": [
+    {
+      "name": "fusedconv-f32-b1c16m32-32x32-k3",
+      "preset": "smoke",
+      "vars": {
+        "batch": 1,
+        "inChannels": 16,
+        "outChannels": 32,
+        "inH": 32,
+        "inW": 32,
+        "kernelH": 3,
+        "kernelW": 3,
+        "strideH": 1,
+        "strideW": 1,
+        "padH": 1,
+        "padW": 1
+      },
+      "attrs": { "activation": "Relu", "strides": [1, 1], "pads": [1, 1, 1, 1] },
+      "inputs": {
+        "x": { "shape": [1, 16, 32, 32], "dtype": "float32", "dist": "normal", "seed": 112, "scale": 0.2 },
+        "w": { "shape": [32, 16, 3, 3], "dtype": "float32", "dist": "normal", "seed": 113, "scale": 0.1 }
+      },
+      "outputs": { "y": { "shape": [1, 32, 32, 32], "dtype": "float32" } },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fusedconv-f32-bias-b1c16m32-32x32-k3",
+      "preset": "smoke",
+      "vars": {
+        "batch": 1,
+        "inChannels": 16,
+        "outChannels": 32,
+        "inH": 32,
+        "inW": 32,
+        "kernelH": 3,
+        "kernelW": 3,
+        "strideH": 1,
+        "strideW": 1,
+        "padH": 1,
+        "padW": 1
+      },
+      "attrs": { "activation": "Relu", "strides": [1, 1], "pads": [1, 1, 1, 1] },
+      "inputs": {
+        "x": { "shape": [1, 16, 32, 32], "dtype": "float32", "dist": "normal", "seed": 113, "scale": 0.2 },
+        "w": { "shape": [32, 16, 3, 3], "dtype": "float32", "dist": "normal", "seed": 114, "scale": 0.1 },
+        "bias": { "shape": [32], "dtype": "float32", "dist": "normal", "seed": 115, "scale": 0.05 }
+      },
+      "outputs": { "y": { "shape": [1, 32, 32, 32], "dtype": "float32" } },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fused-flux2vae-dec-mid-512c-64x64-k3",
+      "preset": "model",
+      "vars": {
+        "batch": 1,
+        "inChannels": 512,
+        "outChannels": 512,
+        "inH": 64,
+        "inW": 64,
+        "kernelH": 3,
+        "kernelW": 3,
+        "strideH": 1,
+        "strideW": 1,
+        "padH": 1,
+        "padW": 1
+      },
+      "attrs": { "activation": "Relu", "strides": [1, 1], "pads": [1, 1, 1, 1] },
+      "inputs": {
+        "x": { "shape": [1, 512, 64, 64], "dtype": "float32", "dist": "normal", "seed": 340, "scale": 0.2 },
+        "w": { "shape": [512, 512, 3, 3], "dtype": "float32", "dist": "normal", "seed": 341, "scale": 0.02 }
+      },
+      "outputs": { "y": { "shape": [1, 512, 64, 64], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fused-flux2vae-dec-up1-512c-128x128-k3",
+      "preset": "model",
+      "vars": {
+        "batch": 1,
+        "inChannels": 512,
+        "outChannels": 512,
+        "inH": 128,
+        "inW": 128,
+        "kernelH": 3,
+        "kernelW": 3,
+        "strideH": 1,
+        "strideW": 1,
+        "padH": 1,
+        "padW": 1
+      },
+      "attrs": { "activation": "Relu", "strides": [1, 1], "pads": [1, 1, 1, 1] },
+      "inputs": {
+        "x": { "shape": [1, 512, 128, 128], "dtype": "float32", "dist": "normal", "seed": 342, "scale": 0.2 },
+        "w": { "shape": [512, 512, 3, 3], "dtype": "float32", "dist": "normal", "seed": 343, "scale": 0.02 }
+      },
+      "outputs": { "y": { "shape": [1, 512, 128, 128], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fused-flux2vae-attn-proj-512c-64x64-k1",
+      "preset": "model",
+      "vars": {
+        "batch": 1,
+        "inChannels": 512,
+        "outChannels": 512,
+        "inH": 64,
+        "inW": 64,
+        "kernelH": 1,
+        "kernelW": 1,
+        "strideH": 1,
+        "strideW": 1,
+        "padH": 0,
+        "padW": 0
+      },
+      "attrs": { "activation": "Relu", "strides": [1, 1], "pads": [0, 0, 0, 0] },
+      "inputs": {
+        "x": { "shape": [1, 512, 64, 64], "dtype": "float32", "dist": "normal", "seed": 344, "scale": 0.2 },
+        "w": { "shape": [512, 512, 1, 1], "dtype": "float32", "dist": "normal", "seed": 345, "scale": 0.02 }
+      },
+      "outputs": { "y": { "shape": [1, 512, 64, 64], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fused-flux2vae-attn-proj-bias-512c-64x64-k1",
+      "preset": "model",
+      "vars": {
+        "batch": 1,
+        "inChannels": 512,
+        "outChannels": 512,
+        "inH": 64,
+        "inW": 64,
+        "kernelH": 1,
+        "kernelW": 1,
+        "strideH": 1,
+        "strideW": 1,
+        "padH": 0,
+        "padW": 0
+      },
+      "attrs": { "activation": "Relu", "strides": [1, 1], "pads": [0, 0, 0, 0] },
+      "inputs": {
+        "x": { "shape": [1, 512, 64, 64], "dtype": "float32", "dist": "normal", "seed": 346, "scale": 0.2 },
+        "w": { "shape": [512, 512, 1, 1], "dtype": "float32", "dist": "normal", "seed": 347, "scale": 0.02 },
+        "bias": { "shape": [512], "dtype": "float32", "dist": "normal", "seed": 348, "scale": 0.05 }
+      },
+      "outputs": { "y": { "shape": [1, 512, 64, 64], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fusedconv-f32-1x1-b1c256m256-32x32",
+      "preset": "smoke",
+      "vars": {
+        "batch": 1,
+        "inChannels": 256,
+        "outChannels": 256,
+        "inH": 32,
+        "inW": 32,
+        "kernelH": 1,
+        "kernelW": 1,
+        "strideH": 1,
+        "strideW": 1,
+        "padH": 0,
+        "padW": 0
+      },
+      "attrs": { "activation": "Relu", "strides": [1, 1], "pads": [0, 0, 0, 0] },
+      "inputs": {
+        "x": { "shape": [1, 256, 32, 32], "dtype": "float32", "dist": "normal", "seed": 350, "scale": 0.2 },
+        "w": { "shape": [256, 256, 1, 1], "dtype": "float32", "dist": "normal", "seed": 351, "scale": 0.02 }
+      },
+      "outputs": { "y": { "shape": [1, 256, 32, 32], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fusedconv-f32-1x1-bias-b1c256m256-32x32",
+      "preset": "smoke",
+      "provenance": {
+        "source": "authored for performance coverage",
+        "notes": "Bias twin of the paired non-bias case at the same 64 register tiles. No bias case sat in the 64-127 tile band, so the bias register-tile floor had no benchmark on either side of it."
+      },
+      "vars": {
+        "batch": 1,
+        "inChannels": 256,
+        "outChannels": 256,
+        "inH": 32,
+        "inW": 32,
+        "kernelH": 1,
+        "kernelW": 1,
+        "strideH": 1,
+        "strideW": 1,
+        "padH": 0,
+        "padW": 0
+      },
+      "attrs": { "activation": "Relu", "strides": [1, 1], "pads": [0, 0, 0, 0] },
+      "inputs": {
+        "x": { "shape": [1, 256, 32, 32], "dtype": "float32", "dist": "normal", "seed": 3521, "scale": 0.2 },
+        "w": { "shape": [256, 256, 1, 1], "dtype": "float32", "dist": "normal", "seed": 3522, "scale": 0.02 },
+        "bias": { "shape": [256], "dtype": "float32", "dist": "normal", "seed": 3523, "scale": 0.05 }
+      },
+      "outputs": { "y": { "shape": [1, 256, 32, 32], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fusedconv-f32-im2col-b1c32m64-32x32-k3",
+      "preset": "smoke",
+      "vars": {
+        "batch": 1,
+        "inChannels": 32,
+        "outChannels": 64,
+        "inH": 32,
+        "inW": 32,
+        "kernelH": 3,
+        "kernelW": 3,
+        "strideH": 1,
+        "strideW": 1,
+        "padH": 1,
+        "padW": 1
+      },
+      "attrs": { "activation": "Relu", "strides": [1, 1], "pads": [1, 1, 1, 1] },
+      "inputs": {
+        "x": { "shape": [1, 32, 32, 32], "dtype": "float32", "dist": "normal", "seed": 352, "scale": 0.2 },
+        "w": { "shape": [64, 32, 3, 3], "dtype": "float32", "dist": "normal", "seed": 353, "scale": 0.02 }
+      },
+      "outputs": { "y": { "shape": [1, 64, 32, 32], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fusedconv-f16-1x1-b1c512m512-64x64",
+      "preset": "model",
+      "vars": {
+        "batch": 1,
+        "inChannels": 512,
+        "outChannels": 512,
+        "inH": 64,
+        "inW": 64,
+        "kernelH": 1,
+        "kernelW": 1,
+        "strideH": 1,
+        "strideW": 1,
+        "padH": 0,
+        "padW": 0
+      },
+      "attrs": { "activation": "Relu", "strides": [1, 1], "pads": [0, 0, 0, 0] },
+      "inputs": {
+        "x": { "shape": [1, 512, 64, 64], "dtype": "float16", "dist": "normal", "seed": 354, "scale": 0.2 },
+        "w": { "shape": [512, 512, 1, 1], "dtype": "float16", "dist": "normal", "seed": 355, "scale": 0.02 }
+      },
+      "outputs": { "y": { "shape": [1, 512, 64, 64], "dtype": "float16" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fusedconv-f16-b1c512m512-64x64-k3",
+      "preset": "model",
+      "vars": {
+        "batch": 1,
+        "inChannels": 512,
+        "outChannels": 512,
+        "inH": 64,
+        "inW": 64,
+        "kernelH": 3,
+        "kernelW": 3,
+        "strideH": 1,
+        "strideW": 1,
+        "padH": 1,
+        "padW": 1
+      },
+      "attrs": { "activation": "Relu", "strides": [1, 1], "pads": [1, 1, 1, 1] },
+      "inputs": {
+        "x": { "shape": [1, 512, 64, 64], "dtype": "float16", "dist": "normal", "seed": 356, "scale": 0.2 },
+        "w": { "shape": [512, 512, 3, 3], "dtype": "float16", "dist": "normal", "seed": 357, "scale": 0.02 }
+      },
+      "outputs": { "y": { "shape": [1, 512, 64, 64], "dtype": "float16" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fusedconv-f32-b1c64m128-56x56-k3s2",
+      "preset": "model",
+      "vars": {
+        "batch": 1,
+        "inChannels": 64,
+        "outChannels": 128,
+        "inH": 56,
+        "inW": 56,
+        "kernelH": 3,
+        "kernelW": 3,
+        "strideH": 2,
+        "strideW": 2,
+        "padH": 1,
+        "padW": 1
+      },
+      "attrs": { "activation": "Relu", "strides": [2, 2], "pads": [1, 1, 1, 1] },
+      "inputs": {
+        "x": { "shape": [1, 64, 56, 56], "dtype": "float32", "dist": "normal", "seed": 358, "scale": 0.2 },
+        "w": { "shape": [128, 64, 3, 3], "dtype": "float32", "dist": "normal", "seed": 359, "scale": 0.02 }
+      },
+      "outputs": { "y": { "shape": [1, 128, 28, 28], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fusedconv-f32-group4-b1c64m64-32x32-k3",
+      "preset": "model",
+      "vars": {
+        "batch": 1,
+        "inChannels": 64,
+        "outChannels": 64,
+        "inH": 32,
+        "inW": 32,
+        "kernelH": 3,
+        "kernelW": 3,
+        "strideH": 1,
+        "strideW": 1,
+        "padH": 1,
+        "padW": 1
+      },
+      "attrs": { "group": 4, "activation": "Relu", "strides": [1, 1], "pads": [1, 1, 1, 1] },
+      "inputs": {
+        "x": { "shape": [1, 64, 32, 32], "dtype": "float32", "dist": "normal", "seed": 360, "scale": 0.2 },
+        "w": { "shape": [64, 16, 3, 3], "dtype": "float32", "dist": "normal", "seed": 361, "scale": 0.02 }
+      },
+      "outputs": { "y": { "shape": [1, 64, 32, 32], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fusedconv-f32-stem-3ch-224-k7s2",
+      "preset": "model",
+      "vars": {
+        "batch": 1,
+        "inChannels": 3,
+        "outChannels": 64,
+        "inH": 224,
+        "inW": 224,
+        "kernelH": 7,
+        "kernelW": 7,
+        "strideH": 2,
+        "strideW": 2,
+        "padH": 3,
+        "padW": 3
+      },
+      "attrs": { "activation": "Relu", "strides": [2, 2], "pads": [3, 3, 3, 3] },
+      "inputs": {
+        "x": { "shape": [1, 3, 224, 224], "dtype": "float32", "dist": "normal", "seed": 1458, "scale": 0.2 },
+        "w": { "shape": [64, 3, 7, 7], "dtype": "float32", "dist": "normal", "seed": 1459, "scale": 0.02 }
+      },
+      "outputs": { "y": { "shape": [1, 64, 112, 112], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fusedconv-f32-resnet-stage-256c-14x14-k3",
+      "preset": "model",
+      "vars": {
+        "batch": 1,
+        "inChannels": 256,
+        "outChannels": 256,
+        "inH": 14,
+        "inW": 14,
+        "kernelH": 3,
+        "kernelW": 3,
+        "strideH": 1,
+        "strideW": 1,
+        "padH": 1,
+        "padW": 1
+      },
+      "attrs": { "activation": "Relu", "strides": [1, 1], "pads": [1, 1, 1, 1] },
+      "inputs": {
+        "x": { "shape": [1, 256, 14, 14], "dtype": "float32", "dist": "normal", "seed": 2458, "scale": 0.2 },
+        "w": { "shape": [256, 256, 3, 3], "dtype": "float32", "dist": "normal", "seed": 2459, "scale": 0.02 }
+      },
+      "outputs": { "y": { "shape": [1, 256, 14, 14], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fusedconv-1x1-tinyN-c256m256-8x8",
+      "preset": "smoke",
+      "vars": {
+        "batch": 1,
+        "inChannels": 256,
+        "outChannels": 256,
+        "inH": 8,
+        "inW": 8,
+        "kernelH": 1,
+        "kernelW": 1,
+        "strideH": 1,
+        "strideW": 1,
+        "padH": 0,
+        "padW": 0
+      },
+      "attrs": { "activation": "Relu", "strides": [1, 1], "pads": [0, 0, 0, 0] },
+      "inputs": {
+        "x": { "shape": [1, 256, 8, 8], "dtype": "float32", "dist": "normal", "seed": 951, "scale": 0.2 },
+        "w": { "shape": [256, 256, 1, 1], "dtype": "float32", "dist": "normal", "seed": 1642, "scale": 0.05 }
+      },
+      "outputs": { "y": { "shape": [1, 256, 8, 8], "dtype": "float32" } },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fusedconv-1x1-bigN-c256m256-64x64",
+      "preset": "smoke",
+      "vars": {
+        "batch": 1,
+        "inChannels": 256,
+        "outChannels": 256,
+        "inH": 64,
+        "inW": 64,
+        "kernelH": 1,
+        "kernelW": 1,
+        "strideH": 1,
+        "strideW": 1,
+        "padH": 0,
+        "padW": 0
+      },
+      "attrs": { "activation": "Relu", "strides": [1, 1], "pads": [0, 0, 0, 0] },
+      "inputs": {
+        "x": { "shape": [1, 256, 64, 64], "dtype": "float32", "dist": "normal", "seed": 7277, "scale": 0.2 },
+        "w": { "shape": [256, 256, 1, 1], "dtype": "float32", "dist": "normal", "seed": 1530, "scale": 0.05 }
+      },
+      "outputs": { "y": { "shape": [1, 256, 64, 64], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fusedconv-3x3-unalignedK-c48m64-32x32",
+      "preset": "smoke",
+      "vars": {
+        "batch": 1,
+        "inChannels": 48,
+        "outChannels": 64,
+        "inH": 32,
+        "inW": 32,
+        "kernelH": 3,
+        "kernelW": 3,
+        "strideH": 1,
+        "strideW": 1,
+        "padH": 1,
+        "padW": 1
+      },
+      "attrs": { "activation": "Relu", "strides": [1, 1], "pads": [1, 1, 1, 1] },
+      "inputs": {
+        "x": { "shape": [1, 48, 32, 32], "dtype": "float32", "dist": "normal", "seed": 3152, "scale": 0.2 },
+        "w": { "shape": [64, 48, 3, 3], "dtype": "float32", "dist": "normal", "seed": 7809, "scale": 0.05 }
+      },
+      "outputs": { "y": { "shape": [1, 64, 32, 32], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fusedconv-3x3-alignedK-c64m64-32x32",
+      "preset": "smoke",
+      "vars": {
+        "batch": 1,
+        "inChannels": 64,
+        "outChannels": 64,
+        "inH": 32,
+        "inW": 32,
+        "kernelH": 3,
+        "kernelW": 3,
+        "strideH": 1,
+        "strideW": 1,
+        "padH": 1,
+        "padW": 1
+      },
+      "attrs": { "activation": "Relu", "strides": [1, 1], "pads": [1, 1, 1, 1] },
+      "inputs": {
+        "x": { "shape": [1, 64, 32, 32], "dtype": "float32", "dist": "normal", "seed": 7770, "scale": 0.2 },
+        "w": { "shape": [64, 64, 3, 3], "dtype": "float32", "dist": "normal", "seed": 2676, "scale": 0.05 }
+      },
+      "outputs": { "y": { "shape": [1, 64, 32, 32], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fusedconv-3x3-depthwise-c256-32x32",
+      "preset": "smoke",
+      "vars": {
+        "batch": 1,
+        "inChannels": 256,
+        "outChannels": 256,
+        "inH": 32,
+        "inW": 32,
+        "kernelH": 3,
+        "kernelW": 3,
+        "strideH": 1,
+        "strideW": 1,
+        "padH": 1,
+        "padW": 1
+      },
+      "attrs": { "group": 256, "activation": "Relu", "strides": [1, 1], "pads": [1, 1, 1, 1] },
+      "inputs": {
+        "x": { "shape": [1, 256, 32, 32], "dtype": "float32", "dist": "normal", "seed": 8532, "scale": 0.2 },
+        "w": { "shape": [256, 1, 3, 3], "dtype": "float32", "dist": "normal", "seed": 7192, "scale": 0.05 }
+      },
+      "outputs": { "y": { "shape": [1, 256, 32, 32], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fusedconv-3x3-dense-c256m256-32x32",
+      "preset": "smoke",
+      "vars": {
+        "batch": 1,
+        "inChannels": 256,
+        "outChannels": 256,
+        "inH": 32,
+        "inW": 32,
+        "kernelH": 3,
+        "kernelW": 3,
+        "strideH": 1,
+        "strideW": 1,
+        "padH": 1,
+        "padW": 1
+      },
+      "attrs": { "activation": "Relu", "strides": [1, 1], "pads": [1, 1, 1, 1] },
+      "inputs": {
+        "x": { "shape": [1, 256, 32, 32], "dtype": "float32", "dist": "normal", "seed": 839, "scale": 0.2 },
+        "w": { "shape": [256, 256, 3, 3], "dtype": "float32", "dist": "normal", "seed": 1375, "scale": 0.05 }
+      },
+      "outputs": { "y": { "shape": [1, 256, 32, 32], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fusedconv-3x3-dense-bias-c256m256-32x32",
+      "preset": "smoke",
+      "provenance": {
+        "source": "authored for performance coverage",
+        "notes": "Bias twin of the paired non-bias case at the same 64 register tiles. No bias case sat in the 64-127 tile band, so the bias register-tile floor had no benchmark on either side of it."
+      },
+      "vars": {
+        "batch": 1,
+        "inChannels": 256,
+        "outChannels": 256,
+        "inH": 32,
+        "inW": 32,
+        "kernelH": 3,
+        "kernelW": 3,
+        "strideH": 1,
+        "strideW": 1,
+        "padH": 1,
+        "padW": 1
+      },
+      "attrs": { "activation": "Relu", "strides": [1, 1], "pads": [1, 1, 1, 1] },
+      "inputs": {
+        "x": { "shape": [1, 256, 32, 32], "dtype": "float32", "dist": "normal", "seed": 3531, "scale": 0.2 },
+        "w": { "shape": [256, 256, 3, 3], "dtype": "float32", "dist": "normal", "seed": 3532, "scale": 0.05 },
+        "bias": { "shape": [256], "dtype": "float32", "dist": "normal", "seed": 3533, "scale": 0.05 }
+      },
+      "outputs": { "y": { "shape": [1, 256, 32, 32], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fusedconv-stem-cin3-k7s2-224",
+      "preset": "smoke",
+      "vars": {
+        "batch": 1,
+        "inChannels": 3,
+        "outChannels": 64,
+        "inH": 224,
+        "inW": 224,
+        "kernelH": 7,
+        "kernelW": 7,
+        "strideH": 2,
+        "strideW": 2,
+        "padH": 3,
+        "padW": 3
+      },
+      "attrs": { "activation": "Relu", "strides": [2, 2], "pads": [3, 3, 3, 3] },
+      "inputs": {
+        "x": { "shape": [1, 3, 224, 224], "dtype": "float32", "dist": "normal", "seed": 1633, "scale": 0.2 },
+        "w": { "shape": [64, 3, 7, 7], "dtype": "float32", "dist": "normal", "seed": 3717, "scale": 0.05 }
+      },
+      "outputs": { "y": { "shape": [1, 64, 112, 112], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fusedconv-mid-cin32-k3-112",
+      "preset": "smoke",
+      "vars": {
+        "batch": 1,
+        "inChannels": 32,
+        "outChannels": 64,
+        "inH": 112,
+        "inW": 112,
+        "kernelH": 3,
+        "kernelW": 3,
+        "strideH": 1,
+        "strideW": 1,
+        "padH": 1,
+        "padW": 1
+      },
+      "attrs": { "activation": "Relu", "strides": [1, 1], "pads": [1, 1, 1, 1] },
+      "inputs": {
+        "x": { "shape": [1, 32, 112, 112], "dtype": "float32", "dist": "normal", "seed": 895, "scale": 0.2 },
+        "w": { "shape": [64, 32, 3, 3], "dtype": "float32", "dist": "normal", "seed": 3583, "scale": 0.05 }
+      },
+      "outputs": { "y": { "shape": [1, 64, 112, 112], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fusedconv-1x1-oddM66-c64-32x32",
+      "preset": "smoke",
+      "vars": {
+        "batch": 1,
+        "inChannels": 64,
+        "outChannels": 66,
+        "inH": 32,
+        "inW": 32,
+        "kernelH": 1,
+        "kernelW": 1,
+        "strideH": 1,
+        "strideW": 1,
+        "padH": 0,
+        "padW": 0
+      },
+      "attrs": { "activation": "Relu", "strides": [1, 1], "pads": [0, 0, 0, 0] },
+      "inputs": {
+        "x": { "shape": [1, 64, 32, 32], "dtype": "float32", "dist": "normal", "seed": 7607, "scale": 0.2 },
+        "w": { "shape": [66, 64, 1, 1], "dtype": "float32", "dist": "normal", "seed": 4643, "scale": 0.05 }
+      },
+      "outputs": { "y": { "shape": [1, 66, 32, 32], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fusedconv-1x1-pow2M64-c64-32x32",
+      "preset": "smoke",
+      "vars": {
+        "batch": 1,
+        "inChannels": 64,
+        "outChannels": 64,
+        "inH": 32,
+        "inW": 32,
+        "kernelH": 1,
+        "kernelW": 1,
+        "strideH": 1,
+        "strideW": 1,
+        "padH": 0,
+        "padW": 0
+      },
+      "attrs": { "activation": "Relu", "strides": [1, 1], "pads": [0, 0, 0, 0] },
+      "inputs": {
+        "x": { "shape": [1, 64, 32, 32], "dtype": "float32", "dist": "normal", "seed": 8186, "scale": 0.2 },
+        "w": { "shape": [64, 64, 1, 1], "dtype": "float32", "dist": "normal", "seed": 7554, "scale": 0.05 }
+      },
+      "outputs": { "y": { "shape": [1, 64, 32, 32], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fusedconv-group2-k7-c64m64-64x64-direct-fallback",
+      "provenance": {
+        "notes": "Named for the direct_unrolled route it was authored to expose. It now selects grouped_large_kernel_w4, the wide-output grouped kernel FusedConv gained from ai.onnx.Conv; kept under its original name as the grouped route regression guard."
+      },
+      "preset": "stress",
+      "vars": {
+        "batch": 1,
+        "inChannels": 64,
+        "outChannels": 64,
+        "inH": 64,
+        "inW": 64,
+        "kernelH": 7,
+        "kernelW": 7,
+        "strideH": 1,
+        "strideW": 1,
+        "padH": 3,
+        "padW": 3,
+        "group": 2
+      },
+      "attrs": { "group": 2, "activation": "Relu", "strides": [1, 1], "pads": [3, 3, 3, 3] },
+      "inputs": {
+        "x": { "shape": [1, 64, 64, 64], "dtype": "float32", "dist": "normal", "seed": 4801, "scale": 0.2 },
+        "w": { "shape": [64, 32, 7, 7], "dtype": "float32", "dist": "normal", "seed": 4802, "scale": 0.02 }
+      },
+      "outputs": { "y": { "shape": [1, 64, 64, 64], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fusedconv-group4-k5-c128m128-48x48-direct-fallback",
+      "provenance": {
+        "notes": "Named for the direct_unrolled route it was authored to expose. It now selects grouped_large_kernel_w4, and its 5x5 kernel also sits in the band that kernel refused before its kernel-size floor was measured."
+      },
+      "preset": "stress",
+      "vars": {
+        "batch": 1,
+        "inChannels": 128,
+        "outChannels": 128,
+        "inH": 48,
+        "inW": 48,
+        "kernelH": 5,
+        "kernelW": 5,
+        "strideH": 1,
+        "strideW": 1,
+        "padH": 2,
+        "padW": 2,
+        "group": 4
+      },
+      "attrs": { "group": 4, "activation": "Relu", "strides": [1, 1], "pads": [2, 2, 2, 2] },
+      "inputs": {
+        "x": { "shape": [1, 128, 48, 48], "dtype": "float32", "dist": "normal", "seed": 4811, "scale": 0.2 },
+        "w": { "shape": [128, 32, 5, 5], "dtype": "float32", "dist": "normal", "seed": 4812, "scale": 0.02 }
+      },
+      "outputs": { "y": { "shape": [1, 128, 48, 48], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * dim(shapes.y, 3) * dim(shapes.w, 1) * args.kernelH * args.kernelW"
+          }
+        ]
+      }
+    },
+    {
+      "name": "fusedconv-conv1d-hardswish-b1c80m512-w3000-k3s1",
+      "preset": "stress",
+      "provenance": {
+        "source": "ONNX Runtime FusedConv provider semantics",
+        "notes": "Large aligned temporal convolution with bias and provider-backed HardSwish on the register-tiled implicit-GEMM path."
+      },
+      "vars": { "batch": 1, "inChannels": 80, "outChannels": 512, "inW": 3000, "kernelW": 3, "strideW": 1, "padW": 1 },
+      "attrs": { "activation": "HardSwish", "strides": [1], "pads": [1, 1] },
+      "inputs": {
+        "x": { "shape": [1, 80, 3000], "dtype": "float32", "dist": "normal", "seed": 481, "scale": 0.2 },
+        "w": { "shape": [512, 80, 3], "dtype": "float32", "dist": "normal", "seed": 482, "scale": 0.1 },
+        "bias": { "shape": [512], "dtype": "float32", "dist": "normal", "seed": 483, "scale": 0.05 }
+      },
+      "outputs": { "y": { "shape": [1, 512, 3000], "dtype": "float32" } },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.outChannels * dim(shapes.y, 2) * args.inChannels * args.kernelW"
+          }
+        ]
+      }
+    }
+  ]
+}

build/webgpu/conv-1x1-gemm-tiled-reg.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,419 @@

+// 1x1 / im2col Conv routed as a batched GEMM:
+// Y[b] (M x N) = W (M x K) * X[b] (K x N), with M = outChannels,
+// K = inChannels * kh * kw after im2col, and N = H * W. Register-blocked
+// workgroup lanes compute TM x TN micro-tiles. The defaults use 16x16 lanes with
+// an 8x4 micro-tile over a 64x64 output tile. Weight A is shared across the batch;
+// input B and the output use the dispatch.z batch offset. Accumulation is f32;
+// the optional epilogue applies bias, residual Z, and activation in the
+// accumulator domain on store. Every dimension is bounds-checked.
+// f16 operands remain packed in workgroup memory and widen only at the FMA site;
+// the f32 accumulation order is retained while shared traffic is halved.
+//
+// Both shared tiles are stored as vec4 groups along the axis each lane consumes
+// contiguously: tileA groups four K values for one output row, tileB groups four
+// N columns for one K. One micro-tile step therefore reads TM + 4 vector words
+// instead of 4 * (TM + TN) scalars, so a shared word feeds four times as many
+// FMAs and the K loop runs four accumulation steps per iteration.
+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+{% set hasActivation = hasActivation is defined and hasActivation %}
+{% if hasActivation %}
+// Apply the fused activation in the f32 accumulator before the single output
+// cast, avoiding an intermediate convolution tensor.
+fn fused_act(v: f32) -> f32 {
+{% if activation == "Relu" %}
+  return max(v, 0.0);
+{% elif activation == "Clip" %}
+  return clamp(v, f32({{ actAlpha }}), f32({{ actBeta }}));
+{% elif activation == "LeakyRelu" %}
+  return select(v * f32({{ actAlpha }}), v, v >= 0.0);
+{% elif activation == "Sigmoid" %}
+  return 1.0 / (1.0 + exp(-v));
+{% elif activation == "Tanh" %}
+  return tanh(v);
+{% elif activation == "HardSigmoid" %}
+  return clamp(f32({{ actAlpha }}) * v + f32({{ actBeta }}), 0.0, 1.0);
+{% else %}
+  return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
+{% endif %}
+}
+{% endif %}
+{% set tileT = "f16" if usesF16 else "f32" -%}
+{%- set GEMM_BK = gemmKTile if gemmKTile is defined else 16 %}
+{% set GEMM_TM = gemmThreadRows if gemmThreadRows is defined else 8 %}
+{% set GEMM_WG_X = gemmWorkgroupX if gemmWorkgroupX is defined else 16 %}
+{% set GEMM_BM = gemmMTile if gemmMTile is defined else 64 %}
+{% set GEMM_BN = gemmNTile if gemmNTile is defined else GEMM_WG_X * 4 %}
+{% set GEMM_TN = gemmThreadColumns if gemmThreadColumns is defined else (GEMM_BN / GEMM_WG_X)|int %}
+{% set GEMM_WG_Y = gemmWorkgroupY if gemmWorkgroupY is defined else (GEMM_BM / GEMM_TM)|int %}
+{% set components = ["x", "y", "z", "w"] %}
+const BK: u32 = {{ GEMM_BK }}u;
+const BM: u32 = {{ GEMM_BM }}u;
+const BN: u32 = {{ GEMM_BN }}u;
+const TM: u32 = {{ GEMM_TM }}u;
+const TN: u32 = {{ GEMM_TN }}u;
+const WG_X: u32 = {{ GEMM_WG_X }}u;
+const WG_SIZE: u32 = {{ GEMM_WG_X * GEMM_WG_Y }}u;
+// Vector words per tile row: A groups K, B groups N.
+const AK_VECS: u32 = BK / 4u;
+const BN_VECS: u32 = BN / 4u;
+// implicitIm2col leaves B unmaterialized. While staging a tile, the kernel
+// decodes each (k, n) column-matrix coordinate into a raw NCHW input address.
+// Overlapping windows may reread input values, trading address arithmetic and
+// cache traffic for the storage and bandwidth of the expanded column matrix.
+{% set implicitIm2col = implicitIm2col is defined and implicitIm2col %}
+{% set fusedNarrowProjection = fusedNarrowProjection is defined and fusedNarrowProjection %}
+{% if fusedNarrowProjection %}
+{% set projectionChannels = projectionOutChannels %}
+{% set projectionInputAct = inputActivation %}
+{% set projectionOutputAct = outputActivation %}
+{% set projectionScaled = hasOutputScale %}
+{% set projectionBiased = hasProjectionBias %}
+{% endif %}
+{% set narrowProjectionTile = "tileB" if GEMM_BK >= GEMM_BM else "projectionTile" %}
+{% set splitKValue = splitK if splitK is defined else 1 %}
+{% set splitKPartial = splitKValue > 1 %}
+{% set gemmKLimit = "kEnd" if splitKPartial else "K" %}
+{% set implicitGatherMode = tunables.IMPLICIT_TILED_GATHER_MODE if tunables.IMPLICIT_TILED_GATHER_MODE is defined else 0 %}
+{% set implicitGatherWgSize = GEMM_WG_X * GEMM_WG_Y %}
+{% set implicitGatherColumnVecs = (GEMM_BN / 4)|int %}
+{% set implicitGatherKChunks = (implicitGatherWgSize / implicitGatherColumnVecs)|int %}
+{% set implicitGatherDivides = implicitGatherWgSize % implicitGatherColumnVecs == 0 and implicitGatherKChunks > 0 and GEMM_BK % implicitGatherKChunks == 0 %}
+{% set useCarriedImplicitGather = implicitIm2col and implicitGatherMode == 1 and implicitGatherDivides %}
+{% if useCarriedImplicitGather %}
+{% set implicitGatherLoadWidth = (GEMM_BK / implicitGatherKChunks)|int %}
+{% endif %}
+{% if implicitIm2col %}
+const CONV_KERNEL_H: u32 = {{ convKernelH }}u;
+const CONV_KERNEL_W: u32 = {{ convKernelW }}u;
+const CONV_KSIZE: u32 = CONV_KERNEL_H * CONV_KERNEL_W;
+const CONV_STRIDE_H: u32 = {{ convStrideH }}u;
+const CONV_STRIDE_W: u32 = {{ convStrideW }}u;
+const CONV_DILATION_H: u32 = {{ convDilationH }}u;
+const CONV_DILATION_W: u32 = {{ convDilationW }}u;
+const CONV_PAD_TOP: i32 = {{ convPadTop }};
+const CONV_PAD_LEFT: i32 = {{ convPadLeft }};
+const CONV_IN_H: u32 = {{ convInH }}u;
+const CONV_IN_W: u32 = {{ convInW }}u;
+const CONV_OUT_W: u32 = {{ convOutW }}u;
+const CONV_IN_CHANNELS: u32 = {{ convInChannels }}u;
+{% endif %}
+{% if fusedNarrowProjection %}
+const PROJECTION_OUT_C: u32 = {{ projectionChannels }}u;
+{% else %}{% set emitConvStoreOut = not splitKPartial %}{% if emitConvStoreOut | default(true) %}fn store_out(m: u32, n: u32, yBase: u32, raw: f32) {
+  if (m >= params.M || n >= params.N) {
+    return;
+  }
+  var v = raw;
+{% if hasBias %}
+  v = v + f32(bias[m]);
+{% endif %}
+{% if hasZ %}
+  v = v + f32(zResidual[yBase + m * params.N + n]);
+{% endif %}
+{% if hasActivation %}
+  v = fused_act(v);
+{% endif %}
+  y[yBase + m * params.N + n] = {{ T }}(v);
+}
+{% endif %}
+{% endif %}
+var<workgroup> tileA: array<array<vec4<{{ tileT }}>, AK_VECS>, BM>;
+var<workgroup> tileB: array<array<vec4<{{ tileT }}>, BN_VECS>, BK>;
+{% if fusedNarrowProjection and GEMM_BK < GEMM_BM %}
+// BK16 cannot reuse the 16-row input tile to publish a BM32 intermediate.
+// The extra 32x64 f32 tile keeps the total at 14 KiB, below WebGPU's
+// guaranteed 16 KiB workgroup-storage floor.
+var<workgroup> projectionTile: array<array<vec4<f32>, BN_VECS>, BM>;
+{% endif %}
+{% macro publish_projection_input(localRow, localVec, component, globalRow, globalColumn, raw) %}
+  if ({{ globalRow }} < M && {{ globalColumn }} < N) {
+    var projectionInput = {{ raw }};
+{% if hasBias %}
+    projectionInput = projectionInput + f32(bias[{{ globalRow }}]);
+{% endif %}
+{% if projectionInputAct == "relu" %}
+    projectionInput = projection_relu(projectionInput);
+{% endif %}
+    {{ narrowProjectionTile }}[{{ localRow }}][{{ localVec }}].{{ component }} = projectionInput;
+  } else {
+    {{ narrowProjectionTile }}[{{ localRow }}][{{ localVec }}].{{ component }} = 0.0;
+  }
+{%- endmacro %}
+{% macro load_a_vec4(rowBase, columnBase, rowLimit, columnLimit) %}
+    for (var linear = li; linear < BM * AK_VECS; linear += WG_SIZE) {
+      let ar = linear / AK_VECS;
+      let ac4 = linear % AK_VECS;
+      let am = {{ rowBase }} + ar;
+      let ak = {{ columnBase }} + ac4 * 4u;
+      var av = vec4<{{ tileT }}>({{ tileT }}(0.0));
+      if (am < {{ rowLimit }}) {
+        let aBase = am * K + ak;
+        if (ak + 3u < {{ columnLimit }}) {
+          av = vec4<{{ tileT }}>({{ tileT }}(w[aBase]), {{ tileT }}(w[aBase + 1u]), {{ tileT }}(w[aBase + 2u]), {{ tileT }}(w[aBase + 3u]));
+        } else {
+          for (var t = 0u; t < 4u; t = t + 1u) {
+            if (ak + t < {{ columnLimit }}) {
+              av[t] = {{ tileT }}(w[aBase + t]);
+            }
+          }
+        }
+      }
+      tileA[ar][ac4] = av;
+    }
+{%- endmacro %}
+{% macro load_b_vec4(rowBase, columnBase, rowLimit, columnLimit) %}
+    for (var linear = li; linear < BK * BN_VECS; linear += WG_SIZE) {
+      let br = linear / BN_VECS;
+      let bc4 = linear % BN_VECS;
+      let bk = {{ rowBase }} + br;
+      let bn = {{ columnBase }} + bc4 * 4u;
+      var bvec = vec4<{{ tileT }}>({{ tileT }}(0.0));
+      if (bk < {{ rowLimit }}) {
+        let bBase = xBatchBase + bk * N + bn;
+        if (bn + 3u < {{ columnLimit }}) {
+          bvec = vec4<{{ tileT }}>({{ tileT }}(xm[bBase]), {{ tileT }}(xm[bBase + 1u]), {{ tileT }}(xm[bBase + 2u]), {{ tileT }}(xm[bBase + 3u]));
+        } else {
+          for (var t = 0u; t < 4u; t = t + 1u) {
+            if (bn + t < {{ columnLimit }}) {
+              bvec[t] = {{ tileT }}(xm[bBase + t]);
+            }
+          }
+        }
+      }
+      tileB[br][bc4] = bvec;
+    }
+{%- endmacro %}
+{% macro load_b_implicit_vec4(rowBase, columnBase, rowLimit, columnLimit) %}
+    for (var linear = li; linear < BK * BN_VECS; linear += WG_SIZE) {
+      let br = linear / BN_VECS;
+      let bc4 = linear % BN_VECS;
+      let bk = {{ rowBase }} + br;
+      var bvec = vec4<{{ tileT }}>({{ tileT }}(0.0));
+      if (bk < {{ rowLimit }}) {
+        // k is the OIHW weight flattening (ic * KH + kh) * KW + kw; n is the output position.
+        let ic = bk / CONV_KSIZE;
+        let kq = bk % CONV_KSIZE;
+        let kh = kq / CONV_KERNEL_W;
+        let kw = kq % CONV_KERNEL_W;
+        let planeBase = (batch * CONV_IN_CHANNELS + ic) * CONV_IN_H;
+        for (var t = 0u; t < 4u; t = t + 1u) {
+          let bn = {{ columnBase }} + bc4 * 4u + t;
+          if (bn < {{ columnLimit }}) {
+            let oh = bn / CONV_OUT_W;
+            let ow = bn % CONV_OUT_W;
+            let ih = i32(oh * CONV_STRIDE_H + kh * CONV_DILATION_H) - CONV_PAD_TOP;
+            let iw = i32(ow * CONV_STRIDE_W + kw * CONV_DILATION_W) - CONV_PAD_LEFT;
+            if (ih >= 0 && ih < i32(CONV_IN_H) && iw >= 0 && iw < i32(CONV_IN_W)) {
+              bvec[t] = {{ tileT }}(xm[(planeBase + u32(ih)) * CONV_IN_W + u32(iw)]);
+            }
+          }
+        }
+      }
+      tileB[br][bc4] = bvec;
+    }
+{%- endmacro %}
+{% macro load_b_implicit_carried(rowBase, columnBase, rowLimit, columnLimit) %}
+    let bColVec = li / {{ implicitGatherKChunks }}u;
+    let bChunk = li % {{ implicitGatherKChunks }}u;
+    let bLocalK = bChunk * {{ implicitGatherLoadWidth }}u;
+    let bStartK = {{ rowBase }} + bLocalK;
+    let bCol0 = {{ columnBase }} + bColVec * 4u;
+    let bCols = vec4<u32>(bCol0, bCol0 + 1u, bCol0 + 2u, bCol0 + 3u);
+    let inColumn = bCols < vec4<u32>({{ columnLimit }});
+    let ih0 = vec4<i32>((bCols / vec4<u32>(CONV_OUT_W)) * CONV_STRIDE_H) - vec4<i32>(CONV_PAD_TOP);
+    let iw0 = vec4<i32>((bCols % vec4<u32>(CONV_OUT_W)) * CONV_STRIDE_W) - vec4<i32>(CONV_PAD_LEFT);
+    var carriedK = bStartK;
+    var carriedIc = carriedK / CONV_KSIZE;
+    let carriedKq = carriedK % CONV_KSIZE;
+    var carriedKh = carriedKq / CONV_KERNEL_W;
+    var carriedKw = carriedKq % CONV_KERNEL_W;
+    var carriedIh = ih0 + vec4<i32>(i32(carriedKh * CONV_DILATION_H));
+    var carriedIw = iw0 + vec4<i32>(i32(carriedKw * CONV_DILATION_W));
+    // Interior words can carry four addresses. Padding, the N tail, and the
+    // final K tile retain the coordinate-checked loop below and publish exact
+    // zeros.
+    let interior = all(inColumn)
+      && bStartK + {{ implicitGatherLoadWidth }}u <= {{ rowLimit }}
+      && all(ih0 >= vec4<i32>(0))
+      && all(ih0 + vec4<i32>(i32((CONV_KERNEL_H - 1u) * CONV_DILATION_H)) < vec4<i32>(i32(CONV_IN_H)))
+      && all(iw0 >= vec4<i32>(0))
+      && all(iw0 + vec4<i32>(i32((CONV_KERNEL_W - 1u) * CONV_DILATION_W)) < vec4<i32>(i32(CONV_IN_W)));
+    if (interior) {
+      let batchBase = batch * CONV_IN_CHANNELS * CONV_IN_H * CONV_IN_W;
+      let colStep = i32(CONV_DILATION_W);
+      let rowStep = i32(CONV_DILATION_H * CONV_IN_W);
+      let planeStep = i32(CONV_IN_H * CONV_IN_W);
+      let kwWrap = colStep * i32(CONV_KERNEL_W) - rowStep;
+      let khWrap = rowStep * i32(CONV_KERNEL_H) - planeStep;
+      var carriedAddress = vec4<i32>(i32(carriedIc) * planeStep) + carriedIh * vec4<i32>(i32(CONV_IN_W)) + carriedIw;
+      for (var i = 0u; i < {{ implicitGatherLoadWidth }}u; i = i + 1u) {
+        tileB[bLocalK + i][bColVec] = vec4<{{ tileT }}>(
+{% for component in components %}
+          {{ tileT }}(xm[batchBase + u32(carriedAddress.{{ component }})]){% if not loop.last %},{% endif %}
+{% endfor %}
+        );
+        carriedAddress = carriedAddress + vec4<i32>(colStep);
+        carriedKw = carriedKw + 1u;
+        if (carriedKw == CONV_KERNEL_W) {
+          carriedKw = 0u;
+          carriedAddress = carriedAddress - vec4<i32>(kwWrap);
+          carriedKh = carriedKh + 1u;
+          if (carriedKh == CONV_KERNEL_H) {
+            carriedKh = 0u;
+            carriedAddress = carriedAddress - vec4<i32>(khWrap);
+          }
+        }
+      }
+    } else {
+      for (var i = 0u; i < {{ implicitGatherLoadWidth }}u; i = i + 1u) {
+        var carriedValue = vec4<{{ tileT }}>({{ tileT }}(0.0));
+        let carriedRow = (batch * CONV_IN_CHANNELS + carriedIc) * CONV_IN_H;
+        let carriedOk = inColumn
+          & (carriedIh >= vec4<i32>(0)) & (carriedIh < vec4<i32>(i32(CONV_IN_H)))
+          & (carriedIw >= vec4<i32>(0)) & (carriedIw < vec4<i32>(i32(CONV_IN_W)));
+        if (carriedK < {{ rowLimit }}) {
+{% for component in components %}
+          if (carriedOk.{{ component }}) {
+            carriedValue.{{ component }} = {{ tileT }}(xm[(carriedRow + u32(carriedIh.{{ component }})) * CONV_IN_W + u32(carriedIw.{{ component }})]);
+          }
+{% endfor %}
+        }
+        tileB[bLocalK + i][bColVec] = carriedValue;
+        carriedK = carriedK + 1u;
+        carriedKw = carriedKw + 1u;
+        carriedIw = carriedIw + vec4<i32>(i32(CONV_DILATION_W));
+        if (carriedKw == CONV_KERNEL_W) {
+          carriedKw = 0u;
+          carriedIw = iw0;
+          carriedKh = carriedKh + 1u;
+          carriedIh = carriedIh + vec4<i32>(i32(CONV_DILATION_H));
+          if (carriedKh == CONV_KERNEL_H) {
+            carriedKh = 0u;
+            carriedIh = ih0;
+            carriedIc = carriedIc + 1u;
+          }
+        }
+      }
+    }
+{%- endmacro %}
+@compute @workgroup_size({{ GEMM_WG_X }}, {{ GEMM_WG_Y }}, 1)
+fn main(
+  @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>
+) {
+  let M = params.M;
+  let K = params.K;
+  let N = params.N;
+  let mBase = wg.y * BM;
+  let nBase = wg.x * BN;
+  let batch = wg.z;
+{% if not implicitIm2col %}
+  let xBatchBase = batch * K * N;
+{% endif %}
+  let li = lid.y * WG_X + lid.x;
+  // Explicit vectors give the backend fixed register indices and expose TM
+  // independent FMA chains.
+{% for row in range(GEMM_TM) %}
+  var acc{{ row }} = vec{{ GEMM_TN }}<f32>(0.0);
+{% endfor %}
+  let numTiles = (K + BK - 1u) / BK;
+  for (var kt: u32 = 0u; kt < numTiles; kt = kt + 1u) {
+    let kBase = kt * BK;
+{{ load_a_vec4("mBase", "kBase", "M", gemmKLimit) }}
+{% if useCarriedImplicitGather %}
+{{ load_b_implicit_carried("kBase", "nBase", gemmKLimit, "N") }}
+{% elif implicitIm2col %}
+{{ load_b_implicit_vec4("kBase", "nBase", gemmKLimit, "N") }}
+{% else %}
+{{ load_b_vec4("kBase", "nBase", gemmKLimit, "N") }}
+{% endif %}
+    workgroupBarrier();
+    // One iteration consumes four K steps: TM vector words of A and four of B
+    // cover TM x TN x 4 FMAs.
+    let aRow = lid.y * TM;
+    for (var kv: u32 = 0u; kv < AK_VECS; kv = kv + 1u) {
+{% for row in range(GEMM_TM) %}
+      let av{{ row }} = vec4<f32>(tileA[aRow + {{ row }}u][kv]);
+{% endfor %}
+{% for sub in range(4) %}
+      let bv{{ sub }} = vec{{ GEMM_TN }}<f32>(tileB[kv * 4u + {{ sub }}u][lid.x]);
+{% for row in range(GEMM_TM) %}
+      acc{{ row }} += vec{{ GEMM_TN }}<f32>(av{{ row }}.{{ components[sub] }}) * bv{{ sub }};
+{% endfor %}
+{% endfor %}
+    }
+    workgroupBarrier();
+  }
+{% if not fusedNarrowProjection %}
+  let yBatchBase = batch * M * N;
+{% endif %}
+  let m0 = mBase + lid.y * TM;
+  let n0 = nBase + lid.x * TN;
+{% for row in range(GEMM_TM) %}
+{% for column in range(GEMM_TN) %}
+{% if fusedNarrowProjection %}
+{{ publish_projection_input("lid.y * TM + " ~ row ~ "u", "lid.x * " ~ ((GEMM_TN / 4)|int) ~ "u + " ~ ((column / 4)|int) ~ "u", components[column % 4], "m0 + " ~ row ~ "u", "n0 + " ~ column ~ "u", "acc" ~ row ~ "." ~ components[column]) }}
+{% else %}
+  store_out(m0 + {{ row }}u, n0 + {{ column }}u, yBatchBase, acc{{ row }}.{{ components[column] }});
+{% endif %}
+{% endfor %}
+{% endfor %}{% if fusedNarrowProjection %}
+  // Conv's final K-tile barrier makes tileB dead before it becomes the
+  // intermediate tile. Every lane publishes its unique micro-tile, then one
+  // lane per spatial vector word consumes all logical M rows in increasing
+  // order, four adjacent columns at a time. Invalid M/N-tail cells are
+  // initialized above, so the barrier is uniform and no lane can observe a
+  // stale input-tile value.
+  workgroupBarrier();
+  if (li < BN_VECS) {
+    let projectionNBase = nBase + li * 4u;
+{% for oc in range(projectionChannels) %}
+    var projectionAcc{{ oc }} = vec4<f32>(0.0);
+{% endfor %}
+    for (var channel = 0u; channel < M; channel = channel + 1u) {
+      let projectionValues = vec4<f32>({{ narrowProjectionTile }}[channel][li]);
+{% for oc in range(projectionChannels) %}
+      projectionAcc{{ oc }} = projectionAcc{{ oc }} + projectionValues * f32(projectionW[{{ oc }}u * M + channel]);
+{% endfor %}
+    }
+    let projectionYBase = batch * PROJECTION_OUT_C * N;
+{% for oc in range(projectionChannels) %}
+    let projected{{ oc }} = projectionAcc{{ oc }}{% if projectionBiased %} + vec4<f32>(f32(projectionBias[{{ oc }}u])){% endif %};
+{% endfor %}
+{% for column in range(4) %}
+    let projectionN{{ column }} = projectionNBase + {{ column }}u;
+    if (projectionN{{ column }} < N) {
+{% for oc in range(projectionChannels) %}
+{% if projectionOutputAct == "relu" %}
+      let activated{{ oc }}_{{ column }} = projection_relu(projected{{ oc }}.{{ components[column] }});
+{% elif projectionOutputAct == "sigmoid" %}
+      let activated{{ oc }}_{{ column }} = sigmoid_safe(projected{{ oc }}.{{ components[column] }});
+{% else %}
+      let activated{{ oc }}_{{ column }} = projected{{ oc }}.{{ components[column] }};
+{% endif %}
+{% if projectionScaled %}
+      y[projectionYBase + {{ oc }}u * N + projectionN{{ column }}] = activated{{ oc }}_{{ column }} * params.outputScale;
+{% else %}
+      y[projectionYBase + {{ oc }}u * N + projectionN{{ column }}] = activated{{ oc }}_{{ column }};
+{% endif %}
+{% endfor %}
+    }
+{% endfor %}
+  }
+{% endif %}
+}

build/webgpu/conv-1x1-gemm-tiled.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,138 @@

+// 1x1 Conv routed as a batched GEMM: Y[b] (M x N) = W (M x K) * X[b] (K x N)
+// with M = outChannels, K = inChannels, N = H*W. For a 1x1 kernel with
+// group == 1, stride 1, no padding, and dilation 1, the NCHW input is already
+// the [K, N] matrix and the OIHW weight is already the [M, K] matrix — no
+// data movement, only index arithmetic.
+//
+// Register-blocked 32x32 shared-memory tile (16x16 threads, each computes a 2x2
+// micro-tile, with a selectable K tile (16 by default), the same structure as
+// the general tiled matmul path — each shared element feeds 2 FMAs
+// and a 32-wide column tile is staged once instead of re-reading W per output
+// column. The weight matrix (A) is shared across the batch (no batch stride);
+// the input (B) and output carry the batch offset via dispatch.z. f32
+// accumulation, cast to the output element type on store; the optional bias is
+// per output channel (M row), added after the full K accumulation per ONNX Conv
+// semantics. f16 operands stay packed in workgroup memory and widen only when
+// consumed. Fully bounds-checked (any M/N/K).
+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+{% set hasActivation = hasActivation is defined and hasActivation %}
+{% if hasActivation %}
+// Apply the fused activation in the f32 accumulator before the single output
+// cast, avoiding an intermediate convolution tensor.
+fn fused_act(v: f32) -> f32 {
+{% if activation == "Relu" %}
+  return max(v, 0.0);
+{% elif activation == "Clip" %}
+  return clamp(v, f32({{ actAlpha }}), f32({{ actBeta }}));
+{% elif activation == "LeakyRelu" %}
+  return select(v * f32({{ actAlpha }}), v, v >= 0.0);
+{% elif activation == "Sigmoid" %}
+  return 1.0 / (1.0 + exp(-v));
+{% elif activation == "Tanh" %}
+  return tanh(v);
+{% elif activation == "HardSigmoid" %}
+  return clamp(f32({{ actAlpha }}) * v + f32({{ actBeta }}), 0.0, 1.0);
+{% else %}
+  return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
+{% endif %}
+}
+{% endif %}
+{% set tileT = "f16" if usesF16 else "f32" -%}
+const BK: u32 = 16u;
+const BM: u32 = 32u;
+const BN: u32 = 32u;
+// Store one output element with the optional bias/residual/activation epilogue in the
+// f32 accumulator domain: Y = activation(conv + bias + Z).
+fn store_out(m: u32, n: u32, yBase: u32, raw: f32) {
+  if (m >= params.M || n >= params.N) {
+    return;
+  }
+  var v = raw;
+{% if hasBias %}
+  v = v + f32(bias[m]);
+{% endif %}
+{% if hasZ %}
+  v = v + f32(zResidual[yBase + m * params.N + n]);
+{% endif %}
+{% if hasActivation %}
+  v = fused_act(v);
+{% endif %}
+  y[yBase + m * params.N + n] = {{ T }}(v);
+}
+var<workgroup> tileA: array<array<{{ tileT }}, BK>, BM>;
+var<workgroup> tileB: array<array<{{ tileT }}, BN>, BK>;
+@compute @workgroup_size(16, 16, 1)
+fn main(
+  @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>
+) {
+  let M = params.M;
+  let K = params.K;
+  let N = params.N;
+  let mBase = wg.y * BM;
+  let nBase = wg.x * BN;
+  let batch = wg.z;
+  let xBatchBase = batch * K * N;
+  let li = lid.y * 16u + lid.x;
+  var acc00: f32 = 0.0;
+  var acc01: f32 = 0.0;
+  var acc10: f32 = 0.0;
+  var acc11: f32 = 0.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(=W) tile + 16x32 B(=X[batch]) tile, 256 threads x 2 each.
+    for (var e: u32 = 0u; e < (BM * BK) / 256u; 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 }}(w[am * K + ak]);
+      } else {
+        tileA[ar][ac] = {{ tileT }}(0.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 }}(xm[xBatchBase + bk * N + bn]);
+      } else {
+        tileB[br][bc] = {{ tileT }}(0.0);
+      }
+    }
+    workgroupBarrier();
+    for (var kk: u32 = 0u; kk < BK; kk = kk + 1u) {
+      let a0 = f32(tileA[lid.y * 2u][kk]);
+      let a1 = f32(tileA[lid.y * 2u + 1u][kk]);
+      let b0 = f32(tileB[kk][lid.x * 2u]);
+      let b1 = f32(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 yBatchBase = batch * M * N;
+  store_out(m0, n0, yBatchBase, acc00);
+  store_out(m0, n1, yBatchBase, acc01);
+  store_out(m1, n0, yBatchBase, acc10);
+  store_out(m1, n1, yBatchBase, acc11);
+}

build/webgpu/conv-1x1-subgroup-matrix.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,353 @@

+// Convolution as subgroup-matrix GEMM. OIHW weights form row-major [M,K]; `xm`
+// is either an NCHW 1x1 view [K,N] or a materialized im2col matrix. Bias is
+// applied per output channel after the full f32 accumulation.
+//
+// Full 8x8 matrix loads require complete K/N tiles. The aligned path enforces
+// that geometry; `padded` zero-fills tails, and `implicitIm2col` gathers and
+// zero-fills directly from NCHW without materializing the column matrix.
+//
+// `splitK` assigns disjoint K ranges to workgroup_id.z. Each slice writes raw
+// sums to tile-padded [splitK,batch,M_PAD,N_PAD] scratch; a later pass combines
+// the slices and applies bias or an epilogue. Padded tail cells remain zero and
+// are never copied to the logical output.
+{% set fusedNarrowProjection = fusedNarrowProjection if fusedNarrowProjection is defined else false %}
+{% set polyphase = polyphaseConvTranspose is defined and polyphaseConvTranspose %}
+{% set splitKValue = splitK if splitK is defined else 1 %}
+{% set splitKPartial = splitKValue > 1 %}
+{% set kLoopVar = "K_LOOP" if padded else "K" %}
+{% set nColsVar = "N_COLS" if padded else "N" %}
+{% 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 useDirectMatrixStore = directMatrixStore is defined and directMatrixStore %}
+{% set tileRowsValue = tileRows if tileRows is defined else 32 %}
+{% set tileColsValue = tileCols if tileCols is defined else 64 %}
+{% set workgroupThreadsValue = workgroupThreads if workgroupThreads is defined else 128 %}
+{% set subgroupRowsValue = subgroupRows if subgroupRows is defined else 2 %}
+{% set subgroupColsValue = subgroupCols if subgroupCols is defined else 2 %}
+{% set subRowsValue = (tileRowsValue / subgroupRowsValue)|int %}
+{% set subColsValue = (tileColsValue / subgroupColsValue)|int %}
+{% set bLoadWidth = (tileColsValue * 32 / workgroupThreadsValue)|int %}
+{% set bKChunks = (32 / bLoadWidth)|int %}
+const M: u32 = {{ M }}u;
+const K: u32 = {{ K }}u;
+const N: u32 = {{ N }}u;
+{% if padded %}
+const K_LOOP: u32 = {{ kPadded }}u;
+{% endif %}
+{% if padded %}
+const N_COLS: u32 = {{ nPadded }}u;
+{% endif %}
+const B_BATCH_STRIDE: u32 = {{ kLoopVar }} * {{ nColsVar }};
+{% if not splitKPartial and not polyphase and not fusedNarrowProjection %}
+const C_BATCH_STRIDE: u32 = M * N;
+{% endif %}
+const TILE_COLS: u32 = {{ tileColsValue }}u;
+const TILE_ROWS: u32 = {{ tileRowsValue }}u;
+const TILE_K: u32 = 32u;
+const SUB_COLS: u32 = {{ subColsValue }}u;
+const SUB_ROWS: u32 = {{ subRowsValue }}u;
+var<workgroup> tile_A: array<{{ operandScalar }}, {{ tileRowsValue }} * 32>;
+var<workgroup> tile_B: array<{{ operandScalar }}, {{ tileColsValue }} * 32>;
+{% if (not useDirectMatrixStore or padded) and not splitKPartial and not fusedNarrowProjection %}
+var<workgroup> scratch: array<array<array<{{ accScalar }}, 64>, 4>, {{ (workgroupThreadsValue / 32)|int }}>;
+{% endif %}
+fn loadSHMA(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 padded %} && k < K{% endif %}) {
+{% set A_PHASE = "phase_base + " if polyphase else "" %}
+{% if operandScalar == "f16" %}
+      tile_A[row * TILE_K + col + col_offset] = f16(w[{{ A_PHASE }}a_global * K + k]);
+{% else %}
+      tile_A[row * TILE_K + col + col_offset] = f32(w[{{ A_PHASE }}a_global * K + k]);
+{% 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 %}
+    }
+  }
+}
+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 * {{ bLoadWidth }}u;
+  for (var i = 0u; i < {{ bLoadWidth }}u; i = i + 1u) {
+    let k = k_idx + col + i;
+{% if operandScalar == "f16" %}
+    tile_B[row * TILE_K + col + i] = f16(xm[b_base + k * {{ nColsVar }} + b_col]);
+{% else %}
+    tile_B[row * TILE_K + col + i] = f32(xm[b_base + k * {{ nColsVar }} + b_col]);
+{% endif %}
+  }
+}
+{% set hasActivation = hasActivation is defined and hasActivation %}
+{% if hasActivation %}
+// Apply the fused activation in the f32 accumulator before the single output
+// cast, avoiding an intermediate convolution tensor.
+fn fused_act(v: f32) -> f32 {
+{% if activation == "Relu" %}
+  return max(v, 0.0);
+{% elif activation == "Clip" %}
+  return clamp(v, f32({{ actAlpha }}), f32({{ actBeta }}));
+{% elif activation == "LeakyRelu" %}
+  return select(v * f32({{ actAlpha }}), v, v >= 0.0);
+{% elif activation == "Sigmoid" %}
+  return 1.0 / (1.0 + exp(-v));
+{% elif activation == "Tanh" %}
+  return tanh(v);
+{% elif activation == "HardSigmoid" %}
+  return clamp(f32({{ actAlpha }}) * v + f32({{ actBeta }}), 0.0, 1.0);
+{% else %}
+  return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
+{% endif %}
+}
+{% endif %}
+{% set biasAdd = " + bv" if hasBias else "" %}
+{% if hasActivation or hasZ %}
+// Fused epilogue: Y = activation(conv + bias + Z), applied at the output store.
+// `raw` is the f32 convolution-plus-bias accumulator; Z shares Y's NCHW layout.
+fn epi(raw: f32{% if hasZ %}, yIndex: u32{% endif %}) -> {{ T }} {
+  var r = raw;
+{% if hasZ %}
+  r = r + f32(zResidual[yIndex]);
+{% endif %}
+{% if hasActivation %}
+  r = fused_act(r);
+{% endif %}
+  return {{ T }}(r);
+}
+{% endif %}
+{% macro store_val(valExpr, idxExpr) %}
+{% if hasActivation or hasZ %}epi({{ valExpr }}{% if hasZ %}, {{ idxExpr }}{% endif %})
+{%- else %}{{ T }}({{ valExpr }})
+{%- endif %}
+{% endmacro %}
+{% if (not useDirectMatrixStore or padded) and not splitKPartial and not fusedNarrowProjection %}
+fn storeOutput(offset: u32, {% if hasBias or polyphase %}row_base: u32, {% endif %}row: u32, col: u32, src_slot: u32, row_limit: i32{% if padded %}, col_anchor: u32{% endif %}) {
+  if (row_limit > 0 && row < u32(row_limit)) {
+    let col2 = col + 1u;
+{% if hasBias %}
+    let bv = {{ accScalar }}(bias[row_base + row]);
+{% endif %}
+{% if padded %}
+    // Padded GEMM: the last column tile runs past the logical N; only write
+    // columns inside the real output. Each subgroup-matrix column maps to a
+    // distinct y column (anchor + local), so the guards are independent.
+    if (col_anchor + col +  0u < N) { y[offset + row * N + col]        = {{ store_val("scratch[src_slot][0][row * 8u + col]" ~ biasAdd, "offset + row * N + col") }}; }
+    if (col_anchor + col +  8u < N) { y[offset + row * N + col +  8u]  = {{ store_val("scratch[src_slot][1][row * 8u + col]" ~ biasAdd, "offset + row * N + col + 8u") }}; }
+    if (col_anchor + col + 16u < N) { y[offset + row * N + col + 16u]  = {{ store_val("scratch[src_slot][2][row * 8u + col]" ~ biasAdd, "offset + row * N + col + 16u") }}; }
+    if (col_anchor + col + 24u < N) { y[offset + row * N + col + 24u]  = {{ store_val("scratch[src_slot][3][row * 8u + col]" ~ biasAdd, "offset + row * N + col + 24u") }}; }
+    if (col_anchor + col2 +  0u < N) { y[offset + row * N + col2]       = {{ store_val("scratch[src_slot][0][row * 8u + col2]" ~ biasAdd, "offset + row * N + col2") }}; }
+    if (col_anchor + col2 +  8u < N) { y[offset + row * N + col2 +  8u] = {{ store_val("scratch[src_slot][1][row * 8u + col2]" ~ biasAdd, "offset + row * N + col2 + 8u") }}; }
+    if (col_anchor + col2 + 16u < N) { y[offset + row * N + col2 + 16u] = {{ store_val("scratch[src_slot][2][row * 8u + col2]" ~ biasAdd, "offset + row * N + col2 + 16u") }}; }
+    if (col_anchor + col2 + 24u < N) { y[offset + row * N + col2 + 24u] = {{ store_val("scratch[src_slot][3][row * 8u + col2]" ~ biasAdd, "offset + row * N + col2 + 24u") }}; }
+{% else %}
+    y[offset + row * N + col]       = {{ store_val("scratch[src_slot][0][row * 8u + col]" ~ biasAdd, "offset + row * N + col") }};
+    y[offset + row * N + col + 8u]  = {{ store_val("scratch[src_slot][1][row * 8u + col]" ~ biasAdd, "offset + row * N + col + 8u") }};
+    y[offset + row * N + col + 16u] = {{ store_val("scratch[src_slot][2][row * 8u + col]" ~ biasAdd, "offset + row * N + col + 16u") }};
+    y[offset + row * N + col + 24u] = {{ store_val("scratch[src_slot][3][row * 8u + col]" ~ biasAdd, "offset + row * N + col + 24u") }};
+    y[offset + row * N + col2]       = {{ store_val("scratch[src_slot][0][row * 8u + col2]" ~ biasAdd, "offset + row * N + col2") }};
+    y[offset + row * N + col2 + 8u]  = {{ store_val("scratch[src_slot][1][row * 8u + col2]" ~ biasAdd, "offset + row * N + col2 + 8u") }};
+    y[offset + row * N + col2 + 16u] = {{ store_val("scratch[src_slot][2][row * 8u + col2]" ~ biasAdd, "offset + row * N + col2 + 16u") }};
+    y[offset + row * N + col2 + 24u] = {{ store_val("scratch[src_slot][3][row * 8u + col2]" ~ biasAdd, "offset + row * N + col2 + 24u") }};
+{% endif %}
+  }
+}
+{% endif %}
+@compute @workgroup_size({{ workgroupThreadsValue }}, 1, 1){{ " @subgroup_size(32)" if pinSubgroupSize32 else "" }}
+fn main(
+  @builtin(workgroup_id) workgroup_id: vec3<u32>,
+  @builtin(local_invocation_index) local_idx: u32,
+{% if (not useDirectMatrixStore or padded) and not splitKPartial and not fusedNarrowProjection %}
+  @builtin(subgroup_invocation_id) sg_id: u32,
+{% endif %}
+  @builtin(subgroup_size) sg_size: u32
+) {
+  let batch = workgroup_id.z;
+  let b_base = batch * B_BATCH_STRIDE;
+{% if not fusedNarrowProjection %}
+  let c_base = batch * C_BATCH_STRIDE;
+{% endif %}
+  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 / {{ subgroupRowsValue }}u;
+  let subtile_idy = subtile_id % {{ subgroupRowsValue }}u;
+  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 < {{ kLoopVar }}; kidx = kidx + TILE_K) {
+    loadSHMA(a_global_base, kidx, local_idx / 4u, local_idx % 4u);
+    loadSHMB(b_base, b_global_base, kidx, local_idx / {{ bKChunks }}u, local_idx % {{ bKChunks }}u);
+    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 fusedNarrowProjection %}
+  // The Conv producer owns exactly one 32-row tile. Its input tile is dead after
+  // the final K iteration, so reuse those 32x64 f32 cells as the only
+  // cross-subgroup handoff into the narrow projection. All collective stores
+  // remain subgroup-uniform, and every fragment has a disjoint destination.
+  let fused_tile_offset = base_A * TILE_COLS + base_B;
+  subgroupMatrixStore(&tile_B, fused_tile_offset + 0u * TILE_COLS +  0u, matC00, false, TILE_COLS);
+  subgroupMatrixStore(&tile_B, fused_tile_offset + 0u * TILE_COLS +  8u, matC01, false, TILE_COLS);
+  subgroupMatrixStore(&tile_B, fused_tile_offset + 0u * TILE_COLS + 16u, matC02, false, TILE_COLS);
+  subgroupMatrixStore(&tile_B, fused_tile_offset + 0u * TILE_COLS + 24u, matC03, false, TILE_COLS);
+  subgroupMatrixStore(&tile_B, fused_tile_offset + 8u * TILE_COLS +  0u, matC10, false, TILE_COLS);
+  subgroupMatrixStore(&tile_B, fused_tile_offset + 8u * TILE_COLS +  8u, matC11, false, TILE_COLS);
+  subgroupMatrixStore(&tile_B, fused_tile_offset + 8u * TILE_COLS + 16u, matC12, false, TILE_COLS);
+  subgroupMatrixStore(&tile_B, fused_tile_offset + 8u * TILE_COLS + 24u, matC13, false, TILE_COLS);
+  workgroupBarrier();
+  // One invocation owns one spatial column and walks the producer channels in
+  // increasing order to preserve the scalar projection's accumulation order.
+  // The inactive half of the workgroup has no remaining collective operation
+  // to reach.
+  let fused_global_col = b_global_base + local_idx;
+  if (local_idx < TILE_COLS && fused_global_col < N) {
+{% for oc in range(projectionOutChannels) %}
+    var projected_acc{{ oc }} = f32(0.0);
+{% endfor %}
+    for (var channel = 0u; channel < M; channel = channel + 1u) {
+      let producer_value = tile_B[channel * TILE_COLS + local_idx]{% if hasBias %} + f32(bias[channel]){% endif %};
+      let activated_value = producer_value;
+{% for oc in range(projectionOutChannels) %}
+      projected_acc{{ oc }} = projected_acc{{ oc }} + activated_value * f32(projectionW[{{ oc }}u * M + channel]);
+{% endfor %}
+    }
+    let fused_y_base = batch * {{ projectionOutChannels }}u * N + fused_global_col;
+{% for oc in range(projectionOutChannels) %}
+    let projected{{ oc }} = projected_acc{{ oc }};
+    let activated{{ oc }} = projected{{ oc }};
+    y[fused_y_base + {{ oc }}u * N] = activated{{ oc }};
+{% endfor %}
+  }
+{% elif useDirectMatrixStore %}
+{% if padded %}
+  // Padded N only needs guarded publication for the final partial-width tile.
+  // Full tiles can use the same direct collective stores as the aligned path;
+  // the branch is uniform because b_global_base depends only on workgroup_id.x.
+  if (b_global_base + TILE_COLS <= N) {
+{% endif %}
+  // Full f32 output tiles can be published by the subgroup-matrix collectives
+  // directly. With no epilogue this removes scratch traffic and three barriers.
+  // When an epilogue is present, publish the raw accumulators first, then apply
+  // bias, residual Z, and activation cooperatively after one storage barrier.
+  let matrix_c_offset = c_base + (a_global_base + base_A) * N + b_global_base + base_B;
+  subgroupMatrixStore(&y, matrix_c_offset + 0u * N +  0u, matC00, false, N);
+  subgroupMatrixStore(&y, matrix_c_offset + 0u * N +  8u, matC01, false, N);
+  subgroupMatrixStore(&y, matrix_c_offset + 0u * N + 16u, matC02, false, N);
+  subgroupMatrixStore(&y, matrix_c_offset + 0u * N + 24u, matC03, false, N);
+  subgroupMatrixStore(&y, matrix_c_offset + 8u * N +  0u, matC10, false, N);
+  subgroupMatrixStore(&y, matrix_c_offset + 8u * N +  8u, matC11, false, N);
+  subgroupMatrixStore(&y, matrix_c_offset + 8u * N + 16u, matC12, false, N);
+  subgroupMatrixStore(&y, matrix_c_offset + 8u * N + 24u, matC13, false, N);
+{% if hasBias or hasActivation or hasZ %}
+  storageBarrier();
+  for (var tile_idx = local_idx; tile_idx < TILE_ROWS * TILE_COLS; tile_idx += {{ workgroupThreadsValue }}u) {
+    let local_row = tile_idx / TILE_COLS;
+    let local_col = tile_idx % TILE_COLS;
+    let global_row = a_global_base + local_row;
+    let global_col = b_global_base + local_col;
+    let y_index = c_base + global_row * N + global_col;
+    var raw = f32(y[y_index]);
+{% if hasBias %}
+    raw += f32(bias[global_row]);
+{% endif %}
+{% if hasActivation or hasZ %}
+    y[y_index] = epi(raw{% if hasZ %}, y_index{% endif %});
+{% else %}
+    y[y_index] = {{ T }}(raw);
+{% endif %}
+  }
+{% endif %}
+{% if padded %}
+    return;
+  }
+{% endif %}
+{% endif %}
+{% if (not useDirectMatrixStore or padded) and not splitKPartial and not fusedNarrowProjection %}
+  // 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 row_base = a_global_base + base_A;
+  var matrix_c_offset = c_base + row_base * N + b_global_base + base_B;
+  var row_limit = i32(M) - i32(row_base);
+  storeOutput(matrix_c_offset, {% if hasBias or polyphase %}row_base, {% endif %}row, col, subtile_id, row_limit{% if padded %}, b_global_base + base_B{% endif %});
+  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(row_base + 8u);
+  storeOutput(matrix_c_offset, {% if hasBias or polyphase %}row_base + 8u, {% endif %}row, col, subtile_id, row_limit{% if padded %}, b_global_base + base_B{% endif %});
+{% endif %}
+}

build/webgpu/conv-direct-nd.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,86 @@

+// Direct {{ source.spatialRank }}D convolution for channels-first tensors. Inputs are
+// accumulated in f32 and narrowed once at the output.
+// A fused epilogue may apply a residual input and activation before the store.
+{% if source.usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+{% set hasActivation = hasActivation is defined and hasActivation %}
+{% if hasActivation %}
+// Apply the fused activation in the f32 accumulator before the single output
+// cast, avoiding an intermediate convolution tensor.
+fn fused_act(v: f32) -> f32 {
+{% if activation == "Relu" %}
+  return max(v, 0.0);
+{% elif activation == "Clip" %}
+  return clamp(v, f32({{ actAlpha }}), f32({{ actBeta }}));
+{% elif activation == "LeakyRelu" %}
+  return select(v * f32({{ actAlpha }}), v, v >= 0.0);
+{% elif activation == "Sigmoid" %}
+  return 1.0 / (1.0 + exp(-v));
+{% elif activation == "Tanh" %}
+  return tanh(v);
+{% elif activation == "HardSigmoid" %}
+  return clamp(f32({{ actAlpha }}) * v + f32({{ actBeta }}), 0.0, 1.0);
+{% else %}
+  return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
+{% endif %}
+}
+{% endif %}
+const WG: u32 = {{ convWorkgroupSize }}u;
+@compute @workgroup_size(WG)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>,
+        @builtin(num_workgroups) nwg: vec3<u32>) {
+  let index = gid.x + gid.y * nwg.x * WG;
+  if (index >= params.count) {
+    return;
+  }
+  let ow = index % params.outW;
+  var q = index / params.outW;
+  let oh = q % params.outH;
+  q /= params.outH;
+  let od = q % params.outD;
+  q /= params.outD;
+  let oc = q % params.outChannels;
+  let batch = q / params.outChannels;
+  let group = oc / params.outChannelsPerGroup;
+  var acc = 0.0;
+{% if source.hasBias %}
+  acc = f32(bias[oc]);
+{% endif %}
+  for (var ic = 0u; ic < params.weightInChannels; ic += 1u) {
+    let inputChannel = group * params.inChannelsPerGroup + ic;
+    for (var kd = 0u; kd < params.kernelD; kd += 1u) {
+      let id = i32(od * params.strideD + kd * params.dilationD) - params.padD;
+      if (id < 0 || id >= i32(params.inD)) {
+        continue;
+      }
+    for (var kh = 0u; kh < params.kernelH; kh += 1u) {
+      let ih = i32(oh * params.strideH + kh * params.dilationH) - params.padH;
+      if (ih < 0 || ih >= i32(params.inH)) {
+        continue;
+      }
+    for (var kw = 0u; kw < params.kernelW; kw += 1u) {
+      let iw = i32(ow * params.strideW + kw * params.dilationW) - params.padW;
+      if (iw < 0 || iw >= i32(params.inW)) {
+        continue;
+      }
+      let xIndex = (((batch * params.inChannels + inputChannel) * params.inD
+                     + u32(id)) * params.inH + u32(ih)) * params.inW + u32(iw);
+      let wIndex = (((oc * params.weightInChannels + ic) * params.kernelD
+                     + kd) * params.kernelH + kh) * params.kernelW + kw;
+      acc += f32(x[xIndex]) * f32(w[wIndex]);
+    }
+    }
+    }
+  }
+{% if hasZ %}
+  // Optional residual epilogue: Y = activation(conv + B + Z); Z shares Y's layout.
+  acc = acc + f32(zResidual[index]);
+{% endif %}
+  y[index] = {{ T }}({% if hasActivation %}fused_act(acc){% else %}acc{% endif %});
+}

build/webgpu/conv-direct-unrolled.wgsl.jinja ADDED Viewed

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

+// Direct NCHW 2D convolution with kernel extent, strides, dilations, and
+// top/left pads compiled into the fully unrolled window below.
+// One thread per output element, same flat index unpack and the same
+// accumulation order (ic outer, then kh, kw ascending) as the
+// scalar NCHW fallback, so f32 results are bit-identical to it.
+// The kh/kw window offsets, kh*KW+kw weight offsets, and the per-channel
+// strides fold to literals; the per-row bounds check is hoisted per kh.
+// Supports grouped convolution exactly like the scalar fallback. f16 inputs
+// are widened to an f32 accumulator and narrowed once at store.
+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+{% set hasActivation = hasActivation is defined and hasActivation %}
+{% if hasActivation %}
+// Apply the fused activation in the f32 accumulator before the single output
+// cast, avoiding an intermediate convolution tensor.
+fn fused_act(v: f32) -> f32 {
+{% if activation == "Relu" %}
+  return max(v, 0.0);
+{% elif activation == "Clip" %}
+  return clamp(v, f32({{ actAlpha }}), f32({{ actBeta }}));
+{% elif activation == "LeakyRelu" %}
+  return select(v * f32({{ actAlpha }}), v, v >= 0.0);
+{% elif activation == "Sigmoid" %}
+  return 1.0 / (1.0 + exp(-v));
+{% elif activation == "Tanh" %}
+  return tanh(v);
+{% elif activation == "HardSigmoid" %}
+  return clamp(f32({{ actAlpha }}) * v + f32({{ actBeta }}), 0.0, 1.0);
+{% else %}
+  return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
+{% endif %}
+}
+{% endif %}
+const KERNEL_AREA: u32 = {{ source.kernelH * source.kernelW }}u;
+const STRIDE_H: u32 = {{ source.strideH }}u;
+const STRIDE_W: u32 = {{ source.strideW }}u;
+const PAD_TOP: i32 = {{ source.padTop }};
+const PAD_LEFT: i32 = {{ source.padLeft }};
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
+  // 2D-folded flat index: gid.y carries the high bits past the
+  // maxComputeWorkgroupsPerDimension dispatch limit (outputs > 16.7M elements).
+  let index = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
+  if (index >= params.count) {
+    return;
+  }
+  let ow = index % params.outW;
+  var t = index / params.outW;
+  let oh = t % params.outH;
+  t = t / params.outH;
+  let oc = t % params.outChannels;
+  let batch = t / params.outChannels;
+  let group = oc / params.outChannelsPerGroup;
+  let ihBase = i32(oh * STRIDE_H) - PAD_TOP;
+  let iwBase = i32(ow * STRIDE_W) - PAD_LEFT;
+  let xChannelStride = params.inH * params.inW;
+  var xBase = (batch * params.inChannels + group * params.inChannelsPerGroup) * xChannelStride;
+  var wBase = oc * params.weightInChannels * KERNEL_AREA;
+  var acc = 0.0;
+{% if hasBias %}
+  acc = f32(bias[oc]);
+{% endif %}
+  for (var ic = 0u; ic < params.weightInChannels; ic = ic + 1u) {
+{% for kh in range(source.kernelH) %}
+    {
+      let ih = ihBase + {{ kh * source.dilationH }}i;
+      if (ih >= 0 && ih < i32(params.inH)) {
+        let xRow = xBase + u32(ih) * params.inW;
+{% for kw in range(source.kernelW) %}
+        {
+          let iw = iwBase + {{ kw * source.dilationW }}i;
+          if (iw >= 0 && iw < i32(params.inW)) {
+            acc = acc + f32(x[xRow + u32(iw)]) * f32(w[wBase + {{ kh * source.kernelW + kw }}u]);
+          }
+        }
+{% endfor %}
+      }
+    }
+{% endfor %}
+    xBase = xBase + xChannelStride;
+    wBase = wBase + KERNEL_AREA;
+  }
+{% if hasZ %}
+  // Optional residual epilogue: Y = activation(conv + B + Z); Z shares Y's layout.
+  acc = acc + f32(zResidual[index]);
+{% endif %}
+  y[index] = {{ T }}({% if hasActivation %}fused_act(acc){% else %}acc{% endif %});
+}

build/webgpu/conv-im2col-nchw.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,73 @@

+// Materialize the convolution column matrix
+//   cols[b][k][pos] = X[b, ic, oh*strideH + kh*dilationH - padTop,
+//                            ow*strideW + kw*dilationW - padLeft]
+// with zero outside the input and k = (ic*KH+kh)*KW+kw. This matches the OIHW
+// weight flattening, so the following GEMM consumes [M,K] x [K,outH*outW].
+// gid.x grid-strides over positions for coalesced reads and writes; gid.y fixes
+// k, allowing its ic/kh/kw decomposition to remain loop-invariant. The f32
+// scratch round-trips f16 inputs exactly when the GEMM narrows its tile loads.
+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+const KERNEL_H: u32 = {{ source.kernelH }}u;
+const KERNEL_W: u32 = {{ source.kernelW }}u;
+const KSIZE: u32 = KERNEL_H * KERNEL_W;
+const STRIDE_H: u32 = {{ source.strideH }}u;
+const STRIDE_W: u32 = {{ source.strideW }}u;
+const DILATION_H: u32 = {{ source.dilationH }}u;
+const DILATION_W: u32 = {{ source.dilationW }}u;
+const PAD_TOP: i32 = {{ source.padTop }};
+const PAD_LEFT: i32 = {{ source.padLeft }};
+{% if padded %}
+// Padded materialization for the subgroup-matrix GEMM: the cols buffer is
+// K -> ceil(K/32)*32 rows by N -> ceil(N/64)*64 columns, with rows >= kRows and
+// columns >= outCount zero-filled, so every GEMM tile load is full and the zero
+// region contributes exact zeros to the dot products / is never stored back.
+const COLS_STRIDE: u32 = {{ nPadded }}u;
+const K_PADDED: u32 = {{ kPadded }}u;
+{% endif %}
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(
+  @builtin(global_invocation_id) gid: vec3<u32>,
+  @builtin(num_workgroups) nwg: vec3<u32>
+) {
+  let k = gid.y;
+  let batch = gid.z;
+  // Loop-invariant: these depend on the k row alone.
+  let ic = k / KSIZE;
+  let kq = k % KSIZE;
+  let kh = kq / KERNEL_W;
+  let kw = kq % KERNEL_W;
+  let planeBase = (batch * params.inChannels + ic) * params.inH;
+{% if padded %}
+  let limit = COLS_STRIDE;
+{% else %}
+  let limit = params.outCount;
+{% endif %}
+  let step = nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
+  var pos = gid.x;
+  loop {
+    if (pos >= limit) { break; }
+    var value = 0.0;
+{% if padded %}
+    if (pos < params.outCount && k < params.kRows) {
+{% endif %}
+    let oh = pos / params.outW;
+    let ow = pos % params.outW;
+    let ih = i32(oh * STRIDE_H + kh * DILATION_H) - PAD_TOP;
+    let iw = i32(ow * STRIDE_W + kw * DILATION_W) - PAD_LEFT;
+    if (ih >= 0 && ih < i32(params.inH) && iw >= 0 && iw < i32(params.inW)) {
+      value = f32(x[(planeBase + u32(ih)) * params.inW + u32(iw)]);
+    }
+{% if padded %}
+    }
+    cols[(batch * K_PADDED + k) * COLS_STRIDE + pos] = value;
+{% else %}
+    cols[(batch * params.kRows + k) * params.outCount + pos] = value;
+{% endif %}
+    pos = pos + step;
+  }
+}

build/webgpu/conv1d-tiled-reg.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,166 @@

+// Implicit-GEMM register-tiled 1-D convolution over NCW: Y[b] (M x N) =
+// W (M x K) @ im2col(X[b]) (K x N) with M = outChannels, K = inChannels *
+// kernelW, N = outW. The im2col operand is never materialized: the X staging
+// loop decodes each GEMM row into (ic, kw) and samples x with stride /
+// dilation / left-pad applied, so one program covers general stride, padding,
+// and dilation (group == 1 only). Workgroup geometry is tunable: a WG_X x WG_Y
+// workgroup computes a (WG_Y*TM) x (WG_X*TN) output tile from BK-deep staged
+// slices, each thread owning a TM x TN register micro-tile. Cooperative loads
+// and stores are fully bounds-checked, so partial tiles on any edge stay
+// correct for every tunable setting. f32 accumulation; bias and the optional
+// activation epilogue are applied in the accumulator domain at
+// store. f16 operands stay packed in workgroup memory and widen at the FMA
+// site. 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 the implicit
+// gather decodes one (ic, kw) tap per word, walking the remaining three. The
+// micro-tile is unrolled into named registers: a dynamically indexed private
+// array of vector words is large enough at the wide-M geometry that it may fail
+// to scalarize and spill the accumulators.
+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+{% set hasActivation = hasActivation is defined and hasActivation %}
+{% if hasActivation %}
+// Apply the fused activation in the f32 accumulator before the single output
+// cast, avoiding an intermediate convolution tensor.
+fn fused_act(v: f32) -> f32 {
+{% if activation == "Relu" %}
+  return max(v, 0.0);
+{% elif activation == "Clip" %}
+  return clamp(v, f32({{ actAlpha }}), f32({{ actBeta }}));
+{% elif activation == "LeakyRelu" %}
+  return select(v * f32({{ actAlpha }}), v, v >= 0.0);
+{% elif activation == "Sigmoid" %}
+  return 1.0 / (1.0 + exp(-v));
+{% elif activation == "Tanh" %}
+  return tanh(v);
+{% elif activation == "HardSigmoid" %}
+  return clamp(f32({{ actAlpha }}) * v + f32({{ actBeta }}), 0.0, 1.0);
+{% else %}
+  return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
+{% endif %}
+}
+{% endif %}
+{% set tileT = "f16" if usesF16 else "f32" %}
+const WG_X: u32 = {{ convWgX }}u;
+const WG_Y: u32 = {{ convWgY }}u;
+const THREADS: u32 = WG_X * WG_Y;
+const BK: u32 = {{ convKTile }}u;
+const TM: u32 = {{ convTileM }}u;
+const TN: u32 = {{ convTileN }}u;
+const BM: u32 = WG_Y * TM;
+const BN: u32 = WG_X * TN;
+fn store_out(oc: u32, ow: u32, yBase: u32, raw: f32) {
+  if (oc >= params.outChannels || ow >= params.outW) {
+    return;
+  }
+  var v = raw;
+{% if hasBias %}
+  v = v + f32(bias[oc]);
+{% endif %}
+{% if hasActivation %}
+  v = fused_act(v);
+{% endif %}
+  y[yBase + oc * params.outW + ow] = {{ T }}(v);
+}
+const K_VECS: u32 = BK / 4u;
+var<workgroup> tileW: array<array<vec4<{{ tileT }}>, K_VECS>, BM>;   // W[oc][k/4]
+var<workgroup> tileX: array<array<vec4<{{ tileT }}>, K_VECS>, BN>;   // im2col[ow][k/4]
+@compute @workgroup_size(WG_X, WG_Y, 1)
+fn main(
+  @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>
+) {
+  let K = params.inChannels * params.kernelW;
+  let ocBase = wg.y * BM;
+  let owBase = wg.x * BN;
+  let batch = wg.z;
+  let li = lid.y * WG_X + lid.x;
+{% for i in range(convTileM) %}
+{% for j in range(convTileN) %}
+  var acc{{ i }}_{{ j }}: f32 = 0.0;
+{% endfor %}
+{% endfor %}
+  let numTiles = (K + BK - 1u) / BK;
+  for (var kt: u32 = 0u; kt < numTiles; kt = kt + 1u) {
+    let kBase = kt * BK;
+    // One vector word per lane per pass: four contiguous K values of one output
+    // channel, and the four K taps of one output position, whose (ic, kw) is
+    // decoded once and then walked.
+    for (var idx: u32 = li; idx < BM * K_VECS; idx = idx + THREADS) {
+      let wr = idx / K_VECS;
+      let wc4 = idx % K_VECS;
+      let oc = ocBase + wr;
+      let k = kBase + wc4 * 4u;
+      var word = vec4<{{ tileT }}>({{ tileT }}(0.0));
+      if (oc < params.outChannels) {
+        let rowBase = oc * K + k;
+{% for component in range(4) %}
+        if (k + {{ component }}u < K) { word[{{ component }}u] = {{ tileT }}(w[rowBase + {{ component }}u]); }
+{% endfor %}
+      }
+      tileW[wr][wc4] = word;
+    }
+    for (var idx: u32 = li; idx < BN * K_VECS; idx = idx + THREADS) {
+      let xc = idx / K_VECS;
+      let xr4 = idx % K_VECS;
+      let ow = owBase + xc;
+      let k0 = kBase + xr4 * 4u;
+      var word = vec4<{{ tileT }}>({{ tileT }}(0.0));
+      if (ow < params.outW) {
+        let owStride = ow * params.strideW;
+        let planeBase = batch * params.inChannels;
+        var ic = k0 / params.kernelW;
+        var kw = k0 % params.kernelW;
+{% for component in range(4) %}
+        if (k0 + {{ component }}u < K) {
+          let iw{{ component }} = i32(owStride + kw * params.dilationW) - params.padW;
+          if (iw{{ component }} >= 0 && iw{{ component }} < i32(params.inW)) {
+            word[{{ component }}u] = {{ tileT }}(x[(planeBase + ic) * params.inW + u32(iw{{ component }})]);
+          }
+        }
+        kw = kw + 1u;
+        if (kw == params.kernelW) {
+          kw = 0u;
+          ic = ic + 1u;
+        }
+{% endfor %}
+      }
+      tileX[xc][xr4] = word;
+    }
+    workgroupBarrier();
+    let tileOc = lid.y * TM;
+    let tileOw = lid.x * TN;
+    for (var kv: u32 = 0u; kv < K_VECS; kv = kv + 1u) {
+{% for i in range(convTileM) %}
+      let wv{{ i }} = vec4<f32>(tileW[tileOc + {{ i }}u][kv]);
+{% endfor %}
+{% for j in range(convTileN) %}
+      let xv{{ j }} = vec4<f32>(tileX[tileOw + {{ j }}u][kv]);
+{% endfor %}
+{% for i in range(convTileM) %}
+{% for j in range(convTileN) %}
+      acc{{ i }}_{{ j }} = acc{{ i }}_{{ j }} + dot(wv{{ i }}, xv{{ j }});
+{% endfor %}
+{% endfor %}
+    }
+    workgroupBarrier();
+  }
+  let yBase = batch * params.outChannels * params.outW;
+  let oc0 = ocBase + lid.y * TM;
+  let ow0 = owBase + lid.x * TN;
+{% for i in range(convTileM) %}
+{% for j in range(convTileN) %}
+  store_out(oc0 + {{ i }}u, ow0 + {{ j }}u, yBase, acc{{ i }}_{{ j }});
+{% endfor %}
+{% endfor %}
+}

build/webgpu/conv2d-grouped-large-w4.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,176 @@

+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+{% set hasActivation = hasActivation is defined and hasActivation %}
+{% if hasActivation %}
+// Apply the fused activation in the f32 accumulator before the single output
+// cast, avoiding an intermediate convolution tensor.
+fn fused_act(v: f32) -> f32 {
+{% if activation == "Relu" %}
+  return max(v, 0.0);
+{% elif activation == "Clip" %}
+  return clamp(v, f32({{ actAlpha }}), f32({{ actBeta }}));
+{% elif activation == "LeakyRelu" %}
+  return select(v * f32({{ actAlpha }}), v, v >= 0.0);
+{% elif activation == "Sigmoid" %}
+  return 1.0 / (1.0 + exp(-v));
+{% elif activation == "Tanh" %}
+  return tanh(v);
+{% elif activation == "HardSigmoid" %}
+  return clamp(f32({{ actAlpha }}) * v + f32({{ actBeta }}), 0.0, 1.0);
+{% else %}
+  return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
+{% endif %}
+}
+{% endif %}
+// Large grouped kernels are dominated by repeatedly fetching the same filter
+// value for neighboring output columns and input window for neighboring output
+// channels. One invocation computes four adjacent columns for OC_TILE channels,
+// sharing each input load across those channels. Shape and window geometry are
+// static, so all indexing divisors and kernel offsets are strength-reduced and
+// unrolled by the shader compiler.
+const COUNT_TILES: u32 = {{ source.countTiles }}u;
+{% if source.tailOutput %}const OUT_W: u32 = max(1u, {{ source.outW }}u);
+{% endif %}
+const OUT_W4: u32 = max(1u, {{ source.outW4 }}u);
+const OUT_H: u32 = max(1u, {{ source.outH }}u);
+const OUT_C: u32 = max(1u, {{ source.outC }}u);
+const OC_TILE: u32 = {{ source.ocTile }}u;
+const OUT_C_TILES: u32 = OUT_C / OC_TILE;
+const OUT_CPG: u32 = {{ source.outCPerGroup }}u;
+const IN_C: u32 = {{ source.inC }}u;
+const IN_CPG: u32 = {{ source.inCPerGroup }}u;
+const IN_H: i32 = {{ source.inH }};
+const IN_W: i32 = {{ source.inW }};
+const IN_H_U: u32 = {{ source.inH }}u;
+const IN_W_U: u32 = {{ source.inW }}u;
+const IN_PLANE: u32 = IN_H_U * IN_W_U;
+{% if not source.tailOutput %}const OUT_PLANE4: u32 = OUT_H * OUT_W4;
+{% endif %}
+const KAREA: u32 = {{ source.kernelH * source.kernelW }}u;
+const STRIDE_H: u32 = {{ source.strideH }}u;
+const STRIDE_W: i32 = {{ source.strideW }};
+const PAD_TOP: i32 = {{ source.padTop }};
+const PAD_LEFT: i32 = {{ source.padLeft }};
+const WG: u32 = {{ source.workgroupSize }}u;
+@compute @workgroup_size({{ source.workgroupSize }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
+  let q = gid.x + gid.y * nwg.x * WG;
+  if (q >= COUNT_TILES) { return; }
+  let ow4 = q % OUT_W4;
+  var t = q / OUT_W4;
+  let oh = t % OUT_H;
+  t /= OUT_H;
+  let ocBase = (t % OUT_C_TILES) * OC_TILE;
+  let batch = t / OUT_C_TILES;
+  let group = ocBase / OUT_CPG;
+  let ihBase = i32(oh * STRIDE_H) - PAD_TOP;
+  let iwBase = i32(ow4 * 4u) * STRIDE_W - PAD_LEFT;
+  var xChannelBase = (batch * IN_C + group * IN_CPG) * IN_PLANE;
+{% for oct in range(source.ocTile) %}
+  var wBase{{ oct }} = (ocBase + {{ oct }}u) * IN_CPG * KAREA;
+  var acc{{ oct }} = vec4<f32>(0.0);
+{% endfor %}
+  for (var ic = 0u; ic < IN_CPG; ic++) {
+{% for kh in range(source.kernelH) %}
+    {
+      let ih = ihBase + {{ kh * source.dilationH }};
+      if (ih >= 0 && ih < IN_H) {
+        let xRow = xChannelBase + u32(ih) * IN_W_U;
+{% if source.span <= source.spanCap %}
+        // The four output columns and the kernel taps overlap heavily: tap kw of
+        // lane l reads input column kw*dilationW + l*strideW, so kernelW*4 reads
+        // land on only SPAN distinct columns. Load the row window once and let
+        // every tap take its value from a register. Out-of-range columns hold
+        // zero, which contributes nothing to the accumulation exactly as the
+        // per-tap bounds test did.
+        //
+        // SPAN grows as (kernelW-1)*dilationW, so a widely dilated kernel would
+        // want more registers than a thread has and spill. Past spanCap the
+        // per-tap form below stays, where the register cost is fixed.
+{% for s in range(source.span) %}
+{% set used = namespace(value=false) %}
+{% for kw in range(source.kernelW) %}
+{% for lane in range(4) %}
+{% if s == kw * source.dilationW + lane * source.strideW %}{% set used.value = true %}{% endif %}
+{% endfor %}
+{% endfor %}
+{% if used.value %}
+        var xw{{ s }} = 0.0;
+{% endif %}
+{% endfor %}
+{% for s in range(source.span) %}
+{% set used = namespace(value=false) %}
+{% for kw in range(source.kernelW) %}
+{% for lane in range(4) %}
+{% if s == kw * source.dilationW + lane * source.strideW %}{% set used.value = true %}{% endif %}
+{% endfor %}
+{% endfor %}
+{% if used.value %}
+        {
+          let iw = iwBase + {{ s }};
+          if (iw >= 0 && iw < IN_W) { xw{{ s }} = f32(x[xRow + u32(iw)]); }
+        }
+{% endif %}
+{% endfor %}
+{% for kw in range(source.kernelW) %}
+        {
+{% for oct in range(source.ocTile) %}
+          let weight{{ oct }} = f32(w[wBase{{ oct }} + {{ kh * source.kernelW + kw }}u]);
+{% endfor %}
+{% for lane in range(4) %}
+{% for oct in range(source.ocTile) %}
+          acc{{ oct }}.{{ ["x", "y", "z", "w"][lane] }} += xw{{ kw * source.dilationW + lane * source.strideW }} * weight{{ oct }};
+{% endfor %}
+{% endfor %}
+        }
+{% endfor %}
+{% else %}
+{% for kw in range(source.kernelW) %}
+        {
+{% for oct in range(source.ocTile) %}
+          let weight{{ oct }} = f32(w[wBase{{ oct }} + {{ kh * source.kernelW + kw }}u]);
+{% endfor %}
+          let iwK = iwBase + {{ kw * source.dilationW }};
+{% for lane in range(4) %}
+          let iwLane{{ lane }} = iwK + {{ lane * source.strideW }};
+          if (iwLane{{ lane }} >= 0 && iwLane{{ lane }} < IN_W) {
+            let xValue = f32(x[xRow + u32(iwLane{{ lane }})]);
+{% for oct in range(source.ocTile) %}
+            acc{{ oct }}.{{ ["x", "y", "z", "w"][lane] }} += xValue * weight{{ oct }};
+{% endfor %}
+          }
+{% endfor %}
+        }
+{% endfor %}
+{% endif %}
+      }
+    }
+{% endfor %}
+    xChannelBase += IN_PLANE;
+{% for oct in range(source.ocTile) %}
+    wBase{{ oct }} += KAREA;
+{% endfor %}
+  }
+{% if source.tailOutput %}
+  let outBase = ((batch * OUT_C + ocBase) * OUT_H + oh) * OUT_W + ow4 * 4u;
+{% for oct in range(source.ocTile) %}
+  let outChannelBase{{ oct }} = outBase + {{ oct }}u * OUT_H * OUT_W;
+{% for lane in range(4) %}
+  if (ow4 * 4u + {{ lane }}u < OUT_W) {
+    y[outChannelBase{{ oct }} + {{ lane }}u] = {{ scalar }}({% if hasActivation %}fused_act(acc{{ oct }}.{{ ["x", "y", "z", "w"][lane] }}){% else %}acc{{ oct }}.{{ ["x", "y", "z", "w"][lane] }}{% endif %});
+  }
+{% endfor %}
+{% endfor %}
+{% else %}
+  let outBase = ((batch * OUT_C + ocBase) * OUT_H + oh) * OUT_W4 + ow4;
+{% for oct in range(source.ocTile) %}
+  y[outBase + {{ oct }}u * OUT_PLANE4] = {{ vectorScalar }}({% if hasActivation %}vec4<f32>(fused_act(acc{{ oct }}.x), fused_act(acc{{ oct }}.y), fused_act(acc{{ oct }}.z), fused_act(acc{{ oct }}.w)){% else %}acc{{ oct }}{% endif %});
+{% endfor %}
+{% endif %}
+}

build/webgpu/manifest.json ADDED Viewed

The diff for this file is too large to render. See raw diff

build/webgpu/metadata.json ADDED Viewed

	@@ -0,0 +1,25 @@

+{
+  "name": "com.microsoft.FusedConv",
+  "id": "_com_microsoft_fusedconv_webgpu_947df1c",
+  "version": 1,
+  "license": "Apache-2.0",
+  "backend": { "type": "webgpu" },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "bench.json": "0T4EaJworpji9PrXCF3znUUqiMBBbSgShzzAaM/dOWY=",
+      "conv-1x1-gemm-tiled-reg.wgsl.jinja": "bDk/sW4LxlVlqVJ6L/hpKGjfky1UKWwYL3C0fdVWjFA=",
+      "conv-1x1-gemm-tiled.wgsl.jinja": "eHdHLa3X2HriPxhMMSmzkVPelPg+sm4xDz9ZJ4qI0hA=",
+      "conv-1x1-subgroup-matrix.wgsl.jinja": "FTHq1ftW5Zk7as0fc8T+b0dvbQFQFmcX0AWARnICV28=",
+      "conv-direct-nd.wgsl.jinja": "yEas2jFgJNfM8NoHjv+4ZafMC1/nwl5MvgQeTFfOl4w=",
+      "conv-direct-unrolled.wgsl.jinja": "84nok2Ie1Lc91UrXiUEl1GFk+cwa4Fu0vEkhRyOmsxs=",
+      "conv-im2col-nchw.wgsl.jinja": "7IUuXo33elZx0pnR9M2vUXzXDjFuw2bJODkZViV/WQg=",
+      "conv1d-tiled-reg.wgsl.jinja": "d9H8LBgSgsqxBbg85o5dCMuOEIW/hOICoEi5gysrvTg=",
+      "conv2d-grouped-large-w4.wgsl.jinja": "jhLKpmvfkcFQPBXXznt3im++rtVwNNxUoRqqgHNpmDM=",
+      "manifest.json": "jMjmEiSVgMOjNheSacCaHlppVU7uWplUp2lv3fvGH1A=",
+      "test.json": "/KodhC9ZJ6IxO1wokLmeGH0zsOfDpds9+mjZLZRnPXw="
+    }
+  },
+  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
+  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.FusedConv" }
+}

build/webgpu/test.json ADDED Viewed

The diff for this file is too large to render. See raw diff