sync 6fdf6301e2bb

Browse files

Files changed (8) hide show

README.md +6 -5
build/webgpu/bench.json +188 -4
build/webgpu/datamove-transpose-2d-tiled-scalar.wgsl.jinja +7 -7
build/webgpu/datamove-transpose-2d-tiled.wgsl.jinja +21 -11
build/webgpu/manifest.json +65 -43
build/webgpu/metadata.json +10 -9
build/webgpu/test.json +457 -5
build/webgpu/transpose.wgsl.jinja +2 -4

README.md CHANGED Viewed

@@ -40,12 +40,13 @@ Attributes and default values (overridable per request):
 | Variable | Allowed dtypes |
 | --- | --- |
-| `T` | `float32`, `float16`, `int32`, `int16`, `uint32`, `uint8`, `int8`, `bool` |
 ## Implementation variants
 One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
 - `tiled_scalar` — Stages ragged two-dimensional transpose planes through a padded workgroup tile so reads and writes remain coalesced when vec4 alignment is unavailable.
 - `inner_vec4` — Vectorizes the contiguous innermost dimension for permutations that leave that axis in place. It assigns multiple vectors per invocation only when enough invocations remain to keep the dispatch populated.
@@ -54,7 +55,7 @@ One implementation is selected per call from the device capabilities, the reques
 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
-- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 - [`datamove-flat-copy.wgsl.jinja`](build/webgpu/datamove-flat-copy.wgsl.jinja)
 - [`datamove-transpose-2d-tiled-scalar.wgsl.jinja`](build/webgpu/datamove-transpose-2d-tiled-scalar.wgsl.jinja)
 - [`datamove-transpose-2d-tiled.wgsl.jinja`](build/webgpu/datamove-transpose-2d-tiled.wgsl.jinja)
@@ -64,7 +65,7 @@ One implementation is selected per call from the device capabilities, the reques
 ## Use with `@huggingface/kernels`
 ```sh
-npm install --save-exact @huggingface/kernels@0.0.1-preview.2
 ```
 Outputs with inferable metadata are allocated automatically. Explicit `outputs` entries request optional results or provide metadata that cannot be inferred from the supplied inputs and attributes.
@@ -83,7 +84,7 @@ import { getKernel } from "@huggingface/kernels";
 const kernel = await getKernel("webgpu-kernels/ai.onnx.Transpose", { version: 1 });
 // Explicit destinations request optional results or supply metadata that cannot be inferred.
-const { y } = await kernel({ x: { data: xData, shape: [] } }, {
-  outputs: { y: { shape: [], dtype: "float32" } },
 });
 ```

 | Variable | Allowed dtypes |
 | --- | --- |
+| `T` | `float32`, `float16`, `int32`, `int16`, `uint32`, `uint8`, `int8`, `bool`, `int64` |
 ## Implementation variants
 One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
+- `strided_planes_vec4` — Tiles independent physical strided planes while preserving the intervening axis. Both global IO directions are contiguous inside the tile. Demoted when padding occupies more than half the tile slots.
 - `tiled_scalar` — Stages ragged two-dimensional transpose planes through a padded workgroup tile so reads and writes remain coalesced when vec4 alignment is unavailable.
 - `inner_vec4` — Vectorizes the contiguous innermost dimension for permutations that leave that axis in place. It assigns multiple vectors per invocation only when enough invocations remain to keep the dispatch populated.
 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark cases
 - [`datamove-flat-copy.wgsl.jinja`](build/webgpu/datamove-flat-copy.wgsl.jinja)
 - [`datamove-transpose-2d-tiled-scalar.wgsl.jinja`](build/webgpu/datamove-transpose-2d-tiled-scalar.wgsl.jinja)
 - [`datamove-transpose-2d-tiled.wgsl.jinja`](build/webgpu/datamove-transpose-2d-tiled.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.3
 ```
 Outputs with inferable metadata are allocated automatically. Explicit `outputs` entries request optional results or provide metadata that cannot be inferred from the supplied inputs and attributes.
 const kernel = await getKernel("webgpu-kernels/ai.onnx.Transpose", { version: 1 });
 // Explicit destinations request optional results or supply metadata that cannot be inferred.
