Xenova HF Staff commited on 18 days ago

Commit

569bc8d

verified ·

1 Parent(s): 630ffd4

sync 2e7068faf55e

Browse files

Files changed (18) hide show

README.md +89 -0
build/webgpu/bench.json +233 -0
build/webgpu/manifest.json +1714 -0
build/webgpu/metadata.json +30 -0
build/webgpu/test.json +1530 -0
build/webgpu/unique-axis-compact-sort.wgsl.jinja +222 -0
build/webgpu/unique-axis-dedup.wgsl.jinja +105 -0
build/webgpu/unique-axis-hash.wgsl.jinja +99 -0
build/webgpu/unique-axis-scatter.wgsl.jinja +46 -0
build/webgpu/unique-axis.wgsl.jinja +559 -0
build/webgpu/unique-compact-sort.wgsl.jinja +177 -0
build/webgpu/unique-dedup.wgsl.jinja +54 -0
build/webgpu/unique-hash-build.wgsl.jinja +61 -0
build/webgpu/unique-hash-collect.wgsl.jinja +104 -0
build/webgpu/unique-hash-init.wgsl.jinja +34 -0
build/webgpu/unique-hash-mark.wgsl.jinja +53 -0
build/webgpu/unique-hash-sort-collected-key-only.wgsl.jinja +83 -0
build/webgpu/unique.wgsl.jinja +194 -0

README.md CHANGED Viewed

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

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# ai.onnx.Unique
+`ai.onnx`  ·  standard ONNX operator  ·  ONNX opset ≥ 11
+## Description
+Finds unique values or subtensors along an optional `axis`. Without an axis, `X` is flattened; results are sorted or retain first-occurrence order. Sub-32-bit integers and booleans use lossless widened 32-bit storage. Metadata outputs remain logical int64 but use lossless uint32 storage because all values are bounded by an addressable tensor extent. Callers supply exact data-dependent output shapes. ONNX-permitted uint16, 64-bit, string, and complex inputs remain unsupported because the runtime lacks matching WebGPU storage.
+See the [ONNX `Unique` spec](https://onnx.ai/onnx/operators/onnx__Unique.html) for the reference semantics.
+## Inputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `X` | `x` | `T` | — | — | The N-D input tensor from which unique values or subtensors are extracted. When `axis` is omitted, tensors of any rank are flattened in row-major order. | required |
+## Outputs
+| Name | Bind key | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- | --- |
+| `Y` | `y` | `T` | runtime-selected; narrow integers and bool use 32-bit slots | derived | — | Tensor containing all unique values or subtensors of X, sorted or in first-occurrence order. | required |
+| `indices` | `indices` | `I` | `uint32` | `1` | — | Optional logical int64 indices of each `Y` value or slice's first occurrence in `X`; stored as bounded uint32 values by WebGPU. | optional |
+| `inverse_indices` | `inverse_indices` | `I` | `uint32` | `1` | — | Optional logical int64 mapping from each flattened input value, or each input-axis slice, to its corresponding index in `Y`; stored as bounded uint32 values by WebGPU. | optional |
+| `counts` | `counts` | `I` | `uint32` | `1` | — | Optional logical int64 occurrence count for each unique value or slice in `Y`; stored as bounded uint32 values by WebGPU. | optional |
+## Attributes
+Attributes and default values (overridable per request):
+| Attribute | Default | Description |
+| --- | --- | --- |
+| `sorted` | `1` | Whether to sort unique elements in ascending order before output; 1 (default) sorts, 0 retains first-occurrence order. |
+| `axis` | — | Optional axis along which unique subtensors are identified. Negative values count from the back; when omitted, the input is flattened. |
+## Type constraints
+| Variable | Allowed dtypes |
+| --- | --- |
+| `T` | `float32`, `float16`, `uint32`, `int32`, `int16`, `uint8`, `int8`, `bool` |
+| `I` | `int64` |
+## 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
+- [`unique-axis-compact-sort.wgsl.jinja`](build/webgpu/unique-axis-compact-sort.wgsl.jinja)
+- [`unique-axis-dedup.wgsl.jinja`](build/webgpu/unique-axis-dedup.wgsl.jinja)
+- [`unique-axis-hash.wgsl.jinja`](build/webgpu/unique-axis-hash.wgsl.jinja)
+- [`unique-axis-scatter.wgsl.jinja`](build/webgpu/unique-axis-scatter.wgsl.jinja)
+- [`unique-axis.wgsl.jinja`](build/webgpu/unique-axis.wgsl.jinja)
+- [`unique-compact-sort.wgsl.jinja`](build/webgpu/unique-compact-sort.wgsl.jinja)
+- [`unique-dedup.wgsl.jinja`](build/webgpu/unique-dedup.wgsl.jinja)
+- [`unique-hash-build.wgsl.jinja`](build/webgpu/unique-hash-build.wgsl.jinja)
+- [`unique-hash-collect.wgsl.jinja`](build/webgpu/unique-hash-collect.wgsl.jinja)
+- [`unique-hash-init.wgsl.jinja`](build/webgpu/unique-hash-init.wgsl.jinja)
+- [`unique-hash-mark.wgsl.jinja`](build/webgpu/unique-hash-mark.wgsl.jinja)
+- [`unique-hash-sort-collected-key-only.wgsl.jinja`](build/webgpu/unique-hash-sort-collected-key-only.wgsl.jinja)
+- [`unique.wgsl.jinja`](build/webgpu/unique.wgsl.jinja)
+## 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.
+```js
+import { getKernel } from "@huggingface/kernels";
+const kernel = await getKernel("webgpu-kernels/ai.onnx.Unique", { version: 1 });
+// Explicit destinations request optional results or supply metadata that cannot be inferred.
+const { y } = await kernel({ x: { data: xData, shape: [1] } }, {
+  outputs: { y: { shape: [1], dtype: "float32" } },
+});
+```

build/webgpu/bench.json ADDED Viewed

	@@ -0,0 +1,233 @@

+{
+  "op": "ai.onnx.Unique",
+  "cases": [
+    {
+      "name": "int32-64k-moderate-distinct-sorted",
+      "preset": "smoke",
+      "inputs": { "x": { "dtype": "int32", "shape": [65536], "dist": "linearMod", "mod": 512 } },
+      "outputs": { "y": { "dtype": "int32", "shape": [512] } },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "name": "hash-init-build-collect-sort traffic",
+            "value": "8 * pow2ceil(2 * numel(shapes.x)) + 12 * numel(shapes.x) + 12 * numel(shapes.y)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "int32-1024-all-equal",
+      "inputs": { "x": { "dtype": "int32", "shape": [1024], "data": { "kind": "constant", "value": 7 } } },
+      "outputs": { "y": { "dtype": "int32", "shape": [1] } }
+    },
+    {
+      "name": "int32-256k-1024-distinct-sorted",
+      "preset": "smoke",
+      "inputs": { "x": { "dtype": "int32", "shape": [262144], "dist": "linearMod", "mod": 1024 } },
+      "outputs": { "y": { "dtype": "int32", "shape": [1024] } },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "name": "hash-init-build-collect-sort traffic",
+            "value": "8 * pow2ceil(2 * numel(shapes.x)) + 12 * numel(shapes.x) + 12 * numel(shapes.y)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "int32-8k-64-distinct-sorted",
+      "preset": "smoke",
+      "inputs": {
+        "x": {
+          "dtype": "int32",
+          "shape": [8192],
+          "data": {
+            "kind": "cycle",
+            "values": [37, -5, 12, 99, -73, 0, 41, 8, -21, 64, 3, -90, 55, 17, -2, 76, 29, -48, 83, 6, -33, 92, 14, -67, 50, 22, -9, 70, 35, -58, 88, 1, -26, 95, 19, -81, 46, 27, -14, 61, 33, -44, 79, 10, -39, 53, 24, -86, 100, -100, 11, -11, 13, -13, 47, -47, 59, -59, 71, -71, 89, -89, 97, -97]
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "int32", "shape": [64] } },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "name": "hash-init-build-collect-sort traffic",
+            "value": "12 * pow2ceil(2 * numel(shapes.x)) + 32 * numel(shapes.x) + 16 * numel(shapes.y)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "float32-262144-8-distinct-serial",
+      "preset": "smoke",
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [262144],
+          "data": { "kind": "cycle", "values": [-4.0, -2.0, -1.0, 0.0, 1.0, 2.0, 3.0, 8.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [8] } },
+      "bench": {
+        "primary": true,
+        "metrics": [{ "type": "bandwidth", "value": "4 * (numel(shapes.x) + numel(shapes.y))" }]
+      }
+    },
+    {
+      "name": "int32-262144-900-distinct-hash",
+      "preset": "smoke",
+      "inputs": { "x": { "dtype": "int32", "shape": [262144], "dist": "linearMod", "mod": 900 } },
+      "outputs": { "y": { "dtype": "int32", "shape": [900] } },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "name": "hash-init-build-collect-sort traffic",
+            "value": "8 * pow2ceil(2 * numel(shapes.x)) + 12 * numel(shapes.x) + 12 * numel(shapes.y)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "int32-262144-2049-distinct-large-output-hash",
+      "preset": "smoke",
+      "inputs": { "x": { "dtype": "int32", "shape": [262144], "dist": "linearMod", "mod": 2049 } },
+      "outputs": { "y": { "dtype": "int32", "shape": [2049] } },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "name": "hash-init-build-collect-sort traffic",
+            "value": "8 * pow2ceil(2 * numel(shapes.x)) + 12 * numel(shapes.x) + 12 * numel(shapes.y)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "axis0-f32-4096x256-all-distinct",
+      "preset": "smoke",
+      "attrs": { "axis": 0, "sorted": 1 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [4096, 256], "dist": "uniform", "seed": 1531, "scale": 8, "offset": 0 }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [4096, 256] } },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "name": "axis-hash-sort-scatter traffic",
+            "value": "4 * numel(shapes.x) + 12 * numel(shapes.y) + 32 * dim(shapes.x, 0) + 8 * pow2ceil(2 * dim(shapes.x, 0)) + 8 * dim(shapes.y, 0) + 888 * pow2ceil(dim(shapes.y, 0))"
+          }
+        ]
+      }
+    },
+    {
+      "name": "axis_hash_split_scatter_unsorted_zero_fills_tail",
+      "preset": "smoke",
+      "attrs": { "axis": 0, "sorted": 0 },
+      "inputs": {
+        "x": { "dtype": "int32", "shape": [2048, 1], "data": { "kind": "cycle", "values": [9, -2, 7, 9, 42, -2] } }
+      },
+      "outputs": { "y": { "dtype": "int32", "shape": [4, 1] } },
+      "bench": {
+        "primary": true,
+        "metrics": [{ "type": "bandwidth", "value": "4 * (numel(shapes.x) + numel(shapes.y))" }]
+      }
+    },
+    {
+      "name": "axis_hash_int32_inner2_duplicate_rows_sorted",
+      "preset": "smoke",
+      "attrs": { "axis": 0, "sorted": 1 },
+      "inputs": {
+        "x": { "dtype": "int32", "shape": [2048, 2], "data": { "kind": "cycle", "values": [2, 1, 0, 3, 2, 1, -1, 4] } }
+      },
+      "outputs": { "y": { "dtype": "int32", "shape": [3, 2] } },
+      "bench": {
+        "primary": true,
+        "metrics": [{ "type": "bandwidth", "value": "4 * (numel(shapes.x) + numel(shapes.y) + dim(shapes.x, 0))" }]
+      }
+    },
+    {
+      "name": "axis0_large_y_int32_unsorted_2100_distinct",
+      "preset": "smoke",
+      "attrs": { "axis": 0, "sorted": 0 },
+      "inputs": {
+        "x": { "dtype": "int32", "shape": [2100, 1], "data": { "kind": "linspace", "start": -4000, "end": 4000 } }
+      },
+      "outputs": { "y": { "dtype": "int32", "shape": [2100, 1] } },
+      "bench": {
+        "primary": true,
+        "metrics": [{ "type": "bandwidth", "value": "4 * (numel(shapes.x) + numel(shapes.y) + dim(shapes.x, 0))" }]
+      }
+    },
+    {
+      "name": "int32-16k-4096-distinct-large-output-hash",
+      "preset": "stress",
+      "inputs": { "x": { "dtype": "int32", "shape": [16384], "dist": "linearMod", "mod": 4096 } },
+      "outputs": { "y": { "dtype": "int32", "shape": [4096] } },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "name": "hash-init-build-collect-sort traffic",
+            "value": "8 * pow2ceil(2 * numel(shapes.x)) + 12 * numel(shapes.x) + 12 * numel(shapes.y)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "axis0-f32-3000x8-all-distinct",
+      "preset": "stress",
+      "attrs": { "axis": 0, "sorted": 1 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [3000, 8], "dist": "uniform", "seed": 2029, "scale": 64, "offset": 0 }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3000, 8] } },
+      "bench": { "primary": true, "metrics": [{ "type": "bandwidth", "value": "4 * numel(shapes.x)" }] }
+    },
+    {
+      "name": "axis0-f32-9000x1-storage-order",
+      "preset": "stress",
+      "attrs": { "axis": 0, "sorted": 1 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [9000, 1], "data": { "kind": "linspace", "start": -4500.0, "end": 4499.0 } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [9000, 1] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "4 * (numel(shapes.x) + dim(shapes.y, 0))" }] }
+    },
+    {
+      "name": "int32-262144-Y4096-token-vocab-sorted",
+      "preset": "stress",
+      "inputs": { "x": { "dtype": "int32", "shape": [262144], "dist": "linearMod", "mod": 4096 } },
+      "outputs": { "y": { "dtype": "int32", "shape": [4096] } },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "name": "hash-init-build-collect-sort traffic",
+            "value": "8 * pow2ceil(2 * numel(shapes.x)) + 12 * numel(shapes.x) + 12 * numel(shapes.y)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "int32-8192-all-distinct-unsorted",
+      "preset": "stress",
+      "attrs": { "sorted": 0 },
+      "inputs": { "x": { "dtype": "int32", "shape": [8192], "dist": "linearMod", "mod": 8192 } },
+      "outputs": { "y": { "dtype": "int32", "shape": [8192] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "4 * (numel(shapes.x) + numel(shapes.y))" }] }
+    }
+  ]
+}

build/webgpu/manifest.json ADDED Viewed

	@@ -0,0 +1,1714 @@

