sync 2e7068faf55e

Browse files

Files changed (7) hide show

README.md +59 -0
build/webgpu/bench.json +26 -0
build/webgpu/dropout-vec4.wgsl.jinja +20 -0
build/webgpu/dropout.wgsl.jinja +46 -0
build/webgpu/manifest.json +152 -0
build/webgpu/metadata.json +19 -0
build/webgpu/test.json +127 -0

README.md CHANGED Viewed

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

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# ai.onnx.Dropout
+`ai.onnx`  ·  standard ONNX operator  ·  ONNX opset ≥ 13
+## Description
+Inference-only Dropout. Copies a floating-point tensor unchanged and optionally emits an all-true boolean mask. The `ratio` and `training_mode` inputs, seeded randomness, and training behavior are not supported by this package.
+See the [ONNX `Dropout` spec](https://onnx.ai/onnx/operators/onnx__Dropout.html) for the reference semantics.
+## Inputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `data` | `x` | `T` | — | — | The input floating-point tensor to copy unchanged. | required |
+## Outputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `output` | `y` | `T` | same as `data` | same as `data` | The output tensor, same shape as the input; equals the input in inference mode. | required |
+| `mask` | `mask` | `M` | same as `data` | same as `data` | Boolean mask indicating which elements were kept (all-ones in inference mode). | optional |
+## Type constraints
+| Variable | Allowed dtypes |
+| --- | --- |
+| `T` | `float32`, `float16` |
+| `M` | `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
+- [`dropout-vec4.wgsl.jinja`](build/webgpu/dropout-vec4.wgsl.jinja)
+- [`dropout.wgsl.jinja`](build/webgpu/dropout.wgsl.jinja)
+## Use with `@huggingface/kernels`
+The loader derives every required output's shape and logical dtype from the manifest contract and this call.
+It then allocates the result tensors automatically.
+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.Dropout", { version: 1 });
+const { y } = await kernel({ x: { data: xData, shape: [] } });
+```

build/webgpu/bench.json ADDED Viewed

	@@ -0,0 +1,26 @@

+{
+  "op": "ai.onnx.Dropout",
+  "cases": [
+    {
+      "name": "f32_16m",
+      "preset": "smoke",
+      "inputs": { "x": { "dtype": "float32", "shape": [16777216], "dist": "normal", "seed": 771, "scale": 1 } },
+      "outputs": { "y": { "dtype": "float32", "shape": [16777216] } },
+      "bench": { "primary": true, "metrics": [{ "type": "bandwidth", "value": "16777216 * 4 * 2" }] }
+    },
+    {
+      "name": "f32_1m",
+      "preset": "smoke",
+      "inputs": { "x": { "dtype": "float32", "shape": [1048576] } },
+      "outputs": { "y": { "dtype": "float32", "shape": [1048576] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "(numel(shapes.x) + numel(shapes.y)) * 4" }] }
+    },
+    {
+      "name": "f32_scalar_path_non_aligned",
+      "preset": "edge",
+      "inputs": { "x": { "dtype": "float32", "shape": [16777213], "dist": "normal", "seed": 883, "scale": 1 } },
+      "outputs": { "y": { "dtype": "float32", "shape": [16777213] } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "16777213 * 4 * 2" }] }
+    }
+  ]
+}

build/webgpu/dropout-vec4.wgsl.jinja ADDED Viewed

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

+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+// Vec4 inference Dropout: pure contiguous copy (Y = X) with mask = 1, 4 elements
+// per thread. A device-capped grid stride covers arbitrary output lengths.
+const COUNT: u32 = {{ source.count }}u;
+const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
+  let stride = nwg.x * WG;
+  for (var i = gid.x; i < COUNT; i = i + stride) {
+    y[i] = x[i];
+{% if hasMask %}
+    mask[i] = vec4<u32>(1u);
+{% endif %}
+  }
+}

build/webgpu/dropout.wgsl.jinja ADDED Viewed

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

+{% 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 %}
+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ 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(note="limit") }}
+  y[i] = x[i];
+{% if hasMask %}
+  mask[i] = 1u;
+{% endif %}
+}

build/webgpu/manifest.json ADDED Viewed

	@@ -0,0 +1,152 @@

