sync 2e7068faf55e

Browse files

Files changed (11) hide show

README.md +77 -0
build/webgpu/bench.json +137 -0
build/webgpu/manifest.json +461 -0
build/webgpu/metadata.json +23 -0
build/webgpu/pool-global-reduction.wgsl.jinja +120 -0
build/webgpu/pool-ncl1d-x4.wgsl.jinja +45 -0
build/webgpu/pool-window-nd.wgsl.jinja +70 -0
build/webgpu/pool-window-unroll.wgsl.jinja +75 -0
build/webgpu/pool2d-nchw-k2s2-vec4.wgsl.jinja +65 -0
build/webgpu/pool2d-nchw-separable.wgsl.jinja +73 -0
build/webgpu/test.json +0 -0

README.md CHANGED Viewed

@@ -1,3 +1,80 @@
 ---
 license: apache-2.0
 ---

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# ai.onnx.LpPool
+`ai.onnx`  ·  standard ONNX operator  ·  ONNX opset ≥ 18
+## Description
+Applies Lp pooling over a spatial input tensor by computing the Lp norm within each kernel window and writing the result to the output. Output spatial dimensions are determined by the kernel size, strides, padding, and `ceil_mode`; `p` controls which norm is used (e.g. `p=1` for sum-of-absolutes, `p=2` for Euclidean).
+See the [ONNX `LpPool` spec](https://onnx.ai/onnx/operators/onnx__LpPool.html) for the reference semantics.
+## Inputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `X` | `x` | `T` | — | — | Input tensor of shape `(N x C x D1 x ... x Dn)`; for images the spatial axes are `H` and `W`. | required |
+## Outputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `Y` | `y` | `T` | same as `X` | derived; see description | Output tensor after Lp pooling; spatial dimensions vary with kernel, stride, and pad settings. | required |
+## Attributes
+Attributes and default values (overridable per request):
+| Attribute | Default | Description |
+| --- | --- | --- |
+| `p` | `2` | The exponent of the Lp norm used for pooling; default `2` gives Euclidean (L2) pooling. |
+| `auto_pad` | `"NOTSET"` | Deprecated auto-padding mode (`NOTSET`, `SAME_UPPER`, `SAME_LOWER`, or `VALID`). It cannot be used together with `pads`. |
+| `ceil_mode` | `0` | When non-zero, uses `ceil` instead of `floor` to compute output spatial dimensions. |
+| `kernel_shape` | — | Required kernel shape, with one positive value per spatial axis. |
+| `strides` | — | Stride along each spatial axis. When omitted, every stride is 1. |
+| `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. |
+| `dilations` | — | Dilation along each spatial axis. When omitted, every dilation is 1. |
+## Type constraints
+| Variable | Allowed dtypes |
+| --- | --- |
+| `T` | `float32`, `float16` |
+## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
+- [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
+- [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
+- [`pool-global-reduction.wgsl.jinja`](build/webgpu/pool-global-reduction.wgsl.jinja)
+- [`pool-ncl1d-x4.wgsl.jinja`](build/webgpu/pool-ncl1d-x4.wgsl.jinja)
+- [`pool-window-nd.wgsl.jinja`](build/webgpu/pool-window-nd.wgsl.jinja)
+- [`pool-window-unroll.wgsl.jinja`](build/webgpu/pool-window-unroll.wgsl.jinja)
+- [`pool2d-nchw-k2s2-vec4.wgsl.jinja`](build/webgpu/pool2d-nchw-k2s2-vec4.wgsl.jinja)
+- [`pool2d-nchw-separable.wgsl.jinja`](build/webgpu/pool2d-nchw-separable.wgsl.jinja)
+## Use with `@huggingface/kernels`
+The loader derives every required output's shape and logical dtype from the manifest contract and this call.
+It then allocates the result tensors automatically.
+The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+Replace each `*Data` placeholder with a typed array containing the corresponding input data.
+```js
+import { getKernel } from "@huggingface/kernels";
+const kernel = await getKernel("webgpu-kernels/ai.onnx.LpPool", { version: 1 });
+const { y } = await kernel({ x: { data: xData, shape: [1, 1, 4] } }, {
+  attrs: { kernel_shape: [3] },
+});
+```

build/webgpu/bench.json ADDED Viewed

	@@ -0,0 +1,137 @@

+{
+  "op": "ai.onnx.LpPool",
+  "cases": [
+    {
+      "name": "nchw_8_32_64_64_k3",
+      "preset": "smoke",
+      "attrs": { "p": 2, "kernel_shape": [3, 3], "strides": [2, 2] },
+      "inputs": { "x": { "dtype": "float32", "shape": [8, 32, 64, 64] } },
+      "outputs": { "y": { "dtype": "float32", "shape": [8, 32, 31, 31] } }
+    },
+    {
+      "name": "nchw_8x64x112x112_k2s2_p2_unrolled",
+      "preset": "smoke",
+      "vars": { "batch": 8, "channels": 64, "inH": 112, "inW": 112 },
+      "attrs": { "p": 2, "kernel_shape": [2, 2], "strides": [2, 2] },
+      "inputs": { "x": { "dtype": "float32", "shape": [8, 64, 112, 112], "dist": "normal", "seed": 206 } },
+      "outputs": { "y": { "dtype": "float32", "shape": [8, 64, 56, 56], "dist": "empty" } },
+      "bench": {
+        "primary": true,
+        "metrics": [{ "type": "bandwidth", "value": "args.batch * args.channels * args.inH * args.inW * 4" }]
+      }
+    },
+    {
+      "name": "nchw_f16_8x64x112x112_k2s2_p2_unrolled",
+      "preset": "smoke",
+      "vars": { "batch": 8, "channels": 64, "inH": 112, "inW": 112 },
+      "attrs": { "p": 2, "kernel_shape": [2, 2], "strides": [2, 2] },
+      "inputs": { "x": { "dtype": "float16", "shape": [8, 64, 112, 112], "dist": "normal", "seed": 207 } },
+      "outputs": { "y": { "dtype": "float16", "shape": [8, 64, 56, 56], "dist": "empty" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "args.batch * args.channels * args.inH * args.inW * 2" }] }
+    },
+    {
+      "name": "nchw_f16_8x64x112x110_k2s2_p2_unaligned_width",
+      "preset": "smoke",
+      "vars": { "batch": 8, "channels": 64, "inH": 112, "inW": 110 },
+      "attrs": { "p": 2, "kernel_shape": [2, 2], "strides": [2, 2] },
+      "inputs": { "x": { "dtype": "float16", "shape": [8, 64, 112, 110], "dist": "normal", "seed": 311 } },
+      "outputs": { "y": { "dtype": "float16", "shape": [8, 64, 56, 55], "dist": "empty" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "args.batch * args.channels * args.inH * args.inW * 2" }] }
+    },
+    {
+      "name": "nchw_f16_8x64x113x112_k2s2_p2_unaligned_height",
+      "preset": "smoke",
+      "vars": { "batch": 8, "channels": 64, "inH": 113, "inW": 112 },
+      "attrs": { "p": 2, "kernel_shape": [2, 2], "strides": [2, 2] },
+      "inputs": { "x": { "dtype": "float16", "shape": [8, 64, 113, 112], "dist": "normal", "seed": 312 } },
+      "outputs": { "y": { "dtype": "float16", "shape": [8, 64, 56, 56], "dist": "empty" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "args.batch * args.channels * args.inH * args.inW * 2" }] }
+    },
+    {
+      "name": "ncl1d_f32_8x256x4000_k4s2_p2_seqpool",
+      "preset": "smoke",
+      "vars": { "batch": 8, "channels": 256, "inW": 4000 },
+      "attrs": { "p": 2, "kernel_shape": [4], "strides": [2] },
+      "inputs": { "x": { "dtype": "float32", "shape": [8, 256, 4000], "dist": "normal", "seed": 313 } },
+      "outputs": { "y": { "dtype": "float32", "shape": [8, 256, 1999], "dist": "empty" } },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "(args.batch * args.channels * args.inW + numel(shapes.y)) * 4" }]
+      }
+    },
+    {
+      "name": "ncdhw3d_f32_2x16x16x16x16_k2s2_p2_volumetric",
+      "preset": "smoke",
+      "vars": { "batch": 2, "channels": 16, "inD": 16, "inH": 16, "inW": 16 },
+      "attrs": { "p": 2, "kernel_shape": [2, 2, 2], "strides": [2, 2, 2] },
+      "inputs": { "x": { "dtype": "float32", "shape": [2, 16, 16, 16, 16], "dist": "normal", "seed": 314 } },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 16, 8, 8, 8], "dist": "empty" } },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "args.batch * args.channels * args.inD * args.inH * args.inW * 4" }]
+      }
+    },
+    {
+      "name": "nchw2d_fallback_bigkernel_lowocc_p2_128x128window",
+      "preset": "stress",
+      "vars": { "batch": 1, "channels": 64, "inH": 128, "inW": 128 },
+      "attrs": { "p": 2, "kernel_shape": [128, 128] },
+      "inputs": { "x": { "dtype": "float32", "shape": [1, 64, 128, 128], "dist": "normal", "seed": 741, "scale": 2 } },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 64, 1, 1], "dist": "empty" } },
+      "bench": {
+        "primary": true,
+        "metrics": [{ "type": "bandwidth", "value": "args.batch * args.channels * args.inH * args.inW * 4" }]
+      }
+    },
+    {
+      "name": "nchw2d_fallback_bigkernel_p3_pow_transcendental_96x96window",
+      "preset": "stress",
+      "vars": { "batch": 1, "channels": 32, "inH": 96, "inW": 96 },
+      "attrs": { "p": 3, "kernel_shape": [96, 96] },
+      "inputs": { "x": { "dtype": "float32", "shape": [1, 32, 96, 96], "dist": "normal", "seed": 742, "scale": 2 } },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 32, 1, 1], "dist": "empty" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "args.batch * args.channels * args.inH * args.inW * 4" }] }
+    },
+    {
+      "name": "ncl1d_bigkernel_lowocc_p2_seqglobalpool",
+      "preset": "stress",
+      "vars": { "batch": 1, "channels": 64, "inW": 16384 },
+      "attrs": { "p": 2, "kernel_shape": [16384] },
+      "inputs": { "x": { "dtype": "float32", "shape": [1, 64, 16384], "dist": "normal", "seed": 743, "scale": 2 } },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 64, 1], "dist": "empty" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "args.batch * args.channels * args.inW * 4" }] }
+    },
+    {
+      "name": "global-lppool-p2-rank5-serial-reduction-pathology",
+      "preset": "stress",
+      "provenance": {
+        "source": "authored for variant coverage",
+        "notes": "Realistic rank-5 global L2 pool that verifies global_window_parallel coverage and its feature-tier fallbacks."
+      },
+      "vars": { "batch": 1, "channels": 32, "inD": 32, "inH": 64, "inW": 64 },
+      "attrs": { "p": 2, "kernel_shape": [32, 64, 64] },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [1, 32, 32, 64, 64], "dist": "normal", "seed": 5103, "scale": 0.2 }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 32, 1, 1, 1], "dist": "empty" } },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "args.batch * args.channels * args.inD * args.inH * args.inW * 4" }]
+      }
+    },
+    {
+      "name": "nchw_f32_1x64x224x224_k7s2p3_generic_bigwindow",
+      "preset": "stress",
+      "attrs": { "p": 2, "kernel_shape": [7, 7], "strides": [2, 2], "pads": [3, 3, 3, 3] },
+      "inputs": { "x": { "dtype": "float32", "shape": [1, 64, 224, 224], "dist": "normal", "seed": 401, "scale": 2 } },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 64, 112, 112], "dist": "empty" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "numel(shapes.x) * 4" }] }
+    },
+    {
+      "name": "nchw_f32_1x64x224x224_k6s2pad2_separable_breakeven",
+      "preset": "stress",
+      "attrs": { "p": 2, "kernel_shape": [6, 6], "strides": [2, 2], "pads": [2, 2, 2, 2] },
+      "inputs": { "x": { "dtype": "float32", "shape": [1, 64, 224, 224], "dist": "normal", "seed": 409, "scale": 2 } },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 64, 112, 112], "dist": "empty" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "numel(shapes.x) * 4" }] }
+    }
+  ],
+  "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] }
+}

build/webgpu/manifest.json ADDED Viewed

	@@ -0,0 +1,461 @@

+{
+  "domain": "ai.onnx",
+  "name": "LpPool",
+  "sinceVersion": 18,
+  "description": "Applies Lp pooling over a spatial input tensor by computing the Lp norm within each kernel window and writing the result to the output. Output spatial dimensions are determined by the kernel size, strides, padding, and `ceil_mode`; `p` controls which norm is used (e.g. `p=1` for sum-of-absolutes, `p=2` for Euclidean).",
+  "inputs": [
+    {
+      "role": "X",
+      "dtype": "T",
+      "description": "Input tensor of shape `(N x C x D1 x ... x Dn)`; for images the spatial axes are `H` and `W`."
+    }
+  ],
+  "outputs": [
+    {
+      "role": "Y",
+      "dtype": "T",
+      "description": "Output tensor after Lp pooling; spatial dimensions vary with kernel, stride, and pad settings.",
+      "rank": "ranks.X",
+      "shape": "poolOutputShape"
+    }
+  ],
+  "attributes": { "p": 2, "auto_pad": "NOTSET", "ceil_mode": 0 },
+  "attributeDescriptions": {
+    "p": "The exponent of the Lp norm used for pooling; default `2` gives Euclidean (L2) pooling.",
+    "auto_pad": "Deprecated auto-padding mode (`NOTSET`, `SAME_UPPER`, `SAME_LOWER`, or `VALID`). It cannot be used together with `pads`.",
+    "kernel_shape": "Required kernel shape, with one positive value per spatial axis.",
+    "strides": "Stride along each spatial axis. When omitted, every stride is 1.",
+    "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.",
+    "dilations": "Dilation along each spatial axis. When omitted, every dilation is 1.",
+    "ceil_mode": "When non-zero, uses `ceil` instead of `floor` to compute output spatial dimensions."
+  },
+  "attributeConstraints": {
+    "auto_pad": { "values": ["NOTSET", "SAME_UPPER", "SAME_LOWER", "VALID"] },
+    "ceil_mode": { "values": [0, 1] },
+    "kernel_shape": { "required": true }
+  },
+  "typeConstraints": { "T": ["float32", "float16"] },
+  "args": {
+    "x": { "kind": "tensor", "semantic": "X", "role": "x" },
+    "y": { "kind": "tensor", "semantic": "Y", "role": "y" }
+  },
+  "derive": {
+    "spatialRank": "ranks.X - 2",
+    "kernelShapeLengthOk": "not has(attrs, \"kernel_shape\") or (attrs.kernel_shape | length) == spatialRank",
+    "stridesLengthOk": "not has(attrs, \"strides\") or (attrs.strides | length) == spatialRank",
+    "dilationsLengthOk": "not has(attrs, \"dilations\") or (attrs.dilations | length) == spatialRank",
+    "padsLengthOk": "not has(attrs, \"pads\") or (attrs.pads | length) == 2 * spatialRank",
+    "kernelD": "attrs.kernel_shape[0] if kernelShapeLengthOk and has(attrs, \"kernel_shape\") and spatialRank == 3 else 1",
+    "kernelH": "attrs.kernel_shape[spatialRank - 2] if kernelShapeLengthOk and has(attrs, \"kernel_shape\") and spatialRank >= 2 else 1",
+    "kernelW": "attrs.kernel_shape[spatialRank - 1] if kernelShapeLengthOk and has(attrs, \"kernel_shape\") and spatialRank >= 1 else 1",
+    "strideD": "attrs.strides[0] if stridesLengthOk and has(attrs, \"strides\") and spatialRank == 3 else 1",
+    "strideH": "attrs.strides[spatialRank - 2] if stridesLengthOk and has(attrs, \"strides\") and spatialRank >= 2 else 1",
+    "strideW": "attrs.strides[spatialRank - 1] if stridesLengthOk and has(attrs, \"strides\") and spatialRank >= 1 else 1",
+    "dilationD": "attrs.dilations[0] if dilationsLengthOk and has(attrs, \"dilations\") and spatialRank == 3 else 1",
+    "dilationH": "attrs.dilations[spatialRank - 2] if dilationsLengthOk and has(attrs, \"dilations\") and spatialRank >= 2 else 1",
+    "dilationW": "attrs.dilations[spatialRank - 1] if dilationsLengthOk and has(attrs, \"dilations\") and spatialRank >= 1 else 1",
+    "padFront": "attrs.pads[0] if padsLengthOk and has(attrs, \"pads\") and spatialRank == 3 else 0",
+    "padTop": "attrs.pads[spatialRank - 2] if padsLengthOk and has(attrs, \"pads\") and spatialRank >= 2 else 0",
+    "padLeft": "attrs.pads[spatialRank - 1] if padsLengthOk and has(attrs, \"pads\") and spatialRank >= 1 else 0",
+    "padBack": "attrs.pads[spatialRank] if padsLengthOk and has(attrs, \"pads\") and spatialRank == 3 else 0",
+    "padBottom": "attrs.pads[2 * spatialRank - 2] if padsLengthOk and has(attrs, \"pads\") and spatialRank >= 2 else 0",
+    "padRight": "attrs.pads[2 * spatialRank - 1] if padsLengthOk and has(attrs, \"pads\") and spatialRank >= 1 else 0",
+    "poolBaseOk": "f16Ok(dtypes.T) and attrs.p > 0 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)",
+    "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",
+    "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",
+    "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",
+    "inferredOutW": "ceil(dim(shapes.X, ranks.X - 1) / strideW) if attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\" else explicitOutW",
+    "inferredOutH": "ceil(dim(shapes.X, ranks.X - 2) / strideH) if attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\" else explicitOutH",
+    "inferredOutD": "ceil(dim(shapes.X, 2) / strideD) if attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\" else explicitOutD",
+    "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])",
+    "autoPadTotalH": "max(0, (dim(shapes.Y, 2) - 1) * strideH + ((kernelH - 1) * dilationH + 1) - dim(shapes.X, 2)) if ranks.X == 4 else 0",
+    "autoPadTotalW": "max(0, (dim(shapes.Y, 3) - 1) * strideW + ((kernelW - 1) * dilationW + 1) - dim(shapes.X, 3)) if ranks.X == 4 else 0",
+    "effectivePadH": "(floor(autoPadTotalH / 2) if attrs.auto_pad == \"SAME_UPPER\" else autoPadTotalH - floor(autoPadTotalH / 2)) if ranks.X == 4 and (attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") else padTop",
+    "effectivePadW": "(floor(autoPadTotalW / 2) if attrs.auto_pad == \"SAME_UPPER\" else autoPadTotalW - floor(autoPadTotalW / 2)) if ranks.X == 4 and (attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") else padLeft",
+    "globalWindowParallelCovered": "poolBaseOk and numel(shapes.Y) == dim(shapes.X, 0) * dim(shapes.X, 1) and numel(shapes.Y) <= device.limits.maxComputeWorkgroupsPerDimension 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)))",
+    "separableTmpElements": "dim(shapes.X, 0) * dim(shapes.X, 1) * dim(shapes.X, 2) * dim(shapes.Y, 3)",
+    "separableFits": "separableTmpElements * 4 <= min(device.limits.maxStorageBufferBindingSize, device.limits.maxBufferSize) and ceilDiv(separableTmpElements, tunables.WORKGROUP_SIZE) <= device.limits.maxComputeWorkgroupsPerDimension * device.limits.maxComputeWorkgroupsPerDimension"
+  },
+  "constants": { "usesF16": "dtypes.T == \"f16\"", "scalar": "dtypes.T" },
+  "bindingSets": {
+    "globalReduction": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$xElem" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": { "name": "Params", "fields": [{ "name": "spatial", "type": "u32", "value": "inner(shapes.X, 1)" }] }
+      }
+    ],
+    "rank2Vec4": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$xVec4" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [{ "name": "tileCount", "type": "u32", "value": "numel(shapes.Y) / 2" }]
+        }
+      }
+    ],
+    "flatDirect": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.Y)" }] }
+      }
+    ],
+    "rank1X4": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$xVec4" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            {
+              "name": "blockCount",
+              "type": "u32",
+              "value": "dim(shapes.X, 0) * dim(shapes.X, 1) * ceilDiv(dim(shapes.Y, 2), 4)"
+            }
+          ]
+        }
+      }
+    ],
+    "rank1Direct": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "channels", "type": "u32", "value": "dim(shapes.X, 1)" },
+            { "name": "inW", "type": "u32", "value": "dim(shapes.X, 2)" },
+            { "name": "outW", "type": "u32", "value": "dim(shapes.Y, 2)" },
+            { "name": "kernelW", "type": "u32", "value": "kernelW" },
+            { "name": "strideW", "type": "u32", "value": "strideW" },
+            { "name": "dilationW", "type": "u32", "value": "dilationW" },
+            { "name": "padW", "type": "i32", "value": "padLeft" },
+            { "name": "p", "type": "f32", "value": "attrs.p" },
+            { "name": "count", "type": "u32", "value": "numel(shapes.Y)" }
+          ]
+        }
+      }
+    ],
+    "rank3Direct": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "channels", "type": "u32", "value": "dim(shapes.X, 1)" },
+            { "name": "inD", "type": "u32", "value": "dim(shapes.X, 2)" },
+            { "name": "inH", "type": "u32", "value": "dim(shapes.X, 3)" },
+            { "name": "inW", "type": "u32", "value": "dim(shapes.X, 4)" },
+            { "name": "outD", "type": "u32", "value": "dim(shapes.Y, 2)" },
+            { "name": "outH", "type": "u32", "value": "dim(shapes.Y, 3)" },
+            { "name": "outW", "type": "u32", "value": "dim(shapes.Y, 4)" },
+            { "name": "kernelD", "type": "u32", "value": "kernelD" },
+            { "name": "kernelH", "type": "u32", "value": "kernelH" },
+            { "name": "kernelW", "type": "u32", "value": "kernelW" },
+            { "name": "strideD", "type": "u32", "value": "strideD" },
+            { "name": "strideH", "type": "u32", "value": "strideH" },
+            { "name": "strideW", "type": "u32", "value": "strideW" },
+            { "name": "dilationD", "type": "u32", "value": "dilationD" },
+            { "name": "dilationH", "type": "u32", "value": "dilationH" },
+            { "name": "dilationW", "type": "u32", "value": "dilationW" },
+            { "name": "padD", "type": "i32", "value": "padFront" },
+            { "name": "padH", "type": "i32", "value": "padTop" },
+            { "name": "padW", "type": "i32", "value": "padLeft" },
+            { "name": "p", "type": "f32", "value": "attrs.p" },
+            { "name": "count", "type": "u32", "value": "numel(shapes.Y)" }
+          ]
+        }
+      }
+    ],
+    "rank2Direct": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "channels", "type": "u32", "value": "dim(shapes.X, 1)" },
+            { "name": "inH", "type": "u32", "value": "dim(shapes.X, 2)" },
+            { "name": "inW", "type": "u32", "value": "dim(shapes.X, 3)" },
+            { "name": "outH", "type": "u32", "value": "dim(shapes.Y, 2)" },
+            { "name": "outW", "type": "u32", "value": "dim(shapes.Y, 3)" },
+            { "name": "kernelH", "type": "u32", "value": "kernelH" },
+            { "name": "kernelW", "type": "u32", "value": "kernelW" },
+            { "name": "strideH", "type": "u32", "value": "strideH" },
+            { "name": "strideW", "type": "u32", "value": "strideW" },
+            { "name": "dilationH", "type": "u32", "value": "dilationH" },
+            { "name": "dilationW", "type": "u32", "value": "dilationW" },
+            { "name": "padH", "type": "i32", "value": "effectivePadH" },
+            { "name": "padW", "type": "i32", "value": "effectivePadW" },
+            { "name": "p", "type": "f32", "value": "attrs.p" },
+            { "name": "count", "type": "u32", "value": "numel(shapes.Y)" }
+          ]
+        }
+      }
+    ],
+    "rank2SeparableW": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "tmp", "semantic": "tmp", "buffer": { "type": "storage" }, "elementType": "f32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "inH", "type": "u32", "value": "dim(shapes.X, 2)" },
+            { "name": "inW", "type": "u32", "value": "dim(shapes.X, 3)" },
+            { "name": "outW", "type": "u32", "value": "dim(shapes.Y, 3)" },
+            { "name": "kernelW", "type": "u32", "value": "kernelW" },
+            { "name": "strideW", "type": "u32", "value": "strideW" },
+            { "name": "dilationW", "type": "u32", "value": "dilationW" },
+            { "name": "padW", "type": "i32", "value": "effectivePadW" },
+            { "name": "p", "type": "f32", "value": "attrs.p" },
+            {
+              "name": "tmpCount",
+              "type": "u32",
+              "value": "dim(shapes.X, 0) * dim(shapes.X, 1) * dim(shapes.X, 2) * dim(shapes.Y, 3)"
+            }
+          ]
+        }
+      }
+    ],
+    "rank2SeparableH": [
+      { "name": "tmp", "semantic": "tmp", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "inH", "type": "u32", "value": "dim(shapes.X, 2)" },
+            { "name": "outH", "type": "u32", "value": "dim(shapes.Y, 2)" },
+            { "name": "outW", "type": "u32", "value": "dim(shapes.Y, 3)" },
+            { "name": "kernelH", "type": "u32", "value": "kernelH" },
+            { "name": "strideH", "type": "u32", "value": "strideH" },
+            { "name": "dilationH", "type": "u32", "value": "dilationH" },
+            { "name": "padH", "type": "i32", "value": "effectivePadH" },
+            { "name": "p", "type": "f32", "value": "attrs.p" },
+            { "name": "count", "type": "u32", "value": "numel(shapes.Y)" }
+          ]
+        }
+      }
+    ]
+  },
+  "tunables": { "WORKGROUP_SIZE": 256, "SEPARABLE_MIN_WINDOW": 42 },
+  "variants": [
+    {
+      "id": "global_window_parallel",
+      "priority": 50,
+      "when": ["globalWindowParallelCovered"],
+      "constants": {
+        "xElem": "(\"vec4<f16>\" if dtypes.T == \"f16\" else \"vec4<f32>\") if inner(shapes.X, 1) % 4 == 0 else dtypes.T"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "LpPool.GlobalWindowParallel",
+          "source": {
+            "shader": "pool-global-reduction.wgsl.jinja",
+            "inputs": {
+              "op": "\"lp\"",
+              "vec4": "inner(shapes.X, 1) % 4 == 0",
+              "p": "attrs.p",
+              "workgroupSize": "tunables.WORKGROUP_SIZE",
+              "useSubgroups": "device.features.has(\"subgroups\")"
+            }
+          },
+          "bindings": "globalReduction",
+          "dispatch": { "x": "numel(shapes.Y)" }
+        }
+      ]
+    },
+    {
+      "id": "nchw2d_k2s2_vec4",
+      "priority": 30,
+      "when": ["(attrs.auto_pad != \"SAME_UPPER\" and attrs.auto_pad != \"SAME_LOWER\")", "poolBaseOk", "ranks.X == 4", "kernelH == 2", "kernelW == 2", "strideH == 2", "strideW == 2", "dilationH == 1", "dilationW == 1", "padTop == 0", "padLeft == 0", "padBottom == 0", "padRight == 0", "attrs.ceil_mode == 0", "dim(shapes.X, 3) % 4 == 0", "dim(shapes.Y, 2) == floor(dim(shapes.X, 2) / 2)", "dim(shapes.Y, 3) == floor(dim(shapes.X, 3) / 2)", "dim(shapes.Y, 2) > 0", "dim(shapes.Y, 3) > 0"],
+      "constants": { "xVec4": "\"vec4<f16>\" if dtypes.T == \"f16\" else \"vec4<f32>\"" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "LpPool.Nchw2dK2s2Vec4",
+          "source": {
+            "shader": "pool2d-nchw-k2s2-vec4.wgsl.jinja",
+            "inputs": {
+              "op": "\"lp\"",
+              "outH": "dim(shapes.Y, 2)",
+              "outW": "dim(shapes.Y, 3)",
+              "outWHalf": "floor(dim(shapes.Y, 3) / 2)",
+              "inWVec4": "floor(dim(shapes.X, 3) / 4)",
+              "inPlaneVec4": "floor(dim(shapes.X, 2) * dim(shapes.X, 3) / 4)",
+              "p": "attrs.p"
+            }
+          },
+          "bindings": "rank2Vec4",
+          "dispatch": { "threads": "numel(shapes.Y) / 2", "workgroupSize": "tunables.WORKGROUP_SIZE" }
+        }
+      ]
+    },
+    {
+      "id": "nchw2d_unrolled",
+      "priority": 25,
+      "when": ["(attrs.auto_pad != \"SAME_UPPER\" and attrs.auto_pad != \"SAME_LOWER\")", "poolBaseOk", "ranks.X == 4", "kernelH > 0", "kernelW > 0", "strideH > 0", "strideW > 0", "kernelH <= 3", "kernelW <= 3", "dilationH >= 1", "dilationW >= 1", "dim(shapes.Y, 2) > 0", "dim(shapes.Y, 3) > 0"],
+      "passes": [
+        {
+          "id": "main",
+          "name": "LpPool.Nchw2dUnrolled",
+          "source": {
+            "shader": "pool-window-unroll.wgsl.jinja",
+            "inputs": {
+              "op": "\"lp\"",
+              "kernelH": "kernelH",
+              "kernelW": "kernelW",
+              "strideH": "strideH",
+              "strideW": "strideW",
+              "dilationH": "dilationH",
+              "dilationW": "dilationW",
+              "padTop": "padTop",
+              "padLeft": "padLeft",
+              "inH": "dim(shapes.X, 2)",
+              "inW": "dim(shapes.X, 3)",
+              "outH": "dim(shapes.Y, 2)",
+              "outW": "dim(shapes.Y, 3)",
+              "p": "attrs.p"
+            }
+          },
+          "bindings": "flatDirect",
+          "dispatch": { "threads": "numel(shapes.Y)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
+        }
+      ]
+    },
+    {
+      "id": "ncl1d_k4s2_p2_x4",
+      "priority": 30,
+      "when": ["poolBaseOk", "ranks.X == 3", "attrs.auto_pad == \"NOTSET\"", "attrs.p == 2", "kernelW == 4", "strideW == 2", "dilationW == 1", "padLeft == 0", "padRight == 0", "attrs.ceil_mode == 0", "dim(shapes.X, 2) % 4 == 0", "dim(shapes.Y, 2) == floor((dim(shapes.X, 2) - 4) / 2) + 1"],
+      "constants": { "xVec4": "\"vec4<f16>\" if dtypes.T == \"f16\" else \"vec4<f32>\"" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "LpPool.Ncl1dK4s2P2X4",
+          "source": {
+            "shader": "pool-ncl1d-x4.wgsl.jinja",
+            "inputs": { "op": "\"lp\"", "inW": "dim(shapes.X, 2)", "outW": "dim(shapes.Y, 2)" }
+          },
+          "bindings": "rank1X4",
+          "dispatch": {
+            "threads": "dim(shapes.X, 0) * dim(shapes.X, 1) * ceilDiv(dim(shapes.Y, 2), 4)",
+            "workgroupSize": "tunables.WORKGROUP_SIZE"
+          }
+        }
+      ]
+    },
+    {
+      "id": "ncl1d",
+      "priority": 10,
+      "when": ["poolBaseOk", "ranks.X == 3", "kernelW > 0", "strideW > 0"],
+      "constants": { "pIs1": "attrs.p == 1", "pIs2": "attrs.p == 2" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "LpPool.NCL1d",
+          "source": { "shader": "pool-window-nd.wgsl.jinja", "inputs": { "op": "\"lp\"", "spatialRank": 1 } },
+          "bindings": "rank1Direct",
+          "dispatch": { "threads": "numel(shapes.Y)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
+        }
+      ]
+    },
+    {
+      "id": "nchw2d_separable",
+      "priority": 22,
+      "when": ["(attrs.auto_pad != \"SAME_UPPER\" and attrs.auto_pad != \"SAME_LOWER\")", "poolBaseOk", "ranks.X == 4", "kernelH > 0", "kernelW > 0", "strideH > 0", "strideW > 0", "kernelH > 1", "kernelW > 1", "kernelH * kernelW >= tunables.SEPARABLE_MIN_WINDOW", "separableFits"],
+      "constants": { "pIs1": "attrs.p == 1", "pIs2": "attrs.p == 2" },
+      "intermediates": [{ "id": "tmp", "dtype": "float32", "shape": "[separableTmpElements]" }],
+      "passes": [
+        {
+          "id": "reduce_w",
+          "name": "LpPool.SeparableW",
+          "source": { "shader": "pool2d-nchw-separable.wgsl.jinja", "inputs": { "op": "\"lp\"", "axis": "\"w\"" } },
+          "bindings": "rank2SeparableW",
+          "dispatch": { "threads": "separableTmpElements", "workgroupSize": "tunables.WORKGROUP_SIZE" }
+        },
+        {
+          "id": "reduce_h",
+          "name": "LpPool.SeparableH",
+          "source": { "shader": "pool2d-nchw-separable.wgsl.jinja", "inputs": { "op": "\"lp\"", "axis": "\"h\"" } },
+          "bindings": "rank2SeparableH",
+          "dispatch": { "threads": "numel(shapes.Y)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
+        }
+      ]
+    },
+    {
+      "id": "nchw2d",
+      "when": ["(attrs.auto_pad != \"SAME_UPPER\" and attrs.auto_pad != \"SAME_LOWER\")", "poolBaseOk", "ranks.X == 4", "kernelH > 0", "kernelW > 0", "strideH > 0", "strideW > 0"],
+      "constants": { "pIs1": "attrs.p == 1", "pIs2": "attrs.p == 2" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "LpPool.NCHW2d",
+          "source": { "shader": "pool-window-nd.wgsl.jinja", "inputs": { "op": "\"lp\"", "spatialRank": 2 } },
+          "bindings": "rank2Direct",
+          "dispatch": { "threads": "numel(shapes.Y)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
+        }
+      ]
+    },
+    {
+      "id": "nchw2d_autopad",
+      "priority": 15,
+      "when": ["(attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\")", "poolBaseOk", "ranks.X == 4", "kernelH > 0", "kernelW > 0", "strideH > 0", "strideW > 0", "dim(shapes.Y, 2) == ceil(dim(shapes.X, 2) / strideH)", "dim(shapes.Y, 3) == ceil(dim(shapes.X, 3) / strideW)"],
+      "constants": { "pIs1": "attrs.p == 1", "pIs2": "attrs.p == 2" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "LpPool.NCHW2dAutoPad",
+          "source": { "shader": "pool-window-nd.wgsl.jinja", "inputs": { "op": "\"lp\"", "spatialRank": 2 } },
+          "bindings": "rank2Direct",
+          "dispatch": { "threads": "numel(shapes.Y)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
+        }
+      ]
+    },
+    {
+      "id": "ncdhw3d",
+      "priority": 10,
+      "supersededBy": ["global_window_parallel"],
+      "when": ["poolBaseOk", "ranks.X == 5", "kernelD > 0", "kernelH > 0", "kernelW > 0", "strideD > 0", "strideH > 0", "strideW > 0"],
+      "passes": [
+        {
+          "id": "main",
+          "name": "LpPool.NCDHW3d",
+          "source": { "shader": "pool-window-nd.wgsl.jinja", "inputs": { "op": "\"lp\"", "spatialRank": 3 } },
+          "bindings": "rank3Direct",
+          "dispatch": { "threads": "numel(shapes.Y)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
+        }
+      ]
+    }
+  ]
+}

build/webgpu/metadata.json ADDED Viewed

	@@ -0,0 +1,23 @@

+{
+  "name": "ai.onnx.LpPool",
+  "id": "_ai_onnx_lppool_webgpu_cd77da5",
+  "version": 1,
+  "license": "Apache-2.0",
+  "backend": { "type": "webgpu" },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "bench.json": "w6aN2KgWABkA/ms17d5eRGnNY4W16SbuGU0Bv4Ua+PM=",
+      "manifest.json": "u27G/WowR9lMUG+lwNF4ur0POWFvxA1Pj1D/why0G3w=",
+      "pool-global-reduction.wgsl.jinja": "6v3XBuPgokHWHa3DRfhk8PyQ2km8soKnuJ9NNZUgSTo=",
+      "pool-ncl1d-x4.wgsl.jinja": "uZW9X2c3ilCWzT0YCg9/fs/L3wNTS5j4aRO31UQpFbA=",
+      "pool-window-nd.wgsl.jinja": "h6hIDb6qxMwApe1Wcw4x9cEnjWY1WpYWk3vR8ecYF/w=",
+      "pool-window-unroll.wgsl.jinja": "HVLK0vaS22ommayIv3VISzdyU3X1lgtFYG9N5WCOUUM=",
+      "pool2d-nchw-k2s2-vec4.wgsl.jinja": "Oq+ar9au4O/oG5ZU8MtSJYLKGXB/GupLNnYHEgmuaUI=",
+      "pool2d-nchw-separable.wgsl.jinja": "0Av3hMIG4Rnrut2XitTZrqU8l7ih1QV0/w01vyAsbe4=",
+      "test.json": "5NKe4iEQZXXCRAfrhel4tIqfaUs+/u9uZCKdPbPU4oI="
+    }
+  },
+  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
+  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.LpPool" }
+}

build/webgpu/pool-global-reduction.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,120 @@

+{% if usesF16 %}
+enable f16;
+{% endif %}
+{% if source.useSubgroups %}
+enable subgroups;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+// Workgroup-parallel global pooling: one workgroup reduces one contiguous
+// (n, c) spatial plane (NC[D]HW layout makes the plane contiguous).
+// Threads stride the plane (vec4 loads when the plane size is divisible by 4),
+// then combine the per-thread partials. With subgroups: a subgroup collective
+// + a shared-memory fold of the per-subgroup partials. Without subgroups, a
+// full workgroup tree keeps workgroup-per-plane parallelism
+// instead of falling to the one-thread-per-plane serial scalar kernel. The
+// epilogue finalizes average / max / Lp. Sum reassociates;
+// max is exact.
+const WG: u32 = {{ source.workgroupSize }}u;
+{% if source.useSubgroups %}
+// WebGPU subgroups are at least 4 wide, so a workgroup holds at most WG / 4 subgroups.
+const MAX_SUBGROUPS: u32 = WG / 4u;
+var<workgroup> wg_partials: array<f32, MAX_SUBGROUPS>;
+{% else %}
+var<workgroup> wg_tree: array<f32, WG>;
+{% endif %}
+{% if source.op == "lp" and source.p != 1 and source.p != 2 %}
+const P: f32 = f32({{ source.p }});
+{% endif %}
+@compute @workgroup_size(WG)
+fn main(@builtin(workgroup_id) wid: vec3<u32>,
+        @builtin(local_invocation_id) lid: vec3<u32>
+{%- if source.useSubgroups %},
+        @builtin(subgroup_size) sg_size: u32
+{%- endif %}) {
+  let plane = wid.x;
+  let tid = lid.x;
+{% if source.useSubgroups %}
+  // Seed cross-subgroup slots with the reduction identity so unwritten slots
+  // (when the workgroup has fewer than MAX_SUBGROUPS subgroups) are harmless.
+  if (tid < MAX_SUBGROUPS) {
+    wg_partials[tid] = 0.0;
+  }
+  workgroupBarrier();
+{% endif %}
+  var partial = 0.0;
+{% if source.vec4 %}
+  let spatialVec = params.spatial / 4u;
+  let base = plane * spatialVec;
+  var accv = vec4<f32>(0.0);
+  for (var i = tid; i < spatialVec; i = i + WG) {
+    let v = vec4<f32>(x[base + i]);
+{% if source.p == 2 %}
+    accv = accv + v * v;
+{% elif source.p == 1 %}
+    accv = accv + abs(v);
+{% else %}
+    accv = accv + pow(abs(v), vec4<f32>(P));
+{% endif %}
+  }
+  partial = accv.x + accv.y + accv.z + accv.w;
+{% else %}
+  let base = plane * params.spatial;
+  for (var i = tid; i < params.spatial; i = i + WG) {
+    let v = f32(x[base + i]);
+{% if source.p == 2 %}
+    partial = partial + v * v;
+{% elif source.p == 1 %}
+    partial = partial + abs(v);
+{% else %}
+    partial = partial + pow(abs(v), P);
+{% endif %}
+  }
+{% endif %}
+{% if source.useSubgroups %}
+  let sg_sum = subgroupAdd(partial);
+  if (subgroupElect()) {
+    wg_partials[tid / sg_size] = sg_sum;
+  }
+  workgroupBarrier();
+  if (tid == 0u) {
+    let nsg = min((WG + sg_size - 1u) / sg_size, MAX_SUBGROUPS);
+    var total = 0.0;
+    for (var i = 0u; i < nsg; i = i + 1u) {
+      total = total + wg_partials[i];
+    }
+{% if source.p == 2 %}
+    y[plane] = {{ scalar }}(sqrt(total));
+{% elif source.p == 1 %}
+    y[plane] = {{ scalar }}(total);
+{% else %}
+    y[plane] = {{ scalar }}(pow(total, 1.0 / P));
+{% endif %}
+  }
+{% else %}
+  // No-subgroups fallback: full workgroup-tree reduction of the per-thread partials.
+  wg_tree[tid] = partial;
+  workgroupBarrier();
+  for (var step = WG >> 1u; step > 0u; step = step >> 1u) {
+    if (tid < step) {
+      wg_tree[tid] = wg_tree[tid] + wg_tree[tid + step];
+    }
+    workgroupBarrier();
+  }
+  if (tid == 0u) {
+    let total = wg_tree[0];
+{% if source.p == 2 %}
+    y[plane] = {{ scalar }}(sqrt(total));
+{% elif source.p == 1 %}
+    y[plane] = {{ scalar }}(total);
+{% else %}
+    y[plane] = {{ scalar }}(pow(total, 1.0 / P));
+{% endif %}
+  }
+{% endif %}
+}

build/webgpu/pool-ncl1d-x4.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,45 @@

+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+// Four adjacent NCL outputs per invocation for non-overlapping max k2/s2 and
+// overlapping L2 k4/s2 windows. Both layouts consume aligned vec4 input blocks.
+// The L2 route reuses shared middle samples instead of loading every window
+// independently.
+const IN_W: u32 = {{ source.inW }}u;
+const OUT_W: u32 = {{ source.outW }}u;
+const BLOCKS_PER_PLANE: u32 = (OUT_W + 3u) / 4u;
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(
+  @builtin(global_invocation_id) gid: vec3<u32>,
+  @builtin(num_workgroups) nwg: vec3<u32>,
+) {
+  let block = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
+  if (block >= params.blockCount) { return; }
+  let plane = block / BLOCKS_PER_PLANE;
+  let outputBlock = block - plane * BLOCKS_PER_PLANE;
+  let ow = outputBlock * 4u;
+  let inputVec = plane * (IN_W / 4u) + ow / 2u;
+  let outputBase = plane * OUT_W + ow;
+  let a = vec4<f32>(x[inputVec]);
+  let b = vec4<f32>(x[inputVec + 1u]);
+  // k4/s2 L2: neighboring windows share their middle pair. The final block
+  // may have only three outputs, in which case no third vec4 read is needed.
+  let aa = a * a;
+  let bb = b * b;
+  y[outputBase] = {{ scalar }}(sqrt(aa.x + aa.y + aa.z + aa.w));
+  if (ow + 1u < OUT_W) {
+    y[outputBase + 1u] = {{ scalar }}(sqrt(aa.z + aa.w + bb.x + bb.y));
+  }
+  if (ow + 2u < OUT_W) {
+    y[outputBase + 2u] = {{ scalar }}(sqrt(bb.x + bb.y + bb.z + bb.w));
+  }
+  if (ow + 3u < OUT_W) {
+    let c = vec4<f32>(x[inputVec + 2u]);
+    let cc = c * c;
+    y[outputBase + 3u] = {{ scalar }}(sqrt(bb.z + bb.w + cc.x + cc.y));
+  }
+}

build/webgpu/pool-window-nd.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,70 @@

+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+// Direct {{ source.spatialRank }}D {{ source.op }}-pool traversal for a channels-first tensor.
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
+  let index = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
+  if (index >= params.count) { return; }
+  let ow = index % params.outW;
+  var t = index / params.outW;
+{% if source.spatialRank >= 2 %}
+  let oh = t % params.outH;
+  t /= params.outH;
+{% endif %}
+{% if source.spatialRank == 3 %}
+  let od = t % params.outD;
+  t /= params.outD;
+{% endif %}
+  let channel = t % params.channels;
+  let batch = t / params.channels;
+  var acc = 0.0;
+{% if source.spatialRank == 3 %}
+  for (var kd = 0u; kd < params.kernelD; kd += 1u) {
+    let id = i32(od * params.strideD + kd * params.dilationD) - params.padD;
+    if (id < 0 || id >= i32(params.inD)) { continue; }
+{% endif %}
+{% if source.spatialRank >= 2 %}
+    for (var kh = 0u; kh < params.kernelH; kh += 1u) {
+      let ih = i32(oh * params.strideH + kh * params.dilationH) - params.padH;
+      if (ih < 0 || ih >= i32(params.inH)) { continue; }
+{% endif %}
+      for (var kw = 0u; kw < params.kernelW; kw += 1u) {
+        let iw = i32(ow * params.strideW + kw * params.dilationW) - params.padW;
+        if (iw < 0 || iw >= i32(params.inW)) { continue; }
+{% if source.spatialRank == 1 %}
+        let xIndex = (batch * params.channels + channel) * params.inW + u32(iw);
+{% elif source.spatialRank == 2 %}
+        let xIndex = ((batch * params.channels + channel) * params.inH + u32(ih)) * params.inW + u32(iw);
+{% else %}
+        let xIndex = (((batch * params.channels + channel) * params.inD + u32(id)) * params.inH + u32(ih)) * params.inW + u32(iw);
+{% endif %}
+{% if pIs2 %}
+        let value = f32(x[xIndex]);
+        acc += value * value;
+{% elif pIs1 %}
+        acc += abs(f32(x[xIndex]));
+{% else %}
+        acc += pow(abs(f32(x[xIndex])), params.p);
+{% endif %}
+      }
+{% if source.spatialRank >= 2 %}
+    }
+{% endif %}
+{% if source.spatialRank == 3 %}
+  }
+{% endif %}
+{% if pIs2 %}
+  y[index] = {{ scalar }}(sqrt(acc));
+{% elif pIs1 %}
+  y[index] = {{ scalar }}(acc);
+{% else %}
+  y[index] = {{ scalar }}(pow(acc, 1.0 / params.p));
+{% endif %}
+}

build/webgpu/pool-window-unroll.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,75 @@

+// Compile-time-unrolled 2D NCHW pooling for small windows. Accumulation remains
+// kh-major/kw-minor, matching the runtime-loop fallbacks.
+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+const OUT_W: u32 = {{ source.outW }}u;
+const OUT_H: u32 = {{ source.outH }}u;
+const IN_H: i32 = {{ source.inH }};
+const IN_W: i32 = {{ source.inW }};
+const IN_W_U: u32 = {{ source.inW }}u;
+const IN_PLANE: u32 = {{ source.inH * source.inW }}u;
+const STRIDE_H: u32 = {{ source.strideH }}u;
+const STRIDE_W: u32 = {{ source.strideW }}u;
+const PAD_TOP: i32 = {{ source.padTop }};
+const PAD_LEFT: i32 = {{ source.padLeft }};
+{% if source.op == "lp" and source.p != 1 and source.p != 2 %}
+const P: f32 = f32({{ source.p }});
+{% endif %}
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(
+  @builtin(global_invocation_id) gid: vec3<u32>,
+  @builtin(num_workgroups) nwg: vec3<u32>
+) {
+  // Fold very large flat outputs into the dispatch's second dimension.
+  let index = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
+  if (index >= params.count) {
+    return;
+  }
+  let ow = index % OUT_W;
+  var plane = index / OUT_W;
+  let oh = plane % OUT_H;
+  plane /= OUT_H;
+  let planeBase = plane * IN_PLANE;
+  let startH = i32(oh * STRIDE_H) - PAD_TOP;
+  let startW = i32(ow * STRIDE_W) - PAD_LEFT;
+{% for kh in range(source.kernelH) %}
+  let ih_{{ kh }} = startH + {{ kh * source.dilationH }};
+{% endfor %}
+{% for kw in range(source.kernelW) %}
+  let iw_{{ kw }} = startW + {{ kw * source.dilationW }};
+{% endfor %}
+  var result = 0.0;
+{% for kh in range(source.kernelH) %}
+  if (ih_{{ kh }} >= 0 && ih_{{ kh }} < IN_H) {
+    let rowBase{{ kh }} = planeBase + u32(ih_{{ kh }}) * IN_W_U;
+{% for kw in range(source.kernelW) %}
+    if (iw_{{ kw }} >= 0 && iw_{{ kw }} < IN_W) {
+      let value_{{ kh }}_{{ kw }} = f32(x[rowBase{{ kh }} + u32(iw_{{ kw }})]);
+{% if source.op == "lp" and source.p == 2 %}
+      result += value_{{ kh }}_{{ kw }} * value_{{ kh }}_{{ kw }};
+{% elif source.op == "lp" and source.p == 1 %}
+      result += abs(value_{{ kh }}_{{ kw }});
+{% else %}
+      result += pow(abs(value_{{ kh }}_{{ kw }}), P);
+{% endif %}
+    }
+{% endfor %}
+  }
+{% endfor %}
+{% if source.op == "lp" and source.p == 2 %}
+  let output = sqrt(result);
+{% elif source.op == "lp" and source.p == 1 %}
+  let output = result;
+{% else %}
+  let output = pow(result, 1.0 / P);
+{% endif %}
+  y[index] = {{ scalar }}(output);
+}

build/webgpu/pool2d-nchw-k2s2-vec4.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,65 @@

+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+// 2D NCHW LpPool specialised for the dense 2x2 / stride-2 / no-pad /
+// dilation-1 window. Each thread produces two adjacent output columns from one
+// aligned vec4 load per input row. A vec4 over
+// input columns [4*tileCol .. 4*tileCol+3] holds the two side-by-side 2-wide
+// windows. The p-norm is accumulated in f32 (sum of
+// |x|^p over the 2x2 window, same four elements / same count as the scalar
+// fallback) and reduced by the p-th root; p=2 specialises to sum-of-squares +
+// sqrt, p=1 to sum-of-abs.
+const OUT_W: u32 = {{ source.outW }}u;        // == IN_W / 2 (no pad, stride 2)
+const OUT_H: u32 = {{ source.outH }}u;
+const OUT_W_HALF: u32 = {{ source.outWHalf }}u;
+const IN_W_VEC4: u32 = {{ source.inWVec4 }}u; // input row length in vec4 units
+const IN_PLANE_VEC4: u32 = {{ source.inPlaneVec4 }}u;
+{% if source.p != 1 and source.p != 2 %}
+const P: f32 = f32({{ source.p }});
+{% endif %}
+{% if source.p == 2 %}
+fn term(v: {{ scalar }}) -> f32 { let f = f32(v); return f * f; }
+{% elif source.p == 1 %}
+fn term(v: {{ scalar }}) -> f32 { return abs(f32(v)); }
+{% else %}
+fn term(v: {{ scalar }}) -> f32 { return pow(abs(f32(v)), P); }
+{% endif %}
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
+  // 2D-folded flat tile index (one tile = two output columns).
+  let tile = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
+  if (tile >= params.tileCount) {
+    return;
+  }
+  let tileCol = tile % OUT_W_HALF;        // pair-of-columns within the row
+  var t = tile / OUT_W_HALF;
+  let oh = t % OUT_H;
+  let plane = t / OUT_H;                   // batch * channels + channel
+  let row0Vec4 = plane * IN_PLANE_VEC4 + (oh * 2u) * IN_W_VEC4 + tileCol;
+  let row1Vec4 = row0Vec4 + IN_W_VEC4;
+  let r0 = x[row0Vec4];
+  let r1 = x[row1Vec4];
+  // Column pair 0 spans vec4 lanes 0,1; column pair 1 spans lanes 2,3.
+  let acc0 = term(r0.x) + term(r0.y) + term(r1.x) + term(r1.y);
+  let acc1 = term(r0.z) + term(r0.w) + term(r1.z) + term(r1.w);
+  let outBase = (plane * OUT_H + oh) * OUT_W + tileCol * 2u;
+{% if source.p == 2 %}
+  y[outBase] = {{ scalar }}(sqrt(acc0));
+  y[outBase + 1u] = {{ scalar }}(sqrt(acc1));
+{% elif source.p == 1 %}
+  y[outBase] = {{ scalar }}(acc0);
+  y[outBase + 1u] = {{ scalar }}(acc1);
+{% else %}
+  y[outBase] = {{ scalar }}(pow(acc0, 1.0 / P));
+  y[outBase + 1u] = {{ scalar }}(pow(acc1, 1.0 / P));
+{% endif %}
+}

build/webgpu/pool2d-nchw-separable.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,73 @@

+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+// Separable NCHW pooling: one spatial axis per pass.
+//
+// Both reductions here are associative, so a kernelH x kernelW window
+// factorizes into a reduction along W followed by one along H. The direct
+// kernel reads kernelH*kernelW inputs per output; separating them reads kernelW
+// into an intermediate cell that neighbouring outputs then share, plus kernelH
+// per output.
+//
+// The intermediate is f32 because the direct kernel also accumulates in f32
+// and casts once at the end, including for i32 input.
+//
+// Max pooling is bit-identical to the direct kernel because max is idempotent
+// as well as associative. Lp pooling is within tolerance but not bit-identical:
+// summing each row first regroups f32 additions relative to the direct kernel.
+//
+// `axis` selects the pass: "w" reduces x into tmp, and "h" reduces tmp into y.
+// AveragePool is not wired here because its denominator would require valid-tap
+// counts to travel alongside the sum.
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
+  let index = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
+{% if source.axis == "w" %}
+  if (index >= params.tmpCount) { return; }
+  // One cell per (batch, channel, input row, output column).
+  let ow = index % params.outW;
+  let t = index / params.outW;
+  let ih = t % params.inH;
+  let nc = t / params.inH;
+  var acc = 0.0;
+  for (var kw = 0u; kw < params.kernelW; kw += 1u) {
+    let iw = i32(ow * params.strideW + kw * params.dilationW) - params.padW;
+    if (iw < 0 || iw >= i32(params.inW)) { continue; }
+    let value = f32(x[(nc * params.inH + ih) * params.inW + u32(iw)]);
+{% if pIs2 %}
+    acc += value * value;
+{% elif pIs1 %}
+    acc += abs(value);
+{% else %}
+    acc += pow(abs(value), params.p);
+{% endif %}
+  }
+  tmp[index] = acc;
+{% else %}
+  if (index >= params.count) { return; }
+  let ow = index % params.outW;
+  let t = index / params.outW;
+  let oh = t % params.outH;
+  let nc = t / params.outH;
+  var acc = 0.0;
+  for (var kh = 0u; kh < params.kernelH; kh += 1u) {
+    let ih = i32(oh * params.strideH + kh * params.dilationH) - params.padH;
+    if (ih < 0 || ih >= i32(params.inH)) { continue; }
+    let partial = tmp[(nc * params.inH + u32(ih)) * params.outW + ow];
+    acc += partial;
+  }
+{% if pIs2 %}
+  y[index] = {{ scalar }}(sqrt(acc));
+{% elif pIs1 %}
+  y[index] = {{ scalar }}(acc);
+{% else %}
+  y[index] = {{ scalar }}(pow(acc, 1.0 / params.p));
+{% endif %}
+{% endif %}
+}

build/webgpu/test.json ADDED Viewed

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