+const { y } = await kernel({ x: { data: xData, shape: [2, 3] } }, {
+  outputs: { y: { shape: [3, 2], dtype: "float32" } },
 });
 ```

build/webgpu/bench.json CHANGED Viewed

@@ -195,7 +195,7 @@
       "preset": "all",
       "provenance": {
         "source": "Depth Anything V2 Small neck reassemble",
-        "notes": "Probe pair (a): the real 37x49 patch grid has 1813 rows and is not vec4-aligned. Together with pair (b), it verifies that extent parity does not determine whether the transpose is staged."
       },
       "vars": { "count": 696192, "dtype": "float32" },
       "attrs": { "perm": [1, 0] },
@@ -220,7 +220,7 @@
       "name": "transpose-f32-rank6-swin-window-partition-perm013245-pathology",
       "preset": "edge",
       "provenance": {
-        "notes": "Rank-6 window partition (swap of the two window-grid axes with a contiguous 96-wide last axis) at a batch-8 Swin 56x56 stage. Rank 6 is outside every tiled and inner-vec4 pattern, so the generic strided kernel is the selected route; the permutation is its own inverse, so one case covers both directions."
       },
       "vars": { "count": 2408448, "dtype": "float32" },
       "attrs": { "perm": [0, 1, 3, 2, 4, 5] },
@@ -246,7 +246,7 @@
       "name": "transpose-f32-rank4-perm0321-8x64x197x64-pathology",
       "preset": "edge",
       "provenance": {
-        "notes": "Rank-4 permutation [0,3,2,1] at a 197-token, 64-head-dim attention shape: axis 3 leaves the last position and the permutation is none of the three tiled rank-4 patterns, so the generic strided kernel is the selected route."
       },
       "vars": { "count": 6455296, "dtype": "float32" },
       "attrs": { "perm": [0, 3, 2, 1] },
@@ -258,13 +258,197 @@
       "name": "transpose-f32-rank4-perm0213-8x64x197x64-control",
       "preset": "edge",
       "provenance": {
-        "notes": "Control for the [0,3,2,1] case: the same tensor under the attention-layout permutation [0,2,1,3] keeps the innermost axis in place and takes the inner-vec4 route."
       },
       "vars": { "count": 6455296, "dtype": "float32" },
       "attrs": { "perm": [0, 2, 1, 3] },
       "inputs": { "x": { "shape": [8, 64, 197, 64], "dtype": "float32", "dist": "normal", "seed": 157, "scale": 0.2 } },
       "outputs": { "y": { "shape": [8, 197, 64, 64], "dtype": "float32" } },
       "bench": { "metrics": [{ "type": "bandwidth", "value": "args.count * dtypeBytes(args.dtype) * 2" }] }
     }
   ]
 }

       "preset": "all",
       "provenance": {
         "source": "Depth Anything V2 Small neck reassemble",
+        "notes": "A 1813x384 transpose (the 37x49 Depth Anything V2 Small patch grid) has a row count that is not a multiple of four; its 1812-row sibling in this file is the vec4-aligned control at the same column width."
       },
       "vars": { "count": 696192, "dtype": "float32" },
       "attrs": { "perm": [1, 0] },
       "name": "transpose-f32-rank6-swin-window-partition-perm013245-pathology",
       "preset": "edge",
       "provenance": {
+        "notes": "Rank-six window partition swaps a within-window axis with a window-grid axis while retaining a contiguous 96-wide final axis at a batch-eight Swin 56x56 stage. The permutation is its own inverse, so one case covers both directions."
       },
       "vars": { "count": 2408448, "dtype": "float32" },
       "attrs": { "perm": [0, 1, 3, 2, 4, 5] },
       "name": "transpose-f32-rank4-perm0321-8x64x197x64-pathology",
       "preset": "edge",
       "provenance": {
+        "notes": "Rank-four permutation [0,3,2,1] at a 197-token, 64-wide attention shape moves the final axis and checks strided output indexing."
       },
       "vars": { "count": 6455296, "dtype": "float32" },
       "attrs": { "perm": [0, 3, 2, 1] },
       "name": "transpose-f32-rank4-perm0213-8x64x197x64-control",
       "preset": "edge",
       "provenance": {
+        "notes": "Control for the [0,3,2,1] case: the same [8,64,197,64] shape under the attention-layout permutation [0,2,1,3] keeps its innermost axis."
       },
       "vars": { "count": 6455296, "dtype": "float32" },
       "attrs": { "perm": [0, 2, 1, 3] },
       "inputs": { "x": { "shape": [8, 64, 197, 64], "dtype": "float32", "dist": "normal", "seed": 157, "scale": 0.2 } },
       "outputs": { "y": { "shape": [8, 197, 64, 64], "dtype": "float32" } },
       "bench": { "metrics": [{ "type": "bandwidth", "value": "args.count * dtypeBytes(args.dtype) * 2" }] }
+    },
+    {
+      "name": "strided_large_float32_95x201x127",
+      "preset": "all",
+      "attrs": { "perm": [2, 1, 0] },
+      "inputs": { "x": { "shape": [95, 201, 127], "dtype": "float32", "dist": "normal", "seed": 842, "scale": 0.2 } },
+      "outputs": { "y": { "shape": [127, 201, 95], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 19400520 }] }
+    },
+    {
+      "name": "strided_large_float32_8x63x197x65",
+      "preset": "all",
+      "attrs": { "perm": [0, 3, 2, 1] },
+      "inputs": { "x": { "shape": [8, 63, 197, 65], "dtype": "float32", "dist": "normal", "seed": 842, "scale": 0.2 } },
+      "outputs": { "y": { "shape": [8, 65, 197, 63], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 51629760 }] }
+    },
+    {
+      "name": "strided_large_float16_95x201x127",
+      "preset": "all",
+      "attrs": { "perm": [2, 1, 0] },
+      "inputs": { "x": { "shape": [95, 201, 127], "dtype": "float16", "dist": "normal", "seed": 842, "scale": 0.2 } },
+      "outputs": { "y": { "shape": [127, 201, 95], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 9700260 }] }
+    },
+    {
+      "name": "strided_large_float16_8x63x197x65",
+      "preset": "all",
+      "attrs": { "perm": [0, 3, 2, 1] },
+      "inputs": { "x": { "shape": [8, 63, 197, 65], "dtype": "float16", "dist": "normal", "seed": 842, "scale": 0.2 } },
+      "outputs": { "y": { "shape": [8, 65, 197, 63], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 25814880 }] }
+    },
+    {
+      "name": "strided_large_float16_8x64x197x64",
+      "preset": "all",
+      "attrs": { "perm": [0, 3, 2, 1] },
+      "inputs": { "x": { "shape": [8, 64, 197, 64], "dtype": "float16", "dist": "normal", "seed": 842, "scale": 0.2 } },
+      "outputs": { "y": { "shape": [8, 64, 197, 64], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 25821184 }] }
+    },
+    {
+      "name": "strided_padding_float32_a8",
+      "preset": "all",
+      "attrs": { "perm": [0, 3, 2, 1] },
+      "inputs": { "x": { "shape": [16, 8, 137, 128], "dtype": "float32", "dist": "normal", "seed": 843, "scale": 0.2 } },
+      "outputs": { "y": { "shape": [16, 128, 137, 8], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 17956864 }] }
+    },
+    {
+      "name": "strided_padding_float32_a16",
+      "preset": "all",
+      "attrs": { "perm": [0, 3, 2, 1] },
+      "inputs": {
+        "x": { "shape": [16, 16, 137, 128], "dtype": "float32", "dist": "normal", "seed": 843, "scale": 0.2 }
+      },
+      "outputs": { "y": { "shape": [16, 128, 137, 16], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 35913728 }] }
+    },
+    {
+      "name": "strided_padding_float32_a20",
+      "preset": "all",
+      "attrs": { "perm": [0, 3, 2, 1] },
+      "inputs": {
+        "x": { "shape": [16, 20, 137, 128], "dtype": "float32", "dist": "normal", "seed": 843, "scale": 0.2 }
+      },
+      "outputs": { "y": { "shape": [16, 128, 137, 20], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 44892160 }] }
+    },
+    {
+      "name": "strided_padding_float32_a15",
+      "preset": "all",
+      "attrs": { "perm": [0, 3, 2, 1] },
+      "inputs": {
+        "x": { "shape": [16, 15, 137, 128], "dtype": "float32", "dist": "normal", "seed": 843, "scale": 0.2 }
+      },
+      "outputs": { "y": { "shape": [16, 128, 137, 15], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 33669120 }] }
+    },
+    {
+      "name": "strided_padding_float32_a17",
+      "preset": "all",
+      "attrs": { "perm": [0, 3, 2, 1] },
+      "inputs": {
+        "x": { "shape": [16, 17, 137, 128], "dtype": "float32", "dist": "normal", "seed": 843, "scale": 0.2 }
+      },
+      "outputs": { "y": { "shape": [16, 128, 137, 17], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 38158336 }] }
+    },
+    {
+      "name": "strided_padding_float16_a8",
+      "preset": "all",
+      "attrs": { "perm": [0, 3, 2, 1] },
+      "inputs": { "x": { "shape": [16, 8, 137, 128], "dtype": "float16", "dist": "normal", "seed": 843, "scale": 0.2 } },
+      "outputs": { "y": { "shape": [16, 128, 137, 8], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 8978432 }] }
+    },
+    {
+      "name": "strided_padding_float16_a15",
+      "preset": "all",
+      "attrs": { "perm": [0, 3, 2, 1] },
+      "inputs": {
+        "x": { "shape": [16, 15, 137, 128], "dtype": "float16", "dist": "normal", "seed": 843, "scale": 0.2 }
+      },
+      "outputs": { "y": { "shape": [16, 128, 137, 15], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 16834560 }] }
+    },
+    {
+      "name": "strided_padding_float16_a16",
+      "preset": "all",
+      "attrs": { "perm": [0, 3, 2, 1] },
+      "inputs": {
+        "x": { "shape": [16, 16, 137, 128], "dtype": "float16", "dist": "normal", "seed": 843, "scale": 0.2 }
+      },
+      "outputs": { "y": { "shape": [16, 128, 137, 16], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 17956864 }] }
+    },
+    {
+      "name": "strided_padding_float16_a17",
+      "preset": "all",
+      "attrs": { "perm": [0, 3, 2, 1] },
+      "inputs": {
+        "x": { "shape": [16, 17, 137, 128], "dtype": "float16", "dist": "normal", "seed": 843, "scale": 0.2 }
+      },
+      "outputs": { "y": { "shape": [16, 128, 137, 17], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 19079168 }] }
+    },
+    {
+      "name": "strided_padding_float16_a20",
+      "preset": "all",
+      "attrs": { "perm": [0, 3, 2, 1] },
+      "inputs": {
+        "x": { "shape": [16, 20, 137, 128], "dtype": "float16", "dist": "normal", "seed": 843, "scale": 0.2 }
+      },
+      "outputs": { "y": { "shape": [16, 128, 137, 20], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 22446080 }] }
+    },
+    {
+      "name": "strided_rect_float32_a47c65",
+      "preset": "all",
+      "attrs": { "perm": [0, 3, 2, 1] },
+      "inputs": { "x": { "shape": [8, 47, 137, 65], "dtype": "float32", "dist": "normal", "seed": 845, "scale": 0.2 } },
+      "outputs": { "y": { "shape": [8, 65, 137, 47], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 26786240 }] }
+    },
+    {
+      "name": "strided_rect_float32_a48c65",
+      "preset": "all",
+      "attrs": { "perm": [0, 3, 2, 1] },
+      "inputs": { "x": { "shape": [8, 48, 137, 65], "dtype": "float32", "dist": "normal", "seed": 845, "scale": 0.2 } },
+      "outputs": { "y": { "shape": [8, 65, 137, 48], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 27356160 }] }
+    },
+    {
+      "name": "strided_fulltile_float32_a33",
+      "preset": "all",
+      "attrs": { "perm": [0, 3, 2, 1] },
+      "inputs": { "x": { "shape": [8, 33, 137, 128], "dtype": "float32", "dist": "normal", "seed": 844, "scale": 0.2 } },
+      "outputs": { "y": { "shape": [8, 128, 137, 33], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 37036032 }] }
+    },
+    {
+      "name": "strided_rect_float16_a47c65",
+      "preset": "all",
+      "attrs": { "perm": [0, 3, 2, 1] },
+      "inputs": { "x": { "shape": [8, 47, 137, 65], "dtype": "float16", "dist": "normal", "seed": 845, "scale": 0.2 } },
+      "outputs": { "y": { "shape": [8, 65, 137, 47], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 13393120 }] }
+    },
+    {
+      "name": "strided_rect_float16_a48c65",
+      "preset": "all",
+      "attrs": { "perm": [0, 3, 2, 1] },
+      "inputs": { "x": { "shape": [8, 48, 137, 65], "dtype": "float16", "dist": "normal", "seed": 845, "scale": 0.2 } },
+      "outputs": { "y": { "shape": [8, 65, 137, 48], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 13678080 }] }
+    },
+    {
+      "name": "strided_fulltile_float16_a33",
+      "preset": "all",
+      "attrs": { "perm": [0, 3, 2, 1] },
+      "inputs": { "x": { "shape": [8, 33, 137, 128], "dtype": "float16", "dist": "normal", "seed": 844, "scale": 0.2 } },
+      "outputs": { "y": { "shape": [8, 128, 137, 33], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": 18518016 }] }
     }
   ]
 }

build/webgpu/datamove-transpose-2d-tiled-scalar.wgsl.jinja CHANGED Viewed

@@ -1,13 +1,13 @@
 // Tiled, optionally batched 2-D transpose over scalar elements.
 //
-// A 32x32 tile enters workgroup memory through consecutive row reads and leaves
 // through consecutive transposed-row writes. Guarded edges make this valid for
 // any extent.
 //
 // Layout contract:
 // - src is row-major [batch, ROWS, COLS]
 // - dst is row-major [batch, COLS, ROWS]
-// - dispatch is x = ceil(COLS / 32), y = ceil(ROWS / 32), z = batch
 //
 // The extra column changes the transposed workgroup-memory stride and reduces
 // bank-conflict risk on banked implementations.
@@ -16,12 +16,12 @@
 const ROWS: u32 = {{ rowCount }}u;
 const COLS: u32 = {{ cols }}u;
 const MAT: u32 = {{ rowCount * cols }}u;
-const TILE: u32 = 32u;
-const ROW_STEP: u32 = 8u;
-var<workgroup> tile: array<array<{{ scalar }}, 33>, 32>;
-@compute @workgroup_size(32, 8, 1)
 fn main(
   @builtin(workgroup_id) wid: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>
@@ -30,7 +30,7 @@ fn main(
   let row_base = wid.y * TILE;
   let col_base = wid.x * TILE;
-  // Read: lid.x runs along the source row, so a subgroup covers 32 consecutive
   // source elements.
   for (var i = 0u; i < TILE; i = i + ROW_STEP) {
     let r = row_base + lid.y + i;

 // Tiled, optionally batched 2-D transpose over scalar elements.
 //
+// A TILE x TILE tile enters workgroup memory through consecutive row reads and leaves
 // through consecutive transposed-row writes. Guarded edges make this valid for
 // any extent.
 //
 // Layout contract:
 // - src is row-major [batch, ROWS, COLS]
 // - dst is row-major [batch, COLS, ROWS]
+// - dispatch is x = ceil(COLS / TILE), y = ceil(ROWS / TILE), z = batch
 //
 // The extra column changes the transposed workgroup-memory stride and reduces
 // bank-conflict risk on banked implementations.
 const ROWS: u32 = {{ rowCount }}u;
 const COLS: u32 = {{ cols }}u;
 const MAT: u32 = {{ rowCount * cols }}u;
+const TILE: u32 = {{ tileSize }}u;
+const ROW_STEP: u32 = {{ tileRowStep }}u;
+var<workgroup> tile: array<array<{{ scalar }}, {{ tileSize + 1 }}>, {{ tileSize }}>;
+@compute @workgroup_size({{ tileSize }}, {{ tileRowStep }}, 1)
 fn main(
   @builtin(workgroup_id) wid: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>
   let row_base = wid.y * TILE;
   let col_base = wid.x * TILE;
+  // Read: lid.x runs along the source row, so a subgroup covers consecutive
   // source elements.
   for (var i = 0u; i < TILE; i = i + ROW_STEP) {
     let r = row_base + lid.y + i;

build/webgpu/datamove-transpose-2d-tiled.wgsl.jinja CHANGED Viewed

@@ -1,38 +1,48 @@
-// Tiled 2D transpose (optionally batched) over vec4 elements.
 //
 // Layout contract:
 // - src is row-major [batch, ROWS, COLS], dst is row-major [batch, COLS, ROWS]
 // - ROWS % 4 == 0 and COLS % 4 == 0, so both buffers bind as vec4<T>
-// - dispatch: x = ceil(COLS / 32), y = ceil(ROWS / 32), z = batch
 //
-// Each 8x32 workgroup moves one 32x32 scalar tile through workgroup memory.
 // Global loads and stores each traverse consecutive vec4 addresses. The extra
 // workgroup-memory column changes the transposed stride and reduces
 // bank-conflict risk on banked implementations.
-{{ env.wgsl.resourceDeclarations }}
 const ROWS: u32 = {{ rowCount }}u;
 const COLS: u32 = {{ cols }}u;
 const ROWS_V: u32 = {{ ((rowCount / 4) | int) }}u;
 const COLS_V: u32 = {{ ((cols / 4) | int) }}u;
-const MAT_V: u32 = {{ (((rowCount * cols) / 4) | int) }}u;
-const TILE: u32 = 32u;
-var<workgroup> tile: array<array<{{ scalar }}, 33>, 32>;
-@compute @workgroup_size(8, 32, 1)
 fn main(
   @builtin(workgroup_id) wid: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>
 ) {
-  let batch_base = wid.z * MAT_V;
   let row_base = wid.y * TILE;
   let col_base = wid.x * TILE;
   let r = row_base + lid.y;
   let cv = (col_base / 4u) + lid.x;
   if (r < ROWS && cv < COLS_V) {
-    let v = src[batch_base + r * COLS_V + cv];
     tile[lid.y][lid.x * 4u + 0u] = v.x;
     tile[lid.y][lid.x * 4u + 1u] = v.y;
     tile[lid.y][lid.x * 4u + 2u] = v.z;
@@ -43,7 +53,7 @@ fn main(
   let out_row = col_base + lid.y;
   let rv = (row_base / 4u) + lid.x;
   if (out_row < COLS && rv < ROWS_V) {
-    dst[batch_base + out_row * ROWS_V + rv] = {{ vectorScalar }}(
       tile[lid.x * 4u + 0u][lid.y],
       tile[lid.x * 4u + 1u][lid.y],
       tile[lid.x * 4u + 2u][lid.y],

+{% if transposePlaneMiddle is defined %}
+// Transpose independent physical planes while retaining the intervening axis.
+// Source [B, ROWS, MIDDLE, COLS] maps to [B, COLS, MIDDLE, ROWS].
+// Dispatch z identifies one (batch, middle) plane; both tile IO directions
+// follow contiguous addresses within that plane.
+{% else %}// Tiled 2D transpose (optionally batched) over vec4 elements.
 //
 // Layout contract:
 // - src is row-major [batch, ROWS, COLS], dst is row-major [batch, COLS, ROWS]
 // - ROWS % 4 == 0 and COLS % 4 == 0, so both buffers bind as vec4<T>
+// - dispatch: x = ceil(COLS / TILE), y = ceil(ROWS / TILE), z = batch
 //
+// Each (TILE / 4) x TILE workgroup moves one TILE x TILE scalar tile through
+// workgroup memory.
 // Global loads and stores each traverse consecutive vec4 addresses. The extra
 // workgroup-memory column changes the transposed stride and reduces
 // bank-conflict risk on banked implementations.
+{% endif %}{{ env.wgsl.resourceDeclarations }}
 const ROWS: u32 = {{ rowCount }}u;
 const COLS: u32 = {{ cols }}u;
 const ROWS_V: u32 = {{ ((rowCount / 4) | int) }}u;
 const COLS_V: u32 = {{ ((cols / 4) | int) }}u;
+const MAT_V: u32 = {{ (((rowCount * cols * (transposePlaneMiddle | default(1))) / 4) | int) }}u;
+const TILE: u32 = {{ tileSize }}u;
+var<workgroup> tile: array<array<{{ scalar }}, {{ tileSize + 1 }}>, {{ tileSize }}>;
+@compute @workgroup_size({{ (tileSize / 4) | int }}, {{ tileSize }}, 1)
 fn main(
   @builtin(workgroup_id) wid: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>
 ) {
+{% if transposePlaneMiddle is defined %}
+  let source_base = (wid.z / {{ transposePlaneMiddle }}u) * MAT_V + (wid.z % {{ transposePlaneMiddle }}u) * COLS_V;
+  let destination_base = (wid.z / {{ transposePlaneMiddle }}u) * MAT_V + (wid.z % {{ transposePlaneMiddle }}u) * ROWS_V;
+{% else %}  let batch_base = wid.z * MAT_V;
+{% endif %}
   let row_base = wid.y * TILE;
   let col_base = wid.x * TILE;
   let r = row_base + lid.y;
   let cv = (col_base / 4u) + lid.x;
   if (r < ROWS && cv < COLS_V) {
+    let v = src[{% if transposePlaneMiddle is defined %}source_base + r * {{ transposePlaneMiddle }}u * COLS_V{% else %}batch_base + r * COLS_V{% endif %} + cv];
     tile[lid.y][lid.x * 4u + 0u] = v.x;
     tile[lid.y][lid.x * 4u + 1u] = v.y;
     tile[lid.y][lid.x * 4u + 2u] = v.z;
   let out_row = col_base + lid.y;
   let rv = (row_base / 4u) + lid.x;
   if (out_row < COLS && rv < ROWS_V) {
+    dst[{% if transposePlaneMiddle is defined %}destination_base + out_row * {{ transposePlaneMiddle }}u * ROWS_V{% else %}batch_base + out_row * ROWS_V{% endif %} + rv] = {{ vectorScalar }}(
       tile[lid.x * 4u + 0u][lid.y],
       tile[lid.x * 4u + 1u][lid.y],
       tile[lid.x * 4u + 2u][lid.y],

build/webgpu/manifest.json CHANGED Viewed

@@ -5,13 +5,14 @@
   "inputs": { "x": { "onnx": "data", "dtype": "T" } },
   "outputs": { "y": { "onnx": "transposed", "dtype": "T", "rank": "ranks.x" } },
   "attributes": { "perm": {} },
-  "typeConstraints": { "T": ["float32", "float16", "int32", "int16", "uint32", "uint8", "int8", "bool"] },
   "tunables": {
     "WORKGROUP_SIZE": { "default": 256 },
     "INNER_VEC4_ITEMS_PER_THREAD": { "default": 4 },
     "INNER_VEC4_MIN_THREADS": { "default": 65536 }
   },
   "derive": {
     "dtypeSupported": "f16Ok(dtypes.T)",
     "rank2IO": "ranks.x == 2 and ranks.y == 2",
     "rank3IO": "ranks.x == 3 and ranks.y == 3",
@@ -31,15 +32,15 @@
     "innerVec4Rank4Ok": "explicitRank4 and attrs.perm[3] == 3 and dim(shapes.x, 3) % 4 == 0 and numel(shapes.y) > 0 and dtypeSupported",
     "innerVec4Ok": "innerVec4Rank3Ok or innerVec4Rank4Ok",
     "innerVec4Perm": "attrs.perm",
-    "tiled2dShapeOk": "(defaultRank2 or (explicitRank2 and attrs.perm[0] == 1 and attrs.perm[1] == 0)) and numel(shapes.x) > 0 and ceil(dim(shapes.x, 0) / 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceil(dim(shapes.x, 1) / 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dtypeSupported",
     "tiled2dOk": "tiled2dShapeOk and dim(shapes.x, 0) % 4 == 0 and dim(shapes.x, 1) % 4 == 0",
-    "tiled3dShapeOk": "explicitRank3 and attrs.perm[0] == 0 and attrs.perm[1] == 2 and attrs.perm[2] == 1 and numel(shapes.x) > 0 and dim(shapes.x, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceil(dim(shapes.x, 1) / 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceil(dim(shapes.x, 2) / 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dtypeSupported",
     "tiled3dOk": "tiled3dShapeOk and dim(shapes.x, 1) % 4 == 0 and dim(shapes.x, 2) % 4 == 0",
-    "tiled4dShapeOk": "explicitRank4 and attrs.perm[0] == 0 and attrs.perm[1] == 1 and attrs.perm[2] == 3 and attrs.perm[3] == 2 and numel(shapes.x) > 0 and dim(shapes.x, 0) * dim(shapes.x, 1) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceil(dim(shapes.x, 2) / 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceil(dim(shapes.x, 3) / 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dtypeSupported",
     "tiled4dOk": "tiled4dShapeOk and dim(shapes.x, 2) % 4 == 0 and dim(shapes.x, 3) % 4 == 0",
-    "tiled5dNcdhwToNdhwcShapeOk": "explicitRank5 and attrs.perm[0] == 0 and attrs.perm[1] == 2 and attrs.perm[2] == 3 and attrs.perm[3] == 4 and attrs.perm[4] == 1 and numel(shapes.x) > 0 and dim(shapes.x, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceil(dim(shapes.x, 1) / 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceil((dim(shapes.x, 2) * dim(shapes.x, 3) * dim(shapes.x, 4)) / 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dtypeSupported",
     "tiled5dNcdhwToNdhwcOk": "tiled5dNcdhwToNdhwcShapeOk and dim(shapes.x, 1) % 4 == 0 and (dim(shapes.x, 2) * dim(shapes.x, 3) * dim(shapes.x, 4)) % 4 == 0",
-    "tiled5dNdhwcToNcdhwShapeOk": "explicitRank5 and attrs.perm[0] == 0 and attrs.perm[1] == 4 and attrs.perm[2] == 1 and attrs.perm[3] == 2 and attrs.perm[4] == 3 and numel(shapes.x) > 0 and dim(shapes.x, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceil(dim(shapes.x, 1) * dim(shapes.x, 2) * dim(shapes.x, 3) / 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceil(dim(shapes.x, 4) / 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dtypeSupported",
     "tiled5dNdhwcToNcdhwOk": "tiled5dNdhwcToNcdhwShapeOk and dim(shapes.x, 1) * dim(shapes.x, 2) * dim(shapes.x, 3) % 4 == 0 and dim(shapes.x, 4) % 4 == 0",
     "memcpyPerm3": "explicitRank3 and (dim(shapes.x, attrs.perm[0]) == 1 or dim(shapes.x, attrs.perm[1]) == 1 or attrs.perm[0] < attrs.perm[1]) and (dim(shapes.x, attrs.perm[0]) == 1 or dim(shapes.x, attrs.perm[2]) == 1 or attrs.perm[0] < attrs.perm[2]) and (dim(shapes.x, attrs.perm[1]) == 1 or dim(shapes.x, attrs.perm[2]) == 1 or attrs.perm[1] < attrs.perm[2])",
     "memcpyPerm4": "explicitRank4 and (dim(shapes.x, attrs.perm[0]) == 1 or dim(shapes.x, attrs.perm[1]) == 1 or attrs.perm[0] < attrs.perm[1]) and (dim(shapes.x, attrs.perm[0]) == 1 or dim(shapes.x, attrs.perm[2]) == 1 or attrs.perm[0] < attrs.perm[2]) and (dim(shapes.x, attrs.perm[0]) == 1 or dim(shapes.x, attrs.perm[3]) == 1 or attrs.perm[0] < attrs.perm[3]) and (dim(shapes.x, attrs.perm[1]) == 1 or dim(shapes.x, attrs.perm[2]) == 1 or attrs.perm[1] < attrs.perm[2]) and (dim(shapes.x, attrs.perm[1]) == 1 or dim(shapes.x, attrs.perm[3]) == 1 or attrs.perm[1] < attrs.perm[3]) and (dim(shapes.x, attrs.perm[2]) == 1 or dim(shapes.x, attrs.perm[3]) == 1 or attrs.perm[2] < attrs.perm[3])",
@@ -47,64 +48,99 @@
     "memcpyPermOk": "memcpyPerm3 or memcpyPerm4 or memcpyPerm5",
     "identityPermutation": "genericTransposeOk and sameShape(genericPerm, reversedAxes(ranks.x) | reverse)",
     "identityCopyOk": "numel(shapes.y) > 0 and numel(shapes.y) % 4 == 0 and (identityPermutation or memcpyPermOk) and dtypeSupported",
-    "tiled4dNchwToNhwcShapeOk": "explicitRank4 and attrs.perm[0] == 0 and attrs.perm[1] == 2 and attrs.perm[2] == 3 and attrs.perm[3] == 1 and numel(shapes.x) > 0 and dim(shapes.x, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceil(dim(shapes.x, 1) / 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceil((dim(shapes.x, 2) * dim(shapes.x, 3)) / 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dtypeSupported",
     "tiled4dNchwToNhwcOk": "tiled4dNchwToNhwcShapeOk and dim(shapes.x, 1) % 4 == 0 and (dim(shapes.x, 2) * dim(shapes.x, 3)) % 4 == 0",
-    "tiled4dNhwcToNchwShapeOk": "explicitRank4 and attrs.perm[0] == 0 and attrs.perm[1] == 3 and attrs.perm[2] == 1 and attrs.perm[3] == 2 and numel(shapes.x) > 0 and dim(shapes.x, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceil((dim(shapes.x, 1) * dim(shapes.x, 2)) / 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceil(dim(shapes.x, 3) / 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dtypeSupported",
     "tiled4dNhwcToNchwOk": "tiled4dNhwcToNchwShapeOk and (dim(shapes.x, 1) * dim(shapes.x, 2)) % 4 == 0 and dim(shapes.x, 3) % 4 == 0",
     "tiledTransposeOk": "tiled2dOk or tiled3dOk or tiled4dOk or tiled5dNcdhwToNdhwcOk or tiled5dNdhwcToNcdhwOk or tiled4dNchwToNhwcOk or tiled4dNhwcToNchwOk",
     "tiledScalarTransposeOk": "tiled2dShapeOk or tiled3dShapeOk or tiled4dShapeOk or tiled5dNcdhwToNdhwcShapeOk or tiled5dNdhwcToNcdhwShapeOk or tiled4dNchwToNhwcShapeOk or tiled4dNhwcToNchwShapeOk",
     "tiledRows": "dim(shapes.x, 0) if tiled2dShapeOk else (dim(shapes.x, 1) if tiled3dShapeOk else (dim(shapes.x, 2) if tiled4dShapeOk else (dim(shapes.x, 1) if tiled5dNcdhwToNdhwcShapeOk else (dim(shapes.x, 1) * dim(shapes.x, 2) * dim(shapes.x, 3) if tiled5dNdhwcToNcdhwShapeOk else (dim(shapes.x, 1) if tiled4dNchwToNhwcShapeOk else (dim(shapes.x, 1) * dim(shapes.x, 2) if tiled4dNhwcToNchwShapeOk else 1))))))",
     "tiledCols": "dim(shapes.x, 1) if tiled2dShapeOk else (dim(shapes.x, 2) if tiled3dShapeOk else (dim(shapes.x, 3) if tiled4dShapeOk else (dim(shapes.x, 2) * dim(shapes.x, 3) * dim(shapes.x, 4) if tiled5dNcdhwToNdhwcShapeOk else (dim(shapes.x, 4) if tiled5dNdhwcToNcdhwShapeOk else (dim(shapes.x, 2) * dim(shapes.x, 3) if tiled4dNchwToNhwcShapeOk else (dim(shapes.x, 3) if tiled4dNhwcToNchwShapeOk else 1))))))",
-    "tiledBatch": "1 if tiled2dShapeOk else (dim(shapes.x, 0) if tiled3dShapeOk else (dim(shapes.x, 0) * dim(shapes.x, 1) if tiled4dShapeOk else (dim(shapes.x, 0) if tiled5dNcdhwToNdhwcShapeOk or tiled5dNdhwcToNcdhwShapeOk or tiled4dNchwToNhwcShapeOk or tiled4dNhwcToNchwShapeOk else 1)))"
   },
   "bindings": {
-    "src": { "arg": "x", "buffer": "read-only-storage", "elementType": "$vectorScalar" },
-    "dst": { "arg": "y", "buffer": "storage", "elementType": "$vectorScalar" },
-    "x": { "buffer": "read-only-storage", "elementType": "$scalar" },
-    "y": { "buffer": "storage", "elementType": "$scalar" }
   },
   "variants": [
     {
       "id": "tiled_vec4",
       "priority": 22,
-      "when": ["tiledTransposeOk"],
-      "derive": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
-          "name": "Transpose.tiled",
           "shader": "datamove-transpose-2d-tiled.wgsl.jinja",
           "derive": { "rowCount": "tiledRows", "cols": "tiledCols" },
           "bindings": ["src", "dst"],
-          "dispatch": { "x": "ceil(tiledCols / 32)", "y": "ceil(tiledRows / 32)", "z": "tiledBatch" }
         }
       ]
     },
     {
       "id": "tiled_scalar",
       "priority": 19,
-      "when": ["tiledScalarTransposeOk", "dtypeSupported"],
-      "derive": { "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
-          "name": "Transpose.tiledScalar",
           "shader": "datamove-transpose-2d-tiled-scalar.wgsl.jinja",
-          "derive": { "rowCount": "tiledRows", "cols": "tiledCols" },
           "bindings": ["x", "y"],
-          "dispatch": { "x": "ceil(tiledCols / 32)", "y": "ceil(tiledRows / 32)", "z": "tiledBatch" }
         }
       ]
     },
     {
       "id": "identity_vec4",
       "priority": 21,
-      "when": ["identityCopyOk"],
-      "derive": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
-          "name": "Transpose.identityVec4",
           "shader": "datamove-flat-copy.wgsl.jinja",
           "derive": { "count": "numel(shapes.y) / 4" },
           "bindings": ["src", "dst"],
@@ -119,20 +155,13 @@
     {
       "id": "inner_vec4",
       "priority": 20,
-      "when": ["innerVec4Ok"],
-      "derive": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
-          "name": "Transpose.innerVec4",
           "shader": "datamove-transpose-vec4.wgsl.jinja",
-          "derive": {
-            "inputShape": "shapes.x",
-            "outputShape": "shapes.y",
-            "rank": "ranks.x",
-            "perm": "innerVec4Perm",
-            "count": "numel(shapes.y) / 4"
-          },
           "bindings": ["src", "dst"],
           "dispatch": {
             "x": "min(ceilDiv((max(min(numel(shapes.y) / 4, tunables.INNER_VEC4_MIN_THREADS), ceilDiv(numel(shapes.y) / 4, tunables.INNER_VEC4_ITEMS_PER_THREAD))), (tunables.WORKGROUP_SIZE)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
@@ -145,19 +174,12 @@
     {
       "id": "generic",
       "when": ["genericTransposeOk", "dtypeSupported"],
-      "derive": { "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "Transpose",
           "shader": "transpose.wgsl.jinja",
-          "derive": {
-            "inputShape": "shapes.x",
-            "outputShape": "shapes.y",
-            "rank": "ranks.x",
-            "perm": "genericPerm",
-            "itemsPerInvocation": 4
-          },
           "bindings": [
             "x",
             "y",

   "inputs": { "x": { "onnx": "data", "dtype": "T" } },
   "outputs": { "y": { "onnx": "transposed", "dtype": "T", "rank": "ranks.x" } },
   "attributes": { "perm": {} },
+  "typeConstraints": { "T": ["float32", "float16", "int32", "int16", "uint32", "uint8", "int8", "bool", "int64"] },
   "tunables": {
     "WORKGROUP_SIZE": { "default": 256 },
     "INNER_VEC4_ITEMS_PER_THREAD": { "default": 4 },
     "INNER_VEC4_MIN_THREADS": { "default": 65536 }
   },
   "derive": {
+    "tileSize": 32,
     "dtypeSupported": "f16Ok(dtypes.T)",
     "rank2IO": "ranks.x == 2 and ranks.y == 2",
     "rank3IO": "ranks.x == 3 and ranks.y == 3",
     "innerVec4Rank4Ok": "explicitRank4 and attrs.perm[3] == 3 and dim(shapes.x, 3) % 4 == 0 and numel(shapes.y) > 0 and dtypeSupported",
     "innerVec4Ok": "innerVec4Rank3Ok or innerVec4Rank4Ok",
     "innerVec4Perm": "attrs.perm",
+    "tiled2dShapeOk": "(defaultRank2 or (explicitRank2 and attrs.perm[0] == 1 and attrs.perm[1] == 0)) and numel(shapes.x) > 0 and ceil(dim(shapes.x, 0) / tileSize) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceil(dim(shapes.x, 1) / tileSize) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dtypeSupported",
     "tiled2dOk": "tiled2dShapeOk and dim(shapes.x, 0) % 4 == 0 and dim(shapes.x, 1) % 4 == 0",
+    "tiled3dShapeOk": "explicitRank3 and attrs.perm[0] == 0 and attrs.perm[1] == 2 and attrs.perm[2] == 1 and numel(shapes.x) > 0 and dim(shapes.x, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceil(dim(shapes.x, 1) / tileSize) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceil(dim(shapes.x, 2) / tileSize) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dtypeSupported",
     "tiled3dOk": "tiled3dShapeOk and dim(shapes.x, 1) % 4 == 0 and dim(shapes.x, 2) % 4 == 0",
+    "tiled4dShapeOk": "explicitRank4 and attrs.perm[0] == 0 and attrs.perm[1] == 1 and attrs.perm[2] == 3 and attrs.perm[3] == 2 and numel(shapes.x) > 0 and dim(shapes.x, 0) * dim(shapes.x, 1) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceil(dim(shapes.x, 2) / tileSize) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceil(dim(shapes.x, 3) / tileSize) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dtypeSupported",
     "tiled4dOk": "tiled4dShapeOk and dim(shapes.x, 2) % 4 == 0 and dim(shapes.x, 3) % 4 == 0",
+    "tiled5dNcdhwToNdhwcShapeOk": "explicitRank5 and attrs.perm[0] == 0 and attrs.perm[1] == 2 and attrs.perm[2] == 3 and attrs.perm[3] == 4 and attrs.perm[4] == 1 and numel(shapes.x) > 0 and dim(shapes.x, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceil(dim(shapes.x, 1) / tileSize) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceil((dim(shapes.x, 2) * dim(shapes.x, 3) * dim(shapes.x, 4)) / tileSize) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dtypeSupported",
     "tiled5dNcdhwToNdhwcOk": "tiled5dNcdhwToNdhwcShapeOk and dim(shapes.x, 1) % 4 == 0 and (dim(shapes.x, 2) * dim(shapes.x, 3) * dim(shapes.x, 4)) % 4 == 0",
+    "tiled5dNdhwcToNcdhwShapeOk": "explicitRank5 and attrs.perm[0] == 0 and attrs.perm[1] == 4 and attrs.perm[2] == 1 and attrs.perm[3] == 2 and attrs.perm[4] == 3 and numel(shapes.x) > 0 and dim(shapes.x, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceil(dim(shapes.x, 1) * dim(shapes.x, 2) * dim(shapes.x, 3) / tileSize) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceil(dim(shapes.x, 4) / tileSize) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dtypeSupported",
     "tiled5dNdhwcToNcdhwOk": "tiled5dNdhwcToNcdhwShapeOk and dim(shapes.x, 1) * dim(shapes.x, 2) * dim(shapes.x, 3) % 4 == 0 and dim(shapes.x, 4) % 4 == 0",
     "memcpyPerm3": "explicitRank3 and (dim(shapes.x, attrs.perm[0]) == 1 or dim(shapes.x, attrs.perm[1]) == 1 or attrs.perm[0] < attrs.perm[1]) and (dim(shapes.x, attrs.perm[0]) == 1 or dim(shapes.x, attrs.perm[2]) == 1 or attrs.perm[0] < attrs.perm[2]) and (dim(shapes.x, attrs.perm[1]) == 1 or dim(shapes.x, attrs.perm[2]) == 1 or attrs.perm[1] < attrs.perm[2])",
     "memcpyPerm4": "explicitRank4 and (dim(shapes.x, attrs.perm[0]) == 1 or dim(shapes.x, attrs.perm[1]) == 1 or attrs.perm[0] < attrs.perm[1]) and (dim(shapes.x, attrs.perm[0]) == 1 or dim(shapes.x, attrs.perm[2]) == 1 or attrs.perm[0] < attrs.perm[2]) and (dim(shapes.x, attrs.perm[0]) == 1 or dim(shapes.x, attrs.perm[3]) == 1 or attrs.perm[0] < attrs.perm[3]) and (dim(shapes.x, attrs.perm[1]) == 1 or dim(shapes.x, attrs.perm[2]) == 1 or attrs.perm[1] < attrs.perm[2]) and (dim(shapes.x, attrs.perm[1]) == 1 or dim(shapes.x, attrs.perm[3]) == 1 or attrs.perm[1] < attrs.perm[3]) and (dim(shapes.x, attrs.perm[2]) == 1 or dim(shapes.x, attrs.perm[3]) == 1 or attrs.perm[2] < attrs.perm[3])",
     "memcpyPermOk": "memcpyPerm3 or memcpyPerm4 or memcpyPerm5",
     "identityPermutation": "genericTransposeOk and sameShape(genericPerm, reversedAxes(ranks.x) | reverse)",
     "identityCopyOk": "numel(shapes.y) > 0 and numel(shapes.y) % 4 == 0 and (identityPermutation or memcpyPermOk) and dtypeSupported",
+    "tiled4dNchwToNhwcShapeOk": "explicitRank4 and attrs.perm[0] == 0 and attrs.perm[1] == 2 and attrs.perm[2] == 3 and attrs.perm[3] == 1 and numel(shapes.x) > 0 and dim(shapes.x, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceil(dim(shapes.x, 1) / tileSize) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceil((dim(shapes.x, 2) * dim(shapes.x, 3)) / tileSize) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dtypeSupported",
     "tiled4dNchwToNhwcOk": "tiled4dNchwToNhwcShapeOk and dim(shapes.x, 1) % 4 == 0 and (dim(shapes.x, 2) * dim(shapes.x, 3)) % 4 == 0",
+    "tiled4dNhwcToNchwShapeOk": "explicitRank4 and attrs.perm[0] == 0 and attrs.perm[1] == 3 and attrs.perm[2] == 1 and attrs.perm[3] == 2 and numel(shapes.x) > 0 and dim(shapes.x, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceil((dim(shapes.x, 1) * dim(shapes.x, 2)) / tileSize) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceil(dim(shapes.x, 3) / tileSize) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dtypeSupported",
     "tiled4dNhwcToNchwOk": "tiled4dNhwcToNchwShapeOk and (dim(shapes.x, 1) * dim(shapes.x, 2)) % 4 == 0 and dim(shapes.x, 3) % 4 == 0",
     "tiledTransposeOk": "tiled2dOk or tiled3dOk or tiled4dOk or tiled5dNcdhwToNdhwcOk or tiled5dNdhwcToNcdhwOk or tiled4dNchwToNhwcOk or tiled4dNhwcToNchwOk",
     "tiledScalarTransposeOk": "tiled2dShapeOk or tiled3dShapeOk or tiled4dShapeOk or tiled5dNcdhwToNdhwcShapeOk or tiled5dNdhwcToNcdhwShapeOk or tiled4dNchwToNhwcShapeOk or tiled4dNhwcToNchwShapeOk",
     "tiledRows": "dim(shapes.x, 0) if tiled2dShapeOk else (dim(shapes.x, 1) if tiled3dShapeOk else (dim(shapes.x, 2) if tiled4dShapeOk else (dim(shapes.x, 1) if tiled5dNcdhwToNdhwcShapeOk else (dim(shapes.x, 1) * dim(shapes.x, 2) * dim(shapes.x, 3) if tiled5dNdhwcToNcdhwShapeOk else (dim(shapes.x, 1) if tiled4dNchwToNhwcShapeOk else (dim(shapes.x, 1) * dim(shapes.x, 2) if tiled4dNhwcToNchwShapeOk else 1))))))",
     "tiledCols": "dim(shapes.x, 1) if tiled2dShapeOk else (dim(shapes.x, 2) if tiled3dShapeOk else (dim(shapes.x, 3) if tiled4dShapeOk else (dim(shapes.x, 2) * dim(shapes.x, 3) * dim(shapes.x, 4) if tiled5dNcdhwToNdhwcShapeOk else (dim(shapes.x, 4) if tiled5dNdhwcToNcdhwShapeOk else (dim(shapes.x, 2) * dim(shapes.x, 3) if tiled4dNchwToNhwcShapeOk else (dim(shapes.x, 3) if tiled4dNhwcToNchwShapeOk else 1))))))",
+    "tiledBatch": "1 if tiled2dShapeOk else (dim(shapes.x, 0) if tiled3dShapeOk else (dim(shapes.x, 0) * dim(shapes.x, 1) if tiled4dShapeOk else (dim(shapes.x, 0) if tiled5dNcdhwToNdhwcShapeOk or tiled5dNdhwcToNcdhwShapeOk or tiled4dNchwToNhwcShapeOk or tiled4dNhwcToNchwShapeOk else 1)))",
+    "stridedPlaneRank3": "rank3IO and genericTransposeOk and genericPerm[0] == 2 and genericPerm[1] == 1 and genericPerm[2] == 0",
+    "stridedPlaneRank4": "rank4IO and genericTransposeOk and genericPerm[0] == 0 and genericPerm[1] == 3 and genericPerm[2] == 2 and genericPerm[3] == 1",
+    "stridedPlaneRows": "dim(shapes.x, 0) if stridedPlaneRank3 else dim(shapes.x, 1) if stridedPlaneRank4 else 1",
+    "stridedPlaneCols": "dim(shapes.x, 2) if stridedPlaneRank3 else dim(shapes.x, 3) if stridedPlaneRank4 else 1",
+    "stridedPlaneMiddle": "dim(shapes.x, 1) if stridedPlaneRank3 else dim(shapes.x, 2) if stridedPlaneRank4 else 1",
+    "stridedPlaneBatch": "1 if stridedPlaneRank3 else dim(shapes.x, 0) if stridedPlaneRank4 else 1",
+    "stridedPlaneGroups": "stridedPlaneBatch * stridedPlaneMiddle",
+    "stridedPlaneOk": "(stridedPlaneRank3 or stridedPlaneRank4) and numel(shapes.x) > 0 and dtypeSupported and ceilDiv(stridedPlaneRows, tileSize) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(stridedPlaneCols, tileSize) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and stridedPlaneGroups <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and device.limits.maxComputeInvocationsPerWorkgroup >= tileSize * tileSize / 4 and device.limits.maxComputeWorkgroupStorageSize >= tileSize * (tileSize + 1) * (2 if dtypes.T == \"f16\" else 4)",
+    "stridedPlaneSparseTile": "ceilDiv(stridedPlaneRows, tileSize) * ceilDiv(stridedPlaneCols, tileSize) * tileSize * tileSize > 2 * stridedPlaneRows * stridedPlaneCols",
+    "scalar": "dtypes.T",
+    "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
+    "inputShape": "shapes.x",
+    "outputShape": "shapes.y",
+    "rank": "ranks.x"
   },
   "bindings": {
+    "src": { "arg": "x", "elementType": "$vectorScalar" },
+    "dst": { "arg": "y", "elementType": "$vectorScalar" },
+    "x": { "elementType": "$scalar" },
+    "y": { "elementType": "$scalar" }
   },
   "variants": [
+    {
+      "id": "strided_planes_vec4",
+      "priority": 22,
+      "when": ["dtypes.T != \"vec2<u32>\"", "stridedPlaneOk", "stridedPlaneRows % 4 == 0 and stridedPlaneCols % 4 == 0", "device.limits.maxComputeWorkgroupSizeX >= tileSize / 4", "device.limits.maxComputeWorkgroupSizeY >= tileSize"],
+      "passes": [
+        {
+          "id": "main",
+          "name": "Transpose.Tiled",
+          "shader": "datamove-transpose-2d-tiled.wgsl.jinja",
+          "derive": {
+            "rowCount": "stridedPlaneRows",
+            "cols": "stridedPlaneCols",
+            "transposePlaneMiddle": "stridedPlaneMiddle"
+          },
+          "bindings": ["src", "dst"],
+          "dispatch": {
+            "x": "ceilDiv(stridedPlaneCols, tileSize)",
+            "y": "ceilDiv(stridedPlaneRows, tileSize)",
+            "z": "stridedPlaneGroups"
+          }
+        }
+      ],
+      "demoteWhen": ["stridedPlaneSparseTile"]
+    },
     {
       "id": "tiled_vec4",
       "priority": 22,
+      "when": ["dtypes.T != \"vec2<u32>\"", "tiledTransposeOk"],
       "passes": [
         {
           "id": "main",
+          "name": "Transpose.Tiled",
           "shader": "datamove-transpose-2d-tiled.wgsl.jinja",
           "derive": { "rowCount": "tiledRows", "cols": "tiledCols" },
           "bindings": ["src", "dst"],
+          "dispatch": { "x": "ceil(tiledCols / tileSize)", "y": "ceil(tiledRows / tileSize)", "z": "tiledBatch" }
         }
       ]
     },
     {
       "id": "tiled_scalar",
       "priority": 19,
+      "when": ["dtypes.T != \"vec2<u32>\"", "tiledScalarTransposeOk", "dtypeSupported"],
       "passes": [
         {
           "id": "main",
+          "name": "Transpose.TiledScalar",
           "shader": "datamove-transpose-2d-tiled-scalar.wgsl.jinja",
+          "derive": { "rowCount": "tiledRows", "cols": "tiledCols", "tileRowStep": 8 },
           "bindings": ["x", "y"],
+          "dispatch": { "x": "ceil(tiledCols / tileSize)", "y": "ceil(tiledRows / tileSize)", "z": "tiledBatch" }
         }
       ]
     },
     {
       "id": "identity_vec4",
       "priority": 21,
+      "when": ["dtypes.T != \"vec2<u32>\"", "identityCopyOk"],
       "passes": [
         {
           "id": "main",
+          "name": "Transpose.IdentityVec4",
           "shader": "datamove-flat-copy.wgsl.jinja",
           "derive": { "count": "numel(shapes.y) / 4" },
           "bindings": ["src", "dst"],
     {
       "id": "inner_vec4",
       "priority": 20,
+      "when": ["dtypes.T != \"vec2<u32>\"", "innerVec4Ok"],
       "passes": [
         {
           "id": "main",
+          "name": "Transpose.InnerVec4",
           "shader": "datamove-transpose-vec4.wgsl.jinja",
+          "derive": { "perm": "innerVec4Perm", "count": "numel(shapes.y) / 4" },
           "bindings": ["src", "dst"],
           "dispatch": {
             "x": "min(ceilDiv((max(min(numel(shapes.y) / 4, tunables.INNER_VEC4_MIN_THREADS), ceilDiv(numel(shapes.y) / 4, tunables.INNER_VEC4_ITEMS_PER_THREAD))), (tunables.WORKGROUP_SIZE)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
     {
       "id": "generic",
       "when": ["genericTransposeOk", "dtypeSupported"],
       "passes": [
         {
           "id": "main",
           "name": "Transpose",
           "shader": "transpose.wgsl.jinja",
+          "derive": { "perm": "genericPerm", "itemsPerInvocation": 4 },
           "bindings": [
             "x",
             "y",

build/webgpu/metadata.json CHANGED Viewed

@@ -1,26 +1,27 @@
 {
   "name": "ai.onnx.Transpose",
-  "id": "_ai_onnx_transpose_webgpu_2a9c5f1",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "ClUFa7Y4uL69eUy46pTK4YciV/+jJxlDxSJFK+/JU6A=",
       "datamove-flat-copy.wgsl.jinja": "g9d62mer2bmHfbScqX5CIT5Zq/Pdinac0W2dH0CD+5s=",
-      "datamove-transpose-2d-tiled-scalar.wgsl.jinja": "ycF9afS0CqHrLIfFi6hAqp+7GIcnHHAyN74hSF+/FrY=",
-      "datamove-transpose-2d-tiled.wgsl.jinja": "GNVRxhagDFVQ0DGjsUGD8BVzktDHsELwe0qw4xRHD60=",
       "datamove-transpose-vec4.wgsl.jinja": "NARpwfOEcl8GRu17lnq62KXK/mgh61Jj3YHfi8Zofik=",
-      "manifest.json": "F4FH/KuY12tkUnIKlE3wBmydDgzgzAbfPva12DYYGBY=",
-      "test.json": "rQerBft2NMuE0I+S7GErvy8rPRoUQxOkI3NvNH5dcKc=",
-      "transpose.wgsl.jinja": "a9TDelnkVrpe1XBA8rTAurSGrTmIQ6BO16ZHRJWiQJw="
     }
   },
-  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
   "webgpu": {
-    "manifestSpec": "2.0",
     "variants": {
       "tiled_vec4": ["datamove-transpose-2d-tiled.wgsl.jinja"],
       "tiled_scalar": ["datamove-transpose-2d-tiled-scalar.wgsl.jinja"],
       "identity_vec4": ["datamove-flat-copy.wgsl.jinja"],

 {
   "name": "ai.onnx.Transpose",
+  "id": "_ai_onnx_transpose_webgpu_4ac33c9",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "Z+G8ub4uWe/9LpOzIYhV6tWU29vMRvN3Ajhc3H+po0Q=",
       "datamove-flat-copy.wgsl.jinja": "g9d62mer2bmHfbScqX5CIT5Zq/Pdinac0W2dH0CD+5s=",
+      "datamove-transpose-2d-tiled-scalar.wgsl.jinja": "JhTsa8OYjMAUAKUDAkFxnxTQyr/f8FpW783nwX6PWZk=",
+      "datamove-transpose-2d-tiled.wgsl.jinja": "T7WQJsuQ85hTO/F+QQ8GOXsvBP80iJmD0s5cRgPHG4U=",
       "datamove-transpose-vec4.wgsl.jinja": "NARpwfOEcl8GRu17lnq62KXK/mgh61Jj3YHfi8Zofik=",
+      "manifest.json": "alujpyF6YtHpSBUjwlWNjNi2z8k0Et7TUtb+yKu+Wm8=",
+      "test.json": "llE4r29U/p+9XIkG4zoy/1KeOXaQ2dk32PpU8GNcrS4=",
+      "transpose.wgsl.jinja": "l11cGbVOuaD1SwWqAhmKpkVs/rWCOm/YO6zFao2YZ4E="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "variants": {
+      "strided_planes_vec4": ["datamove-transpose-2d-tiled.wgsl.jinja"],
       "tiled_vec4": ["datamove-transpose-2d-tiled.wgsl.jinja"],
       "tiled_scalar": ["datamove-transpose-2d-tiled-scalar.wgsl.jinja"],
       "identity_vec4": ["datamove-flat-copy.wgsl.jinja"],

build/webgpu/test.json CHANGED Viewed

@@ -2,7 +2,12 @@
   "fixtureArrays": {
     "onnx_backend_transpose_input_x": [0.54881352186203, 0.7151893377304077, 0.6027633547782898, 0.5448831915855408, 0.42365479469299316, 0.6458941102027893, 0.4375872015953064, 0.891772985458374, 0.9636627435684204, 0.3834415078163147, 0.7917250394821167, 0.5288949012756348, 0.5680445432662964, 0.9255966544151306, 0.07103605568408966, 0.08712930232286453, 0.020218396559357643, 0.832619845867157, 0.7781567573547363, 0.8700121641159058, 0.978618323802948, 0.7991585731506348, 0.4614793658256531, 0.7805292010307312],
     "ort_transpose021_f32_input_x": [1, 2, 3, 4, 5, 6, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3],
-    "ort_nhwc_to_nchw_uint32_input_x": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
   },
   "cases": [
     {
@@ -906,7 +911,7 @@
     {
       "name": "tiled_2d_vec4_naninf_signedzero_payload_preserved",
       "provenance": {
-        "notes": "Aligned [8,8] perm[1,0] f32 selects tiled_vec4 (both dims %4==0, %32 partial). Verifies NaN / +-Inf / signed-zero survive the vec4->scalar->workgroup->vec4 staging bit-exact. Expected values are the [8,8] input transposed (column-major read)."
       },
       "attrs": { "perm": [1, 0] },
       "inputs": {
@@ -935,7 +940,7 @@
     {
       "name": "scalar_rank4_default_reverse_naninf_preserved",
       "provenance": {
-        "notes": "No perm => default full reversal [3,2,1,0] on the generic scalar path. Values placed so index [i,j,k,l] holds 8*i+4*j+2*k+l; output[l,k,j,i] = input[i,j,k,l]. NaN/Inf/-0 payloads at select positions must survive the scalar input_offset gather."
       },
       "inputs": {
         "x": {
@@ -963,7 +968,7 @@
     {
       "name": "rank5_ncdhw_to_ndhwc_tiled_vec4",
       "provenance": {
-        "notes": "Compact route lock for the flattened batched-2D NCDHW-to-NDHWC tile. C=4 and D*H*W=32 exercise one complete spatial tile boundary while preserving exact copy semantics."
       },
       "attrs": { "perm": [0, 2, 3, 4, 1] },
       "inputs": {
@@ -978,7 +983,7 @@
     {
       "name": "rank5_ndhwc_to_ncdhw_tiled_candidate",
       "provenance": {
-        "notes": "Inverse rank-5 layout lock: NDHWC [N,D,H,W,C] becomes NCDHW via perm [0,4,1,2,3]. C=4 and D*H*W=32 match one flattened 2D transpose boundary with bit-exact copy semantics."
       },
       "attrs": { "perm": [0, 4, 1, 2, 3] },
       "inputs": {
@@ -1123,6 +1128,453 @@
         }
       },
       "outputs": { "y": { "dtype": "float32", "shape": [4, 2, 2, 4, 2, 2, 3, 2], "tolerance": 0 } }
     }
   ]
 }

   "fixtureArrays": {
     "onnx_backend_transpose_input_x": [0.54881352186203, 0.7151893377304077, 0.6027633547782898, 0.5448831915855408, 0.42365479469299316, 0.6458941102027893, 0.4375872015953064, 0.891772985458374, 0.9636627435684204, 0.3834415078163147, 0.7917250394821167, 0.5288949012756348, 0.5680445432662964, 0.9255966544151306, 0.07103605568408966, 0.08712930232286453, 0.020218396559357643, 0.832619845867157, 0.7781567573547363, 0.8700121641159058, 0.978618323802948, 0.7991585731506348, 0.4614793658256531, 0.7805292010307312],
     "ort_transpose021_f32_input_x": [1, 2, 3, 4, 5, 6, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3],
+    "ort_nhwc_to_nchw_uint32_input_x": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
+    "strided_f32": [0, 1e-40, -1e-40, 0.125, -3.5, 100000000000000000000, -100000000000000000000, 17.25, -31.5, 0.5, 7.75, -0.75, 2],
+    "strided_f16": [0, 5.960464477539063e-8, -5.960464477539063e-8, 0.125, -3.5, 65504, -65504, 17.25, -31.5, 0.5, 7.75, -0.75, 2],
+    "strided_signed": [-63, -62, -61, -60, -59, -58, -57, -56, -55, -54, -53, -52, -51, -50, -49, -48, -47, -46, -45, -44, -43, -42, -41, -40, -39, -38, -37, -36, -35, -34, -33, -32, -31, -30, -29, -28, -27, -26, -25, -24, -23, -22, -21, -20, -19, -18, -17, -16, -15, -14, -13, -12, -11, -10, -9, -8, -7, -6, -5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63],
+    "strided_unsigned": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126],
+    "strided_bool": [0, 1, 0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 1, 0, 1, 1, 0, 1, 0, 0, 1, 0, 1, 0, 1, 1, 1, 0, 1, 0, 1, 0, 1, 0, 0, 0, 1, 0, 1, 0, 1, 1, 0, 1, 0, 0, 1, 0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 1, 0, 0, 1, 1, 0, 1, 0, 1, 0, 1, 1, 1, 0, 1, 0, 1, 0, 0, 1, 0, 1, 1, 0, 1, 0, 0, 1, 0, 1, 0, 1, 1, 1, 0, 1, 0, 1, 0, 1, 1, 0, 0, 1, 0, 1, 0, 0, 1, 0]
   },
   "cases": [
     {
     {
       "name": "tiled_2d_vec4_naninf_signedzero_payload_preserved",
       "provenance": {
+        "notes": "An aligned [8,8] transpose with perm=[1,0] must carry NaN, both infinities, and zeros exactly to their transposed positions."
       },
       "attrs": { "perm": [1, 0] },
       "inputs": {
     {
       "name": "scalar_rank4_default_reverse_naninf_preserved",
       "provenance": {
+        "notes": "With perm omitted, the default full reversal [3,2,1,0] maps output[l,k,j,i] to input[i,j,k,l]. NaN and both infinities must land exactly at their reversed positions."
       },
       "inputs": {
         "x": {
     {
       "name": "rank5_ncdhw_to_ndhwc_tiled_vec4",
       "provenance": {
+        "notes": "NCDHW-to-NDHWC transpose with four channels and 32 spatial positions checks exact batched copy semantics."
       },
       "attrs": { "perm": [0, 2, 3, 4, 1] },
       "inputs": {
     {
       "name": "rank5_ndhwc_to_ncdhw_tiled_candidate",
       "provenance": {
+        "notes": "A rank-5 NDHWC-to-NCDHW transpose (perm [0,4,1,2,3]) moves the size-4 channel axis to position 1, with the D, H, and W dimensions (2x4x4=32 elements) kept contiguous; distinct linspace values at every position make a wrong index mapping observable."
       },
       "attrs": { "perm": [0, 4, 1, 2, 3] },
       "inputs": {
         }
       },
       "outputs": { "y": { "dtype": "float32", "shape": [4, 2, 2, 4, 2, 2, 3, 2], "tolerance": 0 } }
+    },
+    {
+      "name": "strided_planes_ragged_batched_float32",
+      "inputs": {
+        "x": {
+          "shape": [2, 47, 3, 63],
+          "dtype": "float32",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_f32" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 63, 3, 47], "tolerance": 0, "relTolerance": 0 } },
+      "attrs": { "perm": [0, 3, 2, 1] }
+    },
+    {
+      "name": "strided_planes_ragged_batched_float16",
+      "inputs": {
+        "x": {
+          "shape": [2, 47, 3, 63],
+          "dtype": "float16",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_f16" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [2, 63, 3, 47], "tolerance": 0, "relTolerance": 0 } },
+      "attrs": { "perm": [0, 3, 2, 1] }
+    },
+    {
+      "name": "strided_planes_ragged_batched_int32",
+      "inputs": {
+        "x": {
+          "shape": [2, 47, 3, 63],
+          "dtype": "int32",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_signed" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "int32", "shape": [2, 63, 3, 47], "tolerance": 0, "relTolerance": 0 } },
+      "attrs": { "perm": [0, 3, 2, 1] }
+    },
+    {
+      "name": "strided_planes_ragged_batched_int16",
+      "inputs": {
+        "x": {
+          "shape": [2, 47, 3, 63],
+          "dtype": "int16",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_signed" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "int16", "shape": [2, 63, 3, 47], "tolerance": 0, "relTolerance": 0 } },
+      "attrs": { "perm": [0, 3, 2, 1] }
+    },
+    {
+      "name": "strided_planes_ragged_batched_uint32",
+      "inputs": {
+        "x": {
+          "shape": [2, 47, 3, 63],
+          "dtype": "uint32",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_unsigned" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [2, 63, 3, 47], "tolerance": 0, "relTolerance": 0 } },
+      "attrs": { "perm": [0, 3, 2, 1] }
+    },
+    {
+      "name": "strided_planes_ragged_batched_uint8",
+      "inputs": {
+        "x": {
+          "shape": [2, 47, 3, 63],
+          "dtype": "uint8",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_unsigned" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint8", "shape": [2, 63, 3, 47], "tolerance": 0, "relTolerance": 0 } },
+      "attrs": { "perm": [0, 3, 2, 1] }
+    },
+    {
+      "name": "strided_planes_ragged_batched_int8",
+      "inputs": {
+        "x": {
+          "shape": [2, 47, 3, 63],
+          "dtype": "int8",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_signed" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "int8", "shape": [2, 63, 3, 47], "tolerance": 0, "relTolerance": 0 } },
+      "attrs": { "perm": [0, 3, 2, 1] }
+    },
+    {
+      "name": "strided_planes_ragged_batched_bool",
+      "inputs": {
+        "x": {
+          "shape": [2, 47, 3, 63],
+          "dtype": "bool",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_bool" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "bool", "shape": [2, 63, 3, 47], "tolerance": 0, "relTolerance": 0 } },
+      "attrs": { "perm": [0, 3, 2, 1] }
+    },
+    {
+      "name": "strided_planes_rank3_default_aligned_float32",
+      "inputs": {
+        "x": {
+          "shape": [16, 3, 32],
+          "dtype": "float32",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_f32" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [32, 3, 16], "tolerance": 0, "relTolerance": 0 } }
+    },
+    {
+      "name": "strided_planes_half_full_tile_float32",
+      "inputs": {
+        "x": {
+          "shape": [2, 16, 3, 32],
+          "dtype": "float32",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_f32" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 32, 3, 16], "tolerance": 0, "relTolerance": 0 } },
+      "attrs": { "perm": [0, 3, 2, 1] }
+    },
+    {
+      "name": "strided_planes_padding_demoted_boundary_float32",
+      "inputs": {
+        "x": {
+          "shape": [2, 15, 3, 32],
+          "dtype": "float32",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_f32" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 32, 3, 15], "tolerance": 0, "relTolerance": 0 } },
+      "attrs": { "perm": [0, 3, 2, 1] }
+    },
+    {
+      "name": "strided_planes_skinny_demoted_float32",
+      "inputs": {
+        "x": {
+          "shape": [2, 8, 5, 64],
+          "dtype": "float32",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_f32" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 64, 5, 8], "tolerance": 0, "relTolerance": 0 } },
+      "attrs": { "perm": [0, 3, 2, 1] }
+    },
+    {
+      "name": "strided_planes_partial_scalar_demoted_float32",
+      "inputs": {
+        "x": {
+          "shape": [2, 17, 3, 32],
+          "dtype": "float32",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_f32" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 32, 3, 17], "tolerance": 0, "relTolerance": 0 } },
+      "attrs": { "perm": [0, 3, 2, 1] }
+    },
+    {
+      "name": "strided_planes_rank3_explicit_ragged_float32",
+      "inputs": {
+        "x": {
+          "shape": [47, 3, 63],
+          "dtype": "float32",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_f32" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [63, 3, 47], "tolerance": 0, "relTolerance": 0 } },
+      "attrs": { "perm": [2, 1, 0] }
+    },
+    {
+      "name": "strided_planes_aligned_nan_float32",
+      "inputs": {
+        "x": {
+          "shape": [2, 16, 3, 32],
+          "dtype": "float32",
+          "data": {
+            "kind": "cycle",
+            "values": { "$ref": "#/fixtureArrays/strided_f32" },
+            "nanStart": 3.0,
+            "nanCount": 2.0
+          }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [2, 32, 3, 16], "tolerance": 0, "relTolerance": 0, "allowNaN": true }
+      },
+      "attrs": { "perm": [0, 3, 2, 1] }
+    },
+    {
+      "name": "strided_planes_ragged_nan_float32",
+      "inputs": {
+        "x": {
+          "shape": [2, 47, 3, 63],
+          "dtype": "float32",
+          "data": {
+            "kind": "cycle",
+            "values": { "$ref": "#/fixtureArrays/strided_f32" },
+            "nanStart": 3.0,
+            "nanCount": 2.0
+          }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [2, 63, 3, 47], "tolerance": 0, "relTolerance": 0, "allowNaN": true }
+      },
+      "attrs": { "perm": [0, 3, 2, 1] }
+    },
+    {
+      "name": "strided_planes_rank3_default_aligned_float16",
+      "inputs": {
+        "x": {
+          "shape": [16, 3, 32],
+          "dtype": "float16",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_f16" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [32, 3, 16], "tolerance": 0, "relTolerance": 0 } }
+    },
+    {
+      "name": "strided_planes_half_full_tile_float16",
+      "inputs": {
+        "x": {
+          "shape": [2, 16, 3, 32],
+          "dtype": "float16",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_f16" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [2, 32, 3, 16], "tolerance": 0, "relTolerance": 0 } },
+      "attrs": { "perm": [0, 3, 2, 1] }
+    },
+    {
+      "name": "strided_planes_padding_demoted_boundary_float16",
+      "inputs": {
+        "x": {
+          "shape": [2, 15, 3, 32],
+          "dtype": "float16",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_f16" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [2, 32, 3, 15], "tolerance": 0, "relTolerance": 0 } },
+      "attrs": { "perm": [0, 3, 2, 1] }
+    },
+    {
+      "name": "strided_planes_skinny_demoted_float16",
+      "inputs": {
+        "x": {
+          "shape": [2, 8, 5, 64],
+          "dtype": "float16",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_f16" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [2, 64, 5, 8], "tolerance": 0, "relTolerance": 0 } },
+      "attrs": { "perm": [0, 3, 2, 1] }
+    },
+    {
+      "name": "strided_planes_partial_scalar_demoted_float16",
+      "inputs": {
+        "x": {
+          "shape": [2, 17, 3, 32],
+          "dtype": "float16",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_f16" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [2, 32, 3, 17], "tolerance": 0, "relTolerance": 0 } },
+      "attrs": { "perm": [0, 3, 2, 1] }
+    },
+    {
+      "name": "strided_planes_rank3_explicit_ragged_float16",
+      "inputs": {
+        "x": {
+          "shape": [47, 3, 63],
+          "dtype": "float16",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_f16" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [63, 3, 47], "tolerance": 0, "relTolerance": 0 } },
+      "attrs": { "perm": [2, 1, 0] }
+    },
+    {
+      "name": "strided_planes_aligned_nan_float16",
+      "inputs": {
+        "x": {
+          "shape": [2, 16, 3, 32],
+          "dtype": "float16",
+          "data": {
+            "kind": "cycle",
+            "values": { "$ref": "#/fixtureArrays/strided_f16" },
+            "nanStart": 3.0,
+            "nanCount": 2.0
+          }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float16", "shape": [2, 32, 3, 16], "tolerance": 0, "relTolerance": 0, "allowNaN": true }
+      },
+      "attrs": { "perm": [0, 3, 2, 1] }
+    },
+    {
+      "name": "strided_planes_ragged_nan_float16",
+      "inputs": {
+        "x": {
+          "shape": [2, 47, 3, 63],
+          "dtype": "float16",
+          "data": {
+            "kind": "cycle",
+            "values": { "$ref": "#/fixtureArrays/strided_f16" },
+            "nanStart": 3.0,
+            "nanCount": 2.0
+          }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float16", "shape": [2, 63, 3, 47], "tolerance": 0, "relTolerance": 0, "allowNaN": true }
+      },
+      "attrs": { "perm": [0, 3, 2, 1] }
+    },
+    {
+      "name": "strided_planes_empty_outer_float32",
+      "inputs": {
+        "x": {
+          "shape": [0, 3, 4],
+          "dtype": "float32",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_f32" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [4, 3, 0], "tolerance": 0, "relTolerance": 0 } }
+    },
+    {
+      "name": "strided_planes_empty_middle_float32",
+      "inputs": {
+        "x": {
+          "shape": [2, 16, 0, 32],
+          "dtype": "float32",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_f32" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 32, 0, 16], "tolerance": 0, "relTolerance": 0 } },
+      "attrs": { "perm": [0, 3, 2, 1] }
+    },
+    {
+      "name": "strided_planes_full_scalar_tile_boundary_float32",
+      "inputs": {
+        "x": {
+          "shape": [2, 33, 3, 64],
+          "dtype": "float32",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_f32" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 64, 3, 33], "tolerance": 0, "relTolerance": 0 } },
+      "attrs": { "perm": [0, 3, 2, 1] }
+    },
+    {
+      "name": "strided_planes_full_scalar_tile_boundary_float16",
+      "inputs": {
+        "x": {
+          "shape": [2, 33, 3, 64],
+          "dtype": "float16",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_f16" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [2, 64, 3, 33], "tolerance": 0, "relTolerance": 0 } },
+      "attrs": { "perm": [0, 3, 2, 1] }
+    },
+    {
+      "name": "strided_planes_two_axis_padding_47_float32",
+      "inputs": {
+        "x": {
+          "shape": [2, 47, 3, 65],
+          "dtype": "float32",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_f32" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 65, 3, 47], "tolerance": 0, "relTolerance": 0 } },
+      "attrs": { "perm": [0, 3, 2, 1] }
+    },
+    {
+      "name": "strided_planes_two_axis_padding_48_float32",
+      "inputs": {
+        "x": {
+          "shape": [2, 48, 3, 65],
+          "dtype": "float32",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_f32" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 65, 3, 48], "tolerance": 0, "relTolerance": 0 } },
+      "attrs": { "perm": [0, 3, 2, 1] }
+    },
+    {
+      "name": "strided_planes_two_axis_padding_47_float16",
+      "inputs": {
+        "x": {
+          "shape": [2, 47, 3, 65],
+          "dtype": "float16",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_f16" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [2, 65, 3, 47], "tolerance": 0, "relTolerance": 0 } },
+      "attrs": { "perm": [0, 3, 2, 1] }
+    },
+    {
+      "name": "strided_planes_two_axis_padding_48_float16",
+      "inputs": {
+        "x": {
+          "shape": [2, 48, 3, 65],
+          "dtype": "float16",
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/strided_f16" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [2, 65, 3, 48], "tolerance": 0, "relTolerance": 0 } },
+      "attrs": { "perm": [0, 3, 2, 1] }
+    },
+    {
+      "name": "int64_full_range_f32_subnormal_rank2_preserve_data",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/tensor/transpose_test.cc",
+        "test": "TransposeOpTest.TwoDim",
+        "notes": "Transpose is data movement with reindexing; selected finite subnormal values must survive the permutation."
+      },
+      "attrs": { "perm": [1, 0] },
+      "inputs": {
+        "x": {
+          "dtype": "int64",
+          "shape": [2, 3],
+          "data": {
+            "kind": "cycle",
+            "values": ["0", "4294967296", "8589934593", "-4294967297", "9223372036854775807", "-9223372036854775808", "9223372036854775806", "-1"]
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "int64", "shape": [3, 2], "tolerance": 0, "relTolerance": 0 } },
+      "tolerance": 0,
+      "relTolerance": 0
+    },
+    {
+      "name": "int64_full_range_rank1_default_perm",
+      "inputs": {
+        "x": {
+          "dtype": "int64",
+          "shape": [7],
+          "data": {
+            "kind": "cycle",
+            "values": ["4294967296", "8589934593", "-4294967297", "9223372036854775807", "-9223372036854775808", "9223372036854775806", "-1", "0"]
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "int64", "shape": [7], "tolerance": 0, "relTolerance": 0 } },
+      "tolerance": 0,
+      "relTolerance": 0
     }
   ]
 }

build/webgpu/transpose.wgsl.jinja CHANGED Viewed

@@ -12,9 +12,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   for (var i = begin; i < end; i = i + 1u) {
 {%- endmacro %}
 {% macro flat_tail_close() %}
-  }
-{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 fn input_offset({% if rank > 0 %}out_index: u32{% endif %}) -> u32 {
@@ -42,5 +40,5 @@ fn input_offset({% if rank > 0 %}out_index: u32{% endif %}) -> u32 {
 {{ flat_tail_open() }}
   y[i] = x[input_offset({% if rank > 0 %}i{% endif %})];
-{{ flat_tail_close() -}}
 }

   for (var i = begin; i < end; i = i + 1u) {
 {%- endmacro %}
 {% macro flat_tail_close() %}
+  }{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 fn input_offset({% if rank > 0 %}out_index: u32{% endif %}) -> u32 {
 {{ flat_tail_open() }}
   y[i] = x[input_offset({% if rank > 0 %}i{% endif %})];
+{{ flat_tail_close() }}
 }