Xenova HF Staff commited on
Commit
7b19a1e
·
verified ·
1 Parent(s): 2284b3a

sync 6fdf6301e2bb

Browse files
README.md CHANGED
@@ -53,20 +53,24 @@ Attributes and default values (overridable per request):
53
 
54
  One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
55
 
 
 
 
 
 
 
56
  - `pointwise_channel_reduce` — Handles pointwise unit-stride convolution with few output channels by assigning one invocation to each spatial position and accumulating every output channel in registers. It does not materialize im2col scratch.
57
  - `pointwise_channel_reduce_bias` — Handles pointwise unit-stride convolution with few output channels by assigning one invocation to each spatial position and accumulating every output channel in registers. It does not materialize im2col scratch.
58
  - `gemm_1x1_subgroup_matrix` — Multiplies 1x1 weights directly by NCHW input. Complete output-channel tiles use direct subgroup-matrix loads, eliminating input staging and reduction-loop barriers; partial tiles retain guarded staging. Tile geometry follows the device's workgroup limits.
59
  - `gemm_1x1_subgroup_matrix_bias` — Multiplies 1x1 weights directly by NCHW input. Complete output-channel tiles use direct subgroup-matrix loads, eliminating input staging and reduction-loop barriers; partial tiles retain guarded staging. Tile geometry follows the device's workgroup limits.
60
  - `im2col_gemm_subgroup_matrix` — Materializes convolution windows as a column matrix, then multiplies it by the weights with subgroup-matrix operations. It applies when reduction and spatial extents are tile-aligned and column scratch fits.
61
  - `im2col_gemm_subgroup_matrix_bias` — Materializes convolution windows as a column matrix, then multiplies it by the weights with subgroup-matrix operations. It applies when reduction and spatial extents are tile-aligned and column scratch fits.
62
- - `im2col_direct_f32_subgroup_matrix` — Materialize tile-aligned convolution columns and load weights and columns directly into subgroup matrices, removing operand staging and K-loop barriers. Retain f32 accumulation, existing split-K reduction order, and the implicit gather preference below the measured channel-tile reuse threshold.
63
- - `im2col_direct_f16_subgroup_matrix` — Materialize tile-aligned convolution columns and load weights and columns directly into subgroup matrices, removing operand staging and K-loop barriers. Retain f32 accumulation, existing split-K reduction order, and the implicit gather preference below the measured channel-tile reuse threshold.
64
- - `im2col_direct_f32_subgroup_matrix_bias` — Materialize tile-aligned convolution columns and load weights and columns directly into subgroup matrices, removing operand staging and K-loop barriers. Retain f32 accumulation, existing split-K reduction order, and the implicit gather preference below the measured channel-tile reuse threshold.
65
- - `im2col_direct_f16_subgroup_matrix_bias` — Materialize tile-aligned convolution columns and load weights and columns directly into subgroup matrices, removing operand staging and K-loop barriers. Retain f32 accumulation, existing split-K reduction order, and the implicit gather preference below the measured channel-tile reuse threshold.
66
- - `implicit_im2col_subgroup_matrix` — Gathers each logical im2col tile directly from X instead of materializing a column matrix, trading extra address arithmetic for the scratch and bandwidth that materialization would cost. It applies when the output has few output-channel tiles, so the gather is not repeated across many tiles.
67
- - `implicit_im2col_subgroup_matrix_bias` — Gathers each logical im2col tile directly from X instead of materializing a column matrix, trading extra address arithmetic for the scratch and bandwidth that materialization would cost. It applies when the output has few output-channel tiles, so the gather is not repeated across many tiles.
68
- - `implicit_im2col3d_subgroup_matrix` — Gathers each logical 3-D im2col tile directly from X instead of materializing a column matrix, trading extra address arithmetic for the scratch and bandwidth a volumetric materialization would cost. It applies when the output has few output-channel tiles, so the gather is not repeated across many tiles.
69
- - `implicit_im2col3d_subgroup_matrix_bias` — Gathers each logical 3-D im2col tile directly from X instead of materializing a column matrix, trading extra address arithmetic for the scratch and bandwidth a volumetric materialization would cost. It applies when the output has few output-channel tiles, so the gather is not repeated across many tiles.
70
  - `implicit_im2col_subgroup_matrix_splitk` — Partitions the implicit-im2col reduction across workgroups when the unsplit output grid is too small, then combines raw partial sums from scratch. The combine reassociates floating-point addition relative to the unsplit route.
71
  - `implicit_im2col_subgroup_matrix_bias_splitk` — Partitions the implicit-im2col reduction across workgroups when the unsplit output grid is too small, then combines raw partial sums from scratch. The combine reassociates floating-point addition relative to the unsplit route.
72
  - `im2col_gemm_subgroup_matrix_padded` — Materializes a zero-padded column matrix before a subgroup-matrix multiply, allowing logical reduction and spatial extents that are not tile-aligned. Padded values contribute zero and are excluded from the logical output.
@@ -78,6 +82,14 @@ One implementation is selected per call from the device capabilities, the reques
78
  - `implicit_im2col_tiled_bias_reg` — Gathers logical im2col elements directly from X while filling register-blocked GEMM tiles. It avoids column scratch and provides the portable implicit route without subgroup-matrix support.
79
  - `implicit_im2col_tiled_reg_splitk` — Partitions the register-blocked implicit-im2col reduction across workgroups when the unsplit output grid is too small, then combines raw partial sums. The combine reassociates floating-point addition relative to the unsplit route.
80
  - `implicit_im2col_tiled_bias_reg_splitk` — Partitions the register-blocked implicit-im2col reduction across workgroups when the unsplit output grid is too small, then combines raw partial sums. The combine reassociates floating-point addition relative to the unsplit route.
 
 
 
 
 
 
 
 
81
  - `gemm_1x1_tiled_reg_splitk` — Partitions the register-blocked 1x1 GEMM reduction across additional workgroups when the output grid is small, then combines raw partial sums. The combine reassociates floating-point addition relative to the unsplit route.
82
  - `gemm_1x1_tiled_bias_reg_splitk` — Partitions the register-blocked 1x1 GEMM reduction across additional workgroups when the output grid is small, then combines raw partial sums. The combine reassociates floating-point addition relative to the unsplit route.
83
 
@@ -90,7 +102,7 @@ Some implementation variants require `subgroup-matrix`, `shader-f16`, and `subgr
90
  - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
91
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
92
  - [`test.json`](build/webgpu/test.json) — correctness cases
93
- - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
94
  - [`conv-1x1-channel-reduce.wgsl.jinja`](build/webgpu/conv-1x1-channel-reduce.wgsl.jinja)
95
  - [`conv-1x1-gemm-tiled-reg.wgsl.jinja`](build/webgpu/conv-1x1-gemm-tiled-reg.wgsl.jinja)
96
  - [`conv-1x1-gemm-tiled.wgsl.jinja`](build/webgpu/conv-1x1-gemm-tiled.wgsl.jinja)
@@ -106,7 +118,7 @@ Some implementation variants require `subgroup-matrix`, `shader-f16`, and `subgr
106
  ## Use with `@huggingface/kernels`
107
 
108
  ```sh
109
- npm install --save-exact @huggingface/kernels@0.0.1-preview.2
110
  ```
111
 
112
  Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 
53
 
54
  One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
55
 
56
+ - `grouped_large_kernel_w4` — Shares each input row window across output channels and columns with f32 accumulation. Multi-channel groups loop over kernel rows to limit generated instructions; depthwise groups retain unrolling. Dilation-aware lanes reuse overlapping taps. Unrolled large kernels are demoted when the adapter does not report a fixed 32-wide subgroup range.
57
+ - `grouped_large_kernel_w4_bias` — Shares each input row window across output channels and columns with f32 accumulation. Multi-channel groups loop over kernel rows to limit generated instructions; depthwise groups retain unrolling. Dilation-aware lanes reuse overlapping taps. Unrolled large kernels are demoted when the adapter does not report a fixed 32-wide subgroup range.
58
+ - `grouped_large_kernel_w4_tail` — Shares each input row window across output channels and columns with f32 accumulation. Multi-channel groups loop over kernel rows to limit generated instructions; depthwise groups retain unrolling. Dilation-aware lanes reuse overlapping taps. Unrolled large kernels are demoted when the adapter does not report a fixed 32-wide subgroup range.
59
+ - `grouped_large_kernel_w4_tail_bias` — Shares each input row window across output channels and columns with f32 accumulation. Multi-channel groups loop over kernel rows to limit generated instructions; depthwise groups retain unrolling. Dilation-aware lanes reuse overlapping taps. Unrolled large kernels are demoted when the adapter does not report a fixed 32-wide subgroup range.
60
+ - `grouped_large_kernel_w4_dilated_lanes` — Shares each input row window across output channels and columns with f32 accumulation. Multi-channel groups loop over kernel rows to limit generated instructions; depthwise groups retain unrolling. Dilation-aware lanes reuse overlapping taps. Unrolled large kernels are demoted when the adapter does not report a fixed 32-wide subgroup range.
61
+ - `grouped_large_kernel_w4_dilated_lanes_bias` — Shares each input row window across output channels and columns with f32 accumulation. Multi-channel groups loop over kernel rows to limit generated instructions; depthwise groups retain unrolling. Dilation-aware lanes reuse overlapping taps. Unrolled large kernels are demoted when the adapter does not report a fixed 32-wide subgroup range.
62
  - `pointwise_channel_reduce` — Handles pointwise unit-stride convolution with few output channels by assigning one invocation to each spatial position and accumulating every output channel in registers. It does not materialize im2col scratch.
63
  - `pointwise_channel_reduce_bias` — Handles pointwise unit-stride convolution with few output channels by assigning one invocation to each spatial position and accumulating every output channel in registers. It does not materialize im2col scratch.
64
  - `gemm_1x1_subgroup_matrix` — Multiplies 1x1 weights directly by NCHW input. Complete output-channel tiles use direct subgroup-matrix loads, eliminating input staging and reduction-loop barriers; partial tiles retain guarded staging. Tile geometry follows the device's workgroup limits.
65
  - `gemm_1x1_subgroup_matrix_bias` — Multiplies 1x1 weights directly by NCHW input. Complete output-channel tiles use direct subgroup-matrix loads, eliminating input staging and reduction-loop barriers; partial tiles retain guarded staging. Tile geometry follows the device's workgroup limits.
66
  - `im2col_gemm_subgroup_matrix` — Materializes convolution windows as a column matrix, then multiplies it by the weights with subgroup-matrix operations. It applies when reduction and spatial extents are tile-aligned and column scratch fits.
67
  - `im2col_gemm_subgroup_matrix_bias` — Materializes convolution windows as a column matrix, then multiplies it by the weights with subgroup-matrix operations. It applies when reduction and spatial extents are tile-aligned and column scratch fits.
68
+ - `im2col_direct_f32_subgroup_matrix` — Materializes tile-aligned convolution columns and loads weights and columns directly into subgroup matrices, with no operand staging or K-loop barriers. Accumulates in f32 in the split-K reduction order of the other column-matrix routes; below the channel-tile reuse threshold, implicit gathering is used instead.
69
+ - `im2col_direct_f16_subgroup_matrix` — Materializes tile-aligned convolution columns and loads weights and columns directly into subgroup matrices, with no operand staging or K-loop barriers. Accumulates in f32 in the split-K reduction order of the other column-matrix routes; below the channel-tile reuse threshold, implicit gathering is used instead.
70
+ - `im2col_direct_f32_subgroup_matrix_bias` — Materializes tile-aligned convolution columns and loads weights and columns directly into subgroup matrices, with no operand staging or K-loop barriers. Accumulates in f32 in the split-K reduction order of the other column-matrix routes; below the channel-tile reuse threshold, implicit gathering is used instead.
71
+ - `im2col_direct_f16_subgroup_matrix_bias` — Materializes tile-aligned convolution columns and loads weights and columns directly into subgroup matrices, with no operand staging or K-loop barriers. Accumulates in f32 in the split-K reduction order of the other column-matrix routes; below the channel-tile reuse threshold, implicit gathering is used instead.
72
+ - `implicit_im2col_subgroup_matrix` — Gathers each logical 2-D or 3-D im2col tile directly from X instead of materializing a column matrix, trading extra address arithmetic for the scratch and bandwidth that materialization would cost. It applies when the output has few output-channel tiles, so the gather is not repeated across many tiles.
73
+ - `implicit_im2col_subgroup_matrix_bias` — Gathers each logical 2-D or 3-D im2col tile directly from X instead of materializing a column matrix, trading extra address arithmetic for the scratch and bandwidth that materialization would cost. It applies when the output has few output-channel tiles, so the gather is not repeated across many tiles.
 
 
74
  - `implicit_im2col_subgroup_matrix_splitk` — Partitions the implicit-im2col reduction across workgroups when the unsplit output grid is too small, then combines raw partial sums from scratch. The combine reassociates floating-point addition relative to the unsplit route.
