sync 2e7068faf55e

Browse files

Files changed (6) hide show

README.md +71 -0
build/webgpu/bench.json +20 -0
build/webgpu/if-select.wgsl.jinja +40 -0
build/webgpu/manifest.json +76 -0
build/webgpu/metadata.json +18 -0
build/webgpu/test.json +252 -0

README.md CHANGED Viewed

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

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# ai.onnx.If
+`ai.onnx`  ·  internal tensor lowering (non-standard)  ·  reviewed against ONNX opset 25
+## Description
+Support status: the standard ONNX `If` control-flow operator is not implemented because standalone kernel packages cannot carry or execute its `then_branch` and `else_branch` graph attributes. This internal lowering only selects elementwise between two pre-evaluated, equal-sized tensors from a scalar condition and must not be treated as ONNX `If`.
+See the [standard ONNX `If` spec](https://onnx.ai/onnx/operators/onnx__If.html) for the contract this internal lowering does not implement.
+## Inputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `cond` | `cond` | `B` | — | — | Scalar boolean condition that selects which value tensor to output. | required |
+| `then_value` | `then_value` | `T` | — | — | Values to output when `cond` is true. | required |
+| `else_value` | `else_value` | `T` | — | — | Values to output when `cond` is false. | required |
+## Outputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `y` | `y` | `T` | — | — | Output tensor with the same number of elements as `then_value` and `else_value`. | required |
+## Type constraints
+| Variable | Allowed dtypes |
+| --- | --- |
+| `T` | `float32` |
+| `B` | `uint32`, `bool` |
+## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
+- [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
+- [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
+- [`if-select.wgsl.jinja`](build/webgpu/if-select.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.If", { version: 1 });
+// Explicit destinations request optional results or supply metadata that cannot be inferred.
+const { y } = await kernel({
+  cond: { data: condData, shape: [1] },
+  then_value: { data: then_valueData, shape: [1] },
+  else_value: { data: else_valueData, shape: [1] },
+}, {
+  outputs: { y: { shape: [1], dtype: "float32" } },
+});
+```

build/webgpu/bench.json ADDED Viewed

	@@ -0,0 +1,20 @@

+{
+  "op": "ai.onnx.If",
+  "cases": [
+    {
+      "name": "lowered_select_1m",
+      "preset": "smoke",
+      "inputs": {
+        "cond": { "dtype": "uint32", "shape": [1], "data": { "kind": "values", "values": [1] } },
+        "then_value": { "dtype": "float32", "shape": [1048576] },
+        "else_value": { "dtype": "float32", "shape": [1048576] }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1048576] } },
+      "bench": {
+        "metrics": [
+          { "type": "bandwidth", "value": "4 * (numel(shapes.cond) + numel(shapes.then_value) + numel(shapes.y))" }
+        ]
+      }
+    }
+  ]
+}

build/webgpu/if-select.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,40 @@

+{% 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
+  // maxComputeWorkgroupsPerDimension dispatch limit (outputs > 16.7M elements).
+{% elif note == "limit" %}
+  // 2D-folded flat index: gid.y carries the high bits past the
+  // maxComputeWorkgroupsPerDimension limit.
+{% elif note == "device-axis" %}
+  // The flat dispatch is folded across x/y at the device's per-axis workgroup
+  // limit; 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
+  // maxComputeWorkgroupsPerDimension limit (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
+  // maxComputeWorkgroupsPerDimension limit.
+{% elif note == "dispatch" %}
+  // 2D-folded flat index: gid.y carries the high bits past the
+  // maxComputeWorkgroupsPerDimension dispatch limit.
+{% endif %}
+{% if bound == "" %}
+  let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
+{%- elif guardInline %}
+  let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
+  if ({{ name }} >= {{ bound }}) { return; }
+{%- else %}
+  let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
+  if ({{ name }} >= {{ bound }}) {
+    return;
+  }
+{%- endif %}
+{% endmacro %}
+{{ env.wgsl.resourceDeclarations }}
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
+{{ flat_index_2d(guardInline=true) }}
+  y[i] = select(else_value[i], then_value[i], cond[0] != 0u);
+}

build/webgpu/manifest.json ADDED Viewed

	@@ -0,0 +1,76 @@

+{
+  "domain": "ai.onnx",
+  "name": "If",
+  "conformance": "internal-lowering",
+  "sinceVersion": 25,
+  "description": "Support status: the standard ONNX `If` control-flow operator is not implemented because standalone kernel packages cannot carry or execute its `then_branch` and `else_branch` graph attributes. This internal lowering only selects elementwise between two pre-evaluated, equal-sized tensors from a scalar condition and must not be treated as ONNX `If`.",
+  "inputs": [
+    {
+      "role": "cond",
+      "dtype": "B",
+      "description": "Scalar boolean condition that selects which value tensor to output."
+    },
+    { "role": "then_value", "dtype": "T", "description": "Values to output when `cond` is true." },
+    { "role": "else_value", "dtype": "T", "description": "Values to output when `cond` is false." }
+  ],
+  "outputs": [
+    {
+      "role": "y",
+      "dtype": "T",
+      "description": "Output tensor with the same number of elements as `then_value` and `else_value`."
+    }
+  ],
+  "typeConstraints": { "T": ["float32"], "B": ["uint32", "bool"] },
+  "args": {
+    "cond": { "kind": "tensor", "semantic": "cond", "role": "input" },
+    "then_value": { "kind": "tensor", "semantic": "then_value", "role": "input" },
+    "else_value": { "kind": "tensor", "semantic": "else_value", "role": "input" },
+    "y": { "kind": "tensor", "semantic": "y", "role": "output" }
+  },
+  "tunables": { "WORKGROUP_SIZE": 256 },
+  "variants": [
+    {
+      "id": "lowered_select",
+      "when": ["ranks.cond <= 1", "numel(shapes.cond) == 1", "numel(shapes.then_value) == numel(shapes.y)", "numel(shapes.else_value) == numel(shapes.y)"],
+      "passes": [
+        {
+          "id": "main",
+          "name": "If",
+          "shader": "if-select.wgsl.jinja",
+          "bindings": [
+            {
+              "name": "cond",
+              "arg": "cond",
+              "semantic": "cond",
+              "buffer": { "type": "read-only-storage" },
+              "elementType": "u32",
+              "length": 1
+            },
+            {
+              "name": "then_value",
+              "arg": "then_value",
+              "semantic": "then_value",
+              "buffer": { "type": "read-only-storage" },
+              "elementType": "f32"
+            },
+            {
+              "name": "else_value",
+              "arg": "else_value",
+              "semantic": "else_value",
+              "buffer": { "type": "read-only-storage" },
+              "elementType": "f32"
+            },
+            { "name": "y", "arg": "y", "semantic": "y", "buffer": { "type": "storage" }, "elementType": "f32" },
+            {
+              "name": "params",
+              "semantic": "kernel.params",
+              "buffer": { "type": "uniform" },
+              "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }] }
+            }
+          ],
+          "dispatch": { "threads": "numel(shapes.y)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
+        }
+      ]
+    }
+  ]
+}

build/webgpu/metadata.json ADDED Viewed

	@@ -0,0 +1,18 @@

+{
+  "name": "ai.onnx.If",
+  "id": "_ai_onnx_if_webgpu_2b93a6b",
+  "version": 1,
+  "license": "Apache-2.0",
+  "backend": { "type": "webgpu" },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "bench.json": "aVzy6O2BXjbtp7QD6jSgMGQnCFr9rv3ErzC9HuwA6gE=",
+      "if-select.wgsl.jinja": "SOCO3HJ0GVfXYEOn7wNTEDBq94rDoQ5sGuijBo/4u6Y=",
+      "manifest.json": "qnXbI/F6Tz6YY42S8aJWiA5scMh+fR/ZkPTGfwzYYcc=",
+      "test.json": "wPMNj7icFakVJ6lHXUZxozdTSnpyTakK1is0VuwCvzE="
+    }
+  },
+  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
+  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.If" }
+}

build/webgpu/test.json ADDED Viewed

	@@ -0,0 +1,252 @@

+{
+  "op": "ai.onnx.If",
+  "cases": [
+    {
+      "name": "lowered_dispatch_cliff_select_then",
+      "inputs": {
+        "cond": { "dtype": "uint32", "shape": [1], "data": { "kind": "values", "values": [1] } },
+        "then_value": { "dtype": "float32", "shape": [16777216], "data": { "kind": "constant", "value": 1.5 } },
+        "else_value": { "dtype": "float32", "shape": [16777216], "data": { "kind": "constant", "value": -2.5 } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [16777216], "tolerance": 0 } }
+    },
+    {
+      "name": "lowered_select_then",
+      "inputs": {
+        "cond": { "dtype": "uint32", "shape": [1], "data": { "kind": "values", "values": [1] } },
+        "then_value": {
+          "dtype": "float32",
+          "shape": [2, 2],
+          "data": { "kind": "values", "values": [1.0, 2.0, 3.0, 4.0] }
+        },
+        "else_value": {
+          "dtype": "float32",
+          "shape": [2, 2],
+          "data": { "kind": "values", "values": [-1.0, -2.0, -3.0, -4.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 2] } }
+    },
+    {
+      "name": "lowered_select_else",
+      "inputs": {
+        "cond": { "dtype": "uint32", "shape": [1], "data": { "kind": "values", "values": [0] } },
+        "then_value": { "dtype": "float32", "shape": [3], "data": { "kind": "values", "values": [10.0, 20.0, 30.0] } },
+        "else_value": { "dtype": "float32", "shape": [3], "data": { "kind": "values", "values": [4.0, 5.0, 6.0] } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3] } }
+    },
+    {
+      "name": "lowered_ort_projection_outer_scope_add_then",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/controlflow/if_test.cc",
+        "test": "If.ShapeInMainGraph_NoShapeInSubgraph_True",
+        "notes": "Projection onto the framework's lowered select variant: ORT's then branch computes split_out_0 + if_graph_input_0 = 2."
+      },
+      "inputs": {
+        "cond": { "dtype": "uint32", "shape": [1], "data": { "kind": "values", "values": [1] } },
+        "then_value": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [2.0] } },
+        "else_value": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [11.0] } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1], "tolerance": 0 } }
+    },
+    {
+      "name": "lowered_ort_projection_outer_scope_add_else",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/controlflow/if_test.cc",
+        "test": "If.ShapeInMainGraph_NoShapeInSubgraph_False",
+        "notes": "Projection onto the framework's lowered select variant: ORT's else branch computes split_out_1 + if_graph_input_0 = 11."
+      },
+      "inputs": {
+        "cond": { "dtype": "uint32", "shape": [1], "data": { "kind": "values", "values": [0] } },
+        "then_value": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [2.0] } },
+        "else_value": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [11.0] } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1], "tolerance": 0 } }
+    },
+    {
+      "name": "lowered_ort_projection_constant_then_branch",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/controlflow/if_test.cc",
+        "test": "If.ConditionalBranchesOnlyContainConstantNodes_ThenBranchExecution",
+        "notes": "Projection onto the framework's lowered select variant: branch subgraphs are represented by precomputed branch tensors."
+      },
+      "inputs": {
+        "cond": { "dtype": "uint32", "shape": [1], "data": { "kind": "values", "values": [1] } },
+        "then_value": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [10.0] } },
+        "else_value": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [1000.0] } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1] } }
+    },
+    {
+      "name": "lowered_ort_projection_constant_else_branch",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/controlflow/if_test.cc",
+        "test": "If.ConditionalBranchesOnlyContainConstantNodes_ElseBranchExecution",
+        "notes": "Projection onto the framework's lowered select variant: branch subgraphs are represented by precomputed branch tensors."
+      },
+      "inputs": {
+        "cond": { "dtype": "uint32", "shape": [1], "data": { "kind": "values", "values": [0] } },
+        "then_value": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [10.0] } },
+        "else_value": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [1000.0] } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1] } }
+    },
+    {
+      "name": "lowered_ort_projection_different_branch_shapes_then",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/controlflow/if_test.cc",
+        "test": "If.Opset11ThenAndElseBranchesProduceDifferentOutputShapes",
+        "notes": "Projection onto the framework's lowered select variant: branch tensors use different ranks with equal storage size, and the selected then branch fixes the output shape."
+      },
+      "inputs": {
+        "cond": { "dtype": "uint32", "shape": [1], "data": { "kind": "values", "values": [1] } },
+        "then_value": {
+          "dtype": "float32",
+          "shape": [2, 2],
+          "data": { "kind": "values", "values": [1.0, 2.0, 3.0, 4.0] }
+        },
+        "else_value": {
+          "dtype": "float32",
+          "shape": [4],
+          "data": { "kind": "values", "values": [-1.0, -2.0, -3.0, -4.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 2] } }
+    },
+    {
+      "name": "lowered_ort_projection_different_branch_shapes_else",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/controlflow/if_test.cc",
+        "test": "If.Opset11ThenAndElseBranchesProduceDifferentOutputShapes",
+        "notes": "Projection onto the framework's lowered select variant: branch tensors use different ranks with equal storage size, and the selected else branch fixes the output shape."
+      },
+      "inputs": {
+        "cond": { "dtype": "uint32", "shape": [1], "data": { "kind": "values", "values": [0] } },
+        "then_value": {
+          "dtype": "float32",
+          "shape": [2, 2],
+          "data": { "kind": "values", "values": [1.0, 2.0, 3.0, 4.0] }
+        },
+        "else_value": {
+          "dtype": "float32",
+          "shape": [4],
+          "data": { "kind": "values", "values": [-1.0, -2.0, -3.0, -4.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [4] } }
+    },
+    {
+      "name": "lowered_select_then_scalar",
+      "inputs": {
+        "cond": { "dtype": "uint32", "shape": [1], "data": { "kind": "values", "values": [1] } },
+        "then_value": { "dtype": "float32", "shape": [], "data": { "kind": "values", "values": [42.0] } },
+        "else_value": { "dtype": "float32", "shape": [], "data": { "kind": "values", "values": [-42.0] } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [] } }
+    },
+    {
+      "name": "lowered_scalar_cond_then",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/controlflow/if_test.cc",
+        "test": "If.ConditionalBranchesOnlyContainConstantNodes_ThenBranchExecution",
+        "notes": "ONNX If condition is a scalar bool. The project-lowered select form should accept scalar logical uint32 conditions as well as [1] conditions."
+      },
+      "inputs": {
+        "cond": { "dtype": "uint32", "shape": [], "data": { "kind": "values", "values": [1] } },
+        "then_value": { "dtype": "float32", "shape": [2], "data": { "kind": "values", "values": [9.0, -3.0] } },
+        "else_value": { "dtype": "float32", "shape": [2], "data": { "kind": "values", "values": [4.0, 5.0] } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2], "tolerance": 0 } }
+    },
+    {
+      "name": "lowered_scalar_cond_else",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/controlflow/if_test.cc",
+        "test": "If.ConditionalBranchesOnlyContainConstantNodes_ElseBranchExecution",
+        "notes": "ONNX If condition is a scalar bool. This exercises the false branch with scalar logical uint32 condition storage."
+      },
+      "inputs": {
+        "cond": { "dtype": "uint32", "shape": [], "data": { "kind": "values", "values": [0] } },
+        "then_value": { "dtype": "float32", "shape": [2], "data": { "kind": "values", "values": [9.0, -3.0] } },
+        "else_value": { "dtype": "float32", "shape": [2], "data": { "kind": "values", "values": [4.0, 5.0] } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2], "tolerance": 0 } }
+    },
+    {
+      "name": "lowered_bool_scalar_cond_then",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/controlflow/if_test.cc",
+        "test": "If.ConditionalBranchesOnlyContainConstantNodes_ThenBranchExecution",
+        "notes": "ONNX If conditions are bool tensors. This lowered projection verifies scalar logical bool storage selects the then branch."
+      },
+      "inputs": {
+        "cond": { "dtype": "bool", "shape": [], "data": { "kind": "values", "values": [1] } },
+        "then_value": { "dtype": "float32", "shape": [3], "data": { "kind": "values", "values": [1.25, -2.5, 3.75] } },
+        "else_value": {
+          "dtype": "float32",
+          "shape": [3],
+          "data": { "kind": "values", "values": [-10.0, -20.0, -30.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3], "tolerance": 0 } }
+    },
+    {
+      "name": "lowered_bool_len1_cond_else",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/controlflow/if_test.cc",
+        "test": "If.ConditionalBranchesOnlyContainConstantNodes_ElseBranchExecution",
+        "notes": "Length-1 logical bool condition companion for the lowered select form; false must select the else tensor exactly."
+      },
+      "inputs": {
+        "cond": { "dtype": "bool", "shape": [1], "data": { "kind": "values", "values": [0] } },
+        "then_value": {
+          "dtype": "float32",
+          "shape": [2, 2],
+          "data": { "kind": "values", "values": [9.0, 8.0, 7.0, 6.0] }
+        },
+        "else_value": {
+          "dtype": "float32",
+          "shape": [2, 2],
+          "data": { "kind": "values", "values": [-1.0, -2.0, -3.0, -4.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 2], "tolerance": 0 } }
+    },
+    {
+      "name": "lowered_select_else_zero_sized",
+      "inputs": {
+        "cond": { "dtype": "uint32", "shape": [1], "data": { "kind": "values", "values": [0] } },
+        "then_value": { "dtype": "float32", "shape": [0], "data": { "kind": "values", "values": [] } },
+        "else_value": { "dtype": "float32", "shape": [0], "data": { "kind": "values", "values": [] } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [0] } }
+    },
+    {
+      "name": "lowered_noncanonical_cond_uint32_42_selects_then",
+      "inputs": {
+        "cond": { "dtype": "uint32", "shape": [1], "data": { "kind": "values", "values": [42] } },
+        "then_value": {
+          "dtype": "float32",
+          "shape": [4],
+          "data": { "kind": "values", "values": [3.0, 1.5, -2.5, 0.25] }
+        },
+        "else_value": {
+          "dtype": "float32",
+          "shape": [4],
+          "data": { "kind": "values", "values": [-9.0, 7.0, 4.25, -1.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [4], "tolerance": 0 } }
+    },
+    {
+      "name": "lowered_noncanonical_cond_uint32_max_selects_then",
+      "inputs": {
+        "cond": { "dtype": "uint32", "shape": [1], "data": { "kind": "values", "values": [4294967295] } },
+        "then_value": { "dtype": "float32", "shape": [3], "data": { "kind": "values", "values": [10.0, -5.0, 0.5] } },
+        "else_value": { "dtype": "float32", "shape": [3], "data": { "kind": "values", "values": [0.0, 0.0, 0.0] } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3], "tolerance": 0 } }
+    }
+  ]
+}