sync 91d990483a17

Browse files

Files changed (10) hide show

README.md +19 -9
build/webgpu/bench.json +50 -1
build/webgpu/datamove-flat-copy.wgsl.jinja +1 -4
build/webgpu/datamove-transpose-2d-tiled-scalar.wgsl.jinja +9 -12
build/webgpu/datamove-transpose-2d-tiled.wgsl.jinja +9 -11
build/webgpu/datamove-transpose-vec4.wgsl.jinja +7 -10
build/webgpu/manifest.json +95 -161
build/webgpu/metadata.json +20 -11
build/webgpu/test.json +41 -10
build/webgpu/transpose.wgsl.jinja +15 -17

README.md CHANGED Viewed

@@ -18,15 +18,15 @@ See the [ONNX `Transpose` spec](https://onnx.ai/onnx/operators/onnx__Transpose.h
 ## Inputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `data` | `x` | `T` | — | — | The input tensor to transpose. | required |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `transposed` | `y` | `T` | same as `data` | — | The transposed output tensor with permuted axes. | required |
 ## Attributes
@@ -42,9 +42,16 @@ Attributes and default values (overridable per request):
 | --- | --- |
 | `T` | `float32`, `float16`, `int32`, `int16`, `uint32`, `uint8`, `int8`, `bool` |
 ## 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
@@ -56,15 +63,18 @@ Attributes and default values (overridable per request):
 ## Use with `@huggingface/kernels`
-The loader automatically allocates outputs whose metadata it can derive from the manifest contract and this call.
-The explicit `outputs` entries provide shape and logical dtype metadata for the results listed below:
-- `y`
-Each entry either requests an optional result or supplies metadata that cannot be inferred from the inputs.
 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.

 ## Inputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `x` | `data` | `T` | — | — | The input tensor to transpose. | required |
 ## Outputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `y` | `transposed` | `T` | same as `x` | — | The transposed output tensor with permuted axes. | required |
 ## Attributes
 | --- | --- |
 | `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.
 ## Files
+- [`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
 ## 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.
+This example supplies explicit metadata for:
+- `y`
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

build/webgpu/bench.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.Transpose",
   "tunableSpace": {
     "WORKGROUP_SIZE": [64, 128, 256],
     "INNER_VEC4_ITEMS_PER_THREAD": [1, 2, 4, 8],
@@ -216,6 +215,56 @@
       "inputs": { "x": { "shape": [1812, 384], "dtype": "float32", "dist": "normal", "seed": 9451, "scale": 0.2 } },
       "outputs": { "y": { "shape": [384, 1812], "dtype": "float32" } },
       "bench": { "metrics": [{ "type": "bandwidth", "value": "args.count * dtypeBytes(args.dtype) * 2" }] }
     }
   ]
 }

 {
   "tunableSpace": {
     "WORKGROUP_SIZE": [64, 128, 256],
     "INNER_VEC4_ITEMS_PER_THREAD": [1, 2, 4, 8],
       "inputs": { "x": { "shape": [1812, 384], "dtype": "float32", "dist": "normal", "seed": 9451, "scale": 0.2 } },
       "outputs": { "y": { "shape": [384, 1812], "dtype": "float32" } },
       "bench": { "metrics": [{ "type": "bandwidth", "value": "args.count * dtypeBytes(args.dtype) * 2" }] }
+    },
+    {
+      "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] },
+      "inputs": {
+        "x": { "shape": [8, 8, 7, 8, 7, 96], "dtype": "float32", "dist": "normal", "seed": 154, "scale": 0.2 }
+      },
+      "outputs": { "y": { "shape": [8, 8, 8, 7, 7, 96], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "args.count * dtypeBytes(args.dtype) * 2" }] }
+    },
+    {
+      "name": "transpose-f32-rank6-swin-window-partition-rank4-tiled-control",
+      "preset": "edge",
+      "provenance": {
+        "notes": "Equal-byte tiled control for the rank-6 window-partition case: the same 8x56x56x96 elements as an NHWC-to-NCHW permutation, which takes the vec4 staged transpose tile."
+      },
+      "vars": { "count": 2408448, "dtype": "float32" },
+      "attrs": { "perm": [0, 3, 1, 2] },
+      "inputs": { "x": { "shape": [8, 56, 56, 96], "dtype": "float32", "dist": "normal", "seed": 155, "scale": 0.2 } },
+      "outputs": { "y": { "shape": [8, 96, 56, 56], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "args.count * dtypeBytes(args.dtype) * 2" }] }
+    },
+    {
+      "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] },
+      "inputs": { "x": { "shape": [8, 64, 197, 64], "dtype": "float32", "dist": "normal", "seed": 156, "scale": 0.2 } },
+      "outputs": { "y": { "shape": [8, 64, 197, 64], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "args.count * dtypeBytes(args.dtype) * 2" }] }
+    },
+    {
+      "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" }] }
     }
   ]
 }

build/webgpu/datamove-flat-copy.wgsl.jinja CHANGED Viewed

@@ -1,9 +1,6 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
-const COUNT: u32 = {{ source.count }}u;
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})

 {{ env.wgsl.resourceDeclarations }}
+const COUNT: u32 = {{ count }}u;
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})

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

@@ -1,24 +1,21 @@
 // Tiled, optionally batched 2-D transpose over scalar elements.
 //
-// A 32x32 tile enters workgroup memory through coalesced reads and leaves
-// through coalesced writes. Guarded edges make this valid for any extent,
-// including shapes that cannot use the aligned vec4 path and would otherwise
-// fall back to a generic copy whose reads and writes cannot both be coalesced.
 //
 // 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 avoids bank conflicts on 32-bank GPUs.
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
-const ROWS: u32 = {{ source.rows }}u;
-const COLS: u32 = {{ source.cols }}u;
-const MAT: u32 = {{ source.rows * source.cols }}u;
 const TILE: u32 = 32u;
 const ROW_STEP: u32 = 8u;