+{
+  "domain": "ai.onnx",
+  "name": "Dropout",
+  "sinceVersion": 13,
+  "description": "Inference-only Dropout. Copies a floating-point tensor unchanged and optionally emits an all-true boolean mask. The `ratio` and `training_mode` inputs, seeded randomness, and training behavior are not supported by this package.",
+  "inputs": [{ "role": "data", "dtype": "T", "description": "The input floating-point tensor to copy unchanged." }],
+  "outputs": [
+    {
+      "role": "output",
+      "dtype": "T",
+      "rank": "ranks.data",
+      "description": "The output tensor, same shape as the input; equals the input in inference mode.",
+      "shape": "shapes.data"
+    },
+    {
+      "role": "mask",
+      "dtype": "M",
+      "rank": "ranks.data",
+      "shape": "shapes.data",
+      "optional": true,
+      "description": "Boolean mask indicating which elements were kept (all-ones in inference mode)."
+    }
+  ],
+  "typeConstraints": { "T": ["float32", "float16"], "M": ["bool"] },
+  "args": {
+    "x": { "kind": "tensor", "semantic": "data", "role": "x" },
+    "y": { "kind": "tensor", "semantic": "output", "role": "y" },
+    "mask": { "kind": "tensor", "semantic": "mask", "role": "mask", "required": false }
+  },
+  "tunables": { "WORKGROUP_SIZE": 256 },
+  "constants": { "usesF16": "dtypes.T == \"f16\"", "scalar": "dtypes.T" },
+  "bindingSets": {
+    "inferenceCopyVec4": [
+      {
+        "name": "x",
+        "arg": "x",
+        "semantic": "data",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$vectorScalar"
+      },
+      { "name": "y", "arg": "y", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$vectorScalar" }
+    ],
+    "inferenceCopyMaskVec4": [
+      {
+        "name": "x",
+        "arg": "x",
+        "semantic": "data",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$vectorScalar"
+      },
+      { "name": "y", "arg": "y", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$vectorScalar" },
+      { "name": "mask", "arg": "mask", "semantic": "mask", "buffer": { "type": "storage" }, "elementType": "vec4<u32>" }
+    ]
+  },
+  "variants": [
+    {
+      "id": "inference_copy_vec4",
+      "priority": 20,
+      "when": ["not present.mask", "numel(shapes.data) == numel(shapes.output)", "numel(shapes.output) > 0", "numel(shapes.output) % 4 == 0", "f16Ok(dtypes.T)"],
+      "constants": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"", "hasMask": false },
+      "passes": [
+        {
+          "id": "main",
+          "name": "Dropout.InferenceCopyVec4",
+          "source": { "shader": "dropout-vec4.wgsl.jinja", "inputs": { "count": "numel(shapes.output) / 4" } },
+          "bindings": "inferenceCopyVec4",
+          "dispatch": { "gridStride": "numel(shapes.output) / 4", "workgroupSize": "tunables.WORKGROUP_SIZE" }
+        }
+      ]
+    },
+    {
+      "id": "inference_copy_mask_vec4",
+      "priority": 20,
+      "when": ["present.mask", "numel(shapes.data) == numel(shapes.output)", "numel(shapes.mask) == numel(shapes.output)", "numel(shapes.output) > 0", "numel(shapes.output) % 4 == 0", "f16Ok(dtypes.T)"],
+      "constants": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"", "hasMask": true },
+      "passes": [
+        {
+          "id": "main",
+          "name": "Dropout.InferenceCopyMaskVec4",
+          "source": { "shader": "dropout-vec4.wgsl.jinja", "inputs": { "count": "numel(shapes.output) / 4" } },
+          "bindings": "inferenceCopyMaskVec4",
+          "dispatch": { "gridStride": "numel(shapes.output) / 4", "workgroupSize": "tunables.WORKGROUP_SIZE" }
+        }
+      ]
+    },
+    {
+      "id": "inference_copy",
+      "when": ["not present.mask", "numel(shapes.data) == numel(shapes.output)", "f16Ok(dtypes.T)"],
+      "constants": { "hasMask": false },
+      "passes": [
+        {
+          "id": "main",
+          "name": "Dropout.InferenceCopy",
+          "shader": "dropout.wgsl.jinja",
+          "bindings": [
+            {
+              "name": "x",
+              "arg": "x",
+              "semantic": "data",
+              "buffer": { "type": "read-only-storage" },
+              "elementType": "$scalar"
+            },
+            { "name": "y", "arg": "y", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+            {
+              "name": "params",
+              "semantic": "kernel.params",
+              "buffer": { "type": "uniform" },
+              "struct": {
+                "name": "Params",
+                "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.output)" }]
+              }
+            }
+          ],
+          "dispatch": { "threads": "numel(shapes.output)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
+        }
+      ]
+    },
+    {
+      "id": "inference_copy_mask",
+      "when": ["present.mask", "numel(shapes.data) == numel(shapes.output)", "numel(shapes.mask) == numel(shapes.output)", "f16Ok(dtypes.T)"],
+      "constants": { "hasMask": true },
+      "passes": [
+        {
+          "id": "main",
+          "name": "Dropout.InferenceCopyMask",
+          "shader": "dropout.wgsl.jinja",
+          "bindings": [
+            {
+              "name": "x",
+              "arg": "x",
+              "semantic": "data",
+              "buffer": { "type": "read-only-storage" },
+              "elementType": "$scalar"
+            },
+            { "name": "y", "arg": "y", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$scalar" },
+            { "name": "mask", "arg": "mask", "semantic": "mask", "buffer": { "type": "storage" }, "elementType": "u32" },
+            {
+              "name": "params",
+              "semantic": "kernel.params",
+              "buffer": { "type": "uniform" },
+              "struct": {
+                "name": "Params",
+                "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.output)" }]
+              }
+            }
+          ],
+          "dispatch": { "threads": "numel(shapes.output)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
+        }
+      ]
+    }
+  ]
+}