+{
+  "domain": "ai.onnx",
+  "name": "Unique",
+  "sinceVersion": 11,
+  "description": "Finds unique values or subtensors along an optional `axis`. Without an axis, `X` is flattened; results are sorted or retain first-occurrence order. Sub-32-bit integers and booleans use lossless widened 32-bit storage. Metadata outputs remain logical int64 but use lossless uint32 storage because all values are bounded by an addressable tensor extent. Callers supply exact data-dependent output shapes. ONNX-permitted uint16, 64-bit, string, and complex inputs remain unsupported because the runtime lacks matching WebGPU storage.",
+  "inputs": [
+    {
+      "role": "X",
+      "dtype": "T",
+      "description": "The N-D input tensor from which unique values or subtensors are extracted. When `axis` is omitted, tensors of any rank are flattened in row-major order."
+    }
+  ],
+  "outputs": [
+    {
+      "role": "Y",
+      "dtype": "T",
+      "description": "Tensor containing all unique values or subtensors of X, sorted or in first-occurrence order.",
+      "rank": "ranks.X if has(attrs, \"axis\") else 1"
+    },
+    {
+      "role": "indices",
+      "dtype": "I",
+      "rank": 1,
+      "optional": true,
+      "description": "Optional logical int64 indices of each `Y` value or slice's first occurrence in `X`; stored as bounded uint32 values by WebGPU."
+    },
+    {
+      "role": "inverse_indices",
+      "dtype": "I",
+      "rank": 1,
+      "optional": true,
+      "description": "Optional logical int64 mapping from each flattened input value, or each input-axis slice, to its corresponding index in `Y`; stored as bounded uint32 values by WebGPU."
+    },
+    {
+      "role": "counts",
+      "dtype": "I",
+      "rank": 1,
+      "optional": true,
+      "description": "Optional logical int64 occurrence count for each unique value or slice in `Y`; stored as bounded uint32 values by WebGPU."
+    }
+  ],
+  "attributes": { "sorted": 1 },
+  "attributeDescriptions": {
+    "sorted": "Whether to sort unique elements in ascending order before output; 1 (default) sorts, 0 retains first-occurrence order.",
+    "axis": "Optional axis along which unique subtensors are identified. Negative values count from the back; when omitted, the input is flattened."
+  },
+  "attributeConstraints": { "sorted": { "values": [0, 1] } },
+  "typeConstraints": {
+    "T": ["float32", "float16", "uint32", "int32", "int16", "uint8", "int8", "bool"],
+    "I": ["int64"]
+  },
+  "args": {
+    "x": { "kind": "tensor", "semantic": "X", "role": "input" },
+    "y": { "kind": "tensor", "semantic": "Y", "role": "output" },
+    "indices": { "kind": "tensor", "semantic": "indices", "role": "output", "dtype": "uint32", "required": false },
+    "inverse_indices": {
+      "kind": "tensor",
+      "semantic": "inverse_indices",
+      "role": "output",
+      "dtype": "uint32",
+      "required": false
+    },
+    "counts": { "kind": "tensor", "semantic": "counts", "role": "output", "dtype": "uint32", "required": false }
+  },
+  "tunables": {
+    "WORKGROUP_SIZE": 256,
+    "AXIS_SERIAL_WG": 512,
+    "HASH_MIN_INPUT": 32768,
+    "HASH_MIN_I32_INPUT": 16384,
+    "LOCAL_SORT_CROSSOVER": 2048,
+    "HASH_SORT_MAX_CAPACITY": 4096,
+    "MAX_PARALLEL_INPUT": 1048576,
+    "GLOBAL_MAX_CAPACITY": 65536
+  },
+  "derive": {
+    "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
+    "storageBufferLimit": "min(device.limits.maxStorageBufferBindingSize, device.limits.maxBufferSize)",
+    "axisPresent": "has(attrs, \"axis\")",
+    "normalizedAxis": "attrs.axis + ranks.X if axisPresent and attrs.axis < 0 else (attrs.axis if axisPresent else 0)",
+    "metadataAbsent": "not present.indices and not present.inverse_indices and not present.counts",
+    "metadataShapeContract": "(not present.indices or (ranks.indices == 1 and dim(shapes.indices, 0) == (dim(shapes.Y, normalizedAxis) if axisPresent else numel(shapes.Y)) and tensorDtypes.indices == \"uint32\")) and (not present.inverse_indices or (ranks.inverse_indices == 1 and dim(shapes.inverse_indices, 0) == (dim(shapes.X, normalizedAxis) if axisPresent else numel(shapes.X)) and tensorDtypes.inverse_indices == \"uint32\")) and (not present.counts or (ranks.counts == 1 and dim(shapes.counts, 0) == (dim(shapes.Y, normalizedAxis) if axisPresent else numel(shapes.Y)) and tensorDtypes.counts == \"uint32\"))",
+    "flatLayout": "not axisPresent and f16Ok(dtypes.T) and ranks.Y == 1 and numel(shapes.Y) <= numel(shapes.X)",
+    "axisLayout": "axisPresent and f16Ok(dtypes.T) and ranks.X >= 1 and attrs.axis + ranks.X >= 0 and attrs.axis < ranks.X and ranks.Y == ranks.X and sameShape(prefix(shapes.Y, normalizedAxis), prefix(shapes.X, normalizedAxis)) and sameShape(suffix(shapes.Y, normalizedAxis + 1), suffix(shapes.X, normalizedAxis + 1)) and dim(shapes.Y, normalizedAxis) <= dim(shapes.X, normalizedAxis)",
+    "flatContract": "flatLayout and metadataAbsent",
+    "axisContract": "axisLayout and metadataAbsent",
+    "flatMetadataContract": "flatLayout and not metadataAbsent and metadataShapeContract",
+    "axisMetadataContract": "axisLayout and not metadataAbsent and metadataShapeContract",
+    "flatScanBlocks": "ceilDiv(numel(shapes.X), tunables.WORKGROUP_SIZE)",
+    "flatTableSize": "pow2ceil(2 * numel(shapes.X))",
+    "flatTableBlocks": "ceilDiv(flatTableSize, tunables.WORKGROUP_SIZE)",
+    "flatTableVec4Blocks": "ceilDiv(flatTableSize, 4 * tunables.WORKGROUP_SIZE)",
+    "axisDim": "dim(shapes.X, normalizedAxis) if axisPresent else 0",
+    "axisOutputDim": "dim(shapes.Y, normalizedAxis) if axisPresent else 0",
+    "axisScanBlocks": "ceilDiv(axisDim, tunables.WORKGROUP_SIZE)",
+    "axisTableSize": "pow2ceil(2 * axisDim)",
+    "axisTableBlocks": "ceilDiv(axisTableSize, tunables.WORKGROUP_SIZE)",
+    "flatLocalStorageBytes": "(12 * pow2ceil(numel(shapes.Y)) if attrs.sorted != 0 else 4 * numel(shapes.Y)) + 4 * tunables.WORKGROUP_SIZE + 4",
+    "axisLocalStorageBytes": "(12 * pow2ceil(axisOutputDim) if attrs.sorted != 0 else 4 * pow2ceil(axisOutputDim)) + 4 * tunables.WORKGROUP_SIZE + 4",
+    "flatLocalSortFits": "flatLocalStorageBytes <= device.limits.maxComputeWorkgroupStorageSize",
+    "axisLocalSortFits": "axisLocalStorageBytes <= device.limits.maxComputeWorkgroupStorageSize",
+    "hashKeySortFits": "4 * pow2ceil(numel(shapes.Y)) <= device.limits.maxComputeWorkgroupStorageSize",
+    "flatScanScratchFits": "4 * numel(shapes.X) <= storageBufferLimit",
+    "flatHashScratchFits": "4 * flatTableSize <= storageBufferLimit",
+    "flatOutputScratchFits": "4 * numel(shapes.Y) <= storageBufferLimit",
+    "flatGlobalSortScratchFits": "4 * pow2ceil(min(numel(shapes.X), numel(shapes.Y))) <= storageBufferLimit",
+    "flatScanDispatchFits": "flatScanBlocks <= device.limits.maxComputeWorkgroupsPerDimension",
+    "flatHashDispatchFits": "flatScanDispatchFits and flatTableBlocks <= device.limits.maxComputeWorkgroupsPerDimension",
+    "axisFlagScratchFits": "4 * axisDim <= storageBufferLimit",
+    "axisHashScratchFits": "4 * axisTableSize <= storageBufferLimit",
+    "axisSortScratchFits": "4 * pow2ceil(axisOutputDim) <= storageBufferLimit",
+    "axisSerialWg": "max(1, min(tunables.AXIS_SERIAL_WG, deviceWorkgroupCap))",
+    "axisOrderScratchFits": "4 * max(1, axisOutputDim) <= storageBufferLimit",
+    "axisHeadCacheBytes": "4 * max(1, axisOutputDim) + 64",
+    "axisHeadCacheSlots": "max(1, axisOutputDim) if axisHeadCacheBytes <= device.limits.maxComputeWorkgroupStorageSize else 0",
+    "axisScanDispatchFits": "axisScanBlocks <= device.limits.maxComputeWorkgroupsPerDimension",
+    "axisHashDispatchFits": "axisScanDispatchFits and axisTableBlocks <= device.limits.maxComputeWorkgroupsPerDimension",
+    "narrowPortableExecution": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize < device.adapterInfo.subgroupMaxSize and device.adapterInfo.subgroupMaxSize <= 16",
+    "integerInput": "dtypes.T != \"f16\" and dtypes.T != \"f32\"",
+    "axisHashAvailable": "integerInput and axisOutputDim > 0 and axisOutputDim <= tunables.GLOBAL_MAX_CAPACITY and axisDim >= tunables.LOCAL_SORT_CROSSOVER and axisDim <= tunables.GLOBAL_MAX_CAPACITY and axisFlagScratchFits and axisHashScratchFits and axisSortScratchFits and axisHashDispatchFits",
+    "axisLocalAvailable": "integerInput and axisOutputDim > 0 and axisOutputDim <= tunables.LOCAL_SORT_CROSSOVER and axisLocalSortFits and axisFlagScratchFits and axisScanDispatchFits",
+    "axisGlobalAvailable": "integerInput and axisOutputDim > tunables.LOCAL_SORT_CROSSOVER and axisOutputDim <= tunables.GLOBAL_MAX_CAPACITY and axisFlagScratchFits and axisSortScratchFits and axisScanDispatchFits",
+    "axisParallelAvailable": "axisHashAvailable or axisLocalAvailable or axisGlobalAvailable",
+    "axisSerialSmallShape": "axisDim <= tunables.WORKGROUP_SIZE",
+    "flatHashSortedAvailable": "attrs.sorted != 0 and integerInput and (numel(shapes.X) >= tunables.HASH_MIN_INPUT or (dtypes.T == \"i32\" and numel(shapes.X) >= tunables.HASH_MIN_I32_INPUT and numel(shapes.Y) > tunables.LOCAL_SORT_CROSSOVER)) and numel(shapes.X) <= tunables.MAX_PARALLEL_INPUT and numel(shapes.Y) > 0 and numel(shapes.Y) <= tunables.HASH_SORT_MAX_CAPACITY and hashKeySortFits and flatHashScratchFits and flatOutputScratchFits and flatHashDispatchFits",
+    "flatHashAvailable": "integerInput and numel(shapes.X) >= tunables.HASH_MIN_INPUT and numel(shapes.X) <= tunables.MAX_PARALLEL_INPUT and numel(shapes.Y) > 0 and numel(shapes.Y) <= tunables.LOCAL_SORT_CROSSOVER and flatLocalSortFits and flatScanScratchFits and flatHashScratchFits and flatHashDispatchFits",
+    "flatLocalAvailable": "integerInput and numel(shapes.X) > 0 and numel(shapes.X) < tunables.HASH_MIN_INPUT and numel(shapes.Y) > 0 and numel(shapes.Y) <= tunables.LOCAL_SORT_CROSSOVER and flatLocalSortFits and flatScanScratchFits and flatScanDispatchFits",
+    "flatGlobalAvailable": "integerInput and numel(shapes.X) > 0 and numel(shapes.X) <= tunables.MAX_PARALLEL_INPUT and numel(shapes.Y) > tunables.LOCAL_SORT_CROSSOVER and min(numel(shapes.X), numel(shapes.Y)) <= tunables.GLOBAL_MAX_CAPACITY and flatScanScratchFits and flatGlobalSortScratchFits and flatScanDispatchFits",
+    "flatParallelAvailable": "flatHashSortedAvailable or flatHashAvailable or flatLocalAvailable or flatGlobalAvailable",
+    "flatSerialFallbackNeeded": "not flatParallelAvailable"
+  },
+  "constants": {
+    "scalar": "dtypes.T",
+    "usesF16": "dtypes.T == \"f16\"",
+    "isFloat": "dtypes.T == \"f16\" or dtypes.T == \"f32\"",
+    "isUnsigned": "dtypes.T == \"u32\""
+  },
+  "bindingSets": {
+    "flatMetadataIndices": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "indices",
+        "arg": "indices",
+        "semantic": "indices",
+        "buffer": { "type": "storage" },
+        "elementType": "u32"
+      },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "inputCount", "type": "u32", "value": "numel(shapes.X)" },
+            { "name": "capacity", "type": "u32", "value": "numel(shapes.Y)" }
+          ]
+        }
+      }
+    ],
+    "flatMetadataInverse": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "inverse_indices",
+        "arg": "inverse_indices",
+        "semantic": "inverse_indices",
+        "buffer": { "type": "storage" },
+        "elementType": "u32"
+      },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "inputCount", "type": "u32", "value": "numel(shapes.X)" },
+            { "name": "capacity", "type": "u32", "value": "numel(shapes.Y)" }
+          ]
+        }
+      }
+    ],
+    "flatMetadataCounts": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      { "name": "counts", "arg": "counts", "semantic": "counts", "buffer": { "type": "storage" }, "elementType": "u32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "inputCount", "type": "u32", "value": "numel(shapes.X)" },
+            { "name": "capacity", "type": "u32", "value": "numel(shapes.Y)" }
+          ]
+        }
+      }
+    ],
+    "flatMetadataIndicesInverse": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "indices",
+        "arg": "indices",
+        "semantic": "indices",
+        "buffer": { "type": "storage" },
+        "elementType": "u32"
+      },
+      {
+        "name": "inverse_indices",
+        "arg": "inverse_indices",
+        "semantic": "inverse_indices",
+        "buffer": { "type": "storage" },
+        "elementType": "u32"
+      },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "inputCount", "type": "u32", "value": "numel(shapes.X)" },
+            { "name": "capacity", "type": "u32", "value": "numel(shapes.Y)" }
+          ]
+        }
+      }
+    ],
+    "flatMetadataIndicesCounts": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "indices",
+        "arg": "indices",
+        "semantic": "indices",
+        "buffer": { "type": "storage" },
+        "elementType": "u32"
+      },
+      { "name": "counts", "arg": "counts", "semantic": "counts", "buffer": { "type": "storage" }, "elementType": "u32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "inputCount", "type": "u32", "value": "numel(shapes.X)" },
+            { "name": "capacity", "type": "u32", "value": "numel(shapes.Y)" }
+          ]
+        }
+      }
+    ],
+    "flatMetadataInverseCounts": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "inverse_indices",
+        "arg": "inverse_indices",
+        "semantic": "inverse_indices",
+        "buffer": { "type": "storage" },
+        "elementType": "u32"
+      },
+      { "name": "counts", "arg": "counts", "semantic": "counts", "buffer": { "type": "storage" }, "elementType": "u32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "inputCount", "type": "u32", "value": "numel(shapes.X)" },
+            { "name": "capacity", "type": "u32", "value": "numel(shapes.Y)" }
+          ]
+        }
+      }
+    ],
+    "flatMetadataAll": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "indices",
+        "arg": "indices",
+        "semantic": "indices",
+        "buffer": { "type": "storage" },
+        "elementType": "u32"
+      },
+      {
+        "name": "inverse_indices",
+        "arg": "inverse_indices",
+        "semantic": "inverse_indices",
+        "buffer": { "type": "storage" },
+        "elementType": "u32"
+      },
+      { "name": "counts", "arg": "counts", "semantic": "counts", "buffer": { "type": "storage" }, "elementType": "u32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "inputCount", "type": "u32", "value": "numel(shapes.X)" },
+            { "name": "capacity", "type": "u32", "value": "numel(shapes.Y)" }
+          ]
+        }
+      }
+    ],
+    "axisMetadataIndices": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "indices",
+        "arg": "indices",
+        "semantic": "indices",
+        "buffer": { "type": "storage" },
+        "elementType": "u32"
+      },
+      { "name": "order", "semantic": "order", "buffer": { "type": "storage" }, "elementType": "u32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "outer", "type": "u32", "value": "outer(shapes.X, normalizedAxis)" },
+            { "name": "axisDim", "type": "u32", "value": "dim(shapes.X, normalizedAxis)" },
+            { "name": "inner", "type": "u32", "value": "inner(shapes.X, normalizedAxis)" },
+            { "name": "outputAxisDim", "type": "u32", "value": "dim(shapes.Y, normalizedAxis)" }
+          ]
+        }
+      }
+    ],
+    "axisMetadataInverse": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "inverse_indices",
+        "arg": "inverse_indices",
+        "semantic": "inverse_indices",
+        "buffer": { "type": "storage" },
+        "elementType": "u32"
+      },
+      { "name": "order", "semantic": "order", "buffer": { "type": "storage" }, "elementType": "u32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "outer", "type": "u32", "value": "outer(shapes.X, normalizedAxis)" },
+            { "name": "axisDim", "type": "u32", "value": "dim(shapes.X, normalizedAxis)" },
+            { "name": "inner", "type": "u32", "value": "inner(shapes.X, normalizedAxis)" },
+            { "name": "outputAxisDim", "type": "u32", "value": "dim(shapes.Y, normalizedAxis)" }
+          ]
+        }
+      }
+    ],
+    "axisMetadataCounts": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      { "name": "counts", "arg": "counts", "semantic": "counts", "buffer": { "type": "storage" }, "elementType": "u32" },
+      { "name": "order", "semantic": "order", "buffer": { "type": "storage" }, "elementType": "u32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "outer", "type": "u32", "value": "outer(shapes.X, normalizedAxis)" },
+            { "name": "axisDim", "type": "u32", "value": "dim(shapes.X, normalizedAxis)" },
+            { "name": "inner", "type": "u32", "value": "inner(shapes.X, normalizedAxis)" },
+            { "name": "outputAxisDim", "type": "u32", "value": "dim(shapes.Y, normalizedAxis)" }
+          ]
+        }
+      }
+    ],
+    "axisMetadataIndicesInverse": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "indices",
+        "arg": "indices",
+        "semantic": "indices",
+        "buffer": { "type": "storage" },
+        "elementType": "u32"
+      },
+      {
+        "name": "inverse_indices",
+        "arg": "inverse_indices",
+        "semantic": "inverse_indices",
+        "buffer": { "type": "storage" },
+        "elementType": "u32"
+      },
+      { "name": "order", "semantic": "order", "buffer": { "type": "storage" }, "elementType": "u32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "outer", "type": "u32", "value": "outer(shapes.X, normalizedAxis)" },
+            { "name": "axisDim", "type": "u32", "value": "dim(shapes.X, normalizedAxis)" },
+            { "name": "inner", "type": "u32", "value": "inner(shapes.X, normalizedAxis)" },
+            { "name": "outputAxisDim", "type": "u32", "value": "dim(shapes.Y, normalizedAxis)" }
+          ]
+        }
+      }
+    ],
+    "axisMetadataIndicesCounts": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "indices",
+        "arg": "indices",
+        "semantic": "indices",
+        "buffer": { "type": "storage" },
+        "elementType": "u32"
+      },
+      { "name": "counts", "arg": "counts", "semantic": "counts", "buffer": { "type": "storage" }, "elementType": "u32" },
+      { "name": "order", "semantic": "order", "buffer": { "type": "storage" }, "elementType": "u32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "outer", "type": "u32", "value": "outer(shapes.X, normalizedAxis)" },
+            { "name": "axisDim", "type": "u32", "value": "dim(shapes.X, normalizedAxis)" },
+            { "name": "inner", "type": "u32", "value": "inner(shapes.X, normalizedAxis)" },
+            { "name": "outputAxisDim", "type": "u32", "value": "dim(shapes.Y, normalizedAxis)" }
+          ]
+        }
+      }
+    ],
+    "axisMetadataInverseCounts": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "inverse_indices",
+        "arg": "inverse_indices",
+        "semantic": "inverse_indices",
+        "buffer": { "type": "storage" },
+        "elementType": "u32"
+      },
+      { "name": "counts", "arg": "counts", "semantic": "counts", "buffer": { "type": "storage" }, "elementType": "u32" },
+      { "name": "order", "semantic": "order", "buffer": { "type": "storage" }, "elementType": "u32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "outer", "type": "u32", "value": "outer(shapes.X, normalizedAxis)" },
+            { "name": "axisDim", "type": "u32", "value": "dim(shapes.X, normalizedAxis)" },
+            { "name": "inner", "type": "u32", "value": "inner(shapes.X, normalizedAxis)" },
+            { "name": "outputAxisDim", "type": "u32", "value": "dim(shapes.Y, normalizedAxis)" }
+          ]
+        }
+      }
+    ],
+    "axisMetadataAll": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "indices",
+        "arg": "indices",
+        "semantic": "indices",
+        "buffer": { "type": "storage" },
+        "elementType": "u32"
+      },
+      {
+        "name": "inverse_indices",
+        "arg": "inverse_indices",
+        "semantic": "inverse_indices",
+        "buffer": { "type": "storage" },
+        "elementType": "u32"
+      },
+      { "name": "counts", "arg": "counts", "semantic": "counts", "buffer": { "type": "storage" }, "elementType": "u32" },
+      { "name": "order", "semantic": "order", "buffer": { "type": "storage" }, "elementType": "u32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "outer", "type": "u32", "value": "outer(shapes.X, normalizedAxis)" },
+            { "name": "axisDim", "type": "u32", "value": "dim(shapes.X, normalizedAxis)" },
+            { "name": "inner", "type": "u32", "value": "inner(shapes.X, normalizedAxis)" },
+            { "name": "outputAxisDim", "type": "u32", "value": "dim(shapes.Y, normalizedAxis)" }
+          ]
+        }
+      }
+    ],
+    "hashBuild": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "tableKey", "semantic": "tableKey", "buffer": { "type": "storage" }, "elementType": "atomic<u32>" },
+      { "name": "tableIdx", "semantic": "tableIdx", "buffer": { "type": "storage" }, "elementType": "atomic<u32>" },
+      {
+        "name": "special",
+        "semantic": "special",
+        "buffer": { "type": "storage" },
+        "elementType": "atomic<u32>",
+        "length": 1
+      },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": { "name": "Params", "fields": [{ "name": "inputCount", "type": "u32", "value": "numel(shapes.X)" }] }
+      }
+    ],
+    "hashBuildKeyOnly": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "tableKey", "semantic": "tableKey", "buffer": { "type": "storage" }, "elementType": "atomic<u32>" },
+      {
+        "name": "special",
+        "semantic": "special",
+        "buffer": { "type": "storage" },
+        "elementType": "atomic<u32>",
+        "length": 1
+      },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": { "name": "Params", "fields": [{ "name": "inputCount", "type": "u32", "value": "numel(shapes.X)" }] }
+      }
+    ],
+    "dedup": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "flags", "semantic": "flags", "buffer": { "type": "storage" }, "elementType": "u32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": { "name": "Params", "fields": [{ "name": "inputCount", "type": "u32", "value": "numel(shapes.X)" }] }
+      }
+    ],
+    "compact": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "flags", "semantic": "flags", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": { "name": "Params", "fields": [{ "name": "inputCount", "type": "u32", "value": "numel(shapes.X)" }] }
+      }
+    ],
+    "axisDedup": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "firstFlag", "semantic": "firstFlag", "buffer": { "type": "storage" }, "elementType": "u32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "outer", "type": "u32", "value": "outer(shapes.X, normalizedAxis)" },
+            { "name": "axisDim", "type": "u32", "value": "dim(shapes.X, normalizedAxis)" },
+            { "name": "inner", "type": "u32", "value": "inner(shapes.X, normalizedAxis)" }
+          ]
+        }
+      }
+    ],
+    "axisCompactGlobal": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "firstFlag", "semantic": "firstFlag", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
+      { "name": "slots", "semantic": "slots", "buffer": { "type": "storage" }, "elementType": "u32" },
+      { "name": "sortPad", "semantic": "sortPad", "buffer": { "type": "storage" }, "elementType": "u32" },
+      { "name": "sortKey", "semantic": "sortKey", "buffer": { "type": "storage" }, "elementType": "u32" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "outer", "type": "u32", "value": "outer(shapes.X, normalizedAxis)" },
+            { "name": "axisDim", "type": "u32", "value": "dim(shapes.X, normalizedAxis)" },
+            { "name": "inner", "type": "u32", "value": "inner(shapes.X, normalizedAxis)" },
+            { "name": "outputAxisDim", "type": "u32", "value": "dim(shapes.Y, normalizedAxis)" }
+          ]
+        }
+      }
+    ],
+    "axisCompactGlobalUnsorted": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "firstFlag", "semantic": "firstFlag", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
+      { "name": "slots", "semantic": "slots", "buffer": { "type": "storage" }, "elementType": "u32" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "outer", "type": "u32", "value": "outer(shapes.X, normalizedAxis)" },
+            { "name": "axisDim", "type": "u32", "value": "dim(shapes.X, normalizedAxis)" },
+            { "name": "inner", "type": "u32", "value": "inner(shapes.X, normalizedAxis)" },
+            { "name": "outputAxisDim", "type": "u32", "value": "dim(shapes.Y, normalizedAxis)" }
+          ]
+        }
+      }
+    ],
+    "axisCompactGlobalScratch": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "firstFlag", "semantic": "firstFlag", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
+      { "name": "slots", "semantic": "slots", "buffer": { "type": "storage" }, "elementType": "u32" },
+      { "name": "sortPad", "semantic": "sortPad", "buffer": { "type": "storage" }, "elementType": "u32" },
+      { "name": "sortKey", "semantic": "sortKey", "buffer": { "type": "storage" }, "elementType": "u32" },
+      { "name": "count", "semantic": "count", "buffer": { "type": "storage" }, "elementType": "u32", "length": 1 },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "outer", "type": "u32", "value": "outer(shapes.X, normalizedAxis)" },
+            { "name": "axisDim", "type": "u32", "value": "dim(shapes.X, normalizedAxis)" },
+            { "name": "inner", "type": "u32", "value": "inner(shapes.X, normalizedAxis)" }
+          ]
+        }
+      }
+    ],
+    "axisCompactGlobalScratchUnsortedBase": [
+      { "name": "firstFlag", "semantic": "firstFlag", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
+      { "name": "slots", "semantic": "slots", "buffer": { "type": "storage" }, "elementType": "u32" },
+      { "name": "count", "semantic": "count", "buffer": { "type": "storage" }, "elementType": "u32", "length": 1 }
+    ],
+    "axisCompactGlobalScratchUnsorted": [
+      { "name": "firstFlag", "semantic": "firstFlag", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
+      { "name": "slots", "semantic": "slots", "buffer": { "type": "storage" }, "elementType": "u32" },
+      { "name": "count", "semantic": "count", "buffer": { "type": "storage" }, "elementType": "u32", "length": 1 },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [{ "name": "axisDim", "type": "u32", "value": "dim(shapes.X, normalizedAxis)" }]
+        }
+      }
+    ],
+    "axisScatterGlobal": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "slots", "semantic": "slots", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
+      {
+        "name": "count",
+        "semantic": "count",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "u32",
+        "length": 1
+      },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" }
+    ],
+    "hashInitKeyOnlyWithCount": [
+      { "name": "tableKey", "semantic": "tableKey", "buffer": { "type": "storage" }, "elementType": "vec4<u32>" },
+      {
+        "name": "special",
+        "semantic": "special",
+        "buffer": { "type": "storage" },
+        "elementType": "atomic<u32>",
+        "length": 1
+      },
+      {
+        "name": "count",
+        "semantic": "count",
+        "buffer": { "type": "storage" },
+        "elementType": "atomic<u32>",
+        "length": 1
+      }
+    ],
+    "hashCollectKeyOnly": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "tableKey", "semantic": "tableKey", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
+      {
+        "name": "special",
+        "semantic": "special",
+        "buffer": { "type": "storage" },
+        "elementType": "atomic<u32>",
+        "length": 1
+      },
+      { "name": "compactedBits", "semantic": "compactedBits", "buffer": { "type": "storage" }, "elementType": "u32" },
+      {
+        "name": "count",
+        "semantic": "count",
+        "buffer": { "type": "storage" },
+        "elementType": "atomic<u32>",
+        "length": 1
+      }
+    ],
+    "collectedSort": [
+      {
+        "name": "compactedBits",
+        "semantic": "compactedBits",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "u32"
+      },
+      {
+        "name": "count",
+        "semantic": "count",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "u32",
+        "length": 1
+      },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" }
+    ],
+    "hashInit": [
+      { "name": "tableKey", "semantic": "tableKey", "buffer": { "type": "storage" }, "elementType": "atomic<u32>" },
+      { "name": "tableIdx", "semantic": "tableIdx", "buffer": { "type": "storage" }, "elementType": "atomic<u32>" },
+      {
+        "name": "special",
+        "semantic": "special",
+        "buffer": { "type": "storage" },
+        "elementType": "atomic<u32>",
+        "length": 1
+      }
+    ],
+    "hashMark": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "flags", "semantic": "flags", "buffer": { "type": "storage" }, "elementType": "u32" },
+      { "name": "tableKey", "semantic": "tableKey", "buffer": { "type": "storage" }, "elementType": "atomic<u32>" },
+      { "name": "tableIdx", "semantic": "tableIdx", "buffer": { "type": "storage" }, "elementType": "atomic<u32>" },
+      {
+        "name": "special",
+        "semantic": "special",
+        "buffer": { "type": "storage" },
+        "elementType": "atomic<u32>",
+        "length": 1
+      },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": { "name": "Params", "fields": [{ "name": "inputCount", "type": "u32", "value": "numel(shapes.X)" }] }
+      }
+    ],
+    "flatGlobalCompact": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "flags", "semantic": "flags", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
+      { "name": "sortVal", "semantic": "sortVal", "buffer": { "type": "storage" }, "elementType": "u32" },
+      { "name": "sortKey", "semantic": "sortKey", "buffer": { "type": "storage" }, "elementType": "u32" },
+      { "name": "sortPad", "semantic": "sortPad", "buffer": { "type": "storage" }, "elementType": "u32" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": { "name": "Params", "fields": [{ "name": "inputCount", "type": "u32", "value": "numel(shapes.X)" }] }
+      }
+    ],
+    "flatGlobalCompactUnsorted": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "flags", "semantic": "flags", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
+      { "name": "sortVal", "semantic": "sortVal", "buffer": { "type": "storage" }, "elementType": "u32" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": { "name": "Params", "fields": [{ "name": "inputCount", "type": "u32", "value": "numel(shapes.X)" }] }
+      }
+    ],
+    "flatSerial": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "inputCount", "type": "u32", "value": "numel(shapes.X)" },
+            { "name": "capacity", "type": "u32", "value": "numel(shapes.Y)" }
+          ]
+        }
+      }
+    ],
+    "axisHashInit": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "sliceHash", "semantic": "sliceHash", "buffer": { "type": "storage" }, "elementType": "u32" },
+      { "name": "hashSlot", "semantic": "hashSlot", "buffer": { "type": "storage" }, "elementType": "atomic<u32>" },
+      { "name": "firstFlag", "semantic": "firstFlag", "buffer": { "type": "storage" }, "elementType": "u32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "outer", "type": "u32", "value": "outer(shapes.X, normalizedAxis)" },
+            { "name": "axisDim", "type": "u32", "value": "dim(shapes.X, normalizedAxis)" },
+            { "name": "inner", "type": "u32", "value": "inner(shapes.X, normalizedAxis)" }
+          ]
+        }
+      }
+    ],
+    "axisHashBuild": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "sliceHash", "semantic": "sliceHash", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
+      { "name": "hashSlot", "semantic": "hashSlot", "buffer": { "type": "storage" }, "elementType": "atomic<u32>" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "outer", "type": "u32", "value": "outer(shapes.X, normalizedAxis)" },
+            { "name": "axisDim", "type": "u32", "value": "dim(shapes.X, normalizedAxis)" },
+            { "name": "inner", "type": "u32", "value": "inner(shapes.X, normalizedAxis)" }
+          ]
+        }
+      }
+    ],
+    "axisHashMark": [
+      { "name": "hashSlot", "semantic": "hashSlot", "buffer": { "type": "storage" }, "elementType": "atomic<u32>" },
+      { "name": "firstFlag", "semantic": "firstFlag", "buffer": { "type": "storage" }, "elementType": "u32" }
+    ],
+    "axisCompactLocal": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "firstFlag", "semantic": "firstFlag", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "outer", "type": "u32", "value": "outer(shapes.X, normalizedAxis)" },
+            { "name": "axisDim", "type": "u32", "value": "dim(shapes.X, normalizedAxis)" },
+            { "name": "inner", "type": "u32", "value": "inner(shapes.X, normalizedAxis)" },
+            { "name": "outputAxisDim", "type": "u32", "value": "dim(shapes.Y, normalizedAxis)" }
+          ]
+        }
+      }
+    ],
+    "axisSerial": [
+      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+      { "name": "order", "semantic": "order", "buffer": { "type": "storage" }, "elementType": "u32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "outer", "type": "u32", "value": "outer(shapes.X, normalizedAxis)" },
+            { "name": "axisDim", "type": "u32", "value": "dim(shapes.X, normalizedAxis)" },
+            { "name": "inner", "type": "u32", "value": "inner(shapes.X, normalizedAxis)" },
+            { "name": "outputAxisDim", "type": "u32", "value": "dim(shapes.Y, normalizedAxis)" }
+          ]
+        }
+      }
+    ]
+  },
+  "variants": [
+    {
+      "id": "hash_dedup_sorted_collect",
+      "priority": 30,
+      "when": ["flatContract", "attrs.sorted != 0", "integerInput", "(numel(shapes.X) >= tunables.HASH_MIN_INPUT or (dtypes.T == \"i32\" and numel(shapes.X) >= tunables.HASH_MIN_I32_INPUT and numel(shapes.Y) > tunables.LOCAL_SORT_CROSSOVER))", "numel(shapes.X) <= tunables.MAX_PARALLEL_INPUT", "numel(shapes.Y) > 0", "numel(shapes.Y) <= tunables.HASH_SORT_MAX_CAPACITY", "hashKeySortFits", "flatHashScratchFits", "flatOutputScratchFits", "flatHashDispatchFits"],
+      "derive": {
+        "scanBlocks": "flatScanBlocks",
+        "tableSize": "flatTableSize",
+        "tableBlocks": "flatTableBlocks",
+        "tableVec4Blocks": "flatTableVec4Blocks"
+      },
+      "intermediates": [
+        { "id": "tableKey", "dtype": "uint32", "shape": "[tableSize]" },
+        { "id": "special", "dtype": "uint32", "shape": "[1]" },
+        { "id": "compactedBits", "dtype": "uint32", "shape": "[numel(shapes.Y)]" },
+        { "id": "count", "dtype": "uint32", "shape": "[1]" }
+      ],
+      "passes": [
+        {
+          "id": "init",
+          "name": "Unique.HashInitCollect",
+          "source": {
+            "shader": "unique-hash-init.wgsl.jinja",
+            "inputs": { "tableSize": "tableSize", "keyOnlyVec4": true }
+          },
+          "bindings": "hashInitKeyOnlyWithCount",
+          "dispatch": { "x": "tableVec4Blocks" }
+        },
+        {
+          "id": "build",
+          "name": "Unique.HashBuild",
+          "source": {
+            "shader": "unique-hash-build.wgsl.jinja",
+            "inputs": { "tableSize": "tableSize", "keyOnly": true }
+          },
+          "bindings": "hashBuildKeyOnly",
+          "dispatch": { "x": "scanBlocks" }
+        },
+        {
+          "id": "collect",
+          "name": "Unique.HashCollect",
+          "source": {
+            "shader": "unique-hash-collect.wgsl.jinja",
+            "inputs": {
+              "tableSize": "tableSize",
+              "capacity": "numel(shapes.Y)",
+              "useSubgroups": "device.features.has(\"subgroups\")"
+            }
+          },
+          "bindings": "hashCollectKeyOnly",
+          "dispatch": { "x": "tableBlocks" }
+        },
+        {
+          "id": "sort",
+          "name": "Unique.HashSortCollected",
+          "source": {
+            "shader": "unique-hash-sort-collected-key-only.wgsl.jinja",
+            "inputs": { "capacity": "numel(shapes.Y)", "sortN": "pow2ceil(numel(shapes.Y))" }
+          },
+          "bindings": "collectedSort",
+          "dispatch": { "x": 1 }
+        }
+      ]
+    },
+    {
+      "id": "hash_dedup",
+      "priority": 20,
+      "when": ["flatContract", "integerInput", "numel(shapes.X) >= tunables.HASH_MIN_INPUT", "numel(shapes.X) <= tunables.MAX_PARALLEL_INPUT", "numel(shapes.Y) > 0", "numel(shapes.Y) <= tunables.LOCAL_SORT_CROSSOVER", "flatLocalSortFits", "flatScanScratchFits", "flatHashScratchFits", "flatHashDispatchFits"],
+      "derive": {
+        "scanN": "numel(shapes.X)",
+        "scanBlocks": "flatScanBlocks",
+        "tableSize": "flatTableSize",
+        "tableBlocks": "flatTableBlocks"
+      },
+      "constants": { "sorted": "attrs.sorted != 0" },
+      "intermediates": [
+        { "id": "flags", "dtype": "uint32", "shape": "[scanN]" },
+        { "id": "tableKey", "dtype": "uint32", "shape": "[tableSize]" },
+        { "id": "tableIdx", "dtype": "uint32", "shape": "[tableSize]" },
+        { "id": "special", "dtype": "uint32", "shape": "[1]" }
+      ],
+      "passes": [
+        {
+          "id": "init",
+          "name": "Unique.HashInit",
+          "source": {
+            "shader": "unique-hash-init.wgsl.jinja",
+            "inputs": { "tableSize": "tableSize", "keyOnlyVec4": false }
+          },
+          "bindings": "hashInit",
+          "dispatch": { "x": "tableBlocks" }
+        },
+        {
+          "id": "build",
+          "name": "Unique.HashBuild",
+          "source": {
+            "shader": "unique-hash-build.wgsl.jinja",
+            "inputs": { "tableSize": "tableSize", "keyOnly": false }
+          },
+          "bindings": "hashBuild",
+          "dispatch": { "x": "scanBlocks" }
+        },
+        {
+          "id": "mark",
+          "name": "Unique.HashMark",
+          "source": { "shader": "unique-hash-mark.wgsl.jinja", "inputs": { "tableSize": "tableSize" } },
+          "bindings": "hashMark",
+          "dispatch": { "x": "scanBlocks" }
+        },
+        {
+          "id": "compact",
+          "name": "Unique.CompactSort",
+          "source": {
+            "shader": "unique-compact-sort.wgsl.jinja",
+            "inputs": { "capacity": "numel(shapes.Y)", "sortN": "pow2ceil(numel(shapes.Y))", "globalScratch": false }
+          },
+          "bindings": "compact",
+          "dispatch": { "x": 1 }
+        }
+      ]
+    },
+    {
+      "id": "rank1_bounded",
+      "when": ["flatContract", "flatSerialFallbackNeeded"],
+      "constants": { "sorted": "attrs.sorted != 0" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "Unique",
+          "shader": "unique.wgsl.jinja",
+          "bindings": "flatSerial",
+          "dispatch": { "x": 1 }
+        }
+      ]
+    },
+    {
+      "id": "axis_serial",
+      "when": ["axisContract", "axisOrderScratchFits", "not axisParallelAvailable or axisSerialSmallShape"],
+      "constants": { "sorted": "attrs.sorted != 0", "axisSerialWg": "axisSerialWg" },
+      "intermediates": [{ "id": "order", "dtype": "uint32", "shape": "[max(1, axisOutputDim)]" }],
+      "passes": [
+        {
+          "id": "main",
+          "name": "Unique.Axis",
+          "source": {
+            "shader": "unique-axis.wgsl.jinja",
+            "inputs": { "capacity": "max(1, dim(shapes.Y, normalizedAxis))", "headCacheSlots": "axisHeadCacheSlots" }
+          },
+          "bindings": "axisSerial",
+          "dispatch": { "x": 1 }
+        }
+      ]
+    },
+    {
+      "id": "axis_hash_bounded",
+      "priority": 20,
+      "demoteWhen": ["narrowPortableExecution and (axisLocalAvailable or axisGlobalAvailable)"],
+      "when": ["axisContract", "axisHashAvailable", "attrs.sorted != 0"],
+      "derive": {
+        "axisDim": "dim(shapes.X, normalizedAxis)",
+        "tableSize": "axisTableSize",
+        "axisBlocks": "axisScanBlocks",
+        "tableBlocks": "axisTableBlocks",
+        "sortN": "pow2ceil(dim(shapes.Y, normalizedAxis))"
+      },
+      "constants": { "sorted": true },
+      "intermediates": [
+        { "id": "sliceHash", "dtype": "uint32", "shape": "[axisDim]" },
+        { "id": "hashSlot", "dtype": "uint32", "shape": "[tableSize]" },
+        { "id": "firstFlag", "dtype": "uint32", "shape": "[axisDim]" },
+        { "id": "slots", "dtype": "uint32", "shape": "[sortN]" },
+        { "id": "sortPad", "dtype": "uint32", "shape": "[sortN]" },
+        { "id": "sortKey", "dtype": "uint32", "shape": "[sortN]" },
+        { "id": "count", "dtype": "uint32", "shape": "[1]" }
+      ],
+      "passes": [
+        {
+          "id": "hash_init",
+          "name": "Unique.AxisHashInit",
+          "source": {
+            "shader": "unique-axis-hash.wgsl.jinja",
+            "inputs": { "stage": "\"init\"", "tableSize": "tableSize" }
+          },
+          "bindings": "axisHashInit",
+          "dispatch": { "x": "tableBlocks" }
+        },
+        {
+          "id": "hash_build",
+          "name": "Unique.AxisHashBuild",
+          "source": {
+            "shader": "unique-axis-hash.wgsl.jinja",
+            "inputs": { "stage": "\"build\"", "tableSize": "tableSize" }
+          },
+          "bindings": "axisHashBuild",
+          "dispatch": { "x": "axisBlocks" }
+        },
+        {
+          "id": "hash_mark",
+          "name": "Unique.AxisHashMark",
+          "source": {
+            "shader": "unique-axis-hash.wgsl.jinja",
+            "inputs": { "stage": "\"mark\"", "tableSize": "tableSize" }
+          },
+          "bindings": "axisHashMark",
+          "dispatch": { "x": "tableBlocks" }
+        },
+        {
+          "id": "compact",
+          "name": "Unique.AxisCompactSortGlobal",
+          "source": {
+            "shader": "unique-axis-compact-sort.wgsl.jinja",
+            "inputs": {
+              "capacity": "dim(shapes.Y, normalizedAxis)",
+              "sortN": "sortN",
+              "compactOnly": true,
+              "globalScratch": true,
+              "emitCount": true
+            }
+          },
+          "bindings": "axisCompactGlobalScratch",
+          "dispatch": { "x": 1 }
+        },
+        {
+          "id": "scatter",
+          "name": "Unique.AxisScatter",
+          "source": {
+            "shader": "unique-axis-scatter.wgsl.jinja",
+            "inputs": {
+              "capacity": "dim(shapes.Y, normalizedAxis)",
+              "axisDim": "dim(shapes.X, normalizedAxis)",
+              "inner": "max(1, inner(shapes.X, normalizedAxis))",
+              "outputAxisDim": "dim(shapes.Y, normalizedAxis)",
+              "totalOut": "numel(shapes.Y)",
+              "staticShape": true
+            }
+          },
+          "bindings": "axisScatterGlobal",
+          "dispatch": { "gridStride": "numel(shapes.Y)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
+        }
+      ]
+    },
+    {
+      "id": "axis_hash_bounded_unsorted",
+      "priority": 20,
+      "demoteWhen": ["narrowPortableExecution and (axisLocalAvailable or axisGlobalAvailable)"],
+      "when": ["axisContract", "axisHashAvailable", "attrs.sorted == 0"],
+      "derive": {
+        "axisDim": "dim(shapes.X, normalizedAxis)",
+        "tableSize": "axisTableSize",
+        "axisBlocks": "axisScanBlocks",
+        "tableBlocks": "axisTableBlocks",
+        "sortN": "pow2ceil(dim(shapes.Y, normalizedAxis))"
+      },
+      "constants": { "sorted": false },
+      "intermediates": [
+        { "id": "sliceHash", "dtype": "uint32", "shape": "[axisDim]" },
+        { "id": "hashSlot", "dtype": "uint32", "shape": "[tableSize]" },
+        { "id": "firstFlag", "dtype": "uint32", "shape": "[axisDim]" },
+        { "id": "slots", "dtype": "uint32", "shape": "[dim(shapes.Y, normalizedAxis)]" },
+        { "id": "count", "dtype": "uint32", "shape": "[1]" }
+      ],
+      "passes": [
+        {
+          "id": "hash_init",
+          "name": "Unique.AxisHashInit",
+          "source": {
+            "shader": "unique-axis-hash.wgsl.jinja",
+            "inputs": { "stage": "\"init\"", "tableSize": "tableSize" }
+          },
+          "bindings": "axisHashInit",
+          "dispatch": { "x": "tableBlocks" }
+        },
+        {
+          "id": "hash_build",
+          "name": "Unique.AxisHashBuild",
+          "source": {
+            "shader": "unique-axis-hash.wgsl.jinja",
+            "inputs": { "stage": "\"build\"", "tableSize": "tableSize" }
+          },
+          "bindings": "axisHashBuild",
+          "dispatch": { "x": "axisBlocks" }
+        },
+        {
+          "id": "hash_mark",
+          "name": "Unique.AxisHashMark",
+          "source": {
+            "shader": "unique-axis-hash.wgsl.jinja",
+            "inputs": { "stage": "\"mark\"", "tableSize": "tableSize" }
+          },
+          "bindings": "axisHashMark",
+          "dispatch": { "x": "tableBlocks" }
+        },
+        {
+          "id": "compact",
+          "name": "Unique.AxisCompactGlobal",
+          "source": {
+            "shader": "unique-axis-compact-sort.wgsl.jinja",
+            "inputs": {
+              "capacity": "dim(shapes.Y, normalizedAxis)",
+              "sortN": "sortN",
+              "compactOnly": true,
+              "globalScratch": true,
+              "emitCount": true
+            }
+          },
+          "bindings": "axisCompactGlobalScratchUnsorted",
+          "dispatch": { "x": 1 }
+        },
+        {
+          "id": "scatter",
+          "name": "Unique.AxisScatter",
+          "source": {
+            "shader": "unique-axis-scatter.wgsl.jinja",
+            "inputs": {
+              "capacity": "dim(shapes.Y, normalizedAxis)",
+              "axisDim": "dim(shapes.X, normalizedAxis)",
+              "inner": "max(1, inner(shapes.X, normalizedAxis))",
+              "outputAxisDim": "dim(shapes.Y, normalizedAxis)",
+              "totalOut": "numel(shapes.Y)",
+              "staticShape": true
+            }
+          },
+          "bindings": "axisScatterGlobal",
+          "dispatch": { "gridStride": "numel(shapes.Y)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
+        }
+      ]
+    },
+    {
+      "id": "parallel_dedup",
+      "priority": 10,
+      "when": ["flatContract", "integerInput", "numel(shapes.X) > 0", "numel(shapes.X) < tunables.HASH_MIN_INPUT", "numel(shapes.Y) > 0", "numel(shapes.Y) <= tunables.LOCAL_SORT_CROSSOVER", "flatLocalSortFits", "flatScanScratchFits", "flatScanDispatchFits", "true"],
+      "derive": { "scanN": "numel(shapes.X)", "scanBlocks": "flatScanBlocks" },
+      "constants": { "sorted": "attrs.sorted != 0" },
+      "intermediates": [{ "id": "flags", "dtype": "uint32", "shape": "[scanN]" }],
+      "passes": [
+        {
+          "id": "dedup",
+          "name": "Unique.Dedup",
+          "shader": "unique-dedup.wgsl.jinja",
+          "bindings": "dedup",
+          "dispatch": { "x": "scanBlocks" }
+        },
+        {
+          "id": "compact",
+          "name": "Unique.CompactSort",
+          "source": {
+            "shader": "unique-compact-sort.wgsl.jinja",
+            "inputs": { "capacity": "numel(shapes.Y)", "sortN": "pow2ceil(numel(shapes.Y))", "globalScratch": false }
+          },
+          "bindings": "compact",
+          "dispatch": { "x": 1 }
+        }
+      ]
+    },
+    {
+      "id": "parallel_dedup_large_y",
+      "priority": 5,
+      "when": ["flatContract", "integerInput", "numel(shapes.X) > 0", "numel(shapes.X) <= tunables.MAX_PARALLEL_INPUT", "numel(shapes.Y) > tunables.LOCAL_SORT_CROSSOVER", "min(numel(shapes.X), numel(shapes.Y)) <= tunables.GLOBAL_MAX_CAPACITY", "flatScanScratchFits", "flatGlobalSortScratchFits", "flatScanDispatchFits", "attrs.sorted != 0"],
+      "derive": {
+        "scanN": "numel(shapes.X)",
+        "scanBlocks": "flatScanBlocks",
+        "activeCap": "min(numel(shapes.X), numel(shapes.Y))",
+        "sortN": "pow2ceil(activeCap)"
+      },
+      "constants": { "sorted": true },
+      "intermediates": [
+        { "id": "flags", "dtype": "uint32", "shape": "[scanN]" },
+        { "id": "sortVal", "dtype": "uint32", "shape": "[sortN]" },
+        { "id": "sortKey", "dtype": "uint32", "shape": "[sortN]" },
+        { "id": "sortPad", "dtype": "uint32", "shape": "[sortN]" }
+      ],
+      "passes": [
+        {
+          "id": "dedup",
+          "name": "Unique.Dedup",
+          "shader": "unique-dedup.wgsl.jinja",
+          "bindings": "dedup",
+          "dispatch": { "x": "scanBlocks" }
+        },
+        {
+          "id": "compactsort",
+          "name": "Unique.CompactSortGlobal",
+          "source": {
+            "shader": "unique-compact-sort.wgsl.jinja",
+            "inputs": { "capacity": "numel(shapes.Y)", "sortN": "sortN", "globalScratch": true }
+          },
+          "bindings": "flatGlobalCompact",
+          "dispatch": { "x": 1 }
+        }
+      ]
+    },
+    {
+      "id": "parallel_dedup_large_y_unsorted",
+      "priority": 5,
+      "when": ["flatContract", "integerInput", "numel(shapes.X) > 0", "numel(shapes.X) <= tunables.MAX_PARALLEL_INPUT", "numel(shapes.Y) > tunables.LOCAL_SORT_CROSSOVER", "min(numel(shapes.X), numel(shapes.Y)) <= tunables.GLOBAL_MAX_CAPACITY", "flatScanScratchFits", "flatGlobalSortScratchFits", "flatScanDispatchFits", "attrs.sorted == 0"],
+      "derive": {
+        "scanN": "numel(shapes.X)",
+        "scanBlocks": "flatScanBlocks",
+        "activeCap": "min(numel(shapes.X), numel(shapes.Y))",
+        "sortN": "pow2ceil(activeCap)"
+      },
+      "constants": { "sorted": false },
+      "intermediates": [
+        { "id": "flags", "dtype": "uint32", "shape": "[scanN]" },
+        { "id": "sortVal", "dtype": "uint32", "shape": "[sortN]" }
+      ],
+      "passes": [
+        {
+          "id": "dedup",
+          "name": "Unique.Dedup",
+          "shader": "unique-dedup.wgsl.jinja",
+          "bindings": "dedup",
+          "dispatch": { "x": "scanBlocks" }
+        },
+        {
+          "id": "compactsort",
+          "name": "Unique.CompactGlobal",
+          "source": {
+            "shader": "unique-compact-sort.wgsl.jinja",
+            "inputs": { "capacity": "numel(shapes.Y)", "sortN": "sortN", "globalScratch": true }
+          },
+          "bindings": "flatGlobalCompactUnsorted",
+          "dispatch": { "x": 1 }
+        }
+      ]
+    },
+    {
+      "id": "axis_bounded",
+      "priority": 10,
+      "when": ["axisContract", "axisLocalAvailable", "true"],
+      "derive": { "axisDim": "dim(shapes.X, normalizedAxis)", "dedupBlocks": "axisScanBlocks" },
+      "constants": { "sorted": "attrs.sorted != 0" },
+      "intermediates": [{ "id": "firstFlag", "dtype": "uint32", "shape": "[axisDim]" }],
+      "passes": [
+        {
+          "id": "dedup",
+          "name": "Unique.AxisDedup",
+          "shader": "unique-axis-dedup.wgsl.jinja",
+          "bindings": "axisDedup",
+          "dispatch": { "x": "dedupBlocks" }
+        },
+        {
+          "id": "compact",
+          "name": "Unique.AxisCompactSort",
+          "source": {
+            "shader": "unique-axis-compact-sort.wgsl.jinja",
+            "inputs": {
+              "capacity": "dim(shapes.Y, normalizedAxis)",
+              "sortN": "pow2ceil(dim(shapes.Y, normalizedAxis))",
+              "compactOnly": false,
+              "globalScratch": false
+            }
+          },
+          "bindings": "axisCompactLocal",
+          "dispatch": { "x": 1 }
+        }
+      ]
+    },
+    {
+      "id": "axis_bounded_large_sorted",
+      "priority": 5,
+      "when": ["axisContract", "axisGlobalAvailable", "attrs.sorted != 0"],
+      "derive": {
+        "axisDim": "dim(shapes.X, normalizedAxis)",
+        "dedupBlocks": "axisScanBlocks",
+        "sortN": "pow2ceil(dim(shapes.Y, normalizedAxis))"
+      },
+      "constants": { "sorted": true },
+      "intermediates": [
+        { "id": "firstFlag", "dtype": "uint32", "shape": "[axisDim]" },
+        { "id": "slots", "dtype": "uint32", "shape": "[sortN]" },
+        { "id": "sortPad", "dtype": "uint32", "shape": "[sortN]" },
+        { "id": "sortKey", "dtype": "uint32", "shape": "[sortN]" }
+      ],
+      "passes": [
+        {
+          "id": "dedup",
+          "name": "Unique.AxisDedup",
+          "shader": "unique-axis-dedup.wgsl.jinja",
+          "bindings": "axisDedup",
+          "dispatch": { "x": "dedupBlocks" }
+        },
+        {
+          "id": "compact",
+          "name": "Unique.AxisCompactSortGlobal",
+          "source": {
+            "shader": "unique-axis-compact-sort.wgsl.jinja",
+            "inputs": {
+              "capacity": "dim(shapes.Y, normalizedAxis)",
+              "sortN": "sortN",
+              "compactOnly": false,
+              "globalScratch": true
+            }
+          },
+          "bindings": "axisCompactGlobal",
+          "dispatch": { "x": 1 }
+        }
+      ]
+    },
+    {
+      "id": "axis_bounded_large",
+      "priority": 5,
+      "when": ["axisContract", "axisGlobalAvailable", "attrs.sorted == 0"],
+      "derive": {
+        "axisDim": "dim(shapes.X, normalizedAxis)",
+        "dedupBlocks": "axisScanBlocks",
+        "sortN": "pow2ceil(dim(shapes.Y, normalizedAxis))"
+      },
+      "constants": { "sorted": false },
+      "intermediates": [
+        { "id": "firstFlag", "dtype": "uint32", "shape": "[axisDim]" },
+        { "id": "slots", "dtype": "uint32", "shape": "[sortN]" }
+      ],
+      "passes": [
+        {
+          "id": "dedup",
+          "name": "Unique.AxisDedup",
+          "shader": "unique-axis-dedup.wgsl.jinja",
+          "bindings": "axisDedup",
+          "dispatch": { "x": "dedupBlocks" }
+        },
+        {
+          "id": "compact",
+          "name": "Unique.AxisCompactGlobal",
+          "source": {
+            "shader": "unique-axis-compact-sort.wgsl.jinja",
+            "inputs": {
+              "capacity": "dim(shapes.Y, normalizedAxis)",
+              "sortN": "sortN",
+              "compactOnly": false,
+              "globalScratch": true
+            }
+          },
+          "bindings": "axisCompactGlobalUnsorted",
+          "dispatch": { "x": 1 }
+        }
+      ]
+    },
+    {
+      "id": "flat_metadata_indices",
+      "priority": 40,
+      "when": ["flatMetadataContract", "present.indices", "not present.inverse_indices", "not present.counts"],
+      "constants": { "sorted": "attrs.sorted != 0" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "Unique.Metadata",
+          "source": {
+            "shader": "unique.wgsl.jinja",
+            "inputs": { "hasIndices": true, "hasInverseIndices": false, "hasCounts": false }
+          },
+          "bindings": "flatMetadataIndices",
+          "dispatch": { "x": 1 }
+        }
+      ]
+    },
+    {
+      "id": "axis_metadata_indices",
+      "priority": 40,
+      "when": ["axisMetadataContract", "axisOrderScratchFits", "present.indices", "not present.inverse_indices", "not present.counts"],
+      "constants": { "sorted": "attrs.sorted != 0" },
+      "intermediates": [{ "id": "order", "dtype": "uint32", "shape": "[max(1, axisOutputDim)]" }],
+      "passes": [
+        {
+          "id": "main",
+          "name": "Unique.AxisMetadata",
+          "source": {
+            "shader": "unique-axis.wgsl.jinja",
+            "inputs": {
+              "capacity": "max(1, dim(shapes.Y, normalizedAxis))",
+              "hasIndices": true,
+              "hasInverseIndices": false,
+              "hasCounts": false,
+              "headCacheSlots": "axisHeadCacheSlots"
+            }
+          },
+          "bindings": "axisMetadataIndices",
+          "dispatch": { "x": 1 }
+        }
+      ]
+    },
+    {
+      "id": "flat_metadata_inverse",
+      "priority": 40,
+      "when": ["flatMetadataContract", "not present.indices", "present.inverse_indices", "not present.counts"],
+      "constants": { "sorted": "attrs.sorted != 0" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "Unique.Metadata",
+          "source": {
+            "shader": "unique.wgsl.jinja",
+            "inputs": { "hasIndices": false, "hasInverseIndices": true, "hasCounts": false }
+          },
+          "bindings": "flatMetadataInverse",
+          "dispatch": { "x": 1 }
+        }
+      ]
+    },
+    {
+      "id": "axis_metadata_inverse",
+      "priority": 40,
+      "when": ["axisMetadataContract", "axisOrderScratchFits", "not present.indices", "present.inverse_indices", "not present.counts"],
+      "constants": { "sorted": "attrs.sorted != 0" },
+      "intermediates": [{ "id": "order", "dtype": "uint32", "shape": "[max(1, axisOutputDim)]" }],
+      "passes": [
+        {
+          "id": "main",
+          "name": "Unique.AxisMetadata",
+          "source": {
+            "shader": "unique-axis.wgsl.jinja",
+            "inputs": {
+              "capacity": "max(1, dim(shapes.Y, normalizedAxis))",
+              "hasIndices": false,
+              "hasInverseIndices": true,
+              "hasCounts": false,
+              "headCacheSlots": "axisHeadCacheSlots"
+            }
+          },
+          "bindings": "axisMetadataInverse",
+          "dispatch": { "x": 1 }
+        }
+      ]
+    },
+    {
+      "id": "flat_metadata_counts",
+      "priority": 40,
+      "when": ["flatMetadataContract", "not present.indices", "not present.inverse_indices", "present.counts"],
+      "constants": { "sorted": "attrs.sorted != 0" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "Unique.Metadata",
+          "source": {
+            "shader": "unique.wgsl.jinja",
+            "inputs": { "hasIndices": false, "hasInverseIndices": false, "hasCounts": true }
+          },
+          "bindings": "flatMetadataCounts",
+          "dispatch": { "x": 1 }
+        }
+      ]
+    },
+    {
+      "id": "axis_metadata_counts",
+      "priority": 40,
+      "when": ["axisMetadataContract", "axisOrderScratchFits", "not present.indices", "not present.inverse_indices", "present.counts"],
+      "constants": { "sorted": "attrs.sorted != 0" },
+      "intermediates": [{ "id": "order", "dtype": "uint32", "shape": "[max(1, axisOutputDim)]" }],
+      "passes": [
+        {
+          "id": "main",
+          "name": "Unique.AxisMetadata",
+          "source": {
+            "shader": "unique-axis.wgsl.jinja",
+            "inputs": {
+              "capacity": "max(1, dim(shapes.Y, normalizedAxis))",
+              "hasIndices": false,
+              "hasInverseIndices": false,
+              "hasCounts": true,
+              "headCacheSlots": "axisHeadCacheSlots"
+            }
+          },
+          "bindings": "axisMetadataCounts",
+          "dispatch": { "x": 1 }
+        }
+      ]
+    },
+    {
+      "id": "flat_metadata_indices_inverse",
+      "priority": 40,
+      "when": ["flatMetadataContract", "present.indices", "present.inverse_indices", "not present.counts"],
+      "constants": { "sorted": "attrs.sorted != 0" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "Unique.Metadata",
+          "source": {
+            "shader": "unique.wgsl.jinja",
+            "inputs": { "hasIndices": true, "hasInverseIndices": true, "hasCounts": false }
+          },
+          "bindings": "flatMetadataIndicesInverse",
+          "dispatch": { "x": 1 }
+        }
+      ]
+    },
+    {
+      "id": "axis_metadata_indices_inverse",
+      "priority": 40,
+      "when": ["axisMetadataContract", "axisOrderScratchFits", "present.indices", "present.inverse_indices", "not present.counts"],
+      "constants": { "sorted": "attrs.sorted != 0" },
+      "intermediates": [{ "id": "order", "dtype": "uint32", "shape": "[max(1, axisOutputDim)]" }],
+      "passes": [
+        {
+          "id": "main",
+          "name": "Unique.AxisMetadata",
+          "source": {
+            "shader": "unique-axis.wgsl.jinja",
+            "inputs": {
+              "capacity": "max(1, dim(shapes.Y, normalizedAxis))",
+              "hasIndices": true,
+              "hasInverseIndices": true,
+              "hasCounts": false,
+              "headCacheSlots": "axisHeadCacheSlots"
+            }
+          },
+          "bindings": "axisMetadataIndicesInverse",
+          "dispatch": { "x": 1 }
+        }
+      ]
+    },
+    {
+      "id": "flat_metadata_indices_counts",
+      "priority": 40,
+      "when": ["flatMetadataContract", "present.indices", "not present.inverse_indices", "present.counts"],
+      "constants": { "sorted": "attrs.sorted != 0" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "Unique.Metadata",
+          "source": {
+            "shader": "unique.wgsl.jinja",
+            "inputs": { "hasIndices": true, "hasInverseIndices": false, "hasCounts": true }
+          },
+          "bindings": "flatMetadataIndicesCounts",
+          "dispatch": { "x": 1 }
+        }
+      ]
+    },
+    {
+      "id": "axis_metadata_indices_counts",
+      "priority": 40,
+      "when": ["axisMetadataContract", "axisOrderScratchFits", "present.indices", "not present.inverse_indices", "present.counts"],
+      "constants": { "sorted": "attrs.sorted != 0" },
+      "intermediates": [{ "id": "order", "dtype": "uint32", "shape": "[max(1, axisOutputDim)]" }],
+      "passes": [
+        {
+          "id": "main",
+          "name": "Unique.AxisMetadata",
+          "source": {
+            "shader": "unique-axis.wgsl.jinja",
+            "inputs": {
+              "capacity": "max(1, dim(shapes.Y, normalizedAxis))",
+              "hasIndices": true,
+              "hasInverseIndices": false,
+              "hasCounts": true,
+              "headCacheSlots": "axisHeadCacheSlots"
+            }
+          },
+          "bindings": "axisMetadataIndicesCounts",
+          "dispatch": { "x": 1 }
+        }
+      ]
+    },
+    {
+      "id": "flat_metadata_inverse_counts",
+      "priority": 40,
+      "when": ["flatMetadataContract", "not present.indices", "present.inverse_indices", "present.counts"],
+      "constants": { "sorted": "attrs.sorted != 0" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "Unique.Metadata",
+          "source": {
+            "shader": "unique.wgsl.jinja",
+            "inputs": { "hasIndices": false, "hasInverseIndices": true, "hasCounts": true }
+          },
+          "bindings": "flatMetadataInverseCounts",
+          "dispatch": { "x": 1 }
+        }
+      ]
+    },
+    {
+      "id": "axis_metadata_inverse_counts",
+      "priority": 40,
+      "when": ["axisMetadataContract", "axisOrderScratchFits", "not present.indices", "present.inverse_indices", "present.counts"],
+      "constants": { "sorted": "attrs.sorted != 0" },
+      "intermediates": [{ "id": "order", "dtype": "uint32", "shape": "[max(1, axisOutputDim)]" }],
+      "passes": [
+        {
+          "id": "main",
+          "name": "Unique.AxisMetadata",
+          "source": {
+            "shader": "unique-axis.wgsl.jinja",
+            "inputs": {
+              "capacity": "max(1, dim(shapes.Y, normalizedAxis))",
+              "hasIndices": false,
+              "hasInverseIndices": true,
+              "hasCounts": true,
+              "headCacheSlots": "axisHeadCacheSlots"
+            }
+          },
+          "bindings": "axisMetadataInverseCounts",
+          "dispatch": { "x": 1 }
+        }
+      ]
+    },
+    {
+      "id": "flat_metadata_all",
+      "priority": 40,
+      "when": ["flatMetadataContract", "present.indices", "present.inverse_indices", "present.counts"],
+      "constants": { "sorted": "attrs.sorted != 0" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "Unique.Metadata",
+          "source": {
+            "shader": "unique.wgsl.jinja",
+            "inputs": { "hasIndices": true, "hasInverseIndices": true, "hasCounts": true }
+          },
+          "bindings": "flatMetadataAll",
+          "dispatch": { "x": 1 }
+        }
+      ]
+    },
+    {
+      "id": "axis_metadata_all",
+      "priority": 40,
+      "when": ["axisMetadataContract", "axisOrderScratchFits", "present.indices", "present.inverse_indices", "present.counts"],
+      "constants": { "sorted": "attrs.sorted != 0" },
+      "intermediates": [{ "id": "order", "dtype": "uint32", "shape": "[max(1, axisOutputDim)]" }],
+      "passes": [
+        {
+          "id": "main",
+          "name": "Unique.AxisMetadata",
+          "source": {
+            "shader": "unique-axis.wgsl.jinja",
+            "inputs": {
+              "capacity": "max(1, dim(shapes.Y, normalizedAxis))",
+              "hasIndices": true,
+              "hasInverseIndices": true,
+              "hasCounts": true,
+              "headCacheSlots": "axisHeadCacheSlots"
+            }
+          },
+          "bindings": "axisMetadataAll",
+          "dispatch": { "x": 1 }
+        }
+      ]
+    }
+  ]
+}

