sync 6fdf6301e2bb

Browse files

Files changed (7) hide show

README.md +10 -2
build/webgpu/bench.json +264 -0
build/webgpu/manifest.json +66 -16
build/webgpu/metadata.json +14 -8
build/webgpu/space-depth-collapsed-input.wgsl.jinja +45 -0
build/webgpu/space-depth-permute.wgsl.jinja +25 -32
build/webgpu/test.json +273 -0

README.md CHANGED Viewed

@@ -42,18 +42,26 @@ Attributes and default values (overridable per request):
 | --- | --- |
 | `T` | `float32`, `float16`, `int32`, `int16`, `int8`, `uint32`, `uint8`, `bool` |
 ## Files
 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
-- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 - [`space-depth-permute.wgsl.jinja`](build/webgpu/space-depth-permute.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
-npm install --save-exact @huggingface/kernels@0.0.1-preview.2
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

 | --- | --- |
 | `T` | `float32`, `float16`, `int32`, `int16`, `int8`, `uint32`, `uint8`, `bool` |
+## Implementation variants
+One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
+- `flat_vector` — Groups four flat outputs per invocation, including groups that cross spatial rows or channels. Reuses the shared permutation coordinates and writes one contiguous vector.
+- `collapsed_input_vector` — When one spatial block spans the input plane, reads contiguous vec4 words and writes channel streams. Spatial tiles balance the available channel workgroups; paired f16 channels give each invocation whole output words.
 ## Files
 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark cases
+- [`space-depth-collapsed-input.wgsl.jinja`](build/webgpu/space-depth-collapsed-input.wgsl.jinja)
 - [`space-depth-permute.wgsl.jinja`](build/webgpu/space-depth-permute.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.3
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

build/webgpu/bench.json CHANGED Viewed

@@ -43,6 +43,270 @@
       },
       "outputs": { "output": { "dtype": "float32", "shape": [1, 512, 1024, 3] } },
       "bench": { "metrics": [{ "type": "bandwidth", "value": "args.count * dtypeBytes(args.dtype) * 2" }] }
     }
   ]
 }

       },
       "outputs": { "output": { "dtype": "float32", "shape": [1, 512, 1024, 3] } },
       "bench": { "metrics": [{ "type": "bandwidth", "value": "args.count * dtypeBytes(args.dtype) * 2" }] }