75
  - `implicit_im2col_subgroup_matrix_bias_splitk` — Partitions the implicit-im2col reduction across workgroups when the unsplit output grid is too small, then combines raw partial sums from scratch. The combine reassociates floating-point addition relative to the unsplit route.
76
  - `im2col_gemm_subgroup_matrix_padded` — Materializes a zero-padded column matrix before a subgroup-matrix multiply, allowing logical reduction and spatial extents that are not tile-aligned. Padded values contribute zero and are excluded from the logical output.
 
82
  - `implicit_im2col_tiled_bias_reg` — Gathers logical im2col elements directly from X while filling register-blocked GEMM tiles. It avoids column scratch and provides the portable implicit route without subgroup-matrix support.
83
  - `implicit_im2col_tiled_reg_splitk` — Partitions the register-blocked implicit-im2col reduction across workgroups when the unsplit output grid is too small, then combines raw partial sums. The combine reassociates floating-point addition relative to the unsplit route.
84
  - `implicit_im2col_tiled_bias_reg_splitk` — Partitions the register-blocked implicit-im2col reduction across workgroups when the unsplit output grid is too small, then combines raw partial sums. The combine reassociates floating-point addition relative to the unsplit route.
85
+ - `im2col_gemm_tiled_reg` — Materialized f32 columns feed the configured tiled GEMM with f32 accumulation.
86
+ - `im2col_gemm_tiled` — Materialized f32 columns feed the configured tiled GEMM with f32 accumulation.
87
+ - `im2col_gemm_tiled_bias_reg` — Materialized f32 columns feed the configured tiled GEMM with f32 accumulation. Bias is applied before the final output cast.
88
+ - `im2col_gemm_tiled_bias` — Materialized f32 columns feed the configured tiled GEMM with f32 accumulation. Bias is applied before the final output cast.
89
+ - `im2col_gemm_tiled_reg_f16_columns` — Exact f16 columns feed the register-tiled GEMM with f32 accumulation. Total batched output workgroups determine eligibility. The row tile rounds output channels to 32 and caps at 64, reducing padding for narrow outputs. Actual tile storage and invocation geometry must fit the WebGPU limits.
90
+ - `im2col_gemm_tiled_f16_columns` — Exact f16 columns feed the smaller tiled GEMM with f32 accumulation. Short batched spatial workloads use the f32 column layout when it fits; this half-sized layout serves constrained buffers. The spatial-work crossover is independent of the gather workgroup size.
91
+ - `im2col_gemm_tiled_bias_reg_f16_columns` — Exact f16 columns feed register-tiled GEMM with f32 accumulation. Batched output work determines eligibility; channel tiles round to 32 and cap at 64. Actual storage and invocation geometry must fit WebGPU limits. One bias load per microtile row feeds its columns before the final output cast.
92
+ - `im2col_gemm_tiled_bias_f16_columns` — Exact f16 columns feed the smaller tiled GEMM with f32 accumulation. Short batched spatial workloads use the f32 column layout when it fits; this half-sized layout serves constrained buffers. The spatial-work crossover is independent of the gather workgroup size. Bias is applied before the final output cast.
93
  - `gemm_1x1_tiled_reg_splitk` — Partitions the register-blocked 1x1 GEMM reduction across additional workgroups when the output grid is small, then combines raw partial sums. The combine reassociates floating-point addition relative to the unsplit route.
94
  - `gemm_1x1_tiled_bias_reg_splitk` — Partitions the register-blocked 1x1 GEMM reduction across additional workgroups when the output grid is small, then combines raw partial sums. The combine reassociates floating-point addition relative to the unsplit route.
95
 
 
102
  - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
103
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
104
  - [`test.json`](build/webgpu/test.json) — correctness cases
105
+ - [`bench.json`](build/webgpu/bench.json) — benchmark cases
106
  - [`conv-1x1-channel-reduce.wgsl.jinja`](build/webgpu/conv-1x1-channel-reduce.wgsl.jinja)
107
  - [`conv-1x1-gemm-tiled-reg.wgsl.jinja`](build/webgpu/conv-1x1-gemm-tiled-reg.wgsl.jinja)
108
  - [`conv-1x1-gemm-tiled.wgsl.jinja`](build/webgpu/conv-1x1-gemm-tiled.wgsl.jinja)
 
118
  ## Use with `@huggingface/kernels`
119
 
120
  ```sh
121
+ npm install --save-exact @huggingface/kernels@0.0.1-preview.3
122
  ```
123
 
124
  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
The diff for this file is too large to render. See raw diff
 
build/webgpu/conv-1x1-channel-reduce.wgsl.jinja CHANGED
@@ -1,3 +1,11 @@
 
 
 
 
 
 
 
 
1
  // Pointwise convolution with few output channels is evaluated as a per-position
2
  // channel reduction instead of a tiled GEMM. One invocation walks the input
3
  // channels and accumulates every output channel in a statically named local.
@@ -6,12 +14,6 @@
6
  //
7
  // The output-channel accumulators are named explicitly. Selection caps the
8
  // output-channel count, bounding the invocation-local accumulator storage.
9
- {% set usesF16 = usesF16 is defined and usesF16 %}
10
- {% set batchCount = batchCount | default(0) %}
11
- {% set yScalar = yScalar | default("f32") %}
12
- {% if usesF16 %}
13
- enable f16;
14
- {% endif %}
15
  {{ env.wgsl.resourceDeclarations }}
16
 
17
  const IN_C: u32 = {{ inChannels }}u;
@@ -22,12 +24,7 @@ const WORKGROUP_SIZE: u32 = {{ channelReduceWorkgroupSize }}u;
22
 
23
  @compute @workgroup_size({{ channelReduceWorkgroupSize }}, 1, 1)
24
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
25
- // 2D-folded flat index: gid.y carries the high bits past
26
- // the per-axis dispatch fold width and reduces to the 1D form at y=0.
27
- let idx = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WORKGROUP_SIZE;
28
- if (idx >= COUNT) {
29
- return;
30
- }
31
  let batch = idx / PLANE;
32
  let pos = idx % PLANE;
33
 
 
1
+ {% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
2
+ {% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
3
+ // 2D-folded flat index: gid.y carries the high bits past the dispatch's
4
+ // per-axis workgroup fold width.
5
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
6
+ if ({{ name }} >= {{ bound }}) {
7
+ return;
8
+ }{% endmacro %}
9
  // Pointwise convolution with few output channels is evaluated as a per-position
10
  // channel reduction instead of a tiled GEMM. One invocation walks the input
11
  // channels and accumulates every output channel in a statically named local.
 
14
  //
15
  // The output-channel accumulators are named explicitly. Selection caps the
16
  // output-channel count, bounding the invocation-local accumulator storage.
 
 
 
 
 
 
17
  {{ env.wgsl.resourceDeclarations }}
18
 
19
  const IN_C: u32 = {{ inChannels }}u;
 
24
 
25
  @compute @workgroup_size({{ channelReduceWorkgroupSize }}, 1, 1)
26
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
27
+ {{ flat_index_2d("WORKGROUP_SIZE", "idx", "COUNT") }}
 
 
 
 
 
28
  let batch = idx / PLANE;
29
  let pos = idx % PLANE;
30
 
build/webgpu/conv-1x1-gemm-tiled-reg.wgsl.jinja CHANGED
@@ -14,44 +14,13 @@
14
  // instead of 4 * (TM + TN) scalars, so a shared word feeds four times as many
15
  // FMAs and the K loop runs four accumulation steps per iteration.
16
  {{ env.wgsl.resourceDeclarations }}
17
- {% set hasActivation = hasActivation is defined and hasActivation %}
18
- {% set activation = activation | default("") %}
19
- {% set actAlpha = actAlpha | default(0.0) %}
20
- {% set actBeta = actBeta | default(0.0) %}
21
- {% if hasActivation %}
22
- // Apply the fused activation in the f32 accumulator before the single output
23
- // cast, avoiding an intermediate convolution tensor.
24
- {% macro fused_act_return(mode, alpha, beta) -%}
25
- {% if mode == "Relu" %}
26
- return max(v, 0.0);
27
- {% elif mode == "Clip" %}
28
- return clamp(v, f32({{ alpha }}), f32({{ beta }}));
29
- {% elif mode == "LeakyRelu" %}
30
- return select(v * f32({{ alpha }}), v, v >= 0.0);
31
- {% elif mode == "Sigmoid" %}
32
- return 1.0 / (1.0 + exp(-v));
33
- {% elif mode == "Tanh" %}
34
- // tanh is already +/-1 to f32 precision at the clamp bounds. Clamping before
35
- // the builtin preserves that saturated result for larger accumulators.
36
- return tanh(clamp(v, -10.0, 10.0));
37
- {% elif mode == "HardSigmoid" %}
38
- return clamp(f32({{ alpha }}) * v + f32({{ beta }}), 0.0, 1.0);
39
- {% else %}
40
- return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
41
- {% endif %}
42
- {%- endmacro -%}
43
- fn fused_act(v: f32) -> f32 {
44
- {{ fused_act_return(activation, actAlpha, actBeta) -}}
45
- }
46
- {% endif %}
47
-
48
- {% set tileT = "f16" if usesF16 else "f32" -%}
49
- {%- set GEMM_BK = gemmKTile if gemmKTile is defined else 16 %}
50
  {% set GEMM_TM = gemmThreadRows if gemmThreadRows is defined else 8 %}
51
- {% set GEMM_WG_X = gemmWorkgroupX if gemmWorkgroupX is defined else 16 %}
52
  {% set GEMM_BM = gemmMTile if gemmMTile is defined else 64 %}
53
- {% set GEMM_BN = gemmNTile if gemmNTile is defined else GEMM_WG_X * 4 %}
54
- {% set GEMM_TN = gemmThreadColumns if gemmThreadColumns is defined else (GEMM_BN / GEMM_WG_X)|int %}
55
  {% set GEMM_WG_Y = gemmWorkgroupY if gemmWorkgroupY is defined else (GEMM_BM / GEMM_TM)|int %}
56
  {% set components = ["x", "y", "z", "w"] %}
57
  const BK: u32 = {{ GEMM_BK }}u;
@@ -70,11 +39,6 @@ const BN_VECS: u32 = BN / 4u;
70
  // Overlapping windows may reread input values, trading address arithmetic and
71
  // repeated input loads for writing and rereading an expanded column matrix.
72
  {% set implicitIm2col = implicitIm2col is defined and implicitIm2col %}
73
- {% set fusedNarrowProjection = fusedNarrowProjection is defined and fusedNarrowProjection %}
74
- {% set projectionChannels = projectionOutChannels | default(0) %}
75
- {% set projectionInputAct = inputActivation | default("none") %}
76
- {% set projectionOutputAct = outputActivation | default("none") %}
77
- {% set narrowProjectionTile = "tileB" if GEMM_BK >= GEMM_BM else "projectionTile" %}
78
  {% set splitKValue = splitK if splitK is defined else 1 %}
79
  {% set splitKPartial = splitKValue > 1 %}
80
  {% set kTiles = kTiles | default(0) %}
@@ -84,7 +48,7 @@ const BN_VECS: u32 = BN / 4u;
84
  {% set gemmKLimit = "kEnd" if splitKPartial else "K" %}
85
  {% set implicitGatherMode = tunables.IMPLICIT_TILED_GATHER_MODE if tunables.IMPLICIT_TILED_GATHER_MODE is defined else 0 %}
86
  {% set implicitGatherWgSize = GEMM_WG_X * GEMM_WG_Y %}
87
- {% set implicitGatherColumnVecs = (GEMM_BN / 4)|int %}
88
  {% set implicitGatherKChunks = (implicitGatherWgSize / implicitGatherColumnVecs)|int %}
89
  {% set implicitGatherDivides = implicitGatherWgSize % implicitGatherColumnVecs == 0 and implicitGatherKChunks > 0 and GEMM_BK % implicitGatherKChunks == 0 %}
90
  {% set useCarriedImplicitGather = implicitIm2col and implicitGatherMode == 1 and implicitGatherDivides %}
@@ -116,78 +80,21 @@ const PARTIAL_COLS: u32 = {{ nPadded }}u;
116
  const PARTIAL_BATCH_STRIDE: u32 = PARTIAL_ROWS * PARTIAL_COLS;
117
  const PARTIAL_SLICE_STRIDE: u32 = {{ batchCount }}u * PARTIAL_BATCH_STRIDE;
118
  {% endif %}
119
-
120
- {% if fusedNarrowProjection %}
121
- {% set inputActivation = inputActivation | default("none") %}
122
- {% set outputActivation = outputActivation | default("none") %}
123
- {% if inputActivation == "relu" or outputActivation == "relu" %}
124
- fn is_nan_f32(value: f32) -> bool {
125
- let bits = bitcast<u32>(value);
126
- return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
127
- }
128
-
129
- fn projection_relu(value: f32) -> f32 {
130
- var out = max(value, 0.0);
131
- if (is_nan_f32(value)) {
132
- out = value;
133
- }
134
- return out;
135
- }
136
- {% endif %}
137
- {% if outputActivation == "sigmoid" %}
138
- fn sigmoid_safe(x: f32) -> f32 {
139
- if (x >= 0.0) {
140
- let z = exp(-x);
141
- return 1.0 / (1.0 + z);
142
- }
143
- let z = exp(x);
144
- return z / (1.0 + z);
145
- }
146
- {% endif %}
147
-
148
- const PROJECTION_OUT_C: u32 = {{ projectionChannels }}u;
149
- {% else %}{% set emitConvStoreOut = not splitKPartial %}{% set hasZ = hasZ is defined and hasZ %}
150
- {% if emitConvStoreOut | default(true) %}fn store_out(m: u32, n: u32, yBase: u32, raw: f32) {
151
  if (m >= params.M || n >= params.N) {
152
  return;
153
  }
154
  var v = raw;
155
  {% if hasBias %}
156
- v = v + f32(bias[m]);
157
- {% endif %}
158
- {% if hasZ %}
159
- v = v + f32(zResidual[yBase + m * params.N + n]);
160
- {% endif %}
161
- {% if hasActivation %}
162
- v = fused_act(v);
163
  {% endif %}
164
  y[yBase + m * params.N + n] = {{ T }}(v);
165
  }
166
  {% endif %}
167
- {% endif %}
168
 
169
  var<workgroup> tileA: array<array<vec4<{{ tileT }}>, AK_VECS>, BM>;
170
  var<workgroup> tileB: array<array<vec4<{{ tileT }}>, BN_VECS>, BK>;
171
- {% if fusedNarrowProjection and GEMM_BK < GEMM_BM %}
172
- // BK16 cannot reuse the 16-row input tile to publish a BM32 intermediate.
173
- // The extra 32x64 f32 tile keeps the total at 14 KiB, below WebGPU's
174
- // guaranteed 16 KiB workgroup-storage floor.
175
- var<workgroup> projectionTile: array<array<vec4<f32>, BN_VECS>, BM>;
176
- {% endif %}
177
- {% macro publish_projection_input(localRow, localVec, component, globalRow, globalColumn, raw) %}
178
- if ({{ globalRow }} < M && {{ globalColumn }} < N) {
179
- var projectionInput = {{ raw }};
180
- {% if hasBias %}
181
- projectionInput = projectionInput + f32(bias[{{ globalRow }}]);
182
- {% endif %}
183
- {% if projectionInputAct == "relu" %}
184
- projectionInput = projection_relu(projectionInput);
185
- {% endif %}
186
- {{ narrowProjectionTile }}[{{ localRow }}][{{ localVec }}].{{ component }} = projectionInput;
187
- } else {
188
- {{ narrowProjectionTile }}[{{ localRow }}][{{ localVec }}].{{ component }} = 0.0;
189
- }
190
- {%- endmacro %}
191
  {% macro load_a_vec4(rowBase, columnBase, rowLimit, columnLimit) %}
192
  for (var linear = li; linear < BM * AK_VECS; linear += WG_SIZE) {
193
  let ar = linear / AK_VECS;
@@ -208,8 +115,7 @@ var<workgroup> projectionTile: array<array<vec4<f32>, BN_VECS>, BM>;
208
  }
209
  }
210
  tileA[ar][ac4] = av;
211
- }
212
- {%- endmacro %}
213
  {% macro load_b_vec4(rowBase, columnBase, rowLimit, columnLimit) %}