build/webgpu/metadata.json ADDED Viewed

	@@ -0,0 +1,30 @@

+{
+  "name": "ai.onnx.Unique",
+  "id": "_ai_onnx_unique_webgpu_0d47131",
+  "version": 1,
+  "license": "Apache-2.0",
+  "backend": { "type": "webgpu" },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "bench.json": "uWuMe2FMuN7vfuafXZShcUDGe/C8rp2oVHckKukc+pQ=",
+      "manifest.json": "b9zSmspb5Oo2oSITmFfU4X/mYZnv0Bm1Uu8caVO1jEg=",
+      "test.json": "dcDsfDw6xJR+OIbeTcptXvo7dV21lPAiDnElSLafb0Q=",
+      "unique-axis-compact-sort.wgsl.jinja": "RAgXG0xLG/Yw3PwOPrrJvXOOMINPBQ/REenbvq73Ak0=",
+      "unique-axis-dedup.wgsl.jinja": "O1EWypziZGw+vSmIZWyYo7IBihHTAEEYrUtUSdA+sk4=",
+      "unique-axis-hash.wgsl.jinja": "tLAuHVQcTX2Q+/gIbwupvPWe+85TP9W4oyxAvc1QLcY=",
+      "unique-axis-scatter.wgsl.jinja": "54tyC3ECbMR48qFXnBQ2el6+tMdOipaM+l91vySKeL8=",
+      "unique-axis.wgsl.jinja": "fjQxtH8MjP43bhdPOcxya+xVerIi30R4cgs7URqc71s=",
+      "unique-compact-sort.wgsl.jinja": "jeh9pkkZIYSv2jQA/StxvAaQTVp/sI0eUsekeYSnkvo=",
+      "unique-dedup.wgsl.jinja": "muTd580L/rfaL/p5NjOwMPp2F3ALTR9S7SzVYl9r+wQ=",
+      "unique-hash-build.wgsl.jinja": "HkomLl6SIyeWaOI0dRrXbl2chUwav16ClFwxZ+EDZzY=",
+      "unique-hash-collect.wgsl.jinja": "5+RKxUzJlrcqJ3Z/RfyznHl43WSoVKJSwGlQ/KzwnI4=",
+      "unique-hash-init.wgsl.jinja": "d4tz0KAxmNaRn4ZvRQemr+7wk4GwRZY9tLQguTuIcGo=",
+      "unique-hash-mark.wgsl.jinja": "87mwSXE1BpY/3zHxtaOhLpm4IZAF8hSX6bLvr6GfncE=",
+      "unique-hash-sort-collected-key-only.wgsl.jinja": "Rpuf6QQZwx2XvogNtJoTS5/vsNHbA721YwBqwWt0SLQ=",
+      "unique.wgsl.jinja": "J4xqWcvx4Ca4uW4ZbJ1xG2pfZbvL3cdnYU33fgHg1eo="
+    }
+  },
+  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
+  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.Unique" }
+}