+    },
+    {
+      "name": "collapsed_float32_n2_c65_b2",
+      "attrs": { "blocksize": 2 },
+      "inputs": { "input": { "dtype": "float32", "shape": [2, 65, 2, 2], "dist": "normal", "seed": 1827, "scale": 1 } },
+      "outputs": { "output": { "dtype": "float32", "shape": [2, 260, 1, 1] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "4160" }] }
+    },
+    {
+      "name": "collapsed_float32_n2_c66_b2",
+      "attrs": { "blocksize": 2 },
+      "inputs": { "input": { "dtype": "float32", "shape": [2, 66, 2, 2], "dist": "normal", "seed": 1827, "scale": 1 } },
+      "outputs": { "output": { "dtype": "float32", "shape": [2, 264, 1, 1] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "4224" }] }
+    },
+    {
+      "name": "collapsed_float32_n1_c65537_b4",
+      "attrs": { "blocksize": 4 },
+      "inputs": {
+        "input": { "dtype": "float32", "shape": [1, 65537, 4, 4], "dist": "normal", "seed": 1827, "scale": 1 }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [1, 1048592, 1, 1] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "8388736" }] }
+    },
+    {
+      "name": "collapsed_float32_n1_c65536_b6",
+      "attrs": { "blocksize": 6 },
+      "inputs": {
+        "input": { "dtype": "float32", "shape": [1, 65536, 6, 6], "dist": "normal", "seed": 1827, "scale": 1 }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [1, 2359296, 1, 1] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "18874368" }] }
+    },
+    {
+      "name": "collapsed_float32_n1_c65536_b3",
+      "attrs": { "blocksize": 3 },
+      "inputs": {
+        "input": { "dtype": "float32", "shape": [1, 65536, 3, 3], "dist": "normal", "seed": 1827, "scale": 1 }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [1, 589824, 1, 1] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "4718592" }] }
+    },
+    {
+      "name": "collapsed_float32_n1_c262143_b2",
+      "attrs": { "blocksize": 2 },
+      "inputs": {
+        "input": { "dtype": "float32", "shape": [1, 262143, 2, 2], "dist": "normal", "seed": 1827, "scale": 1 }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [1, 1048572, 1, 1] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "8388576" }] }
+    },
+    {
+      "name": "collapsed_float32_n1_c262144_b2",
+      "attrs": { "blocksize": 2 },
+      "inputs": {
+        "input": { "dtype": "float32", "shape": [1, 262144, 2, 2], "dist": "normal", "seed": 1827, "scale": 1 }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [1, 1048576, 1, 1] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "8388608" }] }
+    },
+    {
+      "name": "flat_float32_n2_c65_6x10_b2",
+      "attrs": { "blocksize": 2 },
+      "inputs": { "input": { "dtype": "float32", "shape": [2, 65, 6, 10], "dist": "normal", "seed": 1828, "scale": 1 } },
+      "outputs": { "output": { "dtype": "float32", "shape": [2, 260, 3, 5] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "62400" }] }
+    },
+    {
+      "name": "flat_float32_n1_c64_1024x10_b2",
+      "attrs": { "blocksize": 2 },
+      "inputs": {
+        "input": { "dtype": "float32", "shape": [1, 64, 1024, 10], "dist": "normal", "seed": 1828, "scale": 1 }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [1, 256, 512, 5] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "5242880" }] }
+    },
+    {
+      "name": "flat_float32_n2_c64_12x15_b3",
+      "attrs": { "blocksize": 3 },
+      "inputs": {
+        "input": { "dtype": "float32", "shape": [2, 64, 12, 15], "dist": "normal", "seed": 1828, "scale": 1 }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [2, 576, 4, 5] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "184320" }] }
+    },
+    {
+      "name": "collapsed_float16_n2_c65_b2",
+      "attrs": { "blocksize": 2 },
+      "inputs": { "input": { "dtype": "float16", "shape": [2, 65, 2, 2], "dist": "normal", "seed": 1827, "scale": 1 } },
+      "outputs": { "output": { "dtype": "float16", "shape": [2, 260, 1, 1] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "2080" }] }
+    },
+    {
+      "name": "collapsed_float16_n2_c66_b2",
+      "attrs": { "blocksize": 2 },
+      "inputs": { "input": { "dtype": "float16", "shape": [2, 66, 2, 2], "dist": "normal", "seed": 1827, "scale": 1 } },
+      "outputs": { "output": { "dtype": "float16", "shape": [2, 264, 1, 1] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "2112" }] }
+    },
+    {
+      "name": "collapsed_float16_n1_c65537_b4",
+      "attrs": { "blocksize": 4 },
+      "inputs": {
+        "input": { "dtype": "float16", "shape": [1, 65537, 4, 4], "dist": "normal", "seed": 1827, "scale": 1 }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [1, 1048592, 1, 1] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "4194368" }] }
+    },
+    {
+      "name": "collapsed_float16_n1_c65536_b6",
+      "attrs": { "blocksize": 6 },
+      "inputs": {
+        "input": { "dtype": "float16", "shape": [1, 65536, 6, 6], "dist": "normal", "seed": 1827, "scale": 1 }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [1, 2359296, 1, 1] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "9437184" }] }
+    },
+    {
+      "name": "collapsed_float16_n1_c65536_b3",
+      "attrs": { "blocksize": 3 },
+      "inputs": {
+        "input": { "dtype": "float16", "shape": [1, 65536, 3, 3], "dist": "normal", "seed": 1827, "scale": 1 }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [1, 589824, 1, 1] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "2359296" }] }
+    },
+    {
+      "name": "collapsed_float16_n1_c262143_b2",
+      "attrs": { "blocksize": 2 },
+      "inputs": {
+        "input": { "dtype": "float16", "shape": [1, 262143, 2, 2], "dist": "normal", "seed": 1827, "scale": 1 }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [1, 1048572, 1, 1] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "4194288" }] }
+    },
+    {
+      "name": "collapsed_float16_n1_c262144_b2",
+      "attrs": { "blocksize": 2 },
+      "inputs": {
+        "input": { "dtype": "float16", "shape": [1, 262144, 2, 2], "dist": "normal", "seed": 1827, "scale": 1 }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [1, 1048576, 1, 1] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "4194304" }] }
+    },
+    {
+      "name": "flat_float16_n2_c65_6x10_b2",
+      "attrs": { "blocksize": 2 },
+      "inputs": { "input": { "dtype": "float16", "shape": [2, 65, 6, 10], "dist": "normal", "seed": 1828, "scale": 1 } },
+      "outputs": { "output": { "dtype": "float16", "shape": [2, 260, 3, 5] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "31200" }] }
+    },
+    {
+      "name": "flat_float16_n1_c64_1024x10_b2",
+      "attrs": { "blocksize": 2 },
+      "inputs": {
+        "input": { "dtype": "float16", "shape": [1, 64, 1024, 10], "dist": "normal", "seed": 1828, "scale": 1 }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [1, 256, 512, 5] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "2621440" }] }
+    },
+    {
+      "name": "flat_float16_n2_c64_12x15_b3",
+      "attrs": { "blocksize": 3 },
+      "inputs": {
+        "input": { "dtype": "float16", "shape": [2, 64, 12, 15], "dist": "normal", "seed": 1828, "scale": 1 }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [2, 576, 4, 5] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "92160" }] }
+    },
+    {
+      "name": "spatial_parallel_float32_n1_c1_b64",
+      "attrs": { "blocksize": 64 },
+      "inputs": { "input": { "dtype": "float32", "shape": [1, 1, 64, 64], "dist": "normal", "seed": 1891, "scale": 1 } },
+      "outputs": { "output": { "dtype": "float32", "shape": [1, 4096, 1, 1] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "32768" }] }
+    },
+    {
+      "name": "spatial_parallel_float32_n2_c3_b64",
+      "attrs": { "blocksize": 64 },
+      "inputs": { "input": { "dtype": "float32", "shape": [2, 3, 64, 64], "dist": "normal", "seed": 1891, "scale": 1 } },
+      "outputs": { "output": { "dtype": "float32", "shape": [2, 12288, 1, 1] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "196608" }] }
+    },
+    {
+      "name": "spatial_parallel_float32_n1_c256_b32",
+      "attrs": { "blocksize": 32 },
+      "inputs": {
+        "input": { "dtype": "float32", "shape": [1, 256, 32, 32], "dist": "normal", "seed": 1891, "scale": 1 }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [1, 262144, 1, 1] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "2097152" }] }
+    },
+    {
+      "name": "spatial_parallel_float32_n1_c16_b64",
+      "attrs": { "blocksize": 64 },
+      "inputs": {
+        "input": { "dtype": "float32", "shape": [1, 16, 64, 64], "dist": "normal", "seed": 1891, "scale": 1 }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [1, 65536, 1, 1] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "524288" }] }
+    },
+    {
+      "name": "spatial_parallel_float16_n1_c1_b64",
+      "attrs": { "blocksize": 64 },
+      "inputs": { "input": { "dtype": "float16", "shape": [1, 1, 64, 64], "dist": "normal", "seed": 1891, "scale": 1 } },
+      "outputs": { "output": { "dtype": "float16", "shape": [1, 4096, 1, 1] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "16384" }] }
+    },
+    {
+      "name": "spatial_parallel_float16_n2_c3_b64",
+      "attrs": { "blocksize": 64 },
+      "inputs": { "input": { "dtype": "float16", "shape": [2, 3, 64, 64], "dist": "normal", "seed": 1891, "scale": 1 } },
+      "outputs": { "output": { "dtype": "float16", "shape": [2, 12288, 1, 1] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "98304" }] }
+    },
+    {
+      "name": "spatial_parallel_float16_n1_c256_b32",
+      "attrs": { "blocksize": 32 },
+      "inputs": {
+        "input": { "dtype": "float16", "shape": [1, 256, 32, 32], "dist": "normal", "seed": 1891, "scale": 1 }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [1, 262144, 1, 1] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "1048576" }] }
+    },
+    {
+      "name": "spatial_parallel_float16_n1_c16_b64",
+      "attrs": { "blocksize": 64 },
+      "inputs": {
+        "input": { "dtype": "float16", "shape": [1, 16, 64, 64], "dist": "normal", "seed": 1891, "scale": 1 }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [1, 65536, 1, 1] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "262144" }] }
+    },
+    {
+      "name": "spatial_partial_float32_n2_c129_b4",
+      "attrs": { "blocksize": 4 },
+      "tunables": { "WORKGROUP_SIZE": 64 },
+      "inputs": { "input": { "dtype": "float32", "shape": [2, 129, 4, 4], "dist": "normal", "seed": 1892, "scale": 1 } },
+      "outputs": { "output": { "dtype": "float32", "shape": [2, 2064, 1, 1] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "33024" }] }
+    },
+    {
+      "name": "spatial_partial_float32_n2_c257_b6",
+      "attrs": { "blocksize": 6 },
+      "tunables": { "WORKGROUP_SIZE": 64 },
+      "inputs": { "input": { "dtype": "float32", "shape": [2, 257, 6, 6], "dist": "normal", "seed": 1892, "scale": 1 } },
+      "outputs": { "output": { "dtype": "float32", "shape": [2, 9252, 1, 1] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "148032" }] }
+    },
+    {
+      "name": "spatial_partial_float16_n2_c258_b4",
+      "attrs": { "blocksize": 4 },
+      "tunables": { "WORKGROUP_SIZE": 64 },
+      "inputs": { "input": { "dtype": "float16", "shape": [2, 258, 4, 4], "dist": "normal", "seed": 1892, "scale": 1 } },
+      "outputs": { "output": { "dtype": "float16", "shape": [2, 4128, 1, 1] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "33024" }] }
+    },
+    {
+      "name": "spatial_partial_float16_n2_c514_b6",
+      "attrs": { "blocksize": 6 },
+      "tunables": { "WORKGROUP_SIZE": 64 },
+      "inputs": { "input": { "dtype": "float16", "shape": [2, 514, 6, 6], "dist": "normal", "seed": 1892, "scale": 1 } },
+      "outputs": { "output": { "dtype": "float16", "shape": [2, 18504, 1, 1] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "148032" }] }
     }
   ]
 }