214
  for (var linear = li; linear < BK * BN_VECS; linear += WG_SIZE) {
215
  let br = linear / BN_VECS;
@@ -230,8 +136,7 @@ var<workgroup> projectionTile: array<array<vec4<f32>, BN_VECS>, BM>;
230
  }
231
  }
232
  tileB[br][bc4] = bvec;
233
- }
234
- {%- endmacro %}
235
  {% macro load_b_implicit_vec4(rowBase, columnBase, rowLimit, columnLimit) %}
236
  for (var linear = li; linear < BK * BN_VECS; linear += WG_SIZE) {
237
  let br = linear / BN_VECS;
@@ -259,8 +164,7 @@ var<workgroup> projectionTile: array<array<vec4<f32>, BN_VECS>, BM>;
259
  }
260
  }
261
  tileB[br][bc4] = bvec;
262
- }
263
- {%- endmacro %}
264
  {% macro load_b_implicit_carried(rowBase, columnBase, rowLimit, columnLimit) %}
265
  let bColVec = li / {{ implicitGatherKChunks }}u;
266
  let bChunk = li % {{ implicitGatherKChunks }}u;
@@ -344,8 +248,7 @@ var<workgroup> projectionTile: array<array<vec4<f32>, BN_VECS>, BM>;
344
  }
345
  }
346
  }
347
- }
348
- {%- endmacro %}
349
 
350
  @compute @workgroup_size({{ GEMM_WG_X }}, {{ GEMM_WG_Y }}, 1)
351
  fn main(
@@ -423,62 +326,20 @@ fn main(
423
  let n0 = nBase + lid.x * TN;
424
  let partialBase = slice * PARTIAL_SLICE_STRIDE + batch * PARTIAL_BATCH_STRIDE;
425
  {% else %}
426
- {% if not fusedNarrowProjection %}
427
  let yBatchBase = batch * M * N;
428
- {% endif %}
429
  let m0 = mBase + lid.y * TM;
430
  let n0 = nBase + lid.x * TN;
431
  {% endif %}
432
  {% for row in range(GEMM_TM) %}
 
 
 
 
433
  {% for column in range(GEMM_TN) %}
434
  {% if splitKPartial %}
435
  y[partialBase + (m0 + {{ row }}u) * PARTIAL_COLS + n0 + {{ column }}u] = acc{{ row }}.{{ components[column] }};
436
- {% elif fusedNarrowProjection %}
437
- {{ 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]) }}
438
  {% else %}
439
- store_out(m0 + {{ row }}u, n0 + {{ column }}u, yBatchBase, acc{{ row }}.{{ components[column] }});
440
  {% endif %}
441
  {% endfor %}
442
- {% endfor %}{% if fusedNarrowProjection %}
443
-
444
- // Conv's final K-tile barrier makes tileB dead before it becomes the
445
- // intermediate tile. Every lane publishes its unique micro-tile, then one
446
- // lane per spatial vector word consumes all logical M rows in increasing
447
- // order, four adjacent columns at a time. Invalid M/N-tail cells are
448
- // initialized above, so the barrier is uniform and no lane can observe a
449
- // stale input-tile value.
450
- workgroupBarrier();
451
- if (li < BN_VECS) {
452
- let projectionNBase = nBase + li * 4u;
453
- {% for oc in range(projectionChannels) %}
454
- var projectionAcc{{ oc }} = vec4<f32>(0.0);
455
- {% endfor %}
456
- for (var channel = 0u; channel < M; channel = channel + 1u) {
457
- let projectionValues = vec4<f32>({{ narrowProjectionTile }}[channel][li]);
458
- {% for oc in range(projectionChannels) %}
459
- projectionAcc{{ oc }} = projectionAcc{{ oc }} + projectionValues * f32(projectionW[{{ oc }}u * M + channel]);
460
- {% endfor %}
461
- }
462
-
463
- let projectionYBase = batch * PROJECTION_OUT_C * N;
464
- {% for oc in range(projectionChannels) %}
465
- let projected{{ oc }} = projectionAcc{{ oc }};
466
- {% endfor %}
467
- {% for column in range(4) %}
468
- let projectionN{{ column }} = projectionNBase + {{ column }}u;
469
- if (projectionN{{ column }} < N) {
470
- {% for oc in range(projectionChannels) %}
471
- {% if projectionOutputAct == "relu" %}
472
- let activated{{ oc }}_{{ column }} = projection_relu(projected{{ oc }}.{{ components[column] }});
473
- {% elif projectionOutputAct == "sigmoid" %}
474
- let activated{{ oc }}_{{ column }} = sigmoid_safe(projected{{ oc }}.{{ components[column] }});
475
- {% else %}
476
- let activated{{ oc }}_{{ column }} = projected{{ oc }}.{{ components[column] }};
477
- {% endif %}
478
- y[projectionYBase + {{ oc }}u * N + projectionN{{ column }}] = activated{{ oc }}_{{ column }};
479
- {% endfor %}
480
- }
481
- {% endfor %}
482
- }
483
- {% endif %}
484
- }
 
14
  // instead of 4 * (TM + TN) scalars, so a shared word feeds four times as many
15
  // FMAs and the K loop runs four accumulation steps per iteration.
16
  {{ env.wgsl.resourceDeclarations }}
17
+ {% set tileT = "f16" if usesF16 else "f32" %}
18
+ {% set GEMM_BK = gemmKTile if gemmKTile is defined else 16 %}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
19
  {% set GEMM_TM = gemmThreadRows if gemmThreadRows is defined else 8 %}
20
+ {% set GEMM_WG_X = 16 %}
21
  {% set GEMM_BM = gemmMTile if gemmMTile is defined else 64 %}
22
+ {% set GEMM_BN = 64 %}
23
+ {% set GEMM_TN = 4 %}
24
  {% set GEMM_WG_Y = gemmWorkgroupY if gemmWorkgroupY is defined else (GEMM_BM / GEMM_TM)|int %}
25
  {% set components = ["x", "y", "z", "w"] %}
26
  const BK: u32 = {{ GEMM_BK }}u;
 
39
  // Overlapping windows may reread input values, trading address arithmetic and
40
  // repeated input loads for writing and rereading an expanded column matrix.
41
  {% set implicitIm2col = implicitIm2col is defined and implicitIm2col %}
 
 
 
 
 
42
  {% set splitKValue = splitK if splitK is defined else 1 %}
43
  {% set splitKPartial = splitKValue > 1 %}
44
  {% set kTiles = kTiles | default(0) %}
 
48
  {% set gemmKLimit = "kEnd" if splitKPartial else "K" %}
49
  {% set implicitGatherMode = tunables.IMPLICIT_TILED_GATHER_MODE if tunables.IMPLICIT_TILED_GATHER_MODE is defined else 0 %}
50
  {% set implicitGatherWgSize = GEMM_WG_X * GEMM_WG_Y %}
51
+ {% set implicitGatherColumnVecs = 16 %}
52
  {% set implicitGatherKChunks = (implicitGatherWgSize / implicitGatherColumnVecs)|int %}
53
  {% set implicitGatherDivides = implicitGatherWgSize % implicitGatherColumnVecs == 0 and implicitGatherKChunks > 0 and GEMM_BK % implicitGatherKChunks == 0 %}
54
  {% set useCarriedImplicitGather = implicitIm2col and implicitGatherMode == 1 and implicitGatherDivides %}
 
80
  const PARTIAL_BATCH_STRIDE: u32 = PARTIAL_ROWS * PARTIAL_COLS;
81
  const PARTIAL_SLICE_STRIDE: u32 = {{ batchCount }}u * PARTIAL_BATCH_STRIDE;
82
  {% endif %}
83
+ {% set emitConvStoreOut = not splitKPartial %}
84
+ {% if emitConvStoreOut | default(true) %}fn store_out(m: u32, n: u32, yBase: u32, raw: f32{% if convPreloadBias | default(false) %}, rowBias: f32{% endif %}) {
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
85
  if (m >= params.M || n >= params.N) {
86
  return;
87
  }
88
  var v = raw;
89
  {% if hasBias %}
90
+ v = v + {% if convPreloadBias | default(false) %}rowBias{% else %}f32(bias[m]){% endif %};
 
 
 
 
 
 
91
  {% endif %}
92
  y[yBase + m * params.N + n] = {{ T }}(v);
93
  }
94
  {% endif %}
 
95
 
96
  var<workgroup> tileA: array<array<vec4<{{ tileT }}>, AK_VECS>, BM>;
97
  var<workgroup> tileB: array<array<vec4<{{ tileT }}>, BN_VECS>, BK>;
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
98
  {% macro load_a_vec4(rowBase, columnBase, rowLimit, columnLimit) %}
99
  for (var linear = li; linear < BM * AK_VECS; linear += WG_SIZE) {
100
  let ar = linear / AK_VECS;
 
115
  }