build/webgpu/test.json ADDED Viewed

	@@ -0,0 +1,1530 @@

+{
+  "op": "ai.onnx.Unique",
+  "fixtureArrays": {
+    "ort_axis1_unsorted_f32_input_x": [1, 1, 0, 1, 2, 1, 0, 1, 1, 1, 0, 1, 2, 1, 0, 1],
+    "serial_sorted_dense_duplicates_int32_input_x": [37, -5, 12, 99, -73, 0, 41, 8, -21, 64, 3, -90, 55, 17, -2, 76, 29, -48, 83, 6, -33, 92, 14, -67, 50, 22, -9, 70, 35, -58, 88, 1, -26, 95, 19, -81, 46, 27, -14, 61, 33, -44, 79, 10, -39, 53, 24, -86]
+  },
+  "cases": [
+    {
+      "name": "sorted_int32",
+      "inputs": {
+        "x": { "dtype": "int32", "shape": [8], "data": { "kind": "values", "values": [3, 1, 3, 2, 1, -1, 2, 4] } }
+      },
+      "outputs": { "y": { "dtype": "int32", "shape": [5] } }
+    },
+    {
+      "name": "first_seen_uint8_overflow",
+      "attrs": { "sorted": 0 },
+      "inputs": {
+        "x": { "dtype": "uint8", "shape": [7], "data": { "kind": "values", "values": [5, 6, 5, 7, 8, 9, 10] } }
+      },
+      "outputs": { "y": { "dtype": "uint8", "shape": [6] } }
+    },
+    {
+      "name": "sorted_f32",
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [6],
+          "data": { "kind": "values", "values": [2.5, -1.0, 2.5, 0.0, -1.0, 4.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [4] } }
+    },
+    {
+      "name": "f32_subnormal_distinct_from_zero",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/tensor/unique_op_test.cc",
+        "test": "Unique.Flatten_Unsorted",
+        "notes": "Subnormal finite values are distinct from zero and preserve first-seen order when sorted=0."
+      },
+      "attrs": { "sorted": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [5],
+          "data": { "kind": "values", "values": [0.0, 1e-40, 0.0, -1e-40, 1e-40] }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float32",
+          "shape": [3],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [0.0, 1e-40, -1e-40] }
+        }
+      }
+    },
+    {
+      "name": "f32_sorted_subnormal_distinct_from_zero",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/tensor/unique_op_test.cc",
+        "test": "Unique.Flatten_Sorted",
+        "notes": "Sorted Unique must keep negative subnormal, zero, and positive subnormal as distinct ordered buckets."
+      },
+      "attrs": { "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [5],
+          "data": { "kind": "values", "values": [0.0, 1e-40, -1e-40, 0.0, 1e-40] }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float32",
+          "shape": [3],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [-1e-40, 0.0, 1e-40] }
+        }
+      }
+    },
+    {
+      "name": "ort_unsorted_nan_equivalence",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/tensor/unique_op_test.cc",
+        "test": "Unique.Flatten_Unsorted",
+        "notes": "Additional ORT CPU probe: ordered-map lower_bound makes NaN comparator-equivalent to its candidate bucket at insertion, so later finite values still deduplicate normally."
+      },
+      "attrs": { "sorted": 0 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [5], "data": { "kind": "values", "values": [2.0, "NaN", 2.0, "NaN", 3.0] } }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [2], "tolerance": 0, "data": { "kind": "values", "values": [2.0, 3.0] } }
+      }
+    },
+    {
+      "name": "ort_backend_unsorted_first_seen_f32",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/tensor/unique_op_test.cc",
+        "test": "Unique.Flatten_Unsorted_MatchOnnxTest"
+      },
+      "attrs": { "sorted": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [6],
+          "data": { "kind": "values", "values": [2.0, 1.0, 1.0, 3.0, 4.0, 3.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [4], "tolerance": 0 } }
+    },
+    {
+      "name": "ort_flatten_unsorted_f32",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/tensor/unique_op_test.cc",
+        "test": "Unique.Flatten_Unsorted"
+      },
+      "attrs": { "sorted": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [6],
+          "data": { "kind": "values", "values": [1.0, 4.0, 1.0, 2.0, 2.0, 0.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [4], "tolerance": 0 } }
+    },
+    {
+      "name": "ort_flatten_sorted_f32",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/tensor/unique_op_test.cc",
+        "test": "Unique.Flatten_Sorted"
+      },
+      "attrs": { "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [6],
+          "data": { "kind": "values", "values": [1.0, 4.0, 1.0, 2.0, 2.0, 0.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [4], "tolerance": 0 } }
+    },
+    {
+      "name": "ort_empty_input_f32",
+      "provenance": { "source": "onnxruntime/test/providers/cpu/tensor/unique_op_test.cc", "test": "Unique.EmptyInput" },
+      "inputs": { "x": { "dtype": "float32", "shape": [0], "data": { "kind": "values", "values": [] } } },
+      "outputs": { "y": { "dtype": "float32", "shape": [0], "tolerance": 0 } }
+    },
+    {
+      "name": "length1_f32_backend",
+      "inputs": { "x": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [0.0] } } },
+      "outputs": { "y": { "dtype": "float32", "shape": [1], "tolerance": 0 } }
+    },
+    {
+      "name": "all_duplicates_exact_output",
+      "attrs": { "sorted": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [5],
+          "data": { "kind": "values", "values": [-7.0, -7.0, -7.0, -7.0, -7.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1], "tolerance": 0 } }
+    },
+    {
+      "name": "int8_signed_sorted_order",
+      "inputs": {
+        "x": {
+          "dtype": "int8",
+          "shape": [8],
+          "data": { "kind": "values", "values": [3, -1, -128, 3, 127, -1, 0, -128] }
+        }
+      },
+      "outputs": { "y": { "dtype": "int8", "shape": [5], "tolerance": 0 } }
+    },
+    {
+      "name": "ort_no_optional_output_int8_sorted",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/tensor/unique_op_test.cc",
+        "test": "Unique.NoOptionalOutput",
+        "notes": "The fixture requests only Y, the standard required output, with its exact data-dependent shape."
+      },
+      "inputs": {
+        "x": { "dtype": "int8", "shape": [8], "data": { "kind": "values", "values": [1, 4, -1, 2, 2, 0, -1, 4] } }
+      },
+      "outputs": { "y": { "dtype": "int8", "shape": [5], "tolerance": 0 } }
+    },
+    {
+      "name": "ort_axis0_unsorted_f32",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/tensor/unique_op_test.cc",
+        "test": "Unique.Axis0_Unsorted",
+        "notes": "The fixture supplies the exact data-dependent Y shape and requests no optional metadata outputs."
+      },
+      "attrs": { "axis": 0, "sorted": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [4, 2],
+          "data": { "kind": "values", "values": [0.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0] }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float32",
+          "shape": [3, 2],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [0.0, 1.0, 1.0, 1.0, 1.0, 0.0] }
+        }
+      }
+    },
+    {
+      "name": "ort_axis0_sorted_f32",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/tensor/unique_op_test.cc",
+        "test": "Unique.Axis0_Sorted",
+        "notes": "The fixture supplies the exact data-dependent Y shape and requests no optional metadata outputs."
+      },
+      "attrs": { "axis": 0, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [4, 2],
+          "data": { "kind": "values", "values": [0.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0] }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float32",
+          "shape": [3, 2],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [0.0, 1.0, 1.0, 0.0, 1.0, 1.0] }
+        }
+      }
+    },
+    {
+      "name": "ort_axis1_unsorted_f32",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/tensor/unique_op_test.cc",
+        "test": "Unique.Axis1_Unsorted",
+        "notes": "The fixture supplies the exact data-dependent Y shape and requests no optional metadata outputs."
+      },
+      "attrs": { "axis": 1, "sorted": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 4, 2],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/ort_axis1_unsorted_f32_input_x" } }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float32",
+          "shape": [2, 3, 2],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [1.0, 1.0, 0.0, 1.0, 2.0, 1.0, 1.0, 1.0, 0.0, 1.0, 2.0, 1.0] }
+        }
+      }
+    },
+    {
+      "name": "ort_axis1_sorted_f32",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/tensor/unique_op_test.cc",
+        "test": "Unique.Axis1_Sorted",
+        "notes": "The fixture supplies the exact data-dependent Y shape and requests no optional metadata outputs."
+      },
+      "attrs": { "axis": 1, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 4, 2],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/ort_axis1_unsorted_f32_input_x" } }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float32",
+          "shape": [2, 3, 2],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [0.0, 1.0, 1.0, 1.0, 2.0, 1.0, 0.0, 1.0, 1.0, 1.0, 2.0, 1.0] }
+        }
+      }
+    },
+    {
+      "name": "ort_axis2_unsorted_f32",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/tensor/unique_op_test.cc",
+        "test": "Unique.Axis2_Unsorted",
+        "notes": "The fixture supplies the exact data-dependent Y shape and requests no optional metadata outputs."
+      },
+      "attrs": { "axis": 2, "sorted": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 2, 4],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/ort_axis1_unsorted_f32_input_x" } }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float32",
+          "shape": [2, 2, 3],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [1.0, 1.0, 0.0, 2.0, 1.0, 0.0, 1.0, 1.0, 0.0, 2.0, 1.0, 0.0] }
+        }
+      }
+    },
+    {
+      "name": "ort_axis2_sorted_f32",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/tensor/unique_op_test.cc",
+        "test": "Unique.Axis2_Sorted",
+        "notes": "The fixture supplies the exact data-dependent Y shape and requests no optional metadata outputs."
+      },
+      "attrs": { "axis": 2, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 2, 4],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/ort_axis1_unsorted_f32_input_x" } }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float32",
+          "shape": [2, 2, 3],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [0.0, 1.0, 1.0, 0.0, 1.0, 2.0, 0.0, 1.0, 1.0, 0.0, 1.0, 2.0] }
+        }
+      }
+    },
+    {
+      "name": "ort_negative_axis_last_sorted_f32",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/tensor/unique_op_test.cc",
+        "test": "Unique.Axis2_Sorted",
+        "notes": "Same semantic case as ORT's positive axis=2 coverage, expressed with the ONNX-valid negative last-axis spelling and an exact data-dependent Y shape."
+      },
+      "attrs": { "axis": -1, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 2, 4],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/ort_axis1_unsorted_f32_input_x" } }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float32",
+          "shape": [2, 2, 3],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [0.0, 1.0, 1.0, 0.0, 1.0, 2.0, 0.0, 1.0, 1.0, 0.0, 1.0, 2.0] }
+        }
+      }
+    },
+    {
+      "name": "onnx_backend_unique_length_1",
+      "attrs": { "sorted": 1 },
+      "inputs": { "x": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [0] } } },
+      "outputs": { "y": { "dtype": "int32", "shape": [1] } },
+      "provenance": {
+        "source": "cmake/external/onnx/onnx/backend/test/data/node/test_unique_length_1",
+        "notes": "The fixture supplies the exact data-dependent Y shape. Optional ONNX int64 metadata outputs are omitted here; requested metadata is represented as uint32 where values are representable."
+      }
+    },
+    {
+      "name": "onnx_backend_unique_not_sorted_without_axis",
+      "attrs": { "sorted": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [6],
+          "data": { "kind": "values", "values": [2.0, 1.0, 1.0, 3.0, 4.0, 3.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [4] } },
+      "provenance": {
+        "source": "cmake/external/onnx/onnx/backend/test/data/node/test_unique_not_sorted_without_axis",
+        "notes": "The fixture supplies the exact data-dependent Y shape. Optional ONNX int64 metadata outputs are omitted here; requested metadata is represented as uint32 where values are representable."
+      }
+    },
+    {
+      "name": "onnx_backend_unique_sorted_without_axis",
+      "attrs": { "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [6],
+          "data": { "kind": "values", "values": [2.0, 1.0, 1.0, 3.0, 4.0, 3.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [4] } },
+      "provenance": {
+        "source": "cmake/external/onnx/onnx/backend/test/data/node/test_unique_sorted_without_axis",
+        "notes": "The fixture supplies the exact data-dependent Y shape. Optional ONNX int64 metadata outputs are omitted here; requested metadata is represented as uint32 where values are representable."
+      }
+    },
+    {
+      "name": "serial_sorted_dense_duplicates_int32",
+      "provenance": {
+        "notes": "2048 elements cycling 48 distinct values exercise the deduplication early exit and sorted exchange sort with an exact 48-element result."
+      },
+      "inputs": {
+        "x": {
+          "dtype": "int32",
+          "shape": [2048],
+          "data": {
+            "kind": "cycle",
+            "values": { "$ref": "#/fixtureArrays/serial_sorted_dense_duplicates_int32_input_x" }
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "int32", "shape": [48], "tolerance": 0 } }
+    },
+    {
+      "name": "serial_first_seen_order_f32_1024",
+      "attrs": { "sorted": 0 },
+      "provenance": {
+        "notes": "First-seen order is preserved over 1024 elements with 32 distinct values and an exact 32-element result."
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1024],
+          "data": {
+            "kind": "cycle",
+            "values": [0.5, -1.25, 3.0, 7.75, -0.5, 2.25, 9.0, -4.5, 1.5, 6.25, -8.0, 0.25, 5.5, -2.75, 4.0, 8.5, -6.25, 1.75, 7.25, -3.5, 2.5, 9.75, -0.75, 5.25, -7.5, 3.25, 6.75, -1.5, 4.75, 8.25, -5.75, 0.75]
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [32], "tolerance": 0 } }
+    },
+    {
+      "name": "serial_f32_subnormal_distinct_from_zero",
+      "attrs": { "sorted": 0 },
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/tensor/unique_op_test.cc",
+        "test": "Unique.Flatten_Unsorted",
+        "notes": "Large-input companion to the scalar subnormal case: the exact serial float comparator keeps finite subnormal values distinct from zero and emits three first-seen buckets."
+      },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [4096], "data": { "kind": "cycle", "values": [0.0, 1e-40, 0.0, -1e-40] } }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float32",
+          "shape": [3],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [0.0, 1e-40, -1e-40] }
+        }
+      }
+    },
+    {
+      "name": "serial_sorted_f32_32k_subnormal",
+      "attrs": { "sorted": 1 },
+      "provenance": {
+        "notes": "32768 f32 elements cycling six distinct values, including positive and negative subnormals, exercise the exact ordered-map float path at scale. Subnormals remain bit-distinct from zero and from each other, and sorted output uses the IEEE total-order key."
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [32768],
+          "data": { "kind": "cycle", "values": [0.0, 1e-40, -1e-40, 2.5, -3.5, 7.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [6], "tolerance": 0 } }
+    },
+    {
+      "name": "serial_unsorted_first_seen_f32_32k_subnormal",
+      "attrs": { "sorted": 0 },
+      "provenance": {
+        "notes": "First-seen-order companion for the exact serial float path: comparator-order buckets are reordered by representative input index for sorted=0."
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [32768],
+          "data": { "kind": "cycle", "values": [0.0, 1e-40, -1e-40, 2.5, -3.5, 7.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [6], "tolerance": 0 } }
+    },
+    {
+      "name": "serial_f32_32k_nan_lower_bound",
+      "attrs": { "sorted": 0 },
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/tensor/unique_op_test.cc",
+        "test": "Unique.Flatten_Unsorted",
+        "notes": "NaN equivalence on the exact ordered-map float path at large input size: each NaN resolves to the lower-bound bucket headed by 2 and forms no new bucket."
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [32768],
+          "data": { "kind": "cycle", "values": [2.0, "NaN", 3.0, "NaN", 5.0] }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float32",
+          "shape": [3],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [2.0, 3.0, 5.0] }
+        }
+      }
+    },
+    {
+      "name": "hash_sorted_int32_32k",
+      "attrs": { "sorted": 1 },
+      "provenance": {
+        "notes": "32768 int32 elements cycling 48 distinct values: exercises the hash-set parallel dedup (numel >= 32768 floor) for the sorted flatten path. Bit-exact vs the CPU reference."
+      },
+      "inputs": {
+        "x": {
+          "dtype": "int32",
+          "shape": [32768],
+          "data": {
+            "kind": "cycle",
+            "values": { "$ref": "#/fixtureArrays/serial_sorted_dense_duplicates_int32_input_x" }
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "int32", "shape": [48], "tolerance": 0 } }
+    },
+    {
+      "name": "hash_unsorted_first_seen_int32_32k",
+      "attrs": { "sorted": 0 },
+      "provenance": {
+        "notes": "32768 int32 elements cycling 48 distinct values, first-seen order: the hash dedup folds in atomicMin(index), so the compacted order matches the appearance order of the parallel/serial paths exactly."
+      },
+      "inputs": {
+        "x": {
+          "dtype": "int32",
+          "shape": [32768],
+          "data": {
+            "kind": "cycle",
+            "values": { "$ref": "#/fixtureArrays/serial_sorted_dense_duplicates_int32_input_x" }
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "int32", "shape": [48], "tolerance": 0 } }
+    },
+    {
+      "name": "hash_sentinel_value_minus_one_int32_64k",
+      "attrs": { "sorted": 1 },
+      "provenance": {
+        "notes": "65536 int32 elements whose distinct set includes -1 (bitcast<u32> == 0xffffffff, the hash table's EMPTY sentinel). Exercises the dedicated `special` min-index slot the hash dedup uses for the one value that cannot be a hash key."
+      },
+      "inputs": {
+        "x": {
+          "dtype": "int32",
+          "shape": [65536],
+          "data": {
+            "kind": "cycle",
+            "values": [-1, 7, -1, 3, 100, -50, -1, 42, 7, 3, -2147483648, 2147483647, 0, -1, 13, 100]
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "int32", "shape": [10], "tolerance": 0 } }
+    },
+    {
+      "name": "hash_uint8_sorted_32k",
+      "attrs": { "sorted": 1 },
+      "provenance": {
+        "notes": "32768 uint8 elements cycling 12 distinct values: exercises the hash dedup's unsigned-key path (dtypes.T == u32 carried width) at the >= 32768 floor."
+      },
+      "inputs": {
+        "x": {
+          "dtype": "uint8",
+          "shape": [32768],
+          "data": { "kind": "cycle", "values": [5, 200, 17, 5, 255, 0, 128, 17, 64, 200, 3, 250, 0, 5] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint8", "shape": [9], "tolerance": 0 } }
+    },
+    {
+      "name": "hash_uint32_few_distinct_64k",
+      "attrs": { "sorted": 1 },
+      "provenance": {
+        "notes": "65536 uint32 elements with only 6 distinct values (heavy duplicates) including 0xffffffff (the EMPTY sentinel): few-distinct stress for the hash dedup atomicMin contention plus the unsigned special-slot path."
+      },
+      "inputs": {
+        "x": {
+          "dtype": "uint32",
+          "shape": [65536],
+          "data": { "kind": "cycle", "values": [4294967295, 0, 7, 4294967295, 1000000, 42, 7, 0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [5], "tolerance": 0 } }
+    },
+    {
+      "name": "axis0_f32_subnormal_collapse_unsorted",
+      "attrs": { "axis": 0, "sorted": 0 },
+      "provenance": {
+        "notes": "Spec-valid axis-mode deduplication over a [N,2] tensor. Raw-bit float comparisons keep finite subnormals distinct from zero even on FTZ GPUs; the pinned output is cross-checked with ORT."
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [5, 2],
+          "data": { "kind": "values", "values": [0.0, 5.0, 1e-40, 5.0, -1e-40, 5.0, 0.0, 5.0, 2.0, 7.0] }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float32",
+          "shape": [4, 2],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [0.0, 5.0, 1e-40, 5.0, -1e-40, 5.0, 2.0, 7.0] }
+        }
+      }
+    },
+    {
+      "name": "axis0_f32_subnormal_collapse_sorted",
+      "attrs": { "axis": 0, "sorted": 1 },
+      "provenance": {
+        "notes": "Sorted companion to the subnormal axis case. Raw-bit equality and total-order keys preserve -1e-40 < 0 < 1e-40 on FTZ GPUs; the pinned output is cross-checked with ORT."
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [5, 2],
+          "data": { "kind": "values", "values": [0.0, 5.0, 1e-40, 5.0, -1e-40, 5.0, 0.0, 5.0, 2.0, 7.0] }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float32",
+          "shape": [4, 2],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [-1e-40, 5.0, 0.0, 5.0, 1e-40, 5.0, 2.0, 7.0] }
+        }
+      }
+    },
+    {
+      "name": "hash_int8_signed_negatives_32k",
+      "attrs": { "sorted": 1 },
+      "provenance": {
+        "notes": "Exercises hash deduplication at 32768 elements with 12 signed int8 values, including -128, 127, and other negatives. The TypeScript reference is the oracle because ORT CPU does not implement Unique for int8."
+      },
+      "inputs": {
+        "x": {
+          "dtype": "int8",
+          "shape": [32768],
+          "data": { "kind": "cycle", "values": [5, -1, -128, 3, 127, -1, 0, -128, 42, -64, 100, -100, 7, 17] }
+        }
+      },
+      "outputs": { "y": { "dtype": "int8", "shape": [12], "tolerance": 0 } }
+    },
+    {
+      "name": "parallel_uint32_sentinel_first_seen_4k",
+      "attrs": { "sorted": 0 },
+      "provenance": {
+        "notes": "Exercises grid-parallel deduplication below the hash threshold with 0xffffffff present in first-seen order. The TypeScript reference is the oracle because ORT CPU does not implement Unique for uint32."
+      },
+      "inputs": {
+        "x": {
+          "dtype": "uint32",
+          "shape": [4096],
+          "data": { "kind": "cycle", "values": [4294967295, 0, 7, 1000000, 4294967295, 42, 0, 13, 7, 999] }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [7], "tolerance": 0 } }
+    },
+    {
+      "name": "flat_f32_all_distinct_4096",
+      "provenance": {
+        "notes": "4096 fully distinct values exercise the exact large-output serial float path and its ordered representative set."
+      },
+      "attrs": { "sorted": 1 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [4096], "data": { "kind": "linspace", "start": -4096.0, "end": 4095.0 } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [4096], "tolerance": 0 } }
+    },
+    {
+      "name": "axis0_f32_all_distinct_2500",
+      "provenance": {
+        "notes": "2500 fully distinct scalar rows exercise exact axis-mode output and bit-preserving slice comparison at a large axis size."
+      },
+      "attrs": { "axis": 0, "sorted": 1 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [2500, 1], "data": { "kind": "linspace", "start": -5000.0, "end": 4999.0 } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2500, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "serial_f32_32k_signed_zero_single_bucket",
+      "provenance": {
+        "notes": "32768 f32 elements containing both +0.0 and -0.0 plus finite values verify that exact serial equality canonicalizes signed zero into one bucket, matching ORT."
+      },
+      "attrs": { "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [32768],
+          "data": { "kind": "cycle", "values": [0.0, 0.0, 2.5, -3.5, 7.0, 0.0, 0.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [4], "tolerance": 0 } }
+    },
+    {
+      "name": "large_y_int32_sorted_2501_distinct",
+      "provenance": {
+        "notes": "6000 int32 values deduplicate to an exact 2501-element sorted result, exercising the large-output global-scratch sort and signed integer order across zero."
+      },
+      "attrs": { "sorted": 1 },
+      "inputs": {
+        "x": { "dtype": "int32", "shape": [6000], "data": { "kind": "linspace", "start": -1200, "end": 1300 } }
+      },
+      "outputs": { "y": { "dtype": "int32", "shape": [2501] } }
+    },
+    {
+      "name": "large_y_int32_unsorted_8192_distinct",
+      "provenance": {
+        "notes": "8192 distinct int32 values exercise the large-output unsorted compaction path without allocating or touching bitonic-sort scratch."
+      },
+      "attrs": { "sorted": 0 },
+      "inputs": {
+        "x": { "dtype": "int32", "shape": [8192], "data": { "kind": "linspace", "start": -4096, "end": 4095 } }
+      },
+      "outputs": { "y": { "dtype": "int32", "shape": [8192] } }
+    },
+    {
+      "name": "f32_sorted_padding_heavy",
+      "provenance": {
+        "notes": "A padding-heavy sorted float case with about 40 distinct values verifies max-key padding and raw-bit float ordering, including negatives and signed zero."
+      },
+      "attrs": { "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [400],
+          "data": {
+            "kind": "cycle",
+            "values": [3.5, -1.0, 0.0, 2.5, -7.25, 100.0, -100.0, 0.5, -0.5, 42.0, -42.0, 1.0, -1.5, 88.75, -88.75, 6.0, -6.0, 13.5, -13.5, 21.0, -21.0, 7.0, -7.0, 55.5, -55.5, 9.0, -9.0, 64.25, -64.25, 4.0, -4.0, 17.0, -17.0, 30.0, -30.0, 2.0, -2.0, 11.0, -11.0, 99.0]
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [40] } }
+    },
+    {
+      "name": "uint32_sorted_unsigned_order",
+      "provenance": {
+        "notes": "The exact 30-element result spans values above 2^31 through 0xffffffff and verifies unsigned ordering rather than a signed interpretation of raw bits."
+      },
+      "attrs": { "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "uint32",
+          "shape": [300],
+          "data": {
+            "kind": "cycle",
+            "values": [10, 4000000000, 5, 2147483648, 0, 3000000000, 100, 2147483647, 42, 4294967295, 7, 1, 2500000000, 99, 2147483649, 3, 500, 4000000001, 8, 2, 123456, 4294967294, 55, 2147483650, 9, 777, 3500000000, 6, 4, 1000000]
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "uint32", "shape": [30] } }
+    },
+    {
+      "name": "int32_unsorted_first_occurrence",
+      "provenance": {
+        "notes": "The exact 30-element unsorted result preserves first-occurrence order while round-tripping signed int32 values through raw-bit scratch."
+      },
+      "attrs": { "sorted": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "int32",
+          "shape": [300],
+          "data": {
+            "kind": "cycle",
+            "values": [37, -5, 12, 99, -73, 0, 41, 8, -21, 64, 3, -90, 55, 17, -2, 76, 29, -48, 83, 6, -33, 92, 14, -67, 50, 22, -9, 70, 35, -58]
+          }
+        }
+      },
+      "outputs": { "y": { "dtype": "int32", "shape": [30] } }
+    },
+    {
+      "name": "axis0_large_y_f32_sorted_2200_distinct",
+      "provenance": {
+        "notes": "The exact serial float axis path handles 2200 distinct scalar slices with its representative order in global storage. This keeps float behavior aligned with ORT's non-transitive NaN comparator without exceeding workgroup-storage limits."
+      },
+      "attrs": { "axis": 0, "sorted": 1 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [2200, 1], "data": { "kind": "linspace", "start": -500.0, "end": 500.0 } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2200, 1] } }
+    },
+    {
+      "name": "axis0_large_y_int32_unsorted_2100_distinct",
+      "provenance": {
+        "notes": "axis=0 output axis dim 2100 > 2048 with sorted=0 selects axis_bounded_large's no-bitonic branch: parallel compaction gathers the first-occurrence slice indices in appearance order and scatters them from GLOBAL slots scratch, skipping the sort network. inner=1 so flat-unique == axis-unique (reference exact). Validates the unsorted large-axis path at multi-chunk scale against the single-lane unique-axis kernel."
+      },
+      "attrs": { "axis": 0, "sorted": 0 },
+      "inputs": {
+        "x": { "dtype": "int32", "shape": [2100, 1], "data": { "kind": "linspace", "start": -4000, "end": 4000 } }
+      },
+      "outputs": { "y": { "dtype": "int32", "shape": [2100, 1] } }
+    },
+    {
+      "name": "axis_hash_split_scatter_sorted_all_distinct",
+      "provenance": {
+        "notes": "The hash-backed axis path compacts, sorts, and scatters all 4096 distinct scalar rows into the exact output shape."
+      },
+      "attrs": { "axis": 0, "sorted": 1 },
+      "inputs": {
+        "x": { "dtype": "int32", "shape": [4096, 1], "data": { "kind": "linspace", "start": -2048, "end": 2047 } }
+      },
+      "outputs": { "y": { "dtype": "int32", "shape": [4096, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "axis_hash_split_scatter_unsorted_four_distinct",
+      "provenance": {
+        "notes": "Unsorted companion for grid-parallel axis scatter: four distinct scalar rows are emitted in first-occurrence order into an exact four-row output."
+      },
+      "attrs": { "axis": 0, "sorted": 0 },
+      "inputs": {
+        "x": { "dtype": "int32", "shape": [2048, 1], "data": { "kind": "cycle", "values": [9, -2, 7, 9, 42, -2] } }
+      },
+      "outputs": { "y": { "dtype": "int32", "shape": [4, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "axis_hash_int32_inner2_duplicate_rows_sorted",
+      "provenance": {
+        "notes": "Routes the >=2048-axis exact hash dedup with vector-valued slices. Repeated rows must share one hash bucket, first-occurrence representatives must survive, and sorted output remains lexicographic over the complete two-element slice. Pinned values avoid relying on the flattened TypeScript Unique reference for axis semantics."
+      },
+      "attrs": { "axis": 0, "sorted": 1 },
+      "inputs": {
+        "x": { "dtype": "int32", "shape": [2048, 2], "data": { "kind": "cycle", "values": [2, 1, 0, 3, 2, 1, -1, 4] } }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "int32",
+          "shape": [3, 2],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [-1, 4, 0, 3, 2, 1] }
+        }
+      }
+    },
+    {
+      "name": "hash_int32_sorted_16_distinct",
+      "provenance": {
+        "notes": "A 32768-element input with 16 distinct signed values exercises hash-backed compaction and sorting with an exact result."
+      },
+      "attrs": { "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "int32",
+          "shape": [32768],
+          "data": { "kind": "cycle", "values": [37, -5, 12, 99, -73, 0, 41, 8, -21, 64, 3, -90, 55, 17, -2, 76] }
+        }
+      },
+      "outputs": { "y": { "dtype": "int32", "shape": [16] } }
+    },
+    {
+      "name": "hash_int32_sorted_16k_4096_distinct",
+      "provenance": {
+        "notes": "A 16K-element input with exactly 4096 distinct integers exercises the narrow int32 large-output hash threshold and signed sorting."
+      },
+      "attrs": { "sorted": 1 },
+      "inputs": { "x": { "dtype": "int32", "shape": [16384], "data": { "kind": "linspace", "start": 0, "end": 4095 } } },
+      "outputs": { "y": { "dtype": "int32", "shape": [4096] } }
+    },
+    {
+      "name": "exact_output_33_unsorted_f32_serial",
+      "provenance": {
+        "notes": "Thirty-three distinct float values exercise exact-output first-occurrence ordering on the serial float path."
+      },
+      "attrs": { "sorted": 0 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [33], "data": { "kind": "linspace", "start": -16.0, "end": 16.0 } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [33], "tolerance": 0 } }
+    },
+    {
+      "name": "rank7_axis_last",
+      "attrs": { "axis": 6, "sorted": 1 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [1, 2, 1, 2, 1, 2, 4], "data": { "kind": "constant", "value": 3.0 } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 2, 1, 2, 1, 2, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "axis0_y1025_storage_order_serial",
+      "provenance": {
+        "notes": "A 1025-row exact float output exercises axis_serial with storage-backed representative order, so its capacity is independent of maxComputeWorkgroupStorageSize."
+      },
+      "attrs": { "axis": 0, "sorted": 1 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [1025, 1], "data": { "kind": "linspace", "start": -1024.0, "end": 1024.0 } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1025, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "axis0_rows65537_over_hash_ceiling_large_y",
+      "provenance": {
+        "notes": "65537 rows exceed the axis hash-table ceiling, while exactly 2049 distinct rows keep the output above the local crossover; this selects the global axis fallback with an exact result."
+      },
+      "attrs": { "axis": 0, "sorted": 1 },
+      "inputs": {
+        "x": { "dtype": "int32", "shape": [65537, 1], "data": { "kind": "linspace", "start": 0, "end": 2048 } }
+      },
+      "outputs": { "y": { "dtype": "int32", "shape": [2049, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "rank8_axis_last",
+      "attrs": { "axis": 7, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 2, 1, 2, 1, 2, 2, 4],
+          "data": { "kind": "cycle", "values": [1.0, 2.0, 3.0, 1.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 2, 1, 2, 1, 2, 2, 3], "tolerance": 0 } }
+    },
+    {
+      "name": "axis0_uint32_sorted_order",
+      "provenance": {
+        "source": "https://onnx.ai/onnx/operators/onnx__Unique.html",
+        "notes": "Exercises unsigned slice ordering through the bounded parallel axis compaction path, including values above int32 range."
+      },
+      "attrs": { "axis": 0, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "uint32",
+          "shape": [6, 1],
+          "data": { "kind": "values", "values": [4000000000, 5, 2147483648, 5, 0, 4000000000] }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "uint32",
+          "shape": [4, 1],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [0, 5, 2147483648, 4000000000] }
+        }
+      }
+    },
+    {
+      "name": "axis_empty_y_only",
+      "provenance": {
+        "source": "https://onnx.ai/onnx/operators/onnx__Unique.html",
+        "notes": "An empty selected axis has zero unique slices and therefore a legal zero-sized Y axis."
+      },
+      "attrs": { "axis": 0, "sorted": 1 },
+      "inputs": { "x": { "dtype": "int32", "shape": [0, 2], "data": { "kind": "values", "values": [] } } },
+      "outputs": { "y": { "dtype": "int32", "shape": [0, 2], "tolerance": 0 } }
+    },
+    {
+      "name": "axis_empty_all_metadata",
+      "provenance": {
+        "source": "https://onnx.ai/onnx/operators/onnx__Unique.html",
+        "notes": "Zero-length axis mode requests every standard optional metadata output with its exact empty shape."
+      },
+      "attrs": { "axis": 1, "sorted": 0 },
+      "inputs": { "x": { "dtype": "int32", "shape": [2, 0, 3], "data": { "kind": "values", "values": [] } } },
+      "outputs": {
+        "y": { "dtype": "int32", "shape": [2, 0, 3], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [0], "tolerance": 0, "data": { "kind": "values", "values": [] } },
+        "inverse_indices": {
+          "dtype": "uint32",
+          "shape": [0],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [] }
+        },
+        "counts": { "dtype": "uint32", "shape": [0], "tolerance": 0, "data": { "kind": "values", "values": [] } }
+      }
+    },
+    {
+      "name": "axis_nan_ordered_map_unsorted_all_metadata",
+      "provenance": {
+        "source": "onnxruntime/core/providers/cpu/tensor/unique.cc",
+        "notes": "ORT's ordered slice comparator stops at the first unequal coordinate. A NaN there makes the key equivalent to the lower-bound bucket and ignores the remaining suffix."
+      },
+      "attrs": { "axis": 0, "sorted": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [4, 2],
+          "data": { "kind": "values", "values": [5.0, 0.0, 1.0, 0.0, "NaN", 9.0, 3.0, 0.0] }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float32",
+          "shape": [3, 2],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [5.0, 0.0, 1.0, 0.0, 3.0, 0.0] }
+        },
+        "indices": {
+          "dtype": "uint32",
+          "shape": [3],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [0, 1, 3] }
+        },
+        "inverse_indices": {
+          "dtype": "uint32",
+          "shape": [4],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [0, 1, 1, 2] }
+        },
+        "counts": { "dtype": "uint32", "shape": [3], "tolerance": 0, "data": { "kind": "values", "values": [1, 2, 1] } }
+      }
+    },
+    {
+      "name": "axis_nan_ordered_map_sorted_all_metadata",
+      "provenance": {
+        "source": "onnxruntime/core/providers/cpu/tensor/unique.cc",
+        "notes": "The same NaN lower-bound equivalence is retained while sorted output follows the ordered map's comparator order."
+      },
+      "attrs": { "axis": 0, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [4, 2],
+          "data": { "kind": "values", "values": [5.0, 0.0, 1.0, 0.0, "NaN", 9.0, 3.0, 0.0] }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float32",
+          "shape": [3, 2],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [1.0, 0.0, 3.0, 0.0, 5.0, 0.0] }
+        },
+        "indices": {
+          "dtype": "uint32",
+          "shape": [3],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [1, 3, 0] }
+        },
+        "inverse_indices": {
+          "dtype": "uint32",
+          "shape": [4],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [2, 0, 0, 1] }
+        },
+        "counts": { "dtype": "uint32", "shape": [3], "tolerance": 0, "data": { "kind": "values", "values": [2, 1, 1] } }
+      }
+    },
+    {
+      "name": "flat_nan_ordered_map_unsorted_all_metadata",
+      "provenance": {
+        "source": "onnxruntime/core/providers/cpu/tensor/unique.cc",
+        "notes": "Flat Unique uses the same stateful lower_bound rule: NaN maps to the smallest bucket present at its insertion point, not unconditionally to output bucket zero."
+      },
+      "attrs": { "sorted": 0 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [4], "data": { "kind": "values", "values": [5.0, 1.0, "NaN", 3.0] } }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float32",
+          "shape": [3],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [5.0, 1.0, 3.0] }
+        },
+        "indices": {
+          "dtype": "uint32",
+          "shape": [3],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [0, 1, 3] }
+        },
+        "inverse_indices": {
+          "dtype": "uint32",
+          "shape": [4],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [0, 1, 1, 2] }
+        },
+        "counts": { "dtype": "uint32", "shape": [3], "tolerance": 0, "data": { "kind": "values", "values": [1, 2, 1] } }
+      }
+    },
+    {
+      "name": "flat_nan_ordered_map_sorted_all_metadata",
+      "provenance": {
+        "source": "onnxruntime/core/providers/cpu/tensor/unique.cc",
+        "notes": "Sorted flat output keeps the ordered-map bucket accounting while emitting comparator order."
+      },
+      "attrs": { "sorted": 1 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [4], "data": { "kind": "values", "values": [5.0, 1.0, "NaN", 3.0] } }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float32",
+          "shape": [3],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [1.0, 3.0, 5.0] }
+        },
+        "indices": {
+          "dtype": "uint32",
+          "shape": [3],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [1, 3, 0] }
+        },
+        "inverse_indices": {
+          "dtype": "uint32",
+          "shape": [4],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [2, 0, 0, 1] }
+        },
+        "counts": { "dtype": "uint32", "shape": [3], "tolerance": 0, "data": { "kind": "values", "values": [2, 1, 1] } }
+      }
+    },
+    {
+      "name": "flat_nan_ordered_map_insertion_state_all_metadata",
+      "provenance": {
+        "source": "onnxruntime/core/providers/cpu/tensor/unique.cc",
+        "notes": "With NaN inserted before the later minimum, lower_bound associates it with 5 rather than retroactively moving it to the later 1 bucket."
+      },
+      "attrs": { "sorted": 0 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [4], "data": { "kind": "values", "values": [5.0, "NaN", 1.0, 3.0] } }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float32",
+          "shape": [3],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [5.0, 1.0, 3.0] }
+        },
+        "indices": {
+          "dtype": "uint32",
+          "shape": [3],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [0, 2, 3] }
+        },
+        "inverse_indices": {
+          "dtype": "uint32",
+          "shape": [4],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [0, 0, 1, 2] }
+        },
+        "counts": { "dtype": "uint32", "shape": [3], "tolerance": 0, "data": { "kind": "values", "values": [2, 1, 1] } }
+      }
+    },
+    {
+      "name": "metadata_flat_all_int32",
+      "provenance": {
+        "source": "https://onnx.ai/onnx/operators/onnx__Unique.html",
+        "notes": "Exercises every optional output on the non-float flattened comparator path."
+      },
+      "attrs": { "sorted": 1 },
+      "inputs": { "x": { "dtype": "int32", "shape": [5], "data": { "kind": "values", "values": [3, 1, 3, 2, 1] } } },
+      "outputs": {
+        "y": { "dtype": "int32", "shape": [3], "tolerance": 0, "data": { "kind": "values", "values": [1, 2, 3] } },
+        "indices": {
+          "dtype": "uint32",
+          "shape": [3],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [1, 3, 0] }
+        },
+        "inverse_indices": {
+          "dtype": "uint32",
+          "shape": [5],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [2, 0, 2, 1, 0] }
+        },
+        "counts": { "dtype": "uint32", "shape": [3], "tolerance": 0, "data": { "kind": "values", "values": [2, 1, 2] } }
+      }
+    },
+    {
+      "name": "metadata_flat_indices_nd_input",
+      "provenance": {
+        "source": "https://onnx.ai/onnx/operators/onnx__Unique.html",
+        "notes": "Exercises the requested standard optional output combination while flattening an N-D input."
+      },
+      "attrs": { "sorted": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "values", "values": [2.0, 1.0, 1.0, 3.0, 4.0, 3.0] }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [4] },
+        "indices": { "dtype": "uint32", "shape": [4], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "metadata_axis_negative_indices",
+      "provenance": {
+        "source": "https://onnx.ai/onnx/operators/onnx__Unique.html",
+        "notes": "Exercises the requested standard optional output combination and normalizes axis=-2 to axis 0 for a rank-2 input."
+      },
+      "attrs": { "axis": -2, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 3],
+          "data": { "kind": "values", "values": [1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 2.0, 3.0, 4.0] }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [2, 3] },
+        "indices": { "dtype": "uint32", "shape": [2], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "metadata_flat_inverse_nd_input",
+      "provenance": {
+        "source": "https://onnx.ai/onnx/operators/onnx__Unique.html",
+        "notes": "Exercises the requested standard optional output combination while flattening an N-D input."
+      },
+      "attrs": { "sorted": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "values", "values": [2.0, 1.0, 1.0, 3.0, 4.0, 3.0] }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [4] },
+        "inverse_indices": { "dtype": "uint32", "shape": [6], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "metadata_axis_negative_inverse",
+      "provenance": {
+        "source": "https://onnx.ai/onnx/operators/onnx__Unique.html",
+        "notes": "Exercises the requested standard optional output combination and normalizes axis=-2 to axis 0 for a rank-2 input."
+      },
+      "attrs": { "axis": -2, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 3],
+          "data": { "kind": "values", "values": [1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 2.0, 3.0, 4.0] }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [2, 3] },
+        "inverse_indices": { "dtype": "uint32", "shape": [3], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "metadata_flat_counts_nd_input",
+      "provenance": {
+        "source": "https://onnx.ai/onnx/operators/onnx__Unique.html",
+        "notes": "Exercises the requested standard optional output combination while flattening an N-D input."
+      },
+      "attrs": { "sorted": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "values", "values": [2.0, 1.0, 1.0, 3.0, 4.0, 3.0] }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [4] },
+        "counts": { "dtype": "uint32", "shape": [4], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "metadata_axis_negative_counts",
+      "provenance": {
+        "source": "https://onnx.ai/onnx/operators/onnx__Unique.html",
+        "notes": "Exercises the requested standard optional output combination and normalizes axis=-2 to axis 0 for a rank-2 input."
+      },
+      "attrs": { "axis": -2, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 3],
+          "data": { "kind": "values", "values": [1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 2.0, 3.0, 4.0] }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [2, 3] },
+        "counts": { "dtype": "uint32", "shape": [2], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "metadata_flat_indices_inverse_nd_input",
+      "provenance": {
+        "source": "https://onnx.ai/onnx/operators/onnx__Unique.html",
+        "notes": "Exercises the requested standard optional output combination while flattening an N-D input."
+      },
+      "attrs": { "sorted": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "values", "values": [2.0, 1.0, 1.0, 3.0, 4.0, 3.0] }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [4] },
+        "indices": { "dtype": "uint32", "shape": [4], "tolerance": 0 },
+        "inverse_indices": { "dtype": "uint32", "shape": [6], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "metadata_axis_negative_indices_inverse",
+      "provenance": {
+        "source": "https://onnx.ai/onnx/operators/onnx__Unique.html",
+        "notes": "Exercises the requested standard optional output combination and normalizes axis=-2 to axis 0 for a rank-2 input."
+      },
+      "attrs": { "axis": -2, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 3],
+          "data": { "kind": "values", "values": [1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 2.0, 3.0, 4.0] }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [2, 3] },
+        "indices": { "dtype": "uint32", "shape": [2], "tolerance": 0 },
+        "inverse_indices": { "dtype": "uint32", "shape": [3], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "metadata_flat_indices_counts_nd_input",
+      "provenance": {
+        "source": "https://onnx.ai/onnx/operators/onnx__Unique.html",
+        "notes": "Exercises the requested standard optional output combination while flattening an N-D input."
+      },
+      "attrs": { "sorted": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "values", "values": [2.0, 1.0, 1.0, 3.0, 4.0, 3.0] }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [4] },
+        "indices": { "dtype": "uint32", "shape": [4], "tolerance": 0 },
+        "counts": { "dtype": "uint32", "shape": [4], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "metadata_axis_negative_indices_counts",
+      "provenance": {
+        "source": "https://onnx.ai/onnx/operators/onnx__Unique.html",
+        "notes": "Exercises the requested standard optional output combination and normalizes axis=-2 to axis 0 for a rank-2 input."
+      },
+      "attrs": { "axis": -2, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 3],
+          "data": { "kind": "values", "values": [1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 2.0, 3.0, 4.0] }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [2, 3] },
+        "indices": { "dtype": "uint32", "shape": [2], "tolerance": 0 },
+        "counts": { "dtype": "uint32", "shape": [2], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "metadata_flat_inverse_counts_nd_input",
+      "provenance": {
+        "source": "https://onnx.ai/onnx/operators/onnx__Unique.html",
+        "notes": "Exercises the requested standard optional output combination while flattening an N-D input."
+      },
+      "attrs": { "sorted": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "values", "values": [2.0, 1.0, 1.0, 3.0, 4.0, 3.0] }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [4] },
+        "inverse_indices": { "dtype": "uint32", "shape": [6], "tolerance": 0 },
+        "counts": { "dtype": "uint32", "shape": [4], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "metadata_axis_negative_inverse_counts",
+      "provenance": {
+        "source": "https://onnx.ai/onnx/operators/onnx__Unique.html",
+        "notes": "Exercises the requested standard optional output combination and normalizes axis=-2 to axis 0 for a rank-2 input."
+      },
+      "attrs": { "axis": -2, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 3],
+          "data": { "kind": "values", "values": [1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 2.0, 3.0, 4.0] }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [2, 3] },
+        "inverse_indices": { "dtype": "uint32", "shape": [3], "tolerance": 0 },
+        "counts": { "dtype": "uint32", "shape": [2], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "metadata_flat_all_nd_input",
+      "provenance": {
+        "source": "https://onnx.ai/onnx/operators/onnx__Unique.html",
+        "notes": "Exercises the requested standard optional output combination while flattening an N-D input."
+      },
+      "attrs": { "sorted": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3],
+          "data": { "kind": "values", "values": [2.0, 1.0, 1.0, 3.0, 4.0, 3.0] }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [4] },
+        "indices": { "dtype": "uint32", "shape": [4], "tolerance": 0 },
+        "inverse_indices": { "dtype": "uint32", "shape": [6], "tolerance": 0 },
+        "counts": { "dtype": "uint32", "shape": [4], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "metadata_axis_negative_all",
+      "provenance": {
+        "source": "https://onnx.ai/onnx/operators/onnx__Unique.html",
+        "notes": "Exercises the requested standard optional output combination and normalizes axis=-2 to axis 0 for a rank-2 input."
+      },
+      "attrs": { "axis": -2, "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 3],
+          "data": { "kind": "values", "values": [1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 2.0, 3.0, 4.0] }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [2, 3] },
+        "indices": { "dtype": "uint32", "shape": [2], "tolerance": 0 },
+        "inverse_indices": { "dtype": "uint32", "shape": [3], "tolerance": 0 },
+        "counts": { "dtype": "uint32", "shape": [2], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "float16_sorted_subnormal_and_signed_zero",
+      "provenance": {
+        "source": "https://onnx.ai/onnx/operators/onnx__Unique.html",
+        "notes": "Covers the standard float16 type with bit-preserving subnormal ordering and signed-zero equality."
+      },
+      "attrs": { "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [5],
+          "data": { "kind": "values", "values": [0.0, -0.0, 5.960464477539063e-8, -5.960464477539063e-8, 0.0] }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float16",
+          "shape": [3],
+          "tolerance": 0,
+          "data": { "kind": "values", "values": [-5.960464477539063e-8, 0.0, 5.960464477539063e-8] }
+        }
+      }
+    },
+    {
+      "name": "float16_nan_ordered_map_all_metadata",
+      "provenance": {
+        "source": "https://onnx.ai/onnx/operators/onnx__Unique.html",
+        "notes": "Exercises float16 NaN lower-bound equivalence and every standard optional metadata output."
+      },
+      "attrs": { "sorted": 0 },
+      "inputs": {
+        "x": { "dtype": "float16", "shape": [4], "data": { "kind": "values", "values": [5.0, "NaN", 1.0, 3.0] } }
+      },
+      "outputs": {
+        "y": { "dtype": "float16", "shape": [3], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [3], "tolerance": 0 },
+        "inverse_indices": { "dtype": "uint32", "shape": [4], "tolerance": 0 },
+        "counts": { "dtype": "uint32", "shape": [3], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "float16_axis_unsorted_duplicate_rows",
+      "provenance": {
+        "source": "https://onnx.ai/onnx/operators/onnx__Unique.html",
+        "notes": "Covers standard float16 axis-mode deduplication and first-occurrence output order."
+      },
+      "attrs": { "axis": 0, "sorted": 0 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [3, 2],
+          "data": { "kind": "values", "values": [2.0, 1.0, 1.0, 3.0, 2.0, 1.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [2, 2], "tolerance": 0 } }
+    },
+    {
+      "name": "int16_sorted_extremes",
+      "provenance": {
+        "source": "https://onnx.ai/onnx/operators/onnx__Unique.html",
+        "notes": "Covers the standard int16 type at both representable extremes."
+      },
+      "attrs": { "sorted": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "int16",
+          "shape": [6],
+          "data": { "kind": "values", "values": [32767, -32768, -1, 0, 32767, -32768] }
+        }
+      },
+      "outputs": { "y": { "dtype": "int16", "shape": [4], "tolerance": 0 } }
+    },
+    {
+      "name": "bool_unsorted_all_metadata",
+      "provenance": {
+        "source": "https://onnx.ai/onnx/operators/onnx__Unique.html",
+        "notes": "Covers the standard bool type, first-occurrence order, and every optional metadata output."
+      },
+      "attrs": { "sorted": 0 },
+      "inputs": { "x": { "dtype": "bool", "shape": [5], "data": { "kind": "values", "values": [1, 0, 1, 1, 0] } } },
+      "outputs": {
+        "y": { "dtype": "bool", "shape": [2], "tolerance": 0 },
+        "indices": { "dtype": "uint32", "shape": [2], "tolerance": 0 },
+        "inverse_indices": { "dtype": "uint32", "shape": [5], "tolerance": 0 },
+        "counts": { "dtype": "uint32", "shape": [2], "tolerance": 0 }
+      }
+    }
+  ]
+}