build/webgpu/metadata.json ADDED Viewed

	@@ -0,0 +1,19 @@

+{
+  "name": "ai.onnx.Dropout",
+  "id": "_ai_onnx_dropout_webgpu_1e25cf4",
+  "version": 1,
+  "license": "Apache-2.0",
+  "backend": { "type": "webgpu" },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "bench.json": "eCAMPumj2uRsPQESWgHKb7zvrb8NiQIfIjsdGBlbO4E=",
+      "dropout-vec4.wgsl.jinja": "UVWQwKy05KoAgvytw7s8nJPWWp/ZSXA6RXQaIPgXfV8=",
+      "dropout.wgsl.jinja": "ihCug0ebE5xa6fNeVS9Tm6h/2ZwezWoYWGJJA73PDk4=",
+      "manifest.json": "v1+Jd5K6LUjgUKpgZy1z/a5x87dauFbr5+AOPRqwICo=",
+      "test.json": "GLjoVvpcDWihcNm/5FeGh5oevXyHsyc/NzyTvmqsyes="
+    }
+  },
+  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
+  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.Dropout" }
+}

build/webgpu/test.json ADDED Viewed

	@@ -0,0 +1,127 @@

+{
+  "op": "ai.onnx.Dropout",
+  "fixtureArrays": {
+    "onnx_backend_input_x": [1.764052391052246, 0.40015721321105957, 0.978738009929657, 2.2408931255340576, 1.8675580024719238, -0.9772778749465942, 0.9500884413719177, -0.15135720372200012, -0.10321885347366333, 0.4105985164642334, 0.14404356479644775, 1.4542734622955322, 0.7610377073287964, 0.12167501449584961, 0.44386324286460876, 0.3336743414402008, 1.4940791130065918, -0.2051582634449005, 0.3130677044391632, -0.8540957570075989, -2.5529897212982178, 0.653618574142456, 0.8644362092018127, -0.7421650290489197, 2.269754648208618, -1.4543657302856445, 0.04575851559638977, -0.18718385696411133, 1.5327792167663574, 1.4693588018417358, 0.154947429895401, 0.37816253304481506, -0.8877857327461243, -1.980796456336975, -0.34791216254234314, 0.15634897351264954, 1.2302906513214111, 1.202379822731018, -0.38732680678367615, -0.302302747964859, -1.0485529899597168, -1.420017957687378, -1.7062702178955078, 1.950775384902954, -0.5096521973609924, -0.4380742907524109, -1.2527953386306763, 0.7774903774261475, -1.6138978004455566, -0.21274028718471527, -0.8954665660858154, 0.38690251111984253, -0.5108051300048828, -1.18063223361969, -0.02818222902715206, 0.4283318817615509, 0.06651721894741058, 0.30247190594673157, -0.6343221068382263, -0.3627411723136902]
+  },
+  "cases": [
+    {
+      "name": "float32_inference_copy",
+      "provenance": { "source": "onnxruntime/test/providers/cpu/nn/dropout_op_test.cc", "test": "Dropout.Opset10" },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "values", "values": [1.0, 2.0, 3.0, 5.0] } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 2] } }
+    },
+    {
+      "name": "ort_opset7_float32_inference_copy",
+      "provenance": { "source": "onnxruntime/test/providers/cpu/nn/dropout_op_test.cc", "test": "Dropout.Opset7" },
+      "inputs": {
+        "x": { "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": "ort_opset13_optional_bool_mask_exact",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/nn/dropout_op_test.cc",
+        "test": "Dropout.WithOptionalOutputOpset10",
+        "notes": "Same optional-mask behavior validated through the current-schema ORT path, but with the ONNX-native bool mask dtype."
+      },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [2, 2], "data": { "kind": "values", "values": [1.0, 2.0, 3.0, 5.0] } }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [2, 2] },
+        "mask": { "dtype": "bool", "shape": [2, 2], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "float16_inference_copy",
+      "inputs": {
+        "x": { "dtype": "float16", "shape": [2, 2], "data": { "kind": "values", "values": [1.0, -2.0, 3.5, 5.0] } }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [2, 2], "tolerance": 0.001 } }
+    },
+    {
+      "name": "float32_scalar_inference_copy",
+      "inputs": { "x": { "dtype": "float32", "shape": [], "data": { "kind": "values", "values": [-4.5] } } },
+      "outputs": { "y": { "dtype": "float32", "shape": [], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "float32_nan_and_infinity_copy",
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [5],
+          "data": { "kind": "values", "values": ["NaN", "-Infinity", 0.0, "Infinity", 3.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [5], "tolerance": 0.000001, "allowNaN": true } }
+    },
+    {
+      "name": "float16_zero_sized_copy",
+      "inputs": { "x": { "dtype": "float16", "shape": [2, 0, 3], "data": { "kind": "values", "values": [] } } },
+      "outputs": { "y": { "dtype": "float16", "shape": [2, 0, 3], "tolerance": 0.001 } }
+    },
+    {
+      "name": "onnx_backend_default_rank3_copy",
+      "provenance": {
+        "source": "cmake/external/onnx/onnx/backend/test/data/node/test_dropout_default",
+        "notes": "Imports the inference-copy output; optional ratio/training-mode inputs and mask output are omitted by this framework variant. The legacy test_dropout_random_old vector projects to the same inference-only request."
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 4, 5],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/onnx_backend_input_x" } }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [3, 4, 5], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "onnx_backend_default_mask_rank3",
+      "provenance": {
+        "source": "cmake/external/onnx/onnx/backend/test/data/node/test_dropout_default_mask",
+        "notes": "Inference mode emits the all-true ONNX bool mask."
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [3, 4, 5],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/onnx_backend_input_x" } }
+        }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [3, 4, 5], "tolerance": 0.000001 },
+        "mask": { "dtype": "bool", "shape": [3, 4, 5], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "onnx_backend_default_old_copy",
+      "provenance": {
+        "source": "cmake/external/onnx/onnx/backend/test/data/node/test_dropout_default_old",
+        "notes": "Older-schema Dropout inference-copy fixture represented with the project inference-only variant."
+      },
+      "inputs": { "x": { "dtype": "float32", "shape": [3], "data": { "kind": "values", "values": [-1.0, 0.0, 1.0] } } },
+      "outputs": { "y": { "dtype": "float32", "shape": [3], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "f32_scalar_mask_path_non_aligned",
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [5], "data": { "kind": "values", "values": [1.0, -2.0, 3.5, 0.0, -0.5] } }
+      },
+      "outputs": {
+        "y": { "dtype": "float32", "shape": [5], "tolerance": 0.000001 },
+        "mask": { "dtype": "bool", "shape": [5], "tolerance": 0 }
+      }
+    },
+    {
+      "name": "f16_scalar_mask_path_numel3",
+      "inputs": { "x": { "dtype": "float16", "shape": [3], "data": { "kind": "values", "values": [1.0, -2.0, 3.5] } } },
+      "outputs": {
+        "y": { "dtype": "float16", "shape": [3], "tolerance": 0.001 },
+        "mask": { "dtype": "bool", "shape": [3], "tolerance": 0 }
+      }
+    }
+  ]
+}