build/webgpu/manifest.json CHANGED Viewed

@@ -15,17 +15,79 @@
   "typeConstraints": { "T": ["float32", "float16", "int32", "int16", "int8", "uint32", "uint8", "bool"] },
   "tunables": { "WORKGROUP_SIZE": { "default": 256 } },
   "derive": {
     "shapeContract": "attrs.blocksize > 0 and ranks.input == 4 and ranks.output == 4 and dim(shapes.output, 0) == dim(shapes.input, 0) and dim(shapes.output, 1) == dim(shapes.input, 1) * attrs.blocksize * attrs.blocksize and dim(shapes.output, 2) * attrs.blocksize == dim(shapes.input, 2) and dim(shapes.output, 3) * attrs.blocksize == dim(shapes.input, 3) and f16Ok(dtypes.T)",
     "reportedWideExecution": "has(device.adapterInfo, \"subgroupMinSize\") and device.adapterInfo.subgroupMinSize >= 32",
-    "f16VectorPermutationPreferred": "dtypes.T != \"f16\" or device.features.has(\"subgroups\") or reportedWideExecution"
   },
   "when": ["shapeContract"],
-  "bindings": { "output": { "buffer": "storage", "elementType": "$outputElement" } },
   "variants": [
     {
       "id": "nchw_vec4",
       "priority": 20,
       "when": ["dim(shapes.output, 3) % 4 == 0", "f16VectorPermutationPreferred"],
       "derive": { "outputElement": "\"vec4<\" ~ dtypes.T ~ \">\"", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
@@ -33,17 +95,11 @@
           "name": "SpaceToDepth",
           "shader": "space-depth-permute.wgsl.jinja",
           "derive": {
-            "direction": "\"spaceToDepth\"",
-            "modeSpec": "\"DCR\"",
             "vectorizedSpec": true,
-            "blocksize": "attrs.blocksize",
             "count": "numel(shapes.output) / 4",
             "outWUnits": "dim(shapes.output, 3) / 4",
             "outH": "dim(shapes.output, 2)",
-            "outC": "dim(shapes.output, 1)",
-            "inC": "dim(shapes.input, 1)",
-            "inH": "dim(shapes.input, 2)",
-            "inW": "dim(shapes.input, 3)"
           },
           "bindings": ["input", "output"],
           "dispatch": {
@@ -64,17 +120,11 @@
           "name": "SpaceToDepth",
           "shader": "space-depth-permute.wgsl.jinja",
           "derive": {
-            "direction": "\"spaceToDepth\"",
-            "modeSpec": "\"DCR\"",
             "vectorizedSpec": false,
-            "blocksize": "attrs.blocksize",
             "count": "numel(shapes.output)",
             "outWUnits": "dim(shapes.output, 3)",
             "outH": "dim(shapes.output, 2)",
-            "outC": "dim(shapes.output, 1)",
-            "inC": "dim(shapes.input, 1)",
-            "inH": "dim(shapes.input, 2)",
-            "inW": "dim(shapes.input, 3)"
           },
           "bindings": ["input", "output"],
           "dispatch": {

   "typeConstraints": { "T": ["float32", "float16", "int32", "int16", "int8", "uint32", "uint8", "bool"] },
   "tunables": { "WORKGROUP_SIZE": { "default": 256 } },
   "derive": {
+    "variableSubgroup16To32": "device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 16 and device.adapterInfo.subgroupMaxSize == 32",
     "shapeContract": "attrs.blocksize > 0 and ranks.input == 4 and ranks.output == 4 and dim(shapes.output, 0) == dim(shapes.input, 0) and dim(shapes.output, 1) == dim(shapes.input, 1) * attrs.blocksize * attrs.blocksize and dim(shapes.output, 2) * attrs.blocksize == dim(shapes.input, 2) and dim(shapes.output, 3) * attrs.blocksize == dim(shapes.input, 3) and f16Ok(dtypes.T)",
     "reportedWideExecution": "has(device.adapterInfo, \"subgroupMinSize\") and device.adapterInfo.subgroupMinSize >= 32",
+    "f16VectorPermutationPreferred": "dtypes.T != \"f16\" or device.features.has(\"subgroups\") or reportedWideExecution",
+    "blocksize": "attrs.blocksize",
+    "inC": "dim(shapes.input, 1)",
+    "inH": "dim(shapes.input, 2)",
+    "inW": "dim(shapes.input, 3)"
   },
   "when": ["shapeContract"],
+  "bindings": { "output": { "elementType": "$outputElement" } },
   "variants": [
+    {
+      "id": "flat_vector",
+      "priority": 25,
+      "when": ["numel(shapes.output) % 4 == 0", "dim(shapes.output, 3) % 4 != 0", "f16VectorPermutationPreferred"],
+      "derive": {
+        "outputElement": "\"vec4<\" ~ dtypes.T ~ \">\"",
+        "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
+        "count": "numel(shapes.output) / 4",
+        "outH": "dim(shapes.output, 2)",
+        "outC": "dim(shapes.output, 1)",
+        "vectorizedSpec": false,
+        "flatVectorSpec": true,
+        "outWUnits": "dim(shapes.output, 3)"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "shader": "space-depth-permute.wgsl.jinja",
+          "bindings": ["input", "output"],
+          "dispatch": {
+            "x": "min(ceilDiv((count), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((count), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          },
+          "name": "SpaceToDepth.FlatVector"
+        }
+      ]
+    },
+    {
+      "id": "collapsed_input_vector",
+      "priority": 30,
+      "when": ["dim(shapes.output, 2) == 1", "dim(shapes.output, 3) == 1", "attrs.blocksize % 2 == 0", "dtypes.T != \"f16\" or dim(shapes.input, 1) % 2 == 0"],
+      "derive": {
+        "channelPack": "2 if dtypes.T == \"f16\" else 1",
+        "outputElement": "\"vec2<f16>\" if dtypes.T == \"f16\" else dtypes.T",
+        "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
+        "channelCount": "dim(shapes.input, 0) * dim(shapes.input, 1) / channelPack",
+        "blockArea": "attrs.blocksize * attrs.blocksize",
+        "inChannels": "dim(shapes.input, 1) / channelPack",
+        "wordTile": "min(blockArea / 4, max(1, ceilDiv(inChannels, tunables.WORKGROUP_SIZE)))",
+        "wordChunks": "ceilDiv(blockArea / 4, wordTile)"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "shader": "space-depth-collapsed-input.wgsl.jinja",
+          "bindings": [{ "arg": "input", "elementType": "$vectorScalar" }, "output"],
+          "dispatch": {
+            "x": "min(ceilDiv((channelCount * wordChunks), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((channelCount * wordChunks), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          },
+          "name": "SpaceToDepth.CollapsedInput"
+        }
+      ]
+    },
     {
       "id": "nchw_vec4",
       "priority": 20,
       "when": ["dim(shapes.output, 3) % 4 == 0", "f16VectorPermutationPreferred"],
+      "demoteWhen": ["variableSubgroup16To32 and dtypes.T == \"f16\""],
       "derive": { "outputElement": "\"vec4<\" ~ dtypes.T ~ \">\"", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "name": "SpaceToDepth",
           "shader": "space-depth-permute.wgsl.jinja",
           "derive": {
             "vectorizedSpec": true,
             "count": "numel(shapes.output) / 4",
             "outWUnits": "dim(shapes.output, 3) / 4",
             "outH": "dim(shapes.output, 2)",
+            "outC": "dim(shapes.output, 1)"
           },
           "bindings": ["input", "output"],
           "dispatch": {
           "name": "SpaceToDepth",
           "shader": "space-depth-permute.wgsl.jinja",
           "derive": {
             "vectorizedSpec": false,
             "count": "numel(shapes.output)",
             "outWUnits": "dim(shapes.output, 3)",
             "outH": "dim(shapes.output, 2)",
+            "outC": "dim(shapes.output, 1)"
           },
           "bindings": ["input", "output"],
           "dispatch": {

build/webgpu/metadata.json CHANGED Viewed

@@ -1,21 +1,27 @@
 {
   "name": "ai.onnx.SpaceToDepth",
-  "id": "_ai_onnx_spacetodepth_webgpu_6bff23a",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "myGa4HIZfSdpBZs8GbVV6ft8zVGVqZEoEu7w6utliEQ=",
-      "manifest.json": "QWMgtc+Yf2KOKYOyuriO1foyZEnYQRhRg+jrBgNSwdI=",
-      "space-depth-permute.wgsl.jinja": "Ju6qlxve31dB+p7oMoI/6Y8IlmMjeQM22fFpFX7FomU=",
-      "test.json": "dCIM8ExjNSAfPfl+rfxsNdlw897HmaSgBr2GgL2Kf9o="
     }
   },
-  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
   "webgpu": {
-    "manifestSpec": "2.0",
-    "variants": { "nchw_vec4": ["space-depth-permute.wgsl.jinja"], "nchw_const": ["space-depth-permute.wgsl.jinja"] }
   }
 }

 {
   "name": "ai.onnx.SpaceToDepth",
+  "id": "_ai_onnx_spacetodepth_webgpu_4eebc5c",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "B+hk3GtNTeUm6v61/3CKUrrttyjtP5WyibudxoqGY2g=",
+      "manifest.json": "g/1ubiwGGpjaR7Ybs7Fbqyzy6xmG6YKQlrNrr+ucDHw=",
+      "space-depth-collapsed-input.wgsl.jinja": "pdLUu9ZKhcDKGa5pY8q3rT/RPobF/9VyvzG4wWrWmes=",
+      "space-depth-permute.wgsl.jinja": "0JmPkn28PRyut+7IECJ/8KCJUVf3PJcvOnNYkUKOlLY=",
+      "test.json": "9sDs1EYNjkx5T7kjJaFkBUN7ghzjdoqHlhotuD+CyTA="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
+    "variants": {
+      "flat_vector": ["space-depth-permute.wgsl.jinja"],
+      "collapsed_input_vector": ["space-depth-collapsed-input.wgsl.jinja"],
+      "nchw_vec4": ["space-depth-permute.wgsl.jinja"],
+      "nchw_const": ["space-depth-permute.wgsl.jinja"]
+    }
   }
 }

build/webgpu/space-depth-collapsed-input.wgsl.jinja ADDED Viewed

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

+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
+  if ({{ name }} >= {{ bound }}) {
+    return;
+  }{% endmacro %}
+{{ env.wgsl.resourceDeclarations }}
+// A channel/block transpose with spatial tiles for few-channel inputs.
+// Pair f16 channels so each invocation owns complete output words.
+const COUNT: u32 = {{ channelCount * wordChunks }}u;
+const CHANNELS: u32 = max(1u, {{ inChannels }}u);
+const AREA: u32 = {{ blockArea }}u;
+{% if wordChunks > 1 %}
+const WORD_TILE: u32 = {{ wordTile }}u;
+const WORD_CHUNKS: u32 = {{ wordChunks }}u;
+{% endif %}
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{% if wordChunks > 1 %}
+{{ flat_index_2d(tunables.WORKGROUP_SIZE, "tile_lane", "COUNT") }}
+  let c = tile_lane % CHANNELS;
+  let tile = tile_lane / CHANNELS;
+  let n = tile / WORD_CHUNKS;
+  let channel = n * CHANNELS + c;
+  let word_begin = (tile % WORD_CHUNKS) * WORD_TILE;
+  for (var word = word_begin; word < min(word_begin + WORD_TILE, AREA / 4u); word++) {
+{% else %}
+{{ flat_index_2d(tunables.WORKGROUP_SIZE, "channel", "COUNT") }}
+  let c = channel % CHANNELS;
+  let n = channel / CHANNELS;
+  for (var word = 0u; word < AREA / 4u; word++) {
+{% endif %}
+{% if channelPack == 2 %}
+    let a = input[channel * 2u * (AREA / 4u) + word];
+    let b = input[(channel * 2u + 1u) * (AREA / 4u) + word];
+{% else %}
+    let values = input[channel * (AREA / 4u) + word];
+{% endif %}
+{% for lane in range(4) %}
+    output[(n * AREA + word * 4u + {{ lane }}u) * CHANNELS + c] ={% if channelPack == 2 %} vec2<f16>(a[{{ lane }}], b[{{ lane }}]){% else %} values[{{ lane }}]{% endif %};
+{% endfor %}
+  }
+}

build/webgpu/space-depth-permute.wgsl.jinja CHANGED Viewed

@@ -1,36 +1,12 @@
-{% macro flat_index_2d(name="i", bound="params.count", guardInline=false, note="dispatch-limit") %}
-{% if note == "dispatch-limit" %}
-  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
-  // per-axis workgroup fold width (outputs > 16.7M elements).
-{% elif note == "limit" %}
-  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
-  // per-axis workgroup fold width.
-{% elif note == "device-axis" %}
-  // The flat dispatch is folded across x/y at a fixed per-axis workgroup
-  // width; gid.y carries the high portion of the output index.
-{% elif note == "vec4-limit" %}
-  // 2D-folded flat vec4 index: gid.y carries the high bits past the dispatch's
-  // per-axis workgroup fold width (the dispatch caps x and spills into y).
-{% elif note == "element-limit" %}
-  // 2D-folded flat element index: gid.y carries the high bits past the
-  // dispatch's per-axis workgroup fold width.
-{% elif note == "dispatch" %}
   // 2D-folded flat index: gid.y carries the high bits past the dispatch's
   // per-axis workgroup fold width.
-{% endif %}
-{% if bound == "" %}
-  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
-{%- elif guardInline %}
-  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
-  if ({{ name }} >= {{ bound }}) { return; }
-{%- else %}
-  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) {
     return;
-  }
-{%- endif %}
-{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 // Specializing all geometry lets division and modulo lower to constant
@@ -47,10 +23,18 @@ const BLOCK_SIZE: u32 = {{ blocksize }}u;
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-{{ flat_index_2d("out_unit", "COUNT", note="") }}
-  let ow_unit = out_unit % OUT_W_UNITS;
-  var remaining = out_unit / OUT_W_UNITS;
   let oh = remaining % OUT_H;
   remaining /= OUT_H;
   let oc = remaining % OUT_C;
@@ -69,9 +53,18 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   let value{{ lane }} = input[input_base + ow{{ lane }} * BLOCK_SIZE];
 {% endfor %}
 {% if vectorizedSpec %}
   output[out_unit] = {{ vectorScalar }}(value0, value1, value2, value3);
 {% else %}
   output[out_unit] = value0;
 {% endif %}
 }

+{% set flatVector = flatVectorSpec is defined and flatVectorSpec %}
+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
   // 2D-folded flat index: gid.y carries the high bits past the dispatch's
   // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
   if ({{ name }} >= {{ bound }}) {
     return;
+  }{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 // Specializing all geometry lets division and modulo lower to constant
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d(tunables.WORKGROUP_SIZE, "out_unit", "COUNT") }}
+{% if flatVector %}
+  var packedValues: {{ vectorScalar }};
+{% endif %}
+{% for flatLane in range(4 if flatVector else 1) %}
+{% if flatVector %}
+  {
+  let element = out_unit * 4u + {{ flatLane }}u;
+{% endif %}
+  let ow_unit = {{ "element" if flatVector else "out_unit" }} % OUT_W_UNITS;
+  var remaining = {{ "element" if flatVector else "out_unit" }} / OUT_W_UNITS;
   let oh = remaining % OUT_H;
   remaining /= OUT_H;
   let oc = remaining % OUT_C;
   let value{{ lane }} = input[input_base + ow{{ lane }} * BLOCK_SIZE];
 {% endfor %}
+{% if flatVector %}
+  packedValues[{{ flatLane }}] = value0;
+  }
+{% else %}
 {% if vectorizedSpec %}
   output[out_unit] = {{ vectorScalar }}(value0, value1, value2, value3);
 {% else %}
   output[out_unit] = value0;
 {% endif %}
+{% endif %}
+{% endfor %}
+{% if flatVector %}
+  output[out_unit] = packedValues;
+{% endif %}
 }

build/webgpu/test.json CHANGED Viewed

@@ -355,6 +355,279 @@
         }
       },
       "outputs": { "output": { "dtype": "bool", "shape": [1, 4, 1, 2], "tolerance": 0 } }
     }
   ]
 }

         }
       },
       "outputs": { "output": { "dtype": "bool", "shape": [1, 4, 1, 2], "tolerance": 0 } }