116
  }
117
  tileA[ar][ac4] = av;
118
+ }{% endmacro %}
 
119
  {% macro load_b_vec4(rowBase, columnBase, rowLimit, columnLimit) %}
120
  for (var linear = li; linear < BK * BN_VECS; linear += WG_SIZE) {
121
  let br = linear / BN_VECS;
 
136
  }
137
  }
138
  tileB[br][bc4] = bvec;
139
+ }{% endmacro %}
 
140
  {% macro load_b_implicit_vec4(rowBase, columnBase, rowLimit, columnLimit) %}
141
  for (var linear = li; linear < BK * BN_VECS; linear += WG_SIZE) {
142
  let br = linear / BN_VECS;
 
164
  }
165
  }
166
  tileB[br][bc4] = bvec;
167
+ }{% endmacro %}
 
168
  {% macro load_b_implicit_carried(rowBase, columnBase, rowLimit, columnLimit) %}
169
  let bColVec = li / {{ implicitGatherKChunks }}u;
170
  let bChunk = li % {{ implicitGatherKChunks }}u;
 
248
  }
249
  }
250
  }
251
+ }{% endmacro %}
 
252
 
253
  @compute @workgroup_size({{ GEMM_WG_X }}, {{ GEMM_WG_Y }}, 1)
254
  fn main(
 
326
  let n0 = nBase + lid.x * TN;
327
  let partialBase = slice * PARTIAL_SLICE_STRIDE + batch * PARTIAL_BATCH_STRIDE;
328
  {% else %}
 
329
  let yBatchBase = batch * M * N;
 
330
  let m0 = mBase + lid.y * TM;
331
  let n0 = nBase + lid.x * TN;
332
  {% endif %}
333
  {% for row in range(GEMM_TM) %}
334
+ {% if convPreloadBias | default(false) %}
335
+ var rowBias{{ row }} = 0.0;
336
+ if (m0 + {{ row }}u < M) { rowBias{{ row }} = f32(bias[m0 + {{ row }}u]); }
337
+ {% endif %}
338
  {% for column in range(GEMM_TN) %}
339
  {% if splitKPartial %}
340
  y[partialBase + (m0 + {{ row }}u) * PARTIAL_COLS + n0 + {{ column }}u] = acc{{ row }}.{{ components[column] }};
 
 
341
  {% else %}
342
+ store_out(m0 + {{ row }}u, n0 + {{ column }}u, yBatchBase, acc{{ row }}.{{ components[column] }}{% if convPreloadBias | default(false) %}, rowBias{{ row }}{% endif %});
343
  {% endif %}
344
  {% endfor %}
345
+ {% endfor %}}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
build/webgpu/conv-1x1-gemm-tiled.wgsl.jinja CHANGED
@@ -6,67 +6,31 @@
6
  // f16 operands remain packed in workgroup memory and widen when consumed. All
7
  // M, N, and K accesses are bounds-checked.
8
  {{ env.wgsl.resourceDeclarations }}
9
- {% set hasActivation = hasActivation is defined and hasActivation %}
10
- {% set activation = activation | default("") %}
11
- {% set actAlpha = actAlpha | default(0.0) %}
12
- {% set actBeta = actBeta | default(0.0) %}
13
- {% if hasActivation %}
14
- // Apply the fused activation in the f32 accumulator before the single output
15
- // cast, avoiding an intermediate convolution tensor.
16
- {% macro fused_act_return(mode, alpha, beta) -%}
17
- {% if mode == "Relu" %}
18
- return max(v, 0.0);
19
- {% elif mode == "Clip" %}
20
- return clamp(v, f32({{ alpha }}), f32({{ beta }}));
21
- {% elif mode == "LeakyRelu" %}
22
- return select(v * f32({{ alpha }}), v, v >= 0.0);
23
- {% elif mode == "Sigmoid" %}
24
- return 1.0 / (1.0 + exp(-v));
25
- {% elif mode == "Tanh" %}
26
- // tanh is already +/-1 to f32 precision at the clamp bounds. Clamping before
27
- // the builtin preserves that saturated result for larger accumulators.
28
- return tanh(clamp(v, -10.0, 10.0));
29
- {% elif mode == "HardSigmoid" %}
30
- return clamp(f32({{ alpha }}) * v + f32({{ beta }}), 0.0, 1.0);
31
- {% else %}
32
- return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
33
- {% endif %}
34
- {%- endmacro -%}
35
- fn fused_act(v: f32) -> f32 {
36
- {{ fused_act_return(activation, actAlpha, actBeta) -}}
37
- }
38
- {% endif %}
39
-
40
- {% set tileT = "f16" if usesF16 else "f32" -%}
41
  const BK: u32 = 16u;
42
  const BM: u32 = 32u;
43
  const BN: u32 = 32u;
44
 
45
  // Store one output element with the configured bias/residual/activation
46
  // epilogue in the f32 accumulator domain.
47
- {% set hasZ = hasZ is defined and hasZ %}
48
- fn store_out(m: u32, n: u32, yBase: u32, raw: f32) {
49
  if (m >= params.M || n >= params.N) {
50
  return;
51
  }
52
  var v = raw;
53
  {% if hasBias %}
54
- v = v + f32(bias[m]);
55
- {% endif %}
56
- {% if hasZ %}
57
- v = v + f32(zResidual[yBase + m * params.N + n]);
58
- {% endif %}
59
- {% if hasActivation %}
60
- v = fused_act(v);
61
  {% endif %}
62
  y[yBase + m * params.N + n] = {{ T }}(v);
63
  }
64
 
65
-
66
  var<workgroup> tileA: array<array<{{ tileT }}, BK>, BM>;
67
  var<workgroup> tileB: array<array<{{ tileT }}, BN>, BK>;
68
 
