sync 6fdf6301e2bb

Browse files

Files changed (13) hide show

README.md +14 -4
build/webgpu/bench.json +0 -0
build/webgpu/conv-1x1-gemm-tiled-reg.wgsl.jinja +46 -125
build/webgpu/conv-1x1-gemm-tiled.wgsl.jinja +59 -21
build/webgpu/conv-1x1-subgroup-matrix.wgsl.jinja +50 -179
build/webgpu/conv-direct-nd.wgsl.jinja +46 -18
build/webgpu/conv-direct-unrolled.wgsl.jinja +46 -17
build/webgpu/conv-splitk-reduce.wgsl.jinja +46 -16
build/webgpu/conv1d-tiled-reg.wgsl.jinja +37 -10
build/webgpu/conv2d-grouped-large-w4.wgsl.jinja +58 -22
build/webgpu/manifest.json +0 -0
build/webgpu/metadata.json +19 -14
build/webgpu/test.json +0 -0

README.md CHANGED Viewed

@@ -37,8 +37,8 @@ Attributes and default values (overridable per request):
 | 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. |
@@ -78,6 +78,16 @@ One implementation is selected per call from the device capabilities, the reques
 - `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.
@@ -96,7 +106,7 @@ Some implementation variants require `subgroup-matrix`, `shader-f16`, and `subgr
 - [`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-1x1-gemm-tiled-reg.wgsl.jinja`](build/webgpu/conv-1x1-gemm-tiled-reg.wgsl.jinja)
 - [`conv-1x1-gemm-tiled.wgsl.jinja`](build/webgpu/conv-1x1-gemm-tiled.wgsl.jinja)
 - [`conv-1x1-subgroup-matrix.wgsl.jinja`](build/webgpu/conv-1x1-subgroup-matrix.wgsl.jinja)
@@ -110,7 +120,7 @@ Some implementation variants require `subgroup-matrix`, `shader-f16`, and `subgr
 ## 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.

 | Attribute | Default | Description |
 | --- | --- | --- |
+| `activation` | — | Optional fused activation name: `Relu`, `LeakyRelu`, `Sigmoid`, `Tanh`, `HardSigmoid`, `HardSwish`, `Clip`, `QuickGelu` (x * sigmoid(alpha * x); alpha 1 is SiLU), `Elu`, `Gelu` (erf form, or the tanh approximation when `activation_params[0]` is nonzero), `FastGelu` (always the tanh approximation), `Softplus`, `ThresholdedRelu`, or `Erf`. Omission applies no activation. The last seven are the set onnxruntime's WebGPU provider fuses into Conv; its CPU provider rejects them. |
+| `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. `QuickGelu`, `Elu` and `ThresholdedRelu` take an optional `[alpha]` (defaults 1.702, 1.0 and 1.0); `Gelu` takes an optional `[approximate]` flag (0 = erf, nonzero = tanh). 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. |
 - `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.
+- `im2col_gemm_tiled_reg` — Materialize f32 columns and run a register-tiled GEMM with f32 accumulation.
+- `im2col_gemm_tiled_bias_reg` — Materialize f32 columns and run a register-tiled GEMM with f32 accumulation. Add bias before activation.
+- `im2col_gemm_tiled_reg_f16_columns` — Materialize f16 inputs directly into f16 columns, then run the register-tiled GEMM with f32 accumulation. This preserves the input values exactly and halves column-buffer traffic and storage; eligibility uses the actual half-precision allocation and WebGPU buffer limits.
+- `im2col_gemm_tiled_bias_reg_f16_columns` — Materialize f16 inputs directly into f16 columns, then run the register-tiled GEMM with f32 accumulation. This preserves the input values exactly and halves column-buffer traffic and storage; eligibility uses the actual half-precision allocation and WebGPU buffer limits. Add bias before activation.
+- `grouped_large_kernel_w4` — Shares input windows across channels and columns with f32 accumulation. Filter area, channel bytes and activation presence choose row looping or unrolling. Bias seeds the accumulators; activation follows the reduction. Small workloads and large unrolled kernels outside a fixed 32-wide subgroup range remain eligible as demoted fallbacks.
+- `grouped_large_kernel_w4_bias` — Shares input windows across channels and columns with f32 accumulation. Filter area, channel bytes and activation presence choose row looping or unrolling. Bias seeds the accumulators; activation follows the reduction. Small workloads and large unrolled kernels outside a fixed 32-wide subgroup range remain eligible as demoted fallbacks.
+- `grouped_large_kernel_w4_tail` — Shares input windows across channels and columns with f32 accumulation. Filter area, channel bytes and activation presence choose row looping or unrolling. Bias seeds the accumulators; activation follows the reduction. Small workloads and large unrolled kernels outside a fixed 32-wide subgroup range remain eligible as demoted fallbacks.
+- `grouped_large_kernel_w4_tail_bias` — Shares input windows across channels and columns with f32 accumulation. Filter area, channel bytes and activation presence choose row looping or unrolling. Bias seeds the accumulators; activation follows the reduction. Small workloads and large unrolled kernels outside a fixed 32-wide subgroup range remain eligible as demoted fallbacks.
+- `grouped_large_kernel_w4_dilated_lanes` — Shares input windows across channels and columns with f32 accumulation. Filter area, channel bytes and activation presence choose row looping or unrolling. Bias seeds the accumulators; activation follows the reduction. Small workloads and large unrolled kernels outside a fixed 32-wide subgroup range remain eligible as demoted fallbacks.
+- `grouped_large_kernel_w4_dilated_lanes_bias` — Shares input windows across channels and columns with f32 accumulation. Filter area, channel bytes and activation presence choose row looping or unrolling. Bias seeds the accumulators; activation follows the reduction. Small workloads and large unrolled kernels outside a fixed 32-wide subgroup range remain eligible as demoted fallbacks.
 - `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.
 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark cases
 - [`conv-1x1-gemm-tiled-reg.wgsl.jinja`](build/webgpu/conv-1x1-gemm-tiled-reg.wgsl.jinja)
 - [`conv-1x1-gemm-tiled.wgsl.jinja`](build/webgpu/conv-1x1-gemm-tiled.wgsl.jinja)
 - [`conv-1x1-subgroup-matrix.wgsl.jinja`](build/webgpu/conv-1x1-subgroup-matrix.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.3
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

build/webgpu/bench.json CHANGED Viewed

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

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

@@ -14,14 +14,11 @@
 // 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" %}
@@ -36,19 +33,49 @@
   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 tileT = "f16" if usesF16 else "f32" -%}
-{%- set GEMM_BK = gemmKTile if gemmKTile is defined else 16 %}
 {% set GEMM_TM = gemmThreadRows if gemmThreadRows is defined else 8 %}
-{% set GEMM_WG_X = gemmWorkgroupX if gemmWorkgroupX is defined else 16 %}
 {% set GEMM_BM = gemmMTile if gemmMTile is defined else 64 %}
 {% set GEMM_BN = gemmNTile if gemmNTile is defined else GEMM_WG_X * 4 %}
 {% set GEMM_TN = gemmThreadColumns if gemmThreadColumns is defined else (GEMM_BN / GEMM_WG_X)|int %}
@@ -70,11 +97,6 @@ const BN_VECS: u32 = BN / 4u;
 // 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) %}
@@ -116,37 +138,7 @@ 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;
@@ -164,30 +156,9 @@ const PROJECTION_OUT_C: u32 = {{ projectionChannels }}u;
   y[yBase + m * params.N + n] = {{ T }}(v);
 }
 {% endif %}
-{% endif %}
 var<workgroup> tileA: array<array<vec4<{{ tileT }}>, AK_VECS>, BM>;
 var<workgroup> tileB: array<array<vec4<{{ tileT }}>, BN_VECS>, BK>;