+    },
+    {
+      "name": "permutation_float32_n2_c3_2x2_block2",
+      "attrs": { "blocksize": 2 },
+      "inputs": {
+        "input": {
+          "dtype": "float32",
+          "shape": [2, 3, 2, 2],
+          "data": { "kind": "linspace", "start": -7.0, "end": 11.0 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [2, 12, 1, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "permutation_float32_n2_c5_4x4_block4",
+      "attrs": { "blocksize": 4 },
+      "inputs": {
+        "input": {
+          "dtype": "float32",
+          "shape": [2, 5, 4, 4],
+          "data": { "kind": "linspace", "start": -7.0, "end": 11.0 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [2, 80, 1, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "permutation_float16_n2_c4_6x6_block6",
+      "attrs": { "blocksize": 6 },
+      "inputs": {
+        "input": {
+          "dtype": "float16",
+          "shape": [2, 4, 6, 6],
+          "data": { "kind": "linspace", "start": -7.0, "end": 11.0 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [2, 144, 1, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "permutation_float16_n2_c3_2x2_block2",
+      "attrs": { "blocksize": 2 },
+      "inputs": {
+        "input": {
+          "dtype": "float16",
+          "shape": [2, 3, 2, 2],
+          "data": { "kind": "linspace", "start": -7.0, "end": 11.0 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [2, 12, 1, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "permutation_float16_n1_c3_3x3_block3",
+      "attrs": { "blocksize": 3 },
+      "inputs": {
+        "input": {
+          "dtype": "float16",
+          "shape": [1, 3, 3, 3],
+          "data": { "kind": "linspace", "start": -7.0, "end": 11.0 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [1, 27, 1, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "permutation_int32_n2_c5_2x2_block2",
+      "attrs": { "blocksize": 2 },
+      "inputs": {
+        "input": { "dtype": "int32", "shape": [2, 5, 2, 2], "data": { "kind": "linspace", "start": -7, "end": 11 } }
+      },
+      "outputs": { "output": { "dtype": "int32", "shape": [2, 20, 1, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "permutation_uint32_n1_c7_4x4_block4",
+      "attrs": { "blocksize": 4 },
+      "inputs": {
+        "input": { "dtype": "uint32", "shape": [1, 7, 4, 4], "data": { "kind": "linspace", "start": 0, "end": 11 } }
+      },
+      "outputs": { "output": { "dtype": "uint32", "shape": [1, 112, 1, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "permutation_float32_n2_c3_6x10_block2",
+      "attrs": { "blocksize": 2 },
+      "inputs": {
+        "input": {
+          "dtype": "float32",
+          "shape": [2, 3, 6, 10],
+          "data": { "kind": "linspace", "start": -7.0, "end": 11.0 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [2, 12, 3, 5], "tolerance": 0 } }
+    },
+    {
+      "name": "permutation_float16_n4_c3_9x9_block3",
+      "attrs": { "blocksize": 3 },
+      "inputs": {
+        "input": {
+          "dtype": "float16",
+          "shape": [4, 3, 9, 9],
+          "data": { "kind": "linspace", "start": -7.0, "end": 11.0 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [4, 27, 3, 3], "tolerance": 0 } }
+    },
+    {
+      "name": "permutation_float32_n0_c4_2x2_block2",
+      "attrs": { "blocksize": 2 },
+      "inputs": {
+        "input": {
+          "dtype": "float32",
+          "shape": [0, 4, 2, 2],
+          "data": { "kind": "linspace", "start": -7.0, "end": 11.0 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [0, 16, 1, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "collapsed_channels_workgroup64_fold_x_to_y",
+      "attrs": { "blocksize": 2 },
+      "inputs": {
+        "input": {
+          "dtype": "float32",
+          "shape": [1, 4194241, 2, 2],
+          "data": { "kind": "linspace", "start": 0.0, "end": 1.0 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [1, 16776964, 1, 1], "tolerance": 0 } },
+      "tunables": { "WORKGROUP_SIZE": 64 }
+    },
+    {
+      "name": "spatial_parallel_float32_n1_c1_b64",
+      "attrs": { "blocksize": 64 },
+      "inputs": {
+        "input": {
+          "dtype": "float32",
+          "shape": [1, 1, 64, 64],
+          "data": { "kind": "linspace", "start": -3.0, "end": 5.0 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [1, 4096, 1, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "spatial_parallel_float32_n2_c3_b64",
+      "attrs": { "blocksize": 64 },
+      "inputs": {
+        "input": {
+          "dtype": "float32",
+          "shape": [2, 3, 64, 64],
+          "data": { "kind": "linspace", "start": -3.0, "end": 5.0 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [2, 12288, 1, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "spatial_parallel_float32_n1_c256_b32",
+      "attrs": { "blocksize": 32 },
+      "inputs": {
+        "input": {
+          "dtype": "float32",
+          "shape": [1, 256, 32, 32],
+          "data": { "kind": "linspace", "start": -3.0, "end": 5.0 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [1, 262144, 1, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "spatial_parallel_float32_n1_c16_b64",
+      "attrs": { "blocksize": 64 },
+      "inputs": {
+        "input": {
+          "dtype": "float32",
+          "shape": [1, 16, 64, 64],
+          "data": { "kind": "linspace", "start": -3.0, "end": 5.0 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [1, 65536, 1, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "spatial_parallel_float16_n1_c1_b64",
+      "attrs": { "blocksize": 64 },
+      "inputs": {
+        "input": {
+          "dtype": "float16",
+          "shape": [1, 1, 64, 64],
+          "data": { "kind": "linspace", "start": -3.0, "end": 5.0 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [1, 4096, 1, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "spatial_parallel_float16_n2_c3_b64",
+      "attrs": { "blocksize": 64 },
+      "inputs": {
+        "input": {
+          "dtype": "float16",
+          "shape": [2, 3, 64, 64],
+          "data": { "kind": "linspace", "start": -3.0, "end": 5.0 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [2, 12288, 1, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "spatial_parallel_float16_n1_c256_b32",
+      "attrs": { "blocksize": 32 },
+      "inputs": {
+        "input": {
+          "dtype": "float16",
+          "shape": [1, 256, 32, 32],
+          "data": { "kind": "linspace", "start": -3.0, "end": 5.0 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [1, 262144, 1, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "spatial_parallel_float16_n1_c16_b64",
+      "attrs": { "blocksize": 64 },
+      "inputs": {
+        "input": {
+          "dtype": "float16",
+          "shape": [1, 16, 64, 64],
+          "data": { "kind": "linspace", "start": -3.0, "end": 5.0 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [1, 65536, 1, 1], "tolerance": 0 } }
+    },
+    {
+      "name": "spatial_partial_float32_n2_c129_b4",
+      "attrs": { "blocksize": 4 },
+      "inputs": {
+        "input": {
+          "dtype": "float32",
+          "shape": [2, 129, 4, 4],
+          "data": { "kind": "linspace", "start": -3.0, "end": 5.0 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [2, 2064, 1, 1], "tolerance": 0 } },
+      "tunables": { "WORKGROUP_SIZE": 64 }
+    },
+    {
+      "name": "spatial_partial_float32_n2_c257_b6",
+      "attrs": { "blocksize": 6 },
+      "inputs": {
+        "input": {
+          "dtype": "float32",
+          "shape": [2, 257, 6, 6],
+          "data": { "kind": "linspace", "start": -3.0, "end": 5.0 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [2, 9252, 1, 1], "tolerance": 0 } },
+      "tunables": { "WORKGROUP_SIZE": 64 }
+    },
+    {
+      "name": "spatial_partial_float16_n2_c258_b4",
+      "attrs": { "blocksize": 4 },
+      "inputs": {
+        "input": {
+          "dtype": "float16",
+          "shape": [2, 258, 4, 4],
+          "data": { "kind": "linspace", "start": -3.0, "end": 5.0 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [2, 4128, 1, 1], "tolerance": 0 } },
+      "tunables": { "WORKGROUP_SIZE": 64 }
+    },
+    {
+      "name": "spatial_partial_float16_n2_c514_b6",
+      "attrs": { "blocksize": 6 },
+      "inputs": {
+        "input": {
+          "dtype": "float16",
+          "shape": [2, 514, 6, 6],
+          "data": { "kind": "linspace", "start": -3.0, "end": 5.0 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float16", "shape": [2, 18504, 1, 1], "tolerance": 0 } },
+      "tunables": { "WORKGROUP_SIZE": 64 }
     }
   ]
 }