build/webgpu/unique-axis-compact-sort.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,222 @@

+{{ env.wgsl.resourceDeclarations }}
+{% set emitCount = source.emitCount | default(false) %}
+// Compact every first-occurrence slice index, optionally bitonic-sort them,
+// then scatter the exact result. Small results keep slots in workgroup memory;
+// larger results use storage-buffer scratch with storage barriers.
+const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
+const CAP: u32 = {{ source.capacity }}u;
+{% if sorted or not source.globalScratch %}
+const SORT_N: u32 = {{ source.sortN }}u;
+{% endif %}
+{% if sorted %}
+fn less_value(a: {{ scalar }}, b: {{ scalar }}) -> bool {
+  return a < b;
+}
+{% endif %}
+{% if sorted or not source.compactOnly %}
+fn slice_at(o: u32, k: u32, n: u32) -> {{ scalar }} {
+  return x[(o * params.axisDim + k) * params.inner + n];
+}
+{% endif %}
+{% if sorted %}
+// True iff slice a < slice b lexicographically over (outer, inner) element order.
+// The serial insertion sort uses the same comparator.
+fn slice_less(a: u32, b: u32) -> bool {
+  for (var o = 0u; o < params.outer; o = o + 1u) {
+    for (var n = 0u; n < params.inner; n = n + 1u) {
+      let va = slice_at(o, a, n);
+      let vb = slice_at(o, b, n);
+      if (less_value(va, vb)) { return true; }
+      if (less_value(vb, va)) { return false; }
+    }
+  }
+  return false;
+}
+// Cache the first element's monotonic integer-order key alongside each slot so
+// the bitonic network normally compares only threadgroup u32s; equal keys still
+// use the exact full-slice comparator.
+fn slice_primary_key(k: u32) -> u32 {
+  if (params.outer == 0u || params.inner == 0u) { return 0u; }
+{% if isUnsigned %}
+  return u32(slice_at(0u, k, 0u));
+{% else %}
+  return bitcast<u32>(slice_at(0u, k, 0u)) ^ 0x80000000u;
+{% endif %}
+}
+{% endif %}
+{% if not source.compactOnly %}
+// Keep exact-result axis scatter identical across static- and dynamic-shape
+// dispatch strategies. The zero branch is a defensive guard for invalid shapes.
+fn unique_axis_zero_value() -> {{ scalar }} {
+  return {{ scalar }}(0);
+}
+fn unique_axis_scatter_element(g: u32, written: u32) {
+  let n = g % params.inner;
+  let tmp = g / params.inner;
+  let p = tmp % params.outputAxisDim;
+  let o = tmp / params.outputAxisDim;
+  if (p < written) {
+    y[g] = slice_at(o, slots[p], n);
+  } else {
+    y[g] = unique_axis_zero_value();
+  }
+}
+{% endif %}
+var<workgroup> wgScan: array<u32, WG>;
+{% if not source.globalScratch %}
+var<workgroup> slots: array<u32, SORT_N>;   // compacted (then sorted) slice indices
+{% endif %}
+var<workgroup> wgCarry: u32;
+{% if sorted and not source.globalScratch %}
+var<workgroup> sortPad: array<u32, SORT_N>;  // 1 = padding slot, sorts after every real
+var<workgroup> sortKey: array<u32, SORT_N>;  // first-element total-order key
+{% endif %}
+@compute @workgroup_size(WG)
+fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
+  let tid = lid.x;
+  if (tid == 0u) { wgCarry = 0u; }
+  workgroupBarrier();
+  // Parallel compaction: chunked Hillis-Steele exclusive scan of firstFlag, gathering
+  // the axis-index k of each first-occurrence slice into slots[appearance-rank].
+  let chunks = (params.axisDim + WG - 1u) / WG;
+  for (var c = 0u; c < chunks; c = c + 1u) {
+    let k = c * WG + tid;
+    var f = 0u;
+    if (k < params.axisDim) {
+      f = firstFlag[k];
+    }
+    wgScan[tid] = f;
+    workgroupBarrier();
+    var stride = 1u;
+    loop {
+      if (stride >= WG) { break; }
+      var add = 0u;
+      if (tid >= stride) {
+        add = wgScan[tid - stride];
+      }
+      workgroupBarrier();
+      wgScan[tid] = wgScan[tid] + add;
+      workgroupBarrier();
+      stride = stride * 2u;
+    }
+    let excl = wgScan[tid] - f;
+    let pos = wgCarry + excl;
+    if (k < params.axisDim && f == 1u && pos < CAP) {
+      slots[pos] = k;
+    }
+    workgroupBarrier();
+    if (tid == 0u) {
+      wgCarry = wgCarry + wgScan[WG - 1u];
+    }
+    workgroupBarrier();
+  }
+  let total = wgCarry;
+{% if sorted or not source.compactOnly %}
+  let written = min(total, CAP);
+{% endif %}
+{% if emitCount %}
+  if (tid == 0u) {
+    count[0] = total;
+  }
+{% endif %}
+{% if source.globalScratch %}
+  storageBarrier();
+{% endif %}
+{% if sorted %}
+  // Pad [written, SORT_N) so padding sorts strictly after every real slice, then
+  // bitonic-sort slots by slice_less (real slices are distinct, so it is a total order).
+  for (var k = tid; k < SORT_N; k = k + WG) {
+    let isPad = select(0u, 1u, k >= written);
+    sortPad[k] = isPad;
+    if (isPad == 1u) {
+      slots[k] = 0u;
+      sortKey[k] = 0u;
+    } else {
+      sortKey[k] = slice_primary_key(slots[k]);
+    }
+  }
+{% if source.globalScratch %}
+  storageBarrier();
+{% else %}
+  workgroupBarrier();
+{% endif %}
+  var size = 2u;
+  loop {
+    if (size > SORT_N) { break; }
+    var stride = size / 2u;
+    loop {
+      if (stride == 0u) { break; }
+      for (var k = tid; k < SORT_N; k = k + WG) {
+        let partner = k ^ stride;
+        if (partner > k) {
+          let ascending = (k & size) == 0u;
+          let pk = sortPad[k];
+          let pp = sortPad[partner];
+          // k should precede partner under (padding-flag asc, slice order asc)
+          var kBeforeP: bool;
+          if (pk != pp) {
+            kBeforeP = pk < pp;
+          } else if (pk == 1u) {
+            kBeforeP = true;
+          } else {
+            let kk = sortKey[k];
+            let kp = sortKey[partner];
+            if (kk != kp) {
+              kBeforeP = kk < kp;
+            } else {
+              kBeforeP = !slice_less(slots[partner], slots[k]);
+            }
+          }
+          let needSwap = select(kBeforeP, !kBeforeP, ascending);
+          if (needSwap) {
+            let tSlot = slots[k];
+            slots[k] = slots[partner];
+            slots[partner] = tSlot;
+            let tPad = sortPad[k];
+            sortPad[k] = sortPad[partner];
+            sortPad[partner] = tPad;
+            let tKey = sortKey[k];
+            sortKey[k] = sortKey[partner];
+            sortKey[partner] = tKey;
+          }
+        }
+      }
+{% if source.globalScratch %}
+      storageBarrier();
+{% else %}
+      workgroupBarrier();
+{% endif %}
+      stride = stride / 2u;
+    }
+    size = size * 2u;
+  }
+{% endif %}
+{% if not source.compactOnly %}
+{% if source.globalScratch %}
+  storageBarrier();
+{% else %}
+  workgroupBarrier();
+{% endif %}
+  // Scatter the distinct slices into the exact Y allocation: one lane per
+  // output element, flat index g = (o*outputAxisDim + p)*inner + n.
+  let totalOut = params.outer * params.outputAxisDim * params.inner;
+  for (var g = tid; g < totalOut; g = g + WG) {
+    unique_axis_scatter_element(g, written);
+  }
+{% endif %}
+}