-{% if fusedNarrowProjection and GEMM_BK < GEMM_BM %}
-// BK16 cannot reuse the 16-row input tile to publish a BM32 intermediate.
-// The extra 32x64 f32 tile keeps the total at 14 KiB, below WebGPU's
-// guaranteed 16 KiB workgroup-storage floor.
-var<workgroup> projectionTile: array<array<vec4<f32>, BN_VECS>, BM>;
-{% endif %}
-{% macro publish_projection_input(localRow, localVec, component, globalRow, globalColumn, raw) %}
-  if ({{ globalRow }} < M && {{ globalColumn }} < N) {
-    var projectionInput = {{ raw }};
-{% if hasBias %}
-    projectionInput = projectionInput + f32(bias[{{ globalRow }}]);
-{% endif %}
-{% if projectionInputAct == "relu" %}
-    projectionInput = projection_relu(projectionInput);
-{% endif %}
-    {{ narrowProjectionTile }}[{{ localRow }}][{{ localVec }}].{{ component }} = projectionInput;
-  } else {
-    {{ narrowProjectionTile }}[{{ localRow }}][{{ localVec }}].{{ component }} = 0.0;
-  }
-{%- endmacro %}
 {% macro load_a_vec4(rowBase, columnBase, rowLimit, columnLimit) %}
     for (var linear = li; linear < BM * AK_VECS; linear += WG_SIZE) {
       let ar = linear / AK_VECS;
@@ -208,8 +179,7 @@ var<workgroup> projectionTile: array<array<vec4<f32>, BN_VECS>, BM>;
         }
       }
       tileA[ar][ac4] = av;
-    }
-{%- endmacro %}
 {% macro load_b_vec4(rowBase, columnBase, rowLimit, columnLimit) %}
     for (var linear = li; linear < BK * BN_VECS; linear += WG_SIZE) {
       let br = linear / BN_VECS;
@@ -230,8 +200,7 @@ var<workgroup> projectionTile: array<array<vec4<f32>, BN_VECS>, BM>;
         }
       }
       tileB[br][bc4] = bvec;
-    }
-{%- endmacro %}
 {% macro load_b_implicit_vec4(rowBase, columnBase, rowLimit, columnLimit) %}
     for (var linear = li; linear < BK * BN_VECS; linear += WG_SIZE) {
       let br = linear / BN_VECS;
@@ -259,8 +228,7 @@ var<workgroup> projectionTile: array<array<vec4<f32>, BN_VECS>, BM>;
         }
       }
       tileB[br][bc4] = bvec;
-    }
-{%- endmacro %}
 {% macro load_b_implicit_carried(rowBase, columnBase, rowLimit, columnLimit) %}
     let bColVec = li / {{ implicitGatherKChunks }}u;
     let bChunk = li % {{ implicitGatherKChunks }}u;
@@ -344,8 +312,7 @@ var<workgroup> projectionTile: array<array<vec4<f32>, BN_VECS>, BM>;
           }
         }
       }
-    }
-{%- endmacro %}
 @compute @workgroup_size({{ GEMM_WG_X }}, {{ GEMM_WG_Y }}, 1)
 fn main(
@@ -423,9 +390,7 @@ fn main(
   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 %}
@@ -433,52 +398,8 @@ fn main(
 {% 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] }});
 {% endif %}
 {% endfor %}
-{% endfor %}{% if fusedNarrowProjection %}
-  // Conv's final K-tile barrier makes tileB dead before it becomes the
-  // intermediate tile. Every lane publishes its unique micro-tile, then one
-  // lane per spatial vector word consumes all logical M rows in increasing
-  // order, four adjacent columns at a time. Invalid M/N-tail cells are
-  // initialized above, so the barrier is uniform and no lane can observe a
-  // stale input-tile value.
-  workgroupBarrier();
-  if (li < BN_VECS) {
-    let projectionNBase = nBase + li * 4u;
-{% for oc in range(projectionChannels) %}
-    var projectionAcc{{ oc }} = vec4<f32>(0.0);
-{% endfor %}
-    for (var channel = 0u; channel < M; channel = channel + 1u) {
-      let projectionValues = vec4<f32>({{ narrowProjectionTile }}[channel][li]);
-{% for oc in range(projectionChannels) %}
-      projectionAcc{{ oc }} = projectionAcc{{ oc }} + projectionValues * f32(projectionW[{{ oc }}u * M + channel]);
-{% endfor %}
-    }
-    let projectionYBase = batch * PROJECTION_OUT_C * N;
-{% for oc in range(projectionChannels) %}
-    let projected{{ oc }} = projectionAcc{{ oc }};
-{% endfor %}
-{% for column in range(4) %}
-    let projectionN{{ column }} = projectionNBase + {{ column }}u;
-    if (projectionN{{ column }} < N) {
-{% for oc in range(projectionChannels) %}
-{% if projectionOutputAct == "relu" %}
-      let activated{{ oc }}_{{ column }} = projection_relu(projected{{ oc }}.{{ components[column] }});
-{% elif projectionOutputAct == "sigmoid" %}
-      let activated{{ oc }}_{{ column }} = sigmoid_safe(projected{{ oc }}.{{ components[column] }});
-{% else %}
-      let activated{{ oc }}_{{ column }} = projected{{ oc }}.{{ components[column] }};
-{% endif %}
-      y[projectionYBase + {{ oc }}u * N + projectionN{{ column }}] = activated{{ oc }}_{{ column }};
-{% endfor %}
-    }
-{% endfor %}
-  }
-{% endif %}
-}

 // 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 }}
 {% if hasActivation %}
 // Apply the fused activation in the f32 accumulator before the single output
 // cast, avoiding an intermediate convolution tensor.
+{% macro fused_act_fn(mode, alpha, beta) %}
+fn fused_act(v: f32) -> f32 {
 {% if mode == "Relu" %}
   return max(v, 0.0);
 {% elif mode == "Clip" %}
   return tanh(clamp(v, -10.0, 10.0));
 {% elif mode == "HardSigmoid" %}
   return clamp(f32({{ alpha }}) * v + f32({{ beta }}), 0.0, 1.0);
+{% elif mode == "QuickGelu" %}
+{% if alpha == 1 %}
+  // alpha == 1 is SiLU: the multiply folds away (the spelling yolov8-class detectors emit).
+  return v * (1.0 / (1.0 + exp(-v)));
+{% else %}
+  return v * (1.0 / (1.0 + exp(-(f32({{ alpha }}) * v))));
+{% endif %}
+{% elif mode == "Elu" %}
+  return select(f32({{ alpha }}) * (exp(v) - 1.0), v, v >= 0.0);
+{% elif mode == "ThresholdedRelu" %}
+  return select(0.0, v, v > f32({{ alpha }}));
+{% elif mode == "Softplus" %}
+  // log(1 + exp(v)) spelled overflow-safe: max(v, 0) + log(1 + exp(-|v|)).
+  return max(v, 0.0) + log(1.0 + exp(-abs(v)));
+{% elif mode == "Erf" %}
+  // Abramowitz & Stegun 7.1.26 (|error| < 1.5e-7), the same polynomial the standalone
+  // Erf/Gelu kernels use, so a fused and an unfused graph agree to f32 rounding.
+  let a = abs(v);
+  let t = 1.0 / (1.0 + 0.3275911 * a);
+  let poly = ((((1.061405429 * t - 1.453152027) * t + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t;
+  return sign(v) * (1.0 - poly * exp(-a * a));
+{% elif mode == "Gelu" and alpha == 0 %}
+  // ONNX Gelu, erf form: 0.5 * v * (1 + erf(v / sqrt(2))); activation_params[0] == 0.
+  let u = v * 0.70710678118654752;
+  let a = abs(u);
+  let t = 1.0 / (1.0 + 0.3275911 * a);
+  let poly = ((((1.061405429 * t - 1.453152027) * t + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t;
+  return 0.5 * v * (1.0 + sign(u) * (1.0 - poly * exp(-a * a)));
+{% elif mode == "Gelu" or mode == "FastGelu" %}
+  // tanh approximation (ONNX Gelu approximate="tanh" arrives as activation_params[0] != 0;
+  // com.microsoft.FastGelu always). tanh saturates to +/-1 well inside the clamp.
+  let inner = v * (0.035677408136300125 * v * v + 0.7978845608028654);
+  return v * (0.5 + 0.5 * tanh(clamp(inner, -10.0, 10.0)));
 {% else %}
   return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
 {% endif %}
+}{% endmacro %}
+{{ fused_act_fn(activation, actAlpha, actBeta) }}
 {% endif %}
+{% set tileT = "f16" if usesF16 else "f32" %}
+{% set GEMM_BK = gemmKTile if gemmKTile is defined else 16 %}
 {% set GEMM_TM = gemmThreadRows if gemmThreadRows is defined else 8 %}
+{% set GEMM_WG_X = 16 %}
 {% set GEMM_BM = gemmMTile if gemmMTile is defined else 64 %}
 {% set GEMM_BN = gemmNTile if gemmNTile is defined else GEMM_WG_X * 4 %}
 {% set GEMM_TN = gemmThreadColumns if gemmThreadColumns is defined else (GEMM_BN / GEMM_WG_X)|int %}
 // 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 splitKValue = splitK if splitK is defined else 1 %}
 {% set splitKPartial = splitKValue > 1 %}
 {% set kTiles = kTiles | default(0) %}
 const PARTIAL_BATCH_STRIDE: u32 = PARTIAL_ROWS * PARTIAL_COLS;
 const PARTIAL_SLICE_STRIDE: u32 = {{ batchCount }}u * PARTIAL_BATCH_STRIDE;
 {% endif %}
+{% 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;
   y[yBase + m * params.N + n] = {{ T }}(v);
 }
 {% endif %}
 var<workgroup> tileA: array<array<vec4<{{ tileT }}>, AK_VECS>, BM>;
 var<workgroup> tileB: array<array<vec4<{{ tileT }}>, BN_VECS>, BK>;
 {% macro load_a_vec4(rowBase, columnBase, rowLimit, columnLimit) %}
     for (var linear = li; linear < BM * AK_VECS; linear += WG_SIZE) {
       let ar = linear / AK_VECS;
         }
       }
       tileA[ar][ac4] = av;
+    }{% endmacro %}
 {% macro load_b_vec4(rowBase, columnBase, rowLimit, columnLimit) %}
     for (var linear = li; linear < BK * BN_VECS; linear += WG_SIZE) {
       let br = linear / BN_VECS;
         }
       }
       tileB[br][bc4] = bvec;
+    }{% endmacro %}
 {% macro load_b_implicit_vec4(rowBase, columnBase, rowLimit, columnLimit) %}
     for (var linear = li; linear < BK * BN_VECS; linear += WG_SIZE) {
       let br = linear / BN_VECS;
         }
       }
       tileB[br][bc4] = bvec;
+    }{% endmacro %}
 {% macro load_b_implicit_carried(rowBase, columnBase, rowLimit, columnLimit) %}
     let bColVec = li / {{ implicitGatherKChunks }}u;
     let bChunk = li % {{ implicitGatherKChunks }}u;
           }
         }
       }
