sync 6fdf6301e2bb
Browse files- README.md +22 -10
- build/webgpu/bench.json +0 -0
- build/webgpu/conv-1x1-channel-reduce.wgsl.jinja +9 -12
- build/webgpu/conv-1x1-gemm-tiled-reg.wgsl.jinja +19 -158
- build/webgpu/conv-1x1-gemm-tiled.wgsl.jinja +13 -50
- build/webgpu/conv-1x1-subgroup-matrix.wgsl.jinja +12 -146
- build/webgpu/conv-direct-nd.wgsl.jinja +10 -44
- build/webgpu/conv-direct-unrolled.wgsl.jinja +10 -43
- build/webgpu/conv-splitk-reduce.wgsl.jinja +9 -43
- build/webgpu/conv1d-tiled-reg.wgsl.jinja +0 -34
- build/webgpu/conv2d-grouped-large-w4.wgsl.jinja +17 -44
- build/webgpu/manifest.json +0 -0
- build/webgpu/metadata.json +19 -17
- build/webgpu/test.json +0 -0
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` —
|
| 63 |
-
- `im2col_direct_f16_subgroup_matrix` —
|
| 64 |
-
- `im2col_direct_f32_subgroup_matrix_bias` —
|
| 65 |
-
- `im2col_direct_f16_subgroup_matrix_bias` —
|
| 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
|
| 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.
|
| 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 |
-
|
| 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
|
| 18 |
-
{% set
|
| 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 =
|
| 52 |
{% set GEMM_BM = gemmMTile if gemmMTile is defined else 64 %}
|
| 53 |
-
{% set GEMM_BN =
|
| 54 |
-
{% set GEMM_TN =
|
| 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 =
|
| 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
|
| 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 %}
|
| 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
|
| 10 |
-
{% set
|
| 11 |
-
{% set
|
| 12 |
-
{% set
|
| 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 |
-
{%
|
| 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(
|
| 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 *
|
| 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) /
|
| 92 |
-
let idx = li + e *
|
| 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 |
-
|
| 119 |
-
|
| 120 |
-
|
| 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 =
|
| 16 |
-
{% set polyphase =
|
| 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 =
|
| 24 |
-
{% set
|
| 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 =
|
| 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 =
|
| 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 = "
|
| 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 %}
|
| 378 |
-
{
|
| 379 |
-
{%
|
| 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
|
| 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 |
-
{%
|
| 5 |
-
|
| 6 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
-
|
| 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 |
-
|
| 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 |
-
|
| 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 |
-
|
| 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 |
-
|
| 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 |
-
|
| 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 |
-
{%
|
|
|
|
| 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 |
-
|
|
|
|
| 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 }}(
|
| 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 }}(
|
| 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 }}(
|
| 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": "
|
| 4 |
"version": 1,
|
| 5 |
"license": "Apache-2.0",
|
| 6 |
"backend": { "type": "webgpu" },
|
| 7 |
"digest": {
|
| 8 |
"algorithm": "sha256",
|
| 9 |
"files": {
|
| 10 |
-
"bench.json": "
|
| 11 |
-
"conv-1x1-channel-reduce.wgsl.jinja": "
|
| 12 |
-
"conv-1x1-gemm-tiled-reg.wgsl.jinja": "
|
| 13 |
-
"conv-1x1-gemm-tiled.wgsl.jinja": "
|
| 14 |
-
"conv-1x1-subgroup-matrix.wgsl.jinja": "
|
| 15 |
-
"conv-direct-nd.wgsl.jinja": "
|
| 16 |
-
"conv-direct-unrolled.wgsl.jinja": "
|
| 17 |
"conv-im2col-ncdhw.wgsl.jinja": "vQtnzeEKC2wm23+kwcaWoX3VSaJSqno18ACe3w3trhI=",
|
| 18 |
"conv-im2col-nchw.wgsl.jinja": "IRpmtpeVQsLE6ud+Sf1230XaZVS1dxDtTZuSDwU3P8s=",
|
| 19 |
-
"conv-splitk-reduce.wgsl.jinja": "
|
| 20 |
-
"conv1d-tiled-reg.wgsl.jinja": "
|
| 21 |
-
"conv2d-grouped-large-w4.wgsl.jinja": "
|
| 22 |
-
"manifest.json": "
|
| 23 |
-
"test.json": "
|
| 24 |
}
|
| 25 |
},
|
| 26 |
-
"provenance": { "kernel": { "sha": "
|
| 27 |
"webgpu": {
|
| 28 |
-
"manifestSpec": "2.
|
| 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
|
|
|