sync 91d990483a17

Browse files

Files changed (14) hide show

README.md +53 -13
build/webgpu/bench.json +460 -8
build/webgpu/conv-1x1-gemm-tiled-reg.wgsl.jinja +104 -42
build/webgpu/conv-1x1-gemm-tiled.wgsl.jinja +28 -34
build/webgpu/conv-1x1-subgroup-matrix.wgsl.jinja +302 -29
build/webgpu/conv-direct-nd.wgsl.jinja +24 -19
build/webgpu/conv-direct-unrolled.wgsl.jinja +36 -34
build/webgpu/conv-im2col-nchw.wgsl.jinja +13 -14
build/webgpu/conv-splitk-reduce.wgsl.jinja +83 -0
build/webgpu/conv1d-tiled-reg.wgsl.jinja +29 -34
build/webgpu/conv2d-grouped-large-w4.wgsl.jinja +128 -77
build/webgpu/manifest.json +0 -0
build/webgpu/metadata.json +80 -14
build/webgpu/test.json +0 -0

README.md CHANGED Viewed

@@ -18,18 +18,18 @@ See the [ONNX Runtime `FusedConv` contrib-operator spec](https://github.com/micr
 ## 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
@@ -37,11 +37,11 @@ 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. |
@@ -52,13 +52,48 @@ Attributes and default values (overridable per request):
 | --- | --- |
 | `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
@@ -68,15 +103,20 @@ Some implementation variants require `subgroup-matrix` and `subgroups`. These ar
 - [`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.

 ## Inputs
+| Name | Upstream name | 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 |
+| `bias` | `B` | `T` | `1` | — | Optional 1-D bias tensor of length `out_channels`, broadcast-added to each output channel. | optional |
+| `zResidual` | `Z` | `T` | same as `x` | — | Optional residual tensor with the same shape as the output `Y`, added before the activation. | optional |
 ## Outputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `y` | `Y` | `T` | same as `x` | derived | Output feature map tensor after convolution, optional bias/residual addition, and the fused activation. | required |
 ## Attributes
 | Attribute | Default | Description |
 | --- | --- | --- |
 | `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. |
+| `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. |
 | `dilations` | — | Optional dilation factors, one positive integer per spatial axis. Omission means all ones. |
+| `group` | `1` | Number of groups that input and output channels are split into; defaults to 1. |
 | `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. |
 | --- | --- |
 | `T` | `float32`, `float16` |
+## Implementation variants
+One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
+- `implicit_im2col_tiled_reg_splitk` — Partitions the register-blocked implicit-im2col reduction across workgroups when the unsplit output grid is too small, then combines raw partial sums and applies the bias and fused activation once. The combine reassociates floating-point addition relative to the unsplit route.
+- `implicit_im2col_tiled_bias_reg_splitk` — Partitions the register-blocked implicit-im2col reduction across workgroups when the unsplit output grid is too small, then combines raw partial sums and applies the bias and fused activation once. The combine reassociates floating-point addition relative to the unsplit route.
+- `implicit_im2col_tiled_reg_splitk_preferred` — Partitions the register-blocked implicit-im2col reduction across workgroups when the unsplit output grid is too small, then combines raw partial sums and applies the bias and fused activation once. The combine reassociates floating-point addition relative to the unsplit route.
+- `implicit_im2col_tiled_bias_reg_splitk_preferred` — Partitions the register-blocked implicit-im2col reduction across workgroups when the unsplit output grid is too small, then combines raw partial sums and applies the bias and fused activation once. The combine reassociates floating-point addition relative to the unsplit route.
+- `implicit_im2col_tiled_reg` — Gathers logical im2col elements directly from `X` into register-blocked GEMM tiles and applies the fused activation in the epilogue, avoiding a materialized column scratch buffer.
+- `implicit_im2col_tiled_bias_reg` — Gathers logical im2col elements directly from `X` into register-blocked GEMM tiles and applies bias and the fused activation in the epilogue, avoiding a materialized column scratch buffer.
+- `gemm_1x1_subgroup_matrix` — Multiplies 1x1 weights directly by NCHW input, then applies the selected bias, residual and activation epilogue. Complete output-channel tiles bypass input staging and reduction-loop barriers; partial tiles retain guarded staging. Tile geometry follows the device's workgroup limits.
+- `gemm_1x1_subgroup_matrix_z` — Multiplies 1x1 weights directly by NCHW input, then applies the selected bias, residual and activation epilogue. Complete output-channel tiles bypass input staging and reduction-loop barriers; partial tiles retain guarded staging. Tile geometry follows the device's workgroup limits.
+- `gemm_1x1_subgroup_matrix_bias` — Multiplies 1x1 weights directly by NCHW input, then applies the selected bias, residual and activation epilogue. Complete output-channel tiles bypass input staging and reduction-loop barriers; partial tiles retain guarded staging. Tile geometry follows the device's workgroup limits.
+- `gemm_1x1_subgroup_matrix_bias_z` — Multiplies 1x1 weights directly by NCHW input, then applies the selected bias, residual and activation epilogue. Complete output-channel tiles bypass input staging and reduction-loop barriers; partial tiles retain guarded staging. Tile geometry follows the device's workgroup limits.
+- `im2col_gemm_subgroup_matrix` — Multiply aligned materialized columns through workgroup-staged subgroup matrices with f32 accumulation and the fused epilogue.
+- `im2col_direct_inputs_subgroup_matrix` — Multiply aligned materialized f32 columns through direct subgroup-matrix loads, retaining the f32 accumulation order and fused epilogue while omitting operand staging and K-loop barriers.
+- `im2col_gemm_subgroup_matrix_z` — Multiply aligned materialized columns through workgroup-staged subgroup matrices with f32 accumulation and the fused epilogue.
+- `im2col_direct_inputs_subgroup_matrix_z` — Multiply aligned materialized f32 columns through direct subgroup-matrix loads, retaining the f32 accumulation order and fused epilogue while omitting operand staging and K-loop barriers.
+- `im2col_gemm_subgroup_matrix_bias` — Multiply aligned materialized columns through workgroup-staged subgroup matrices with f32 accumulation and the fused epilogue.
+- `im2col_direct_inputs_subgroup_matrix_bias` — Multiply aligned materialized f32 columns through direct subgroup-matrix loads, retaining the f32 accumulation order and fused epilogue while omitting operand staging and K-loop barriers.
+- `im2col_gemm_subgroup_matrix_bias_z` — Multiply aligned materialized columns through workgroup-staged subgroup matrices with f32 accumulation and the fused epilogue.
+- `im2col_direct_inputs_subgroup_matrix_bias_z` — Multiply aligned materialized f32 columns through direct subgroup-matrix loads, retaining the f32 accumulation order and fused epilogue while omitting operand staging and K-loop barriers.
+- `im2col_half_direct_subgroup_matrix` — Materialize aligned f16 convolution columns without widening their storage, then load weights and columns directly into subgroup matrices. Preserve f32 accumulation and the existing bias, residual, activation and f16 output rounding while removing operand staging and K-loop barriers.
+- `im2col_half_direct_subgroup_matrix_z` — Materialize aligned f16 convolution columns without widening their storage, then load weights and columns directly into subgroup matrices. Preserve f32 accumulation and the existing bias, residual, activation and f16 output rounding while removing operand staging and K-loop barriers.
+- `im2col_half_direct_subgroup_matrix_bias` — Materialize aligned f16 convolution columns without widening their storage, then load weights and columns directly into subgroup matrices. Preserve f32 accumulation and the existing bias, residual, activation and f16 output rounding while removing operand staging and K-loop barriers.
+- `im2col_half_direct_subgroup_matrix_bias_z` — Materialize aligned f16 convolution columns without widening their storage, then load weights and columns directly into subgroup matrices. Preserve f32 accumulation and the existing bias, residual, activation and f16 output rounding while removing operand staging and K-loop barriers.
+- `implicit_im2col_subgroup_matrix` — Gathers convolution input tiles directly into subgroup-matrix operands and applies the existing bias, residual, and activation epilogue, avoiding a materialized column matrix when output-channel reuse permits or both column layouts exceed device allocation limits.
+- `implicit_im2col_subgroup_matrix_z` — Gathers convolution input tiles directly into subgroup-matrix operands and applies the existing bias, residual, and activation epilogue, avoiding a materialized column matrix when output-channel reuse permits or both column layouts exceed device allocation limits.
+- `implicit_im2col_subgroup_matrix_bias` — Gathers convolution input tiles directly into subgroup-matrix operands and applies the existing bias, residual, and activation epilogue, avoiding a materialized column matrix when output-channel reuse permits or both column layouts exceed device allocation limits.
+- `implicit_im2col_subgroup_matrix_bias_z` — Gathers convolution input tiles directly into subgroup-matrix operands and applies the existing bias, residual, and activation epilogue, avoiding a materialized column matrix when output-channel reuse permits or both column layouts exceed device allocation limits.
+- `implicit_im2col_subgroup_matrix_f16` — Gathers convolution input tiles directly into subgroup-matrix operands and applies the existing bias, residual, and activation epilogue, avoiding a materialized column matrix when output-channel reuse permits or both column layouts exceed device allocation limits.
+- `implicit_im2col_subgroup_matrix_z_f16` — Gathers convolution input tiles directly into subgroup-matrix operands and applies the existing bias, residual, and activation epilogue, avoiding a materialized column matrix when output-channel reuse permits or both column layouts exceed device allocation limits.
+- `implicit_im2col_subgroup_matrix_bias_f16` — Gathers convolution input tiles directly into subgroup-matrix operands and applies the existing bias, residual, and activation epilogue, avoiding a materialized column matrix when output-channel reuse permits or both column layouts exceed device allocation limits.
+- `implicit_im2col_subgroup_matrix_bias_z_f16` — Gathers convolution input tiles directly into subgroup-matrix operands and applies the existing bias, residual, and activation epilogue, avoiding a materialized column matrix when output-channel reuse permits or both column layouts exceed device allocation limits.
 ## Device requirements
+Some implementation variants require `subgroup-matrix`, `shader-f16`, and `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
 ## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 - [`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)
+- [`conv-splitk-reduce.wgsl.jinja`](build/webgpu/conv-splitk-reduce.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`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.2
+```
+Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

build/webgpu/bench.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "com.microsoft.FusedConv",
   "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
   "cases": [
     {
@@ -227,8 +226,8 @@
       "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,
@@ -667,8 +666,8 @@
       "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,
@@ -826,7 +825,7 @@
     {
       "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": {
@@ -861,7 +860,7 @@
     {
       "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": {
@@ -898,7 +897,7 @@
       "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] },
@@ -917,6 +916,459 @@
           }
         ]
       }
     }
   ]
 }

 {
   "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
   "cases": [
     {
       "name": "fusedconv-f32-1x1-bias-b1c256m256-32x32",
       "preset": "smoke",
       "provenance": {
+        "source": "synthetic",
+        "notes": "A biased 1x1 convolution produces 64 register tiles, measuring the register-tiled route at its admission floor."
       },
       "vars": {
         "batch": 1,
       "name": "fusedconv-3x3-dense-bias-c256m256-32x32",
       "preset": "smoke",
       "provenance": {
+        "source": "synthetic",
+        "notes": "A biased 3x3 convolution produces 64 register tiles, measuring the register-tiled im2col route at its admission floor."
       },
       "vars": {
         "batch": 1,
     {
       "name": "fusedconv-group2-k7-c64m64-64x64-direct-fallback",
       "provenance": {
+        "notes": "A group-2, 7x7 convolution with wide output exercises the four-wide grouped large-kernel route."
       },
       "preset": "stress",
       "vars": {
     {
       "name": "fusedconv-group4-k5-c128m128-48x48-direct-fallback",
       "provenance": {
+        "notes": "A group-4, 5x5 convolution exercises the four-wide grouped large-kernel route near its kernel-size boundary."
       },
       "preset": "stress",
       "vars": {
       "preset": "stress",
       "provenance": {
         "source": "ONNX Runtime FusedConv provider semantics",
+        "notes": "A large aligned temporal convolution with bias and fused HardSwish exercises 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] },
           }
         ]
       }
+    },
+    {
+      "name": "fusedconv-splitk-grid32-b1m128-32x32",
+      "preset": "smoke",
+      "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": [128, 64, 3, 3], "dtype": "float32", "dist": "normal", "seed": 2676, "scale": 0.05 }
+      },
+      "outputs": { "y": { "shape": [1, 128, 32, 32], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          { "type": "gflops", "value": "2 * numel(shapes.y) * dim(shapes.w, 1) * dim(shapes.w, 2) * dim(shapes.w, 3)" }
+        ]
+      },
+      "provenance": {
+        "source": "synthetic",
+        "notes": "Measures the 32-workgroup implicit split-K preference boundary with output channels, spatial extent, or batch contributing to the grid."
+      }
+    },
+    {
+      "name": "fusedconv-splitk-grid32-b1m64-32x64",
+      "preset": "smoke",
+      "attrs": { "activation": "Relu", "strides": [1, 1], "pads": [1, 1, 1, 1] },
+      "inputs": {
+        "x": { "shape": [1, 64, 32, 64], "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, 64], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          { "type": "gflops", "value": "2 * numel(shapes.y) * dim(shapes.w, 1) * dim(shapes.w, 2) * dim(shapes.w, 3)" }
+        ]
+      },
+      "provenance": {
+        "source": "synthetic",
+        "notes": "Measures the 32-workgroup implicit split-K preference boundary with output channels, spatial extent, or batch contributing to the grid."
+      }
+    },
+    {
+      "name": "fusedconv-splitk-grid32-b2m64-32x32",
+      "preset": "smoke",
+      "attrs": { "activation": "Relu", "strides": [1, 1], "pads": [1, 1, 1, 1] },
+      "inputs": {
+        "x": { "shape": [2, 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": [2, 64, 32, 32], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          { "type": "gflops", "value": "2 * numel(shapes.y) * dim(shapes.w, 1) * dim(shapes.w, 2) * dim(shapes.w, 3)" }
+        ]
+      },
+      "provenance": {
+        "source": "synthetic",
+        "notes": "Measures the 32-workgroup implicit split-K preference boundary with output channels, spatial extent, or batch contributing to the grid."
+      }
+    },
+    {
+      "name": "fusedconv-direct-inputs-row-band-m32-k288-n256",
+      "preset": "smoke",
+      "vars": {
+        "batch": 1,
+        "inChannels": 32,
+        "outChannels": 32,
+        "inH": 16,
+        "inW": 16,
+        "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, 16, 16], "dtype": "float32", "dist": "normal", "seed": 839, "scale": 0.2 },
+        "w": { "shape": [32, 32, 3, 3], "dtype": "float32", "dist": "normal", "seed": 1375, "scale": 0.05 }
+      },
+      "outputs": { "y": { "shape": [1, 32, 16, 16], "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-direct-inputs-row-band-m64-k288-n256",
+      "preset": "smoke",
+      "vars": {
+        "batch": 1,
+        "inChannels": 32,
+        "outChannels": 64,
+        "inH": 16,
+        "inW": 16,
+        "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, 16, 16], "dtype": "float32", "dist": "normal", "seed": 839, "scale": 0.2 },
+        "w": { "shape": [64, 32, 3, 3], "dtype": "float32", "dist": "normal", "seed": 1375, "scale": 0.05 }
+      },
+      "outputs": { "y": { "shape": [1, 64, 16, 16], "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-half-direct-inputs-row-band-m32-k288-n256",
+      "preset": "smoke",
+      "vars": {
+        "batch": 1,
+        "inChannels": 32,
+        "outChannels": 32,
+        "inH": 16,
+        "inW": 16,
+        "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, 16, 16], "dtype": "float16", "dist": "normal", "seed": 839, "scale": 0.2 },
+        "w": { "shape": [32, 32, 3, 3], "dtype": "float16", "dist": "normal", "seed": 1375, "scale": 0.05 }
+      },
+      "outputs": { "y": { "shape": [1, 32, 16, 16], "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-half-direct-inputs-row-band-m64-k288-n256",
+      "preset": "smoke",
+      "vars": {
+        "batch": 1,
+        "inChannels": 32,
+        "outChannels": 64,
+        "inH": 16,
+        "inW": 16,
+        "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, 16, 16], "dtype": "float16", "dist": "normal", "seed": 839, "scale": 0.2 },
+        "w": { "shape": [64, 32, 3, 3], "dtype": "float16", "dist": "normal", "seed": 1375, "scale": 0.05 }
+      },
+      "outputs": { "y": { "shape": [1, 64, 16, 16], "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-half-direct-inputs-row-band-m96-k288-n256",
+      "preset": "smoke",
+      "vars": {
+        "batch": 1,
+        "inChannels": 32,
+        "outChannels": 96,
+        "inH": 16,
+        "inW": 16,
+        "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, 16, 16], "dtype": "float16", "dist": "normal", "seed": 839, "scale": 0.2 },
+        "w": { "shape": [96, 32, 3, 3], "dtype": "float16", "dist": "normal", "seed": 1375, "scale": 0.05 }
+      },
+      "outputs": { "y": { "shape": [1, 96, 16, 16], "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-half-direct-inputs-row-band-m128-k288-n256",
+      "preset": "smoke",
+      "vars": {
+        "batch": 1,
+        "inChannels": 32,
+        "outChannels": 128,
+        "inH": 16,
+        "inW": 16,
+        "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, 16, 16], "dtype": "float16", "dist": "normal", "seed": 839, "scale": 0.2 },
+        "w": { "shape": [128, 32, 3, 3], "dtype": "float16", "dist": "normal", "seed": 1375, "scale": 0.05 }
+      },
+      "outputs": { "y": { "shape": [1, 128, 16, 16], "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": "direct-1x1-float32-m63-k512-n4096",
+      "preset": "model",
+      "vars": {
+        "batch": 1,
+        "inChannels": 512,
+        "outChannels": 63,
+        "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": 354, "scale": 0.2 },
+        "w": { "shape": [63, 512, 1, 1], "dtype": "float32", "dist": "normal", "seed": 355, "scale": 0.02 }
+      },
+      "outputs": { "y": { "shape": [1, 63, 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"
+          }
+        ]
+      },
+      "provenance": {
+        "source": "synthetic",
+        "notes": "Row-tile boundary sweep: full, partial, and device-dependent complete row bands share the same reduction and spatial extent."
+      }
+    },
+    {
+      "name": "direct-1x1-float32-m64-k512-n4096",
+      "preset": "model",
+      "vars": {
+        "batch": 1,
+        "inChannels": 512,
+        "outChannels": 64,
+        "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": 354, "scale": 0.2 },
+        "w": { "shape": [64, 512, 1, 1], "dtype": "float32", "dist": "normal", "seed": 355, "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"
+          }
+        ]
+      },
+      "provenance": {
+        "source": "synthetic",
+        "notes": "Row-tile boundary sweep: full, partial, and device-dependent complete row bands share the same reduction and spatial extent."
+      }
+    },
+    {
+      "name": "direct-1x1-float32-m96-k512-n4096",
+      "preset": "model",
+      "vars": {
+        "batch": 1,
+        "inChannels": 512,
+        "outChannels": 96,
+        "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": 354, "scale": 0.2 },
+        "w": { "shape": [96, 512, 1, 1], "dtype": "float32", "dist": "normal", "seed": 355, "scale": 0.02 }
+      },
+      "outputs": { "y": { "shape": [1, 96, 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"
+          }
+        ]
+      },
+      "provenance": {
+        "source": "synthetic",
+        "notes": "Row-tile boundary sweep: full, partial, and device-dependent complete row bands share the same reduction and spatial extent."
+      }
+    },
+    {
+      "name": "direct-1x1-float16-m63-k512-n4096",
+      "preset": "model",
+      "vars": {
+        "batch": 1,
+        "inChannels": 512,
+        "outChannels": 63,
+        "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": [63, 512, 1, 1], "dtype": "float16", "dist": "normal", "seed": 355, "scale": 0.02 }
+      },
+      "outputs": { "y": { "shape": [1, 63, 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"
+          }
+        ]
+      },
+      "provenance": {
+        "source": "synthetic",
+        "notes": "Row-tile boundary sweep: full, partial, and device-dependent complete row bands share the same reduction and spatial extent."
+      }
+    },
+    {
+      "name": "direct-1x1-float16-m64-k512-n4096",
+      "preset": "model",
+      "vars": {
+        "batch": 1,
+        "inChannels": 512,
+        "outChannels": 64,
+        "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": [64, 512, 1, 1], "dtype": "float16", "dist": "normal", "seed": 355, "scale": 0.02 }
+      },
+      "outputs": { "y": { "shape": [1, 64, 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"
+          }
+        ]
+      },
+      "provenance": {
+        "source": "synthetic",
+        "notes": "Row-tile boundary sweep: full, partial, and device-dependent complete row bands share the same reduction and spatial extent."
+      }
+    },
+    {
+      "name": "direct-1x1-float16-m96-k512-n4096",
+      "preset": "model",
+      "vars": {
+        "batch": 1,
+        "inChannels": 512,
+        "outChannels": 96,
+        "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": [96, 512, 1, 1], "dtype": "float16", "dist": "normal", "seed": 355, "scale": 0.02 }
+      },
+      "outputs": { "y": { "shape": [1, 96, 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"
+          }
+        ]
+      },
+      "provenance": {
+        "source": "synthetic",
+        "notes": "Row-tile boundary sweep: full, partial, and device-dependent complete row bands share the same reduction and spatial extent."
+      }
     }
   ]
 }

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

@@ -1,11 +1,10 @@
-// 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.
 //
@@ -14,33 +13,35 @@
 // 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" %}
-  // Clamped for the reason given in the GEMM activation: an unbounded accumulator
-  // overflows the hardware tanh's internal e^2v to Inf/Inf = NaN, and tanh is
-  // already +/-1 to full f32 precision by |v| ~ 9.
   return tanh(clamp(v, -10.0, 10.0));
-{% 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 %}
@@ -67,19 +68,19 @@ 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 %}
@@ -105,11 +106,48 @@ 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;
   }
@@ -201,7 +239,7 @@ var<workgroup> projectionTile: array<array<vec4<f32>, BN_VECS>, BM>;
       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;
@@ -319,21 +357,41 @@ fn main(
   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") }}
@@ -360,14 +418,22 @@ fn main(
     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] }});
@@ -396,7 +462,7 @@ fn main(
     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;
@@ -409,11 +475,7 @@ fn main(
 {% 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 %}

+// Register-blocked batched matrix product:
+// C[b] (M x N) = A (M x K) * B[b] (K x N). Workgroup lanes compute TM x TN
+// micro-tiles. The defaults use a 16x8 lane grid with an 8x4 micro-tile over a
+// 64x64 output tile. A is shared across the batch; B and C use the dispatch.z
+// batch offset. Accumulation is f32; an optional epilogue applies configured
+// bias, residual, and activation terms in the accumulator domain on store.
+// Every matrix 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.
 //
 // 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.
 {{ env.wgsl.resourceDeclarations }}
 {% set hasActivation = hasActivation is defined and hasActivation %}
+{% set activation = activation | default("") %}
+{% set actAlpha = actAlpha | default(0.0) %}
+{% set actBeta = actBeta | default(0.0) %}
 {% if hasActivation %}
 // Apply the fused activation in the f32 accumulator before the single output
 // cast, avoiding an intermediate convolution tensor.
+{% macro fused_act_return(mode, alpha, beta) -%}
+{% if mode == "Relu" %}
   return max(v, 0.0);
+{% elif mode == "Clip" %}
+  return clamp(v, f32({{ alpha }}), f32({{ beta }}));
+{% elif mode == "LeakyRelu" %}
+  return select(v * f32({{ alpha }}), v, v >= 0.0);
+{% elif mode == "Sigmoid" %}
   return 1.0 / (1.0 + exp(-v));
+{% elif mode == "Tanh" %}
+  // tanh is already +/-1 to f32 precision at the clamp bounds. Clamping before
+  // the builtin preserves that saturated result for larger accumulators.
   return tanh(clamp(v, -10.0, 10.0));
+{% elif mode == "HardSigmoid" %}
+  return clamp(f32({{ alpha }}) * v + f32({{ beta }}), 0.0, 1.0);
 {% else %}
   return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
 {% endif %}
+{%- endmacro -%}
+fn fused_act(v: f32) -> f32 {
+{{ fused_act_return(activation, actAlpha, actBeta) -}}
 }
 {% endif %}
 // 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
+// repeated input loads for writing and rereading an expanded column matrix.
 {% set implicitIm2col = implicitIm2col is defined and implicitIm2col %}
 {% set fusedNarrowProjection = fusedNarrowProjection is defined and fusedNarrowProjection %}
+{% set projectionChannels = projectionOutChannels | default(0) %}
+{% set projectionInputAct = inputActivation | default("none") %}
+{% set projectionOutputAct = outputActivation | default("none") %}
 {% set narrowProjectionTile = "tileB" if GEMM_BK >= GEMM_BM else "projectionTile" %}
 {% set splitKValue = splitK if splitK is defined else 1 %}
 {% set splitKPartial = splitKValue > 1 %}
+{% set kTiles = kTiles | default(0) %}
+{% set mPadded = mPadded | default(0) %}
+{% set nPadded = nPadded | default(0) %}
+{% set batchCount = batchCount | default(0) %}
 {% 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 %}
 const CONV_OUT_W: u32 = {{ convOutW }}u;
 const CONV_IN_CHANNELS: u32 = {{ convInChannels }}u;
 {% endif %}
+{% if splitKPartial %}
+// Split-K publishes raw f32 partials. A second pass owns the bias, optional
+// epilogue, output cast, and logical-shape store.
+const SPLIT_K: u32 = {{ splitKValue }}u;
+const K_TILES: u32 = {{ kTiles }}u;
+const PARTIAL_ROWS: u32 = {{ mPadded }}u;
+const PARTIAL_COLS: u32 = {{ nPadded }}u;
+const PARTIAL_BATCH_STRIDE: u32 = PARTIAL_ROWS * PARTIAL_COLS;
+const PARTIAL_SLICE_STRIDE: u32 = {{ batchCount }}u * PARTIAL_BATCH_STRIDE;
+{% endif %}
 {% if fusedNarrowProjection %}
+{% set inputActivation = inputActivation | default("none") %}
+{% set outputActivation = outputActivation | default("none") %}
+{% if inputActivation == "relu" or outputActivation == "relu" %}
+fn is_nan_f32(value: f32) -> bool {
+  let bits = bitcast<u32>(value);
+  return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
+}
+fn projection_relu(value: f32) -> f32 {
+  var out = max(value, 0.0);
+  if (is_nan_f32(value)) {
+    out = value;
+  }
+  return out;
+}
+{% endif %}
+{% if outputActivation == "sigmoid" %}
+fn sigmoid_safe(x: f32) -> f32 {
+  if (x >= 0.0) {
+    let z = exp(-x);
+    return 1.0 / (1.0 + z);
+  }
+  let z = exp(x);
+  return z / (1.0 + z);
+}
+{% endif %}
 const PROJECTION_OUT_C: u32 = {{ projectionChannels }}u;
+{% else %}{% set emitConvStoreOut = not splitKPartial %}{% set hasZ = hasZ is defined and hasZ %}
+{% if emitConvStoreOut | default(true) %}fn store_out(m: u32, n: u32, yBase: u32, raw: f32) {
   if (m >= params.M || n >= params.N) {
     return;
   }
       let bk = {{ rowBase }} + br;
       var bvec = vec4<{{ tileT }}>({{ tileT }}(0.0));
       if (bk < {{ rowLimit }}) {
+        // k flattens the input-channel and spatial-tap axes; n is the output position.
         let ic = bk / CONV_KSIZE;
         let kq = bk % CONV_KSIZE;
         let kh = kq / CONV_KERNEL_W;
   let N = params.N;
   let mBase = wg.y * BM;
   let nBase = wg.x * BN;
+{% if splitKPartial %}
+  // z carries a K slice within a batch image, matching the subgroup-matrix
+  // Conv split-K layout and the shared reduction pass.
+  let slice = wg.z % SPLIT_K;
+  let batch = wg.z / SPLIT_K;
+{% else %}
   let batch = wg.z;
+{% endif %}
 {% if not implicitIm2col %}
   let xBatchBase = batch * K * N;
 {% endif %}
   let li = lid.y * WG_X + lid.x;
+  // The statically named accumulator vectors form TM independent FMA chains.
 {% for row in range(GEMM_TM) %}
   var acc{{ row }} = vec{{ GEMM_TN }}<f32>(0.0);
 {% endfor %}
+{% if splitKPartial %}
+  // Partition whole BK-wide tiles by quotient boundaries. This keeps every
+  // slice within one tile of every other slice and creates no empty tail slice
+  // whenever K_TILES >= SPLIT_K (this route requires a stronger floor).
+  let kBegin = (slice * K_TILES / SPLIT_K) * BK;
+  let kEnd = min(((slice + 1u) * K_TILES / SPLIT_K) * BK, K);
+  let kCount = kEnd - kBegin;
+  let numTiles = (kCount + BK - 1u) / BK;
+{% else %}
   let numTiles = (K + BK - 1u) / BK;
+{% endif %}
   for (var kt: u32 = 0u; kt < numTiles; kt = kt + 1u) {
+{% if splitKPartial %}
+    let kBase = kBegin + kt * BK;
+{% else %}
     let kBase = kt * BK;
+{% endif %}
 {{ load_a_vec4("mBase", "kBase", "M", gemmKLimit) }}
 {% if useCarriedImplicitGather %}
 {{ load_b_implicit_carried("kBase", "nBase", gemmKLimit, "N") }}
     workgroupBarrier();
   }
+{% if splitKPartial %}
+  let m0 = mBase + lid.y * TM;
+  let n0 = nBase + lid.x * TN;
+  let partialBase = slice * PARTIAL_SLICE_STRIDE + batch * PARTIAL_BATCH_STRIDE;
+{% else %}
 {% if not fusedNarrowProjection %}
   let yBatchBase = batch * M * N;
 {% endif %}
   let m0 = mBase + lid.y * TM;
   let n0 = nBase + lid.x * TN;
+{% endif %}
 {% for row in range(GEMM_TM) %}
 {% for column in range(GEMM_TN) %}
+{% if splitKPartial %}
+  y[partialBase + (m0 + {{ row }}u) * PARTIAL_COLS + n0 + {{ column }}u] = acc{{ row }}.{{ components[column] }};
+{% elif 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] }});
     let projectionYBase = batch * PROJECTION_OUT_C * N;
 {% for oc in range(projectionChannels) %}
+    let projected{{ oc }} = projectionAcc{{ oc }};
 {% endfor %}
 {% for column in range(4) %}
     let projectionN{{ column }} = projectionNBase + {{ column }}u;
 {% else %}
       let activated{{ oc }}_{{ column }} = projected{{ oc }}.{{ components[column] }};
 {% endif %}
       y[projectionYBase + {{ oc }}u * N + projectionN{{ column }}] = activated{{ oc }}_{{ column }};
 {% endfor %}
     }
 {% endfor %}

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

@@ -1,46 +1,39 @@
-// 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" %}
-  // Clamped for the reason given in the GEMM activation: an unbounded accumulator
-  // overflows the hardware tanh's internal e^2v to Inf/Inf = NaN, and tanh is
-  // already +/-1 to full f32 precision by |v| ~ 9.
   return tanh(clamp(v, -10.0, 10.0));
-{% 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 %}
@@ -49,8 +42,9 @@ 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;

+// Tiled batched matrix product C[b] (M x N) = A (M x K) * B[b] (K x N).
+// A 16x16 workgroup computes a 32x32 output tile with a 2x2 micro-tile per
+// thread and a 16-element K tile. A is shared across the batch; B and C use the
+// dispatch.z batch offset. Accumulation is f32 and is cast to the output type on
+// store. A configured epilogue may add bias, residual, and activation terms.
+// f16 operands remain packed in workgroup memory and widen when consumed. All
+// M, N, and K accesses are bounds-checked.
 {{ env.wgsl.resourceDeclarations }}
 {% set hasActivation = hasActivation is defined and hasActivation %}
+{% set activation = activation | default("") %}
+{% set actAlpha = actAlpha | default(0.0) %}
+{% set actBeta = actBeta | default(0.0) %}
 {% if hasActivation %}
 // Apply the fused activation in the f32 accumulator before the single output
 // cast, avoiding an intermediate convolution tensor.
+{% macro fused_act_return(mode, alpha, beta) -%}
+{% if mode == "Relu" %}
   return max(v, 0.0);
+{% elif mode == "Clip" %}
+  return clamp(v, f32({{ alpha }}), f32({{ beta }}));
+{% elif mode == "LeakyRelu" %}
+  return select(v * f32({{ alpha }}), v, v >= 0.0);
+{% elif mode == "Sigmoid" %}
   return 1.0 / (1.0 + exp(-v));
+{% elif mode == "Tanh" %}
+  // tanh is already +/-1 to f32 precision at the clamp bounds. Clamping before
+  // the builtin preserves that saturated result for larger accumulators.
   return tanh(clamp(v, -10.0, 10.0));
+{% elif mode == "HardSigmoid" %}
+  return clamp(f32({{ alpha }}) * v + f32({{ beta }}), 0.0, 1.0);
 {% else %}
   return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
 {% endif %}
+{%- endmacro -%}
+fn fused_act(v: f32) -> f32 {
+{{ fused_act_return(activation, actAlpha, actBeta) -}}
 }
 {% endif %}
 const BM: u32 = 32u;
 const BN: u32 = 32u;
+// Store one output element with the configured bias/residual/activation
+// epilogue in the f32 accumulator domain.
+{% set hasZ = hasZ is defined and hasZ %}
 fn store_out(m: u32, n: u32, yBase: u32, raw: f32) {
   if (m >= params.M || n >= params.N) {
     return;

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

@@ -1,6 +1,7 @@
-// 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
@@ -10,15 +11,41 @@
 // 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;
@@ -26,7 +53,35 @@ enable subgroup_size_control;
 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
 {% set operandScalar = fScalar %}
 {% set accScalar = "f32" %}
@@ -43,29 +98,78 @@ diagnostic(off, chromium.subgroup_matrix_uniformity);
 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;
@@ -91,6 +195,126 @@ fn loadSHMA(tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
 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" %}
@@ -99,31 +323,38 @@ fn loadSHMB(b_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
     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" %}
-  // Clamped for the reason given in the GEMM activation: an unbounded accumulator
-  // overflows the hardware tanh's internal e^2v to Inf/Inf = NaN, and tanh is
-  // already +/-1 to full f32 precision by |v| ~ 9.
   return tanh(clamp(v, -10.0, 10.0));
-{% 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 %}
@@ -191,10 +422,18 @@ fn main(
 {% 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;
@@ -213,21 +452,30 @@ fn main(
   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>, row_major>(&tile_A, matrix_a_offset, TILE_K);
-      var matA1: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset + 8u * TILE_K, 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>, col_major>(&tile_B, matrix_b_offset, TILE_K);
-      var matB1: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 8u * TILE_K, TILE_K);
-      var matB2: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 16u * TILE_K, TILE_K);
-      var matB3: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 24u * TILE_K, TILE_K);
       matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
       matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);
@@ -238,7 +486,9 @@ fn main(
       matC12 = subgroupMatrixMultiplyAccumulate(matA1, matB2, matC12);
       matC13 = subgroupMatrixMultiplyAccumulate(matA1, matB3, matC13);
     }
     workgroupBarrier();
   }
 {% if fusedNarrowProjection %}
@@ -268,7 +518,11 @@ fn main(
 {% 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 %}
@@ -276,11 +530,30 @@ fn main(
     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.

+// Batched subgroup-matrix product C[b] (M x N) = A (M x K) * B[b] (K x N).
+// A is row-major; B may be a direct matrix view, a materialized column matrix,
+// or an on-demand gather. Accumulation is f32. A configured epilogue may apply
+// bias or other output transforms after the full K reduction.
 //
 // Full 8x8 matrix loads require complete K/N tiles. The aligned path enforces
 // that geometry; `padded` zero-fills tails, and `implicitIm2col` gathers and
 // 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 directMatrixInputs = directMatrixInputs is defined and directMatrixInputs %}
 {% set fusedNarrowProjection = fusedNarrowProjection if fusedNarrowProjection is defined else false %}
 {% set polyphase = polyphaseConvTranspose is defined and polyphaseConvTranspose %}
+{% set implicit = implicitIm2col is defined and implicitIm2col %}
+{% set implicit3d = implicit and (convSpatialDims is defined and convSpatialDims == 3) %}
 {% 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" %}
+{% set hasZ = hasZ is defined and hasZ %}
+{% set projectionOutChannels = projectionOutChannels | default(0) %}
+{% set hasProjectionBias = hasProjectionBias is defined and hasProjectionBias %}
+{% set hasOutputScale = hasOutputScale is defined and hasOutputScale %}
+{% set convKernelH = convKernelH | default(0) %}
+{% set convKernelW = convKernelW | default(0) %}
+{% set convStrideH = convStrideH | default(0) %}
+{% set convStrideW = convStrideW | default(0) %}
+{% set convDilationH = convDilationH | default(0) %}
+{% set convDilationW = convDilationW | default(0) %}
+{% set convPadTop = convPadTop | default(0) %}
+{% set convPadLeft = convPadLeft | default(0) %}
+{% set convInH = convInH | default(0) %}
+{% set convInW = convInW | default(0) %}
+{% set convOutW = convOutW | default(0) %}
+{% set convInChannels = convInChannels | default(0) %}
+{% set convKernelD = convKernelD | default(0) %}
+{% set convStrideD = convStrideD | default(0) %}
+{% set convDilationD = convDilationD | default(0) %}
+{% set convPadFront = convPadFront | default(0) %}
+{% set convInD = convInD | default(0) %}
+{% set convOutH = convOutH | default(0) %}
+{% set nPadded = nPadded | default(0) %}
+{% set mPadded = mPadded | default(0) %}
+{% set kChunk = kChunk | default(0) %}
+{% set batchCount = batchCount | default(0) %}
 enable subgroups;
 {% if pinSubgroupSize32 %}
 enable subgroup_size_control;
 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
+{% if fusedNarrowProjection %}{% set inputActivation = inputActivation | default("none") %}
+{% set outputActivation = outputActivation | default("none") %}
+{% if inputActivation == "relu" or outputActivation == "relu" %}
+fn is_nan_f32(value: f32) -> bool {
+  let bits = bitcast<u32>(value);
+  return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
+}
+fn projection_relu(value: f32) -> f32 {
+  var out = max(value, 0.0);
+  if (is_nan_f32(value)) {
+    out = value;
+  }
+  return out;
+}
+{% endif %}
+{% if outputActivation == "sigmoid" %}
+fn sigmoid_safe(x: f32) -> f32 {
+  if (x >= 0.0) {
+    let z = exp(-x);
+    return 1.0 / (1.0 + z);
+  }
+  let z = exp(x);
+  return z / (1.0 + z);
+}
+{% endif %}
+{% endif %}
 {% set operandScalar = fScalar %}
 {% set accScalar = "f32" %}
 const M: u32 = {{ M }}u;
 const K: u32 = {{ K }}u;
+{% if not (splitKPartial and padded) or implicit %}
 const N: u32 = {{ N }}u;
+{% endif %}
 {% if padded %}
 const K_LOOP: u32 = {{ kPadded }}u;
 {% endif %}
+{% if implicit %}
+const CONV_KERNEL_H: u32 = {{ convKernelH }}u;
+const CONV_KERNEL_W: u32 = {{ convKernelW }}u;
+{% if implicit3d %}
+const CONV_KERNEL_D: u32 = {{ convKernelD }}u;
+const CONV_KSIZE_HW: u32 = CONV_KERNEL_H * CONV_KERNEL_W;
+const CONV_KSIZE: u32 = CONV_KERNEL_D * CONV_KSIZE_HW;
+const CONV_STRIDE_D: u32 = {{ convStrideD }}u;
+const CONV_DILATION_D: u32 = {{ convDilationD }}u;
+const CONV_PAD_FRONT: i32 = {{ convPadFront }};
+const CONV_IN_D: i32 = {{ convInD }};
+const CONV_OUT_H: u32 = {{ convOutH }}u;
+{% else %}
+const CONV_KSIZE: u32 = CONV_KERNEL_H * CONV_KERNEL_W;
+{% endif %}
+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: i32 = {{ convInH }};
+const CONV_IN_W: i32 = {{ convInW }};
+const CONV_OUT_W: u32 = {{ convOutW }}u;
+{% if implicit3d %}
+// X is [batch, inChannels, inD, inH, inW]; the tile gather indexes it directly.
+const B_BATCH_STRIDE: u32 = {{ convInChannels }}u * u32(CONV_IN_D) * u32(CONV_IN_H) * u32(CONV_IN_W);
+{% else %}
+// X is [batch, inChannels, inH, inW]; the tile gather indexes it directly.
+const B_BATCH_STRIDE: u32 = {{ convInChannels }}u * u32(CONV_IN_H) * u32(CONV_IN_W);
+{% endif %}
+{% else %}
 {% if padded %}
 const N_COLS: u32 = {{ nPadded }}u;
 {% endif %}
 const B_BATCH_STRIDE: u32 = {{ kLoopVar }} * {{ nColsVar }};
+{% endif %}
 {% if not splitKPartial and not polyphase and not fusedNarrowProjection %}
 const C_BATCH_STRIDE: u32 = M * N;
 {% endif %}
+{% if splitKPartial %}
+const SPLIT_K: u32 = {{ splitKValue }}u;
+// Whole 32-wide K tiles per slice, rounded up, with the last slice clamped to
+// K_LOOP. This avoids padding K to a multiple of SPLIT_K times the tile width.
+// A slice entirely past K_LOOP runs zero iterations and stores zero
+// accumulators, which is the correct partial for an empty K range.
+const K_CHUNK: u32 = {{ kChunk }}u;
+const PARTIAL_ROWS: u32 = {{ mPadded }}u;
+const PARTIAL_COLS: u32 = {{ nPadded }}u;
+const PARTIAL_SLICE_STRIDE: u32 = {{ batchCount }}u * PARTIAL_ROWS * PARTIAL_COLS;
+{% 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;
+{% if not directMatrixInputs %}
 var<workgroup> tile_A: array<{{ operandScalar }}, {{ tileRowsValue }} * 32>;
 var<workgroup> tile_B: array<{{ operandScalar }}, {{ tileColsValue }} * 32>;
+{% endif %}
 {% if (not useDirectMatrixStore or padded) and not splitKPartial and not fusedNarrowProjection %}
 var<workgroup> scratch: array<array<array<{{ accScalar }}, 64>, 4>, {{ (workgroupThreadsValue / 32)|int }}>;
 {% endif %}
+{% if not directMatrixInputs %}
 fn loadSHMA(tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
   let a_global = tile_base + row;
   let col = c_idx * 8u;
 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;
+{% if implicit %}
+  // One output position per tile column, so its window origin is loop-invariant.
+{% if implicit3d %}
+  let id0 = i32((b_col / (CONV_OUT_W * CONV_OUT_H)) * CONV_STRIDE_D) - CONV_PAD_FRONT;
+  let ih0 = i32(((b_col / CONV_OUT_W) % CONV_OUT_H) * CONV_STRIDE_H) - CONV_PAD_TOP;
+{% else %}
+  let ih0 = i32((b_col / CONV_OUT_W) * CONV_STRIDE_H) - CONV_PAD_TOP;
+{% endif %}
+  let iw0 = i32((b_col % CONV_OUT_W) * CONV_STRIDE_W) - CONV_PAD_LEFT;
+  let in_column = b_col < N;
+  // k advances by exactly one per element, so the (in-channel, tap-row, tap-col)
+  // decomposition is CARRIED across the loop instead of recomputed: two divisions
+  // and two moduli per loaded element become one increment and one compare, with
+  // the source coordinates advanced by the dilation in step. Only the entry values
+  // divide, and their divisors are compile-time literals.
+  var k = k_idx + col;
+  var ic = k / CONV_KSIZE;
+  var kq = k % CONV_KSIZE;
+{% if implicit3d %}
+  var kd = kq / CONV_KSIZE_HW;
+  var khw = kq % CONV_KSIZE_HW;
+  var kh = khw / CONV_KERNEL_W;
+  var kw = khw % CONV_KERNEL_W;
+  var id = id0 + i32(kd * CONV_DILATION_D);
+{% else %}
+  var kh = kq / CONV_KERNEL_W;
+  var kw = kq % CONV_KERNEL_W;
+{% endif %}
+  var ih = ih0 + i32(kh * CONV_DILATION_H);
+  var iw = iw0 + i32(kw * CONV_DILATION_W);
+  // A tile column whose whole kernel window and K span are in bounds cannot
+  // take any of the element guards below. Hoist them into one uniform test so
+  // the interior loop carries a plain address instead of recomputing coordinates.
+  let interior = in_column
+    && k_idx + col + {{ bLoadWidth }}u <= K
+{% if implicit3d %}
+    && id0 >= 0 && id0 + i32((CONV_KERNEL_D - 1u) * CONV_DILATION_D) < CONV_IN_D
+{% endif %}
+    && ih0 >= 0 && ih0 + i32((CONV_KERNEL_H - 1u) * CONV_DILATION_H) < CONV_IN_H
+    && iw0 >= 0 && iw0 + i32((CONV_KERNEL_W - 1u) * CONV_DILATION_W) < CONV_IN_W;
+  if (interior) {
+    let colStep = i32(CONV_DILATION_W);
+    let rowStep = i32(CONV_DILATION_H) * CONV_IN_W;
+{% if implicit3d %}
+    let depthStep = i32(CONV_DILATION_D) * CONV_IN_H * CONV_IN_W;
+    let planeStep = CONV_IN_D * CONV_IN_H * CONV_IN_W;
+{% else %}
+    let planeStep = CONV_IN_H * CONV_IN_W;
+{% endif %}
+    // Advancing k already added one colStep, so each wrap rewinds that plus the taps it walked.
+    let kwWrap = colStep * i32(CONV_KERNEL_W) - rowStep;
+{% if implicit3d %}
+    let khWrap = rowStep * i32(CONV_KERNEL_H) - depthStep;
+    let kdWrap = depthStep * i32(CONV_KERNEL_D) - planeStep;
+    var addr = i32(b_base) + i32(ic) * planeStep + (id * CONV_IN_H + ih) * CONV_IN_W + iw;
+{% else %}
+    let khWrap = rowStep * i32(CONV_KERNEL_H) - planeStep;
+    var addr = i32(b_base) + i32(ic) * planeStep + ih * CONV_IN_W + iw;
+{% endif %}
+    for (var i = 0u; i < {{ bLoadWidth }}u; i = i + 1u) {
+      tile_B[row * TILE_K + col + i] = {{ operandScalar }}(xm[u32(addr)]);
+      addr = addr + colStep;
+      kw = kw + 1u;
+      if (kw == CONV_KERNEL_W) {
+        kw = 0u;
+        addr = addr - kwWrap;
+        kh = kh + 1u;
+        if (kh == CONV_KERNEL_H) {
+          kh = 0u;
+          addr = addr - khWrap;
+{% if implicit3d %}
+          kd = kd + 1u;
+          if (kd == CONV_KERNEL_D) {
+            kd = 0u;
+            addr = addr - kdWrap;
+          }
+{% endif %}
+        }
+      }
+    }
+    return;
+  }
+  for (var i = 0u; i < {{ bLoadWidth }}u; i = i + 1u) {
+    var value = {{ operandScalar }}(0.0);
+    // The K and N tails read as exact zeros, exactly as the materialized padded
+    // cols buffer does, so they contribute nothing to the dot product.
+    if (in_column && k < K{% if implicit3d %} && id >= 0 && id < CONV_IN_D{% endif %} && ih >= 0 && ih < CONV_IN_H && iw >= 0 && iw < CONV_IN_W) {
+{% if implicit3d %}
+      value = {{ operandScalar }}(xm[b_base + ((ic * u32(CONV_IN_D) + u32(id)) * u32(CONV_IN_H) + u32(ih)) * u32(CONV_IN_W) + u32(iw)]);
+{% else %}
+      value = {{ operandScalar }}(xm[b_base + (ic * u32(CONV_IN_H) + u32(ih)) * u32(CONV_IN_W) + u32(iw)]);
+{% endif %}
+    }
+    tile_B[row * TILE_K + col + i] = value;
+    k = k + 1u;
+    kw = kw + 1u;
+    iw = iw + i32(CONV_DILATION_W);
+    if (kw == CONV_KERNEL_W) {
+      kw = 0u;
+      iw = iw0;
+      kh = kh + 1u;
+      ih = ih + i32(CONV_DILATION_H);
+      if (kh == CONV_KERNEL_H) {
+        kh = 0u;
+        ih = ih0;
+{% if implicit3d %}
+        kd = kd + 1u;
+        id = id + i32(CONV_DILATION_D);
+        if (kd == CONV_KERNEL_D) {
+          kd = 0u;
+          id = id0;
+          ic = ic + 1u;
+        }
+{% else %}
+        ic = ic + 1u;
+{% endif %}
+      }
+    }
+  }
+{% else %}
   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] = f32(xm[b_base + k * {{ nColsVar }} + b_col]);
 {% endif %}
   }
+{% endif %}
 }
+{% endif %}
 {% set hasActivation = hasActivation is defined and hasActivation %}
+{% set activation = activation | default("") %}
+{% set actAlpha = actAlpha | default(0.0) %}
+{% set actBeta = actBeta | default(0.0) %}
 {% if hasActivation %}
 // Apply the fused activation in the f32 accumulator before the single output
 // cast, avoiding an intermediate convolution tensor.
+{% macro fused_act_return(mode, alpha, beta) -%}
+{% if mode == "Relu" %}
   return max(v, 0.0);
+{% elif mode == "Clip" %}
+  return clamp(v, f32({{ alpha }}), f32({{ beta }}));
+{% elif mode == "LeakyRelu" %}
+  return select(v * f32({{ alpha }}), v, v >= 0.0);
+{% elif mode == "Sigmoid" %}
   return 1.0 / (1.0 + exp(-v));
+{% elif mode == "Tanh" %}
+  // tanh is already +/-1 to f32 precision at the clamp bounds. Clamping before
+  // the builtin preserves that saturated result for larger accumulators.
   return tanh(clamp(v, -10.0, 10.0));
+{% elif mode == "HardSigmoid" %}
+  return clamp(f32({{ alpha }}) * v + f32({{ beta }}), 0.0, 1.0);
 {% else %}
   return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
 {% endif %}
+{%- endmacro -%}
+fn fused_act(v: f32) -> f32 {
+{{ fused_act_return(activation, actAlpha, actBeta) -}}
 }
 {% endif %}
 {% endif %}
   @builtin(subgroup_size) sg_size: u32
 ) {
+{% if splitKPartial %}
+  // z carries both axes, with split slices contiguous inside each batch image.
+  let slice = workgroup_id.z % SPLIT_K;
+  let batch = workgroup_id.z / SPLIT_K;
+  let b_base = batch * B_BATCH_STRIDE;
+  let c_base = slice * PARTIAL_SLICE_STRIDE + batch * PARTIAL_ROWS * PARTIAL_COLS;
+{% else %}
   let batch = workgroup_id.z;
   let b_base = batch * B_BATCH_STRIDE;
 {% if not fusedNarrowProjection %}
   let c_base = batch * C_BATCH_STRIDE;
+{% endif %}
 {% endif %}
   let a_global_base = workgroup_id.y * TILE_ROWS;
   let b_global_base = workgroup_id.x * TILE_COLS;
   var matC12: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
   var matC13: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
+{% if splitKPartial %}
+  let k_begin = slice * K_CHUNK;
+  let k_end = min(k_begin + K_CHUNK, {{ kLoopVar }});
+  for (var kidx = k_begin; kidx < k_end; kidx = kidx + TILE_K) {
+{% else %}
   for (var kidx = 0u; kidx < {{ kLoopVar }}; kidx = kidx + TILE_K) {
+{% endif %}
+{% if not directMatrixInputs %}
     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();
+{% endif %}
     for (var step = 0u; step < TILE_K; step = step + 8u) {
+{% set directInputs = directMatrixInputs is defined and directMatrixInputs %}
+      let matrix_a_offset = {% if directInputs %}(a_global_base + subtile_idy * SUB_ROWS) * K + kidx + step{% else %}subtile_idy * SUB_ROWS * TILE_K + step{% endif %};
+{% for r in range(2) %}
+      var matA{{ r }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&{{ "w" if directInputs else "tile_A" }}, matrix_a_offset{% if r > 0 %} + 8u * {{ "K" if directInputs else "TILE_K" }}{% endif %}, {{ "K" if directInputs else "TILE_K" }});
+{% endfor %}
+      let matrix_b_offset = {% if directInputs %}b_base + (kidx + step) * N + b_global_base + subtile_idx * SUB_COLS{% else %}subtile_idx * SUB_COLS * TILE_K + step{% endif %};
+{% for c in range(4) %}
+      var matB{{ c }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, {% if directInputs %}row_major{% else %}col_major{% endif %}>(&{{ "xm" if directInputs else "tile_B" }}, matrix_b_offset{% if c > 0 %} + {{ c * 8 }}u{% if not directInputs %} * TILE_K{% endif %}{% endif %}, {{ "N" if directInputs else "TILE_K" }});
+{% endfor %}
       matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
       matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);
       matC12 = subgroupMatrixMultiplyAccumulate(matA1, matB2, matC12);
       matC13 = subgroupMatrixMultiplyAccumulate(matA1, matB3, matC13);
     }
+{% if not directMatrixInputs %}
     workgroupBarrier();
+{% endif %}
   }
 {% if fusedNarrowProjection %}
 {% 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 %};
+{% if inputActivation == "relu" %}
+      let activated_value = projection_relu(producer_value);
+{% else %}
       let activated_value = producer_value;
+{% endif %}
 {% 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 }}{% if hasProjectionBias %} + f32(projectionBias[{{ oc }}u]){% endif %};
+{% if outputActivation == "relu" %}
+    let activated{{ oc }} = projection_relu(projected{{ oc }});
+{% elif outputActivation == "sigmoid" %}
+    let activated{{ oc }} = sigmoid_safe(projected{{ oc }});
+{% else %}
     let activated{{ oc }} = projected{{ oc }};
+{% endif %}
+    y[fused_y_base + {{ oc }}u * N] = activated{{ oc }}{% if hasOutputScale %} * params.outputScale{% endif %};
 {% endfor %}
   }
+{% elif splitKPartial %}
+  // Raw partials, no bias and no epilogue — the reduce pass owns both. The
+  // scratch is whole tiles in both axes, so this needs none of the column or
+  // row guards the real output store carries.
+  let partial_offset = c_base + (a_global_base + base_A) * PARTIAL_COLS + b_global_base + base_B;
+  subgroupMatrixStore<row_major>(&y, partial_offset + 0u * PARTIAL_COLS +  0u, matC00, PARTIAL_COLS);
+  subgroupMatrixStore<row_major>(&y, partial_offset + 0u * PARTIAL_COLS +  8u, matC01, PARTIAL_COLS);
+  subgroupMatrixStore<row_major>(&y, partial_offset + 0u * PARTIAL_COLS + 16u, matC02, PARTIAL_COLS);
+  subgroupMatrixStore<row_major>(&y, partial_offset + 0u * PARTIAL_COLS + 24u, matC03, PARTIAL_COLS);
+  subgroupMatrixStore<row_major>(&y, partial_offset + 8u * PARTIAL_COLS +  0u, matC10, PARTIAL_COLS);
+  subgroupMatrixStore<row_major>(&y, partial_offset + 8u * PARTIAL_COLS +  8u, matC11, PARTIAL_COLS);
+  subgroupMatrixStore<row_major>(&y, partial_offset + 8u * PARTIAL_COLS + 16u, matC12, PARTIAL_COLS);
+  subgroupMatrixStore<row_major>(&y, partial_offset + 8u * PARTIAL_COLS + 24u, matC13, PARTIAL_COLS);
 {% elif useDirectMatrixStore %}
 {% if padded %}
   // Padded N only needs guarded publication for the final partial-width tile.

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

@@ -1,42 +1,47 @@
-// 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" %}
-  // Clamped for the reason given in the GEMM activation: an unbounded accumulator
-  // overflows the hardware tanh's internal e^2v to Inf/Inf = NaN, and tanh is
-  // already +/-1 to full f32 precision by |v| ~ 9.
   return tanh(clamp(v, -10.0, 10.0));
-{% 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;
   }
@@ -52,7 +57,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>,
   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) {

+// Direct N-dimensional 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 usesF16Spec %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
 {% set hasActivation = hasActivation is defined and hasActivation %}
+{% set activation = activation | default("") %}
+{% set actAlpha = actAlpha | default(0.0) %}
+{% set actBeta = actBeta | default(0.0) %}
 {% if hasActivation %}
 // Apply the fused activation in the f32 accumulator before the single output
 // cast, avoiding an intermediate convolution tensor.
+{% macro fused_act_return(mode, alpha, beta) -%}
+{% if mode == "Relu" %}
   return max(v, 0.0);
+{% elif mode == "Clip" %}
+  return clamp(v, f32({{ alpha }}), f32({{ beta }}));
+{% elif mode == "LeakyRelu" %}
+  return select(v * f32({{ alpha }}), v, v >= 0.0);
+{% elif mode == "Sigmoid" %}
   return 1.0 / (1.0 + exp(-v));
+{% elif mode == "Tanh" %}
+  // tanh is already +/-1 to f32 precision at the clamp bounds. Clamping before
+  // the builtin preserves that saturated result for larger accumulators.
   return tanh(clamp(v, -10.0, 10.0));
+{% elif mode == "HardSigmoid" %}
+  return clamp(f32({{ alpha }}) * v + f32({{ beta }}), 0.0, 1.0);
 {% else %}
   return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
 {% endif %}
+{%- endmacro -%}
+fn fused_act(v: f32) -> f32 {
+{{ fused_act_return(activation, actAlpha, actBeta) -}}
 }
 {% endif %}
+{% set hasZ = hasZ is defined and hasZ %}
 const WG: u32 = {{ convWorkgroupSize }}u;
 @compute @workgroup_size(WG)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
   if (index >= params.count) {
     return;
   }
   let group = oc / params.outChannelsPerGroup;
   var acc = 0.0;
+{% if hasBias %}
   acc = f32(bias[oc]);
 {% endif %}
   for (var ic = 0u; ic < params.weightInChannels; ic += 1u) {

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

@@ -1,53 +1,55 @@
 // 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" %}
-  // Clamped for the reason given in the GEMM activation: an unbounded accumulator
-  // overflows the hardware tanh's internal e^2v to Inf/Inf = NaN, and tanh is
-  // already +/-1 to full f32 precision by |v| ~ 9.
   return tanh(clamp(v, -10.0, 10.0));
-{% 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;
   }
@@ -71,16 +73,16 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
   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 %}

 // Direct NCHW 2D convolution with kernel extent, strides, dilations, and
 // top/left pads compiled into the fully unrolled window below.
+// One thread owns each output element and accumulates with input channels outer,
+// then kernel rows and columns in ascending order.
 // 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.
+// Group offsets are applied to both input channels and weights. f16 inputs are
+// widened to an f32 accumulator and narrowed once at store.
 {{ env.wgsl.resourceDeclarations }}
 {% set hasActivation = hasActivation is defined and hasActivation %}
+{% set activation = activation | default("") %}
+{% set actAlpha = actAlpha | default(0.0) %}
+{% set actBeta = actBeta | default(0.0) %}
 {% if hasActivation %}
 // Apply the fused activation in the f32 accumulator before the single output
 // cast, avoiding an intermediate convolution tensor.
+{% macro fused_act_return(mode, alpha, beta) -%}
+{% if mode == "Relu" %}
   return max(v, 0.0);
+{% elif mode == "Clip" %}
+  return clamp(v, f32({{ alpha }}), f32({{ beta }}));
+{% elif mode == "LeakyRelu" %}
+  return select(v * f32({{ alpha }}), v, v >= 0.0);
+{% elif mode == "Sigmoid" %}
   return 1.0 / (1.0 + exp(-v));
+{% elif mode == "Tanh" %}
+  // tanh is already +/-1 to f32 precision at the clamp bounds. Clamping before
+  // the builtin preserves that saturated result for larger accumulators.
   return tanh(clamp(v, -10.0, 10.0));
+{% elif mode == "HardSigmoid" %}
+  return clamp(f32({{ alpha }}) * v + f32({{ beta }}), 0.0, 1.0);
 {% else %}
   return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
 {% endif %}
+{%- endmacro -%}
+fn fused_act(v: f32) -> f32 {
+{{ fused_act_return(activation, actAlpha, actBeta) -}}
 }
 {% endif %}
+{% set hasZ = hasZ is defined and hasZ %}
+const KERNEL_AREA: u32 = {{ kernelHSpec * kernelWSpec }}u;
+const STRIDE_H: u32 = {{ strideHSpec }}u;
+const STRIDE_W: u32 = {{ strideWSpec }}u;
+const PAD_TOP: i32 = {{ padTopSpec }};
+const PAD_LEFT: i32 = {{ padLeftSpec }};
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   // 2D-folded flat index: gid.y carries the high bits past the
+  // per-axis dispatch fold width (outputs > 16.7M elements).
+  let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
   if (index >= params.count) {
     return;
   }
   acc = f32(bias[oc]);
 {% endif %}
   for (var ic = 0u; ic < params.weightInChannels; ic = ic + 1u) {
+{% for kh in range(kernelHSpec) %}
     {
+      let ih = ihBase + {{ kh * dilationHSpec }}i;
       if (ih >= 0 && ih < i32(params.inH)) {
         let xRow = xBase + u32(ih) * params.inW;
+{% for kw in range(kernelWSpec) %}
         {
+          let iw = iwBase + {{ kw * dilationWSpec }}i;
           if (iw >= 0 && iw < i32(params.inW)) {
+            acc = acc + f32(x[xRow + u32(iw)]) * f32(w[wBase + {{ kh * kernelWSpec + kw }}u]);
           }
         }
 {% endfor %}

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

@@ -6,20 +6,19 @@
 // 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
@@ -64,10 +63,10 @@ fn main(
     }
 {% 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;
   }
 }

 // 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.
+// A plan for f16 inputs may retain f16 columns to avoid that widening.
+{% set colsScalar = colsScalar | default("f32") %}
 {{ env.wgsl.resourceDeclarations }}
+const KERNEL_H: u32 = {{ kernelHSpec }}u;
+const KERNEL_W: u32 = {{ kernelWSpec }}u;
 const KSIZE: u32 = KERNEL_H * KERNEL_W;
+const STRIDE_H: u32 = {{ strideHSpec }}u;
+const STRIDE_W: u32 = {{ strideWSpec }}u;
+const DILATION_H: u32 = {{ dilationHSpec }}u;
+const DILATION_W: u32 = {{ dilationWSpec }}u;
+const PAD_TOP: i32 = {{ padTopSpec }};
+const PAD_LEFT: i32 = {{ padLeftSpec }};
 {% 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
     }
 {% if padded %}
     }
 {% endif %}
+{% set columnRows = "K_PADDED" if padded else "params.kRows" %}
+{% set columnStride = "COLS_STRIDE" if padded else "params.outCount" %}
+    cols[(batch * {{ columnRows }} + k) * {{ columnStride }} + pos] = {% if colsScalar == "f16" %}f16(value){% else %}value{% endif %};
     pos = pos + step;
   }
 }

build/webgpu/conv-splitk-reduce.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,83 @@

+// Pass 2 of the split-K Conv GEMM: sum the SPLIT_K raw partial slices that
+// a Conv GEMM partial pass wrote, then apply the bias and epilogue
+// that pass deliberately skipped, and publish the logical [batch, M, N] output.
+//
+// The partial scratch is [SPLIT_K, batch, PARTIAL_ROWS, PARTIAL_COLS] with both
+// trailing axes padded to whole GEMM tiles; this pass indexes only the logical
+// (batch, m, n), so the padded rows and columns are read by nobody and their
+// contents never matter. One output element per thread, walking n fastest so
+// consecutive threads touch consecutive addresses in every slice.
+//
+// Summing slices in index order changes the f32 association relative to one
+// uninterrupted K loop, so the two orders need not be bit-identical.
+{{ env.wgsl.resourceDeclarations }}
+{% set applyActivation = hasActivation is defined and hasActivation %}
+{% if applyActivation %}
+{% set hasActivation = hasActivation is defined and hasActivation %}
+{% set activation = activation | default("") %}
+{% set actAlpha = actAlpha | default(0.0) %}
+{% set actBeta = actBeta | default(0.0) %}
+{% if hasActivation %}
+// Apply the fused activation in the f32 accumulator before the single output
+// cast, avoiding an intermediate convolution tensor.
+{% macro fused_act_return(mode, alpha, beta) -%}
+{% if mode == "Relu" %}
+  return max(v, 0.0);
+{% elif mode == "Clip" %}
+  return clamp(v, f32({{ alpha }}), f32({{ beta }}));
+{% elif mode == "LeakyRelu" %}
+  return select(v * f32({{ alpha }}), v, v >= 0.0);
+{% elif mode == "Sigmoid" %}
+  return 1.0 / (1.0 + exp(-v));
+{% elif mode == "Tanh" %}
+  // tanh is already +/-1 to f32 precision at the clamp bounds. Clamping before
+  // the builtin preserves that saturated result for larger accumulators.
+  return tanh(clamp(v, -10.0, 10.0));
+{% elif mode == "HardSigmoid" %}
+  return clamp(f32({{ alpha }}) * v + f32({{ beta }}), 0.0, 1.0);
+{% else %}
+  return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
+{% endif %}
+{%- endmacro -%}
+fn fused_act(v: f32) -> f32 {
+{{ fused_act_return(activation, actAlpha, actBeta) -}}
+}
+{% endif %}
+{% endif %}
+const M: u32 = {{ M }}u;
+const N: u32 = {{ N }}u;
+const SPLIT_K: u32 = {{ splitK }}u;
+const PARTIAL_ROWS: u32 = {{ mPadded }}u;
+const PARTIAL_COLS: u32 = {{ nPadded }}u;
+const PARTIAL_SLICE_STRIDE: u32 = {{ batchCount }}u * PARTIAL_ROWS * PARTIAL_COLS;
+const COUNT: u32 = {{ batchCount }}u * M * N;
+const WORKGROUP_SIZE: u32 = {{ reduceWorkgroupSize }}u;
+@compute @workgroup_size({{ reduceWorkgroupSize }}, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  // 2D-folded flat index: gid.y carries the high bits past
+  // the per-axis dispatch fold width and reduces to the 1D form at y=0.
+  let idx = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WORKGROUP_SIZE;
+  if (idx >= COUNT) {
+    return;
+  }
+  let col = idx % N;
+  let row = (idx / N) % M;
+  let batch = idx / (M * N);
+  var src = (batch * PARTIAL_ROWS + row) * PARTIAL_COLS + col;
+  var acc = 0.0;
+  for (var z = 0u; z < SPLIT_K; z = z + 1u) {
+    acc = acc + partial[src];
+    src = src + PARTIAL_SLICE_STRIDE;
+  }
+{% if hasBias %}
+  acc = acc + f32(bias[row]);
+{% endif %}
+{% if applyActivation %}
+  acc = fused_act(acc);
+{% endif %}
+  y[idx] = {{ T }}(acc);
+}

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

@@ -1,48 +1,43 @@
-// 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" %}
-  // Clamped for the reason given in the GEMM activation: an unbounded accumulator
-  // overflows the hardware tanh's internal e^2v to Inf/Inf = NaN, and tanh is
-  // already +/-1 to full f32 precision by |v| ~ 9.
   return tanh(clamp(v, -10.0, 10.0));
-{% 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 %}

+// Register-tiled implicit GEMM for group-1 NCW convolution:
+// Y[b] (M x N) = W (M x K) @ im2col(X[b]) (K x N), where M is outChannels, K
+// is inChannels * kernelW, and N is outW. The X tile is gathered directly from
+// NCW with the configured stride, dilation, and left padding; no column matrix
+// is materialized.
+//
+// A WG_X x WG_Y workgroup computes a (WG_Y * TM) x (WG_X * TN) output tile
+// from BK-deep slices, with one TM x TN register micro-tile per thread. Weight
+// and input tiles group four K values per workgroup-memory word and accumulate
+// through dot(). Bounds checks cover every partial tile. Accumulation, bias,
+// and the optional activation epilogue use f32; f16 operands widen at the FMA.
 {{ env.wgsl.resourceDeclarations }}
 {% set hasActivation = hasActivation is defined and hasActivation %}
+{% set activation = activation | default("") %}
+{% set actAlpha = actAlpha | default(0.0) %}
+{% set actBeta = actBeta | default(0.0) %}
 {% if hasActivation %}
 // Apply the fused activation in the f32 accumulator before the single output
 // cast, avoiding an intermediate convolution tensor.
+{% macro fused_act_return(mode, alpha, beta) -%}
+{% if mode == "Relu" %}
   return max(v, 0.0);
+{% elif mode == "Clip" %}
+  return clamp(v, f32({{ alpha }}), f32({{ beta }}));
+{% elif mode == "LeakyRelu" %}
+  return select(v * f32({{ alpha }}), v, v >= 0.0);
+{% elif mode == "Sigmoid" %}
   return 1.0 / (1.0 + exp(-v));
+{% elif mode == "Tanh" %}
+  // tanh is already +/-1 to f32 precision at the clamp bounds. Clamping before
+  // the builtin preserves that saturated result for larger accumulators.
   return tanh(clamp(v, -10.0, 10.0));
+{% elif mode == "HardSigmoid" %}
+  return clamp(f32({{ alpha }}) * v + f32({{ beta }}), 0.0, 1.0);
 {% else %}
   return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
 {% endif %}
+{%- endmacro -%}
+fn fused_act(v: f32) -> f32 {
+{{ fused_act_return(activation, actAlpha, actBeta) -}}
 }
 {% endif %}

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

@@ -1,30 +1,32 @@
-{% 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" %}
-  // Clamped for the reason given in the GEMM activation: an unbounded accumulator
-  // overflows the hardware tanh's internal e^2v to Inf/Inf = NaN, and tanh is
-  // already +/-1 to full f32 precision by |v| ~ 9.
   return tanh(clamp(v, -10.0, 10.0));
-{% 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 %}
@@ -32,86 +34,125 @@ fn fused_act(v: f32) -> f32 {
 // 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 %}
@@ -121,31 +162,31 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
         }
 {% 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 %}
@@ -156,13 +197,23 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
     }
 {% 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) {
@@ -172,7 +223,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
 {% 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 %}

 {{ env.wgsl.resourceDeclarations }}
 {% set hasActivation = hasActivation is defined and hasActivation %}
+{% set activation = activation | default("") %}
+{% set actAlpha = actAlpha | default(0.0) %}
+{% set actBeta = actBeta | default(0.0) %}
 {% if hasActivation %}
 // Apply the fused activation in the f32 accumulator before the single output
 // cast, avoiding an intermediate convolution tensor.
+{% macro fused_act_return(mode, alpha, beta) -%}
+{% if mode == "Relu" %}
   return max(v, 0.0);
+{% elif mode == "Clip" %}
+  return clamp(v, f32({{ alpha }}), f32({{ beta }}));
+{% elif mode == "LeakyRelu" %}
+  return select(v * f32({{ alpha }}), v, v >= 0.0);
+{% elif mode == "Sigmoid" %}
   return 1.0 / (1.0 + exp(-v));
+{% elif mode == "Tanh" %}
+  // tanh is already +/-1 to f32 precision at the clamp bounds. Clamping before
+  // the builtin preserves that saturated result for larger accumulators.
   return tanh(clamp(v, -10.0, 10.0));
+{% elif mode == "HardSigmoid" %}
+  return clamp(f32({{ alpha }}) * v + f32({{ beta }}), 0.0, 1.0);
 {% else %}
   return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
 {% endif %}
+{%- endmacro -%}
+fn fused_act(v: f32) -> f32 {
+{{ fused_act_return(activation, actAlpha, actBeta) -}}
 }
 {% endif %}
 // 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 indexing divisors and kernel offsets are shader constants.
+{% set dilatedLanes = dilatedLanes | default(false) %}
+{% set lanes = lanes | default(4) %}
+{% set quads = quads | default(1) %}
+{% set laneStep = dilationWSpec if dilatedLanes else strideWSpec %}
+{% set comps = ["x", "y", "z", "w"] %}
+{% macro acc_ref(oct, lane) %}{% if dilatedLanes %}acc{{ oct }}_{{ lane }}{% else %}acc{{ oct }}.{{ comps[lane] }}{% endif %}{% endmacro %}
+{% set usedCount = namespace(value=0) %}
+{% for s in range(span) %}
+{% set used = namespace(value=false) %}
+{% for kw in range(kernelWSpec) %}
+{% for lane in range(lanes) %}
+{% if s == kw * dilationWSpec + lane * laneStep %}{% set used.value = true %}{% endif %}
+{% endfor %}
+{% endfor %}
+{% if used.value %}{% set usedCount.value = usedCount.value + 1 %}{% endif %}
+{% endfor %}
+{% set registerForm = (usedCount.value <= spanCap) if dilatedLanes else (span <= spanCap) %}
+const COUNT_TILES: u32 = {{ countTiles }}u;
+{% if tailOutput %}const OUT_W: u32 = max(1u, {{ outW }}u);
+{% endif %}
+{% if dilatedLanes %}
+const OUT_W: u32 = max(1u, {{ outW }}u);
+const LANES: u32 = {{ lanes }}u;
+const RESIDUES: u32 = {{ dilationWSpec }}u;
+const TILES_PER_ROW: u32 = RESIDUES * {{ quads }}u;
+{% else %}
+const OUT_W4: u32 = max(1u, {{ outW4 }}u);
 {% endif %}
+const OUT_H: u32 = max(1u, {{ outH }}u);
+const OUT_C: u32 = max(1u, {{ outC }}u);
+const OC_TILE: u32 = {{ ocTile }}u;
 const OUT_C_TILES: u32 = OUT_C / OC_TILE;
+const OUT_CPG: u32 = {{ outCPerGroup }}u;
+const IN_C: u32 = {{ inC }}u;
+const IN_CPG: u32 = {{ inCPerGroup }}u;
+const IN_H: i32 = {{ inH }};
+const IN_W: i32 = {{ inW }};
+const IN_H_U: u32 = {{ inH }}u;
+const IN_W_U: u32 = {{ inW }}u;
 const IN_PLANE: u32 = IN_H_U * IN_W_U;
+{% if not tailOutput and not dilatedLanes %}const OUT_PLANE4: u32 = OUT_H * OUT_W4;
 {% endif %}
+const KAREA: u32 = {{ kernelHSpec * kernelWSpec }}u;
+const STRIDE_H: u32 = {{ strideHSpec }}u;
+const STRIDE_W: i32 = {{ strideWSpec }};
+const PAD_TOP: i32 = {{ padTopSpec }};
+const PAD_LEFT: i32 = {{ padLeftSpec }};
+const WG: u32 = {{ workgroupSizeSpec }}u;
+@compute @workgroup_size({{ workgroupSizeSpec }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let q = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
   if (q >= COUNT_TILES) { return; }
+{% if dilatedLanes %}
+  let tile = q % TILES_PER_ROW;
+  var t = q / TILES_PER_ROW;
+{% else %}
   let ow4 = q % OUT_W4;
   var t = q / OUT_W4;
+{% endif %}
   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;
+{% if dilatedLanes %}
+  // Lane l owns output column ow0 + l * dilationW: the residue class first,
+  // then the LANES-wide run within it. Runs past the row store nothing.
+  let ow0 = tile % RESIDUES + (tile / RESIDUES) * (LANES * RESIDUES);
+  let iwBase = i32(ow0) * STRIDE_W - PAD_LEFT;
+{% else %}
   let iwBase = i32(ow4 * 4u) * STRIDE_W - PAD_LEFT;
+{% endif %}
   var xChannelBase = (batch * IN_C + group * IN_CPG) * IN_PLANE;
+{% for oct in range(ocTile) %}
   var wBase{{ oct }} = (ocBase + {{ oct }}u) * IN_CPG * KAREA;
+{% if dilatedLanes %}
+{% for lane in range(lanes) %}
+  var acc{{ oct }}_{{ lane }} = 0.0;
+{% endfor %}
+{% else %}
   var acc{{ oct }} = vec4<f32>(0.0);
+{% endif %}
 {% endfor %}
   for (var ic = 0u; ic < IN_CPG; ic++) {
+{% for kh in range(kernelHSpec) %}
     {
+      let ih = ihBase + {{ kh * dilationHSpec }};
       if (ih >= 0 && ih < IN_H) {
         let xRow = xChannelBase + u32(ih) * IN_W_U;
+{% if registerForm %}
         // 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 statically named local. Out-of-range
+        // columns contain zero and contribute nothing to the accumulation.
         //
+        // SPAN grows as (kernelW-1)*dilationW. `spanCap` bounds the number of
+        // staged row-window values; wider spans use the fixed-size per-tap form below.
+{% for s in range(span) %}
 {% set used = namespace(value=false) %}
+{% for kw in range(kernelWSpec) %}
+{% for lane in range(lanes) %}
+{% if s == kw * dilationWSpec + lane * laneStep %}{% set used.value = true %}{% endif %}
 {% endfor %}
 {% endfor %}
 {% if used.value %}
         var xw{{ s }} = 0.0;
 {% endif %}
 {% endfor %}
+{% for s in range(span) %}
 {% set used = namespace(value=false) %}
+{% for kw in range(kernelWSpec) %}
+{% for lane in range(lanes) %}
+{% if s == kw * dilationWSpec + lane * laneStep %}{% set used.value = true %}{% endif %}
 {% endfor %}
 {% endfor %}
 {% if used.value %}
         }
 {% endif %}
 {% endfor %}
+{% for kw in range(kernelWSpec) %}
         {
+{% for oct in range(ocTile) %}
+          let weight{{ oct }} = f32(w[wBase{{ oct }} + {{ kh * kernelWSpec + kw }}u]);
 {% endfor %}
+{% for lane in range(lanes) %}
+{% for oct in range(ocTile) %}
+          {{ acc_ref(oct, lane) }} += xw{{ kw * dilationWSpec + lane * laneStep }} * weight{{ oct }};
 {% endfor %}
 {% endfor %}
         }
 {% endfor %}
 {% else %}
+{% for kw in range(kernelWSpec) %}
         {
+{% for oct in range(ocTile) %}
+          let weight{{ oct }} = f32(w[wBase{{ oct }} + {{ kh * kernelWSpec + kw }}u]);
 {% endfor %}
+          let iwK = iwBase + {{ kw * dilationWSpec }};
+{% for lane in range(lanes) %}
+          let iwLane{{ lane }} = iwK + {{ lane * laneStep }};
           if (iwLane{{ lane }} >= 0 && iwLane{{ lane }} < IN_W) {
             let xValue = f32(x[xRow + u32(iwLane{{ lane }})]);
+{% for oct in range(ocTile) %}
+            {{ acc_ref(oct, lane) }} += xValue * weight{{ oct }};
 {% endfor %}
           }
 {% endfor %}
     }
 {% endfor %}
     xChannelBase += IN_PLANE;
+{% for oct in range(ocTile) %}
     wBase{{ oct }} += KAREA;
 {% endfor %}
   }
+{% if dilatedLanes %}
+  let outBase = ((batch * OUT_C + ocBase) * OUT_H + oh) * OUT_W;
+{% for oct in range(ocTile) %}
+  let outChannelBase{{ oct }} = outBase + {{ oct }}u * OUT_H * OUT_W;
+{% for lane in range(lanes) %}
+  if (ow0 + {{ lane * dilationWSpec }}u < OUT_W) {
+    y[outChannelBase{{ oct }} + ow0 + {{ lane * dilationWSpec }}u] = {{ scalar }}({% if hasActivation %}fused_act(acc{{ oct }}_{{ lane }}){% else %}acc{{ oct }}_{{ lane }}{% endif %});
+  }
+{% endfor %}
+{% endfor %}
+{% elif tailOutput %}
   let outBase = ((batch * OUT_C + ocBase) * OUT_H + oh) * OUT_W + ow4 * 4u;
+{% for oct in range(ocTile) %}
   let outChannelBase{{ oct }} = outBase + {{ oct }}u * OUT_H * OUT_W;
 {% for lane in range(4) %}
   if (ow4 * 4u + {{ lane }}u < OUT_W) {
 {% endfor %}
 {% else %}
   let outBase = ((batch * OUT_C + ocBase) * OUT_H + oh) * OUT_W4 + ow4;
+{% for oct in range(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 CHANGED Viewed

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

build/webgpu/metadata.json CHANGED Viewed

@@ -1,25 +1,91 @@
 {
   "name": "com.microsoft.FusedConv",
-  "id": "_com_microsoft_fusedconv_webgpu_e37405a",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "0T4EaJworpji9PrXCF3znUUqiMBBbSgShzzAaM/dOWY=",
-      "conv-1x1-gemm-tiled-reg.wgsl.jinja": "n0CRMXMLwCsmBCTq9ruaSshgRFN+QcWfYHmLLT/nYDg=",
-      "conv-1x1-gemm-tiled.wgsl.jinja": "tAifcSzQFhljZg7xs7W7xJ1t92ZjRlCCfET/7sxsrzI=",
-      "conv-1x1-subgroup-matrix.wgsl.jinja": "4M+d7l/v8cK31b+6c2QaD20PSrT3EcpMNT5LV14H+wg=",
-      "conv-direct-nd.wgsl.jinja": "+rZdtrCmC5XWDKzniyQgM8ZuEPelUn15Wy8XASMqRoI=",
-      "conv-direct-unrolled.wgsl.jinja": "Z8nzgvMffIyDdiORizHznKPUpi70ZcWPZvmwhnuE+PQ=",
-      "conv-im2col-nchw.wgsl.jinja": "7IUuXo33elZx0pnR9M2vUXzXDjFuw2bJODkZViV/WQg=",
-      "conv1d-tiled-reg.wgsl.jinja": "f6XtRoLCTQvmBsoTWjymS9jK1ewlj9R2t+DWBa7vv/0=",
-      "conv2d-grouped-large-w4.wgsl.jinja": "HMab+h8mlTI14B1EdMAkYmwFRFkRipfdT2ZOx36X2dY=",
-      "manifest.json": "jMjmEiSVgMOjNheSacCaHlppVU7uWplUp2lv3fvGH1A=",
-      "test.json": "/KodhC9ZJ6IxO1wokLmeGH0zsOfDpds9+mjZLZRnPXw="
     }
   },
-  "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.FusedConv" }
 }

 {
   "name": "com.microsoft.FusedConv",
+  "id": "_com_microsoft_fusedconv_webgpu_e4dcf04",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "rDPb1+xln8/qfupw2ajiwaBUEKJvlo+NuJEs81cAgTk=",
+      "conv-1x1-gemm-tiled-reg.wgsl.jinja": "9yFpwXS1HzmWmm4W72n+aOIryLqc9+e88awZ1ogTe+Q=",
+      "conv-1x1-gemm-tiled.wgsl.jinja": "uqG5mRVvcHiUo0LDCUcVgqElVTvRNknx19Gm4oIALiU=",
+      "conv-1x1-subgroup-matrix.wgsl.jinja": "iDJs4RJk4iJarOZCjY7/5MFrHgR7CcqI6x5uwv2dYaA=",
+      "conv-direct-nd.wgsl.jinja": "5UA6BigLbtPRMre3LRfNShldvupB5yTGn8EzX2tBzP8=",
+      "conv-direct-unrolled.wgsl.jinja": "vrGUPwYkbYE1slII7nuK5wbERRcG5u0X1cW1kMc2sQs=",
+      "conv-im2col-nchw.wgsl.jinja": "IRpmtpeVQsLE6ud+Sf1230XaZVS1dxDtTZuSDwU3P8s=",
+      "conv-splitk-reduce.wgsl.jinja": "LbVU5DRL0Q51fvLjThmxJQZH8aPbZkvNzij4nY/cHmo=",
+      "conv1d-tiled-reg.wgsl.jinja": "ERqOPl+1RbnJ+CFqZvLSYw45qYMKgB7Je0teew5gyOM=",
+      "conv2d-grouped-large-w4.wgsl.jinja": "qBG2cfOsQ4FMwtNT8wBOqLJFDOs6w0wJCmfcZuxbho8=",
+      "manifest.json": "EpZ3lQZ0Ar7ZY4xDkUlcQ3m6lyw6x6+cp9KoBC47TSE=",
+      "test.json": "g7doPXd+oCU97EVAbYReA8qInOswdv5wrLcMYCX7CcM="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "implicit_im2col_tiled_reg_splitk": ["conv-1x1-gemm-tiled-reg.wgsl.jinja", "conv-splitk-reduce.wgsl.jinja"],
+      "implicit_im2col_tiled_bias_reg_splitk": ["conv-1x1-gemm-tiled-reg.wgsl.jinja", "conv-splitk-reduce.wgsl.jinja"],
+      "implicit_im2col_tiled_reg_splitk_preferred": ["conv-1x1-gemm-tiled-reg.wgsl.jinja", "conv-splitk-reduce.wgsl.jinja"],
+      "implicit_im2col_tiled_bias_reg_splitk_preferred": ["conv-1x1-gemm-tiled-reg.wgsl.jinja", "conv-splitk-reduce.wgsl.jinja"],
+      "implicit_im2col_tiled_reg": ["conv-1x1-gemm-tiled-reg.wgsl.jinja"],
+      "implicit_im2col_tiled_bias_reg": ["conv-1x1-gemm-tiled-reg.wgsl.jinja"],
+      "gemm_1x1_subgroup_matrix": ["conv-1x1-subgroup-matrix.wgsl.jinja"],
+      "gemm_1x1_subgroup_matrix_z": ["conv-1x1-subgroup-matrix.wgsl.jinja"],
+      "gemm_1x1_subgroup_matrix_bias": ["conv-1x1-subgroup-matrix.wgsl.jinja"],
+      "gemm_1x1_subgroup_matrix_bias_z": ["conv-1x1-subgroup-matrix.wgsl.jinja"],
+      "im2col_gemm_subgroup_matrix": ["conv-1x1-subgroup-matrix.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
+      "im2col_direct_inputs_subgroup_matrix": ["conv-1x1-subgroup-matrix.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
+      "im2col_gemm_subgroup_matrix_z": ["conv-1x1-subgroup-matrix.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
+      "im2col_direct_inputs_subgroup_matrix_z": ["conv-1x1-subgroup-matrix.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
+      "im2col_gemm_subgroup_matrix_bias": ["conv-1x1-subgroup-matrix.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
+      "im2col_direct_inputs_subgroup_matrix_bias": ["conv-1x1-subgroup-matrix.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
+      "im2col_gemm_subgroup_matrix_bias_z": ["conv-1x1-subgroup-matrix.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
+      "im2col_direct_inputs_subgroup_matrix_bias_z": ["conv-1x1-subgroup-matrix.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
+      "im2col_half_direct_subgroup_matrix": ["conv-1x1-subgroup-matrix.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
+      "im2col_half_direct_subgroup_matrix_z": ["conv-1x1-subgroup-matrix.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
+      "im2col_half_direct_subgroup_matrix_bias": ["conv-1x1-subgroup-matrix.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
+      "im2col_half_direct_subgroup_matrix_bias_z": ["conv-1x1-subgroup-matrix.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
+      "im2col_gemm_subgroup_matrix_padded": ["conv-1x1-subgroup-matrix.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
+      "im2col_gemm_subgroup_matrix_padded_z": ["conv-1x1-subgroup-matrix.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
+      "im2col_gemm_subgroup_matrix_padded_bias": ["conv-1x1-subgroup-matrix.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
+      "im2col_gemm_subgroup_matrix_padded_bias_z": ["conv-1x1-subgroup-matrix.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
+      "im2col_gemm_tiled": ["conv-1x1-gemm-tiled.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
+      "im2col_gemm_tiled_z": ["conv-1x1-gemm-tiled.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
+      "im2col_gemm_tiled_bias": ["conv-1x1-gemm-tiled.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
+      "im2col_gemm_tiled_bias_z": ["conv-1x1-gemm-tiled.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
+      "im2col_gemm_tiled_reg": ["conv-1x1-gemm-tiled-reg.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
+      "im2col_gemm_tiled_bias_reg": ["conv-1x1-gemm-tiled-reg.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
+      "gemm_1x1_tiled": ["conv-1x1-gemm-tiled.wgsl.jinja"],
+      "gemm_1x1_tiled_z": ["conv-1x1-gemm-tiled.wgsl.jinja"],
+      "gemm_1x1_tiled_bias": ["conv-1x1-gemm-tiled.wgsl.jinja"],
+      "gemm_1x1_tiled_bias_z": ["conv-1x1-gemm-tiled.wgsl.jinja"],
+      "gemm_1x1_tiled_reg": ["conv-1x1-gemm-tiled-reg.wgsl.jinja"],
+      "gemm_1x1_tiled_bias_reg": ["conv-1x1-gemm-tiled-reg.wgsl.jinja"],
+      "conv1d_direct": ["conv-direct-unrolled.wgsl.jinja"],
+      "direct_unrolled": ["conv-direct-unrolled.wgsl.jinja"],
+      "direct_unrolled_z": ["conv-direct-unrolled.wgsl.jinja"],
+      "conv1d_direct_bias": ["conv-direct-unrolled.wgsl.jinja"],
+      "direct_unrolled_bias": ["conv-direct-unrolled.wgsl.jinja"],
+      "direct_unrolled_bias_z": ["conv-direct-unrolled.wgsl.jinja"],
+      "conv1d_tiled_reg": ["conv1d-tiled-reg.wgsl.jinja"],
+      "conv1d_tiled_bias_reg": ["conv1d-tiled-reg.wgsl.jinja"],
+      "ncdhw3d": ["conv-direct-nd.wgsl.jinja"],
+      "ncdhw3d_z": ["conv-direct-nd.wgsl.jinja"],
+      "ncdhw3d_bias": ["conv-direct-nd.wgsl.jinja"],
+      "ncdhw3d_bias_z": ["conv-direct-nd.wgsl.jinja"],
+      "grouped_large_kernel_w4": ["conv2d-grouped-large-w4.wgsl.jinja"],
+      "grouped_large_kernel_w4_tail": ["conv2d-grouped-large-w4.wgsl.jinja"],
+      "grouped_large_kernel_w4_dilated_lanes": ["conv2d-grouped-large-w4.wgsl.jinja"],
+      "implicit_im2col_subgroup_matrix": ["conv-1x1-subgroup-matrix.wgsl.jinja"],
+      "implicit_im2col_subgroup_matrix_z": ["conv-1x1-subgroup-matrix.wgsl.jinja"],
+      "implicit_im2col_subgroup_matrix_bias": ["conv-1x1-subgroup-matrix.wgsl.jinja"],
+      "implicit_im2col_subgroup_matrix_bias_z": ["conv-1x1-subgroup-matrix.wgsl.jinja"],
+      "implicit_im2col_subgroup_matrix_f16": ["conv-1x1-subgroup-matrix.wgsl.jinja"],
+      "implicit_im2col_subgroup_matrix_z_f16": ["conv-1x1-subgroup-matrix.wgsl.jinja"],
+      "implicit_im2col_subgroup_matrix_bias_f16": ["conv-1x1-subgroup-matrix.wgsl.jinja"],
+      "implicit_im2col_subgroup_matrix_bias_z_f16": ["conv-1x1-subgroup-matrix.wgsl.jinja"]
+    }
+  }
 }

build/webgpu/test.json CHANGED Viewed

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