+    }{% endmacro %}
 @compute @workgroup_size({{ GEMM_WG_X }}, {{ GEMM_WG_Y }}, 1)
 fn main(
   let n0 = nBase + lid.x * TN;
   let partialBase = slice * PARTIAL_SLICE_STRIDE + batch * PARTIAL_BATCH_STRIDE;
 {% else %}
   let yBatchBase = batch * M * N;
   let m0 = mBase + lid.y * TM;
   let n0 = nBase + lid.x * TN;
 {% endif %}
 {% for column in range(GEMM_TN) %}
 {% if splitKPartial %}
   y[partialBase + (m0 + {{ row }}u) * PARTIAL_COLS + n0 + {{ column }}u] = acc{{ row }}.{{ components[column] }};
 {% else %}
   store_out(m0 + {{ row }}u, n0 + {{ column }}u, yBatchBase, acc{{ row }}.{{ components[column] }});
 {% endif %}
 {% endfor %}
+{% endfor %}}

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

@@ -6,14 +6,11 @@
 // 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" %}
@@ -28,23 +25,55 @@
   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 tileT = "f16" if usesF16 else "f32" -%}
 const BK: u32 = 16u;
 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;
@@ -62,11 +91,10 @@ fn store_out(m: u32, n: u32, yBase: u32, raw: f32) {
   y[yBase + m * params.N + n] = {{ T }}(v);
 }
 var<workgroup> tileA: array<array<{{ tileT }}, BK>, BM>;
 var<workgroup> tileB: array<array<{{ tileT }}, BN>, BK>;