69
- @compute @workgroup_size(16, 16, 1)
70
  fn main(
71
  @builtin(workgroup_id) wg: vec3<u32>,
72
  @builtin(local_invocation_id) lid: vec3<u32>
@@ -78,7 +42,7 @@ fn main(
78
  let nBase = wg.x * BN;
79
  let batch = wg.z;
80
  let xBatchBase = batch * K * N;
81
- let li = lid.y * 16u + lid.x;
82
 
83
  var acc00: f32 = 0.0;
84
  var acc01: f32 = 0.0;
@@ -88,8 +52,8 @@ fn main(
88
  for (var kt: u32 = 0u; kt < numTiles; kt = kt + 1u) {
89
  let kBase = kt * BK;
90
  // Cooperative load: 32x16 A(=W) tile + 16x32 B(=X[batch]) tile, 256 threads x 2 each.
91
- for (var e: u32 = 0u; e < (BM * BK) / 256u; e = e + 1u) {
92
- let idx = li + e * 256u;
93
  let ar = idx / BK;
94
  let ac = idx % BK;
95
  let am = mBase + ar;
@@ -115,10 +79,9 @@ fn main(
115
  let a1 = f32(tileA[lid.y * 2u + 1u][kk]);
116
  let b0 = f32(tileB[kk][lid.x * 2u]);
117
  let b1 = f32(tileB[kk][lid.x * 2u + 1u]);
118
- acc00 = acc00 + a0 * b0;
119
- acc01 = acc01 + a0 * b1;
120
- acc10 = acc10 + a1 * b0;
121
- acc11 = acc11 + a1 * b1;
122
  }
123
  workgroupBarrier();
124
  }
 
6
  // f16 operands remain packed in workgroup memory and widen when consumed. All
7
  // M, N, and K accesses are bounds-checked.
8
  {{ env.wgsl.resourceDeclarations }}
9
+ {% set accumulatorPrefix = "acc" %}
10
+ {% set tileT = "f16" if usesF16 else "f32" %}
11
+ {% set wgX = 16 %}
12
+ {% set wgY = 16 %}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
13
  const BK: u32 = 16u;
14
  const BM: u32 = 32u;
15
  const BN: u32 = 32u;
16
 
17
  // Store one output element with the configured bias/residual/activation
18
  // epilogue in the f32 accumulator domain.
19
+ fn store_out(m: u32, n: u32, yBase: u32, raw: f32{% if convPreloadBias | default(false) %}, rowBias: f32{% endif %}) {
 
20
  if (m >= params.M || n >= params.N) {
21
  return;
22
  }
23
  var v = raw;
24
  {% if hasBias %}
25
+ v = v + {% if convPreloadBias | default(false) %}rowBias{% else %}f32(bias[m]){% endif %};
 
 
 
 
 
 
26
  {% endif %}
27
  y[yBase + m * params.N + n] = {{ T }}(v);
28
  }
29
 
 
30
  var<workgroup> tileA: array<array<{{ tileT }}, BK>, BM>;
31
  var<workgroup> tileB: array<array<{{ tileT }}, BN>, BK>;
32
 
33
+ @compute @workgroup_size({{ wgX }}, {{ wgY }}, 1)
34
  fn main(
35
  @builtin(workgroup_id) wg: vec3<u32>,
36
  @builtin(local_invocation_id) lid: vec3<u32>
 
42
  let nBase = wg.x * BN;
43
  let batch = wg.z;
44
  let xBatchBase = batch * K * N;
45
+ let li = lid.y * {{ wgX }}u + lid.x;
46
 
47
  var acc00: f32 = 0.0;
48
  var acc01: f32 = 0.0;
 
52
  for (var kt: u32 = 0u; kt < numTiles; kt = kt + 1u) {
53
  let kBase = kt * BK;
54
  // Cooperative load: 32x16 A(=W) tile + 16x32 B(=X[batch]) tile, 256 threads x 2 each.
55
+ for (var e: u32 = 0u; e < (BM * BK) / {{ wgX * wgY }}u; e = e + 1u) {
56
+ let idx = li + e * {{ wgX * wgY }}u;
57
  let ar = idx / BK;
58
  let ac = idx % BK;
59
  let am = mBase + ar;
 
79
  let a1 = f32(tileA[lid.y * 2u + 1u][kk]);
80
  let b0 = f32(tileB[kk][lid.x * 2u]);
81
  let b1 = f32(tileB[kk][lid.x * 2u + 1u]);
82
+ {% for i in range(2) %}{% for j in range(2) %}
83
+ {{ accumulatorPrefix }}{{ i }}{{ j }} = {{ accumulatorPrefix }}{{ i }}{{ j }} + a{{ i }} * b{{ j }};
84
+ {% endfor %}{% endfor %}
 
85
  }
86
  workgroupBarrier();
87
  }
build/webgpu/conv-1x1-subgroup-matrix.wgsl.jinja CHANGED
@@ -12,18 +12,16 @@
12
  // the slices and applies bias or an epilogue. Padded tail cells remain zero and
13
  // are never copied to the logical output.
14
  {% set directMatrixInputs = directMatrixInputs is defined and directMatrixInputs %}
15
- {% set fusedNarrowProjection = fusedNarrowProjection if fusedNarrowProjection is defined else false %}
16
- {% set polyphase = polyphaseConvTranspose is defined and polyphaseConvTranspose %}
17
  {% set implicit = implicitIm2col is defined and implicitIm2col %}
18
  {% set implicit3d = implicit and (convSpatialDims is defined and convSpatialDims == 3) %}
19
  {% set splitKValue = splitK if splitK is defined else 1 %}
20
  {% set splitKPartial = splitKValue > 1 %}
21
  {% set kLoopVar = "K_LOOP" if padded else "K" %}
22
  {% set nColsVar = "N_COLS" if padded else "N" %}
23
- {% set hasZ = hasZ is defined and hasZ %}
24
- {% set projectionOutChannels = projectionOutChannels | default(0) %}
25
- {% set hasProjectionBias = hasProjectionBias is defined and hasProjectionBias %}
26
- {% set hasOutputScale = hasOutputScale is defined and hasOutputScale %}
27
  {% set convKernelH = convKernelH | default(0) %}
28
  {% set convKernelW = convKernelW | default(0) %}
29
  {% set convStrideH = convStrideH | default(0) %}
@@ -53,46 +51,16 @@ enable subgroup_size_control;
53
  enable chromium_experimental_subgroup_matrix;
54
  diagnostic(off, chromium.subgroup_matrix_uniformity);
55
 
56
-
57
  {{ env.wgsl.resourceDeclarations }}
58
- {% if fusedNarrowProjection %}{% set inputActivation = inputActivation | default("none") %}
59
- {% set outputActivation = outputActivation | default("none") %}
60
- {% if inputActivation == "relu" or outputActivation == "relu" %}
61
- fn is_nan_f32(value: f32) -> bool {
62
- let bits = bitcast<u32>(value);
63
- return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
64
- }
65
-
66
- fn projection_relu(value: f32) -> f32 {
67
- var out = max(value, 0.0);
68
- if (is_nan_f32(value)) {
69
- out = value;
70
- }
71
- return out;
72
- }
73
- {% endif %}
74
- {% if outputActivation == "sigmoid" %}
75
- fn sigmoid_safe(x: f32) -> f32 {
76
- if (x >= 0.0) {
77
- let z = exp(-x);
78
- return 1.0 / (1.0 + z);
79
- }
80
- let z = exp(x);
81
- return z / (1.0 + z);
82
- }
83
- {% endif %}
84
- {% endif %}
85
-
86
  {% set operandScalar = fScalar %}
87
  {% set accScalar = "f32" %}
88
  {% set useDirectMatrixStore = directMatrixStore is defined and directMatrixStore %}
89
  {% set tileRowsValue = tileRows if tileRows is defined else 32 %}
90
- {% set tileColsValue = tileCols if tileCols is defined else 64 %}
91
  {% set workgroupThreadsValue = workgroupThreads if workgroupThreads is defined else 128 %}
92
  {% set subgroupRowsValue = subgroupRows if subgroupRows is defined else 2 %}
93
- {% set subgroupColsValue = subgroupCols if subgroupCols is defined else 2 %}
94
  {% set subRowsValue = (tileRowsValue / subgroupRowsValue)|int %}
95
- {% set subColsValue = (tileColsValue / subgroupColsValue)|int %}
96
  {% set bLoadWidth = (tileColsValue * 32 / workgroupThreadsValue)|int %}
97
  {% set bKChunks = (32 / bLoadWidth)|int %}
98
 
@@ -176,7 +144,7 @@ fn loadSHMA(tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
176
  for (var col_offset = 0u; col_offset < 8u; col_offset = col_offset + 1u) {
177
  let k = k_idx + col + col_offset;
178
  if (a_global < M{% if padded %} && k < K{% endif %}) {
179
- {% set A_PHASE = "phase_base + " if polyphase else "" %}
180
  {% if operandScalar == "f16" %}
181
  tile_A[row * TILE_K + col + col_offset] = f16(w[{{ A_PHASE }}a_global * K + k]);
182
  {% else %}
@@ -325,59 +293,13 @@ fn loadSHMB(b_base: u32, tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
325
  }
326
  {% endif %}
327
  }
328
-
329
- {% endif %}
330
- {% set hasActivation = hasActivation is defined and hasActivation %}
331
- {% set activation = activation | default("") %}
332
- {% set actAlpha = actAlpha | default(0.0) %}
333
- {% set actBeta = actBeta | default(0.0) %}
334
- {% if hasActivation %}
335
- // Apply the fused activation in the f32 accumulator before the single output
336
- // cast, avoiding an intermediate convolution tensor.
337
- {% macro fused_act_return(mode, alpha, beta) -%}
338
- {% if mode == "Relu" %}
339
- return max(v, 0.0);
340
- {% elif mode == "Clip" %}
341
- return clamp(v, f32({{ alpha }}), f32({{ beta }}));
342
- {% elif mode == "LeakyRelu" %}
343
- return select(v * f32({{ alpha }}), v, v >= 0.0);
344
- {% elif mode == "Sigmoid" %}
345
- return 1.0 / (1.0 + exp(-v));
346
- {% elif mode == "Tanh" %}
347
- // tanh is already +/-1 to f32 precision at the clamp bounds. Clamping before
348
- // the builtin preserves that saturated result for larger accumulators.
349
- return tanh(clamp(v, -10.0, 10.0));
350
- {% elif mode == "HardSigmoid" %}
351
- return clamp(f32({{ alpha }}) * v + f32({{ beta }}), 0.0, 1.0);
352
- {% else %}
353
- return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
354
- {% endif %}
355
- {%- endmacro -%}
356
- fn fused_act(v: f32) -> f32 {
357
- {{ fused_act_return(activation, actAlpha, actBeta) -}}
358
- }
359
  {% endif %}
360
-
361
  {% set biasAdd = " + bv" if hasBias else "" %}
362
- {% if hasActivation or hasZ %}
363
- // Fused epilogue: Y = activation(conv + bias + Z), applied at the output store.
364
- // `raw` is the f32 convolution-plus-bias accumulator; Z shares Y's NCHW layout.
365
- fn epi(raw: f32{% if hasZ %}, yIndex: u32{% endif %}) -> {{ T }} {
366
- var r = raw;
367
- {% if hasZ %}
368
- r = r + f32(zResidual[yIndex]);
369
- {% endif %}
370
- {% if hasActivation %}
371
- r = fused_act(r);
372
- {% endif %}
373
- return {{ T }}(r);
374
- }
375
- {% endif %}
376
  {% macro store_val(valExpr, idxExpr) %}
377
- {% if hasActivation or hasZ %}epi({{ valExpr }}{% if hasZ %}, {{ idxExpr }}{% endif %})
378
- {%- else %}{{ T }}({{ valExpr }})
379
- {%- endif %}
380
- {% endmacro %}
381
 
382
  {% if (not useDirectMatrixStore or padded) and not splitKPartial and not fusedNarrowProjection %}
383
  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 +353,7 @@ fn main(
431
  {% else %}
432
  let batch = workgroup_id.z;
433
  let b_base = batch * B_BATCH_STRIDE;
434
- {% if not fusedNarrowProjection %}
435
  let c_base = batch * C_BATCH_STRIDE;
436
- {% endif %}
437
  {% endif %}
438
  let a_global_base = workgroup_id.y * TILE_ROWS;
439
  let b_global_base = workgroup_id.x * TILE_COLS;
@@ -491,57 +411,7 @@ fn main(
491
  {% endif %}
492
  }
493
 
494
- {% if fusedNarrowProjection %}
495
- // The Conv producer owns exactly one 32-row tile. Its input tile is dead after
496
- // the final K iteration, so reuse those 32x64 f32 cells as the only
497
- // cross-subgroup handoff into the narrow projection. All collective stores
498
- // remain subgroup-uniform, and every fragment has a disjoint destination.
499
- let fused_tile_offset = base_A * TILE_COLS + base_B;
500
- subgroupMatrixStore<row_major>(&tile_B, fused_tile_offset + 0u * TILE_COLS + 0u, matC00, TILE_COLS);
501
- subgroupMatrixStore<row_major>(&tile_B, fused_tile_offset + 0u * TILE_COLS + 8u, matC01, TILE_COLS);
502
- subgroupMatrixStore<row_major>(&tile_B, fused_tile_offset + 0u * TILE_COLS + 16u, matC02, TILE_COLS);
503
- subgroupMatrixStore<row_major>(&tile_B, fused_tile_offset + 0u * TILE_COLS + 24u, matC03, TILE_COLS);
504
- subgroupMatrixStore<row_major>(&tile_B, fused_tile_offset + 8u * TILE_COLS + 0u, matC10, TILE_COLS);
505
- subgroupMatrixStore<row_major>(&tile_B, fused_tile_offset + 8u * TILE_COLS + 8u, matC11, TILE_COLS);
506
- subgroupMatrixStore<row_major>(&tile_B, fused_tile_offset + 8u * TILE_COLS + 16u, matC12, TILE_COLS);
507
- subgroupMatrixStore<row_major>(&tile_B, fused_tile_offset + 8u * TILE_COLS + 24u, matC13, TILE_COLS);
508
- workgroupBarrier();
509
-
510
- // One invocation owns one spatial column and walks the producer channels in
511
- // increasing order to preserve the scalar projection's accumulation order.
512
- // The inactive half of the workgroup has no remaining collective operation
513
- // to reach.
514
- let fused_global_col = b_global_base + local_idx;
515
- if (local_idx < TILE_COLS && fused_global_col < N) {
516
- {% for oc in range(projectionOutChannels) %}
517
- var projected_acc{{ oc }} = f32(0.0);
518
- {% endfor %}
519
- for (var channel = 0u; channel < M; channel = channel + 1u) {
520
- let producer_value = tile_B[channel * TILE_COLS + local_idx]{% if hasBias %} + f32(bias[channel]){% endif %};
521
- {% if inputActivation == "relu" %}
522
- let activated_value = projection_relu(producer_value);
523
- {% else %}
524
- let activated_value = producer_value;
525
- {% endif %}
526
- {% for oc in range(projectionOutChannels) %}
527
- projected_acc{{ oc }} = projected_acc{{ oc }} + activated_value * f32(projectionW[{{ oc }}u * M + channel]);
528
- {% endfor %}
529
- }
530
-
531
- let fused_y_base = batch * {{ projectionOutChannels }}u * N + fused_global_col;
532
- {% for oc in range(projectionOutChannels) %}
533
- let projected{{ oc }} = projected_acc{{ oc }}{% if hasProjectionBias %} + f32(projectionBias[{{ oc }}u]){% endif %};
534
- {% if outputActivation == "relu" %}
535
- let activated{{ oc }} = projection_relu(projected{{ oc }});
536
- {% elif outputActivation == "sigmoid" %}
537
- let activated{{ oc }} = sigmoid_safe(projected{{ oc }});
538
- {% else %}
539
- let activated{{ oc }} = projected{{ oc }};
540
- {% endif %}
541
- y[fused_y_base + {{ oc }}u * N] = activated{{ oc }}{% if hasOutputScale %} * params.outputScale{% endif %};
542
- {% endfor %}
543
- }
544
- {% elif splitKPartial %}
545
  // Raw partials, no bias and no epilogue — the reduce pass owns both. The
546
  // scratch is whole tiles in both axes, so this needs none of the column or
547
  // row guards the real output store carries.
@@ -586,11 +456,7 @@ fn main(
586
  {% if hasBias %}
587
  raw += f32(bias[global_row]);
588
  {% endif %}
589
- {% if hasActivation or hasZ %}
590
- y[y_index] = epi(raw{% if hasZ %}, y_index{% endif %});
591
- {% else %}
592
  y[y_index] = {{ T }}(raw);
593
- {% endif %}
594
  }
595
  {% endif %}
596
  {% if padded %}
 
12
  // the slices and applies bias or an epilogue. Padded tail cells remain zero and
13
  // are never copied to the logical output.
14
  {% set directMatrixInputs = directMatrixInputs is defined and directMatrixInputs %}
15
+ {% set fusedNarrowProjection = false %}
16
+ {% set polyphase = false %}
17
  {% set implicit = implicitIm2col is defined and implicitIm2col %}
18
  {% set implicit3d = implicit and (convSpatialDims is defined and convSpatialDims == 3) %}
19
  {% set splitKValue = splitK if splitK is defined else 1 %}
20
  {% set splitKPartial = splitKValue > 1 %}
21
  {% set kLoopVar = "K_LOOP" if padded else "K" %}
22
  {% set nColsVar = "N_COLS" if padded else "N" %}
23
+ {% set hasZ = false %}
24
+ {% set hasActivation = false %}
 
 
25
  {% set convKernelH = convKernelH | default(0) %}
26
  {% set convKernelW = convKernelW | default(0) %}
27
  {% set convStrideH = convStrideH | default(0) %}
 
51
  enable chromium_experimental_subgroup_matrix;
52
  diagnostic(off, chromium.subgroup_matrix_uniformity);
53
 
 
54
  {{ env.wgsl.resourceDeclarations }}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
55
  {% set operandScalar = fScalar %}
56
  {% set accScalar = "f32" %}
57
  {% set useDirectMatrixStore = directMatrixStore is defined and directMatrixStore %}
58
  {% set tileRowsValue = tileRows if tileRows is defined else 32 %}
59
+ {% set tileColsValue = 64 %}
60
  {% set workgroupThreadsValue = workgroupThreads if workgroupThreads is defined else 128 %}
61
  {% set subgroupRowsValue = subgroupRows if subgroupRows is defined else 2 %}
 
62
  {% set subRowsValue = (tileRowsValue / subgroupRowsValue)|int %}
63
+ {% set subColsValue = 32 %}
64
  {% set bLoadWidth = (tileColsValue * 32 / workgroupThreadsValue)|int %}
65
  {% set bKChunks = (32 / bLoadWidth)|int %}
66
 
 
144
  for (var col_offset = 0u; col_offset < 8u; col_offset = col_offset + 1u) {
145
  let k = k_idx + col + col_offset;
146
  if (a_global < M{% if padded %} && k < K{% endif %}) {
147
+ {% set A_PHASE = "" %}
148
  {% if operandScalar == "f16" %}
149
  tile_A[row * TILE_K + col + col_offset] = f16(w[{{ A_PHASE }}a_global * K + k]);
150
  {% else %}
 
293
  }
294
  {% endif %}
295
  }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
296
  {% endif %}
297
+ {% set hasActivation = false %}
298
  {% set biasAdd = " + bv" if hasBias else "" %}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
299
  {% macro store_val(valExpr, idxExpr) %}
300
+ {% if hasActivation or hasZ %}
301
+ epi({{ valExpr }}){% else %}
302
+ {{ T }}({{ valExpr }}){% endif %}{% endmacro %}
 
303
 
304
  {% if (not useDirectMatrixStore or padded) and not splitKPartial and not fusedNarrowProjection %}
305
  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 %}) {
 
353
  {% else %}
354
  let batch = workgroup_id.z;
355
  let b_base = batch * B_BATCH_STRIDE;
 
356
  let c_base = batch * C_BATCH_STRIDE;
 
357
  {% endif %}
358
  let a_global_base = workgroup_id.y * TILE_ROWS;
359
  let b_global_base = workgroup_id.x * TILE_COLS;
 
411
  {% endif %}
412
  }
413
 
414
+ {% if splitKPartial %}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
415
  // Raw partials, no bias and no epilogue — the reduce pass owns both. The
416
  // scratch is whole tiles in both axes, so this needs none of the column or
417
  // row guards the real output store carries.
 
456
  {% if hasBias %}
457
  raw += f32(bias[global_row]);
458
  {% endif %}
 
 
 
459
  y[y_index] = {{ T }}(raw);
 
460
  }
461
  {% endif %}
462
  {% if padded %}
build/webgpu/conv-direct-nd.wgsl.jinja CHANGED
@@ -1,50 +1,20 @@
1
  // Direct N-dimensional convolution for channels-first tensors. Inputs are
2
  // accumulated in f32 and narrowed once at the output.
3
  // A fused epilogue may apply a residual input and activation before the store.
4
- {% if usesF16Spec %}
5
- enable f16;
6
- {% endif %}
 
 
 
 
 
7
  {{ env.wgsl.resourceDeclarations }}
8
- {% set hasActivation = hasActivation is defined and hasActivation %}
9
- {% set activation = activation | default("") %}
10
- {% set actAlpha = actAlpha | default(0.0) %}
11
- {% set actBeta = actBeta | default(0.0) %}
12
- {% if hasActivation %}
13
- // Apply the fused activation in the f32 accumulator before the single output
14
- // cast, avoiding an intermediate convolution tensor.
15
- {% macro fused_act_return(mode, alpha, beta) -%}
16
- {% if mode == "Relu" %}
17
- return max(v, 0.0);
18
- {% elif mode == "Clip" %}
19
- return clamp(v, f32({{ alpha }}), f32({{ beta }}));
20
- {% elif mode == "LeakyRelu" %}
21
- return select(v * f32({{ alpha }}), v, v >= 0.0);
22
- {% elif mode == "Sigmoid" %}
23
- return 1.0 / (1.0 + exp(-v));
24
- {% elif mode == "Tanh" %}
25
- // tanh is already +/-1 to f32 precision at the clamp bounds. Clamping before
26
- // the builtin preserves that saturated result for larger accumulators.
27
- return tanh(clamp(v, -10.0, 10.0));
28
- {% elif mode == "HardSigmoid" %}
29
- return clamp(f32({{ alpha }}) * v + f32({{ beta }}), 0.0, 1.0);
30
- {% else %}
31
- return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
32
- {% endif %}
33
- {%- endmacro -%}
34
- fn fused_act(v: f32) -> f32 {
35
- {{ fused_act_return(activation, actAlpha, actBeta) -}}
36
- }
37
- {% endif %}
38
-
39
- {% set hasZ = hasZ is defined and hasZ %}
40
  const WG: u32 = {{ convWorkgroupSize }}u;
41
 
42
  @compute @workgroup_size(WG)
43
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
44
- let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
45
- if (index >= params.count) {
46
- return;
47
- }
48
 
49
  let ow = index % params.outW;
50
  var q = index / params.outW;
@@ -108,9 +78,5 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
108
  }
109
  {% endif %}
110
  }
111
- {% if hasZ %}
112
- // Optional residual epilogue: Y = activation(conv + B + Z); Z shares Y's layout.
113
- acc = acc + f32(zResidual[index]);
114
- {% endif %}
115
- y[index] = {{ T }}({% if hasActivation %}fused_act(acc){% else %}acc{% endif %});
116
  }
 
1
  // Direct N-dimensional convolution for channels-first tensors. Inputs are
2
  // accumulated in f32 and narrowed once at the output.
3
  // A fused epilogue may apply a residual input and activation before the store.
4
+ {% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
5
+ {% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
6
+ // 2D-folded flat index: gid.y carries the high bits past the dispatch's
7
+ // per-axis workgroup fold width.
8
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
9
+ if ({{ name }} >= {{ bound }}) {
10
+ return;
11
+ }{% endmacro %}
12
  {{ env.wgsl.resourceDeclarations }}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
13
  const WG: u32 = {{ convWorkgroupSize }}u;
14
 
15
  @compute @workgroup_size(WG)
16
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
17
+ {{ flat_index_2d("WG", "index") }}
 
 
 
18
 
19
  let ow = index % params.outW;
20
  var q = index / params.outW;
 
78
  }
79
  {% endif %}
80
  }
81
+ y[index] = {{ T }}(acc);
 
 
 
 
82
  }
build/webgpu/conv-direct-unrolled.wgsl.jinja CHANGED
@@ -6,39 +6,15 @@
6
  // strides fold to literals; the per-row bounds check is hoisted per kh.
7
  // Group offsets are applied to both input channels and weights. f16 inputs are
8
  // widened to an f32 accumulator and narrowed once at store.
 
 
 
 
 
 
 
 
9
  {{ env.wgsl.resourceDeclarations }}
10
- {% set hasActivation = hasActivation is defined and hasActivation %}
11
- {% set activation = activation | default("") %}
12
- {% set actAlpha = actAlpha | default(0.0) %}
13
- {% set actBeta = actBeta | default(0.0) %}
14
- {% if hasActivation %}
15
- // Apply the fused activation in the f32 accumulator before the single output
16
- // cast, avoiding an intermediate convolution tensor.
17
- {% macro fused_act_return(mode, alpha, beta) -%}
18
- {% if mode == "Relu" %}
19
- return max(v, 0.0);
20
- {% elif mode == "Clip" %}
21
- return clamp(v, f32({{ alpha }}), f32({{ beta }}));
22
- {% elif mode == "LeakyRelu" %}
23
- return select(v * f32({{ alpha }}), v, v >= 0.0);
24
- {% elif mode == "Sigmoid" %}
25
- return 1.0 / (1.0 + exp(-v));
26
- {% elif mode == "Tanh" %}
27
- // tanh is already +/-1 to f32 precision at the clamp bounds. Clamping before
28
- // the builtin preserves that saturated result for larger accumulators.
29
- return tanh(clamp(v, -10.0, 10.0));
30
- {% elif mode == "HardSigmoid" %}
31
- return clamp(f32({{ alpha }}) * v + f32({{ beta }}), 0.0, 1.0);
32
- {% else %}
33
- return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
34
- {% endif %}
35
- {%- endmacro -%}
36
- fn fused_act(v: f32) -> f32 {
37
- {{ fused_act_return(activation, actAlpha, actBeta) -}}
38
- }
39
- {% endif %}
40
-
41
- {% set hasZ = hasZ is defined and hasZ %}
42
  const KERNEL_AREA: u32 = {{ kernelHSpec * kernelWSpec }}u;
43
  const STRIDE_H: u32 = {{ strideHSpec }}u;
44
  const STRIDE_W: u32 = {{ strideWSpec }}u;
@@ -47,12 +23,7 @@ const PAD_LEFT: i32 = {{ padLeftSpec }};
47
 
48
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
49
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
50
- // 2D-folded flat index: gid.y carries the high bits past the
51
- // per-axis dispatch fold width (outputs > 16.7M elements).
52
- let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
53
- if (index >= params.count) {
54
- return;
55
- }
56
 
57
  let ow = index % params.outW;
58
  var t = index / params.outW;
@@ -92,9 +63,5 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
92
  xBase = xBase + xChannelStride;
93
  wBase = wBase + KERNEL_AREA;
94
  }
