sync 6fdf6301e2bb

Browse files

Files changed (6) hide show

README.md +6 -6
build/webgpu/bench.json +2 -2
build/webgpu/manifest.json +8 -12
build/webgpu/metadata.json +7 -7
build/webgpu/test.json +105 -3
build/webgpu/tile.wgsl.jinja +9 -6

README.md CHANGED Viewed

@@ -33,7 +33,7 @@ See the [ONNX `Tile` spec](https://onnx.ai/onnx/operators/onnx__Tile.html) for t
 | Variable | Allowed dtypes |
 | --- | --- |
-| `T` | `float32`, `float16`, `uint32`, `int32`, `int16`, `uint8`, `int8`, `bool` |
 | `S` | `int64` |
 ## Files
@@ -41,14 +41,14 @@ See the [ONNX `Tile` spec](https://onnx.ai/onnx/operators/onnx__Tile.html) for t
 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
-- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 - [`datamove-tile-vec4.wgsl.jinja`](build/webgpu/datamove-tile-vec4.wgsl.jinja)
 - [`tile.wgsl.jinja`](build/webgpu/tile.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
-npm install --save-exact @huggingface/kernels@0.0.1-preview.2
 ```
 Outputs with inferable metadata are allocated automatically. Explicit `outputs` entries request optional results or provide metadata that cannot be inferred from the supplied inputs and attributes.
@@ -68,9 +68,9 @@ import { getKernel } from "@huggingface/kernels";
 const kernel = await getKernel("webgpu-kernels/ai.onnx.Tile", { version: 1 });
 // Explicit destinations request optional results or supply metadata that cannot be inferred.
 const { output } = await kernel({
-  input: { data: inputData, shape: [2, 1, 3] },
-  repeats: { data: repeatsData, shape: [3] },
 }, {
-  outputs: { output: { shape: [2, 1, 3], dtype: "float32" } },
 });
 ```

 | Variable | Allowed dtypes |
 | --- | --- |
+| `T` | `float32`, `float16`, `uint32`, `int32`, `int16`, `uint8`, `int8`, `bool`, `int64` |
 | `S` | `int64` |
 ## 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
 - [`datamove-tile-vec4.wgsl.jinja`](build/webgpu/datamove-tile-vec4.wgsl.jinja)
 - [`tile.wgsl.jinja`](build/webgpu/tile.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.3
 ```
 Outputs with inferable metadata are allocated automatically. Explicit `outputs` entries request optional results or provide metadata that cannot be inferred from the supplied inputs and attributes.
 const kernel = await getKernel("webgpu-kernels/ai.onnx.Tile", { version: 1 });
 // Explicit destinations request optional results or supply metadata that cannot be inferred.
 const { output } = await kernel({
+  input: { data: inputData, shape: [3] },
+  repeats: { data: repeatsData, shape: [1] },
 }, {
+  outputs: { output: { shape: [6], dtype: "int64" } },
 });
 ```

build/webgpu/bench.json CHANGED Viewed

@@ -93,7 +93,7 @@
       "name": "tile-f32-8x1x197x197-repeat-heads-generic-pathology",
       "preset": "stress",
       "provenance": {
-        "notes": "A [B,1,S,S] attention mask tiled over 8 heads at the ViT-B/16 token count: the trailing unrepeated block is 197*197 = 38809 elements, not a multiple of four, so tileVec4Ok fails and the 2.5M-element copy runs the generic per-element route. Control: tile-f32-8x1x200x200-repeat-heads-control."
       },
       "vars": { "dtype": "float32", "outCount": 2483776 },
       "inputs": {
@@ -107,7 +107,7 @@
       "name": "tile-f32-8x1x200x200-repeat-heads-control",
       "preset": "stress",
       "provenance": {
-        "notes": "Nearest multiple-of-4 token count for tile-f32-8x1x197x197-repeat-heads-generic-pathology (3% more bytes): the trailing block 200*200 folds to a vec4 tail and inner_vec4 is selected."
       },
       "vars": { "dtype": "float32", "outCount": 2560000 },
       "inputs": {

       "name": "tile-f32-8x1x197x197-repeat-heads-generic-pathology",
       "preset": "stress",
       "provenance": {
+        "notes": "A [B,1,S,S] attention mask repeated over eight heads at the ViT-B/16 token count: the trailing block has 197*197 elements, not divisible by four. The paired 200-token case checks an aligned block."
       },
       "vars": { "dtype": "float32", "outCount": 2483776 },
       "inputs": {
       "name": "tile-f32-8x1x200x200-repeat-heads-control",
       "preset": "stress",
       "provenance": {
+        "notes": "The next multiple-of-four token count above the paired 197-token case, with about 3% more bytes."
       },
       "vars": { "dtype": "float32", "outCount": 2560000 },
       "inputs": {

build/webgpu/manifest.json CHANGED Viewed

@@ -8,7 +8,7 @@
   },
   "outputs": { "output": { "dtype": "T", "rank": "ranks.input" } },
   "typeConstraints": {
-    "T": ["float32", "float16", "uint32", "int32", "int16", "uint8", "int8", "bool"],
     "S": ["int64"]
   },
   "tunables": { "WORKGROUP_SIZE": { "default": 256 } },
@@ -16,27 +16,24 @@
     "tileTailAxes": "pick([[numel(suffix(shapes.input, ranks.input - 8)) == numel(suffix(shapes.output, ranks.input - 8)), 8], [numel(suffix(shapes.input, ranks.input - 7)) == numel(suffix(shapes.output, ranks.input - 7)), 7], [numel(suffix(shapes.input, ranks.input - 6)) == numel(suffix(shapes.output, ranks.input - 6)), 6], [numel(suffix(shapes.input, ranks.input - 5)) == numel(suffix(shapes.output, ranks.input - 5)), 5], [numel(suffix(shapes.input, ranks.input - 4)) == numel(suffix(shapes.output, ranks.input - 4)), 4], [numel(suffix(shapes.input, ranks.input - 3)) == numel(suffix(shapes.output, ranks.input - 3)), 3], [numel(suffix(shapes.input, ranks.input - 2)) == numel(suffix(shapes.output, ranks.input - 2)), 2], [numel(suffix(shapes.input, ranks.input - 1)) == numel(suffix(shapes.output, ranks.input - 1)), 1]], 0)",
     "tileFoldAxes": "min(tileTailAxes, ranks.input)",
     "tileTail": "numel(suffix(shapes.input, ranks.input - tileFoldAxes)) if tileFoldAxes > 0 else (dim(shapes.input, ranks.input - 1) if ranks.input > 0 else 0)",
-    "tileVec4Ok": "ranks.input >= 1 and numel(shapes.output) > 0 and tileTail % 4 == 0"
   },
   "when": ["ranks.input == ranks.output", "ranks.repeats == 1", "dim(shapes.repeats, 0) == ranks.input", "f16Ok(dtypes.T)"],
   "variants": [
     {
       "id": "inner_vec4",
       "priority": 15,
-      "when": ["ranks.input >= 1", "numel(shapes.output) > 0", "tileVec4Ok"],
       "derive": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
-          "name": "Tile.innerVec4",
           "shader": "datamove-tile-vec4.wgsl.jinja",
-          "derive": {
-            "inputShape": "shapes.input",
-            "outputShape": "shapes.output",
-            "rank": "ranks.input",
-            "count": "numel(shapes.output) / 4",
-            "foldAxes": "tileFoldAxes"
-          },
           "bindings": [
             { "arg": "input", "elementType": "$vectorScalar" },
             { "arg": "output", "elementType": "$vectorScalar" }
@@ -57,7 +54,6 @@
           "id": "main",
           "name": "Tile",
           "shader": "tile.wgsl.jinja",
-          "derive": { "inputShape": "shapes.input", "outputShape": "shapes.output", "rank": "ranks.input" },
           "bindings": [
             { "arg": "input", "elementType": "$scalar" },
             { "arg": "output", "elementType": "$scalar" },

   },
   "outputs": { "output": { "dtype": "T", "rank": "ranks.input" } },
   "typeConstraints": {
+    "T": ["float32", "float16", "uint32", "int32", "int16", "uint8", "int8", "bool", "int64"],
     "S": ["int64"]
   },
   "tunables": { "WORKGROUP_SIZE": { "default": 256 } },
     "tileTailAxes": "pick([[numel(suffix(shapes.input, ranks.input - 8)) == numel(suffix(shapes.output, ranks.input - 8)), 8], [numel(suffix(shapes.input, ranks.input - 7)) == numel(suffix(shapes.output, ranks.input - 7)), 7], [numel(suffix(shapes.input, ranks.input - 6)) == numel(suffix(shapes.output, ranks.input - 6)), 6], [numel(suffix(shapes.input, ranks.input - 5)) == numel(suffix(shapes.output, ranks.input - 5)), 5], [numel(suffix(shapes.input, ranks.input - 4)) == numel(suffix(shapes.output, ranks.input - 4)), 4], [numel(suffix(shapes.input, ranks.input - 3)) == numel(suffix(shapes.output, ranks.input - 3)), 3], [numel(suffix(shapes.input, ranks.input - 2)) == numel(suffix(shapes.output, ranks.input - 2)), 2], [numel(suffix(shapes.input, ranks.input - 1)) == numel(suffix(shapes.output, ranks.input - 1)), 1]], 0)",
     "tileFoldAxes": "min(tileTailAxes, ranks.input)",
     "tileTail": "numel(suffix(shapes.input, ranks.input - tileFoldAxes)) if tileFoldAxes > 0 else (dim(shapes.input, ranks.input - 1) if ranks.input > 0 else 0)",
+    "tileVec4Ok": "ranks.input >= 1 and numel(shapes.output) > 0 and tileTail % 4 == 0",
+    "inputShape": "shapes.input",
+    "outputShape": "shapes.output",
+    "rank": "ranks.input"
   },
   "when": ["ranks.input == ranks.output", "ranks.repeats == 1", "dim(shapes.repeats, 0) == ranks.input", "f16Ok(dtypes.T)"],
   "variants": [
     {
       "id": "inner_vec4",
       "priority": 15,
+      "when": ["dtypes.T != \"vec2<u32>\"", "ranks.input >= 1", "numel(shapes.output) > 0", "tileVec4Ok"],
       "derive": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
+          "name": "Tile.InnerVec4",
           "shader": "datamove-tile-vec4.wgsl.jinja",
+          "derive": { "count": "numel(shapes.output) / 4", "foldAxes": "tileFoldAxes" },
           "bindings": [
             { "arg": "input", "elementType": "$vectorScalar" },
             { "arg": "output", "elementType": "$vectorScalar" }
           "id": "main",
           "name": "Tile",
           "shader": "tile.wgsl.jinja",
           "bindings": [
             { "arg": "input", "elementType": "$scalar" },
             { "arg": "output", "elementType": "$scalar" },

build/webgpu/metadata.json CHANGED Viewed

@@ -1,22 +1,22 @@
 {
   "name": "ai.onnx.Tile",
-  "id": "_ai_onnx_tile_webgpu_f5f5692",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "YfCc57C/DdUbEVmhWy5t2xoP2Le9Yi9K1tINlEYgeLw=",
       "datamove-tile-vec4.wgsl.jinja": "2+EFl6e/twKm/3oj2RBh+sx0EhYl/NNGGloW33ZprLs=",
-      "manifest.json": "YC55CwcGcRBjXpT8bLDiPA2dt7Ml7TBb2Z1w4Bmilos=",
-      "test.json": "b8XhMm/6K0JPfgbWB718034tNom6lm9mc2dOE/ctr3U=",
-      "tile.wgsl.jinja": "JTi9v0O2HGR8u6QvjZFSHW2cycDILumledR4r/GmkA8="
     }
   },
-  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
   "webgpu": {
-    "manifestSpec": "2.0",
     "variants": { "inner_vec4": ["datamove-tile-vec4.wgsl.jinja"], "generic": ["tile.wgsl.jinja"] }
   }
 }

 {
   "name": "ai.onnx.Tile",
+  "id": "_ai_onnx_tile_webgpu_32e0f56",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "1eM36L7PfUEX11p0lx+TyZ25BWbeAJD+H3gHG3sSRZM=",
       "datamove-tile-vec4.wgsl.jinja": "2+EFl6e/twKm/3oj2RBh+sx0EhYl/NNGGloW33ZprLs=",
+      "manifest.json": "1KJxVLG3UNqIFOiQACZOvKRxTS8R046eRe9Hw0jYJUQ=",
+      "test.json": "Ur/BnRjKBHpjTa1kPqGIHWhfmOX42bRXne1iRh3xGas=",
+      "tile.wgsl.jinja": "/Se423/ugSXxY/IqEB7Ktbc7G3xf9sVsE13TSM419hE="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "variants": { "inner_vec4": ["datamove-tile-vec4.wgsl.jinja"], "generic": ["tile.wgsl.jinja"] }
   }
 }

build/webgpu/test.json CHANGED Viewed

@@ -414,8 +414,8 @@
     {
       "name": "rank7_unit_input_alternating_repeats",
       "provenance": {
-        "source": "ONNX Tile spec (no rank limit) + onnxruntime CPU Tile kernel",
-        "notes": "Spec-valid rank-7 Tile. ORT CPU computes it via its generic strided copy."
       },
       "inputs": {
         "input": {
@@ -536,7 +536,7 @@
     {
       "name": "inner_vec4_lastaxis_repeat_mul4",
       "provenance": {
-        "notes": "The innermost axis itself is repeated; its input extent (8) is a multiple of four so every output vec4 still lies inside one copy and the vec4 route applies the modulo on the last axis."
       },
       "inputs": {
         "input": {
@@ -611,6 +611,108 @@
         "repeats": { "dtype": "uint32", "shape": [3], "data": { "kind": "values", "values": [2, 1, 1] } }
       },
       "outputs": { "output": { "dtype": "float32", "shape": [4, 3, 3], "tolerance": 0 } }
     }
   ]
 }

     {
       "name": "rank7_unit_input_alternating_repeats",
       "provenance": {
+        "source": "ONNX Tile spec (no rank limit) and the Tile kernel in ONNX Runtime's CPU provider",
+        "notes": "Spec-valid rank-7 Tile. ONNX Runtime's CPU provider computes it with its generic strided copy."
       },
       "inputs": {
         "input": {
     {
       "name": "inner_vec4_lastaxis_repeat_mul4",
       "provenance": {
+        "notes": "The innermost axis has extent eight and repeats in four-aligned copies, checking modulo wrap at each copy boundary."
       },
       "inputs": {
         "input": {
         "repeats": { "dtype": "uint32", "shape": [3], "data": { "kind": "values", "values": [2, 1, 1] } }
       },
       "outputs": { "output": { "dtype": "float32", "shape": [4, 3, 3], "tolerance": 0 } }
+    },
+    {
+      "name": "int64_full_range_rank2",
+      "inputs": {
+        "input": {
+          "dtype": "int64",
+          "shape": [2, 2],
+          "data": {
+            "kind": "cycle",
+            "values": ["0", "4294967296", "8589934593", "-4294967297", "9223372036854775807", "-9223372036854775808", "9223372036854775806", "-1"]
+          }
+        },
+        "repeats": { "dtype": "uint32", "shape": [2], "data": { "kind": "values", "values": [2, 3] } }
+      },
+      "outputs": { "output": { "dtype": "int64", "shape": [4, 6], "tolerance": 0, "relTolerance": 0 } },
+      "tolerance": 0,
+      "relTolerance": 0
+    },
+    {
+      "name": "int64_full_range_ort_float_2d_two_axes_full_repeat",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/tensor/tile_op_test.cc",
+        "test": "TensorOpTest.TileFloatType",
+        "notes": "RunTest<float>({2, 2}, {2, 2}) from ORT's shared Tile wrapper."
+      },
+      "inputs": {
+        "input": {
+          "dtype": "int64",
+          "shape": [2, 2],
+          "data": {
+            "kind": "cycle",
+            "values": ["4294967296", "8589934593", "-4294967297", "9223372036854775807", "-9223372036854775808", "9223372036854775806", "-1", "0"]
+          }
+        },
+        "repeats": { "dtype": "uint32", "shape": [2], "data": { "kind": "values", "values": [2, 2] } }
+      },
+      "outputs": { "output": { "dtype": "int64", "shape": [4, 4], "tolerance": 0, "relTolerance": 0 } },
+      "tolerance": 0,
+      "relTolerance": 0
+    },
+    {
+      "name": "ort_int64_1d",
+      "inputs": {
+        "input": {
+          "dtype": "int64",
+          "shape": [3],
+          "data": { "kind": "values", "values": ["281483566841860", "1407400653815816", "-1"] }
+        },
+        "repeats": { "dtype": "uint32", "shape": [1], "data": { "kind": "values", "values": [2] } }
+      },
+      "outputs": {
+        "output": {
+          "dtype": "int64",
+          "shape": [6],
+          "data": {
+            "kind": "values",
+            "values": ["281483566841860", "1407400653815816", "-1", "281483566841860", "1407400653815816", "-1"]
+          }
+        }
+      }
+    },
+    {
+      "name": "ort_int64_2d",
+      "inputs": {
+        "input": {
+          "dtype": "int64",
+          "shape": [2, 2],
+          "data": { "kind": "values", "values": ["281483566841860", "1407400653815816", "-1", "281483566841860"] }
+        },
+        "repeats": { "dtype": "uint32", "shape": [2], "data": { "kind": "values", "values": [2, 1] } }
+      },
+      "outputs": {
+        "output": {
+          "dtype": "int64",
+          "shape": [4, 2],
+          "data": {
+            "kind": "values",
+            "values": ["281483566841860", "1407400653815816", "-1", "281483566841860", "281483566841860", "1407400653815816", "-1", "281483566841860"]
+          }
+        }
+      }
+    },
+    {
+      "name": "ort_int64_3d",
+      "inputs": {
+        "input": {
+          "dtype": "int64",
+          "shape": [1, 1, 3],
+          "data": { "kind": "values", "values": ["281483566841860", "1407400653815816", "-1"] }
+        },
+        "repeats": { "dtype": "uint32", "shape": [3], "data": { "kind": "values", "values": [1, 2, 1] } }
+      },
+      "outputs": {
+        "output": {
+          "dtype": "int64",
+          "shape": [1, 2, 3],
+          "data": {
+            "kind": "values",
+            "values": ["281483566841860", "1407400653815816", "-1", "281483566841860", "1407400653815816", "-1"]
+          }
+        }
+      }
     }
   ]
 }

build/webgpu/tile.wgsl.jinja CHANGED Viewed

@@ -1,3 +1,11 @@
 {{ env.wgsl.resourceDeclarations }}
 fn input_offset({% if rank > 0 %}out_index: u32{% endif %}) -> u32 {
@@ -25,11 +33,6 @@ fn input_offset({% if rank > 0 %}out_index: u32{% endif %}) -> u32 {
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-  // 2D-folded flat index: gid.y carries the high bits past the
-  // per-axis dispatch fold width.
-  let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
-  if (i >= params.count) {
-    return;
-  }
   output[i] = input[input_offset({% if rank > 0 %}i{% endif %})];
 }

+{% 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 }}
 fn input_offset({% if rank > 0 %}out_index: u32{% endif %}) -> u32 {
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d(tunables.WORKGROUP_SIZE) }}
   output[i] = input[input_offset({% if rank > 0 %}i{% endif %})];
 }