build/webgpu/unique-axis-dedup.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,105 @@

+{{ env.wgsl.resourceDeclarations }}
+// Pass 1 of parallel axis-Unique. One thread per slice k writes
+// firstFlag[k] = 1 iff no earlier slice j < k is element-wise equal to k.
+//
+// A "slice" along the axis is the vector x[(o*axisDim + k)*inner + n] over all
+// (o in outer, n in inner). This parallel path compares only integer storage
+// values.
+//
+// The quadratic scan's common case only needs to prove that two slices differ.
+// Loading the first element of every candidate directly from X made that scan a
+// large, cache-hostile gather (the axis stride can be kilobytes). Instead, every
+// workgroup stages a chunk of small sampled slice fingerprints in threadgroup
+// memory. A fingerprint mismatch proves inequality; collisions fall through to
+// the exact element-wise comparator, so this is an optimization only and cannot
+// change Unique semantics.
+fn slice_at(o: u32, k: u32, n: u32) -> {{ scalar }} {
+  return x[(o * params.axisDim + k) * params.inner + n];
+}
+fn value_key(value: {{ scalar }}) -> u32 {
+  // Integer conversion is injective for the 8/32-bit integer storage types.
+  return u32(value);
+}
+fn mix_key(hash: u32, value: u32) -> u32 {
+  var h = (hash ^ value) * 0x9e3779b1u;
+  h = h ^ (h >> 16u);
+  return h;
+}
+fn slice_key(k: u32) -> u32 {
+  let elements = params.outer * params.inner;
+  if (elements == 0u) { return 0u; }
+  // Four evenly-spaced samples make the prefilter useful for structured rows
+  // whose leading value is shared, while keeping the staging pass tiny.
+  let last = elements - 1u;
+  var hash = 0x811c9dc5u;
+  for (var sample = 0u; sample < 4u; sample = sample + 1u) {
+    let logical = (sample * last) / 3u;
+    let o = logical / params.inner;
+    let n = logical % params.inner;
+    hash = mix_key(hash, value_key(slice_at(o, k, n)));
+  }
+  return hash;
+}
+fn slice_eq(a: u32, b: u32) -> bool {
+  for (var o = 0u; o < params.outer; o = o + 1u) {
+    for (var n = 0u; n < params.inner; n = n + 1u) {
+      if (slice_at(o, a, n) != slice_at(o, b, n)) { return false; }
+    }
+  }
+  return true;
+}
+const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
+var<workgroup> keyCache: array<u32, WG>;
+@compute @workgroup_size(WG)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>,
+        @builtin(workgroup_id) wid: vec3<u32>,
+        @builtin(local_invocation_id) lid: vec3<u32>) {
+  let k = gid.x;
+  let laneValid = k < params.axisDim;
+  var kKey = 0u;
+  if (laneValid) {
+    kKey = slice_key(k);
+  }
+  var seen = false;
+  // Workgroup b only owns k in chunk b, so no lane can need a j from a later
+  // chunk. Keeping this loop bound workgroup-uniform preserves legal barriers.
+  let chunks = min(wid.x + 1u, (params.axisDim + WG - 1u) / WG);
+  for (var chunk = 0u; chunk < chunks; chunk = chunk + 1u) {
+    let chunkBase = chunk * WG;
+    let loadK = chunkBase + lid.x;
+    var loadedKey = 0u;
+    if (loadK < params.axisDim) {
+      loadedKey = slice_key(loadK);
+    }
+    keyCache[lid.x] = loadedKey;
+    workgroupBarrier();
+    if (laneValid && !seen) {
+      if (chunkBase < k) {
+        let end = min(WG, k - chunkBase);
+        for (var offset = 0u; offset < end; offset = offset + 1u) {
+          let j = chunkBase + offset;
+          if (keyCache[offset] == kKey && slice_eq(j, k)) {
+            seen = true;
+            break;
+          }
+        }
+      }
+    }
+    // Every lane must finish reading this chunk before the next one overwrites it.
+    workgroupBarrier();
+  }
+  if (laneValid) {
+    firstFlag[k] = select(1u, 0u, seen);
+  }
+}

build/webgpu/unique-axis-hash.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,99 @@

+// Hash-backed axis Unique front end. Initialization computes one full-slice
+// hash per axis index and clears the open-addressed table. Build inserts
+// exact-equality buckets and atomically retains the earliest representative.
+// Mark converts occupied buckets into first-occurrence flags for compaction.
+// Hash collisions always use the exact slice comparator, so hashes never
+// affect semantics.
+{{ env.wgsl.resourceDeclarations }}
+const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
+const TABLE_SIZE: u32 = {{ source.tableSize }}u;
+{% if source.stage == "build" %}
+const TABLE_MASK: u32 = TABLE_SIZE - 1u;
+{% endif %}
+const EMPTY: u32 = 0xffffffffu;
+{% if source.stage == "init" %}
+fn value_key(value: {{ scalar }}) -> u32 {
+  return u32(value);
+}
+fn slice_at(o: u32, k: u32, n: u32) -> {{ scalar }} {
+  return x[(o * params.axisDim + k) * params.inner + n];
+}
+fn slice_hash(k: u32) -> u32 {
+  // FNV-1a over the complete slice. The axis index is deliberately excluded:
+  // equal slices must hash identically regardless of where they occur.
+  var hash = 0x811c9dc5u;
+  for (var o = 0u; o < params.outer; o += 1u) {
+    for (var n = 0u; n < params.inner; n += 1u) {
+      hash = (hash ^ value_key(slice_at(o, k, n))) * 0x01000193u;
+    }
+  }
+  // One final avalanche reduces clustering for short structured slices.
+  hash ^= hash >> 16u;
+  hash *= 0x7feb352du;
+  hash ^= hash >> 15u;
+  return hash;
+}
+@compute @workgroup_size(WG)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let i = gid.x;
+  if (i < TABLE_SIZE) {
+    atomicStore(&hashSlot[i], EMPTY);
+  }
+  if (i < params.axisDim) {
+    firstFlag[i] = 0u;
+    sliceHash[i] = slice_hash(i);
+  }
+}
+{% elif source.stage == "build" %}
+fn slice_at(o: u32, k: u32, n: u32) -> {{ scalar }} {
+  return x[(o * params.axisDim + k) * params.inner + n];
+}
+fn slice_eq(a: u32, b: u32) -> bool {
+  for (var o = 0u; o < params.outer; o += 1u) {
+    for (var n = 0u; n < params.inner; n += 1u) {
+      if (slice_at(o, a, n) != slice_at(o, b, n)) { return false; }
+    }
+  }
+  return true;
+}
+@compute @workgroup_size(WG)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let k = gid.x;
+  if (k >= params.axisDim) { return; }
+  let key = sliceHash[k];
+  var slot = key & TABLE_MASK;
+  for (var probe = 0u; probe < TABLE_SIZE; probe += 1u) {
+    var representative = atomicLoad(&hashSlot[slot]);
+    if (representative == EMPTY) {
+      let inserted = atomicCompareExchangeWeak(&hashSlot[slot], EMPTY, k);
+      if (inserted.exchanged) { return; }
+      representative = inserted.old_value;
+      // compare-exchange may fail spuriously while the slot is still empty.
+      if (representative == EMPTY) { continue; }
+    }
+    if (sliceHash[representative] == key && slice_eq(representative, k)) {
+      atomicMin(&hashSlot[slot], k);
+      return;
+    }
+    slot = (slot + 1u) & TABLE_MASK;
+  }
+}
+{% else %}
+@compute @workgroup_size(WG)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let slot = gid.x;
+  if (slot >= TABLE_SIZE) { return; }
+  let representative = atomicLoad(&hashSlot[slot]);
+  if (representative != EMPTY) {
+    firstFlag[representative] = 1u;
+  }
+}
+{% endif %}

build/webgpu/unique-axis-scatter.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,46 @@

+{{ env.wgsl.resourceDeclarations }}
+// Grid-parallel output stage for axis Unique. Compact/sort must remain a single
+// globally synchronized workgroup, but copying its selected slice indices into a
+// wide Y does not: distribute that bandwidth-heavy tail across the device.
+const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
+const CAP: u32 = {{ source.capacity }}u;
+const AXIS_DIM: u32 = {{ source.axisDim }}u;
+const INNER: u32 = {{ source.inner }}u;
+const OUTPUT_AXIS_DIM: u32 = {{ source.outputAxisDim }}u;
+const TOTAL_OUT: u32 = {{ source.totalOut }}u;
+fn slice_at(o: u32, k: u32, n: u32) -> {{ scalar }} {
+  return x[(o * AXIS_DIM + k) * INNER + n];
+}
+// Keep exact-result axis scatter identical across static- and dynamic-shape
+// dispatch strategies. The zero branch is a defensive guard for invalid shapes.
+fn unique_axis_zero_value() -> {{ scalar }} {
+  return {{ scalar }}(0);
+}
+fn unique_axis_scatter_element(g: u32, written: u32) {
+  let n = g % INNER;
+  let tmp = g / INNER;
+  let p = tmp % OUTPUT_AXIS_DIM;
+  let o = tmp / OUTPUT_AXIS_DIM;
+  if (p < written) {
+    y[g] = slice_at(o, slots[p], n);
+  } else {
+    y[g] = unique_axis_zero_value();
+  }
+}
+@compute @workgroup_size(WG)
+fn main(
+  @builtin(global_invocation_id) gid: vec3<u32>,
+  @builtin(num_workgroups) nwg: vec3<u32>
+) {
+  let written = min(count[0], CAP);
+  let stride = nwg.x * WG;
+  for (var g = gid.x; g < TOTAL_OUT; g += stride) {
+    unique_axis_scatter_element(g, written);
+  }
+}