95
- {% if hasZ %}
96
- // Optional residual epilogue: Y = activation(conv + B + Z); Z shares Y's layout.
97
- acc = acc + f32(zResidual[index]);
98
- {% endif %}
99
- y[index] = {{ T }}({% if hasActivation %}fused_act(acc){% else %}acc{% endif %});
100
  }
 
6
  // strides fold to literals; the per-row bounds check is hoisted per kh.
7
  // Group offsets are applied to both input channels and weights. f16 inputs are
8
  // widened to an f32 accumulator and narrowed once at store.
9
+ {% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
10
+ {% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
11
+ // 2D-folded flat index: gid.y carries the high bits past the dispatch's
12
+ // per-axis workgroup fold width.
13
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
14
+ if ({{ name }} >= {{ bound }}) {
15
+ return;
16
+ }{% endmacro %}
17
  {{ env.wgsl.resourceDeclarations }}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
18
  const KERNEL_AREA: u32 = {{ kernelHSpec * kernelWSpec }}u;
19
  const STRIDE_H: u32 = {{ strideHSpec }}u;
20
  const STRIDE_W: u32 = {{ strideWSpec }}u;
 
23
 
24
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
25
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
26
+ {{ flat_index_2d(tunables.WORKGROUP_SIZE, "index") }}
 
 
 
 
 
27
 
28
  let ow = index % params.outW;
29
  var t = index / params.outW;
 
63
  xBase = xBase + xChannelStride;
64
  wBase = wBase + KERNEL_AREA;
65
  }
66
+ y[index] = {{ T }}(acc);
 
 
 
 
67
  }
build/webgpu/conv-splitk-reduce.wgsl.jinja CHANGED
@@ -10,41 +10,15 @@
10
  //
11
  // Summing slices in index order changes the f32 association relative to one
12
  // uninterrupted K loop, so the two orders need not be bit-identical.
 
 
 
 
 
 
 
 
13
  {{ env.wgsl.resourceDeclarations }}
14
- {% set applyActivation = hasActivation is defined and hasActivation %}
15
- {% if applyActivation %}
16
- {% set hasActivation = hasActivation is defined and hasActivation %}
17
- {% set activation = activation | default("") %}
18
- {% set actAlpha = actAlpha | default(0.0) %}
19
- {% set actBeta = actBeta | default(0.0) %}
20
- {% if hasActivation %}
21
- // Apply the fused activation in the f32 accumulator before the single output
22
- // cast, avoiding an intermediate convolution tensor.
23
- {% macro fused_act_return(mode, alpha, beta) -%}
24
- {% if mode == "Relu" %}
25
- return max(v, 0.0);
26
- {% elif mode == "Clip" %}
27
- return clamp(v, f32({{ alpha }}), f32({{ beta }}));
28
- {% elif mode == "LeakyRelu" %}
29
- return select(v * f32({{ alpha }}), v, v >= 0.0);
30
- {% elif mode == "Sigmoid" %}
31
- return 1.0 / (1.0 + exp(-v));
32
- {% elif mode == "Tanh" %}
33
- // tanh is already +/-1 to f32 precision at the clamp bounds. Clamping before
34
- // the builtin preserves that saturated result for larger accumulators.
35
- return tanh(clamp(v, -10.0, 10.0));
36
- {% elif mode == "HardSigmoid" %}
37
- return clamp(f32({{ alpha }}) * v + f32({{ beta }}), 0.0, 1.0);
38
- {% else %}
39
- return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
40
- {% endif %}
41
- {%- endmacro -%}
42
- fn fused_act(v: f32) -> f32 {
43
- {{ fused_act_return(activation, actAlpha, actBeta) -}}
44
- }
45
- {% endif %}
46
-
47
- {% endif %}
48
 
49
  const M: u32 = {{ M }}u;
50
  const N: u32 = {{ N }}u;
@@ -57,12 +31,7 @@ const WORKGROUP_SIZE: u32 = {{ reduceWorkgroupSize }}u;
57
 
58
  @compute @workgroup_size({{ reduceWorkgroupSize }}, 1, 1)
59
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
60
- // 2D-folded flat index: gid.y carries the high bits past
61
- // the per-axis dispatch fold width and reduces to the 1D form at y=0.
62
- let idx = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WORKGROUP_SIZE;
63
- if (idx >= COUNT) {
64
- return;
65
- }
66
  let col = idx % N;
67
  let row = (idx / N) % M;
68
  let batch = idx / (M * N);
@@ -75,9 +44,6 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
75
  }
76
  {% if hasBias %}
77
  acc = acc + f32(bias[row]);
78
- {% endif %}
79
- {% if applyActivation %}
80
- acc = fused_act(acc);
81
  {% endif %}