@@ -46,7 +43,7 @@ fn main(
   // Write: lid.x runs along the destination row, which is the source column
   // direction — the transpose happens in the tile index, not in the memory
-  // access, so this side is coalesced too.
   for (var i = 0u; i < TILE; i = i + ROW_STEP) {
     let out_row = col_base + lid.y + i;
     let out_col = row_base + lid.x;

 // 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.
 {{ env.wgsl.resourceDeclarations }}
+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;
   // Write: lid.x runs along the destination row, which is the source column
   // direction — the transpose happens in the tile index, not in the memory
+  // access, so this side also traverses consecutive addresses.
   for (var i = 0u; i < TILE; i = i + ROW_STEP) {
     let out_row = col_base + lid.y + i;
     let out_col = row_base + lid.x;

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

@@ -5,19 +5,17 @@
 // - 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, so
-// both global loads and global stores are coalesced vec4 accesses. The +1
-// column pad avoids bank conflicts on 32-bank GPUs.
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
-const ROWS: u32 = {{ source.rows }}u;
-const COLS: u32 = {{ source.cols }}u;
-const ROWS_V: u32 = {{ ((source.rows / 4) | int) }}u;
-const COLS_V: u32 = {{ ((source.cols / 4) | int) }}u;
-const MAT_V: u32 = {{ (((source.rows * source.cols) / 4) | int) }}u;
 const TILE: u32 = 32u;
 var<workgroup> tile: array<array<{{ scalar }}, 33>, 32>;

 // - 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>;

build/webgpu/datamove-transpose-vec4.wgsl.jinja CHANGED Viewed

@@ -2,27 +2,24 @@
 // axis in place. Its dimension is a multiple of four, so both buffers bind as
 // vec4<T>. The index fold uses effective shapes whose last dimension is divided
 // by four, enabling 128-bit loads and stores.
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
-const COUNT: u32 = {{ source.count }}u;
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
 fn input_offset(out_index: u32) -> u32 {
   var rem = out_index;
   var offset = 0u;
-{% for out_axis in range(source.rank) %}
 {% set out_stride = namespace(value=1) %}
-{% for j in range(out_axis + 1, source.rank) %}
-{% set dj = ((source.outputShape[j] / 4) | int) if j == source.rank - 1 else source.outputShape[j] %}
 {% set out_stride.value = out_stride.value * dj %}
 {% endfor %}
-{% set in_axis = source.perm[out_axis] %}
 {% set in_stride = namespace(value=1) %}
-{% for j in range(in_axis + 1, source.rank) %}
-{% set dj = ((source.inputShape[j] / 4) | int) if j == source.rank - 1 else source.inputShape[j] %}
 {% set in_stride.value = in_stride.value * dj %}
 {% endfor %}
   let coord{{ out_axis }} = rem / {{ out_stride.value }}u;

 // axis in place. Its dimension is a multiple of four, so both buffers bind as
 // vec4<T>. The index fold uses effective shapes whose last dimension is divided
 // by four, enabling 128-bit loads and stores.
 {{ env.wgsl.resourceDeclarations }}
+const COUNT: u32 = {{ count }}u;
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
 fn input_offset(out_index: u32) -> u32 {
   var rem = out_index;
   var offset = 0u;
+{% for out_axis in range(rank) %}
 {% set out_stride = namespace(value=1) %}
+{% for j in range(out_axis + 1, rank) %}
+{% set dj = ((outputShape[j] / 4) | int) if j == rank - 1 else outputShape[j] %}
 {% set out_stride.value = out_stride.value * dj %}
 {% endfor %}
+{% set in_axis = perm[out_axis] %}
 {% set in_stride = namespace(value=1) %}
+{% for j in range(in_axis + 1, rank) %}
+{% set dj = ((inputShape[j] / 4) | int) if j == rank - 1 else inputShape[j] %}
 {% set in_stride.value = in_stride.value * dj %}
 {% endfor %}
   let coord{{ out_axis }} = rem / {{ out_stride.value }}u;

build/webgpu/manifest.json CHANGED Viewed

@@ -2,171 +2,96 @@
   "domain": "ai.onnx",
   "name": "Transpose",
   "sinceVersion": 13,
-  "description": "Transposes the input tensor by permuting its axes according to the `perm` attribute. Axis `i` of the output corresponds to axis `perm[i]` of the input; if `perm` is omitted, the axes are reversed (`n-1, ..., 0`).",
-  "inputs": [{ "role": "data", "dtype": "T", "description": "The input tensor to transpose." }],
-  "outputs": [
-    {
-      "role": "transposed",
-      "dtype": "T",
-      "description": "The transposed output tensor with permuted axes.",
-      "rank": "ranks.data"
-    }
-  ],
-  "attributes": {},
-  "attributeDescriptions": {
-    "perm": "Optional permutation of the input axes. It must contain every axis from 0 through rank - 1 exactly once. When omitted, the axes are reversed."
-  },
   "typeConstraints": { "T": ["float32", "float16", "int32", "int16", "uint32", "uint8", "int8", "bool"] },
-  "args": {
-    "x": { "kind": "tensor", "semantic": "data", "role": "input" },
-    "y": { "kind": "tensor", "semantic": "transposed", "role": "output" }
-  },
-  "tunables": { "WORKGROUP_SIZE": 256, "INNER_VEC4_ITEMS_PER_THREAD": 4, "INNER_VEC4_MIN_THREADS": 65536 },
-  "tunableDescriptions": {
-    "INNER_VEC4_ITEMS_PER_THREAD": "Sets the target maximum number of vec4 items assigned to each invocation by the inner-axis vector route.",
-    "INNER_VEC4_MIN_THREADS": "Sets the minimum logical thread count retained when shrinking the inner-axis vector dispatch."
   },
   "derive": {
     "dtypeSupported": "f16Ok(dtypes.T)",
-    "rank0IO": "ranks.data == 0 and ranks.transposed == 0",
-    "rank1IO": "ranks.data == 1 and ranks.transposed == 1",
-    "rank2IO": "ranks.data == 2 and ranks.transposed == 2",
-    "rank3IO": "ranks.data == 3 and ranks.transposed == 3",
-    "rank4IO": "ranks.data == 4 and ranks.transposed == 4",
-    "rank5IO": "ranks.data == 5 and ranks.transposed == 5",
-    "rank6IO": "ranks.data == 6 and ranks.transposed == 6",
-    "rank7IO": "ranks.data == 7 and ranks.transposed == 7",
     "permPresent": "has(attrs, \"perm\")",
-    "permContract": "not permPresent or ((attrs.perm | length) == ranks.data and isUniqueIntList(attrs.perm, 0, ranks.data))",
-    "rank0Permutation": "rank0IO and permContract",
-    "rank1Permutation": "rank1IO and permContract and dim(shapes.transposed, 0) == dim(shapes.data, 0)",
-    "explicitRank2": "rank2IO and permContract and permPresent and dim(shapes.transposed, 0) == dim(shapes.data, attrs.perm[0]) and dim(shapes.transposed, 1) == dim(shapes.data, attrs.perm[1])",
-    "explicitRank3": "rank3IO and permContract and permPresent and dim(shapes.transposed, 0) == dim(shapes.data, attrs.perm[0]) and dim(shapes.transposed, 1) == dim(shapes.data, attrs.perm[1]) and dim(shapes.transposed, 2) == dim(shapes.data, attrs.perm[2])",
-    "explicitRank4": "rank4IO and permContract and permPresent and dim(shapes.transposed, 0) == dim(shapes.data, attrs.perm[0]) and dim(shapes.transposed, 1) == dim(shapes.data, attrs.perm[1]) and dim(shapes.transposed, 2) == dim(shapes.data, attrs.perm[2]) and dim(shapes.transposed, 3) == dim(shapes.data, attrs.perm[3])",
-    "explicitRank5": "rank5IO and permContract and permPresent and dim(shapes.transposed, 0) == dim(shapes.data, attrs.perm[0]) and dim(shapes.transposed, 1) == dim(shapes.data, attrs.perm[1]) and dim(shapes.transposed, 2) == dim(shapes.data, attrs.perm[2]) and dim(shapes.transposed, 3) == dim(shapes.data, attrs.perm[3]) and dim(shapes.transposed, 4) == dim(shapes.data, attrs.perm[4])",
-    "explicitRank6": "rank6IO and permContract and permPresent and dim(shapes.transposed, 0) == dim(shapes.data, attrs.perm[0]) and dim(shapes.transposed, 1) == dim(shapes.data, attrs.perm[1]) and dim(shapes.transposed, 2) == dim(shapes.data, attrs.perm[2]) and dim(shapes.transposed, 3) == dim(shapes.data, attrs.perm[3]) and dim(shapes.transposed, 4) == dim(shapes.data, attrs.perm[4]) and dim(shapes.transposed, 5) == dim(shapes.data, attrs.perm[5])",
-    "explicitRank7": "rank7IO and permContract and permPresent and dim(shapes.transposed, 0) == dim(shapes.data, attrs.perm[0]) and dim(shapes.transposed, 1) == dim(shapes.data, attrs.perm[1]) and dim(shapes.transposed, 2) == dim(shapes.data, attrs.perm[2]) and dim(shapes.transposed, 3) == dim(shapes.data, attrs.perm[3]) and dim(shapes.transposed, 4) == dim(shapes.data, attrs.perm[4]) and dim(shapes.transposed, 5) == dim(shapes.data, attrs.perm[5]) and dim(shapes.transposed, 6) == dim(shapes.data, attrs.perm[6])",
-    "defaultRank2": "rank2IO and permContract and not permPresent and dim(shapes.transposed, 0) == dim(shapes.data, 1) and dim(shapes.transposed, 1) == dim(shapes.data, 0)",
-    "defaultRank3": "rank3IO and permContract and not permPresent and dim(shapes.transposed, 0) == dim(shapes.data, 2) and dim(shapes.transposed, 1) == dim(shapes.data, 1) and dim(shapes.transposed, 2) == dim(shapes.data, 0)",
-    "defaultRank4": "rank4IO and permContract and not permPresent and dim(shapes.transposed, 0) == dim(shapes.data, 3) and dim(shapes.transposed, 1) == dim(shapes.data, 2) and dim(shapes.transposed, 2) == dim(shapes.data, 1) and dim(shapes.transposed, 3) == dim(shapes.data, 0)",
-    "defaultRank5": "rank5IO and permContract and not permPresent and dim(shapes.transposed, 0) == dim(shapes.data, 4) and dim(shapes.transposed, 1) == dim(shapes.data, 3) and dim(shapes.transposed, 2) == dim(shapes.data, 2) and dim(shapes.transposed, 3) == dim(shapes.data, 1) and dim(shapes.transposed, 4) == dim(shapes.data, 0)",
-    "defaultRank6": "rank6IO and permContract and not permPresent and dim(shapes.transposed, 0) == dim(shapes.data, 5) and dim(shapes.transposed, 1) == dim(shapes.data, 4) and dim(shapes.transposed, 2) == dim(shapes.data, 3) and dim(shapes.transposed, 3) == dim(shapes.data, 2) and dim(shapes.transposed, 4) == dim(shapes.data, 1) and dim(shapes.transposed, 5) == dim(shapes.data, 0)",
-    "defaultRank7": "rank7IO and permContract and not permPresent and dim(shapes.transposed, 0) == dim(shapes.data, 6) and dim(shapes.transposed, 1) == dim(shapes.data, 5) and dim(shapes.transposed, 2) == dim(shapes.data, 4) and dim(shapes.transposed, 3) == dim(shapes.data, 3) and dim(shapes.transposed, 4) == dim(shapes.data, 2) and dim(shapes.transposed, 5) == dim(shapes.data, 1) and dim(shapes.transposed, 6) == dim(shapes.data, 0)",
-    "genericTransposeOk": "rank0Permutation or rank1Permutation or explicitRank2 or defaultRank2 or explicitRank3 or defaultRank3 or explicitRank4 or defaultRank4 or explicitRank5 or defaultRank5 or explicitRank6 or defaultRank6 or explicitRank7 or defaultRank7",
-    "defaultPerm": "pick([[rank0IO, []], [rank1IO, [0]], [rank2IO, [1, 0]], [rank3IO, [2, 1, 0]], [rank4IO, [3, 2, 1, 0]], [rank5IO, [4, 3, 2, 1, 0]], [rank6IO, [5, 4, 3, 2, 1, 0]], [rank7IO, [6, 5, 4, 3, 2, 1, 0]]], [])",
     "genericPerm": "pick([[permPresent, attrs.perm]], defaultPerm)",
-    "innerVec4Rank3Ok": "explicitRank3 and attrs.perm[2] == 2 and dim(shapes.data, 2) % 4 == 0 and numel(shapes.transposed) > 0 and dtypeSupported",
-    "innerVec4Rank4Ok": "explicitRank4 and attrs.perm[3] == 3 and dim(shapes.data, 3) % 4 == 0 and numel(shapes.transposed) > 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.data) > 0 and ceil(dim(shapes.data, 0) / 32) <= device.limits.maxComputeWorkgroupsPerDimension and ceil(dim(shapes.data, 1) / 32) <= device.limits.maxComputeWorkgroupsPerDimension and dtypeSupported",
-    "tiled2dOk": "tiled2dShapeOk and dim(shapes.data, 0) % 4 == 0 and dim(shapes.data, 1) % 4 == 0",
-    "tiled3dShapeOk": "explicitRank3 and attrs.perm[0] == 0 and attrs.perm[1] == 2 and attrs.perm[2] == 1 and numel(shapes.data) > 0 and dim(shapes.data, 0) <= device.limits.maxComputeWorkgroupsPerDimension and ceil(dim(shapes.data, 1) / 32) <= device.limits.maxComputeWorkgroupsPerDimension and ceil(dim(shapes.data, 2) / 32) <= device.limits.maxComputeWorkgroupsPerDimension and dtypeSupported",
-    "tiled3dOk": "tiled3dShapeOk and dim(shapes.data, 1) % 4 == 0 and dim(shapes.data, 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.data) > 0 and dim(shapes.data, 0) * dim(shapes.data, 1) <= device.limits.maxComputeWorkgroupsPerDimension and ceil(dim(shapes.data, 2) / 32) <= device.limits.maxComputeWorkgroupsPerDimension and ceil(dim(shapes.data, 3) / 32) <= device.limits.maxComputeWorkgroupsPerDimension and dtypeSupported",
-    "tiled4dOk": "tiled4dShapeOk and dim(shapes.data, 2) % 4 == 0 and dim(shapes.data, 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.data) > 0 and dim(shapes.data, 0) <= device.limits.maxComputeWorkgroupsPerDimension and ceil(dim(shapes.data, 1) / 32) <= device.limits.maxComputeWorkgroupsPerDimension and ceil((dim(shapes.data, 2) * dim(shapes.data, 3) * dim(shapes.data, 4)) / 32) <= device.limits.maxComputeWorkgroupsPerDimension and dtypeSupported",
-    "tiled5dNcdhwToNdhwcOk": "tiled5dNcdhwToNdhwcShapeOk and dim(shapes.data, 1) % 4 == 0 and (dim(shapes.data, 2) * dim(shapes.data, 3) * dim(shapes.data, 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.data) > 0 and dim(shapes.data, 0) <= device.limits.maxComputeWorkgroupsPerDimension and ceil(dim(shapes.data, 1) * dim(shapes.data, 2) * dim(shapes.data, 3) / 32) <= device.limits.maxComputeWorkgroupsPerDimension and ceil(dim(shapes.data, 4) / 32) <= device.limits.maxComputeWorkgroupsPerDimension and dtypeSupported",
-    "tiled5dNdhwcToNcdhwOk": "tiled5dNdhwcToNcdhwShapeOk and dim(shapes.data, 1) * dim(shapes.data, 2) * dim(shapes.data, 3) % 4 == 0 and dim(shapes.data, 4) % 4 == 0",
-    "memcpyPerm3": "explicitRank3 and (dim(shapes.data, attrs.perm[0]) == 1 or dim(shapes.data, attrs.perm[1]) == 1 or attrs.perm[0] < attrs.perm[1]) and (dim(shapes.data, attrs.perm[0]) == 1 or dim(shapes.data, attrs.perm[2]) == 1 or attrs.perm[0] < attrs.perm[2]) and (dim(shapes.data, attrs.perm[1]) == 1 or dim(shapes.data, attrs.perm[2]) == 1 or attrs.perm[1] < attrs.perm[2])",
-    "memcpyPerm4": "explicitRank4 and (dim(shapes.data, attrs.perm[0]) == 1 or dim(shapes.data, attrs.perm[1]) == 1 or attrs.perm[0] < attrs.perm[1]) and (dim(shapes.data, attrs.perm[0]) == 1 or dim(shapes.data, attrs.perm[2]) == 1 or attrs.perm[0] < attrs.perm[2]) and (dim(shapes.data, attrs.perm[0]) == 1 or dim(shapes.data, attrs.perm[3]) == 1 or attrs.perm[0] < attrs.perm[3]) and (dim(shapes.data, attrs.perm[1]) == 1 or dim(shapes.data, attrs.perm[2]) == 1 or attrs.perm[1] < attrs.perm[2]) and (dim(shapes.data, attrs.perm[1]) == 1 or dim(shapes.data, attrs.perm[3]) == 1 or attrs.perm[1] < attrs.perm[3]) and (dim(shapes.data, attrs.perm[2]) == 1 or dim(shapes.data, attrs.perm[3]) == 1 or attrs.perm[2] < attrs.perm[3])",
-    "memcpyPerm5": "explicitRank5 and (dim(shapes.data, attrs.perm[0]) == 1 or dim(shapes.data, attrs.perm[1]) == 1 or attrs.perm[0] < attrs.perm[1]) and (dim(shapes.data, attrs.perm[0]) == 1 or dim(shapes.data, attrs.perm[2]) == 1 or attrs.perm[0] < attrs.perm[2]) and (dim(shapes.data, attrs.perm[0]) == 1 or dim(shapes.data, attrs.perm[3]) == 1 or attrs.perm[0] < attrs.perm[3]) and (dim(shapes.data, attrs.perm[0]) == 1 or dim(shapes.data, attrs.perm[4]) == 1 or attrs.perm[0] < attrs.perm[4]) and (dim(shapes.data, attrs.perm[1]) == 1 or dim(shapes.data, attrs.perm[2]) == 1 or attrs.perm[1] < attrs.perm[2]) and (dim(shapes.data, attrs.perm[1]) == 1 or dim(shapes.data, attrs.perm[3]) == 1 or attrs.perm[1] < attrs.perm[3]) and (dim(shapes.data, attrs.perm[1]) == 1 or dim(shapes.data, attrs.perm[4]) == 1 or attrs.perm[1] < attrs.perm[4]) and (dim(shapes.data, attrs.perm[2]) == 1 or dim(shapes.data, attrs.perm[3]) == 1 or attrs.perm[2] < attrs.perm[3]) and (dim(shapes.data, attrs.perm[2]) == 1 or dim(shapes.data, attrs.perm[4]) == 1 or attrs.perm[2] < attrs.perm[4]) and (dim(shapes.data, attrs.perm[3]) == 1 or dim(shapes.data, attrs.perm[4]) == 1 or attrs.perm[3] < attrs.perm[4])",
     "memcpyPermOk": "memcpyPerm3 or memcpyPerm4 or memcpyPerm5",
-    "identityPermutation": "rank1Permutation or (explicitRank2 and attrs.perm[0] == 0 and attrs.perm[1] == 1) or (explicitRank3 and attrs.perm[0] == 0 and attrs.perm[1] == 1 and attrs.perm[2] == 2) or (explicitRank4 and attrs.perm[0] == 0 and attrs.perm[1] == 1 and attrs.perm[2] == 2 and attrs.perm[3] == 3) or (explicitRank5 and attrs.perm[0] == 0 and attrs.perm[1] == 1 and attrs.perm[2] == 2 and attrs.perm[3] == 3 and attrs.perm[4] == 4) or (explicitRank6 and attrs.perm[0] == 0 and attrs.perm[1] == 1 and attrs.perm[2] == 2 and attrs.perm[3] == 3 and attrs.perm[4] == 4 and attrs.perm[5] == 5)",
-    "identityCopyOk": "numel(shapes.transposed) > 0 and numel(shapes.transposed) % 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.data) > 0 and dim(shapes.data, 0) <= device.limits.maxComputeWorkgroupsPerDimension and ceil(dim(shapes.data, 1) / 32) <= device.limits.maxComputeWorkgroupsPerDimension and ceil((dim(shapes.data, 2) * dim(shapes.data, 3)) / 32) <= device.limits.maxComputeWorkgroupsPerDimension and dtypeSupported",
-    "tiled4dNchwToNhwcOk": "tiled4dNchwToNhwcShapeOk and dim(shapes.data, 1) % 4 == 0 and (dim(shapes.data, 2) * dim(shapes.data, 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.data) > 0 and dim(shapes.data, 0) <= device.limits.maxComputeWorkgroupsPerDimension and ceil((dim(shapes.data, 1) * dim(shapes.data, 2)) / 32) <= device.limits.maxComputeWorkgroupsPerDimension and ceil(dim(shapes.data, 3) / 32) <= device.limits.maxComputeWorkgroupsPerDimension and dtypeSupported",
-    "tiled4dNhwcToNchwOk": "tiled4dNhwcToNchwShapeOk and (dim(shapes.data, 1) * dim(shapes.data, 2)) % 4 == 0 and dim(shapes.data, 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.data, 0) if tiled2dShapeOk else (dim(shapes.data, 1) if tiled3dShapeOk else (dim(shapes.data, 2) if tiled4dShapeOk else (dim(shapes.data, 1) if tiled5dNcdhwToNdhwcShapeOk else (dim(shapes.data, 1) * dim(shapes.data, 2) * dim(shapes.data, 3) if tiled5dNdhwcToNcdhwShapeOk else (dim(shapes.data, 1) if tiled4dNchwToNhwcShapeOk else (dim(shapes.data, 1) * dim(shapes.data, 2) if tiled4dNhwcToNchwShapeOk else 1))))))",
-    "tiledCols": "dim(shapes.data, 1) if tiled2dShapeOk else (dim(shapes.data, 2) if tiled3dShapeOk else (dim(shapes.data, 3) if tiled4dShapeOk else (dim(shapes.data, 2) * dim(shapes.data, 3) * dim(shapes.data, 4) if tiled5dNcdhwToNdhwcShapeOk else (dim(shapes.data, 4) if tiled5dNdhwcToNcdhwShapeOk else (dim(shapes.data, 2) * dim(shapes.data, 3) if tiled4dNchwToNhwcShapeOk else (dim(shapes.data, 3) if tiled4dNhwcToNchwShapeOk else 1))))))",
-    "tiledBatch": "1 if tiled2dShapeOk else (dim(shapes.data, 0) if tiled3dShapeOk else (dim(shapes.data, 0) * dim(shapes.data, 1) if tiled4dShapeOk else (dim(shapes.data, 0) if tiled5dNcdhwToNdhwcShapeOk or tiled5dNdhwcToNcdhwShapeOk or tiled4dNchwToNhwcShapeOk or tiled4dNhwcToNchwShapeOk else 1)))"
   },
-  "constants": { "usesF16": "dtypes.T == \"f16\"" },
-  "bindingSets": {
-    "vectorData": [
-      {
-        "name": "src",
-        "arg": "x",
-        "semantic": "data",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$vectorScalar"
-      },
-      {
-        "name": "dst",
-        "arg": "y",
-        "semantic": "transposed",
-        "buffer": { "type": "storage" },
-        "elementType": "$vectorScalar"
-      }
-    ],
-    "scalarData": [
-      {
-        "name": "x",
-        "arg": "x",
-        "semantic": "data",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      { "name": "y", "arg": "y", "semantic": "transposed", "buffer": { "type": "storage" }, "elementType": "$scalar" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.transposed)" }]
-        }
-      }
-    ],
-    "scalarDataNoParams": [
-      {
-        "name": "x",
-        "arg": "x",
-        "semantic": "data",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      { "name": "y", "arg": "y", "semantic": "transposed", "buffer": { "type": "storage" }, "elementType": "$scalar" }
-    ]
   },
   "variants": [
     {
       "id": "tiled_vec4",
       "priority": 22,
       "when": ["tiledTransposeOk"],
-      "constants": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Transpose.tiled",
-          "source": {
-            "shader": "datamove-transpose-2d-tiled.wgsl.jinja",
-            "inputs": { "rows": "tiledRows", "cols": "tiledCols" }
-          },
-          "bindings": "vectorData",
           "dispatch": { "x": "ceil(tiledCols / 32)", "y": "ceil(tiledRows / 32)", "z": "tiledBatch" }
         }
       ]
     },
     {
       "id": "tiled_scalar",
-      "description": "Stages ragged two-dimensional transpose planes through a padded workgroup tile so reads and writes remain coalesced when vec4 alignment is unavailable.",
       "priority": 19,
       "when": ["tiledScalarTransposeOk", "dtypeSupported"],
-      "constants": { "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "Transpose.tiledScalar",
-          "source": {
-            "shader": "datamove-transpose-2d-tiled-scalar.wgsl.jinja",
-            "inputs": { "rows": "tiledRows", "cols": "tiledCols" }
-          },
-          "bindings": "scalarDataNoParams",
           "dispatch": { "x": "ceil(tiledCols / 32)", "y": "ceil(tiledRows / 32)", "z": "tiledBatch" }
         }
       ]
@@ -175,41 +100,44 @@
       "id": "identity_vec4",
       "priority": 21,
       "when": ["identityCopyOk"],
-      "constants": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Transpose.identityVec4",
-          "source": { "shader": "datamove-flat-copy.wgsl.jinja", "inputs": { "count": "numel(shapes.transposed) / 4" } },
-          "bindings": "vectorData",
-          "dispatch": { "gridStride": "numel(shapes.transposed) / 4", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "inner_vec4",
-      "description": "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.",
       "priority": 20,
       "when": ["innerVec4Ok"],
-      "constants": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Transpose.innerVec4",
-          "source": {
-            "shader": "datamove-transpose-vec4.wgsl.jinja",
-            "inputs": {
-              "inputShape": "shapes.data",
-              "outputShape": "shapes.transposed",
-              "rank": "ranks.data",
-              "perm": "innerVec4Perm",
-              "count": "numel(shapes.transposed) / 4"
-            }
           },
-          "bindings": "vectorData",
           "dispatch": {
-            "gridStride": "max(min(numel(shapes.transposed) / 4, tunables.INNER_VEC4_MIN_THREADS), ceilDiv(numel(shapes.transposed) / 4, tunables.INNER_VEC4_ITEMS_PER_THREAD))",
-            "workgroupSize": "tunables.WORKGROUP_SIZE"
           }
         }
       ]
@@ -217,23 +145,29 @@
     {
       "id": "generic",
       "when": ["genericTransposeOk", "dtypeSupported"],
-      "constants": { "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "Transpose",
-          "source": {
-            "shader": "transpose.wgsl.jinja",
-            "inputs": {
-              "inputShape": "shapes.data",
-              "outputShape": "shapes.transposed",
-              "rank": "ranks.data",
-              "perm": "genericPerm",
-              "itemsPerInvocation": 4
-            }
           },
-          "bindings": "scalarData",
-          "dispatch": { "threads": "ceilDiv(numel(shapes.transposed), 4)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     }

   "domain": "ai.onnx",
   "name": "Transpose",
   "sinceVersion": 13,
+  "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",
+    "rank4IO": "ranks.x == 4 and ranks.y == 4",
+    "rank5IO": "ranks.x == 5 and ranks.y == 5",
     "permPresent": "has(attrs, \"perm\")",
+    "permContract": "not permPresent or ((attrs.perm | length) == ranks.x and isUniqueIntList(attrs.perm, 0, ranks.x))",
+    "explicitRank2": "rank2IO and permContract and permPresent and dim(shapes.y, 0) == dim(shapes.x, attrs.perm[0]) and dim(shapes.y, 1) == dim(shapes.x, attrs.perm[1])",
+    "explicitRank3": "rank3IO and permContract and permPresent and dim(shapes.y, 0) == dim(shapes.x, attrs.perm[0]) and dim(shapes.y, 1) == dim(shapes.x, attrs.perm[1]) and dim(shapes.y, 2) == dim(shapes.x, attrs.perm[2])",
+    "explicitRank4": "rank4IO and permContract and permPresent and dim(shapes.y, 0) == dim(shapes.x, attrs.perm[0]) and dim(shapes.y, 1) == dim(shapes.x, attrs.perm[1]) and dim(shapes.y, 2) == dim(shapes.x, attrs.perm[2]) and dim(shapes.y, 3) == dim(shapes.x, attrs.perm[3])",
+    "explicitRank5": "rank5IO and permContract and permPresent and dim(shapes.y, 0) == dim(shapes.x, attrs.perm[0]) and dim(shapes.y, 1) == dim(shapes.x, attrs.perm[1]) and dim(shapes.y, 2) == dim(shapes.x, attrs.perm[2]) and dim(shapes.y, 3) == dim(shapes.x, attrs.perm[3]) and dim(shapes.y, 4) == dim(shapes.x, attrs.perm[4])",
+    "defaultRank2": "rank2IO and permContract and not permPresent and dim(shapes.y, 0) == dim(shapes.x, 1) and dim(shapes.y, 1) == dim(shapes.x, 0)",
+    "defaultPerm": "reversedAxes(ranks.x)",
     "genericPerm": "pick([[permPresent, attrs.perm]], defaultPerm)",
+    "genericTransposeOk": "ranks.x == ranks.y and permContract and sameShape(shapes.y, permuteShape(shapes.x, genericPerm))",
+    "innerVec4Rank3Ok": "explicitRank3 and attrs.perm[2] == 2 and dim(shapes.x, 2) % 4 == 0 and numel(shapes.y) > 0 and dtypeSupported",
+    "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])",
+    "memcpyPerm5": "explicitRank5 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[0]) == 1 or dim(shapes.x, attrs.perm[4]) == 1 or attrs.perm[0] < attrs.perm[4]) 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[1]) == 1 or dim(shapes.x, attrs.perm[4]) == 1 or attrs.perm[1] < attrs.perm[4]) and (dim(shapes.x, attrs.perm[2]) == 1 or dim(shapes.x, attrs.perm[3]) == 1 or attrs.perm[2] < attrs.perm[3]) and (dim(shapes.x, attrs.perm[2]) == 1 or dim(shapes.x, attrs.perm[4]) == 1 or attrs.perm[2] < attrs.perm[4]) and (dim(shapes.x, attrs.perm[3]) == 1 or dim(shapes.x, attrs.perm[4]) == 1 or attrs.perm[3] < attrs.perm[4])",
     "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"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
+            "y": 1,
+            "z": 1
+          }
         }
       ]
     },
     {
       "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))",
+            "y": 1,
+            "z": 1
           }
         }
       ]
     {
       "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",
+            { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }] }
+          ],
+          "dispatch": {
+            "x": "min(ceilDiv((ceilDiv(numel(shapes.y), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.y), 4)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,22 +1,31 @@
 {
   "name": "ai.onnx.Transpose",
-  "id": "_ai_onnx_transpose_webgpu_713616c",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "wT5OzWQiJ0OSKnc+JuSj0vOoFzY1Q/JzgJW1jpIbv5Q=",
-      "datamove-flat-copy.wgsl.jinja": "rEpeYMvIZQuQnB2f3lt51jrrklKICsqUTzCJPRTknNk=",
-      "datamove-transpose-2d-tiled-scalar.wgsl.jinja": "zy7rlN8Se+mn8RB9akbBTHq1RK8ot6NBboxG3vX1Xdg=",
-      "datamove-transpose-2d-tiled.wgsl.jinja": "AtmY7NlBNCV1unIf2epvpBO3hfnLBmgGKGB7WiWMCQk=",
-      "datamove-transpose-vec4.wgsl.jinja": "RgwjTbu7SHDVsUW7W2f5S81gv1UJBpT93udQX/MZdJ0=",
-      "manifest.json": "Av+HRbQxG96oiNzrlouEobs2tS/VmfX0tCkNdy8rWL4=",
-      "test.json": "acNlC/DcnpoXb3ni/Z4QMbpYD7C1YR16iHXX20mNp0Y=",
-      "transpose.wgsl.jinja": "lYYw/Xqp/wTaYlWr8iB/vn1ZxacE3jOes9gYViMlhlo="
     }
   },