build/webgpu/unique-axis.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,559 @@

+// IEEE total-order comparators for Unique. Floating-point equality and ordering
+// use raw bits because GPUs may flush subnormals in float comparisons, which
+// would otherwise collapse distinct values. Equality canonicalizes -0 to +0;
+// ordering uses a monotonic u32 key over the real line. Integers use == and <.
+{% macro float_bits_def() %}
+fn float_bits(v: {{ scalar }}) -> u32 {
+{% if usesF16 %}
+  // WGSL has no scalar u16 type. Packing v into the low component preserves
+  // its binary16 representation while producing a bitcast-compatible 32 bits.
+  return bitcast<u32>(vec2<f16>(v, 0.0h)) & 0xffffu;
+{% else %}
+  return bitcast<u32>(v);
+{% endif %}
+}
+{%- endmacro -%}
+{% macro eq_value_def() %}
+fn eq_value(a: {{ scalar }}, b: {{ scalar }}) -> bool {
+{% if isFloat %}
+  var ba = float_bits(a);
+  var bb = float_bits(b);
+{% if usesF16 %}
+  if (ba == 0x8000u) { ba = 0u; }  // -0 -> +0
+  if (bb == 0x8000u) { bb = 0u; }
+{% else %}
+  if (ba == 0x80000000u) { ba = 0u; }  // -0 -> +0
+  if (bb == 0x80000000u) { bb = 0u; }
+{% endif %}
+  return ba == bb;
+{% else %}
+  return a == b;
+{% endif %}
+}
+{%- endmacro -%}
+{%- macro less_value_def() %}
+fn less_value(a: {{ scalar }}, b: {{ scalar }}) -> bool {
+{% if isFloat %}
+  let ba = float_bits(a);
+  let bb = float_bits(b);
+{% if usesF16 %}
+  let ka = select(ba | 0x8000u, (~ba) & 0xffffu, (ba & 0x8000u) != 0u);
+  let kb = select(bb | 0x8000u, (~bb) & 0xffffu, (bb & 0x8000u) != 0u);
+{% else %}
+  let ka = select(ba | 0x80000000u, ~ba, (ba & 0x80000000u) != 0u);
+  let kb = select(bb | 0x80000000u, ~bb, (bb & 0x80000000u) != 0u);
+{% endif %}
+  return ka < kb;
+{% else %}
+  return a < b;
+{% endif %}
+}
+{%- endmacro -%}
+{%- macro zero_value_def() %}
+fn zero_value() -> {{ scalar }} {
+  return {{ scalar }}(0);
+}
+{%- endmacro %}
+{% set emitIndices = source.hasIndices | default(false) %}
+{% set emitInverseIndices = source.hasInverseIndices | default(false) %}
+{% set emitCounts = source.hasCounts | default(false) %}
+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+// Axis-mode Unique views the input as outer x axisDim x inner. Two axis slices
+// are equal iff every (outer, inner) element matches. The first occurrence of
+// each distinct slice is retained, optionally sorted lexicographically, then
+// scattered into the data-dependent output allocation.
+//
+// Float equality and ordering use the same bit-preserving rules as the other
+// Unique kernels, including canonicalized signed zero and FTZ-safe subnormals.
+const CAP: u32 = {{ source.capacity }}u;
+// The walk is one sequential step per slice and every step carries several
+// workgroup barriers, so this width is barrier latency, not throughput: the
+// head test settles in one round on real data and the chunked scan and shift
+// both adapt to any width. Tuned separately from the op-wide WORKGROUP_SIZE.
+const WG: u32 = {{ axisSerialWg }}u;
+{% if isFloat %}
+// Sentinel for the cooperative lower-bound reduction: no lane in the current
+// chunk found a candidate whose comparison fails. It is above every reachable
+// slice index because a chunk only tests indices below `written <= CAP`.
+const NO_HIT: u32 = 0xffffffffu;
+var<workgroup> lowerBoundHit: atomic<u32>;
+var<workgroup> hitIndex: u32;
+var<workgroup> diffAt: atomic<u32>;
+var<workgroup> cmpOut: u32;
+{% endif %}
+{% set headCache = isFloat and (source.headCacheSlots | default(0)) > 0 %}
+{% set floatSliceLess = isFloat and (emitInverseIndices or emitCounts) %}
+{% set headOf = "headCache[candidate]" if headCache else "head_bits(order[candidate])" %}
+{% set signBit = "0x8000u" if usesF16 else "0x80000000u" %}
+{% set keyMask = "0xffffu" if usesF16 else "0xffffffffu" %}
+{% set expMask = "0x7c00u" if usesF16 else "0x7f800000u" %}
+{% set mantMask = "0x03ffu" if usesF16 else "0x007fffffu" %}
+{% if isFloat %}
+// `is_nan_bits` and `less_value` over a bit pattern that has already been
+// loaded, so the comparators below can hoist one `float_bits` per value.
+fn is_nan_head(b: u32) -> bool {
+  return (b & {{ expMask }}) == {{ expMask }} && (b & {{ mantMask }}) != 0u;
+}
+// Same monotonic key as `less_value`. Folding -0 to +0 moves its key by one
+// step onto +0's, and no other value's key lies between them, so ordering
+// against every other value is unchanged.
+fn less_head(ba: u32, bb: u32) -> bool {
+  let ka = select(ba | {{ signBit }}, (~ba) & {{ keyMask }}, (ba & {{ signBit }}) != 0u);
+  let kb = select(bb | {{ signBit }}, (~bb) & {{ keyMask }}, (bb & {{ signBit }}) != 0u);
+  return ka < kb;
+}
+{% endif %}
+{% if headCache %}
+// Leading element of each written representative, kept beside `order`. The scan
+// below compares one slice against every representative, and reading those
+// leading elements out of X is a gather whose stride is the slice width, which
+// can be a kilobyte, so every candidate lands on its own cache line. Widening
+// the workgroup makes that worse rather than better, which is the tell: the
+// limit is cache capacity, not latency. Caching the leading elements densely
+// means the common case, where two slices already differ at element zero, never
+// touches X. The integer dedup path solves the same problem the same way, with
+// sampled slice fingerprints.
+var<workgroup> headCache: array<u32, {{ source.headCacheSlots }}>;
+{% endif %}
+{% if isFloat %}
+// Canonicalized leading-element bits: -0 folds to +0 exactly as `eq_value`
+// does, so equal head bits mean the leading elements compare equal.
+fn head_bits(k: u32) -> u32 {
+  if (params.outer == 0u || params.inner == 0u) { return 0u; }
+  let b = float_bits(slice_at(0u, k, 0u));
+  return select(b, 0u, b == {{ signBit }});
+}
+{% endif %}
+{% if isFloat %}
+// Whether `cand` could be the lower bound, judged from leading elements alone:
+// exactly `!slice_less(cand, k)` when the leading elements differ, and
+// optimistically true when they match, which `slice_cmp_coop` then settles.
+// Over-reporting only ever lowers the candidate, and a candidate that fails the
+// exact test is skipped and the scan resumes after it -- so the index this
+// converges on is the first truly qualifying one, which is what the serial
+// scan stops at.
+fn head_qualifies(hc: u32, hk: u32) -> bool {
+  if (is_nan_head(hc) || is_nan_head(hk)) { return true; }
+  if (hc == hk) { return true; }
+  return !less_head(hc, hk);
+}
+// Three-way slice comparison across the whole workgroup: 0 when the slices are
+// equivalent (identical, or decided at a coordinate holding a NaN, which
+// `slice_less` reports as neither-less), 1 when a < b, 2 when b < a.
+//
+// The serial `slice_less` walks a slice one element at a time on a single lane,
+// and its exit is data dependent, so a pair of EQUAL slices -- what a duplicate
+// input produces, and the reason this comparison exists -- runs the full slice
+// length at one memory latency per element, and that walk dominates the kernel.
+// Here every lane takes one coordinate and the first difference falls out of a
+// min-reduction, so an equal pair costs one round instead of one per element.
+fn slice_cmp_coop(a: u32, b: u32, lane: u32) -> u32 {
+  if (lane == 0u) { atomicStore(&diffAt, NO_HIT); }
+  workgroupBarrier();
+  let total = params.outer * params.inner;
+  for (var c = lane; c < total; c = c + WG) {
+    // A difference already found below this coordinate settles the comparison;
+    // this lane's remaining coordinates are all above it.
+    if (atomicLoad(&diffAt) < c) { break; }
+    let o = c / params.inner;
+    let n = c - o * params.inner;
+    let ba = float_bits(slice_at(o, a, n));
+    let bb = float_bits(slice_at(o, b, n));
+    if (is_nan_head(ba) || is_nan_head(bb)
+        || select(ba, 0u, ba == {{ signBit }}) != select(bb, 0u, bb == {{ signBit }})) {
+      atomicMin(&diffAt, c);
+    }
+  }
+  workgroupBarrier();
+  if (lane == 0u) {
+    let c = atomicLoad(&diffAt);
+    var verdict = 0u;
+    if (c != NO_HIT) {
+      let o = c / params.inner;
+      let n = c - o * params.inner;
+      let ba = float_bits(slice_at(o, a, n));
+      let bb = float_bits(slice_at(o, b, n));
+      if (!(is_nan_head(ba) || is_nan_head(bb))) {
+        verdict = select(2u, 1u,
+          less_head(select(ba, 0u, ba == {{ signBit }}), select(bb, 0u, bb == {{ signBit }})));
+      }
+    }
+    cmpOut = verdict;
+  }
+  return workgroupUniformLoad(&cmpOut);
+}
+{% endif %}
+fn slice_at(o: u32, k: u32, n: u32) -> {{ scalar }} {
+  return x[(o * params.axisDim + k) * params.inner + n];
+}
+// FTZ-safe element comparators use raw IEEE bits because float ALUs may flush
+// subnormals and collapse distinct rows. Signed zero is canonicalized before
+// applying an IEEE total-order key; non-float types use native ==/<.
+{% if isFloat %}
+{{ float_bits_def() }}
+{% endif %}
+{% if not isFloat %}
+{{ eq_value_def() }}
+{% if sorted %}
+{{ less_value_def() }}
+{% endif %}
+{% endif %}
+{% if floatSliceLess %}
+// Lexicographic comparison stops at the first unequal coordinate. If either
+// value there is NaN, neither slice is less and ordered insertion treats the
+// keys as equivalent.
+fn slice_less(a: u32, b: u32) -> bool {
+  for (var o = 0u; o < params.outer; o = o + 1u) {
+    let baseA = (o * params.axisDim + a) * params.inner;
+    let baseB = (o * params.axisDim + b) * params.inner;
+    for (var n = 0u; n < params.inner; n = n + 1u) {
+      // Same three tests as `is_nan_bits` / `eq_value` / `less_value`, over one
+      // pair of bit patterns instead of six independent `float_bits` calls --
+      // which on f16 is a vec2 pack and mask each time. Two equal slices, the
+      // case a duplicate input hits, walk to the end of the slice, so this loop
+      // body runs once per element and was the kernel's dominant cost.
+      let ba = float_bits(x[baseA + n]);
+      let bb = float_bits(x[baseB + n]);
+      if (is_nan_head(ba) || is_nan_head(bb)) { return false; }
+      // -0 folds to +0 exactly as `eq_value` does.
+      let ca = select(ba, 0u, ba == {{ signBit }});
+      let cb = select(bb, 0u, bb == {{ signBit }});
+      if (ca == cb) { continue; }
+      return less_head(ca, cb);
+    }
+  }
+  return false;
+}
+{% elif sorted and not isFloat %}
+// true iff slice a < slice b lexicographically over (outer, inner) order.
+fn slice_less(a: u32, b: u32) -> bool {
+  for (var o = 0u; o < params.outer; o = o + 1u) {
+    for (var n = 0u; n < params.inner; n = n + 1u) {
+      let va = slice_at(o, a, n);
+      let vb = slice_at(o, b, n);
+      if (less_value(va, vb)) { return true; }
+      if (less_value(vb, va)) { return false; }
+    }
+  }
+  return false;
+}
+{% endif %}
+{% if not isFloat %}
+// true iff slice a equals slice b element-wise.
+fn slice_eq(a: u32, b: u32) -> bool {
+  for (var o = 0u; o < params.outer; o = o + 1u) {
+    for (var n = 0u; n < params.inner; n = n + 1u) {
+      if (!eq_value(slice_at(o, a, n), slice_at(o, b, n))) { return false; }
+    }
+  }
+  return true;
+}
+{% endif %}
+{{ zero_value_def() }}
+{% if sorted and not isFloat %}
+fn sort_order(written: u32) {
+  for (var a = 1u; a < written; a = a + 1u) {
+    let key = order[a];
+    var b = a;
+    loop {
+      if (b == 0u) { break; }
+      if (!slice_less(key, order[b - 1u])) { break; }
+      order[b] = order[b - 1u];
+      b = b - 1u;
+    }
+    order[b] = key;
+  }
+}
+{% endif %}
+{% if isFloat and not sorted %}
+// The map is built in comparator order. Unsorted ONNX output instead uses
+// bucket insertion order, which is ascending representative input index.
+fn sort_by_first_occurrence(written: u32) {
+  for (var a = 1u; a < written; a = a + 1u) {
+    let key = order[a];
+    var b = a;
+    loop {
+      if (b == 0u || order[b - 1u] < key) { break; }
+      order[b] = order[b - 1u];
+      b = b - 1u;
+    }
+    order[b] = key;
+  }
+}
+{% endif %}
+{% if isFloat and (emitInverseIndices or emitCounts) %}
+// Find the bucket selected when input slice k was inserted. Representatives
+// with a later first occurrence did not exist yet and must not influence the
+// lower_bound result; this preserves stateful NaN equivalence after output is
+// reordered.
+fn find_bucket(k: u32, written: u32) -> u32 {
+  var candidate = written;
+  for (var p = 0u; p < written; p = p + 1u) {
+    let representative = order[p];
+    if (representative > k || slice_less(representative, k)) { continue; }
+    if (candidate == written || slice_less(representative, order[candidate])) {
+      candidate = p;
+    }
+  }
+  if (candidate < written && !slice_less(k, order[candidate])) {
+    return candidate;
+  }
+  return 0u;
+}
+{% endif %}
+@compute @workgroup_size(WG)
+fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
+{% if isFloat %}
+  // Reproduce ordered lower-bound insertion. NaN makes the comparator
+  // non-transitive, so pairwise equality and conventional hash tables cannot
+  // preserve its stateful bucket behavior.
+  //
+  // Every representative written so far is a lower-bound candidate, so the scan
+  // over them is the whole cost, and each step is a strided pair of slice loads
+  // that pays full memory latency. Scanning on one lane left the rest of the
+  // workgroup idle and issued those loads one at a time. Here the workgroup
+  // tests a chunk of candidates at once and keeps the SMALLEST index whose
+  // comparison fails -- the index the serial scan stops at, since it stops at
+  // the first failure too. Testing the remainder of a chunk is redundant, never
+  // observable: `slice_less` reads x and writes nothing.
+  //
+  // Both inner loops are barrier-free, so their bounds may depend on the live
+  // count; every barrier sits at the top level of this loop, whose bound comes
+  // from the uniform buffer, which is what WGSL's uniformity analysis requires.
+  var written = 0u;
+  for (var k = 0u; k < params.axisDim; k = k + 1u) {
+    let headK = head_bits(k);
+    // Walk the representatives for the first one that is not less than this
+    // slice. The head-only test is exact whenever the leading elements differ,
+    // so on real data this settles in a single round; a leading-element tie is
+    // resolved by one exact comparison, and the scan resumes past it only if
+    // that representative really does sort first.
+    var searchFrom = 0u;
+    var lowerBound = written;
+    var equivalent = false;
+    loop {
+      if (lid.x == 0u) { atomicStore(&lowerBoundHit, NO_HIT); }
+      workgroupBarrier();
+      for (var base = searchFrom; base < written; base = base + WG) {
+        // Chunks ascend, so an index already recorded is smaller than anything
+        // this chunk holds. Reading the atomic without a barrier can only see
+        // the hit late, which costs a redundant chunk and never a wrong bound.
+        if (atomicLoad(&lowerBoundHit) != NO_HIT) { break; }
+        let candidate = base + lid.x;
+        if (candidate >= searchFrom && candidate < written
+            && head_qualifies({{ headOf }}, headK)) {
+          atomicMin(&lowerBoundHit, candidate);
+        }
+      }
+      workgroupBarrier();
+      // Republished so the analysis accepts it as uniform: the exact comparison
+      // below carries barriers, so everything gating it has to be uniform.
+      if (lid.x == 0u) { hitIndex = atomicLoad(&lowerBoundHit); }
+      let hit = workgroupUniformLoad(&hitIndex);
+      if (hit == NO_HIT) {
+        // Nothing qualifies, so the serial scan would have run off the end.
+        lowerBound = written;
+        equivalent = false;
+        break;
+      }
+      let verdict = slice_cmp_coop(k, order[hit], lid.x);
+      if (verdict == 0u) { lowerBound = hit; equivalent = true; break; }
+      if (verdict == 1u) { lowerBound = hit; equivalent = false; break; }
+      // This representative sorts before the slice after all, so it is not the
+      // lower bound; resume the scan above it.
+      searchFrom = hit + 1u;
+    }
+    let insert = !equivalent && written < CAP;
+    // Open a slot at `lowerBound` by moving [lowerBound, written) up one place.
+    // Splitting that range into one contiguous run per lane keeps every move
+    // private except at a run's low end, which the run below overwrites -- so
+    // each lane preloads that single element and the whole shift needs one
+    // barrier instead of one per chunk.
+    let span = select(0u, written - lowerBound, insert);
+    let perLane = (span + WG - 1u) / WG;
+    let runLo = lowerBound + min(span, lid.x * perLane);
+    let runHi = lowerBound + min(span, (lid.x + 1u) * perLane);
+    var carry = 0u;
+{% if headCache %}
+    var carryHead = 0u;
+{% endif %}
+    if (runLo < runHi) {
+      carry = order[runLo];
+{% if headCache %}
+      carryHead = headCache[runLo];
+{% endif %}
+    }
+    storageBarrier();
+{% if headCache %}
+    workgroupBarrier();
+{% endif %}
+    if (runLo < runHi) {
+      // Top-down within the run, so a slot is read before it is written.
+      var p = runHi - 1u;
+      loop {
+        if (p <= runLo) { break; }
+        order[p + 1u] = order[p];
+{% if headCache %}
+        headCache[p + 1u] = headCache[p];
+{% endif %}
+        p = p - 1u;
+      }
+      order[runLo + 1u] = carry;
+{% if headCache %}
+      headCache[runLo + 1u] = carryHead;
+{% endif %}
+    }
+    if (insert && lid.x == 0u) {
+      order[lowerBound] = k;
+{% if headCache %}
+      headCache[lowerBound] = headK;
+{% endif %}
+    }
+    storageBarrier();
+{% if headCache %}
+    workgroupBarrier();
+{% endif %}
+    if (insert) { written = written + 1u; }
+  }
+{% if not sorted %}
+  if (lid.x == 0u) {
+    sort_by_first_occurrence(written);
+  }
+{% endif %}
+{% else %}
+  if (params.axisDim <= CAP) {
+    // When every input slice can be represented in the order scratch, test
+    // first-occurrence status independently. Duplicate-heavy inputs usually
+    // exit after only a handful of comparisons, and all lanes participate.
+    for (var k = lid.x; k < params.axisDim; k = k + WG) {
+      var seen = false;
+      for (var j = 0u; j < k; j = j + 1u) {
+        if (slice_eq(j, k)) { seen = true; break; }
+      }
+      order[k] = select(1u, 0u, seen);
+    }
+    storageBarrier();
+    if (lid.x == 0u) {
+      var written = 0u;
+      for (var k = 0u; k < params.axisDim; k = k + 1u) {
+        let first = order[k];
+        if (first != 0u) {
+          // Compacting in increasing k order preserves first-occurrence order.
+          // `written <= k`, so this never overwrites an unread flag.
+          order[written] = k;
+          written = written + 1u;
+        }
+      }
+{% if sorted %}
+      sort_order(written);
+{% endif %}
+    }
+  } else if (lid.x == 0u) {
+    // When axisDim exceeds the exact output capacity, discover slices serially
+    // while retaining one representative for every distinct result.
+    var written = 0u;
+    for (var k = 0u; k < params.axisDim; k = k + 1u) {
+      var seen = false;
+      for (var p = 0u; p < written; p = p + 1u) {
+        if (slice_eq(order[p], k)) { seen = true; break; }
+      }
+      if (!seen) {
+        if (written < CAP) {
+          order[written] = k;
+          written = written + 1u;
+        }
+      }
+    }
+{% if sorted %}
+    sort_order(written);
+{% endif %}
+  }
+{% endif %}
+  // `order` is storage-backed so legal data-dependent capacities are not
+  // constrained by maxComputeWorkgroupStorageSize.
+  storageBarrier();
+  // The public contract requires the caller-provided Y axis dimension to equal
+  // the exact unique count, so it is also the cross-lane written count.
+  let writtenCount = params.outputAxisDim;
+  // Scattering the selected slices is independent once `order` is ready, so
+  // spread both data and metadata stores across the workgroup.
+  let totalOut = params.outer * params.outputAxisDim * params.inner;
+  for (var out_i = lid.x; out_i < totalOut; out_i = out_i + WG) {
+    let n = out_i % params.inner;
+    let axisOuter = out_i / params.inner;
+    let p = axisOuter % params.outputAxisDim;
+    let o = axisOuter / params.outputAxisDim;
+    if (p < writtenCount) {
+      y[out_i] = slice_at(o, order[p], n);
+    } else {
+      y[out_i] = zero_value();
+    }
+  }
+{% if emitIndices %}
+  for (var p = lid.x; p < writtenCount; p = p + WG) {
+    indices[p] = order[p];
+  }
+{% endif %}
+{% if emitInverseIndices %}
+  for (var k = lid.x; k < params.axisDim; k = k + WG) {
+{% if isFloat %}
+    inverse_indices[k] = find_bucket(k, writtenCount);
+{% else %}
+    var uniqueIndex = 0u;
+    for (var p = 0u; p < writtenCount; p = p + 1u) {
+      if (slice_eq(order[p], k)) {
+        uniqueIndex = p;
+        break;
+      }
+    }
+    inverse_indices[k] = uniqueIndex;
+{% endif %}
+  }
+{% endif %}
+{% if emitCounts %}
+  for (var p = lid.x; p < writtenCount; p = p + WG) {
+    var occurrenceCount = 0u;
+    for (var k = 0u; k < params.axisDim; k = k + 1u) {
+{% if isFloat %}
+      if (find_bucket(k, writtenCount) == p) {
+{% else %}
+      if (slice_eq(order[p], k)) {
+{% endif %}
+        occurrenceCount = occurrenceCount + 1u;
+      }
+    }
+    counts[p] = occurrenceCount;
+  }
+{% endif %}
+}

build/webgpu/unique-compact-sort.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,177 @@