82
  y[idx] = {{ T }}(acc);
83
  }
 
10
  //
11
  // Summing slices in index order changes the f32 association relative to one
12
  // uninterrupted K loop, so the two orders need not be bit-identical.
13
+ {% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
14
+ {% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
15
+ // 2D-folded flat index: gid.y carries the high bits past the dispatch's
16
+ // per-axis workgroup fold width.
17
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
18
+ if ({{ name }} >= {{ bound }}) {
19
+ return;
20
+ }{% endmacro %}
21
  {{ env.wgsl.resourceDeclarations }}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
22
 
23
  const M: u32 = {{ M }}u;
24
  const N: u32 = {{ N }}u;
 
31
 
32
  @compute @workgroup_size({{ reduceWorkgroupSize }}, 1, 1)
33
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
34
+ {{ flat_index_2d("WORKGROUP_SIZE", "idx", "COUNT") }}
 
 
 
 
 
35
  let col = idx % N;
36
  let row = (idx / N) % M;
37
  let batch = idx / (M * N);
 
44
  }
45
  {% if hasBias %}
46
  acc = acc + f32(bias[row]);
 
 
 
47
  {% endif %}
48
  y[idx] = {{ T }}(acc);
49
  }
build/webgpu/conv1d-tiled-reg.wgsl.jinja CHANGED
@@ -10,37 +10,6 @@
10
  // through dot(). Bounds checks cover every partial tile. Accumulation, bias,
11
  // and the optional activation epilogue use f32; f16 operands widen at the FMA.
12
  {{ env.wgsl.resourceDeclarations }}
13
- {% set hasActivation = hasActivation is defined and hasActivation %}
14
- {% set activation = activation | default("") %}
15
- {% set actAlpha = actAlpha | default(0.0) %}
16
- {% set actBeta = actBeta | default(0.0) %}
17
- {% if hasActivation %}
18
- // Apply the fused activation in the f32 accumulator before the single output
19
- // cast, avoiding an intermediate convolution tensor.
20
- {% macro fused_act_return(mode, alpha, beta) -%}
21
- {% if mode == "Relu" %}
22
- return max(v, 0.0);
23
- {% elif mode == "Clip" %}
24
- return clamp(v, f32({{ alpha }}), f32({{ beta }}));
25
- {% elif mode == "LeakyRelu" %}
26
- return select(v * f32({{ alpha }}), v, v >= 0.0);
27
- {% elif mode == "Sigmoid" %}
28
- return 1.0 / (1.0 + exp(-v));
29
- {% elif mode == "Tanh" %}
30
- // tanh is already +/-1 to f32 precision at the clamp bounds. Clamping before
31
- // the builtin preserves that saturated result for larger accumulators.
32
- return tanh(clamp(v, -10.0, 10.0));
33
- {% elif mode == "HardSigmoid" %}
34
- return clamp(f32({{ alpha }}) * v + f32({{ beta }}), 0.0, 1.0);
35
- {% else %}
36
- return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
37
- {% endif %}
38
- {%- endmacro -%}
39
- fn fused_act(v: f32) -> f32 {
40
- {{ fused_act_return(activation, actAlpha, actBeta) -}}
41
- }
42
- {% endif %}
43
-
44
  {% set tileT = "f16" if usesF16 else "f32" %}
45
 
46
  const WG_X: u32 = {{ convWgX }}u;
@@ -59,9 +28,6 @@ fn store_out(oc: u32, ow: u32, yBase: u32, raw: f32) {
59
  var v = raw;
60
  {% if hasBias %}
61
  v = v + f32(bias[oc]);
62
- {% endif %}
63
- {% if hasActivation %}
64
- v = fused_act(v);
65
  {% endif %}
66
  y[yBase + oc * params.outW + ow] = {{ T }}(v);
67
  }
 
10
  // through dot(). Bounds checks cover every partial tile. Accumulation, bias,
11
  // and the optional activation epilogue use f32; f16 operands widen at the FMA.
12
  {{ env.wgsl.resourceDeclarations }}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
13
  {% set tileT = "f16" if usesF16 else "f32" %}
14
 
15
  const WG_X: u32 = {{ convWgX }}u;
 
28
  var v = raw;
29
  {% if hasBias %}
30
  v = v + f32(bias[oc]);
 
 
 
31
  {% endif %}
32
  y[yBase + oc * params.outW + ow] = {{ T }}(v);
33
  }
build/webgpu/conv2d-grouped-large-w4.wgsl.jinja CHANGED
@@ -1,43 +1,15 @@
 
 
 
 
 
 
1
  {{ env.wgsl.resourceDeclarations }}
2
- {% set hasActivation = hasActivation is defined and hasActivation %}
3
- {% set activation = activation | default("") %}
4
- {% set actAlpha = actAlpha | default(0.0) %}
5
- {% set actBeta = actBeta | default(0.0) %}
6
- {% if hasActivation %}
7
- // Apply the fused activation in the f32 accumulator before the single output
8
- // cast, avoiding an intermediate convolution tensor.
9
- {% macro fused_act_return(mode, alpha, beta) -%}
10
- {% if mode == "Relu" %}
11
- return max(v, 0.0);
12
- {% elif mode == "Clip" %}
13
- return clamp(v, f32({{ alpha }}), f32({{ beta }}));
14
- {% elif mode == "LeakyRelu" %}
15
- return select(v * f32({{ alpha }}), v, v >= 0.0);
16
- {% elif mode == "Sigmoid" %}
17
- return 1.0 / (1.0 + exp(-v));
18
- {% elif mode == "Tanh" %}
19
- // tanh is already +/-1 to f32 precision at the clamp bounds. Clamping before
20
- // the builtin preserves that saturated result for larger accumulators.
21
- return tanh(clamp(v, -10.0, 10.0));
22
- {% elif mode == "HardSigmoid" %}
23
- return clamp(f32({{ alpha }}) * v + f32({{ beta }}), 0.0, 1.0);
24
- {% else %}
25
- return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
26
- {% endif %}
27
- {%- endmacro -%}
28
- fn fused_act(v: f32) -> f32 {
29
- {{ fused_act_return(activation, actAlpha, actBeta) -}}
30
- }
31
- {% endif %}
32
-
33
  // Large grouped kernels are dominated by repeatedly fetching the same filter
34
  // value for neighboring output columns and input window for neighboring output
35
  // channels. One invocation computes four adjacent columns for OC_TILE channels,
36
  // sharing each input load across those channels. Shape and window geometry are
37
  // static, so indexing divisors and kernel offsets are shader constants.
38
- {% set dilatedLanes = dilatedLanes | default(false) %}
39
- {% set lanes = lanes | default(4) %}
40
- {% set quads = quads | default(1) %}
41
  {% set laneStep = dilationWSpec if dilatedLanes else strideWSpec %}
42
  {% set comps = ["x", "y", "z", "w"] %}
43
  {% macro acc_ref(oct, lane) %}{% if dilatedLanes %}acc{{ oct }}_{{ lane }}{% else %}acc{{ oct }}.{{ comps[lane] }}{% endif %}{% endmacro %}
@@ -86,8 +58,7 @@ const WG: u32 = {{ workgroupSizeSpec }}u;
86
 
87
  @compute @workgroup_size({{ workgroupSizeSpec }})
88
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
89
- let q = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
90
- if (q >= COUNT_TILES) { return; }
91
 
92
  {% if dilatedLanes %}
93
  let tile = q % TILES_PER_ROW;
@@ -128,9 +99,10 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
128
  {% endfor %}
129
 