-  "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.Transpose" }
 }

 {
   "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"],
+      "inner_vec4": ["datamove-transpose-vec4.wgsl.jinja"],
+      "generic": ["transpose.wgsl.jinja"]
+    }
+  }
 }

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.Transpose",
   "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],
@@ -189,7 +188,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/tensor/transpose_test.cc",
         "test": "TransposeOpTest.TwoDimNoAttr",
-        "notes": "Same two-dimensional data-movement fixture with an ONNX-valid bool payload."
       },
       "inputs": {
         "x": { "dtype": "bool", "shape": [2, 3], "data": { "kind": "values", "values": [1, 0, 1, 0, 1, 0] } }
@@ -249,7 +248,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/tensor/transpose_test.cc",
         "test": "TransposeOpTest.TwoDim_int32",
-        "notes": "ORT covers int32; this framework covers the same 2D explicit transpose with supported uint32 data."
       },
       "attrs": { "perm": [1, 0] },
       "inputs": {
@@ -401,7 +400,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/tensor/transpose_test.cc",
         "test": "TransposeOpTest.SingleAxisMovingInwardsBlockCopy",
-        "notes": "ORT uses uint64 payload values; this framework stores the representable test values as uint32."
       },
       "attrs": { "perm": [1, 2, 0, 3] },
       "inputs": {
@@ -807,7 +806,7 @@
     {
       "name": "rank2_scalar_odd_width_5x7",
       "provenance": {
-        "notes": "Compact correctness sibling for the odd-width rank2 scalar transpose benchmark; the dimensions avoid vec4/tiled alignment while preserving the same rank2 permutation."
       },
       "attrs": { "perm": [1, 0] },
       "inputs": {
@@ -824,7 +823,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/tensor/transpose_test.cc",
         "test": "TransposeOpTest (generic N-dim)",
-        "notes": "Rank-7 explicit-permutation coverage for the generic transpose path. The manifest supplies the rank and permutation as template inputs instead of maintaining a rank-specific variant."
       },
       "attrs": { "perm": [1, 0, 3, 2, 5, 4, 6] },
       "inputs": {
@@ -844,7 +843,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/tensor/transpose_test.cc",
         "test": "TransposeOpTest.TwoDim (bool payload)",
-        "notes": "bool is in typeConstraints T and scalarType maps it to u32, so an aligned [4,4] bool transpose selects tiled_vec4 (gate is only f16Ok, true for u32). The pre-existing ort_two_dim_default_bool_gpu_gap fixture is stale-named; this NON-gpu-gapped case gives bool its first real GPU validation on the vec4 path."
       },
       "attrs": { "perm": [1, 0] },
       "inputs": {
@@ -861,7 +860,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/tensor/transpose_test.cc",
         "test": "TransposeOpTest.TransposeReshape (scrambled rank5)",
-        "notes": "Genuinely scrambled rank-5 explicit perm [4,2,0,3,1] exercising the generic scalar-gather path with rank-specific template inputs; prior rank-5 coverage only exercised a default reversal and a near-reshape perm."
       },
       "attrs": { "perm": [4, 2, 0, 3, 1] },
       "inputs": {
@@ -893,7 +892,7 @@
     {
       "name": "rank1_2d_fold_boundary_last_element_guard",
       "provenance": {
-        "notes": "numel = 2^24+1 = 16777217 forces ceilDiv(numel,256)=65536 > maxComputeWorkgroupsPerDimension(65535), engaging the scalar path's y-axis 2D fold (i = gid.x + gid.y*nwg.x*WG). rank1 numel%4 != 0 bypasses identity_vec4 and lands on the generic scalar path, isolating the fold last-row guard. Referenceless: TS reference supplies truth so a corrupted/dropped tail element in the second y-row fails the compare."
       },
       "inputs": {
         "x": {
@@ -1081,7 +1080,7 @@
       "provenance": {
         "source": "onnxruntime/core/providers/webgpu/tensor/transpose.cc",
         "test": "unit-dim permutations lower to a flat copy",
-        "notes": "A permutation that only moves size-1 axes preserves the row-major layout, so the flat vec4 copy is exact."
       },
       "attrs": { "perm": [0, 2, 1, 3] },
       "inputs": {
@@ -1092,6 +1091,38 @@
         }
       },
       "outputs": { "y": { "dtype": "float32", "shape": [2, 4, 3, 4], "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],
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/tensor/transpose_test.cc",
         "test": "TransposeOpTest.TwoDimNoAttr",
+        "notes": "A two-dimensional transpose with an ONNX-valid bool payload verifies bit-preserving data movement."
       },
       "inputs": {
         "x": { "dtype": "bool", "shape": [2, 3], "data": { "kind": "values", "values": [1, 0, 1, 0, 1, 0] } }
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/tensor/transpose_test.cc",
         "test": "TransposeOpTest.TwoDim_int32",
+        "notes": "A uint32 payload exercises an explicit two-dimensional transpose."
       },
       "attrs": { "perm": [1, 0] },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/tensor/transpose_test.cc",
         "test": "TransposeOpTest.SingleAxisMovingInwardsBlockCopy",
+        "notes": "Projects representable uint64 source values to uint32 storage."
       },
       "attrs": { "perm": [1, 2, 0, 3] },
       "inputs": {
     {
       "name": "rank2_scalar_odd_width_5x7",
       "provenance": {
+        "notes": "A 5-by-7 transpose avoids vector and tiled alignment, exercising scalar rank-2 address mapping."
       },
       "attrs": { "perm": [1, 0] },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/tensor/transpose_test.cc",
         "test": "TransposeOpTest (generic N-dim)",
+        "notes": "A rank-7 explicit permutation exercises generic high-rank index decomposition."
       },
       "attrs": { "perm": [1, 0, 3, 2, 5, 4, 6] },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/tensor/transpose_test.cc",
         "test": "TransposeOpTest.TwoDim (bool payload)",
+        "notes": "An aligned 4-by-4 bool tensor uses u32 storage and exercises vectorized tiled transpose without requiring subnormal preservation."
       },
       "attrs": { "perm": [1, 0] },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/tensor/transpose_test.cc",
         "test": "TransposeOpTest.TransposeReshape (scrambled rank5)",
+        "notes": "The rank-5 permutation [4,2,0,3,1] scrambles all nontrivial strides and exercises generic scalar gathering."
       },
       "attrs": { "perm": [4, 2, 0, 3, 1] },
       "inputs": {
     {
       "name": "rank1_2d_fold_boundary_last_element_guard",
       "provenance": {
+        "notes": "A rank-1 identity transpose with 16,777,217 elements requires 65,536 workgroups and therefore a two-row scalar dispatch. The odd element count exercises the final-row bound and makes a dropped tail element observable."
       },
       "inputs": {
         "x": {
       "provenance": {
         "source": "onnxruntime/core/providers/webgpu/tensor/transpose.cc",
         "test": "unit-dim permutations lower to a flat copy",
+        "notes": "Permutation [0,2,1,3] swaps non-unit extents 3 and 4, so it must perform a true transpose rather than a flat copy."
       },
       "attrs": { "perm": [0, 2, 1, 3] },
       "inputs": {
         }
       },
       "outputs": { "y": { "dtype": "float32", "shape": [2, 4, 3, 4], "tolerance": 0 } }
+    },
+    {
+      "name": "rank8_explicit_perm_regroups_two_middle_axes",
+      "provenance": {
+        "source": "ONNX Transpose-13 specification",
+        "notes": "A rank-8 staging permutation that moves one middle axis past two others. A distinct value per position makes any axis-stride mistake visible at the position it corrupts."
+      },
+      "attrs": { "perm": [0, 1, 2, 3, 5, 6, 4, 7] },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3, 2, 2, 4, 2, 2, 4],
+          "data": { "kind": "linspace", "start": 0.0, "end": 1535.0 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 3, 2, 2, 2, 2, 4, 4], "tolerance": 0 } }
+    },
+    {
+      "name": "rank8_default_perm_reverses_every_axis",
+      "provenance": {
+        "source": "ONNX Transpose-13 specification",
+        "notes": "With `perm` omitted on a rank-8 tensor, the operator reverses every axis."
+      },
+      "attrs": {},
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3, 2, 2, 4, 2, 2, 4],
+          "data": { "kind": "linspace", "start": 0.0, "end": 1535.0 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [4, 2, 2, 4, 2, 2, 3, 2], "tolerance": 0 } }
     }
   ]
 }

