sync 91d990483a17
Browse files- README.md +15 -10
- build/webgpu/average-pool2d-nchw-horizontal-reuse.wgsl.jinja +21 -24
- build/webgpu/average-pool2d-nchw-w3s1-reuse.wgsl.jinja +12 -15
- build/webgpu/bench.json +33 -3
- build/webgpu/manifest.json +366 -339
- build/webgpu/metadata.json +27 -11
- build/webgpu/pool-global-reduction.wgsl.jinja +38 -37
- build/webgpu/pool-window-nd.wgsl.jinja +17 -20
- build/webgpu/pool-window-unroll.wgsl.jinja +28 -29
- build/webgpu/pool2d-nchw-separable.wgsl.jinja +48 -0
- build/webgpu/test.json +0 -0
README.md
CHANGED
|
@@ -18,15 +18,15 @@ See the [ONNX `AveragePool` spec](https://onnx.ai/onnx/operators/onnx__AveragePo
|
|
| 18 |
|
| 19 |
## Inputs
|
| 20 |
|
| 21 |
-
| Name |
|
| 22 |
| --- | --- | --- | --- | --- | --- | --- |
|
| 23 |
-
| `
|
| 24 |
|
| 25 |
## Outputs
|
| 26 |
|
| 27 |
-
| Name |
|
| 28 |
| --- | --- | --- | --- | --- | --- | --- |
|
| 29 |
-
| `
|
| 30 |
|
| 31 |
## Attributes
|
| 32 |
|
|
@@ -35,12 +35,12 @@ Attributes and default values (overridable per request):
|
|
| 35 |
| Attribute | Default | Description |
|
| 36 |
| --- | --- | --- |
|
| 37 |
| `auto_pad` | `"NOTSET"` | Deprecated auto-padding mode: `NOTSET` (use explicit pads), `SAME_UPPER`, `SAME_LOWER` (pad so output size is `ceil(input / stride)`), or `VALID` (no padding). |
|
| 38 |
-
| `count_include_pad` | `0` | When non-zero, pad pixels are counted in the divisor when computing the average; defaults to 0 (exclude pad). |
|
| 39 |
| `ceil_mode` | `0` | When non-zero, uses ceiling instead of floor when computing the output spatial shape; defaults to 0. |
|
|
|
|
|
|
|
| 40 |
| `kernel_shape` | — | Required kernel shape, with one positive value per spatial axis. |
|
| 41 |
-
| `strides` | — | Stride along each spatial axis. When omitted, every stride is 1. |
|
| 42 |
| `pads` | — | Padding at the beginning and end of each spatial axis, ordered as `[begin_0, ..., begin_n, end_0, ..., end_n]`. When omitted, every pad is 0. |
|
| 43 |
-
| `
|
| 44 |
|
| 45 |
## Type constraints
|
| 46 |
|
|
@@ -50,7 +50,7 @@ Attributes and default values (overridable per request):
|
|
| 50 |
|
| 51 |
## Files
|
| 52 |
|
| 53 |
-
- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
|
| 54 |
- [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
|
| 55 |
- [`test.json`](build/webgpu/test.json) — correctness cases
|
| 56 |
- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
|
|
@@ -59,13 +59,18 @@ Attributes and default values (overridable per request):
|
|
| 59 |
- [`pool-global-reduction.wgsl.jinja`](build/webgpu/pool-global-reduction.wgsl.jinja)
|
| 60 |
- [`pool-window-nd.wgsl.jinja`](build/webgpu/pool-window-nd.wgsl.jinja)
|
| 61 |
- [`pool-window-unroll.wgsl.jinja`](build/webgpu/pool-window-unroll.wgsl.jinja)
|
|
|
|
| 62 |
|
| 63 |
## Use with `@huggingface/kernels`
|
| 64 |
|
| 65 |
-
|
| 66 |
-
|
|
|
|
|
|
|
|
|
|
| 67 |
|
| 68 |
The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
|
|
|
|
| 69 |
|
| 70 |
Replace each `*Data` placeholder with a typed array containing the corresponding input data.
|
| 71 |
|
|
|
|
| 18 |
|
| 19 |
## Inputs
|
| 20 |
|
| 21 |
+
| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
|
| 22 |
| --- | --- | --- | --- | --- | --- | --- |
|
| 23 |
+
| `x` | `X` | `T` | — | — | Input data tensor of shape `(N x C x D1 x D2 ... Dn)`, where `N` is the batch size and `C` is the number of channels. | required |
|
| 24 |
|
| 25 |
## Outputs
|
| 26 |
|
| 27 |
+
| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
|
| 28 |
| --- | --- | --- | --- | --- | --- | --- |
|
| 29 |
+
| `y` | `Y` | `T` | same as `x` | derived | Output tensor with pooled values; spatial dimensions vary based on kernel, stride, dilation, and pad settings. | required |
|
| 30 |
|
| 31 |
## Attributes
|
| 32 |
|
|
|
|
| 35 |
| Attribute | Default | Description |
|
| 36 |
| --- | --- | --- |
|
| 37 |
| `auto_pad` | `"NOTSET"` | Deprecated auto-padding mode: `NOTSET` (use explicit pads), `SAME_UPPER`, `SAME_LOWER` (pad so output size is `ceil(input / stride)`), or `VALID` (no padding). |
|
|
|
|
| 38 |
| `ceil_mode` | `0` | When non-zero, uses ceiling instead of floor when computing the output spatial shape; defaults to 0. |
|
| 39 |
+
| `count_include_pad` | `0` | When non-zero, pad pixels are counted in the divisor when computing the average; defaults to 0 (exclude pad). |
|
| 40 |
+
| `dilations` | — | Dilation along each spatial axis. When omitted, every dilation is 1. |
|
| 41 |
| `kernel_shape` | — | Required kernel shape, with one positive value per spatial axis. |
|
|
|
|
| 42 |
| `pads` | — | Padding at the beginning and end of each spatial axis, ordered as `[begin_0, ..., begin_n, end_0, ..., end_n]`. When omitted, every pad is 0. |
|
| 43 |
+
| `strides` | — | Stride along each spatial axis. When omitted, every stride is 1. |
|
| 44 |
|
| 45 |
## Type constraints
|
| 46 |
|
|
|
|
| 50 |
|
| 51 |
## Files
|
| 52 |
|
| 53 |
+
- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
|
| 54 |
- [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
|
| 55 |
- [`test.json`](build/webgpu/test.json) — correctness cases
|
| 56 |
- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
|
|
|
|
| 59 |
- [`pool-global-reduction.wgsl.jinja`](build/webgpu/pool-global-reduction.wgsl.jinja)
|
| 60 |
- [`pool-window-nd.wgsl.jinja`](build/webgpu/pool-window-nd.wgsl.jinja)
|
| 61 |
- [`pool-window-unroll.wgsl.jinja`](build/webgpu/pool-window-unroll.wgsl.jinja)
|
| 62 |
+
- [`pool2d-nchw-separable.wgsl.jinja`](build/webgpu/pool2d-nchw-separable.wgsl.jinja)
|
| 63 |
|
| 64 |
## Use with `@huggingface/kernels`
|
| 65 |
|
| 66 |
+
```sh
|
| 67 |
+
npm install --save-exact @huggingface/kernels@0.0.1-preview.2
|
| 68 |
+
```
|
| 69 |
+
|
| 70 |
+
Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
|
| 71 |
|
| 72 |
The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
|
| 73 |
+
It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
|
| 74 |
|
| 75 |
Replace each `*Data` placeholder with a typed array containing the corresponding input data.
|
| 76 |
|
build/webgpu/average-pool2d-nchw-horizontal-reuse.wgsl.jinja
CHANGED
|
@@ -1,21 +1,18 @@
|
|
| 1 |
-
{% if usesF16 %}
|
| 2 |
-
enable f16;
|
| 3 |
-
{% endif %}
|
| 4 |
{{ env.wgsl.resourceDeclarations }}
|
| 5 |
|
| 6 |
// Four adjacent windows share source columns. For a K-wide, dilation-D window
|
| 7 |
// at stride S, their union spans (K-1)*D+1+3*S columns per row. Each output
|
| 8 |
// retains its own accumulation order.
|
| 9 |
-
const COUNT4: u32 = {{
|
| 10 |
-
const OUT_W4: u32 = max(1u, {{
|
| 11 |
-
const OUT_H: u32 = max(1u, {{
|
| 12 |
-
{% if
|
| 13 |
-
const OUT_W: u32 = {{
|
| 14 |
{% endif %}
|
| 15 |
-
const IN_H: i32 = {{
|
| 16 |
-
const IN_W: i32 = {{
|
| 17 |
-
const IN_W_U: u32 = {{
|
| 18 |
-
const IN_PLANE: u32 = {{
|
| 19 |
const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
|
| 20 |
|
| 21 |
fn load_value(validRow: bool, base: u32, col: i32) -> f32 {
|
|
@@ -24,8 +21,8 @@ fn load_value(validRow: bool, base: u32, col: i32) -> f32 {
|
|
| 24 |
}
|
| 25 |
|
| 26 |
@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
|
| 27 |
-
fn main(@builtin(global_invocation_id) gid: vec3<u32>
|
| 28 |
-
let q = gid.x + gid.y *
|
| 29 |
if (q >= COUNT4) { return; }
|
| 30 |
|
| 31 |
let ow0 = (q % OUT_W4) * 4u;
|
|
@@ -33,7 +30,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
|
|
| 33 |
let oh = t % OUT_H;
|
| 34 |
let plane = t / OUT_H;
|
| 35 |
let planeBase = plane * IN_PLANE;
|
| 36 |
-
let startW = i32(ow0 * {{
|
| 37 |
|
| 38 |
var acc0 = 0.0; var acc1 = 0.0; var acc2 = 0.0; var acc3 = 0.0;
|
| 39 |
// Window validity is separable: every valid kernel row combines with every
|
|
@@ -42,27 +39,27 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
|
|
| 42 |
var validRows = 0u;
|
| 43 |
{% for lane in range(4) %}
|
| 44 |
var validCols{{ lane }} = 0u;
|
| 45 |
-
{% for kw in range(
|
| 46 |
validCols{{ lane }} += select(
|
| 47 |
0u,
|
| 48 |
1u,
|
| 49 |
-
startW + {{ lane *
|
| 50 |
-
startW + {{ lane *
|
| 51 |
);
|
| 52 |
{% endfor %}
|
| 53 |
{% endfor %}
|
| 54 |
-
{% for kh in range(
|
| 55 |
{
|
| 56 |
-
let ih = i32(oh * {{
|
| 57 |
let validRow = ih >= 0 && ih < IN_H;
|
| 58 |
validRows += select(0u, 1u, validRow);
|
| 59 |
let rowBase = planeBase + u32(max(ih, 0)) * IN_W_U;
|
| 60 |
-
{% for j in range((
|
| 61 |
let v{{ kh }}_{{ j }} = load_value(validRow, rowBase, startW + {{ j }});
|
| 62 |
{% endfor %}
|
| 63 |
{% for lane in range(4) %}
|
| 64 |
-
{% for kw in range(
|
| 65 |
-
acc{{ lane }} += v{{ kh }}_{{ lane *
|
| 66 |
{% endfor %}
|
| 67 |
{% endfor %}
|
| 68 |
}
|
|
@@ -71,7 +68,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
|
|
| 71 |
let cnt{{ lane }} = validRows * validCols{{ lane }};
|
| 72 |
{% endfor %}
|
| 73 |
|
| 74 |
-
{% if
|
| 75 |
let outBase = (plane * OUT_H + oh) * OUT_W + ow0;
|
| 76 |
y[outBase] = {{ scalar }}(select(0.0, acc0 / f32(cnt0), cnt0 > 0u));
|
| 77 |
if (ow0 + 1u < OUT_W) { y[outBase + 1u] = {{ scalar }}(select(0.0, acc1 / f32(cnt1), cnt1 > 0u)); }
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
{{ env.wgsl.resourceDeclarations }}
|
| 2 |
|
| 3 |
// Four adjacent windows share source columns. For a K-wide, dilation-D window
|
| 4 |
// at stride S, their union spans (K-1)*D+1+3*S columns per row. Each output
|
| 5 |
// retains its own accumulation order.
|
| 6 |
+
const COUNT4: u32 = {{ count4 }}u;
|
| 7 |
+
const OUT_W4: u32 = max(1u, {{ outW4 }}u);
|
| 8 |
+
const OUT_H: u32 = max(1u, {{ outH }}u);
|
| 9 |
+
{% if scalarOutput is defined and scalarOutput %}
|
| 10 |
+
const OUT_W: u32 = {{ outW }}u;
|
| 11 |
{% endif %}
|
| 12 |
+
const IN_H: i32 = {{ inH }};
|
| 13 |
+
const IN_W: i32 = {{ inW }};
|
| 14 |
+
const IN_W_U: u32 = {{ inW }}u;
|
| 15 |
+
const IN_PLANE: u32 = {{ inPlane }}u;
|
| 16 |
const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
|
| 17 |
|
| 18 |
fn load_value(validRow: bool, base: u32, col: i32) -> f32 {
|
|
|
|
| 21 |
}
|
| 22 |
|
| 23 |
@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
|
| 24 |
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
| 25 |
+
let q = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
|
| 26 |
if (q >= COUNT4) { return; }
|
| 27 |
|
| 28 |
let ow0 = (q % OUT_W4) * 4u;
|
|
|
|
| 30 |
let oh = t % OUT_H;
|
| 31 |
let plane = t / OUT_H;
|
| 32 |
let planeBase = plane * IN_PLANE;
|
| 33 |
+
let startW = i32(ow0 * {{ strideWSpec }}u) - {{ padLeftSpec }};
|
| 34 |
|
| 35 |
var acc0 = 0.0; var acc1 = 0.0; var acc2 = 0.0; var acc3 = 0.0;
|
| 36 |
// Window validity is separable: every valid kernel row combines with every
|
|
|
|
| 39 |
var validRows = 0u;
|
| 40 |
{% for lane in range(4) %}
|
| 41 |
var validCols{{ lane }} = 0u;
|
| 42 |
+
{% for kw in range(kernelWSpec) %}
|
| 43 |
validCols{{ lane }} += select(
|
| 44 |
0u,
|
| 45 |
1u,
|
| 46 |
+
startW + {{ lane * strideWSpec + kw * dilationWSpec }} >= 0 &&
|
| 47 |
+
startW + {{ lane * strideWSpec + kw * dilationWSpec }} < IN_W,
|
| 48 |
);
|
| 49 |
{% endfor %}
|
| 50 |
{% endfor %}
|
| 51 |
+
{% for kh in range(kernelHSpec) %}
|
| 52 |
{
|
| 53 |
+
let ih = i32(oh * {{ strideHSpec }}u) + {{ kh * dilationHSpec - padTopSpec }};
|
| 54 |
let validRow = ih >= 0 && ih < IN_H;
|
| 55 |
validRows += select(0u, 1u, validRow);
|
| 56 |
let rowBase = planeBase + u32(max(ih, 0)) * IN_W_U;
|
| 57 |
+
{% for j in range((kernelWSpec - 1) * dilationWSpec + 1 + 3 * strideWSpec) %}
|
| 58 |
let v{{ kh }}_{{ j }} = load_value(validRow, rowBase, startW + {{ j }});
|
| 59 |
{% endfor %}
|
| 60 |
{% for lane in range(4) %}
|
| 61 |
+
{% for kw in range(kernelWSpec) %}
|
| 62 |
+
acc{{ lane }} += v{{ kh }}_{{ lane * strideWSpec + kw * dilationWSpec }};
|
| 63 |
{% endfor %}
|
| 64 |
{% endfor %}
|
| 65 |
}
|
|
|
|
| 68 |
let cnt{{ lane }} = validRows * validCols{{ lane }};
|
| 69 |
{% endfor %}
|
| 70 |
|
| 71 |
+
{% if scalarOutput is defined and scalarOutput %}
|
| 72 |
let outBase = (plane * OUT_H + oh) * OUT_W + ow0;
|
| 73 |
y[outBase] = {{ scalar }}(select(0.0, acc0 / f32(cnt0), cnt0 > 0u));
|
| 74 |
if (ow0 + 1u < OUT_W) { y[outBase + 1u] = {{ scalar }}(select(0.0, acc1 / f32(cnt1), cnt1 > 0u)); }
|
build/webgpu/average-pool2d-nchw-w3s1-reuse.wgsl.jinja
CHANGED
|
@@ -1,6 +1,3 @@
|
|
| 1 |
-
{% if usesF16 %}
|
| 2 |
-
enable f16;
|
| 3 |
-
{% endif %}
|
| 4 |
{{ env.wgsl.resourceDeclarations }}
|
| 5 |
|
| 6 |
// 2D NCHW AveragePool specialised for a 3-wide / stride-1 / pad-1 / dilation-1
|
|
@@ -10,14 +7,14 @@ enable f16;
|
|
| 10 |
// instead loads six contiguous columns per kernel row and slides the window in
|
| 11 |
// registers. Accumulation remains kh-major then kw-major, and excludes padding
|
| 12 |
// from the denominator.
|
| 13 |
-
const OUT_W: u32 = {{
|
| 14 |
-
const OUT_H: u32 = {{
|
| 15 |
-
const OUT_W_Q: u32 = {{
|
| 16 |
-
const IN_H: i32 = {{
|
| 17 |
-
const IN_W: i32 = {{
|
| 18 |
-
const IN_W_U: u32 = {{
|
| 19 |
-
const IN_PLANE: u32 = {{
|
| 20 |
-
const PAD_TOP: i32 = {{
|
| 21 |
|
| 22 |
// Load input element at (rowBase, col), returning 0 outside [0, IN_W). The row
|
| 23 |
// is already gated, so only the horizontal bound is checked here.
|
|
@@ -34,8 +31,8 @@ fn col_valid(rowValid: bool, col: i32) -> u32 {
|
|
| 34 |
}
|
| 35 |
|
| 36 |
@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
|
| 37 |
-
fn main(@builtin(global_invocation_id) gid: vec3<u32>
|
| 38 |
-
let tile = gid.x + gid.y *
|
| 39 |
if (tile >= params.tileCount) {
|
| 40 |
return;
|
| 41 |
}
|
|
@@ -53,9 +50,9 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
|
|
| 53 |
var acc0 = 0.0; var acc1 = 0.0; var acc2 = 0.0; var acc3 = 0.0;
|
| 54 |
var cnt0 = 0u; var cnt1 = 0u; var cnt2 = 0u; var cnt3 = 0u;
|
| 55 |
|
| 56 |
-
{% for kh in range(
|
| 57 |
{
|
| 58 |
-
let ih = startH + {{ kh *
|
| 59 |
let rowValid = ih >= 0 && ih < IN_H;
|
| 60 |
let rowBase = planeBase + u32(max(ih, 0)) * IN_W_U;
|
| 61 |
// Six contiguous input columns feeding the four 3-wide windows.
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
{{ env.wgsl.resourceDeclarations }}
|
| 2 |
|
| 3 |
// 2D NCHW AveragePool specialised for a 3-wide / stride-1 / pad-1 / dilation-1
|
|
|
|
| 7 |
// instead loads six contiguous columns per kernel row and slides the window in
|
| 8 |
// registers. Accumulation remains kh-major then kw-major, and excludes padding
|
| 9 |
// from the denominator.
|
| 10 |
+
const OUT_W: u32 = {{ outW }}u;
|
| 11 |
+
const OUT_H: u32 = {{ outH }}u;
|
| 12 |
+
const OUT_W_Q: u32 = {{ outWQuarter }}u; // OUT_W / 4 (columns-of-four)
|
| 13 |
+
const IN_H: i32 = {{ inH }};
|
| 14 |
+
const IN_W: i32 = {{ inW }};
|
| 15 |
+
const IN_W_U: u32 = {{ inW }}u;
|
| 16 |
+
const IN_PLANE: u32 = {{ inPlane }}u;
|
| 17 |
+
const PAD_TOP: i32 = {{ padTopSpec }};
|
| 18 |
|
| 19 |
// Load input element at (rowBase, col), returning 0 outside [0, IN_W). The row
|
| 20 |
// is already gated, so only the horizontal bound is checked here.
|
|
|
|
| 31 |
}
|
| 32 |
|
| 33 |
@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
|
| 34 |
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
| 35 |
+
let tile = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
|
| 36 |
if (tile >= params.tileCount) {
|
| 37 |
return;
|
| 38 |
}
|
|
|
|
| 50 |
var acc0 = 0.0; var acc1 = 0.0; var acc2 = 0.0; var acc3 = 0.0;
|
| 51 |
var cnt0 = 0u; var cnt1 = 0u; var cnt2 = 0u; var cnt3 = 0u;
|
| 52 |
|
| 53 |
+
{% for kh in range(kernelHSpec) %}
|
| 54 |
{
|
| 55 |
+
let ih = startH + {{ kh * dilationHSpec }};
|
| 56 |
let rowValid = ih >= 0 && ih < IN_H;
|
| 57 |
let rowBase = planeBase + u32(max(ih, 0)) * IN_W_U;
|
| 58 |
// Six contiguous input columns feeding the four 3-wide windows.
|
build/webgpu/bench.json
CHANGED
|
@@ -1,5 +1,4 @@
|
|
| 1 |
{
|
| 2 |
-
"op": "ai.onnx.AveragePool",
|
| 3 |
"tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
|
| 4 |
"cases": [
|
| 5 |
{
|
|
@@ -151,8 +150,8 @@
|
|
| 151 |
"name": "global-averagepool-rank5-serial-reduction-pathology",
|
| 152 |
"preset": "stress",
|
| 153 |
"provenance": {
|
| 154 |
-
"source": "
|
| 155 |
-
"notes": "Rank-5 global pooling over
|
| 156 |
},
|
| 157 |
"vars": { "batch": 1, "channels": 32, "inD": 32, "inH": 64, "inW": 64 },
|
| 158 |
"attrs": {
|
|
@@ -168,6 +167,37 @@
|
|
| 168 |
"bench": {
|
| 169 |
"metrics": [{ "type": "bandwidth", "value": "args.batch * args.channels * args.inD * args.inH * args.inW * 4" }]
|
| 170 |
}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 171 |
}
|
| 172 |
]
|
| 173 |
}
|
|
|
|
| 1 |
{
|
|
|
|
| 2 |
"tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
|
| 3 |
"cases": [
|
| 4 |
{
|
|
|
|
| 150 |
"name": "global-averagepool-rank5-serial-reduction-pathology",
|
| 151 |
"preset": "stress",
|
| 152 |
"provenance": {
|
| 153 |
+
"source": "synthetic benchmark",
|
| 154 |
+
"notes": "Rank-5 global pooling over 131,072 spatial values per plane exercises the parallel global-window reduction and its feature-tier alternatives."
|
| 155 |
},
|
| 156 |
"vars": { "batch": 1, "channels": 32, "inD": 32, "inH": 64, "inW": 64 },
|
| 157 |
"attrs": {
|
|
|
|
| 167 |
"bench": {
|
| 168 |
"metrics": [{ "type": "bandwidth", "value": "args.batch * args.channels * args.inD * args.inH * args.inW * 4" }]
|
| 169 |
}
|
| 170 |
+
},
|
| 171 |
+
{
|
| 172 |
+
"name": "nchw_8x64x112x112_k3s1_same_upper_spelling",
|
| 173 |
+
"preset": "smoke",
|
| 174 |
+
"vars": { "batch": 8, "channels": 64, "inH": 112, "inW": 112 },
|
| 175 |
+
"attrs": { "kernel_shape": [3, 3], "strides": [1, 1], "auto_pad": "SAME_UPPER" },
|
| 176 |
+
"inputs": { "x": { "dtype": "float32", "shape": [8, 64, 112, 112], "dist": "normal", "seed": 205 } },
|
| 177 |
+
"outputs": { "y": { "dtype": "float32", "shape": [8, 64, 112, 112], "dist": "empty" } },
|
| 178 |
+
"bench": {
|
| 179 |
+
"primary": true,
|
| 180 |
+
"metrics": [{ "type": "bandwidth", "value": "args.batch * args.channels * args.inH * args.inW * 4 * 2" }]
|
| 181 |
+
},
|
| 182 |
+
"provenance": {
|
| 183 |
+
"notes": "SAME_UPPER resolves to symmetric padding for this shape, so every pooling window has the same geometry as the explicit-padding form."
|
| 184 |
+
}
|
| 185 |
+
},
|
| 186 |
+
{
|
| 187 |
+
"name": "average-pool-f32-b8c64-56x56-k7s1-valid-no-pad",
|
| 188 |
+
"preset": "edge",
|
| 189 |
+
"attrs": { "kernel_shape": [7, 7], "strides": [1, 1] },
|
| 190 |
+
"inputs": { "x": { "dtype": "float32", "shape": [8, 64, 56, 56] } },
|
| 191 |
+
"outputs": { "y": { "dtype": "float32", "shape": [8, 64, 50, 50] } },
|
| 192 |
+
"bench": { "metrics": [{ "type": "bandwidth", "value": "(numel(shapes.x) + numel(shapes.y)) * 4" }] }
|
| 193 |
+
},
|
| 194 |
+
{
|
| 195 |
+
"name": "average-pool-f32-b8c64-56x56-k7s7-valid-no-pad",
|
| 196 |
+
"preset": "edge",
|
| 197 |
+
"attrs": { "kernel_shape": [7, 7], "strides": [7, 7] },
|
| 198 |
+
"inputs": { "x": { "dtype": "float32", "shape": [8, 64, 56, 56] } },
|
| 199 |
+
"outputs": { "y": { "dtype": "float32", "shape": [8, 64, 8, 8] } },
|
| 200 |
+
"bench": { "metrics": [{ "type": "bandwidth", "value": "(numel(shapes.x) + numel(shapes.y)) * 4" }] }
|
| 201 |
}
|
| 202 |
]
|
| 203 |
}
|
build/webgpu/manifest.json
CHANGED
|
@@ -2,32 +2,16 @@
|
|
| 2 |
"domain": "ai.onnx",
|
| 3 |
"name": "AveragePool",
|
| 4 |
"sinceVersion": 19,
|
| 5 |
-
"
|
| 6 |
-
"
|
| 7 |
-
|
| 8 |
-
|
| 9 |
-
|
| 10 |
-
|
| 11 |
-
}
|
| 12 |
-
|
| 13 |
-
|
| 14 |
-
{
|
| 15 |
-
"role": "Y",
|
| 16 |
-
"dtype": "T",
|
| 17 |
-
"description": "Output tensor with pooled values; spatial dimensions vary based on kernel, stride, dilation, and pad settings.",
|
| 18 |
-
"rank": "ranks.X",
|
| 19 |
-
"shape": "poolOutputShape"
|
| 20 |
-
}
|
| 21 |
-
],
|
| 22 |
-
"attributes": { "auto_pad": "NOTSET", "count_include_pad": 0, "ceil_mode": 0 },
|
| 23 |
-
"attributeDescriptions": {
|
| 24 |
-
"auto_pad": "Deprecated auto-padding mode: `NOTSET` (use explicit pads), `SAME_UPPER`, `SAME_LOWER` (pad so output size is `ceil(input / stride)`), or `VALID` (no padding).",
|
| 25 |
-
"kernel_shape": "Required kernel shape, with one positive value per spatial axis.",
|
| 26 |
-
"strides": "Stride along each spatial axis. When omitted, every stride is 1.",
|
| 27 |
-
"pads": "Padding at the beginning and end of each spatial axis, ordered as `[begin_0, ..., begin_n, end_0, ..., end_n]`. When omitted, every pad is 0.",
|
| 28 |
-
"dilations": "Dilation along each spatial axis. When omitted, every dilation is 1.",
|
| 29 |
-
"count_include_pad": "When non-zero, pad pixels are counted in the divisor when computing the average; defaults to 0 (exclude pad).",
|
| 30 |
-
"ceil_mode": "When non-zero, uses ceiling instead of floor when computing the output spatial shape; defaults to 0."
|
| 31 |
},
|
| 32 |
"attributeConstraints": {
|
| 33 |
"auto_pad": { "values": ["NOTSET", "SAME_UPPER", "SAME_LOWER", "VALID"] },
|
|
@@ -36,12 +20,9 @@
|
|
| 36 |
"kernel_shape": { "required": true }
|
| 37 |
},
|
| 38 |
"typeConstraints": { "T": ["float32", "float16"] },
|
| 39 |
-
"
|
| 40 |
-
"x": { "kind": "tensor", "semantic": "X", "role": "input" },
|
| 41 |
-
"y": { "kind": "tensor", "semantic": "Y", "role": "output" }
|
| 42 |
-
},
|
| 43 |
"derive": {
|
| 44 |
-
"spatialRank": "ranks.
|
| 45 |
"narrowSubgroupRange": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize < device.adapterInfo.subgroupMaxSize and device.adapterInfo.subgroupMaxSize <= 16",
|
| 46 |
"kernelShapeLengthOk": "not has(attrs, \"kernel_shape\") or (attrs.kernel_shape | length) == spatialRank",
|
| 47 |
"stridesLengthOk": "not has(attrs, \"strides\") or (attrs.strides | length) == spatialRank",
|
|
@@ -62,217 +43,93 @@
|
|
| 62 |
"padBack": "attrs.pads[spatialRank] if padsLengthOk and has(attrs, \"pads\") and spatialRank == 3 else 0",
|
| 63 |
"padBottom": "attrs.pads[2 * spatialRank - 2] if padsLengthOk and has(attrs, \"pads\") and spatialRank >= 2 else 0",
|
| 64 |
"padRight": "attrs.pads[2 * spatialRank - 1] if padsLengthOk and has(attrs, \"pads\") and spatialRank >= 1 else 0",
|
| 65 |
-
"poolBaseOk": "f16Ok(dtypes.T) and ranks.
|
| 66 |
-
"explicitOutW": "min(ceil((dim(shapes.
|
| 67 |
-
"explicitOutH": "min(ceil((dim(shapes.
|
| 68 |
-
"explicitOutD": "(min(ceil((dim(shapes.
|
| 69 |
-
"inferredOutW": "ceil(dim(shapes.
|
| 70 |
-
"inferredOutH": "ceil(dim(shapes.
|
| 71 |
-
"inferredOutD": "ceil(dim(shapes.
|
| 72 |
-
"poolOutputShape": "[dim(shapes.
|
| 73 |
-
"
|
| 74 |
-
"
|
| 75 |
-
"
|
| 76 |
-
"
|
| 77 |
-
"
|
| 78 |
-
"
|
|
|
|
|
|
|
|
|
|
| 79 |
"narrowF32W3UnrolledPreferred": "narrowSubgroupRange and dtypes.T == \"f32\"",
|
| 80 |
-
"globalWindowParallelCovered": "poolBaseOk and numel(shapes.
|
| 81 |
-
|
| 82 |
-
|
| 83 |
-
|
| 84 |
-
"
|
| 85 |
-
{ "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$xElem" },
|
| 86 |
-
{ "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
|
| 87 |
-
{
|
| 88 |
-
"name": "params",
|
| 89 |
-
"semantic": "kernel.params",
|
| 90 |
-
"buffer": { "type": "uniform" },
|
| 91 |
-
"struct": { "name": "Params", "fields": [{ "name": "spatial", "type": "u32", "value": "inner(shapes.X, 1)" }] }
|
| 92 |
-
}
|
| 93 |
-
],
|
| 94 |
-
"horizontal": [
|
| 95 |
-
{ "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
|
| 96 |
-
{ "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$outputElement" }
|
| 97 |
-
],
|
| 98 |
-
"w3Reuse": [
|
| 99 |
-
{ "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
|
| 100 |
-
{ "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$T" },
|
| 101 |
-
{
|
| 102 |
-
"name": "params",
|
| 103 |
-
"semantic": "kernel.params",
|
| 104 |
-
"buffer": { "type": "uniform" },
|
| 105 |
-
"struct": {
|
| 106 |
-
"name": "Params",
|
| 107 |
-
"fields": [{ "name": "tileCount", "type": "u32", "value": "numel(shapes.Y) / 4" }]
|
| 108 |
-
}
|
| 109 |
-
}
|
| 110 |
-
],
|
| 111 |
-
"flatDirect": [
|
| 112 |
-
{ "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
|
| 113 |
-
{ "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$T" },
|
| 114 |
-
{
|
| 115 |
-
"name": "params",
|
| 116 |
-
"semantic": "kernel.params",
|
| 117 |
-
"buffer": { "type": "uniform" },
|
| 118 |
-
"struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.Y)" }] }
|
| 119 |
-
}
|
| 120 |
-
],
|
| 121 |
-
"rank1Direct": [
|
| 122 |
-
{ "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
|
| 123 |
-
{ "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$T" },
|
| 124 |
-
{
|
| 125 |
-
"name": "params",
|
| 126 |
-
"semantic": "kernel.params",
|
| 127 |
-
"buffer": { "type": "uniform" },
|
| 128 |
-
"struct": {
|
| 129 |
-
"name": "Params",
|
| 130 |
-
"fields": [
|
| 131 |
-
{ "name": "channels", "type": "u32", "value": "dim(shapes.X, 1)" },
|
| 132 |
-
{ "name": "inW", "type": "u32", "value": "dim(shapes.X, 2)" },
|
| 133 |
-
{ "name": "outW", "type": "u32", "value": "dim(shapes.Y, 2)" },
|
| 134 |
-
{ "name": "kernelW", "type": "u32", "value": "kernelW" },
|
| 135 |
-
{ "name": "strideW", "type": "u32", "value": "strideW" },
|
| 136 |
-
{ "name": "dilationW", "type": "u32", "value": "dilationW" },
|
| 137 |
-
{ "name": "padLeft", "type": "i32", "value": "padLeft" },
|
| 138 |
-
{ "name": "padRight", "type": "i32", "value": "padRight" },
|
| 139 |
-
{ "name": "countIncludePad", "type": "u32", "value": "attrs.count_include_pad" },
|
| 140 |
-
{ "name": "count", "type": "u32", "value": "numel(shapes.Y)" }
|
| 141 |
-
]
|
| 142 |
-
}
|
| 143 |
-
}
|
| 144 |
-
],
|
| 145 |
-
"rank3Direct": [
|
| 146 |
-
{ "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
|
| 147 |
-
{ "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$T" },
|
| 148 |
-
{
|
| 149 |
-
"name": "params",
|
| 150 |
-
"semantic": "kernel.params",
|
| 151 |
-
"buffer": { "type": "uniform" },
|
| 152 |
-
"struct": {
|
| 153 |
-
"name": "Params",
|
| 154 |
-
"fields": [
|
| 155 |
-
{ "name": "channels", "type": "u32", "value": "dim(shapes.X, 1)" },
|
| 156 |
-
{ "name": "inD", "type": "u32", "value": "dim(shapes.X, 2)" },
|
| 157 |
-
{ "name": "inH", "type": "u32", "value": "dim(shapes.X, 3)" },
|
| 158 |
-
{ "name": "inW", "type": "u32", "value": "dim(shapes.X, 4)" },
|
| 159 |
-
{ "name": "outD", "type": "u32", "value": "dim(shapes.Y, 2)" },
|
| 160 |
-
{ "name": "outH", "type": "u32", "value": "dim(shapes.Y, 3)" },
|
| 161 |
-
{ "name": "outW", "type": "u32", "value": "dim(shapes.Y, 4)" },
|
| 162 |
-
{ "name": "kernelD", "type": "u32", "value": "kernelD" },
|
| 163 |
-
{ "name": "kernelH", "type": "u32", "value": "kernelH" },
|
| 164 |
-
{ "name": "kernelW", "type": "u32", "value": "kernelW" },
|
| 165 |
-
{ "name": "strideD", "type": "u32", "value": "strideD" },
|
| 166 |
-
{ "name": "strideH", "type": "u32", "value": "strideH" },
|
| 167 |
-
{ "name": "strideW", "type": "u32", "value": "strideW" },
|
| 168 |
-
{ "name": "dilationD", "type": "u32", "value": "dilationD" },
|
| 169 |
-
{ "name": "dilationH", "type": "u32", "value": "dilationH" },
|
| 170 |
-
{ "name": "dilationW", "type": "u32", "value": "dilationW" },
|
| 171 |
-
{ "name": "padFront", "type": "i32", "value": "padFront" },
|
| 172 |
-
{ "name": "padTop", "type": "i32", "value": "padTop" },
|
| 173 |
-
{ "name": "padLeft", "type": "i32", "value": "padLeft" },
|
| 174 |
-
{ "name": "padBack", "type": "i32", "value": "padBack" },
|
| 175 |
-
{ "name": "padBottom", "type": "i32", "value": "padBottom" },
|
| 176 |
-
{ "name": "padRight", "type": "i32", "value": "padRight" },
|
| 177 |
-
{ "name": "countIncludePad", "type": "u32", "value": "attrs.count_include_pad" },
|
| 178 |
-
{ "name": "count", "type": "u32", "value": "numel(shapes.Y)" }
|
| 179 |
-
]
|
| 180 |
-
}
|
| 181 |
-
}
|
| 182 |
-
],
|
| 183 |
-
"rank2Direct": [
|
| 184 |
-
{ "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
|
| 185 |
-
{ "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$T" },
|
| 186 |
-
{
|
| 187 |
-
"name": "params",
|
| 188 |
-
"semantic": "kernel.params",
|
| 189 |
-
"buffer": { "type": "uniform" },
|
| 190 |
-
"struct": {
|
| 191 |
-
"name": "Params",
|
| 192 |
-
"fields": [
|
| 193 |
-
{ "name": "channels", "type": "u32", "value": "dim(shapes.X, 1)" },
|
| 194 |
-
{ "name": "inH", "type": "u32", "value": "dim(shapes.X, 2)" },
|
| 195 |
-
{ "name": "inW", "type": "u32", "value": "dim(shapes.X, 3)" },
|
| 196 |
-
{ "name": "outH", "type": "u32", "value": "dim(shapes.Y, 2)" },
|
| 197 |
-
{ "name": "outW", "type": "u32", "value": "dim(shapes.Y, 3)" },
|
| 198 |
-
{ "name": "kernelH", "type": "u32", "value": "kernelH" },
|
| 199 |
-
{ "name": "kernelW", "type": "u32", "value": "kernelW" },
|
| 200 |
-
{ "name": "strideH", "type": "u32", "value": "strideH" },
|
| 201 |
-
{ "name": "strideW", "type": "u32", "value": "strideW" },
|
| 202 |
-
{ "name": "dilationH", "type": "u32", "value": "dilationH" },
|
| 203 |
-
{ "name": "dilationW", "type": "u32", "value": "dilationW" },
|
| 204 |
-
{ "name": "padTop", "type": "i32", "value": "effectivePadTop" },
|
| 205 |
-
{ "name": "padLeft", "type": "i32", "value": "effectivePadLeft" },
|
| 206 |
-
{ "name": "padBottom", "type": "i32", "value": "effectivePadBottom" },
|
| 207 |
-
{ "name": "padRight", "type": "i32", "value": "effectivePadRight" },
|
| 208 |
-
{ "name": "countIncludePad", "type": "u32", "value": "attrs.count_include_pad" },
|
| 209 |
-
{ "name": "count", "type": "u32", "value": "numel(shapes.Y)" }
|
| 210 |
-
]
|
| 211 |
-
}
|
| 212 |
-
}
|
| 213 |
-
]
|
| 214 |
},
|
| 215 |
-
"
|
| 216 |
"variants": [
|
| 217 |
{
|
| 218 |
"id": "global_window_parallel",
|
| 219 |
"priority": 50,
|
| 220 |
"when": ["globalWindowParallelCovered"],
|
| 221 |
-
"
|
| 222 |
-
"xElem": "(\"vec4<f16>\" if dtypes.T == \"f16\" else \"vec4<f32>\") if inner(shapes.
|
| 223 |
},
|
| 224 |
"passes": [
|
| 225 |
{
|
| 226 |
"id": "main",
|
| 227 |
"name": "AveragePool.GlobalWindowParallel",
|
| 228 |
-
"
|
| 229 |
-
|
| 230 |
-
|
| 231 |
-
|
| 232 |
-
|
| 233 |
-
|
| 234 |
-
|
| 235 |
-
}
|
| 236 |
},
|
| 237 |
-
"bindings":
|
| 238 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 239 |
}
|
| 240 |
]
|
| 241 |
},
|
| 242 |
{
|
| 243 |
"id": "nchw2d_k7s1_reuse_tail",
|
| 244 |
"priority": 34,
|
| 245 |
-
"when": ["poolBaseOk", "ranks.
|
| 246 |
-
"
|
| 247 |
"passes": [
|
| 248 |
{
|
| 249 |
"id": "main",
|
| 250 |
"name": "AveragePool.Nchw2dK7s1ReuseTail",
|
| 251 |
-
"
|
| 252 |
-
|
| 253 |
-
"
|
| 254 |
-
|
| 255 |
-
|
| 256 |
-
|
| 257 |
-
|
| 258 |
-
|
| 259 |
-
|
| 260 |
-
|
| 261 |
-
|
| 262 |
-
|
| 263 |
-
|
| 264 |
-
|
| 265 |
-
|
| 266 |
-
|
| 267 |
-
|
| 268 |
-
|
| 269 |
-
"scalarOutput": true
|
| 270 |
-
}
|
| 271 |
},
|
| 272 |
-
"bindings": "
|
| 273 |
"dispatch": {
|
| 274 |
-
"
|
| 275 |
-
"
|
|
|
|
| 276 |
}
|
| 277 |
}
|
| 278 |
]
|
|
@@ -280,65 +137,135 @@
|
|
| 280 |
{
|
| 281 |
"id": "nchw2d_w3s1_reuse",
|
| 282 |
"priority": 30,
|
|
|
|
| 283 |
"demoteWhen": ["narrowF32W3UnrolledPreferred"],
|
| 284 |
-
"when": ["poolBaseOk", "ranks.X == 4", "kernelW == 3", "strideW == 1", "dilationW == 1", "padLeft == 1", "padRight == 1", "kernelH >= 1", "kernelH <= 3", "strideH == 1", "dilationH >= 1", "attrs.count_include_pad == 0", "attrs.ceil_mode == 0", "dim(shapes.X, 3) % 4 == 0", "dim(shapes.Y, 2) == dim(shapes.X, 2) + padTop + padBottom - ((kernelH - 1) * dilationH + 1) + 1", "dim(shapes.Y, 3) == dim(shapes.X, 3)", "dim(shapes.Y, 2) > 0", "dim(shapes.Y, 3) > 0"],
|
| 285 |
"passes": [
|
| 286 |
{
|
| 287 |
"id": "main",
|
| 288 |
"name": "AveragePool.Nchw2dW3s1Reuse",
|
| 289 |
-
"
|
| 290 |
-
|
| 291 |
-
"
|
| 292 |
-
|
| 293 |
-
|
| 294 |
-
|
| 295 |
-
|
| 296 |
-
|
| 297 |
-
|
| 298 |
-
|
| 299 |
-
|
| 300 |
-
"outWQuarter": "floor(dim(shapes.Y, 3) / 4)"
|
| 301 |
-
}
|
| 302 |
},
|
| 303 |
-
"bindings":
|
| 304 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 305 |
}
|
| 306 |
]
|
| 307 |
},
|
| 308 |
{
|
| 309 |
"id": "nchw2d_w3s1_reuse_tail",
|
| 310 |
"priority": 29,
|
| 311 |
-
"when": ["poolBaseOk", "ranks.
|
| 312 |
-
"
|
| 313 |
"passes": [
|
| 314 |
{
|
| 315 |
"id": "main",
|
| 316 |
"name": "AveragePool.Nchw2dW3s1ReuseTail",
|
| 317 |
-
"
|
| 318 |
-
|
| 319 |
-
"
|
| 320 |
-
|
| 321 |
-
|
| 322 |
-
|
| 323 |
-
|
| 324 |
-
|
| 325 |
-
|
| 326 |
-
|
| 327 |
-
|
| 328 |
-
|
| 329 |
-
|
| 330 |
-
|
| 331 |
-
|
| 332 |
-
|
| 333 |
-
|
| 334 |
-
|
| 335 |
-
"scalarOutput": true
|
| 336 |
-
}
|
| 337 |
},
|
| 338 |
-
"bindings": "
|
| 339 |
"dispatch": {
|
| 340 |
-
"
|
| 341 |
-
"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 342 |
}
|
| 343 |
}
|
| 344 |
]
|
|
@@ -346,144 +273,244 @@
|
|
| 346 |
{
|
| 347 |
"id": "nchw2d_unrolled",
|
| 348 |
"priority": 25,
|
| 349 |
-
"when": ["poolBaseOk", "ranks.
|
| 350 |
"passes": [
|
| 351 |
{
|
| 352 |
"id": "main",
|
| 353 |
"name": "AveragePool.Nchw2dUnrolled",
|
| 354 |
-
"
|
| 355 |
-
|
| 356 |
-
"
|
| 357 |
-
|
| 358 |
-
|
| 359 |
-
|
| 360 |
-
|
| 361 |
-
|
| 362 |
-
|
| 363 |
-
|
| 364 |
-
|
| 365 |
-
|
| 366 |
-
|
| 367 |
-
|
| 368 |
-
|
| 369 |
-
|
| 370 |
-
|
| 371 |
-
|
| 372 |
-
"countIncludePad": "attrs.count_include_pad != 0"
|
| 373 |
-
}
|
| 374 |
},
|
| 375 |
-
"bindings":
|
| 376 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 377 |
}
|
| 378 |
]
|
| 379 |
},
|
| 380 |
{
|
| 381 |
"id": "ncl1d",
|
| 382 |
"priority": 10,
|
| 383 |
-
"when": ["poolBaseOk", "ranks.
|
| 384 |
"passes": [
|
| 385 |
{
|
| 386 |
"id": "main",
|
| 387 |
"name": "AveragePool.Ncl1d",
|
| 388 |
-
"
|
| 389 |
-
"
|
| 390 |
-
"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 391 |
}
|
| 392 |
]
|
| 393 |
},
|
| 394 |
{
|
| 395 |
"id": "nchw2d",
|
| 396 |
-
"when": ["poolBaseOk", "ranks.
|
| 397 |
"passes": [
|
| 398 |
{
|
| 399 |
"id": "main",
|
| 400 |
"name": "nchw2d",
|
| 401 |
-
"
|
| 402 |
-
"
|
| 403 |
-
"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 404 |
}
|
| 405 |
]
|
| 406 |
},
|
| 407 |
{
|
| 408 |
"id": "ncdhw3d",
|
| 409 |
"priority": 10,
|
|
|
|
| 410 |
"supersededBy": ["global_window_parallel"],
|
| 411 |
-
"when": ["poolBaseOk", "ranks.X == 5", "dim(shapes.Y, 2) == explicitOutD", "dim(shapes.Y, 3) == explicitOutH", "dim(shapes.Y, 4) == explicitOutW"],
|
| 412 |
"passes": [
|
| 413 |
{
|
| 414 |
"id": "main",
|
| 415 |
"name": "AveragePool.Ncdhw3d",
|
| 416 |
-
"
|
| 417 |
-
"
|
| 418 |
-
"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 419 |
}
|
| 420 |
]
|
| 421 |
},
|
| 422 |
{
|
| 423 |
"id": "nchw2d_k7s2_reuse",
|
| 424 |
"priority": 36,
|
| 425 |
-
"when": ["poolBaseOk", "ranks.
|
| 426 |
-
"
|
| 427 |
"passes": [
|
| 428 |
{
|
| 429 |
"id": "main",
|
| 430 |
"name": "AveragePool.Nchw2dK7s2Reuse",
|
| 431 |
-
"
|
| 432 |
-
|
| 433 |
-
"
|
| 434 |
-
|
| 435 |
-
|
| 436 |
-
|
| 437 |
-
|
| 438 |
-
|
| 439 |
-
|
| 440 |
-
|
| 441 |
-
|
| 442 |
-
|
| 443 |
-
|
| 444 |
-
|
| 445 |
-
|
| 446 |
-
|
| 447 |
-
|
| 448 |
-
"count4": "numel(shapes.Y) / 4"
|
| 449 |
-
}
|
| 450 |
},
|
| 451 |
-
"bindings": "
|
| 452 |
-
"dispatch": {
|
|
|
|
|
|
|
|
|
|
|
|
|
| 453 |
}
|
| 454 |
]
|
| 455 |
},
|
| 456 |
{
|
| 457 |
"id": "nchw2d_k7s1_reuse",
|
| 458 |
"priority": 35,
|
| 459 |
-
"when": ["poolBaseOk", "ranks.
|
| 460 |
-
"
|
| 461 |
"passes": [
|
| 462 |
{
|
| 463 |
"id": "main",
|
| 464 |
"name": "AveragePool.Nchw2dK7s1Reuse",
|
| 465 |
-
"
|
| 466 |
-
|
| 467 |
-
"
|
| 468 |
-
|
| 469 |
-
|
| 470 |
-
|
| 471 |
-
|
| 472 |
-
|
| 473 |
-
|
| 474 |
-
|
| 475 |
-
|
| 476 |
-
|
| 477 |
-
|
| 478 |
-
|
| 479 |
-
|
| 480 |
-
|
| 481 |
-
|
| 482 |
-
"count4": "numel(shapes.Y) / 4"
|
| 483 |
-
}
|
| 484 |
},
|
| 485 |
-
"bindings": "
|
| 486 |
-
"dispatch": {
|
|
|
|
|
|
|
|
|
|
|
|
|
| 487 |
}
|
| 488 |
]
|
| 489 |
}
|
|
|
|
| 2 |
"domain": "ai.onnx",
|
| 3 |
"name": "AveragePool",
|
| 4 |
"sinceVersion": 19,
|
| 5 |
+
"inputs": { "x": { "onnx": "X", "dtype": "T" } },
|
| 6 |
+
"outputs": { "y": { "onnx": "Y", "dtype": "T", "rank": "ranks.x", "shape": "poolOutputShape" } },
|
| 7 |
+
"attributes": {
|
| 8 |
+
"auto_pad": { "default": "NOTSET" },
|
| 9 |
+
"count_include_pad": { "default": 0 },
|
| 10 |
+
"ceil_mode": { "default": 0 },
|
| 11 |
+
"kernel_shape": {},
|
| 12 |
+
"strides": {},
|
| 13 |
+
"pads": {},
|
| 14 |
+
"dilations": {}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 15 |
},
|
| 16 |
"attributeConstraints": {
|
| 17 |
"auto_pad": { "values": ["NOTSET", "SAME_UPPER", "SAME_LOWER", "VALID"] },
|
|
|
|
| 20 |
"kernel_shape": { "required": true }
|
| 21 |
},
|
| 22 |
"typeConstraints": { "T": ["float32", "float16"] },
|
| 23 |
+
"tunables": { "WORKGROUP_SIZE": { "default": 256 }, "SEPARABLE_MIN_WINDOW": { "default": 25 } },
|
|
|
|
|
|
|
|
|
|
| 24 |
"derive": {
|
| 25 |
+
"spatialRank": "ranks.x - 2",
|
| 26 |
"narrowSubgroupRange": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize < device.adapterInfo.subgroupMaxSize and device.adapterInfo.subgroupMaxSize <= 16",
|
| 27 |
"kernelShapeLengthOk": "not has(attrs, \"kernel_shape\") or (attrs.kernel_shape | length) == spatialRank",
|
| 28 |
"stridesLengthOk": "not has(attrs, \"strides\") or (attrs.strides | length) == spatialRank",
|
|
|
|
| 43 |
"padBack": "attrs.pads[spatialRank] if padsLengthOk and has(attrs, \"pads\") and spatialRank == 3 else 0",
|
| 44 |
"padBottom": "attrs.pads[2 * spatialRank - 2] if padsLengthOk and has(attrs, \"pads\") and spatialRank >= 2 else 0",
|
| 45 |
"padRight": "attrs.pads[2 * spatialRank - 1] if padsLengthOk and has(attrs, \"pads\") and spatialRank >= 1 else 0",
|
| 46 |
+
"poolBaseOk": "f16Ok(dtypes.T) and ranks.x >= 3 and ranks.x <= 5 and ranks.y == ranks.x and has(attrs, \"kernel_shape\") and kernelShapeLengthOk and stridesLengthOk and dilationsLengthOk and padsLengthOk and kernelD > 0 and kernelH > 0 and kernelW > 0 and strideD > 0 and strideH > 0 and strideW > 0 and dilationD > 0 and dilationH > 0 and dilationW > 0 and padFront >= 0 and padTop >= 0 and padLeft >= 0 and padBack >= 0 and padBottom >= 0 and padRight >= 0 and (attrs.auto_pad == \"NOTSET\" or not has(attrs, \"pads\")) and dim(shapes.y, 0) == dim(shapes.x, 0) and dim(shapes.y, 1) == dim(shapes.x, 1)",
|
| 47 |
+
"explicitOutW": "min(ceil((dim(shapes.x, ranks.x - 1) + padLeft + padRight - ((kernelW - 1) * dilationW + 1)) / strideW) + 1, ceil((dim(shapes.x, ranks.x - 1) + padLeft) / strideW)) if attrs.ceil_mode != 0 else floor((dim(shapes.x, ranks.x - 1) + padLeft + padRight - ((kernelW - 1) * dilationW + 1)) / strideW) + 1",
|
| 48 |
+
"explicitOutH": "min(ceil((dim(shapes.x, ranks.x - 2) + padTop + padBottom - ((kernelH - 1) * dilationH + 1)) / strideH) + 1, ceil((dim(shapes.x, ranks.x - 2) + padTop) / strideH)) if attrs.ceil_mode != 0 else floor((dim(shapes.x, ranks.x - 2) + padTop + padBottom - ((kernelH - 1) * dilationH + 1)) / strideH) + 1",
|
| 49 |
+
"explicitOutD": "(min(ceil((dim(shapes.x, 2) + padFront + padBack - ((kernelD - 1) * dilationD + 1)) / strideD) + 1, ceil((dim(shapes.x, 2) + padFront) / strideD)) if attrs.ceil_mode != 0 else floor((dim(shapes.x, 2) + padFront + padBack - ((kernelD - 1) * dilationD + 1)) / strideD) + 1) if ranks.x == 5 else 0",
|
| 50 |
+
"inferredOutW": "ceil(dim(shapes.x, ranks.x - 1) / strideW) if attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\" else explicitOutW",
|
| 51 |
+
"inferredOutH": "ceil(dim(shapes.x, ranks.x - 2) / strideH) if attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\" else explicitOutH",
|
| 52 |
+
"inferredOutD": "ceil(dim(shapes.x, 2) / strideD) if attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\" else explicitOutD",
|
| 53 |
+
"poolOutputShape": "[dim(shapes.x, 0), dim(shapes.x, 1), inferredOutW] if ranks.x == 3 else ([dim(shapes.x, 0), dim(shapes.x, 1), inferredOutH, inferredOutW] if ranks.x == 4 else [dim(shapes.x, 0), dim(shapes.x, 1), inferredOutD, inferredOutH, inferredOutW])",
|
| 54 |
+
"autoPadTotalD": "max(0, (dim(shapes.y, 2) - 1) * strideD + ((kernelD - 1) * dilationD + 1) - dim(shapes.x, 2)) if ranks.x == 5 else 0",
|
| 55 |
+
"autoPadTotalH": "max(0, (dim(shapes.y, ranks.x - 2) - 1) * strideH + ((kernelH - 1) * dilationH + 1) - dim(shapes.x, ranks.x - 2)) if ranks.x >= 4 else 0",
|
| 56 |
+
"autoPadTotalW": "max(0, (dim(shapes.y, ranks.x - 1) - 1) * strideW + ((kernelW - 1) * dilationW + 1) - dim(shapes.x, ranks.x - 1))",
|
| 57 |
+
"effectivePadFront": "(floor(autoPadTotalD / 2) if attrs.auto_pad == \"SAME_UPPER\" else autoPadTotalD - floor(autoPadTotalD / 2)) if (attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") else padFront",
|
| 58 |
+
"effectivePadBack": "autoPadTotalD - effectivePadFront if (attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") else padBack",
|
| 59 |
+
"effectivePadTop": "(floor(autoPadTotalH / 2) if attrs.auto_pad == \"SAME_UPPER\" else autoPadTotalH - floor(autoPadTotalH / 2)) if (attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") else padTop",
|
| 60 |
+
"effectivePadLeft": "(floor(autoPadTotalW / 2) if attrs.auto_pad == \"SAME_UPPER\" else autoPadTotalW - floor(autoPadTotalW / 2)) if (attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") else padLeft",
|
| 61 |
+
"effectivePadBottom": "autoPadTotalH - effectivePadTop if (attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") else padBottom",
|
| 62 |
+
"effectivePadRight": "autoPadTotalW - effectivePadLeft if (attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") else padRight",
|
| 63 |
"narrowF32W3UnrolledPreferred": "narrowSubgroupRange and dtypes.T == \"f32\"",
|
| 64 |
+
"globalWindowParallelCovered": "poolBaseOk and numel(shapes.y) == dim(shapes.x, 0) * dim(shapes.x, 1) and numel(shapes.y) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and attrs.auto_pad == \"NOTSET\" and strideD == 1 and strideH == 1 and strideW == 1 and dilationD == 1 and dilationH == 1 and dilationW == 1 and padFront == 0 and padTop == 0 and padLeft == 0 and padBack == 0 and padBottom == 0 and padRight == 0 and ((ranks.x == 3 and kernelW == dim(shapes.x, 2)) or (ranks.x == 4 and kernelH == dim(shapes.x, 2) and kernelW == dim(shapes.x, 3)) or (ranks.x == 5 and kernelD == dim(shapes.x, 2) and kernelH == dim(shapes.x, 3) and kernelW == dim(shapes.x, 4))) and inner(shapes.x, 1) >= 256",
|
| 65 |
+
"separableTmpElements": "dim(shapes.x, 0) * dim(shapes.x, 1) * dim(shapes.x, 2) * dim(shapes.y, 3)",
|
| 66 |
+
"separableFits": "separableTmpElements * 4 <= min(device.limits.maxStorageBufferBindingSize, device.limits.maxBufferSize) and ceilDiv(separableTmpElements, tunables.WORKGROUP_SIZE) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
|
| 67 |
+
"wholeWindowDenominator": "attrs.ceil_mode == 0 and (attrs.count_include_pad == 1 or (effectivePadTop == 0 and effectivePadLeft == 0 and effectivePadBottom == 0 and effectivePadRight == 0))",
|
| 68 |
+
"scalar": "dtypes.T"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 69 |
},
|
| 70 |
+
"bindings": { "y_2": { "name": "y", "buffer": "storage", "elementType": "$outputElement" } },
|
| 71 |
"variants": [
|
| 72 |
{
|
| 73 |
"id": "global_window_parallel",
|
| 74 |
"priority": 50,
|
| 75 |
"when": ["globalWindowParallelCovered"],
|
| 76 |
+
"derive": {
|
| 77 |
+
"xElem": "(\"vec4<f16>\" if dtypes.T == \"f16\" else \"vec4<f32>\") if inner(shapes.x, 1) % 4 == 0 else dtypes.T"
|
| 78 |
},
|
| 79 |
"passes": [
|
| 80 |
{
|
| 81 |
"id": "main",
|
| 82 |
"name": "AveragePool.GlobalWindowParallel",
|
| 83 |
+
"shader": "pool-global-reduction.wgsl.jinja",
|
| 84 |
+
"subgroupCollectivesWidth": "portable",
|
| 85 |
+
"derive": {
|
| 86 |
+
"op": "\"average\"",
|
| 87 |
+
"vec4": "inner(shapes.x, 1) % 4 == 0",
|
| 88 |
+
"workgroupSizeSpec": "tunables.WORKGROUP_SIZE",
|
| 89 |
+
"useSubgroups": "device.features.has(\"subgroups\")"
|
|
|
|
| 90 |
},
|
| 91 |
+
"bindings": [
|
| 92 |
+
{ "arg": "x", "elementType": "$xElem" },
|
| 93 |
+
{ "arg": "y", "elementType": "$scalar" },
|
| 94 |
+
{ "name": "params", "struct": [{ "name": "spatial", "type": "u32", "value": "inner(shapes.x, 1)" }] }
|
| 95 |
+
],
|
| 96 |
+
"dispatch": { "x": "numel(shapes.y)" }
|
| 97 |
}
|
| 98 |
]
|
| 99 |
},
|
| 100 |
{
|
| 101 |
"id": "nchw2d_k7s1_reuse_tail",
|
| 102 |
"priority": 34,
|
| 103 |
+
"when": ["poolBaseOk", "ranks.x == 4", "kernelH == 7", "kernelW == 7", "strideH == 1", "strideW == 1", "dilationH == 1", "dilationW == 1", "effectivePadTop == 3", "effectivePadBottom == 3", "effectivePadLeft == 3", "effectivePadRight == 3", "attrs.count_include_pad == 0", "attrs.ceil_mode == 0", "dim(shapes.x, 3) % 4 != 0", "dim(shapes.y, 2) == dim(shapes.x, 2)", "dim(shapes.y, 3) == dim(shapes.x, 3)", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) > 0"],
|
| 104 |
+
"derive": { "outputElement": "dtypes.T" },
|
| 105 |
"passes": [
|
| 106 |
{
|
| 107 |
"id": "main",
|
| 108 |
"name": "AveragePool.Nchw2dK7s1ReuseTail",
|
| 109 |
+
"shader": "average-pool2d-nchw-horizontal-reuse.wgsl.jinja",
|
| 110 |
+
"derive": {
|
| 111 |
+
"kernelHSpec": 7,
|
| 112 |
+
"kernelWSpec": 7,
|
| 113 |
+
"strideHSpec": 1,
|
| 114 |
+
"strideWSpec": 1,
|
| 115 |
+
"dilationHSpec": 1,
|
| 116 |
+
"dilationWSpec": 1,
|
| 117 |
+
"padTopSpec": 3,
|
| 118 |
+
"padLeftSpec": 3,
|
| 119 |
+
"inH": "dim(shapes.x, 2)",
|
| 120 |
+
"inW": "dim(shapes.x, 3)",
|
| 121 |
+
"inPlane": "dim(shapes.x, 2) * dim(shapes.x, 3)",
|
| 122 |
+
"outH": "dim(shapes.y, 2)",
|
| 123 |
+
"outW": "dim(shapes.y, 3)",
|
| 124 |
+
"outW4": "ceilDiv(dim(shapes.y, 3), 4)",
|
| 125 |
+
"count4": "dim(shapes.y, 0) * dim(shapes.y, 1) * dim(shapes.y, 2) * ceilDiv(dim(shapes.y, 3), 4)",
|
| 126 |
+
"scalarOutput": true
|
|
|
|
|
|
|
| 127 |
},
|
| 128 |
+
"bindings": ["x", "y_2"],
|
| 129 |
"dispatch": {
|
| 130 |
+
"x": "min(ceilDiv((dim(shapes.y, 0) * dim(shapes.y, 1) * dim(shapes.y, 2) * ceilDiv(dim(shapes.y, 3), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 131 |
+
"y": "ceilDiv(ceilDiv((dim(shapes.y, 0) * dim(shapes.y, 1) * dim(shapes.y, 2) * ceilDiv(dim(shapes.y, 3), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 132 |
+
"z": 1
|
| 133 |
}
|
| 134 |
}
|
| 135 |
]
|
|
|
|
| 137 |
{
|
| 138 |
"id": "nchw2d_w3s1_reuse",
|
| 139 |
"priority": 30,
|
| 140 |
+
"when": ["poolBaseOk", "ranks.x == 4", "kernelW == 3", "strideW == 1", "dilationW == 1", "effectivePadLeft == 1", "effectivePadRight == 1", "kernelH >= 1", "kernelH <= 3", "strideH == 1", "dilationH >= 1", "attrs.count_include_pad == 0", "attrs.ceil_mode == 0", "dim(shapes.x, 3) % 4 == 0", "dim(shapes.y, 2) == dim(shapes.x, 2) + effectivePadTop + effectivePadBottom - ((kernelH - 1) * dilationH + 1) + 1", "dim(shapes.y, 3) == dim(shapes.x, 3)", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) > 0"],
|
| 141 |
"demoteWhen": ["narrowF32W3UnrolledPreferred"],
|
|
|
|
| 142 |
"passes": [
|
| 143 |
{
|
| 144 |
"id": "main",
|
| 145 |
"name": "AveragePool.Nchw2dW3s1Reuse",
|
| 146 |
+
"shader": "average-pool2d-nchw-w3s1-reuse.wgsl.jinja",
|
| 147 |
+
"derive": {
|
| 148 |
+
"kernelHSpec": "kernelH",
|
| 149 |
+
"dilationHSpec": "dilationH",
|
| 150 |
+
"padTopSpec": "effectivePadTop",
|
| 151 |
+
"inH": "dim(shapes.x, 2)",
|
| 152 |
+
"inW": "dim(shapes.x, 3)",
|
| 153 |
+
"inPlane": "dim(shapes.x, 2) * dim(shapes.x, 3)",
|
| 154 |
+
"outH": "dim(shapes.y, 2)",
|
| 155 |
+
"outW": "dim(shapes.y, 3)",
|
| 156 |
+
"outWQuarter": "floor(dim(shapes.y, 3) / 4)"
|
|
|
|
|
|
|
| 157 |
},
|
| 158 |
+
"bindings": [
|
| 159 |
+
"x",
|
| 160 |
+
"y",
|
| 161 |
+
{ "name": "params", "struct": [{ "name": "tileCount", "type": "u32", "value": "numel(shapes.y) / 4" }] }
|
| 162 |
+
],
|
| 163 |
+
"dispatch": {
|
| 164 |
+
"x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 165 |
+
"y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 166 |
+
"z": 1
|
| 167 |
+
}
|
| 168 |
}
|
| 169 |
]
|
| 170 |
},
|
| 171 |
{
|
| 172 |
"id": "nchw2d_w3s1_reuse_tail",
|
| 173 |
"priority": 29,
|
| 174 |
+
"when": ["poolBaseOk", "ranks.x == 4", "kernelW == 3", "strideW == 1", "dilationW == 1", "effectivePadLeft == 1", "effectivePadRight == 1", "kernelH >= 1", "kernelH <= 3", "strideH == 1", "dilationH >= 1", "attrs.count_include_pad == 0", "attrs.ceil_mode == 0", "dim(shapes.x, 3) % 4 != 0", "dim(shapes.y, 2) == dim(shapes.x, 2) + effectivePadTop + effectivePadBottom - ((kernelH - 1) * dilationH + 1) + 1", "dim(shapes.y, 3) == dim(shapes.x, 3)", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) > 0"],
|
| 175 |
+
"derive": { "outputElement": "dtypes.T" },
|
| 176 |
"passes": [
|
| 177 |
{
|
| 178 |
"id": "main",
|
| 179 |
"name": "AveragePool.Nchw2dW3s1ReuseTail",
|
| 180 |
+
"shader": "average-pool2d-nchw-horizontal-reuse.wgsl.jinja",
|
| 181 |
+
"derive": {
|
| 182 |
+
"kernelHSpec": "kernelH",
|
| 183 |
+
"kernelWSpec": 3,
|
| 184 |
+
"strideHSpec": 1,
|
| 185 |
+
"strideWSpec": 1,
|
| 186 |
+
"dilationHSpec": "dilationH",
|
| 187 |
+
"dilationWSpec": 1,
|
| 188 |
+
"padTopSpec": "effectivePadTop",
|
| 189 |
+
"padLeftSpec": 1,
|
| 190 |
+
"inH": "dim(shapes.x, 2)",
|
| 191 |
+
"inW": "dim(shapes.x, 3)",
|
| 192 |
+
"inPlane": "dim(shapes.x, 2) * dim(shapes.x, 3)",
|
| 193 |
+
"outH": "dim(shapes.y, 2)",
|
| 194 |
+
"outW": "dim(shapes.y, 3)",
|
| 195 |
+
"outW4": "ceilDiv(dim(shapes.y, 3), 4)",
|
| 196 |
+
"count4": "dim(shapes.y, 0) * dim(shapes.y, 1) * dim(shapes.y, 2) * ceilDiv(dim(shapes.y, 3), 4)",
|
| 197 |
+
"scalarOutput": true
|
|
|
|
|
|
|
| 198 |
},
|
| 199 |
+
"bindings": ["x", "y_2"],
|
| 200 |
"dispatch": {
|
| 201 |
+
"x": "min(ceilDiv((dim(shapes.y, 0) * dim(shapes.y, 1) * dim(shapes.y, 2) * ceilDiv(dim(shapes.y, 3), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 202 |
+
"y": "ceilDiv(ceilDiv((dim(shapes.y, 0) * dim(shapes.y, 1) * dim(shapes.y, 2) * ceilDiv(dim(shapes.y, 3), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 203 |
+
"z": 1
|
| 204 |
+
}
|
| 205 |
+
}
|
| 206 |
+
]
|
| 207 |
+
},
|
| 208 |
+
{
|
| 209 |
+
"id": "nchw2d_separable",
|
| 210 |
+
"priority": 22,
|
| 211 |
+
"when": ["poolBaseOk", "ranks.x == 4", "ranks.y == 4", "dim(shapes.y, 2) == inferredOutH", "dim(shapes.y, 3) == inferredOutW", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) > 0", "kernelH > 1", "kernelW > 1", "kernelH * kernelW >= tunables.SEPARABLE_MIN_WINDOW", "strideH < kernelH", "wholeWindowDenominator", "separableFits"],
|
| 212 |
+
"derive": { "scalar": "dtypes.T", "pIs1": false, "pIs2": false },
|
| 213 |
+
"intermediates": [{ "id": "tmp", "dtype": "float32", "shape": "[separableTmpElements]" }],
|
| 214 |
+
"passes": [
|
| 215 |
+
{
|
| 216 |
+
"id": "reduce_w",
|
| 217 |
+
"name": "AveragePool.SeparableW",
|
| 218 |
+
"shader": "pool2d-nchw-separable.wgsl.jinja",
|
| 219 |
+
"derive": { "axisSpec": "\"w\"", "op": "\"avg\"", "windowArea": "kernelH * kernelW" },
|
| 220 |
+
"bindings": [
|
| 221 |
+
"x",
|
| 222 |
+
{ "name": "tmp", "buffer": "storage", "elementType": "f32" },
|
| 223 |
+
{
|
| 224 |
+
"name": "params",
|
| 225 |
+
"struct": [
|
| 226 |
+
{ "name": "inH", "type": "u32", "value": "dim(shapes.x, 2)" },
|
| 227 |
+
{ "name": "inW", "type": "u32", "value": "dim(shapes.x, 3)" },
|
| 228 |
+
{ "name": "outW", "type": "u32", "value": "dim(shapes.y, 3)" },
|
| 229 |
+
{ "name": "kernelW", "type": "u32", "value": "kernelW" },
|
| 230 |
+
{ "name": "strideW", "type": "u32", "value": "strideW" },
|
| 231 |
+
{ "name": "dilationW", "type": "u32", "value": "dilationW" },
|
| 232 |
+
{ "name": "padW", "type": "i32", "value": "effectivePadLeft" },
|
| 233 |
+
{ "name": "tmpCount", "type": "u32", "value": "separableTmpElements" }
|
| 234 |
+
]
|
| 235 |
+
}
|
| 236 |
+
],
|
| 237 |
+
"dispatch": {
|
| 238 |
+
"x": "min(ceilDiv((separableTmpElements), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 239 |
+
"y": "ceilDiv(ceilDiv((separableTmpElements), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 240 |
+
"z": 1
|
| 241 |
+
}
|
| 242 |
+
},
|
| 243 |
+
{
|
| 244 |
+
"id": "reduce_h",
|
| 245 |
+
"name": "AveragePool.SeparableH",
|
| 246 |
+
"shader": "pool2d-nchw-separable.wgsl.jinja",
|
| 247 |
+
"derive": { "axisSpec": "\"h\"", "op": "\"avg\"", "windowArea": "kernelH * kernelW" },
|
| 248 |
+
"bindings": [
|
| 249 |
+
{ "name": "tmp", "buffer": "read-only-storage", "elementType": "f32" },
|
| 250 |
+
"y",
|
| 251 |
+
{
|
| 252 |
+
"name": "params",
|
| 253 |
+
"struct": [
|
| 254 |
+
{ "name": "inH", "type": "u32", "value": "dim(shapes.x, 2)" },
|
| 255 |
+
{ "name": "outH", "type": "u32", "value": "dim(shapes.y, 2)" },
|
| 256 |
+
{ "name": "outW", "type": "u32", "value": "dim(shapes.y, 3)" },
|
| 257 |
+
{ "name": "kernelH", "type": "u32", "value": "kernelH" },
|
| 258 |
+
{ "name": "strideH", "type": "u32", "value": "strideH" },
|
| 259 |
+
{ "name": "dilationH", "type": "u32", "value": "dilationH" },
|
| 260 |
+
{ "name": "padH", "type": "i32", "value": "effectivePadTop" },
|
| 261 |
+
{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }
|
| 262 |
+
]
|
| 263 |
+
}
|
| 264 |
+
],
|
| 265 |
+
"dispatch": {
|
| 266 |
+
"x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 267 |
+
"y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 268 |
+
"z": 1
|
| 269 |
}
|
| 270 |
}
|
| 271 |
]
|
|
|
|
| 273 |
{
|
| 274 |
"id": "nchw2d_unrolled",
|
| 275 |
"priority": 25,
|
| 276 |
+
"when": ["poolBaseOk", "ranks.x == 4", "dim(shapes.y, 2) == inferredOutH", "dim(shapes.y, 3) == inferredOutW", "kernelH >= 1", "kernelH <= 3", "kernelW >= 1", "kernelW <= 3", "strideH >= 1", "strideW >= 1", "dilationH >= 1", "dilationW >= 1", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) > 0"],
|
| 277 |
"passes": [
|
| 278 |
{
|
| 279 |
"id": "main",
|
| 280 |
"name": "AveragePool.Nchw2dUnrolled",
|
| 281 |
+
"shader": "pool-window-unroll.wgsl.jinja",
|
| 282 |
+
"derive": {
|
| 283 |
+
"op": "\"average\"",
|
| 284 |
+
"kernelHSpec": "kernelH",
|
| 285 |
+
"kernelWSpec": "kernelW",
|
| 286 |
+
"strideHSpec": "strideH",
|
| 287 |
+
"strideWSpec": "strideW",
|
| 288 |
+
"dilationHSpec": "dilationH",
|
| 289 |
+
"dilationWSpec": "dilationW",
|
| 290 |
+
"padTopSpec": "effectivePadTop",
|
| 291 |
+
"padLeftSpec": "effectivePadLeft",
|
| 292 |
+
"inH": "dim(shapes.x, 2)",
|
| 293 |
+
"inW": "dim(shapes.x, 3)",
|
| 294 |
+
"outH": "dim(shapes.y, 2)",
|
| 295 |
+
"outW": "dim(shapes.y, 3)",
|
| 296 |
+
"padBottomSpec": "effectivePadBottom",
|
| 297 |
+
"padRightSpec": "effectivePadRight",
|
| 298 |
+
"countIncludePad": "attrs.count_include_pad != 0"
|
|
|
|
|
|
|
| 299 |
},
|
| 300 |
+
"bindings": [
|
| 301 |
+
"x",
|
| 302 |
+
"y",
|
| 303 |
+
{ "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }] }
|
| 304 |
+
],
|
| 305 |
+
"dispatch": {
|
| 306 |
+
"x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 307 |
+
"y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 308 |
+
"z": 1
|
| 309 |
+
}
|
| 310 |
}
|
| 311 |
]
|
| 312 |
},
|
| 313 |
{
|
| 314 |
"id": "ncl1d",
|
| 315 |
"priority": 10,
|
| 316 |
+
"when": ["poolBaseOk", "ranks.x == 3", "dim(shapes.y, 2) == inferredOutW"],
|
| 317 |
"passes": [
|
| 318 |
{
|
| 319 |
"id": "main",
|
| 320 |
"name": "AveragePool.Ncl1d",
|
| 321 |
+
"shader": "pool-window-nd.wgsl.jinja",
|
| 322 |
+
"derive": { "op": "\"average\"", "spatialRankSpec": 1 },
|
| 323 |
+
"bindings": [
|
| 324 |
+
"x",
|
| 325 |
+
"y",
|
| 326 |
+
{
|
| 327 |
+
"name": "params",
|
| 328 |
+
"struct": [
|
| 329 |
+
{ "name": "channels", "type": "u32", "value": "dim(shapes.x, 1)" },
|
| 330 |
+
{ "name": "inW", "type": "u32", "value": "dim(shapes.x, 2)" },
|
| 331 |
+
{ "name": "outW", "type": "u32", "value": "dim(shapes.y, 2)" },
|
| 332 |
+
{ "name": "kernelW", "type": "u32", "value": "kernelW" },
|
| 333 |
+
{ "name": "strideW", "type": "u32", "value": "strideW" },
|
| 334 |
+
{ "name": "dilationW", "type": "u32", "value": "dilationW" },
|
| 335 |
+
{ "name": "padLeft", "type": "i32", "value": "effectivePadLeft" },
|
| 336 |
+
{ "name": "padRight", "type": "i32", "value": "effectivePadRight" },
|
| 337 |
+
{ "name": "countIncludePad", "type": "u32", "value": "attrs.count_include_pad" },
|
| 338 |
+
{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }
|
| 339 |
+
]
|
| 340 |
+
}
|
| 341 |
+
],
|
| 342 |
+
"dispatch": {
|
| 343 |
+
"x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 344 |
+
"y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 345 |
+
"z": 1
|
| 346 |
+
}
|
| 347 |
}
|
| 348 |
]
|
| 349 |
},
|
| 350 |
{
|
| 351 |
"id": "nchw2d",
|
| 352 |
+
"when": ["poolBaseOk", "ranks.x == 4", "((dim(shapes.y, 2) == explicitOutH and dim(shapes.y, 3) == explicitOutW) or ((attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") and dim(shapes.y, 2) == ceil(dim(shapes.x, 2) / strideH) and dim(shapes.y, 3) == ceil(dim(shapes.x, 3) / strideW)))"],
|
| 353 |
"passes": [
|
| 354 |
{
|
| 355 |
"id": "main",
|
| 356 |
"name": "nchw2d",
|
| 357 |
+
"shader": "pool-window-nd.wgsl.jinja",
|
| 358 |
+
"derive": { "op": "\"average\"", "spatialRankSpec": 2 },
|
| 359 |
+
"bindings": [
|
| 360 |
+
"x",
|
| 361 |
+
"y",
|
| 362 |
+
{
|
| 363 |
+
"name": "params",
|
| 364 |
+
"struct": [
|
| 365 |
+
{ "name": "channels", "type": "u32", "value": "dim(shapes.x, 1)" },
|
| 366 |
+
{ "name": "inH", "type": "u32", "value": "dim(shapes.x, 2)" },
|
| 367 |
+
{ "name": "inW", "type": "u32", "value": "dim(shapes.x, 3)" },
|
| 368 |
+
{ "name": "outH", "type": "u32", "value": "dim(shapes.y, 2)" },
|
| 369 |
+
{ "name": "outW", "type": "u32", "value": "dim(shapes.y, 3)" },
|
| 370 |
+
{ "name": "kernelH", "type": "u32", "value": "kernelH" },
|
| 371 |
+
{ "name": "kernelW", "type": "u32", "value": "kernelW" },
|
| 372 |
+
{ "name": "strideH", "type": "u32", "value": "strideH" },
|
| 373 |
+
{ "name": "strideW", "type": "u32", "value": "strideW" },
|
| 374 |
+
{ "name": "dilationH", "type": "u32", "value": "dilationH" },
|
| 375 |
+
{ "name": "dilationW", "type": "u32", "value": "dilationW" },
|
| 376 |
+
{ "name": "padTop", "type": "i32", "value": "effectivePadTop" },
|
| 377 |
+
{ "name": "padLeft", "type": "i32", "value": "effectivePadLeft" },
|
| 378 |
+
{ "name": "padBottom", "type": "i32", "value": "effectivePadBottom" },
|
| 379 |
+
{ "name": "padRight", "type": "i32", "value": "effectivePadRight" },
|
| 380 |
+
{ "name": "countIncludePad", "type": "u32", "value": "attrs.count_include_pad" },
|
| 381 |
+
{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }
|
| 382 |
+
]
|
| 383 |
+
}
|
| 384 |
+
],
|
| 385 |
+
"dispatch": {
|
| 386 |
+
"x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 387 |
+
"y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 388 |
+
"z": 1
|
| 389 |
+
}
|
| 390 |
}
|
| 391 |
]
|
| 392 |
},
|
| 393 |
{
|
| 394 |
"id": "ncdhw3d",
|
| 395 |
"priority": 10,
|
| 396 |
+
"when": ["poolBaseOk", "ranks.x == 5", "dim(shapes.y, 2) == inferredOutD", "dim(shapes.y, 3) == inferredOutH", "dim(shapes.y, 4) == inferredOutW"],
|
| 397 |
"supersededBy": ["global_window_parallel"],
|
|
|
|
| 398 |
"passes": [
|
| 399 |
{
|
| 400 |
"id": "main",
|
| 401 |
"name": "AveragePool.Ncdhw3d",
|
| 402 |
+
"shader": "pool-window-nd.wgsl.jinja",
|
| 403 |
+
"derive": { "op": "\"average\"", "spatialRankSpec": 3 },
|
| 404 |
+
"bindings": [
|
| 405 |
+
"x",
|
| 406 |
+
"y",
|
| 407 |
+
{
|
| 408 |
+
"name": "params",
|
| 409 |
+
"struct": [
|
| 410 |
+
{ "name": "channels", "type": "u32", "value": "dim(shapes.x, 1)" },
|
| 411 |
+
{ "name": "inD", "type": "u32", "value": "dim(shapes.x, 2)" },
|
| 412 |
+
{ "name": "inH", "type": "u32", "value": "dim(shapes.x, 3)" },
|
| 413 |
+
{ "name": "inW", "type": "u32", "value": "dim(shapes.x, 4)" },
|
| 414 |
+
{ "name": "outD", "type": "u32", "value": "dim(shapes.y, 2)" },
|
| 415 |
+
{ "name": "outH", "type": "u32", "value": "dim(shapes.y, 3)" },
|
| 416 |
+
{ "name": "outW", "type": "u32", "value": "dim(shapes.y, 4)" },
|
| 417 |
+
{ "name": "kernelD", "type": "u32", "value": "kernelD" },
|
| 418 |
+
{ "name": "kernelH", "type": "u32", "value": "kernelH" },
|
| 419 |
+
{ "name": "kernelW", "type": "u32", "value": "kernelW" },
|
| 420 |
+
{ "name": "strideD", "type": "u32", "value": "strideD" },
|
| 421 |
+
{ "name": "strideH", "type": "u32", "value": "strideH" },
|
| 422 |
+
{ "name": "strideW", "type": "u32", "value": "strideW" },
|
| 423 |
+
{ "name": "dilationD", "type": "u32", "value": "dilationD" },
|
| 424 |
+
{ "name": "dilationH", "type": "u32", "value": "dilationH" },
|
| 425 |
+
{ "name": "dilationW", "type": "u32", "value": "dilationW" },
|
| 426 |
+
{ "name": "padFront", "type": "i32", "value": "effectivePadFront" },
|
| 427 |
+
{ "name": "padTop", "type": "i32", "value": "effectivePadTop" },
|
| 428 |
+
{ "name": "padLeft", "type": "i32", "value": "effectivePadLeft" },
|
| 429 |
+
{ "name": "padBack", "type": "i32", "value": "effectivePadBack" },
|
| 430 |
+
{ "name": "padBottom", "type": "i32", "value": "effectivePadBottom" },
|
| 431 |
+
{ "name": "padRight", "type": "i32", "value": "effectivePadRight" },
|
| 432 |
+
{ "name": "countIncludePad", "type": "u32", "value": "attrs.count_include_pad" },
|
| 433 |
+
{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }
|
| 434 |
+
]
|
| 435 |
+
}
|
| 436 |
+
],
|
| 437 |
+
"dispatch": {
|
| 438 |
+
"x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 439 |
+
"y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 440 |
+
"z": 1
|
| 441 |
+
}
|
| 442 |
}
|
| 443 |
]
|
| 444 |
},
|
| 445 |
{
|
| 446 |
"id": "nchw2d_k7s2_reuse",
|
| 447 |
"priority": 36,
|
| 448 |
+
"when": ["poolBaseOk", "ranks.x == 4", "kernelH == 7", "kernelW == 7", "strideH == 2", "strideW == 2", "dilationH == 1", "dilationW == 1", "effectivePadTop == 3", "effectivePadBottom == 3", "effectivePadLeft == 3", "effectivePadRight == 3", "attrs.count_include_pad == 0", "attrs.ceil_mode == 0", "dim(shapes.y, 2) == inferredOutH", "dim(shapes.y, 3) == inferredOutW", "dim(shapes.y, 3) % 4 == 0", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) > 0"],
|
| 449 |
+
"derive": { "outputElement": "\"vec4<\" ~ dtypes.T ~ \">\"", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
|
| 450 |
"passes": [
|
| 451 |
{
|
| 452 |
"id": "main",
|
| 453 |
"name": "AveragePool.Nchw2dK7s2Reuse",
|
| 454 |
+
"shader": "average-pool2d-nchw-horizontal-reuse.wgsl.jinja",
|
| 455 |
+
"derive": {
|
| 456 |
+
"kernelHSpec": 7,
|
| 457 |
+
"kernelWSpec": 7,
|
| 458 |
+
"strideHSpec": 2,
|
| 459 |
+
"strideWSpec": 2,
|
| 460 |
+
"dilationHSpec": 1,
|
| 461 |
+
"dilationWSpec": 1,
|
| 462 |
+
"padTopSpec": 3,
|
| 463 |
+
"padLeftSpec": 3,
|
| 464 |
+
"inH": "dim(shapes.x, 2)",
|
| 465 |
+
"inW": "dim(shapes.x, 3)",
|
| 466 |
+
"inPlane": "dim(shapes.x, 2) * dim(shapes.x, 3)",
|
| 467 |
+
"outH": "dim(shapes.y, 2)",
|
| 468 |
+
"outW": "dim(shapes.y, 3)",
|
| 469 |
+
"outW4": "dim(shapes.y, 3) / 4",
|
| 470 |
+
"count4": "numel(shapes.y) / 4"
|
|
|
|
|
|
|
| 471 |
},
|
| 472 |
+
"bindings": ["x", "y_2"],
|
| 473 |
+
"dispatch": {
|
| 474 |
+
"x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 475 |
+
"y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 476 |
+
"z": 1
|
| 477 |
+
}
|
| 478 |
}
|
| 479 |
]
|
| 480 |
},
|
| 481 |
{
|
| 482 |
"id": "nchw2d_k7s1_reuse",
|
| 483 |
"priority": 35,
|
| 484 |
+
"when": ["poolBaseOk", "ranks.x == 4", "kernelH == 7", "kernelW == 7", "strideH == 1", "strideW == 1", "dilationH == 1", "dilationW == 1", "effectivePadTop == 3", "effectivePadBottom == 3", "effectivePadLeft == 3", "effectivePadRight == 3", "attrs.count_include_pad == 0", "attrs.ceil_mode == 0", "dim(shapes.x, 3) % 4 == 0", "dim(shapes.y, 2) == dim(shapes.x, 2)", "dim(shapes.y, 3) == dim(shapes.x, 3)", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) > 0"],
|
| 485 |
+
"derive": { "outputElement": "\"vec4<\" ~ dtypes.T ~ \">\"", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
|
| 486 |
"passes": [
|
| 487 |
{
|
| 488 |
"id": "main",
|
| 489 |
"name": "AveragePool.Nchw2dK7s1Reuse",
|
| 490 |
+
"shader": "average-pool2d-nchw-horizontal-reuse.wgsl.jinja",
|
| 491 |
+
"derive": {
|
| 492 |
+
"kernelHSpec": 7,
|
| 493 |
+
"kernelWSpec": 7,
|
| 494 |
+
"strideHSpec": 1,
|
| 495 |
+
"strideWSpec": 1,
|
| 496 |
+
"dilationHSpec": 1,
|
| 497 |
+
"dilationWSpec": 1,
|
| 498 |
+
"padTopSpec": 3,
|
| 499 |
+
"padLeftSpec": 3,
|
| 500 |
+
"inH": "dim(shapes.x, 2)",
|
| 501 |
+
"inW": "dim(shapes.x, 3)",
|
| 502 |
+
"inPlane": "dim(shapes.x, 2) * dim(shapes.x, 3)",
|
| 503 |
+
"outH": "dim(shapes.y, 2)",
|
| 504 |
+
"outW": "dim(shapes.y, 3)",
|
| 505 |
+
"outW4": "dim(shapes.y, 3) / 4",
|
| 506 |
+
"count4": "numel(shapes.y) / 4"
|
|
|
|
|
|
|
| 507 |
},
|
| 508 |
+
"bindings": ["x", "y_2"],
|
| 509 |
+
"dispatch": {
|
| 510 |
+
"x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 511 |
+
"y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
|
| 512 |
+
"z": 1
|
| 513 |
+
}
|
| 514 |
}
|
| 515 |
]
|
| 516 |
}
|
build/webgpu/metadata.json
CHANGED
|
@@ -1,22 +1,38 @@
|
|
| 1 |
{
|
| 2 |
"name": "ai.onnx.AveragePool",
|
| 3 |
-
"id": "
|
| 4 |
"version": 1,
|
| 5 |
"license": "Apache-2.0",
|
| 6 |
"backend": { "type": "webgpu" },
|
| 7 |
"digest": {
|
| 8 |
"algorithm": "sha256",
|
| 9 |
"files": {
|
| 10 |
-
"average-pool2d-nchw-horizontal-reuse.wgsl.jinja": "
|
| 11 |
-
"average-pool2d-nchw-w3s1-reuse.wgsl.jinja": "
|
| 12 |
-
"bench.json": "
|
| 13 |
-
"manifest.json": "
|
| 14 |
-
"pool-global-reduction.wgsl.jinja": "
|
| 15 |
-
"pool-window-nd.wgsl.jinja": "
|
| 16 |
-
"pool-window-unroll.wgsl.jinja": "
|
| 17 |
-
"
|
|
|
|
| 18 |
}
|
| 19 |
},
|
| 20 |
-
"provenance": { "kernel": { "sha": "
|
| 21 |
-
"webgpu": {
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 22 |
}
|
|
|
|
| 1 |
{
|
| 2 |
"name": "ai.onnx.AveragePool",
|
| 3 |
+
"id": "_ai_onnx_averagepool_webgpu_1a52fa9",
|
| 4 |
"version": 1,
|
| 5 |
"license": "Apache-2.0",
|
| 6 |
"backend": { "type": "webgpu" },
|
| 7 |
"digest": {
|
| 8 |
"algorithm": "sha256",
|
| 9 |
"files": {
|
| 10 |
+
"average-pool2d-nchw-horizontal-reuse.wgsl.jinja": "WhpEal2kC3XudqPetAZeL6v/IGK1olc+Fb4HzzDj890=",
|
| 11 |
+
"average-pool2d-nchw-w3s1-reuse.wgsl.jinja": "IsI2gMcZbY0eemYs5z0pTC86cgaDw605Xfn2iQDBJkw=",
|
| 12 |
+
"bench.json": "7vXmd7ZA5xBCZjqXJYtMzO7vq/0qgzf5gNZ1S6J5Zps=",
|
| 13 |
+
"manifest.json": "a0sXaOaq8613yIAtQBB1R+xrbvA48WTKgkPo14ruXhA=",
|
| 14 |
+
"pool-global-reduction.wgsl.jinja": "SKnTSUq9nU838U0t7VyvC7hgz4h5L+Ndv1rQCM9Vks8=",
|
| 15 |
+
"pool-window-nd.wgsl.jinja": "2shuAamdoFjQCOq2LK2XC6pKsa10zad6bDhqY1h2zBI=",
|
| 16 |
+
"pool-window-unroll.wgsl.jinja": "iSbZ3lQkW4r/CZzsBitQBpHgBLRJ83MekQxyVqbInwI=",
|
| 17 |
+
"pool2d-nchw-separable.wgsl.jinja": "GgafQWV6Pn5SqSn5gMns0EJEHGqk2WkdCdVurt3MgA8=",
|
| 18 |
+
"test.json": "2Xnd9YqkiREFwBJdjyqFIbqnT2KGPNdXYpCezSLDXkA="
|
| 19 |
}
|
| 20 |
},
|
| 21 |
+
"provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
|
| 22 |
+
"webgpu": {
|
| 23 |
+
"manifestSpec": "2.0",
|
| 24 |
+
"variants": {
|
| 25 |
+
"global_window_parallel": ["pool-global-reduction.wgsl.jinja"],
|
| 26 |
+
"nchw2d_k7s1_reuse_tail": ["average-pool2d-nchw-horizontal-reuse.wgsl.jinja"],
|
| 27 |
+
"nchw2d_w3s1_reuse": ["average-pool2d-nchw-w3s1-reuse.wgsl.jinja"],
|
| 28 |
+
"nchw2d_w3s1_reuse_tail": ["average-pool2d-nchw-horizontal-reuse.wgsl.jinja"],
|
| 29 |
+
"nchw2d_separable": ["pool2d-nchw-separable.wgsl.jinja"],
|
| 30 |
+
"nchw2d_unrolled": ["pool-window-unroll.wgsl.jinja"],
|
| 31 |
+
"ncl1d": ["pool-window-nd.wgsl.jinja"],
|
| 32 |
+
"nchw2d": ["pool-window-nd.wgsl.jinja"],
|
| 33 |
+
"ncdhw3d": ["pool-window-nd.wgsl.jinja"],
|
| 34 |
+
"nchw2d_k7s2_reuse": ["average-pool2d-nchw-horizontal-reuse.wgsl.jinja"],
|
| 35 |
+
"nchw2d_k7s1_reuse": ["average-pool2d-nchw-horizontal-reuse.wgsl.jinja"]
|
| 36 |
+
}
|
| 37 |
+
}
|
| 38 |
}
|
build/webgpu/pool-global-reduction.wgsl.jinja
CHANGED
|
@@ -1,7 +1,4 @@
|
|
| 1 |
-
{% if
|
| 2 |
-
enable f16;
|
| 3 |
-
{% endif %}
|
| 4 |
-
{% if source.useSubgroups %}
|
| 5 |
enable subgroups;
|
| 6 |
{% endif %}
|
| 7 |
{{ env.wgsl.resourceDeclarations }}
|
|
@@ -9,17 +6,18 @@ enable subgroups;
|
|
| 9 |
// Workgroup-parallel global pooling: one workgroup reduces one contiguous
|
| 10 |
// (n, c) spatial plane (NC[D]HW layout makes the plane contiguous).
|
| 11 |
// Threads stride the plane (vec4 loads when the plane size is divisible by 4),
|
| 12 |
-
// then combine the per-thread partials. With subgroups: a subgroup collective
|
| 13 |
-
//
|
| 14 |
-
//
|
| 15 |
-
//
|
| 16 |
-
// epilogue
|
| 17 |
-
//
|
| 18 |
-
const WG: u32 = {{
|
| 19 |
-
{% if
|
| 20 |
-
//
|
| 21 |
-
|
| 22 |
-
|
|
|
|
| 23 |
{% else %}
|
| 24 |
var<workgroup> wg_tree: array<f32, WG>;
|
| 25 |
{% endif %}
|
|
@@ -28,23 +26,13 @@ var<workgroup> wg_tree: array<f32, WG>;
|
|
| 28 |
@compute @workgroup_size(WG)
|
| 29 |
fn main(@builtin(workgroup_id) wid: vec3<u32>,
|
| 30 |
@builtin(local_invocation_id) lid: vec3<u32>
|
| 31 |
-
{%- if
|
| 32 |
@builtin(subgroup_size) sg_size: u32
|
| 33 |
{%- endif %}) {
|
| 34 |
let plane = wid.x;
|
| 35 |
let tid = lid.x;
|
| 36 |
-
{% if source.useSubgroups %}
|
| 37 |
-
|
| 38 |
-
// Seed cross-subgroup slots with the reduction identity so unwritten slots
|
| 39 |
-
// (when the workgroup has fewer than MAX_SUBGROUPS subgroups) are harmless.
|
| 40 |
-
if (tid < MAX_SUBGROUPS) {
|
| 41 |
-
wg_partials[tid] = 0.0;
|
| 42 |
-
}
|
| 43 |
-
workgroupBarrier();
|
| 44 |
-
|
| 45 |
-
{% endif %}
|
| 46 |
var partial = 0.0;
|
| 47 |
-
{% if
|
| 48 |
let spatialVec = params.spatial / 4u;
|
| 49 |
let base = plane * spatialVec;
|
| 50 |
var accv = vec4<f32>(0.0);
|
|
@@ -61,18 +49,31 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
|
|
| 61 |
}
|
| 62 |
{% endif %}
|
| 63 |
|
| 64 |
-
{% if
|
| 65 |
-
|
| 66 |
-
|
| 67 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 68 |
}
|
| 69 |
-
workgroupBarrier();
|
| 70 |
if (tid == 0u) {
|
| 71 |
-
let nsg = min((WG + sg_size - 1u) / sg_size, MAX_SUBGROUPS);
|
| 72 |
-
var total = 0.0;
|
| 73 |
-
for (var i = 0u; i < nsg; i = i + 1u) {
|
| 74 |
-
total = total + wg_partials[i];
|
| 75 |
-
}
|
| 76 |
y[plane] = {{ scalar }}(total / f32(params.spatial));
|
| 77 |
}
|
| 78 |
{% else %}
|
|
|
|
| 1 |
+
{% if useSubgroups %}
|
|
|
|
|
|
|
|
|
|
| 2 |
enable subgroups;
|
| 3 |
{% endif %}
|
| 4 |
{{ env.wgsl.resourceDeclarations }}
|
|
|
|
| 6 |
// Workgroup-parallel global pooling: one workgroup reduces one contiguous
|
| 7 |
// (n, c) spatial plane (NC[D]HW layout makes the plane contiguous).
|
| 8 |
// Threads stride the plane (vec4 loads when the plane size is divisible by 4),
|
| 9 |
+
// then combine the per-thread partials. With subgroups: a subgroup collective,
|
| 10 |
+
// then one shared-memory slot per invocation (the subgroup partial from its
|
| 11 |
+
// elected lane, the identity from every other lane) that every subgroup folds.
|
| 12 |
+
// Without subgroups, a full workgroup tree combines the per-thread partials.
|
| 13 |
+
// The epilogue divides the f32 sum by the spatial extent. The tree changes the
|
| 14 |
+
// association of additions.
|
| 15 |
+
const WG: u32 = {{ workgroupSizeSpec }}u;
|
| 16 |
+
{% if useSubgroups %}
|
| 17 |
+
// One slot per invocation: the cross-subgroup fold below indexes by
|
| 18 |
+
// local_invocation_id, never by a subgroup ordinal, so no subgroup width or
|
| 19 |
+
// partition can overflow or alias it.
|
| 20 |
+
var<workgroup> wg_partials: array<f32, WG>;
|
| 21 |
{% else %}
|
| 22 |
var<workgroup> wg_tree: array<f32, WG>;
|
| 23 |
{% endif %}
|
|
|
|
| 26 |
@compute @workgroup_size(WG)
|
| 27 |
fn main(@builtin(workgroup_id) wid: vec3<u32>,
|
| 28 |
@builtin(local_invocation_id) lid: vec3<u32>
|
| 29 |
+
{%- if useSubgroups %},
|
| 30 |
@builtin(subgroup_size) sg_size: u32
|
| 31 |
{%- endif %}) {
|
| 32 |
let plane = wid.x;
|
| 33 |
let tid = lid.x;
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 34 |
var partial = 0.0;
|
| 35 |
+
{% if vec4 %}
|
| 36 |
let spatialVec = params.spatial / 4u;
|
| 37 |
let base = plane * spatialVec;
|
| 38 |
var accv = vec4<f32>(0.0);
|
|
|
|
| 49 |
}
|
| 50 |
{% endif %}
|
| 51 |
|
| 52 |
+
{% if useSubgroups %}
|
| 53 |
+
var total = subgroupAdd(partial);
|
| 54 |
+
// Cross-subgroup fold that assumes nothing about which invocations share a
|
| 55 |
+
// subgroup, how many subgroups there are, or which of a subgroup's lanes are
|
| 56 |
+
// active: every invocation owns the slot at its own index, the elected lane
|
| 57 |
+
// publishes its subgroup partial there and every other lane publishes the
|
| 58 |
+
// reduction identity. Each subgroup then folds all WG slots — lane `rank`,
|
| 59 |
+
// its dense position among the active lanes, walks slots rank, rank + count,
|
| 60 |
+
// ... — and one more collective merges the lane partials, so every slot is
|
| 61 |
+
// merged exactly once at any legal width and partition. A one-subgroup
|
| 62 |
+
// workgroup skips shared memory entirely; that test reads the `subgroup_size`
|
| 63 |
+
// builtin, which is uniform, because a collective's result is not uniform to
|
| 64 |
+
// WGSL's analysis and may not guard a barrier.
|
| 65 |
+
if (sg_size != WG) {
|
| 66 |
+
let rank = subgroupExclusiveAdd(1u);
|
| 67 |
+
let count = subgroupAdd(1u);
|
| 68 |
+
wg_partials[tid] = select(0.0, total, rank == 0u);
|
| 69 |
+
workgroupBarrier();
|
| 70 |
+
var acc = 0.0;
|
| 71 |
+
for (var i = rank; i < WG; i = i + count) {
|
| 72 |
+
acc = acc + wg_partials[i];
|
| 73 |
+
}
|
| 74 |
+
total = subgroupAdd(acc);
|
| 75 |
}
|
|
|
|
| 76 |
if (tid == 0u) {
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 77 |
y[plane] = {{ scalar }}(total / f32(params.spatial));
|
| 78 |
}
|
| 79 |
{% else %}
|
build/webgpu/pool-window-nd.wgsl.jinja
CHANGED
|
@@ -1,43 +1,40 @@
|
|
| 1 |
-
{% if usesF16 %}
|
| 2 |
-
enable f16;
|
| 3 |
-
{% endif %}
|
| 4 |
{{ env.wgsl.resourceDeclarations }}
|
| 5 |
|
| 6 |
-
// Direct
|
| 7 |
@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
|
| 8 |
-
fn main(@builtin(global_invocation_id) gid: vec3<u32>
|
| 9 |
-
let index = gid.x + gid.y *
|
| 10 |
if (index >= params.count) { return; }
|
| 11 |
|
| 12 |
let ow = index % params.outW;
|
| 13 |
var t = index / params.outW;
|
| 14 |
-
{% if
|
| 15 |
let oh = t % params.outH;
|
| 16 |
t /= params.outH;
|
| 17 |
{% endif %}
|
| 18 |
-
{% if
|
| 19 |
let od = t % params.outD;
|
| 20 |
t /= params.outD;
|
| 21 |
{% endif %}
|
| 22 |
let channel = t % params.channels;
|
| 23 |
let batch = t / params.channels;
|
| 24 |
|
| 25 |
-
{% if
|
| 26 |
let startD = i32(od * params.strideD) - params.padFront;
|
| 27 |
{% endif %}
|
| 28 |
-
{% if
|
| 29 |
let startH = i32(oh * params.strideH) - params.padTop;
|
| 30 |
{% endif %}
|
| 31 |
let startW = i32(ow * params.strideW) - params.padLeft;
|
| 32 |
var acc = 0.0;
|
| 33 |
var seen = 0u;
|
| 34 |
|
| 35 |
-
{% if
|
| 36 |
for (var kd = 0u; kd < params.kernelD; kd += 1u) {
|
| 37 |
let id = startD + i32(kd * params.dilationD);
|
| 38 |
if (id < 0 || id >= i32(params.inD)) { continue; }
|
| 39 |
{% endif %}
|
| 40 |
-
{% if
|
| 41 |
for (var kh = 0u; kh < params.kernelH; kh += 1u) {
|
| 42 |
let ih = startH + i32(kh * params.dilationH);
|
| 43 |
if (ih < 0 || ih >= i32(params.inH)) { continue; }
|
|
@@ -45,9 +42,9 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
|
|
| 45 |
for (var kw = 0u; kw < params.kernelW; kw += 1u) {
|
| 46 |
let iw = startW + i32(kw * params.dilationW);
|
| 47 |
if (iw < 0 || iw >= i32(params.inW)) { continue; }
|
| 48 |
-
{% if
|
| 49 |
let xIndex = (batch * params.channels + channel) * params.inW + u32(iw);
|
| 50 |
-
{% elif
|
| 51 |
let xIndex = ((batch * params.channels + channel) * params.inH + u32(ih)) * params.inW + u32(iw);
|
| 52 |
{% else %}
|
| 53 |
let xIndex = (((batch * params.channels + channel) * params.inD + u32(id)) * params.inH + u32(ih)) * params.inW + u32(iw);
|
|
@@ -55,22 +52,22 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
|
|
| 55 |
acc += f32(x[xIndex]);
|
| 56 |
seen += 1u;
|
| 57 |
}
|
| 58 |
-
{% if
|
| 59 |
}
|
| 60 |
{% endif %}
|
| 61 |
-
{% if
|
| 62 |
}
|
| 63 |
{% endif %}
|
| 64 |
|
| 65 |
var denom = seen;
|
| 66 |
if (params.countIncludePad != 0u) {
|
| 67 |
denom = 0u;
|
| 68 |
-
{% if
|
| 69 |
for (var kd = 0u; kd < params.kernelD; kd += 1u) {
|
| 70 |
let id = startD + i32(kd * params.dilationD);
|
| 71 |
if (id < -params.padFront || id >= i32(params.inD) + params.padBack) { continue; }
|
| 72 |
{% endif %}
|
| 73 |
-
{% if
|
| 74 |
for (var kh = 0u; kh < params.kernelH; kh += 1u) {
|
| 75 |
let ih = startH + i32(kh * params.dilationH);
|
| 76 |
if (ih < -params.padTop || ih >= i32(params.inH) + params.padBottom) { continue; }
|
|
@@ -80,10 +77,10 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
|
|
| 80 |
if (iw < -params.padLeft || iw >= i32(params.inW) + params.padRight) { continue; }
|
| 81 |
denom += 1u;
|
| 82 |
}
|
| 83 |
-
{% if
|
| 84 |
}
|
| 85 |
{% endif %}
|
| 86 |
-
{% if
|
| 87 |
}
|
| 88 |
{% endif %}
|
| 89 |
}
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
{{ env.wgsl.resourceDeclarations }}
|
| 2 |
|
| 3 |
+
// Direct spatial-pooling traversal for a channels-first tensor.
|
| 4 |
@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
|
| 5 |
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
| 6 |
+
let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
|
| 7 |
if (index >= params.count) { return; }
|
| 8 |
|
| 9 |
let ow = index % params.outW;
|
| 10 |
var t = index / params.outW;
|
| 11 |
+
{% if spatialRankSpec >= 2 %}
|
| 12 |
let oh = t % params.outH;
|
| 13 |
t /= params.outH;
|
| 14 |
{% endif %}
|
| 15 |
+
{% if spatialRankSpec == 3 %}
|
| 16 |
let od = t % params.outD;
|
| 17 |
t /= params.outD;
|
| 18 |
{% endif %}
|
| 19 |
let channel = t % params.channels;
|
| 20 |
let batch = t / params.channels;
|
| 21 |
|
| 22 |
+
{% if spatialRankSpec == 3 %}
|
| 23 |
let startD = i32(od * params.strideD) - params.padFront;
|
| 24 |
{% endif %}
|
| 25 |
+
{% if spatialRankSpec >= 2 %}
|
| 26 |
let startH = i32(oh * params.strideH) - params.padTop;
|
| 27 |
{% endif %}
|
| 28 |
let startW = i32(ow * params.strideW) - params.padLeft;
|
| 29 |
var acc = 0.0;
|
| 30 |
var seen = 0u;
|
| 31 |
|
| 32 |
+
{% if spatialRankSpec == 3 %}
|
| 33 |
for (var kd = 0u; kd < params.kernelD; kd += 1u) {
|
| 34 |
let id = startD + i32(kd * params.dilationD);
|
| 35 |
if (id < 0 || id >= i32(params.inD)) { continue; }
|
| 36 |
{% endif %}
|
| 37 |
+
{% if spatialRankSpec >= 2 %}
|
| 38 |
for (var kh = 0u; kh < params.kernelH; kh += 1u) {
|
| 39 |
let ih = startH + i32(kh * params.dilationH);
|
| 40 |
if (ih < 0 || ih >= i32(params.inH)) { continue; }
|
|
|
|
| 42 |
for (var kw = 0u; kw < params.kernelW; kw += 1u) {
|
| 43 |
let iw = startW + i32(kw * params.dilationW);
|
| 44 |
if (iw < 0 || iw >= i32(params.inW)) { continue; }
|
| 45 |
+
{% if spatialRankSpec == 1 %}
|
| 46 |
let xIndex = (batch * params.channels + channel) * params.inW + u32(iw);
|
| 47 |
+
{% elif spatialRankSpec == 2 %}
|
| 48 |
let xIndex = ((batch * params.channels + channel) * params.inH + u32(ih)) * params.inW + u32(iw);
|
| 49 |
{% else %}
|
| 50 |
let xIndex = (((batch * params.channels + channel) * params.inD + u32(id)) * params.inH + u32(ih)) * params.inW + u32(iw);
|
|
|
|
| 52 |
acc += f32(x[xIndex]);
|
| 53 |
seen += 1u;
|
| 54 |
}
|
| 55 |
+
{% if spatialRankSpec >= 2 %}
|
| 56 |
}
|
| 57 |
{% endif %}
|
| 58 |
+
{% if spatialRankSpec == 3 %}
|
| 59 |
}
|
| 60 |
{% endif %}
|
| 61 |
|
| 62 |
var denom = seen;
|
| 63 |
if (params.countIncludePad != 0u) {
|
| 64 |
denom = 0u;
|
| 65 |
+
{% if spatialRankSpec == 3 %}
|
| 66 |
for (var kd = 0u; kd < params.kernelD; kd += 1u) {
|
| 67 |
let id = startD + i32(kd * params.dilationD);
|
| 68 |
if (id < -params.padFront || id >= i32(params.inD) + params.padBack) { continue; }
|
| 69 |
{% endif %}
|
| 70 |
+
{% if spatialRankSpec >= 2 %}
|
| 71 |
for (var kh = 0u; kh < params.kernelH; kh += 1u) {
|
| 72 |
let ih = startH + i32(kh * params.dilationH);
|
| 73 |
if (ih < -params.padTop || ih >= i32(params.inH) + params.padBottom) { continue; }
|
|
|
|
| 77 |
if (iw < -params.padLeft || iw >= i32(params.inW) + params.padRight) { continue; }
|
| 78 |
denom += 1u;
|
| 79 |
}
|
| 80 |
+
{% if spatialRankSpec >= 2 %}
|
| 81 |
}
|
| 82 |
{% endif %}
|
| 83 |
+
{% if spatialRankSpec == 3 %}
|
| 84 |
}
|
| 85 |
{% endif %}
|
| 86 |
}
|
build/webgpu/pool-window-unroll.wgsl.jinja
CHANGED
|
@@ -1,32 +1,31 @@
|
|
| 1 |
// Compile-time-unrolled 2D NCHW pooling for small windows. Accumulation remains
|
| 2 |
-
// kh-major/kw-minor, matching the
|
| 3 |
-
{%
|
| 4 |
-
|
| 5 |
-
{%
|
| 6 |
{{ env.wgsl.resourceDeclarations }}
|
| 7 |
|
| 8 |
-
const OUT_W: u32 = {{
|
| 9 |
-
const OUT_H: u32 = {{
|
| 10 |
-
const IN_H: i32 = {{
|
| 11 |
-
const IN_W: i32 = {{
|
| 12 |
-
const IN_W_U: u32 = {{
|
| 13 |
-
const IN_PLANE: u32 = {{
|
| 14 |
-
const STRIDE_H: u32 = {{
|
| 15 |
-
const STRIDE_W: u32 = {{
|
| 16 |
-
const PAD_TOP: i32 = {{
|
| 17 |
-
const PAD_LEFT: i32 = {{
|
| 18 |
-
{% if
|
| 19 |
-
const PAD_BOTTOM: i32 = {{
|
| 20 |
-
const PAD_RIGHT: i32 = {{
|
| 21 |
{% endif %}
|
| 22 |
|
| 23 |
@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
|
| 24 |
fn main(
|
| 25 |
-
@builtin(global_invocation_id) gid: vec3<u32>
|
| 26 |
-
@builtin(num_workgroups) nwg: vec3<u32>
|
| 27 |
) {
|
| 28 |
// Fold very large flat outputs into the dispatch's second dimension.
|
| 29 |
-
let index = gid.x + gid.y *
|
| 30 |
if (index >= params.count) {
|
| 31 |
return;
|
| 32 |
}
|
|
@@ -38,19 +37,19 @@ fn main(
|
|
| 38 |
let planeBase = plane * IN_PLANE;
|
| 39 |
let startH = i32(oh * STRIDE_H) - PAD_TOP;
|
| 40 |
let startW = i32(ow * STRIDE_W) - PAD_LEFT;
|
| 41 |
-
{% for kh in range(
|
| 42 |
-
let ih_{{ kh }} = startH + {{ kh *
|
| 43 |
{% endfor %}
|
| 44 |
-
{% for kw in range(
|
| 45 |
-
let iw_{{ kw }} = startW + {{ kw *
|
| 46 |
{% endfor %}
|
| 47 |
|
| 48 |
var result = 0.0;
|
| 49 |
var seen = 0u;
|
| 50 |
-
{% for kh in range(
|
| 51 |
if (ih_{{ kh }} >= 0 && ih_{{ kh }} < IN_H) {
|
| 52 |
let rowBase{{ kh }} = planeBase + u32(ih_{{ kh }}) * IN_W_U;
|
| 53 |
-
{% for kw in range(
|
| 54 |
if (iw_{{ kw }} >= 0 && iw_{{ kw }} < IN_W) {
|
| 55 |
let value_{{ kh }}_{{ kw }} = f32(x[rowBase{{ kh }} + u32(iw_{{ kw }})]);
|
| 56 |
result += value_{{ kh }}_{{ kw }};
|
|
@@ -60,11 +59,11 @@ fn main(
|
|
| 60 |
}
|
| 61 |
{% endfor %}
|
| 62 |
|
| 63 |
-
{% if
|
| 64 |
var denominator = 0u;
|
| 65 |
-
{% for kh in range(
|
| 66 |
if (ih_{{ kh }} >= -PAD_TOP && ih_{{ kh }} < IN_H + PAD_BOTTOM) {
|
| 67 |
-
{% for kw in range(
|
| 68 |
if (iw_{{ kw }} >= -PAD_LEFT && iw_{{ kw }} < IN_W + PAD_RIGHT) {
|
| 69 |
denominator += 1u;
|
| 70 |
}
|
|
|
|
| 1 |
// Compile-time-unrolled 2D NCHW pooling for small windows. Accumulation remains
|
| 2 |
+
// kh-major/kw-minor, matching the loop-based implementations.
|
| 3 |
+
{% set countIncludePad = countIncludePad is defined and countIncludePad %}
|
| 4 |
+
{% set padBottomSpec = padBottomSpec | default(0) %}
|
| 5 |
+
{% set padRightSpec = padRightSpec | default(0) %}
|
| 6 |
{{ env.wgsl.resourceDeclarations }}
|
| 7 |
|
| 8 |
+
const OUT_W: u32 = {{ outW }}u;
|
| 9 |
+
const OUT_H: u32 = {{ outH }}u;
|
| 10 |
+
const IN_H: i32 = {{ inH }};
|
| 11 |
+
const IN_W: i32 = {{ inW }};
|
| 12 |
+
const IN_W_U: u32 = {{ inW }}u;
|
| 13 |
+
const IN_PLANE: u32 = {{ inH * inW }}u;
|
| 14 |
+
const STRIDE_H: u32 = {{ strideHSpec }}u;
|
| 15 |
+
const STRIDE_W: u32 = {{ strideWSpec }}u;
|
| 16 |
+
const PAD_TOP: i32 = {{ padTopSpec }};
|
| 17 |
+
const PAD_LEFT: i32 = {{ padLeftSpec }};
|
| 18 |
+
{% if op == "average" and countIncludePad %}
|
| 19 |
+
const PAD_BOTTOM: i32 = {{ padBottomSpec }};
|
| 20 |
+
const PAD_RIGHT: i32 = {{ padRightSpec }};
|
| 21 |
{% endif %}
|
| 22 |
|
| 23 |
@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
|
| 24 |
fn main(
|
| 25 |
+
@builtin(global_invocation_id) gid: vec3<u32>
|
|
|
|
| 26 |
) {
|
| 27 |
// Fold very large flat outputs into the dispatch's second dimension.
|
| 28 |
+
let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
|
| 29 |
if (index >= params.count) {
|
| 30 |
return;
|
| 31 |
}
|
|
|
|
| 37 |
let planeBase = plane * IN_PLANE;
|
| 38 |
let startH = i32(oh * STRIDE_H) - PAD_TOP;
|
| 39 |
let startW = i32(ow * STRIDE_W) - PAD_LEFT;
|
| 40 |
+
{% for kh in range(kernelHSpec) %}
|
| 41 |
+
let ih_{{ kh }} = startH + {{ kh * dilationHSpec }};
|
| 42 |
{% endfor %}
|
| 43 |
+
{% for kw in range(kernelWSpec) %}
|
| 44 |
+
let iw_{{ kw }} = startW + {{ kw * dilationWSpec }};
|
| 45 |
{% endfor %}
|
| 46 |
|
| 47 |
var result = 0.0;
|
| 48 |
var seen = 0u;
|
| 49 |
+
{% for kh in range(kernelHSpec) %}
|
| 50 |
if (ih_{{ kh }} >= 0 && ih_{{ kh }} < IN_H) {
|
| 51 |
let rowBase{{ kh }} = planeBase + u32(ih_{{ kh }}) * IN_W_U;
|
| 52 |
+
{% for kw in range(kernelWSpec) %}
|
| 53 |
if (iw_{{ kw }} >= 0 && iw_{{ kw }} < IN_W) {
|
| 54 |
let value_{{ kh }}_{{ kw }} = f32(x[rowBase{{ kh }} + u32(iw_{{ kw }})]);
|
| 55 |
result += value_{{ kh }}_{{ kw }};
|
|
|
|
| 59 |
}
|
| 60 |
{% endfor %}
|
| 61 |
|
| 62 |
+
{% if op == "average" and countIncludePad %}
|
| 63 |
var denominator = 0u;
|
| 64 |
+
{% for kh in range(kernelHSpec) %}
|
| 65 |
if (ih_{{ kh }} >= -PAD_TOP && ih_{{ kh }} < IN_H + PAD_BOTTOM) {
|
| 66 |
+
{% for kw in range(kernelWSpec) %}
|
| 67 |
if (iw_{{ kw }} >= -PAD_LEFT && iw_{{ kw }} < IN_W + PAD_RIGHT) {
|
| 68 |
denominator += 1u;
|
| 69 |
}
|
build/webgpu/pool2d-nchw-separable.wgsl.jinja
ADDED
|
@@ -0,0 +1,48 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{{ env.wgsl.resourceDeclarations }}
|
| 2 |
+
|
| 3 |
+
// Separable NCHW pooling: one spatial axis per pass.
|
| 4 |
+
//
|
| 5 |
+
// A kernelH x kernelW window factorizes into a width reduction followed by a
|
| 6 |
+
// height reduction. The width pass writes one f32 intermediate per input row
|
| 7 |
+
// and output column; the height pass completes the selected pooling operation.
|
| 8 |
+
// `axis == "w"` writes `tmp`, and `axis == "h"` writes `y`.
|
| 9 |
+
// Both passes sum in f32. This specialization requires the whole window as the
|
| 10 |
+
// denominator, whose baked area is applied once after the height fold.
|
| 11 |
+
|
| 12 |
+
@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
|
| 13 |
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
| 14 |
+
let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
|
| 15 |
+
{% if axisSpec == "w" %}
|
| 16 |
+
if (index >= params.tmpCount) { return; }
|
| 17 |
+
// One cell per (batch, channel, input row, output column).
|
| 18 |
+
let ow = index % params.outW;
|
| 19 |
+
let t = index / params.outW;
|
| 20 |
+
let ih = t % params.inH;
|
| 21 |
+
let nc = t / params.inH;
|
| 22 |
+
|
| 23 |
+
var acc = 0.0;
|
| 24 |
+
for (var kw = 0u; kw < params.kernelW; kw += 1u) {
|
| 25 |
+
let iw = i32(ow * params.strideW + kw * params.dilationW) - params.padW;
|
| 26 |
+
if (iw < 0 || iw >= i32(params.inW)) { continue; }
|
| 27 |
+
let value = f32(x[(nc * params.inH + ih) * params.inW + u32(iw)]);
|
| 28 |
+
acc += value;
|
| 29 |
+
}
|
| 30 |
+
tmp[index] = acc;
|
| 31 |
+
{% else %}
|
| 32 |
+
if (index >= params.count) { return; }
|
| 33 |
+
let ow = index % params.outW;
|
| 34 |
+
let t = index / params.outW;
|
| 35 |
+
let oh = t % params.outH;
|
| 36 |
+
let nc = t / params.outH;
|
| 37 |
+
|
| 38 |
+
var acc = 0.0;
|
| 39 |
+
for (var kh = 0u; kh < params.kernelH; kh += 1u) {
|
| 40 |
+
let ih = i32(oh * params.strideH + kh * params.dilationH) - params.padH;
|
| 41 |
+
if (ih < 0 || ih >= i32(params.inH)) { continue; }
|
| 42 |
+
let partial = tmp[(nc * params.inH + u32(ih)) * params.outW + ow];
|
| 43 |
+
acc += partial;
|
| 44 |
+
}
|
| 45 |
+
|
| 46 |
+
y[index] = {{ scalar }}(acc / {{ windowArea }}.0);
|
| 47 |
+
{% endif %}
|
| 48 |
+
}
|
build/webgpu/test.json
CHANGED
|
The diff for this file is too large to render.
See raw diff
|
|
|