130
  for (var ic = 0u; ic < IN_CPG; ic++) {
131
- {% for kh in range(kernelHSpec) %}
 
132
  {
133
- let ih = ihBase + {{ kh * dilationHSpec }};
134
  if (ih >= 0 && ih < IN_H) {
135
  let xRow = xChannelBase + u32(ih) * IN_W_U;
136
  {% if registerForm %}
@@ -170,7 +142,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
170
  {% for kw in range(kernelWSpec) %}
171
  {
172
  {% for oct in range(ocTile) %}
173
- let weight{{ oct }} = f32(w[wBase{{ oct }} + {{ kh * kernelWSpec + kw }}u]);
174
  {% endfor %}
175
  {% for lane in range(lanes) %}
176
  {% for oct in range(ocTile) %}
@@ -183,7 +155,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
183
  {% for kw in range(kernelWSpec) %}
184
  {
185
  {% for oct in range(ocTile) %}
186
- let weight{{ oct }} = f32(w[wBase{{ oct }} + {{ kh * kernelWSpec + kw }}u]);
187
  {% endfor %}
188
  let iwK = iwBase + {{ kw * dilationWSpec }};
189
  {% for lane in range(lanes) %}
@@ -201,7 +173,8 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
201
  }
202
  }
203
  {% endfor %}
204
- xChannelBase += IN_PLANE;
 
205
  {% for oct in range(ocTile) %}
206
  wBase{{ oct }} += KAREA;
207
  {% endfor %}
@@ -212,7 +185,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
212
  let outChannelBase{{ oct }} = outBase + {{ oct }}u * OUT_H * OUT_W;
213
  {% for lane in range(lanes) %}
214
  if (ow0 + {{ lane * dilationWSpec }}u < OUT_W) {
215
- y[outChannelBase{{ oct }} + ow0 + {{ lane * dilationWSpec }}u] = {{ scalar }}({% if hasActivation %}fused_act(acc{{ oct }}_{{ lane }}){% else %}acc{{ oct }}_{{ lane }}{% endif %});
216
  }
217
  {% endfor %}
218
  {% endfor %}
@@ -222,14 +195,14 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
222
  let outChannelBase{{ oct }} = outBase + {{ oct }}u * OUT_H * OUT_W;
223
  {% for lane in range(4) %}
224
  if (ow4 * 4u + {{ lane }}u < OUT_W) {
225
- y[outChannelBase{{ oct }} + {{ lane }}u] = {{ scalar }}({% if hasActivation %}fused_act(acc{{ oct }}.{{ ["x", "y", "z", "w"][lane] }}){% else %}acc{{ oct }}.{{ ["x", "y", "z", "w"][lane] }}{% endif %});
226
  }
227
  {% endfor %}
228
  {% endfor %}
229
  {% else %}
230
  let outBase = ((batch * OUT_C + ocBase) * OUT_H + oh) * OUT_W4 + ow4;
231
  {% for oct in range(ocTile) %}
232
- y[outBase + {{ oct }}u * OUT_PLANE4] = {{ vectorScalar }}({% if hasActivation %}vec4<f32>(fused_act(acc{{ oct }}.x), fused_act(acc{{ oct }}.y), fused_act(acc{{ oct }}.z), fused_act(acc{{ oct }}.w)){% else %}acc{{ oct }}{% endif %});
233
  {% endfor %}
234
  {% endif %}
235
  }
 
1
+ {% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
2
+ {% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
3
+ // 2D-folded flat index: gid.y carries the high bits past the dispatch's
4
+ // per-axis workgroup fold width.
5
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
6
+ if ({{ name }} >= {{ bound }}) { return; }{% endmacro %}
7
  {{ env.wgsl.resourceDeclarations }}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8
  // Large grouped kernels are dominated by repeatedly fetching the same filter
9
  // value for neighboring output columns and input window for neighboring output
10
  // channels. One invocation computes four adjacent columns for OC_TILE channels,
11
  // sharing each input load across those channels. Shape and window geometry are
12
  // static, so indexing divisors and kernel offsets are shader constants.
 
 
 
13
  {% set laneStep = dilationWSpec if dilatedLanes else strideWSpec %}
14
  {% set comps = ["x", "y", "z", "w"] %}
15
  {% macro acc_ref(oct, lane) %}{% if dilatedLanes %}acc{{ oct }}_{{ lane }}{% else %}acc{{ oct }}.{{ comps[lane] }}{% endif %}{% endmacro %}
 
58
 
59
  @compute @workgroup_size({{ workgroupSizeSpec }})
60
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
61
+ {{ flat_index_2d("WG", "q", "COUNT_TILES", guardInline=true) }}
 
62
 
63
  {% if dilatedLanes %}
64
  let tile = q % TILES_PER_ROW;
 
99
  {% endfor %}
100
 
101
  for (var ic = 0u; ic < IN_CPG; ic++) {
102
+ {% if rollKernelRows %} for (var kernelRow = 0u; kernelRow < {{ kernelHSpec }}u; kernelRow++) {
103
+ {% endif %}{% for kh in range(1 if rollKernelRows else kernelHSpec) %}
104
  {
105
+ let ih = ihBase + {% if rollKernelRows %}i32(kernelRow) * {{ dilationHSpec }}{% else %}{{ kh * dilationHSpec }}{% endif %};
106
  if (ih >= 0 && ih < IN_H) {
107
  let xRow = xChannelBase + u32(ih) * IN_W_U;
108
  {% if registerForm %}
 
142
  {% for kw in range(kernelWSpec) %}
143
  {
144
  {% for oct in range(ocTile) %}
145
+ let weight{{ oct }} = f32(w[wBase{{ oct }} + {% if rollKernelRows %}kernelRow * {{ kernelWSpec }}u + {{ kw }}u{% else %}{{ kh * kernelWSpec + kw }}u{% endif %}]);
146
  {% endfor %}
147
  {% for lane in range(lanes) %}
148
  {% for oct in range(ocTile) %}
 
155
  {% for kw in range(kernelWSpec) %}
156
  {
157
  {% for oct in range(ocTile) %}
158
+ let weight{{ oct }} = f32(w[wBase{{ oct }} + {% if rollKernelRows %}kernelRow * {{ kernelWSpec }}u + {{ kw }}u{% else %}{{ kh * kernelWSpec + kw }}u{% endif %}]);
159
  {% endfor %}
160
  let iwK = iwBase + {{ kw * dilationWSpec }};
161
  {% for lane in range(lanes) %}
 
173
  }
174
  }
175
  {% endfor %}
176
+ {% if rollKernelRows %} }
177
+ {% endif %} xChannelBase += IN_PLANE;
178
  {% for oct in range(ocTile) %}
179
  wBase{{ oct }} += KAREA;
180
  {% endfor %}
 
185
  let outChannelBase{{ oct }} = outBase + {{ oct }}u * OUT_H * OUT_W;
186
  {% for lane in range(lanes) %}
187
  if (ow0 + {{ lane * dilationWSpec }}u < OUT_W) {
188
+ y[outChannelBase{{ oct }} + ow0 + {{ lane * dilationWSpec }}u] = {{ scalar }}(acc{{ oct }}_{{ lane }});
189
  }
190
  {% endfor %}
191
  {% endfor %}
 
195
  let outChannelBase{{ oct }} = outBase + {{ oct }}u * OUT_H * OUT_W;
196
  {% for lane in range(4) %}
197
  if (ow4 * 4u + {{ lane }}u < OUT_W) {
198
+ y[outChannelBase{{ oct }} + {{ lane }}u] = {{ scalar }}(acc{{ oct }}.{{ ["x", "y", "z", "w"][lane] }});
199
  }
200
  {% endfor %}
201
  {% endfor %}
202
  {% else %}
203
  let outBase = ((batch * OUT_C + ocBase) * OUT_H + oh) * OUT_W4 + ow4;
204
  {% for oct in range(ocTile) %}
205
+ y[outBase + {{ oct }}u * OUT_PLANE4] = {{ vectorScalar }}(acc{{ oct }});
206
  {% endfor %}
207
  {% endif %}
208
  }
build/webgpu/manifest.json CHANGED
The diff for this file is too large to render. See raw diff
 
build/webgpu/metadata.json CHANGED
@@ -1,31 +1,31 @@
1
  {
2
  "name": "ai.onnx.Conv",
3
- "id": "_ai_onnx_conv_webgpu_010d4b0",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
- "bench.json": "H1K2v5/L5nbBfjqTwogslpt8SrxTH6se0gj31ONSJeU=",
11
- "conv-1x1-channel-reduce.wgsl.jinja": "Zjbu1oi+w5guLw54Y9C6TVJo0SntZDjgtsiie3o5md0=",
12
- "conv-1x1-gemm-tiled-reg.wgsl.jinja": "9yFpwXS1HzmWmm4W72n+aOIryLqc9+e88awZ1ogTe+Q=",
13
- "conv-1x1-gemm-tiled.wgsl.jinja": "uqG5mRVvcHiUo0LDCUcVgqElVTvRNknx19Gm4oIALiU=",
14
- "conv-1x1-subgroup-matrix.wgsl.jinja": "iDJs4RJk4iJarOZCjY7/5MFrHgR7CcqI6x5uwv2dYaA=",
15
- "conv-direct-nd.wgsl.jinja": "FqwfAKVCF1BYBAtU7Y3rDcGrmZ5AYkH4FuTzdRlT2Rg=",
16
- "conv-direct-unrolled.wgsl.jinja": "vrGUPwYkbYE1slII7nuK5wbERRcG5u0X1cW1kMc2sQs=",
17
  "conv-im2col-ncdhw.wgsl.jinja": "vQtnzeEKC2wm23+kwcaWoX3VSaJSqno18ACe3w3trhI=",
18
  "conv-im2col-nchw.wgsl.jinja": "IRpmtpeVQsLE6ud+Sf1230XaZVS1dxDtTZuSDwU3P8s=",
19
- "conv-splitk-reduce.wgsl.jinja": "LbVU5DRL0Q51fvLjThmxJQZH8aPbZkvNzij4nY/cHmo=",
20
- "conv1d-tiled-reg.wgsl.jinja": "ERqOPl+1RbnJ+CFqZvLSYw45qYMKgB7Je0teew5gyOM=",
21
- "conv2d-grouped-large-w4.wgsl.jinja": "v8aAMQ2saiOgbQwLZAiVr3IeDT3h5YLO/v5tGL1x8pk=",
22
- "manifest.json": "3/3/k9ZVj/vLw3xZ66XDMGWb86wzdZCoqDS0tKg9ziE=",
23
- "test.json": "1OUEVPx0U/BTLtlxoSCFOWVgMfrlOfQb/TUY3VivfGw="
24
  }
25
  },
26
- "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
27
  "webgpu": {
28
- "manifestSpec": "2.0",
29
  "variants": {
30
  "ncw1d": ["conv-direct-nd.wgsl.jinja"],
31
  "ncw1d_bias": ["conv-direct-nd.wgsl.jinja"],
@@ -51,8 +51,6 @@
51
  "im2col_direct_f16_subgroup_matrix_bias": ["conv-1x1-subgroup-matrix.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
52
  "implicit_im2col_subgroup_matrix": ["conv-1x1-subgroup-matrix.wgsl.jinja"],
53
  "implicit_im2col_subgroup_matrix_bias": ["conv-1x1-subgroup-matrix.wgsl.jinja"],
54
- "implicit_im2col3d_subgroup_matrix": ["conv-1x1-subgroup-matrix.wgsl.jinja"],
55
- "implicit_im2col3d_subgroup_matrix_bias": ["conv-1x1-subgroup-matrix.wgsl.jinja"],
56
  "implicit_im2col_subgroup_matrix_splitk": ["conv-1x1-subgroup-matrix.wgsl.jinja", "conv-splitk-reduce.wgsl.jinja"],
57
  "implicit_im2col_subgroup_matrix_bias_splitk": ["conv-1x1-subgroup-matrix.wgsl.jinja", "conv-splitk-reduce.wgsl.jinja"],
58
  "im2col3d_gemm_subgroup_matrix": ["conv-1x1-subgroup-matrix.wgsl.jinja", "conv-im2col-ncdhw.wgsl.jinja"],
@@ -74,6 +72,10 @@
74
  "im2col_gemm_tiled": ["conv-1x1-gemm-tiled.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
75
  "im2col_gemm_tiled_bias_reg": ["conv-1x1-gemm-tiled-reg.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
76
  "im2col_gemm_tiled_bias": ["conv-1x1-gemm-tiled.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
 
 
 
 
77
  "gemm_1x1_tiled_reg": ["conv-1x1-gemm-tiled-reg.wgsl.jinja"],
78
  "gemm_1x1_tiled_bias_reg": ["conv-1x1-gemm-tiled-reg.wgsl.jinja"],
79
  "gemm_1x1_tiled": ["conv-1x1-gemm-tiled.wgsl.jinja"],
 
1
  {
2
  "name": "ai.onnx.Conv",
3
+ "id": "_ai_onnx_conv_webgpu_425dc7b",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
+ "bench.json": "qQlQzLmos636YZMqlseiLZNw3VMjVExXhvQhIkhNmso=",
11
+ "conv-1x1-channel-reduce.wgsl.jinja": "Ug2iTybBtLS8dnr7KOGbxUYKidE0/sb7LGLeAP3xka8=",
12
+ "conv-1x1-gemm-tiled-reg.wgsl.jinja": "fqy1AlbukRB6+EsP4NT8IsBbiAYWtxE7o8z87QN7nHI=",
13
+ "conv-1x1-gemm-tiled.wgsl.jinja": "mQWnnhk9Xlj20YcMc+6sQBuy67ZhdkgPUSdYgtauDaY=",
14
+ "conv-1x1-subgroup-matrix.wgsl.jinja": "/M/DUaLDk2rBCojLxieYQxQhpMei79mv4iM0shwJbl8=",
15
+ "conv-direct-nd.wgsl.jinja": "qooloadvj03r5vD7sSuCe1xGrWeuCxNx7B4lUjilOGQ=",
16
+ "conv-direct-unrolled.wgsl.jinja": "GsmGTy9UvercDDC0vNc3OOve9xcRtg2Hj9rSpx9YOSc=",
17
  "conv-im2col-ncdhw.wgsl.jinja": "vQtnzeEKC2wm23+kwcaWoX3VSaJSqno18ACe3w3trhI=",
18
  "conv-im2col-nchw.wgsl.jinja": "IRpmtpeVQsLE6ud+Sf1230XaZVS1dxDtTZuSDwU3P8s=",
19
+ "conv-splitk-reduce.wgsl.jinja": "jPCiZlPWzoQEoa9+ZwfZdSOHH+/hSYknvRFHKMvgkKM=",
20
+ "conv1d-tiled-reg.wgsl.jinja": "HvWuECHIzxqPrIyvfdlljhdOXpwzZo0jvAYuYDY+3kM=",
21
+ "conv2d-grouped-large-w4.wgsl.jinja": "mt2+nqnbFfJMac9s3sAsjTZlySD9VJFE+gHnkToVwt8=",
22
+ "manifest.json": "RtFRWBQdJ7bpJaP2G26O3a2GIKy1QaRpfZQqTFFmoYs=",
23
+ "test.json": "Y5+ejG2R11fFhIYo5ExoKSNeCtn278UgO3/JQgG1bY0="
24
  }
25
  },
26
+ "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
27
  "webgpu": {
28
+ "manifestSpec": "2.1",
29
  "variants": {
30
  "ncw1d": ["conv-direct-nd.wgsl.jinja"],
31
  "ncw1d_bias": ["conv-direct-nd.wgsl.jinja"],
 
51
  "im2col_direct_f16_subgroup_matrix_bias": ["conv-1x1-subgroup-matrix.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
52
  "implicit_im2col_subgroup_matrix": ["conv-1x1-subgroup-matrix.wgsl.jinja"],
53
  "implicit_im2col_subgroup_matrix_bias": ["conv-1x1-subgroup-matrix.wgsl.jinja"],
 
 
54
  "implicit_im2col_subgroup_matrix_splitk": ["conv-1x1-subgroup-matrix.wgsl.jinja", "conv-splitk-reduce.wgsl.jinja"],
55
  "implicit_im2col_subgroup_matrix_bias_splitk": ["conv-1x1-subgroup-matrix.wgsl.jinja", "conv-splitk-reduce.wgsl.jinja"],
56
  "im2col3d_gemm_subgroup_matrix": ["conv-1x1-subgroup-matrix.wgsl.jinja", "conv-im2col-ncdhw.wgsl.jinja"],
 
72
  "im2col_gemm_tiled": ["conv-1x1-gemm-tiled.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
73
  "im2col_gemm_tiled_bias_reg": ["conv-1x1-gemm-tiled-reg.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
74
  "im2col_gemm_tiled_bias": ["conv-1x1-gemm-tiled.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
75
+ "im2col_gemm_tiled_reg_f16_columns": ["conv-1x1-gemm-tiled-reg.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
76
+ "im2col_gemm_tiled_f16_columns": ["conv-1x1-gemm-tiled.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
77
+ "im2col_gemm_tiled_bias_reg_f16_columns": ["conv-1x1-gemm-tiled-reg.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
78
+ "im2col_gemm_tiled_bias_f16_columns": ["conv-1x1-gemm-tiled.wgsl.jinja", "conv-im2col-nchw.wgsl.jinja"],
79
  "gemm_1x1_tiled_reg": ["conv-1x1-gemm-tiled-reg.wgsl.jinja"],
80
  "gemm_1x1_tiled_bias_reg": ["conv-1x1-gemm-tiled-reg.wgsl.jinja"],
81
  "gemm_1x1_tiled": ["conv-1x1-gemm-tiled.wgsl.jinja"],
build/webgpu/test.json CHANGED
The diff for this file is too large to render. See raw diff