build/webgpu/transpose.wgsl.jinja CHANGED Viewed

@@ -1,39 +1,37 @@
 {% macro flat_tail_open() %}
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
-fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
   // 2D-folded flat index: gid.y carries the high bits when the element count exceeds the
-  // maxComputeWorkgroupsPerDimension limit (the dispatch caps x and spills the rest into y).
-  let invocation = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
   // Tail-safe scalar x4 keeps vector-like dispatch density without requiring
   // the logical tensor length (or its storage binding) to be vec4 aligned.
-  let begin = invocation * {{ source.itemsPerInvocation }}u;
-  let end = min(begin + {{ source.itemsPerInvocation }}u, params.count);
   for (var i = begin; i < end; i = i + 1u) {
 {%- endmacro %}
 {% macro flat_tail_close() %}
   }
 {% endmacro %}
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
-fn input_offset({% if source.rank > 0 %}out_index: u32{% endif %}) -> u32 {
-{% if source.rank > 0 %}
   var rem = out_index;
 {% endif %}
   var offset = 0u;
-{% for out_axis in range(source.rank) %}
 {% set out_stride = namespace(value=1) %}
-{% for j in range(out_axis + 1, source.rank) %}
-{% set out_stride.value = out_stride.value * source.outputShape[j] %}
 {% endfor %}
 {% set out_stride_divisor = out_stride.value if out_stride.value != 0 else 1 %}
-{% set in_axis = source.perm[out_axis] %}
 {% set in_stride = namespace(value=1) %}
-{% for j in range(in_axis + 1, source.rank) %}
-{% set in_stride.value = in_stride.value * source.inputShape[j] %}
 {% endfor %}
   let coord{{ out_axis }} = rem / {{ out_stride_divisor }}u;
   rem = rem % {{ out_stride_divisor }}u;
@@ -43,6 +41,6 @@ fn input_offset({% if source.rank > 0 %}out_index: u32{% endif %}) -> u32 {
 }
 {{ flat_tail_open() }}
-  y[i] = x[input_offset({% if source.rank > 0 %}i{% endif %})];
 {{ flat_tail_close() -}}
 }

 {% macro flat_tail_open() %}
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   // 2D-folded flat index: gid.y carries the high bits when the element count exceeds the
+  // dispatch's per-axis workgroup fold width (the dispatch caps x and spills the rest into y).
+  let invocation = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
   // Tail-safe scalar x4 keeps vector-like dispatch density without requiring
   // the logical tensor length (or its storage binding) to be vec4 aligned.
+{% set itemsPerInvocation = itemsPerInvocation if itemsPerInvocation is defined else 4 %}
+  let begin = invocation * {{ itemsPerInvocation }}u;
+  let end = min(begin + {{ itemsPerInvocation }}u, params.count);
   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 {
+{% if rank > 0 %}
   var rem = out_index;
 {% endif %}
   var offset = 0u;
+{% for out_axis in range(rank) %}
 {% set out_stride = namespace(value=1) %}
+{% for j in range(out_axis + 1, rank) %}
+{% set out_stride.value = out_stride.value * outputShape[j] %}
 {% endfor %}
 {% set out_stride_divisor = out_stride.value if out_stride.value != 0 else 1 %}
+{% set in_axis = perm[out_axis] %}
 {% set in_stride = namespace(value=1) %}
+{% for j in range(in_axis + 1, rank) %}
+{% set in_stride.value = in_stride.value * inputShape[j] %}
 {% endfor %}
   let coord{{ out_axis }} = rem / {{ out_stride_divisor }}u;
   rem = rem % {{ out_stride_divisor }}u;
 }
 {{ flat_tail_open() }}
+  y[i] = x[input_offset({% if rank > 0 %}i{% endif %})];
 {{ flat_tail_close() -}}
 }