-@compute @workgroup_size(16, 16, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>
@@ -78,7 +106,7 @@ fn main(
   let nBase = wg.x * BN;
   let batch = wg.z;
   let xBatchBase = batch * K * N;
-  let li = lid.y * 16u + lid.x;
   var acc00: f32 = 0.0;
   var acc01: f32 = 0.0;
@@ -88,8 +116,8 @@ fn main(
   for (var kt: u32 = 0u; kt < numTiles; kt = kt + 1u) {
     let kBase = kt * BK;
     // Cooperative load: 32x16 A(=W) tile + 16x32 B(=X[batch]) tile, 256 threads x 2 each.
-    for (var e: u32 = 0u; e < (BM * BK) / 256u; e = e + 1u) {
-      let idx = li + e * 256u;
       let ar = idx / BK;
       let ac = idx % BK;
       let am = mBase + ar;
@@ -110,16 +138,26 @@ fn main(
       }
     }
     workgroupBarrier();
     for (var kk: u32 = 0u; kk < BK; kk = kk + 1u) {
       let a0 = f32(tileA[lid.y * 2u][kk]);
       let a1 = f32(tileA[lid.y * 2u + 1u][kk]);
       let b0 = f32(tileB[kk][lid.x * 2u]);
       let b1 = f32(tileB[kk][lid.x * 2u + 1u]);
-      acc00 = acc00 + a0 * b0;
-      acc01 = acc01 + a0 * b1;
-      acc10 = acc10 + a1 * b0;
-      acc11 = acc11 + a1 * b1;
     }
     workgroupBarrier();
   }

 // f16 operands remain packed in workgroup memory and widen when consumed. All
 // M, N, and K accesses are bounds-checked.
 {{ env.wgsl.resourceDeclarations }}
 {% if hasActivation %}
 // Apply the fused activation in the f32 accumulator before the single output
 // cast, avoiding an intermediate convolution tensor.
+{% macro fused_act_fn(mode, alpha, beta) %}
+fn fused_act(v: f32) -> f32 {
 {% if mode == "Relu" %}
   return max(v, 0.0);
 {% elif mode == "Clip" %}
   return tanh(clamp(v, -10.0, 10.0));
 {% elif mode == "HardSigmoid" %}
   return clamp(f32({{ alpha }}) * v + f32({{ beta }}), 0.0, 1.0);
+{% elif mode == "QuickGelu" %}
+{% if alpha == 1 %}
+  // alpha == 1 is SiLU: the multiply folds away (the spelling yolov8-class detectors emit).
+  return v * (1.0 / (1.0 + exp(-v)));
+{% else %}
+  return v * (1.0 / (1.0 + exp(-(f32({{ alpha }}) * v))));
+{% endif %}
+{% elif mode == "Elu" %}
+  return select(f32({{ alpha }}) * (exp(v) - 1.0), v, v >= 0.0);
+{% elif mode == "ThresholdedRelu" %}
+  return select(0.0, v, v > f32({{ alpha }}));
+{% elif mode == "Softplus" %}
+  // log(1 + exp(v)) spelled overflow-safe: max(v, 0) + log(1 + exp(-|v|)).
+  return max(v, 0.0) + log(1.0 + exp(-abs(v)));
+{% elif mode == "Erf" %}
+  // Abramowitz & Stegun 7.1.26 (|error| < 1.5e-7), the same polynomial the standalone
+  // Erf/Gelu kernels use, so a fused and an unfused graph agree to f32 rounding.
+  let a = abs(v);
+  let t = 1.0 / (1.0 + 0.3275911 * a);
+  let poly = ((((1.061405429 * t - 1.453152027) * t + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t;
+  return sign(v) * (1.0 - poly * exp(-a * a));
+{% elif mode == "Gelu" and alpha == 0 %}
+  // ONNX Gelu, erf form: 0.5 * v * (1 + erf(v / sqrt(2))); activation_params[0] == 0.
+  let u = v * 0.70710678118654752;
+  let a = abs(u);
+  let t = 1.0 / (1.0 + 0.3275911 * a);
+  let poly = ((((1.061405429 * t - 1.453152027) * t + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t;
+  return 0.5 * v * (1.0 + sign(u) * (1.0 - poly * exp(-a * a)));
+{% elif mode == "Gelu" or mode == "FastGelu" %}
+  // tanh approximation (ONNX Gelu approximate="tanh" arrives as activation_params[0] != 0;
+  // com.microsoft.FastGelu always). tanh saturates to +/-1 well inside the clamp.
+  let inner = v * (0.035677408136300125 * v * v + 0.7978845608028654);
+  return v * (0.5 + 0.5 * tanh(clamp(inner, -10.0, 10.0)));
 {% else %}
   return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
 {% endif %}
+}{% endmacro %}
+{{ fused_act_fn(activation, actAlpha, actBeta) }}
 {% endif %}
+{% set accumulatorPrefix = "sum" if tilePartialAccumulation else "acc" %}
+{% set tileT = "f16" if usesF16 else "f32" %}
+{% set wgX = 16 %}
+{% set wgY = 16 %}
 const BK: u32 = 16u;
 const BM: u32 = 32u;
 const BN: u32 = 32u;
 // Store one output element with the configured bias/residual/activation
 // epilogue in the f32 accumulator domain.
 fn store_out(m: u32, n: u32, yBase: u32, raw: f32) {
   if (m >= params.M || n >= params.N) {
     return;
   y[yBase + m * params.N + n] = {{ T }}(v);
 }
 var<workgroup> tileA: array<array<{{ tileT }}, BK>, BM>;
 var<workgroup> tileB: array<array<{{ tileT }}, BN>, BK>;
+@compute @workgroup_size({{ wgX }}, {{ wgY }}, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>
   let nBase = wg.x * BN;
   let batch = wg.z;
   let xBatchBase = batch * K * N;
+  let li = lid.y * {{ wgX }}u + lid.x;
   var acc00: f32 = 0.0;
   var acc01: f32 = 0.0;
   for (var kt: u32 = 0u; kt < numTiles; kt = kt + 1u) {
     let kBase = kt * BK;
     // Cooperative load: 32x16 A(=W) tile + 16x32 B(=X[batch]) tile, 256 threads x 2 each.
+    for (var e: u32 = 0u; e < (BM * BK) / {{ wgX * wgY }}u; e = e + 1u) {
+      let idx = li + e * {{ wgX * wgY }}u;
       let ar = idx / BK;
       let ac = idx % BK;
       let am = mBase + ar;
       }
     }
     workgroupBarrier();
+{% if tilePartialAccumulation %}
+    // Bound each serial dot-product chain by the K tile before merging it.
+{% for i in range(2) %}{% for j in range(2) %}
+    var sum{{ i }}{{ j }}: f32 = 0.0;
+{% endfor %}{% endfor %}
+{% endif %}
     for (var kk: u32 = 0u; kk < BK; kk = kk + 1u) {
       let a0 = f32(tileA[lid.y * 2u][kk]);
       let a1 = f32(tileA[lid.y * 2u + 1u][kk]);
       let b0 = f32(tileB[kk][lid.x * 2u]);
       let b1 = f32(tileB[kk][lid.x * 2u + 1u]);
+{% for i in range(2) %}{% for j in range(2) %}
+      {{ accumulatorPrefix }}{{ i }}{{ j }} = {{ accumulatorPrefix }}{{ i }}{{ j }} + a{{ i }} * b{{ j }};
+{% endfor %}{% endfor %}
     }
+{% if tilePartialAccumulation %}
+{% for i in range(2) %}{% for j in range(2) %}
+    acc{{ i }}{{ j }} = acc{{ i }}{{ j }} + sum{{ i }}{{ j }};
+{% endfor %}{% endfor %}
+{% endif %}
     workgroupBarrier();
   }

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

@@ -12,18 +12,13 @@
 // 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) %}
@@ -36,15 +31,9 @@
 {% 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 %}
@@ -53,46 +42,16 @@ 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" %}
 {% set useDirectMatrixStore = directMatrixStore is defined and directMatrixStore %}
 {% set tileRowsValue = tileRows if tileRows is defined else 32 %}
-{% set tileColsValue = tileCols if tileCols is defined else 64 %}
 {% set workgroupThreadsValue = workgroupThreads if workgroupThreads is defined else 128 %}
 {% set subgroupRowsValue = subgroupRows if subgroupRows is defined else 2 %}
-{% set subgroupColsValue = subgroupCols if subgroupCols is defined else 2 %}
 {% set subRowsValue = (tileRowsValue / subgroupRowsValue)|int %}
-{% set subColsValue = (tileColsValue / subgroupColsValue)|int %}
 {% set bLoadWidth = (tileColsValue * 32 / workgroupThreadsValue)|int %}
 {% set bKChunks = (32 / bLoadWidth)|int %}
@@ -107,18 +66,7 @@ const K_LOOP: u32 = {{ kPadded }}u;
 {% 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;
@@ -128,13 +76,8 @@ 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;
@@ -176,7 +119,7 @@ fn loadSHMA(tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
   for (var col_offset = 0u; col_offset < 8u; col_offset = col_offset + 1u) {
     let k = k_idx + col + col_offset;
     if (a_global < M{% if padded %} && k < K{% endif %}) {
-{% set A_PHASE = "phase_base + " if polyphase else "" %}
 {% if operandScalar == "f16" %}
       tile_A[row * TILE_K + col + col_offset] = f16(w[{{ A_PHASE }}a_global * K + k]);
 {% else %}
@@ -197,12 +140,7 @@ fn loadSHMB(b_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
   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)
@@ -213,16 +151,8 @@ fn loadSHMB(b_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
   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
@@ -230,30 +160,16 @@ fn loadSHMB(b_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
   // 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;
@@ -265,13 +181,6 @@ fn loadSHMB(b_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
         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 %}
         }
       }
     }
@@ -281,12 +190,8 @@ fn loadSHMB(b_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
     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;
@@ -300,17 +205,7 @@ fn loadSHMB(b_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
       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 %}
       }
     }
   }
@@ -325,16 +220,15 @@ fn loadSHMB(b_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
   }
 {% 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" %}
@@ -349,15 +243,45 @@ fn loadSHMB(b_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
   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 biasAdd = " + bv" if hasBias else "" %}
 {% if hasActivation or hasZ %}
 // Fused epilogue: Y = activation(conv + bias + Z), applied at the output store.
@@ -374,10 +298,9 @@ fn epi(raw: f32{% if hasZ %}, yIndex: u32{% endif %}) -> {{ T }} {
 }
 {% endif %}
 {% macro store_val(valExpr, idxExpr) %}
-{% if hasActivation or hasZ %}epi({{ valExpr }}{% if hasZ %}, {{ idxExpr }}{% endif %})
-{%- else %}{{ T }}({{ valExpr }})
-{%- endif %}
-{% endmacro %}
 {% if (not useDirectMatrixStore or padded) and not splitKPartial and not fusedNarrowProjection %}
 fn storeOutput(offset: u32, {% if hasBias or polyphase %}row_base: u32, {% endif %}row: u32, col: u32, src_slot: u32, row_limit: i32{% if padded %}, col_anchor: u32{% endif %}) {
@@ -431,9 +354,7 @@ fn main(
 {% 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;
@@ -491,57 +412,7 @@ fn main(
 {% endif %}
   }
-{% if fusedNarrowProjection %}
-  // The Conv producer owns exactly one 32-row tile. Its input tile is dead after
-  // the final K iteration, so reuse those 32x64 f32 cells as the only
-  // cross-subgroup handoff into the narrow projection. All collective stores
-  // remain subgroup-uniform, and every fragment has a disjoint destination.
-  let fused_tile_offset = base_A * TILE_COLS + base_B;
-  subgroupMatrixStore<row_major>(&tile_B, fused_tile_offset + 0u * TILE_COLS +  0u, matC00, TILE_COLS);
-  subgroupMatrixStore<row_major>(&tile_B, fused_tile_offset + 0u * TILE_COLS +  8u, matC01, TILE_COLS);
-  subgroupMatrixStore<row_major>(&tile_B, fused_tile_offset + 0u * TILE_COLS + 16u, matC02, TILE_COLS);
-  subgroupMatrixStore<row_major>(&tile_B, fused_tile_offset + 0u * TILE_COLS + 24u, matC03, TILE_COLS);
-  subgroupMatrixStore<row_major>(&tile_B, fused_tile_offset + 8u * TILE_COLS +  0u, matC10, TILE_COLS);
-  subgroupMatrixStore<row_major>(&tile_B, fused_tile_offset + 8u * TILE_COLS +  8u, matC11, TILE_COLS);
-  subgroupMatrixStore<row_major>(&tile_B, fused_tile_offset + 8u * TILE_COLS + 16u, matC12, TILE_COLS);
-  subgroupMatrixStore<row_major>(&tile_B, fused_tile_offset + 8u * TILE_COLS + 24u, matC13, TILE_COLS);
-  workgroupBarrier();
-  // One invocation owns one spatial column and walks the producer channels in
-  // increasing order to preserve the scalar projection's accumulation order.
-  // The inactive half of the workgroup has no remaining collective operation
-  // to reach.
-  let fused_global_col = b_global_base + local_idx;
-  if (local_idx < TILE_COLS && fused_global_col < N) {
-{% for oc in range(projectionOutChannels) %}
-    var projected_acc{{ oc }} = f32(0.0);
-{% endfor %}
-    for (var channel = 0u; channel < M; channel = channel + 1u) {
-      let producer_value = tile_B[channel * TILE_COLS + local_idx]{% if hasBias %} + f32(bias[channel]){% endif %};
-{% 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.

 // 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 = false %}
+{% set polyphase = false %}
 {% set implicit = implicitIm2col is defined and implicitIm2col %}
 {% 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 convKernelH = convKernelH | default(0) %}
 {% set convKernelW = convKernelW | default(0) %}
 {% set convStrideH = convStrideH | default(0) %}
 {% set convInW = convInW | default(0) %}
 {% set convOutW = convOutW | default(0) %}
 {% set convInChannels = convInChannels | default(0) %}
 {% set nPadded = nPadded | default(0) %}
 {% set mPadded = mPadded | default(0) %}
+{% set kChunk = 0 %}
 {% set batchCount = batchCount | default(0) %}
 enable subgroups;
 {% if pinSubgroupSize32 %}
 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
 {% set operandScalar = fScalar %}
 {% set accScalar = "f32" %}
 {% set useDirectMatrixStore = directMatrixStore is defined and directMatrixStore %}
 {% set tileRowsValue = tileRows if tileRows is defined else 32 %}
+{% set tileColsValue = 64 %}
 {% set workgroupThreadsValue = workgroupThreads if workgroupThreads is defined else 128 %}
 {% set subgroupRowsValue = subgroupRows if subgroupRows is defined else 2 %}
 {% set subRowsValue = (tileRowsValue / subgroupRowsValue)|int %}
+{% set subColsValue = 32 %}
 {% set bLoadWidth = (tileColsValue * 32 / workgroupThreadsValue)|int %}
 {% set bKChunks = (32 / bLoadWidth)|int %}
 {% if implicit %}
 const CONV_KERNEL_H: u32 = {{ convKernelH }}u;
 const CONV_KERNEL_W: u32 = {{ convKernelW }}u;
 const CONV_KSIZE: u32 = CONV_KERNEL_H * CONV_KERNEL_W;
 const CONV_STRIDE_H: u32 = {{ convStrideH }}u;
 const CONV_STRIDE_W: u32 = {{ convStrideW }}u;
 const CONV_DILATION_H: u32 = {{ convDilationH }}u;
 const CONV_IN_H: i32 = {{ convInH }};
 const CONV_IN_W: i32 = {{ convInW }};
 const CONV_OUT_W: u32 = {{ convOutW }}u;
 // 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);
 {% else %}
 {% if padded %}
 const N_COLS: u32 = {{ nPadded }}u;
   for (var col_offset = 0u; col_offset < 8u; col_offset = col_offset + 1u) {
     let k = k_idx + col + col_offset;
     if (a_global < M{% if padded %} && k < K{% endif %}) {
+{% set A_PHASE = "" %}
 {% if operandScalar == "f16" %}
       tile_A[row * TILE_K + col + col_offset] = f16(w[{{ A_PHASE }}a_global * K + k]);
 {% else %}
   let col = c_idx * {{ bLoadWidth }}u;
 {% if implicit %}
   // One output position per tile column, so its window origin is loop-invariant.
   let ih0 = i32((b_col / CONV_OUT_W) * CONV_STRIDE_H) - CONV_PAD_TOP;
   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)
   var k = k_idx + col;
   var ic = k / CONV_KSIZE;
   var kq = k % CONV_KSIZE;
   var kh = kq / CONV_KERNEL_W;
   var kw = kq % CONV_KERNEL_W;
   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
   // the interior loop carries a plain address instead of recomputing coordinates.
   let interior = in_column
     && k_idx + col + {{ bLoadWidth }}u <= K
     && 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;
     let planeStep = CONV_IN_H * CONV_IN_W;
     // Advancing k already added one colStep, so each wrap rewinds that plus the taps it walked.
     let kwWrap = colStep * i32(CONV_KERNEL_W) - rowStep;
     let khWrap = rowStep * i32(CONV_KERNEL_H) - planeStep;
     var addr = i32(b_base) + i32(ic) * planeStep + ih * CONV_IN_W + iw;
     for (var i = 0u; i < {{ bLoadWidth }}u; i = i + 1u) {
       tile_B[row * TILE_K + col + i] = {{ operandScalar }}(xm[u32(addr)]);
       addr = addr + colStep;
         if (kh == CONV_KERNEL_H) {
           kh = 0u;
           addr = addr - khWrap;
         }
       }
     }
     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 && ih >= 0 && ih < CONV_IN_H && iw >= 0 && iw < CONV_IN_W) {
       value = {{ operandScalar }}(xm[b_base + (ic * u32(CONV_IN_H) + u32(ih)) * u32(CONV_IN_W) + u32(iw)]);
     }
     tile_B[row * TILE_K + col + i] = value;
     k = k + 1u;
       if (kh == CONV_KERNEL_H) {
         kh = 0u;
         ih = ih0;
         ic = ic + 1u;
       }
     }
   }
   }
 {% endif %}
 }
+{% if hasActivation or hasZ %}
 {% endif %}
+{% endif %}
 {% if hasActivation %}
 // Apply the fused activation in the f32 accumulator before the single output
 // cast, avoiding an intermediate convolution tensor.
+{% macro fused_act_fn(mode, alpha, beta) %}
+fn fused_act(v: f32) -> f32 {
 {% if mode == "Relu" %}
   return max(v, 0.0);
 {% elif mode == "Clip" %}
   return tanh(clamp(v, -10.0, 10.0));
 {% elif mode == "HardSigmoid" %}
   return clamp(f32({{ alpha }}) * v + f32({{ beta }}), 0.0, 1.0);
+{% elif mode == "QuickGelu" %}
+{% if alpha == 1 %}
+  // alpha == 1 is SiLU: the multiply folds away (the spelling yolov8-class detectors emit).
+  return v * (1.0 / (1.0 + exp(-v)));
+{% else %}
+  return v * (1.0 / (1.0 + exp(-(f32({{ alpha }}) * v))));
+{% endif %}
+{% elif mode == "Elu" %}
+  return select(f32({{ alpha }}) * (exp(v) - 1.0), v, v >= 0.0);
+{% elif mode == "ThresholdedRelu" %}
+  return select(0.0, v, v > f32({{ alpha }}));
+{% elif mode == "Softplus" %}
+  // log(1 + exp(v)) spelled overflow-safe: max(v, 0) + log(1 + exp(-|v|)).
+  return max(v, 0.0) + log(1.0 + exp(-abs(v)));
+{% elif mode == "Erf" %}
+  // Abramowitz & Stegun 7.1.26 (|error| < 1.5e-7), the same polynomial the standalone
+  // Erf/Gelu kernels use, so a fused and an unfused graph agree to f32 rounding.
+  let a = abs(v);
+  let t = 1.0 / (1.0 + 0.3275911 * a);
+  let poly = ((((1.061405429 * t - 1.453152027) * t + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t;
+  return sign(v) * (1.0 - poly * exp(-a * a));
+{% elif mode == "Gelu" and alpha == 0 %}
+  // ONNX Gelu, erf form: 0.5 * v * (1 + erf(v / sqrt(2))); activation_params[0] == 0.
+  let u = v * 0.70710678118654752;
+  let a = abs(u);
+  let t = 1.0 / (1.0 + 0.3275911 * a);
+  let poly = ((((1.061405429 * t - 1.453152027) * t + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t;
+  return 0.5 * v * (1.0 + sign(u) * (1.0 - poly * exp(-a * a)));
+{% elif mode == "Gelu" or mode == "FastGelu" %}
+  // tanh approximation (ONNX Gelu approximate="tanh" arrives as activation_params[0] != 0;
+  // com.microsoft.FastGelu always). tanh saturates to +/-1 well inside the clamp.
+  let inner = v * (0.035677408136300125 * v * v + 0.7978845608028654);
+  return v * (0.5 + 0.5 * tanh(clamp(inner, -10.0, 10.0)));
 {% else %}
   return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
 {% endif %}
+}{% endmacro %}
+{{ fused_act_fn(activation, actAlpha, actBeta) }}
 {% endif %}
 {% set biasAdd = " + bv" if hasBias else "" %}
 {% if hasActivation or hasZ %}
 // Fused epilogue: Y = activation(conv + bias + Z), applied at the output store.
 }
 {% endif %}
 {% macro store_val(valExpr, idxExpr) %}
+{% if hasActivation or hasZ %}
+epi({{ valExpr }}{% if hasZ %}, {{ idxExpr }}{% endif %}){% else %}
+{{ T }}({{ valExpr }}){% endif %}{% endmacro %}
 {% if (not useDirectMatrixStore or padded) and not splitKPartial and not fusedNarrowProjection %}
 fn storeOutput(offset: u32, {% if hasBias or polyphase %}row_base: u32, {% endif %}row: u32, col: u32, src_slot: u32, row_limit: i32{% if padded %}, col_anchor: u32{% endif %}) {
 {% else %}
   let batch = workgroup_id.z;
   let b_base = batch * B_BATCH_STRIDE;
   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;
 {% endif %}
   }
+{% if 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.

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

@@ -1,18 +1,20 @@
 // 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" %}
@@ -27,24 +29,50 @@ enable f16;
   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 ow = index % params.outW;
   var q = index / params.outW;

 // 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.
+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
+  if ({{ name }} >= {{ bound }}) {
+    return;
+  }{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 {% if hasActivation %}
 // Apply the fused activation in the f32 accumulator before the single output
 // cast, avoiding an intermediate convolution tensor.
+{% macro fused_act_fn(mode, alpha, beta) %}
+fn fused_act(v: f32) -> f32 {
 {% if mode == "Relu" %}
   return max(v, 0.0);
 {% elif mode == "Clip" %}
   return tanh(clamp(v, -10.0, 10.0));
 {% elif mode == "HardSigmoid" %}
   return clamp(f32({{ alpha }}) * v + f32({{ beta }}), 0.0, 1.0);
+{% elif mode == "QuickGelu" %}
+{% if alpha == 1 %}
+  // alpha == 1 is SiLU: the multiply folds away (the spelling yolov8-class detectors emit).
+  return v * (1.0 / (1.0 + exp(-v)));
+{% else %}
+  return v * (1.0 / (1.0 + exp(-(f32({{ alpha }}) * v))));
+{% endif %}
+{% elif mode == "Elu" %}
+  return select(f32({{ alpha }}) * (exp(v) - 1.0), v, v >= 0.0);
+{% elif mode == "ThresholdedRelu" %}
+  return select(0.0, v, v > f32({{ alpha }}));
+{% elif mode == "Softplus" %}
+  // log(1 + exp(v)) spelled overflow-safe: max(v, 0) + log(1 + exp(-|v|)).
+  return max(v, 0.0) + log(1.0 + exp(-abs(v)));
+{% elif mode == "Erf" %}
+  // Abramowitz & Stegun 7.1.26 (|error| < 1.5e-7), the same polynomial the standalone
+  // Erf/Gelu kernels use, so a fused and an unfused graph agree to f32 rounding.
+  let a = abs(v);
+  let t = 1.0 / (1.0 + 0.3275911 * a);
+  let poly = ((((1.061405429 * t - 1.453152027) * t + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t;
+  return sign(v) * (1.0 - poly * exp(-a * a));
+{% elif mode == "Gelu" and alpha == 0 %}
+  // ONNX Gelu, erf form: 0.5 * v * (1 + erf(v / sqrt(2))); activation_params[0] == 0.
+  let u = v * 0.70710678118654752;
+  let a = abs(u);
+  let t = 1.0 / (1.0 + 0.3275911 * a);
+  let poly = ((((1.061405429 * t - 1.453152027) * t + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t;
+  return 0.5 * v * (1.0 + sign(u) * (1.0 - poly * exp(-a * a)));
+{% elif mode == "Gelu" or mode == "FastGelu" %}
+  // tanh approximation (ONNX Gelu approximate="tanh" arrives as activation_params[0] != 0;
+  // com.microsoft.FastGelu always). tanh saturates to +/-1 well inside the clamp.
+  let inner = v * (0.035677408136300125 * v * v + 0.7978845608028654);
+  return v * (0.5 + 0.5 * tanh(clamp(inner, -10.0, 10.0)));
 {% else %}
   return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
 {% endif %}
+}{% endmacro %}
+{{ fused_act_fn(activation, actAlpha, actBeta) }}
 {% endif %}
 const WG: u32 = {{ convWorkgroupSize }}u;
 @compute @workgroup_size(WG)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d("WG", "index") }}
   let ow = index % params.outW;
   var q = index / params.outW;

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

@@ -6,15 +6,20 @@
 // 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" %}
@@ -29,16 +34,45 @@
   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;
@@ -47,12 +81,7 @@ 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;
-  }
   let ow = index % params.outW;
   var t = index / params.outW;

 // 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.
+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
+  if ({{ name }} >= {{ bound }}) {
+    return;
+  }{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 {% if hasActivation %}
 // Apply the fused activation in the f32 accumulator before the single output
 // cast, avoiding an intermediate convolution tensor.
+{% macro fused_act_fn(mode, alpha, beta) %}
+fn fused_act(v: f32) -> f32 {
 {% if mode == "Relu" %}
   return max(v, 0.0);
 {% elif mode == "Clip" %}
   return tanh(clamp(v, -10.0, 10.0));
 {% elif mode == "HardSigmoid" %}
   return clamp(f32({{ alpha }}) * v + f32({{ beta }}), 0.0, 1.0);
+{% elif mode == "QuickGelu" %}
+{% if alpha == 1 %}
+  // alpha == 1 is SiLU: the multiply folds away (the spelling yolov8-class detectors emit).
+  return v * (1.0 / (1.0 + exp(-v)));
+{% else %}
+  return v * (1.0 / (1.0 + exp(-(f32({{ alpha }}) * v))));
+{% endif %}
+{% elif mode == "Elu" %}
+  return select(f32({{ alpha }}) * (exp(v) - 1.0), v, v >= 0.0);
+{% elif mode == "ThresholdedRelu" %}
+  return select(0.0, v, v > f32({{ alpha }}));
+{% elif mode == "Softplus" %}
+  // log(1 + exp(v)) spelled overflow-safe: max(v, 0) + log(1 + exp(-|v|)).
+  return max(v, 0.0) + log(1.0 + exp(-abs(v)));
+{% elif mode == "Erf" %}
+  // Abramowitz & Stegun 7.1.26 (|error| < 1.5e-7), the same polynomial the standalone
+  // Erf/Gelu kernels use, so a fused and an unfused graph agree to f32 rounding.
+  let a = abs(v);
+  let t = 1.0 / (1.0 + 0.3275911 * a);
+  let poly = ((((1.061405429 * t - 1.453152027) * t + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t;
+  return sign(v) * (1.0 - poly * exp(-a * a));
+{% elif mode == "Gelu" and alpha == 0 %}
+  // ONNX Gelu, erf form: 0.5 * v * (1 + erf(v / sqrt(2))); activation_params[0] == 0.
+  let u = v * 0.70710678118654752;
+  let a = abs(u);
+  let t = 1.0 / (1.0 + 0.3275911 * a);
+  let poly = ((((1.061405429 * t - 1.453152027) * t + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t;
+  return 0.5 * v * (1.0 + sign(u) * (1.0 - poly * exp(-a * a)));
+{% elif mode == "Gelu" or mode == "FastGelu" %}
+  // tanh approximation (ONNX Gelu approximate="tanh" arrives as activation_params[0] != 0;
+  // com.microsoft.FastGelu always). tanh saturates to +/-1 well inside the clamp.
+  let inner = v * (0.035677408136300125 * v * v + 0.7978845608028654);
+  return v * (0.5 + 0.5 * tanh(clamp(inner, -10.0, 10.0)));
 {% else %}
   return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
 {% endif %}
+}{% endmacro %}
+{{ fused_act_fn(activation, actAlpha, actBeta) }}
 {% endif %}
 const KERNEL_AREA: u32 = {{ kernelHSpec * kernelWSpec }}u;
 const STRIDE_H: u32 = {{ strideHSpec }}u;
 const STRIDE_W: u32 = {{ strideWSpec }}u;
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d(tunables.WORKGROUP_SIZE, "index") }}
   let ow = index % params.outW;
   var t = index / params.outW;

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

@@ -10,17 +10,22 @@
 //
 // 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" %}
@@ -35,15 +40,45 @@
   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;
@@ -57,12 +92,7 @@ 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);

 //
 // Summing slices in index order changes the f32 association relative to one
 // uninterrupted K loop, so the two orders need not be bit-identical.
+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
+  if ({{ name }} >= {{ bound }}) {
+    return;
+  }{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 {% set applyActivation = hasActivation is defined and hasActivation %}
 {% if applyActivation %}
 {% if hasActivation %}
 // Apply the fused activation in the f32 accumulator before the single output
 // cast, avoiding an intermediate convolution tensor.
+{% macro fused_act_fn(mode, alpha, beta) %}
+fn fused_act(v: f32) -> f32 {
 {% if mode == "Relu" %}
   return max(v, 0.0);
 {% elif mode == "Clip" %}
   return tanh(clamp(v, -10.0, 10.0));
 {% elif mode == "HardSigmoid" %}
   return clamp(f32({{ alpha }}) * v + f32({{ beta }}), 0.0, 1.0);
+{% elif mode == "QuickGelu" %}
+{% if alpha == 1 %}
+  // alpha == 1 is SiLU: the multiply folds away (the spelling yolov8-class detectors emit).
+  return v * (1.0 / (1.0 + exp(-v)));
+{% else %}
+  return v * (1.0 / (1.0 + exp(-(f32({{ alpha }}) * v))));
+{% endif %}
+{% elif mode == "Elu" %}
+  return select(f32({{ alpha }}) * (exp(v) - 1.0), v, v >= 0.0);
+{% elif mode == "ThresholdedRelu" %}
+  return select(0.0, v, v > f32({{ alpha }}));
+{% elif mode == "Softplus" %}
+  // log(1 + exp(v)) spelled overflow-safe: max(v, 0) + log(1 + exp(-|v|)).
+  return max(v, 0.0) + log(1.0 + exp(-abs(v)));
+{% elif mode == "Erf" %}
+  // Abramowitz & Stegun 7.1.26 (|error| < 1.5e-7), the same polynomial the standalone
+  // Erf/Gelu kernels use, so a fused and an unfused graph agree to f32 rounding.
+  let a = abs(v);
+  let t = 1.0 / (1.0 + 0.3275911 * a);
+  let poly = ((((1.061405429 * t - 1.453152027) * t + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t;
+  return sign(v) * (1.0 - poly * exp(-a * a));
+{% elif mode == "Gelu" and alpha == 0 %}
+  // ONNX Gelu, erf form: 0.5 * v * (1 + erf(v / sqrt(2))); activation_params[0] == 0.
+  let u = v * 0.70710678118654752;
+  let a = abs(u);
+  let t = 1.0 / (1.0 + 0.3275911 * a);
+  let poly = ((((1.061405429 * t - 1.453152027) * t + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t;
+  return 0.5 * v * (1.0 + sign(u) * (1.0 - poly * exp(-a * a)));
+{% elif mode == "Gelu" or mode == "FastGelu" %}
+  // tanh approximation (ONNX Gelu approximate="tanh" arrives as activation_params[0] != 0;
+  // com.microsoft.FastGelu always). tanh saturates to +/-1 well inside the clamp.
+  let inner = v * (0.035677408136300125 * v * v + 0.7978845608028654);
+  return v * (0.5 + 0.5 * tanh(clamp(inner, -10.0, 10.0)));
 {% else %}
   return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
 {% endif %}
+}{% endmacro %}
+{{ fused_act_fn(activation, actAlpha, actBeta) }}
 {% endif %}
 {% endif %}
 const M: u32 = {{ M }}u;
 @compute @workgroup_size({{ reduceWorkgroupSize }}, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d("WORKGROUP_SIZE", "idx", "COUNT") }}
   let col = idx % N;
   let row = (idx / N) % M;
   let batch = idx / (M * N);

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

@@ -10,14 +10,11 @@
 // 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" %}
@@ -32,15 +29,45 @@
   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 tileT = "f16" if usesF16 else "f32" %}
 const WG_X: u32 = {{ convWgX }}u;

 // 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 }}
 {% if hasActivation %}
 // Apply the fused activation in the f32 accumulator before the single output
 // cast, avoiding an intermediate convolution tensor.
+{% macro fused_act_fn(mode, alpha, beta) %}
+fn fused_act(v: f32) -> f32 {
 {% if mode == "Relu" %}
   return max(v, 0.0);
 {% elif mode == "Clip" %}
   return tanh(clamp(v, -10.0, 10.0));
 {% elif mode == "HardSigmoid" %}
   return clamp(f32({{ alpha }}) * v + f32({{ beta }}), 0.0, 1.0);
+{% elif mode == "QuickGelu" %}
+{% if alpha == 1 %}
+  // alpha == 1 is SiLU: the multiply folds away (the spelling yolov8-class detectors emit).
+  return v * (1.0 / (1.0 + exp(-v)));
+{% else %}
+  return v * (1.0 / (1.0 + exp(-(f32({{ alpha }}) * v))));
+{% endif %}
+{% elif mode == "Elu" %}
+  return select(f32({{ alpha }}) * (exp(v) - 1.0), v, v >= 0.0);
+{% elif mode == "ThresholdedRelu" %}
+  return select(0.0, v, v > f32({{ alpha }}));
+{% elif mode == "Softplus" %}
+  // log(1 + exp(v)) spelled overflow-safe: max(v, 0) + log(1 + exp(-|v|)).
+  return max(v, 0.0) + log(1.0 + exp(-abs(v)));
+{% elif mode == "Erf" %}
+  // Abramowitz & Stegun 7.1.26 (|error| < 1.5e-7), the same polynomial the standalone
+  // Erf/Gelu kernels use, so a fused and an unfused graph agree to f32 rounding.
+  let a = abs(v);
+  let t = 1.0 / (1.0 + 0.3275911 * a);
+  let poly = ((((1.061405429 * t - 1.453152027) * t + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t;
+  return sign(v) * (1.0 - poly * exp(-a * a));
+{% elif mode == "Gelu" and alpha == 0 %}
+  // ONNX Gelu, erf form: 0.5 * v * (1 + erf(v / sqrt(2))); activation_params[0] == 0.
+  let u = v * 0.70710678118654752;
+  let a = abs(u);
+  let t = 1.0 / (1.0 + 0.3275911 * a);
+  let poly = ((((1.061405429 * t - 1.453152027) * t + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t;
+  return 0.5 * v * (1.0 + sign(u) * (1.0 - poly * exp(-a * a)));
+{% elif mode == "Gelu" or mode == "FastGelu" %}
+  // tanh approximation (ONNX Gelu approximate="tanh" arrives as activation_params[0] != 0;
+  // com.microsoft.FastGelu always). tanh saturates to +/-1 well inside the clamp.
+  let inner = v * (0.035677408136300125 * v * v + 0.7978845608028654);
+  return v * (0.5 + 0.5 * tanh(clamp(inner, -10.0, 10.0)));
 {% else %}
   return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
 {% endif %}
+}{% endmacro %}
+{{ fused_act_fn(activation, actAlpha, actBeta) }}
 {% endif %}
 {% set tileT = "f16" if usesF16 else "f32" %}
 const WG_X: u32 = {{ convWgX }}u;

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

@@ -1,12 +1,15 @@
 {{ 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" %}
@@ -21,23 +24,50 @@
   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 %}
 // Large grouped kernels are dominated by repeatedly fetching the same filter
 // value for neighboring output columns and input window for neighboring output
 // channels. One invocation computes four adjacent columns for OC_TILE channels,
 // sharing each input load across those channels. Shape and window geometry are
 // static, so 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 %}
@@ -86,8 +116,7 @@ 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;
@@ -113,19 +142,25 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   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 %}
@@ -165,7 +200,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 {% 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) %}
@@ -178,7 +213,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 {% 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) %}
@@ -196,7 +231,8 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
       }
     }
 {% endfor %}
-    xChannelBase += IN_PLANE;
 {% for oct in range(ocTile) %}
     wBase{{ oct }} += KAREA;
 {% endfor %}

+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
+  if ({{ name }} >= {{ bound }}) { return; }{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 {% if hasActivation %}
 // Apply the fused activation in the f32 accumulator before the single output
 // cast, avoiding an intermediate convolution tensor.
+{% macro fused_act_fn(mode, alpha, beta) %}
+fn fused_act(v: f32) -> f32 {
 {% if mode == "Relu" %}
   return max(v, 0.0);
 {% elif mode == "Clip" %}
   return tanh(clamp(v, -10.0, 10.0));
 {% elif mode == "HardSigmoid" %}
   return clamp(f32({{ alpha }}) * v + f32({{ beta }}), 0.0, 1.0);
+{% elif mode == "QuickGelu" %}
+{% if alpha == 1 %}
+  // alpha == 1 is SiLU: the multiply folds away (the spelling yolov8-class detectors emit).
+  return v * (1.0 / (1.0 + exp(-v)));
+{% else %}
+  return v * (1.0 / (1.0 + exp(-(f32({{ alpha }}) * v))));
+{% endif %}
+{% elif mode == "Elu" %}
+  return select(f32({{ alpha }}) * (exp(v) - 1.0), v, v >= 0.0);
+{% elif mode == "ThresholdedRelu" %}
+  return select(0.0, v, v > f32({{ alpha }}));
+{% elif mode == "Softplus" %}
+  // log(1 + exp(v)) spelled overflow-safe: max(v, 0) + log(1 + exp(-|v|)).
+  return max(v, 0.0) + log(1.0 + exp(-abs(v)));
+{% elif mode == "Erf" %}
+  // Abramowitz & Stegun 7.1.26 (|error| < 1.5e-7), the same polynomial the standalone
+  // Erf/Gelu kernels use, so a fused and an unfused graph agree to f32 rounding.
+  let a = abs(v);
+  let t = 1.0 / (1.0 + 0.3275911 * a);
+  let poly = ((((1.061405429 * t - 1.453152027) * t + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t;
+  return sign(v) * (1.0 - poly * exp(-a * a));
+{% elif mode == "Gelu" and alpha == 0 %}
+  // ONNX Gelu, erf form: 0.5 * v * (1 + erf(v / sqrt(2))); activation_params[0] == 0.
+  let u = v * 0.70710678118654752;
+  let a = abs(u);
+  let t = 1.0 / (1.0 + 0.3275911 * a);
+  let poly = ((((1.061405429 * t - 1.453152027) * t + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t;
+  return 0.5 * v * (1.0 + sign(u) * (1.0 - poly * exp(-a * a)));
+{% elif mode == "Gelu" or mode == "FastGelu" %}
+  // tanh approximation (ONNX Gelu approximate="tanh" arrives as activation_params[0] != 0;
+  // com.microsoft.FastGelu always). tanh saturates to +/-1 well inside the clamp.
+  let inner = v * (0.035677408136300125 * v * v + 0.7978845608028654);
+  return v * (0.5 + 0.5 * tanh(clamp(inner, -10.0, 10.0)));
 {% else %}
   return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
 {% endif %}
+}{% endmacro %}
+{{ fused_act_fn(activation, actAlpha, actBeta) }}
 {% endif %}
 // Large grouped kernels are dominated by repeatedly fetching the same filter
 // value for neighboring output columns and input window for neighboring output
 // channels. One invocation computes four adjacent columns for OC_TILE channels,
 // sharing each input load across those channels. Shape and window geometry are
 // static, so indexing divisors and kernel offsets are shader constants.
 {% 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 %}
 @compute @workgroup_size({{ workgroupSizeSpec }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d("WG", "q", "COUNT_TILES", guardInline=true) }}
 {% if dilatedLanes %}
   let tile = q % TILES_PER_ROW;
   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 hasBias %}
+  // Seeding the accumulator with the channel bias replaces the zero seed and
+  // incorporates the bias before the reduction.
+  let bias{{ oct }} = f32(bias[ocBase + {{ oct }}u]);
+{% endif %}
 {% if dilatedLanes %}
 {% for lane in range(lanes) %}
+  var acc{{ oct }}_{{ lane }} = {% if hasBias %}bias{{ oct }}{% else %}0.0{% endif %};
 {% endfor %}
 {% else %}
+  var acc{{ oct }} = vec4<f32>({% if hasBias %}bias{{ oct }}{% else %}0.0{% endif %});
 {% endif %}
 {% endfor %}
   for (var ic = 0u; ic < IN_CPG; ic++) {
+{% if rollKernelRows %}  for (var kernelRow = 0u; kernelRow < {{ kernelHSpec }}u; kernelRow++) {
+{% endif %}{% for kh in range(1 if rollKernelRows else kernelHSpec) %}
     {
+      let ih = ihBase + {% if rollKernelRows %}i32(kernelRow) * {{ dilationHSpec }}{% else %}{{ kh * dilationHSpec }}{% endif %};
       if (ih >= 0 && ih < IN_H) {
         let xRow = xChannelBase + u32(ih) * IN_W_U;
 {% if registerForm %}
 {% for kw in range(kernelWSpec) %}
         {
 {% for oct in range(ocTile) %}
+          let weight{{ oct }} = f32(w[wBase{{ oct }} + {% if rollKernelRows %}kernelRow * {{ kernelWSpec }}u + {{ kw }}u{% else %}{{ kh * kernelWSpec + kw }}u{% endif %}]);
 {% endfor %}
 {% for lane in range(lanes) %}
 {% for oct in range(ocTile) %}
 {% for kw in range(kernelWSpec) %}
         {
 {% for oct in range(ocTile) %}
+          let weight{{ oct }} = f32(w[wBase{{ oct }} + {% if rollKernelRows %}kernelRow * {{ kernelWSpec }}u + {{ kw }}u{% else %}{{ kh * kernelWSpec + kw }}u{% endif %}]);
 {% endfor %}
           let iwK = iwBase + {{ kw * dilationWSpec }};
 {% for lane in range(lanes) %}
       }
     }
 {% endfor %}
+{% if rollKernelRows %}  }
+{% endif %}    xChannelBase += IN_PLANE;
 {% for oct in range(ocTile) %}
     wBase{{ oct }} += KAREA;
 {% endfor %}

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,29 +1,29 @@
 {
   "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"],
@@ -57,6 +57,8 @@
       "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"],
@@ -76,8 +78,11 @@
       "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"],

 {
   "name": "com.microsoft.FusedConv",
+  "id": "_com_microsoft_fusedconv_webgpu_a0813e7",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "h7j7tWEEOwdOjHUeFKfE46+kvmVgwWNFJTEcVB6dJK8=",
+      "conv-1x1-gemm-tiled-reg.wgsl.jinja": "J6MUpDFROmGCbCEln+xq1imr5+In47OPstc5DBcqv8U=",
+      "conv-1x1-gemm-tiled.wgsl.jinja": "eFyqUTNykx9plj5o3Se9dEHR4W3H/lFjrTgmQ7N/81A=",
+      "conv-1x1-subgroup-matrix.wgsl.jinja": "1f8hTKEnvhrxDuYMBFnXRjDoxKxr9u4ToeanV/Hjk5s=",
+      "conv-direct-nd.wgsl.jinja": "7CVIjupPj0kd1lioOiY4PZY2JLcHYiMZ2CnykIoPnpQ=",
+      "conv-direct-unrolled.wgsl.jinja": "efRWnTZ8xYUWP+Y8JpGN7w5r7log+H+qTZDTkIyH6gE=",
       "conv-im2col-nchw.wgsl.jinja": "IRpmtpeVQsLE6ud+Sf1230XaZVS1dxDtTZuSDwU3P8s=",
+      "conv-splitk-reduce.wgsl.jinja": "rxt8ttb0FFw6RpO6hfDbpneMPt5AqD1A7LwA1kiBG+w=",
+      "conv1d-tiled-reg.wgsl.jinja": "+SsmtbIw2VkES8J32j1SJ6tGZgHon0Z1ZySA8Ga1iPk=",
+      "conv2d-grouped-large-w4.wgsl.jinja": "2KdxhCkpdSFYVRX4LIRnTLHBYW0oCx6MbQmzvauz61Y=",
+      "manifest.json": "Kwp+pHg8GOMiDPWE3CeSMu6UAzi+RAOtEA4wninC8Yk=",
+      "test.json": "YY8eK1tSLM+NV1kDm5q+BugcjSDAeOS0mpny06NELas="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "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"],
       "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"],
+      "im2col_gemm_tiled_reg_f16_columns": ["conv-1x1-gemm-tiled-reg.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
+      "im2col_gemm_tiled_bias_reg_f16_columns": ["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"],
       "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_bias": ["conv2d-grouped-large-w4.wgsl.jinja"],
       "grouped_large_kernel_w4_tail": ["conv2d-grouped-large-w4.wgsl.jinja"],
+      "grouped_large_kernel_w4_tail_bias": ["conv2d-grouped-large-w4.wgsl.jinja"],
       "grouped_large_kernel_w4_dilated_lanes": ["conv2d-grouped-large-w4.wgsl.jinja"],
+      "grouped_large_kernel_w4_dilated_lanes_bias": ["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"],

build/webgpu/test.json CHANGED Viewed

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