+{{ env.wgsl.resourceDeclarations }}
+// Pass 2 of parallel Unique: compact flagged first occurrences in appearance
+// order, optionally bitonic-sort them, and write the exact result. Small results
+// keep scratch in workgroup memory, while larger results use storage buffers
+// selected from device limits.
+const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
+const CAP: u32 = {{ source.capacity }}u;
+{% if sorted %}
+// The power-of-two padded tail sorts after every real value.
+const SORT_N: u32 = {{ source.sortN }}u;
+{% endif %}
+var<workgroup> wgScan: array<u32, WG>;
+var<workgroup> wgCarry: u32;
+{% if not source.globalScratch %}
+{% if sorted %}
+var<workgroup> sortKey: array<u32, SORT_N>;
+var<workgroup> sortVal: array<{{ scalar }}, SORT_N>;
+var<workgroup> sortPad: array<u32, SORT_N>;  // 1 = padding, 0 = real unique
+{% else %}
+var<workgroup> compacted: array<{{ scalar }}, CAP>;
+{% endif %}
+{% endif %}
+fn zero_value() -> {{ scalar }} {
+  return {{ scalar }}(0);
+}
+{% if sorted %}
+// Monotonic key matching unsigned or signed integer order. Global scratch
+// already stores raw bits.
+fn sort_key(v: {% if source.globalScratch %}u32{% else %}{{ scalar }}{% endif %}) -> u32 {
+{% if source.globalScratch %}
+  let b = v;
+{% else %}
+  let b = bitcast<u32>(v);
+{% endif %}
+{% if isUnsigned %}
+  return b;
+{% else %}
+  return b ^ 0x80000000u;
+{% endif %}
+}
+{% endif %}
+{% macro scratch_barrier() %}
+{% if source.globalScratch %}storageBarrier();{% else %}workgroupBarrier();{% endif %}
+{% endmacro %}
+@compute @workgroup_size(WG)
+fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
+  let tid = lid.x;
+  if (tid == 0u) {
+    wgCarry = 0u;
+  }
+  workgroupBarrier();
+  // Chunked Hillis-Steele exclusive scan maps flags to appearance-order ranks.
+  let chunks = (params.inputCount + WG - 1u) / WG;
+  for (var c = 0u; c < chunks; c = c + 1u) {
+    let i = c * WG + tid;
+    var f = 0u;
+    if (i < params.inputCount) {
+      f = flags[i];
+    }
+    wgScan[tid] = f;
+    workgroupBarrier();
+    var stride = 1u;
+    loop {
+      if (stride >= WG) { break; }
+      var add = 0u;
+      if (tid >= stride) {
+        add = wgScan[tid - stride];
+      }
+      workgroupBarrier();
+      wgScan[tid] = wgScan[tid] + add;
+      workgroupBarrier();
+      stride = stride * 2u;
+    }
+    let excl = wgScan[tid] - f;
+    let pos = wgCarry + excl;
+    if (i < params.inputCount && f == 1u && pos < CAP) {
+{% if source.globalScratch %}
+      sortVal[pos] = bitcast<u32>(x[i]);
+{% elif sorted %}
+      sortVal[pos] = x[i];
+{% else %}
+      compacted[pos] = x[i];
+{% endif %}
+    }
+    workgroupBarrier();
+    if (tid == 0u) {
+      wgCarry = wgCarry + wgScan[WG - 1u];
+    }
+    workgroupBarrier();
+  }
+  let total = wgCarry;
+  let written = min(total, CAP);
+{% if source.globalScratch %}
+  storageBarrier();
+{% endif %}
+{% if sorted %}
+  // Padding flag is the primary key, so even a real maximum-valued key stays
+  // before padding. Unique values are distinct, making the bit key a total order.
+  for (var k = tid; k < SORT_N; k = k + WG) {
+    let isPad = select(0u, 1u, k >= written);
+    sortPad[k] = isPad;
+    if (isPad == 0u) {
+      sortKey[k] = sort_key(sortVal[k]);
+    } else {
+{% if source.globalScratch %}
+      sortVal[k] = 0u;
+{% else %}
+      sortVal[k] = zero_value();
+{% endif %}
+      sortKey[k] = 0xffffffffu;
+    }
+  }
+  {{ scratch_barrier() }}
+  // Batcher bitonic network. Each lane owns one side of a compare-exchange.
+  var size = 2u;
+  loop {
+    if (size > SORT_N) { break; }
+    var stride = size / 2u;
+    loop {
+      if (stride == 0u) { break; }
+      for (var k = tid; k < SORT_N; k = k + WG) {
+        let partner = k ^ stride;
+        if (partner > k) {
+          let ascending = (k & size) == 0u;
+          let pk = sortPad[k];
+          let pp = sortPad[partner];
+          let kk = sortKey[k];
+          let kp = sortKey[partner];
+          let kBeforeP = (pk < pp) || (pk == pp && kk <= kp);
+          let needSwap = select(kBeforeP, !kBeforeP, ascending);
+          if (needSwap) {
+            let tKey = sortKey[k];
+            sortKey[k] = sortKey[partner];
+            sortKey[partner] = tKey;
+            let tVal = sortVal[k];
+            sortVal[k] = sortVal[partner];
+            sortVal[partner] = tVal;
+            let tPad = sortPad[k];
+            sortPad[k] = sortPad[partner];
+            sortPad[partner] = tPad;
+          }
+        }
+      }
+      {{ scratch_barrier() }}
+      stride = stride / 2u;
+    }
+    size = size * 2u;
+  }
+{% endif %}
+{% if not source.globalScratch %}
+  workgroupBarrier();
+{% endif %}
+  for (var k = tid; k < CAP; k = k + WG) {
+    if (k < written) {
+{% if source.globalScratch %}
+      y[k] = bitcast<{{ scalar }}>(sortVal[k]);
+{% elif sorted %}
+      y[k] = sortVal[k];
+{% else %}
+      y[k] = compacted[k];
+{% endif %}
+    } else {
+      y[k] = zero_value();
+    }
+  }
+}

build/webgpu/unique-dedup.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,54 @@

+{{ env.wgsl.resourceDeclarations }}
+// Pass 1 of parallel Unique: grid-parallel first-occurrence test. One thread per
+// input element, dispatched across all cores (ceil(n/WG) workgroups), so the
+// O(n^2) "have I seen this value before" scan runs in parallel rather than on a
+// single lane. Each workgroup cooperatively stages earlier values in a shared
+// tile: WG lanes consume each global load instead of independently rereading
+// the same prefix. flags[i] = 1 iff x[i] is the first occurrence of its value.
+//
+// This parallel equality path handles integer storage values only.
+const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
+var<workgroup> valueTile: array<{{ scalar }}, WG>;
+@compute @workgroup_size(WG)
+fn main(
+  @builtin(global_invocation_id) gid: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>,
+  @builtin(workgroup_id) workgroup: vec3<u32>,
+) {
+  let i = gid.x;
+  let laneActive = i < params.inputCount;
+  var value: {{ scalar }};
+  if (laneActive) {
+    value = x[i];
+  }
+  var seen = !laneActive;
+  // Every lane executes every barrier, including the inactive tail lanes of the
+  // final workgroup. groupEnd is uniform within a workgroup, so the tile loop is
+  // uniform as well. The current workgroup's tile is the last one any lane needs.
+  let groupEnd = min((workgroup.x + 1u) * WG, params.inputCount);
+  for (var tileBase = 0u; tileBase < groupEnd; tileBase += WG) {
+    let sourceIndex = tileBase + lid.x;
+    if (sourceIndex < params.inputCount) {
+      valueTile[lid.x] = x[sourceIndex];
+    }
+    workgroupBarrier();
+    if (laneActive && !seen) {
+      for (var k = 0u; k < WG; k += 1u) {
+        let earlier = tileBase + k;
+        if (earlier >= i) { break; }
+        if (valueTile[k] == value) {
+          seen = true;
+          break;
+        }
+      }
+    }
+    workgroupBarrier();
+  }
+  if (laneActive) {
+    flags[i] = select(0u, 1u, !seen);
+  }
+}

build/webgpu/unique-hash-build.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,61 @@

+{{ env.wgsl.resourceDeclarations }}
+// One thread per input element inserts its value into an open-addressing hash
+// table. Unsorted Unique also folds in each value's minimum input index. Sorted
+// Unique only needs the keys themselves; its final pass sorts their values, so
+// omitting tableIdx removes one atomic operation and one table-sized buffer.
+//
+// Race-freedom: the claim-and-publish of a slot's key is a single atomic
+// compareExchange (EMPTY -> key). There is no window where a slot is claimed but
+// its key is unpublished, so a concurrent probe either sees EMPTY (slot still
+// free, keep this slot via the CAS) or sees the final key value. Linear probing
+// on collision; the table is sized to >= 2*inputCount, so an empty slot always
+// exists and the probe terminates. Raw u32 value bits form the key. Integer
+// 0xffffffff cannot occupy a table slot because it is the EMPTY sentinel, so
+// its minimum index uses `special`.
+// Open-addressing table constants and key/hash helpers.
+const EMPTY: u32 = 0xffffffffu;
+const MASK: u32 = {{ source.tableSize }}u - 1u;
+fn key_bits(v: {{ scalar }}) -> u32 {
+{% if isUnsigned %}
+  return v;
+{% else %}
+  return bitcast<u32>(v);
+{% endif %}
+}
+// Integer hash (Wang-style mix) to spread keys across the table.
+fn hash_key(k: u32) -> u32 {
+  var x = k;
+  x = x ^ (x >> 16u);
+  x = x * 0x7feb352du;
+  x = x ^ (x >> 15u);
+  x = x * 0x846ca68bu;
+  x = x ^ (x >> 16u);
+  return x;
+}
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let i = gid.x;
+  if (i >= params.inputCount) {
+    return;
+  }
+  let k = key_bits(x[i]);
+  if (k == EMPTY) {
+    atomicMin(&special[0], i);
+    return;
+  }
+  var h = hash_key(k) & MASK;
+  loop {
+    let res = atomicCompareExchangeWeak(&tableKey[h], EMPTY, k);
+    if (res.exchanged || res.old_value == k) {
+{% if not source.keyOnly %}
+      // This slot now holds k (we just placed it, or it already held k).
+      atomicMin(&tableIdx[h], i);
+{% endif %}
+      break;
+    }
+    h = (h + 1u) & MASK;
+  }
+}

build/webgpu/unique-hash-collect.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,104 @@

+{% if source.useSubgroups %}
+enable subgroups;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+// Sorted Unique does not need appearance order before its final sort. Collect
+// occupied hash keys directly and reserve output positions once per workgroup.
+// This avoids serializing every distinct value through one global atomic.
+const EMPTY: u32 = 0xffffffffu;
+const TABLE_SIZE: u32 = {{ source.tableSize }}u;
+const CAP: u32 = {{ source.capacity }}u;
+{% if source.useSubgroups %}
+const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
+{% endif %}
+var<workgroup> groupBase: u32;
+{% if source.useSubgroups %}
+// The first subgroup lane writes its total. Lane zero converts these totals to
+// subgroup offsets and performs the workgroup's sole global count reservation.
+var<workgroup> subgroupOffsets: array<u32, WG>;
+{% else %}
+// Without subgroup operations, a workgroup-local atomic provides the same
+// compaction. Its contention never leaves shared memory and is bounded by WG.
+var<workgroup> localCount: atomic<u32>;
+{% endif %}
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>,
+        @builtin(local_invocation_id) lid: vec3<u32>{% if source.useSubgroups %},
+        @builtin(subgroup_invocation_id) sgLane: u32,
+        @builtin(subgroup_size) sgSize: u32{% endif %}) {
+  let h = gid.x;
+  var itemCount = 0u;
+  var item0 = 0u;
+  var item1 = 0u;
+  if (h < TABLE_SIZE) {
+    let key = tableKey[h];
+    if (key != EMPTY) {
+      item0 = key;
+      itemCount = 1u;
+    }
+    if (h == 0u) {
+      let specialIndex = atomicLoad(&special[0]);
+      if (specialIndex != EMPTY) {
+        let specialBits = bitcast<u32>(x[specialIndex]);
+        if (itemCount == 0u) {
+          item0 = specialBits;
+        } else {
+          item1 = specialBits;
+        }
+        itemCount += 1u;
+      }
+    }
+  }
+{% if source.useSubgroups %}
+  let subgroupPrefix = subgroupExclusiveAdd(itemCount);
+  let subgroupTotal = subgroupAdd(itemCount);
+  let subgroupId = lid.x / sgSize;
+  let subgroupCount = (WG + sgSize - 1u) / sgSize;
+  if (sgLane == 0u) {
+    subgroupOffsets[subgroupId] = subgroupTotal;
+  }
+  workgroupBarrier();
+  if (lid.x == 0u) {
+    var running = 0u;
+    for (var s = 0u; s < subgroupCount; s += 1u) {
+      let subgroupItems = subgroupOffsets[s];
+      subgroupOffsets[s] = running;
+      running += subgroupItems;
+    }
+    groupBase = 0u;
+    if (running != 0u) {
+      groupBase = atomicAdd(&count[0], running);
+    }
+  }
+  workgroupBarrier();
+  let localPrefix = subgroupOffsets[subgroupId] + subgroupPrefix;
+{% else %}
+  if (lid.x == 0u) {
+    atomicStore(&localCount, 0u);
+  }
+  workgroupBarrier();
+  let localPrefix = atomicAdd(&localCount, itemCount);
+  workgroupBarrier();
+  if (lid.x == 0u) {
+    let workgroupItems = atomicLoad(&localCount);
+    groupBase = 0u;
+    if (workgroupItems != 0u) {
+      groupBase = atomicAdd(&count[0], workgroupItems);
+    }
+  }
+  workgroupBarrier();
+{% endif %}
+  let outputIndex = groupBase + localPrefix;
+  if (itemCount >= 1u && outputIndex < CAP) {
+    compactedBits[outputIndex] = item0;
+  }
+  if (itemCount >= 2u && outputIndex + 1u < CAP) {
+    compactedBits[outputIndex + 1u] = item1;
+  }
+}

build/webgpu/unique-hash-init.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,34 @@

+{{ env.wgsl.resourceDeclarations }}
+// Initialize the Unique hash table. Sorted collection only needs the key table,
+// so that path clears four keys per invocation with ordinary vector stores.
+// The first-occurrence path keeps both atomic tables because its build pass
+// folds each key's minimum input index into tableIdx, and it leaves the output
+// counter alone -- only the sorted-collect path counts from zero here.
+const EMPTY: u32 = 0xffffffffu;
+const TABLE_SIZE: u32 = {{ source.tableSize }}u;
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{% if source.keyOnlyVec4 %}
+  let h4 = gid.x;
+  if (h4 >= TABLE_SIZE / 4u) {
+    return;
+  }
+  tableKey[h4] = vec4<u32>(EMPTY);
+  if (h4 == 0u) {
+    atomicStore(&special[0], EMPTY);
+    atomicStore(&count[0], 0u);
+  }
+{% else %}
+  let h = gid.x;
+  if (h >= TABLE_SIZE) {
+    return;
+  }
+  atomicStore(&tableKey[h], EMPTY);
+  atomicStore(&tableIdx[h], EMPTY);
+  if (h == 0u) {
+    atomicStore(&special[0], EMPTY);
+  }
+{% endif %}
+}

build/webgpu/unique-hash-mark.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,53 @@

+{{ env.wgsl.resourceDeclarations }}
+// Pass 3 of the hash-set parallel Unique (integer dtypes only). One thread per
+// input element re-probes the (now fully built) hash table for its value and
+// reads the stored minimum index. flags[i] = 1 iff i is that minimum, i.e. i is
+// the first occurrence of its value — exactly the predicate the O(n^2) scan
+// produced, but in O(n) probes. The slot is guaranteed to exist (this value was
+// inserted in the build pass), so the probe always finds a matching key before
+// hitting an EMPTY slot.
+// Open-addressing table constants and key/hash helpers.
+const EMPTY: u32 = 0xffffffffu;
+const MASK: u32 = {{ source.tableSize }}u - 1u;
+fn key_bits(v: {{ scalar }}) -> u32 {
+{% if isUnsigned %}
+  return v;
+{% else %}
+  return bitcast<u32>(v);
+{% endif %}
+}
+fn hash_key(k: u32) -> u32 {
+  var x = k;
+  x = x ^ (x >> 16u);
+  x = x * 0x7feb352du;
+  x = x ^ (x >> 15u);
+  x = x * 0x846ca68bu;
+  x = x ^ (x >> 16u);
+  return x;
+}
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let i = gid.x;
+  if (i >= params.inputCount) {
+    return;
+  }
+  let k = key_bits(x[i]);
+  if (k == EMPTY) {
+    flags[i] = select(0u, 1u, atomicLoad(&special[0]) == i);
+    return;
+  }
+  var h = hash_key(k) & MASK;
+  loop {
+    let stored = atomicLoad(&tableKey[h]);
+    if (stored == k) {
+      flags[i] = select(0u, 1u, atomicLoad(&tableIdx[h]) == i);
+      return;
+    }
+    // stored == EMPTY would mean the value is absent, which is impossible here
+    // (it was inserted in the build pass); keep probing on any other key.
+    h = (h + 1u) & MASK;
+  }
+}

build/webgpu/unique-hash-sort-collected-key-only.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,83 @@

+{{ env.wgsl.resourceDeclarations }}
+const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
+const CAP: u32 = {{ source.capacity }}u;
+const SORT_N: u32 = {{ source.sortN }}u;
+// The monotonic key transform is invertible, so the workgroup sort only needs
+// one u32 per slot. Padding uses the maximum key; if a real maximum key exists,
+// it decodes to the same value, so no separate padding flag/value arrays are
+// needed. SORT_N is specialized so this array fits the device's workgroup-
+// storage limit.
+var<workgroup> keys: array<u32, SORT_N>;
+fn key_from_bits(bits: u32) -> u32 {
+{% if isUnsigned %}
+  return bits;
+{% else %}
+  return bits ^ 0x80000000u;
+{% endif %}
+}
+fn bits_from_key(key: u32) -> u32 {
+{% if isUnsigned %}
+  return key;
+{% else %}
+  return key ^ 0x80000000u;
+{% endif %}
+}
+fn zero_value() -> {{ scalar }} {
+  return {{ scalar }}(0);
+}
+@compute @workgroup_size(WG)
+fn main(@builtin(local_invocation_id) local: vec3<u32>) {
+  let tid = local.x;
+  let written = min(count[0], CAP);
+  for (var index = tid; index < SORT_N; index = index + WG) {
+    var key = 0xffffffffu;
+    if (index < written) {
+      key = key_from_bits(compactedBits[index]);
+    }
+    keys[index] = key;
+  }
+  workgroupBarrier();
+  var size = 2u;
+  loop {
+    if (size > SORT_N) {
+      break;
+    }
+    var stride = size / 2u;
+    loop {
+      if (stride == 0u) {
+        break;
+      }
+      for (var index = tid; index < SORT_N; index = index + WG) {
+        let partner = index ^ stride;
+        if (partner > index) {
+          let a = keys[index];
+          let b = keys[partner];
+          let ascending = (index & size) == 0u;
+          let swap = select((a < b), (a > b), ascending);
+          if (swap) {
+            keys[index] = b;
+            keys[partner] = a;
+          }
+        }
+      }
+      workgroupBarrier();
+      stride = stride / 2u;
+    }
+    size = size * 2u;
+  }
+  for (var index = tid; index < CAP; index = index + WG) {
+    if (index < written) {
+      y[index] = bitcast<{{ scalar }}>(bits_from_key(keys[index]));
+    } else {
+      y[index] = zero_value();
+    }
+  }
+}

build/webgpu/unique.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,194 @@

+// IEEE total-order comparators for Unique. Floating-point equality and ordering
+// use raw bits because GPUs may flush subnormals in float comparisons, which
+// would otherwise collapse distinct values. Equality canonicalizes -0 to +0;
+// ordering uses a monotonic u32 key over the real line. Integers use == and <.
+{% macro float_bits_def() %}
+fn float_bits(v: {{ scalar }}) -> u32 {
+{% if usesF16 %}
+  // WGSL has no scalar u16 type. Packing v into the low component preserves
+  // its binary16 representation while producing a bitcast-compatible 32 bits.
+  return bitcast<u32>(vec2<f16>(v, 0.0h)) & 0xffffu;
+{% else %}
+  return bitcast<u32>(v);
+{% endif %}
+}
+{%- endmacro -%}
+{% macro eq_value_def() %}
+fn eq_value(a: {{ scalar }}, b: {{ scalar }}) -> bool {
+{% if isFloat %}
+  var ba = float_bits(a);
+  var bb = float_bits(b);
+{% if usesF16 %}
+  if (ba == 0x8000u) { ba = 0u; }  // -0 -> +0
+  if (bb == 0x8000u) { bb = 0u; }
+{% else %}
+  if (ba == 0x80000000u) { ba = 0u; }  // -0 -> +0
+  if (bb == 0x80000000u) { bb = 0u; }
+{% endif %}
+  return ba == bb;
+{% else %}
+  return a == b;
+{% endif %}
+}
+{%- endmacro -%}
+{%- macro less_value_def() %}
+fn less_value(a: {{ scalar }}, b: {{ scalar }}) -> bool {
+{% if isFloat %}
+  let ba = float_bits(a);
+  let bb = float_bits(b);
+{% if usesF16 %}
+  let ka = select(ba | 0x8000u, (~ba) & 0xffffu, (ba & 0x8000u) != 0u);
+  let kb = select(bb | 0x8000u, (~bb) & 0xffffu, (bb & 0x8000u) != 0u);
+{% else %}
+  let ka = select(ba | 0x80000000u, ~ba, (ba & 0x80000000u) != 0u);
+  let kb = select(bb | 0x80000000u, ~bb, (bb & 0x80000000u) != 0u);
+{% endif %}
+  return ka < kb;
+{% else %}
+  return a < b;
+{% endif %}
+}
+{%- endmacro -%}
+{%- macro zero_value_def() %}
+fn zero_value() -> {{ scalar }} {
+  return {{ scalar }}(0);
+}
+{%- endmacro -%}
+{%- macro is_nan_bits_def() %}
+fn is_nan_bits(v: {{ scalar }}) -> bool {
+  let b = float_bits(v);
+{% if usesF16 %}
+  return (b & 0x7c00u) == 0x7c00u && (b & 0x03ffu) != 0u;
+{% else %}
+  return (b & 0x7f800000u) == 0x7f800000u && (b & 0x007fffffu) != 0u;
+{% endif %}
+}
+{%- endmacro %}
+{% set emitIndices = source.hasIndices | default(false) %}
+{% set emitInverseIndices = source.hasInverseIndices | default(false) %}
+{% set emitCounts = source.hasCounts | default(false) %}
+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+// Floating-point equality/order use the raw IEEE bit pattern, not the float
+// ALU: GPUs may flush subnormals to zero in comparisons and collapse distinct
+// values. Equality canonicalizes -0 to +0; ordering uses an IEEE total-order
+// key. Integers use ==/< with their native signedness.
+{% if isFloat %}
+{{ float_bits_def() }}
+{% endif %}
+{{ eq_value_def() }}
+{{ less_value_def() }}
+{{ zero_value_def() }}
+{% if isFloat %}
+{{ is_nan_bits_def() }}
+{% endif %}
+// NaN is unequal to every value, but neither comparison is less, so ordered
+// insertion treats it as equivalent to the current candidate. The candidate
+// therefore depends on insertion history.
+fn ordered_less(a: {{ scalar }}, b: {{ scalar }}) -> bool {
+{% if isFloat %}
+  if (is_nan_bits(a) || is_nan_bits(b)) { return false; }
+{% endif %}
+  if (eq_value(a, b)) { return false; }
+  return less_value(a, b);
+}
+{% if emitIndices or emitInverseIndices or emitCounts %}
+fn representative_input_index(output_i: u32) -> u32 {
+  let representative = y[output_i];
+  for (var input_i = 0u; input_i < params.inputCount; input_i = input_i + 1u) {
+    if (eq_value(x[input_i], representative)) { return input_i; }
+  }
+  return params.inputCount;
+}
+{% endif %}
+{% if emitInverseIndices or emitCounts %}
+// Resolve the bucket selected when input_i was inserted. A representative that
+// first appeared later did not exist at that point and cannot change the
+// stateful lower_bound result.
+fn find_bucket(input_i: u32, written: u32) -> u32 {
+  let value = x[input_i];
+  var candidate = written;
+  for (var output_i = 0u; output_i < written; output_i = output_i + 1u) {
+    if (representative_input_index(output_i) > input_i || ordered_less(y[output_i], value)) { continue; }
+    if (candidate == written || ordered_less(y[output_i], y[candidate])) {
+      candidate = output_i;
+    }
+  }
+  if (candidate < written && !ordered_less(value, y[candidate])) {
+    return candidate;
+  }
+  return 0u;
+}
+{% endif %}
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  if (gid.x != 0u) { return; }
+  var written = 0u;
+  for (var i = 0u; i < params.inputCount; i = i + 1u) {
+    let value = x[i];
+    var candidate = written;
+    for (var output_i = 0u; output_i < written; output_i = output_i + 1u) {
+      if (!ordered_less(y[output_i], value)) {
+        if (candidate == written || ordered_less(y[output_i], y[candidate])) {
+          candidate = output_i;
+        }
+      }
+    }
+    let seen = candidate < written && !ordered_less(value, y[candidate]);
+    if (!seen) {
+      if (written < params.capacity) {
+        y[written] = value;
+        written = written + 1u;
+      }
+    }
+  }
+  {% if sorted %}
+  for (var i = 0u; i < written; i = i + 1u) {
+    for (var j = i + 1u; j < written; j = j + 1u) {
+      if (ordered_less(y[j], y[i])) {
+        let tmp = y[i];
+        y[i] = y[j];
+        y[j] = tmp;
+      }
+    }
+  }
+  {% endif %}
+  for (var i = written; i < params.capacity; i = i + 1u) {
+    y[i] = zero_value();
+  }
+{% if emitIndices %}
+  for (var unique_i = 0u; unique_i < params.capacity; unique_i = unique_i + 1u) {
+    indices[unique_i] = representative_input_index(unique_i);
+  }
+{% endif %}
+{% if emitCounts %}
+  for (var unique_i = 0u; unique_i < params.capacity; unique_i = unique_i + 1u) {
+    counts[unique_i] = 0u;
+  }
+{% endif %}
+{% if emitInverseIndices or emitCounts %}
+  // Resolve metadata against final Y order while retaining the prefix state
+  // that selected each NaN-equivalent bucket.
+  for (var input_i = 0u; input_i < params.inputCount; input_i = input_i + 1u) {
+    let unique_i = find_bucket(input_i, written);
+{% if emitInverseIndices %}
+    inverse_indices[input_i] = unique_i;
+{% endif %}
+{% if emitCounts %}
+    counts[unique_i] = counts[unique_i] + 1u;
+{% endif %}
+  }
+{% endif %}
+}