sync 91d990483a17

Browse files

Files changed (7) hide show

README.md +12 -8
build/webgpu/bench.json +2 -3
build/webgpu/manifest.json +60 -168
build/webgpu/metadata.json +16 -8
build/webgpu/mod-vec4.wgsl.jinja +26 -28
build/webgpu/mod.wgsl.jinja +30 -26
build/webgpu/test.json +9 -7

README.md CHANGED Viewed

@@ -18,16 +18,16 @@ See the [ONNX `Mod` spec](https://onnx.ai/onnx/operators/onnx__Mod.html) for the
 ## Inputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `A` | `a` | `T` | — | — | Dividend tensor. | required |
-| `B` | `b` | `T` | — | — | Divisor tensor. | required |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `C` | `c` | `T` | derived | broadcast result of `A` and `B` | Remainder tensor; same shape as the broadcast result of A and B. | required |
 ## Attributes
@@ -45,7 +45,7 @@ Default values (overridable per request):
 ## 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
@@ -54,10 +54,14 @@ Default values (overridable per request):
 ## 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.

 ## Inputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `a` | `A` | `T` | — | — | Dividend tensor. | required |
+| `b` | `B` | `T` | — | — | Divisor tensor. | required |
 ## Outputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `c` | `C` | `T` | derived | broadcast result of `a` and `b` | Remainder tensor; same shape as the broadcast result of A and B. | required |
 ## Attributes
 ## 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
 ## 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.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

build/webgpu/bench.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.Mod",
   "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
   "cases": [
     {
@@ -58,7 +57,7 @@
       "preset": "smoke",
       "vars": { "dtype": "int32", "count": 1048576 },
       "inputs": {
-        "a": { "dtype": "int32", "shape": [1048576], "dist": "normal", "seed": 713, "scale": 1000 },
         "b": { "dtype": "int32", "shape": [1048576], "data": { "kind": "constant", "value": 7 } }
       },
       "outputs": { "c": { "dtype": "int32", "shape": [1048576] } },
@@ -69,7 +68,7 @@
       "preset": "stress",
       "attrs": { "fmod": 0 },
       "inputs": {
-        "a": { "dtype": "int32", "shape": [2049, 2051], "dist": "normal", "seed": 801, "scale": 1000 },
         "b": { "dtype": "int32", "shape": [2051], "data": { "kind": "constant", "value": 7 } }
       },
       "outputs": { "c": { "dtype": "int32", "shape": [2049, 2051] } },

 {
   "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
   "cases": [
     {
       "preset": "smoke",
       "vars": { "dtype": "int32", "count": 1048576 },
       "inputs": {
+        "a": { "dtype": "int32", "shape": [1048576], "dist": "uniform", "seed": 713, "min": -1000, "max": 1001 },
         "b": { "dtype": "int32", "shape": [1048576], "data": { "kind": "constant", "value": 7 } }
       },
       "outputs": { "c": { "dtype": "int32", "shape": [1048576] } },
       "preset": "stress",
       "attrs": { "fmod": 0 },
       "inputs": {
+        "a": { "dtype": "int32", "shape": [2049, 2051], "dist": "uniform", "seed": 801, "min": -1000, "max": 1001 },
         "b": { "dtype": "int32", "shape": [2051], "data": { "kind": "constant", "value": 7 } }
       },
       "outputs": { "c": { "dtype": "int32", "shape": [2049, 2051] } },

build/webgpu/manifest.json CHANGED Viewed

@@ -2,222 +2,114 @@
   "domain": "ai.onnx",
   "name": "Mod",
   "sinceVersion": 13,
-  "description": "Performs elementwise binary modulo on tensors `A` and `B` with multidirectional broadcasting. When `fmod` is `0` (default), applies Python-style `%` with the sign of the divisor; when `fmod` is `1`, applies C-style `fmod` with the sign of the dividend.",
-  "inputs": [
-    { "role": "A", "dtype": "T", "description": "Dividend tensor." },
-    { "role": "B", "dtype": "T", "description": "Divisor tensor." }
-  ],
-  "outputs": [
-    {
-      "role": "C",
-      "dtype": "T",
-      "rank": "max(ranks.A, ranks.B)",
-      "description": "Remainder tensor; same shape as the broadcast result of A and B.",
-      "shape": "broadcastShape(shapes.A, shapes.B)"
-    }
-  ],
-  "attributes": { "fmod": 0 },
-  "attributeDescriptions": {
-    "fmod": "Controls the modulo mode: `0` (default) uses Python-style integer mod (sign of divisor); `1` uses C-style `fmod` (sign of dividend, floating-point types only)."
   },
   "attributeConstraints": { "fmod": { "values": [0, 1] } },
   "typeConstraints": { "T": ["float32", "float16", "int32", "uint32", "int16", "int8", "uint8"] },
-  "args": {
-    "a": { "kind": "tensor", "semantic": "A", "role": "input" },
-    "b": { "kind": "tensor", "semantic": "B", "role": "input" },
-    "c": { "kind": "tensor", "semantic": "C", "role": "output" }
-  },
-  "tunables": { "WORKGROUP_SIZE": 256 },
-  "derive": { "floatFmodContract": "attrs.fmod != 0 or not (dtypes.T == \"f32\" or dtypes.T == \"f16\")" },
-  "constants": {
     "scalar": "dtypes.T",
-    "usesF16": "dtypes.T == \"f16\"",
     "isU32": "dtypes.T == \"u32\"",
     "isI32": "dtypes.T == \"i32\"",
     "fmod": "attrs.fmod != 0"
   },
   "variants": [
     {
       "id": "last_axis_broadcast_vec4",
       "priority": 25,
-      "when": ["ranks.C >= 1", "sameShape(shapes.A, shapes.C)", "ranks.B == 1", "dim(shapes.B, 0) == dim(shapes.C, ranks.C - 1)", "numel(shapes.C) > 0", "numel(shapes.C) % 4 == 0", "f16Ok(dtypes.T)", "floatFmodContract"],
-      "constants": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Mod.lastAxisBroadcastVec4",
-          "source": { "shader": "mod-vec4.wgsl.jinja", "inputs": { "perAxisB": true, "axisDim": "dim(shapes.B, 0)" } },
-          "bindings": [
-            {
-              "name": "a",
-              "arg": "a",
-              "semantic": "A",
-              "buffer": { "type": "read-only-storage" },
-              "elementType": "$vectorScalar"
-            },
-            {
-              "name": "b",
-              "arg": "b",
-              "semantic": "B",
-              "buffer": { "type": "read-only-storage" },
-              "elementType": "$scalar"
-            },
-            {
-              "name": "c",
-              "arg": "c",
-              "semantic": "C",
-              "buffer": { "type": "storage" },
-              "elementType": "$vectorScalar"
-            },
-            {
-              "name": "params",
-              "semantic": "kernel.params",
-              "buffer": { "type": "uniform" },
-              "struct": {
-                "name": "Params",
-                "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.C) / 4" }]
-              }
-            }
-          ],
-          "dispatch": { "threads": "numel(shapes.C) / 4", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "scalar_b_vec4",
       "priority": 30,
-      "when": ["sameShape(shapes.A, shapes.C)", "numel(shapes.B) == 1", "numel(shapes.C) > 0", "numel(shapes.C) % 4 == 0", "f16Ok(dtypes.T)", "floatFmodContract"],
-      "constants": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Mod.scalarBVec4",
-          "source": { "shader": "mod-vec4.wgsl.jinja", "inputs": { "scalarB": true } },
-          "bindings": [
-            {
-              "name": "a",
-              "arg": "a",
-              "semantic": "A",
-              "buffer": { "type": "read-only-storage" },
-              "elementType": "$vectorScalar"
-            },
-            {
-              "name": "b",
-              "arg": "b",
-              "semantic": "B",
-              "buffer": { "type": "read-only-storage" },
-              "elementType": "$scalar",
-              "length": 1
-            },
-            {
-              "name": "c",
-              "arg": "c",
-              "semantic": "C",
-              "buffer": { "type": "storage" },
-              "elementType": "$vectorScalar"
-            },
-            {
-              "name": "params",
-              "semantic": "kernel.params",
-              "buffer": { "type": "uniform" },
-              "struct": {
-                "name": "Params",
-                "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.C) / 4" }]
-              }
-            }
-          ],
-          "dispatch": { "threads": "numel(shapes.C) / 4", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "same_shape_vec4",
       "priority": 26,
-      "when": ["sameShape(shapes.A, shapes.C)", "sameShape(shapes.B, shapes.C)", "numel(shapes.C) > 0", "numel(shapes.C) % 4 == 0", "f16Ok(dtypes.T)", "floatFmodContract"],
-      "constants": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Mod.vec4",
-          "source": { "shader": "mod-vec4.wgsl.jinja", "inputs": {} },
-          "bindings": [
-            {
-              "name": "a",
-              "arg": "a",
-              "semantic": "A",
-              "buffer": { "type": "read-only-storage" },
-              "elementType": "$vectorScalar"
-            },
-            {
-              "name": "b",
-              "arg": "b",
-              "semantic": "B",
-              "buffer": { "type": "read-only-storage" },
-              "elementType": "$vectorScalar"
-            },
-            {
-              "name": "c",
-              "arg": "c",
-              "semantic": "C",
-              "buffer": { "type": "storage" },
-              "elementType": "$vectorScalar"
-            },
-            {
-              "name": "params",
-              "semantic": "kernel.params",
-              "buffer": { "type": "uniform" },
-              "struct": {
-                "name": "Params",
-                "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.C) / 4" }]
-              }
-            }
-          ],
-          "dispatch": { "threads": "numel(shapes.C) / 4", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "broadcast",
-      "when": ["ranks.A >= 0", "ranks.B >= 0", "ranks.C >= 0", "ranks.A <= ranks.C", "ranks.B <= ranks.C", "f16Ok(dtypes.T)", "floatFmodContract"],
       "passes": [
         {
           "id": "main",
           "name": "Mod",
-          "source": {
-            "shader": "mod.wgsl.jinja",
-            "inputs": {
-              "aShape": "shapes.A",
-              "bShape": "shapes.B",
-              "cShape": "shapes.C",
-              "aRank": "ranks.A",
-              "bRank": "ranks.B",
-              "cRank": "ranks.C"
-            }
           },
           "bindings": [
-            {
-              "name": "a",
-              "arg": "a",
-              "semantic": "A",
-              "buffer": { "type": "read-only-storage" },
-              "elementType": "$scalar"
-            },
-            {
-              "name": "b",
-              "arg": "b",
-              "semantic": "B",
-              "buffer": { "type": "read-only-storage" },
-              "elementType": "$scalar"
-            },
-            { "name": "c", "arg": "c", "semantic": "C", "buffer": { "type": "storage" }, "elementType": "$scalar" },
-            {
-              "name": "params",
-              "semantic": "kernel.params",
-              "buffer": { "type": "uniform" },
-              "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.C)" }] }
-            }
           ],
-          "dispatch": { "threads": "numel(shapes.C)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     }

   "domain": "ai.onnx",
   "name": "Mod",
   "sinceVersion": 13,
+  "inputs": { "a": { "onnx": "A", "dtype": "T" }, "b": { "onnx": "B", "dtype": "T" } },
+  "outputs": {
+    "c": { "onnx": "C", "dtype": "T", "rank": "max(ranks.a, ranks.b)", "shape": "broadcastShape(shapes.a, shapes.b)" }
   },
+  "attributes": { "fmod": { "default": 0 } },
   "attributeConstraints": { "fmod": { "values": [0, 1] } },
   "typeConstraints": { "T": ["float32", "float16", "int32", "uint32", "int16", "int8", "uint8"] },
+  "tunables": { "WORKGROUP_SIZE": { "default": 256 } },
+  "derive": {
+    "floatFmodContract": "attrs.fmod != 0 or not (dtypes.T == \"f32\" or dtypes.T == \"f16\")",
     "scalar": "dtypes.T",
     "isU32": "dtypes.T == \"u32\"",
     "isI32": "dtypes.T == \"i32\"",
     "fmod": "attrs.fmod != 0"
   },
+  "when": ["f16Ok(dtypes.T)", "floatFmodContract"],
+  "bindings": {
+    "a": { "buffer": "read-only-storage", "elementType": "$vectorScalar" },
+    "c": { "buffer": "storage", "elementType": "$vectorScalar" },
+    "params": { "buffer": "uniform", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.c) / 4" }] }
+  },
   "variants": [
     {
       "id": "last_axis_broadcast_vec4",
       "priority": 25,
+      "when": ["ranks.c >= 1", "sameShape(shapes.a, shapes.c)", "ranks.b == 1", "dim(shapes.b, 0) == dim(shapes.c, ranks.c - 1)", "numel(shapes.c) > 0", "numel(shapes.c) % 4 == 0"],
+      "derive": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Mod.lastAxisBroadcastVec4",
+          "shader": "mod-vec4.wgsl.jinja",
+          "derive": { "perAxisB": true, "axisDimSpec": "dim(shapes.b, 0)" },
+          "bindings": ["a", "b", "c", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.c) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.c) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "scalar_b_vec4",
       "priority": 30,
+      "when": ["sameShape(shapes.a, shapes.c)", "numel(shapes.b) == 1", "numel(shapes.c) > 0", "numel(shapes.c) % 4 == 0"],
+      "derive": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Mod.scalarBVec4",
+          "shader": "mod-vec4.wgsl.jinja",
+          "derive": { "scalarB": true },
+          "bindings": ["a", { "arg": "b", "elementType": "$scalar", "length": 1 }, "c", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.c) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.c) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "same_shape_vec4",
       "priority": 26,
+      "when": ["sameShape(shapes.a, shapes.c)", "sameShape(shapes.b, shapes.c)", "numel(shapes.c) > 0", "numel(shapes.c) % 4 == 0"],
+      "derive": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Mod.vec4",
+          "shader": "mod-vec4.wgsl.jinja",
+          "bindings": ["a", { "arg": "b", "elementType": "$vectorScalar" }, "c", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.c) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.c) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "broadcast",
+      "when": ["ranks.a >= 0", "ranks.b >= 0", "ranks.c >= 0", "ranks.a <= ranks.c", "ranks.b <= ranks.c"],
       "passes": [
         {
           "id": "main",
           "name": "Mod",
+          "shader": "mod.wgsl.jinja",
+          "derive": {
+            "aShape": "shapes.a",
+            "bShape": "shapes.b",
+            "cShape": "shapes.c",
+            "aRank": "ranks.a",
+            "bRank": "ranks.b",
+            "cRank": "ranks.c"
           },
           "bindings": [
+            { "arg": "a" },
+            "b",
+            { "arg": "c" },
+            { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.c)" }] }
           ],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.c)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.c)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,19 +1,27 @@
 {
   "name": "ai.onnx.Mod",
-  "id": "_ai_onnx_mod_webgpu_4a7f8df",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "QdfVC09wyipoUAu6hPHgQsOE7mtlr2T9OZJffvCh8YQ=",
-      "manifest.json": "EWrTl0E4Q3pRgojKEehDCETQSaGUYd6LvDy+6DtWFCI=",
-      "mod-vec4.wgsl.jinja": "OoOWti3TFUtnneaDtonReshC71gzRLiUQOOK4IZ9P7w=",
-      "mod.wgsl.jinja": "R3tnz/Y9NdD2Ut/jujKAUl2H2Y2CYSu9LpKXdUYJsz4=",
-      "test.json": "3XigfrMBzVhWpcZHr5WWptcT2N637Wb7iEjzEValb4A="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.Mod" }
 }

 {
   "name": "ai.onnx.Mod",
+  "id": "_ai_onnx_mod_webgpu_821fe84",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "u7NGAyuRdKHN1czKMTxOofOGp/1+B7oS8d3EHcUjAcc=",
+      "manifest.json": "RLwkzGX5SssKyYjn+cZRb9VMTal+0BdSoSyxjtZhBXc=",
+      "mod-vec4.wgsl.jinja": "m5J+jy22HBFcUnuKvocYEGJdjpQ1xnOKyR/q7H9dxew=",
+      "mod.wgsl.jinja": "cp6hTSRn1C61dnn8jXx4YtvYu/4qnW3Vu9ccZLct2YE=",
+      "test.json": "1LykXvYWjgy9h3UgENoljeihrg0vihh13G6x4EVH8Cg="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "last_axis_broadcast_vec4": ["mod-vec4.wgsl.jinja"],
+      "scalar_b_vec4": ["mod-vec4.wgsl.jinja"],
+      "same_shape_vec4": ["mod-vec4.wgsl.jinja"],
+      "broadcast": ["mod.wgsl.jinja"]
+    }
+  }
 }

build/webgpu/mod-vec4.wgsl.jinja CHANGED Viewed

@@ -1,39 +1,36 @@
 {% 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 }}
 {% if not isU32 and not isI32 %}
@@ -46,7 +43,8 @@ enable f16;
 fn fmod_trunc(a: f32, b: f32) -> f32 {
   let aa = abs(a);
   let ab = abs(b);
-  // b == 0, or a non-finite: defer to the hardware formula for IEEE NaN/Inf.
   if (!(ab > 0.0) || bitcast<u32>(aa) >= 0x7f800000u) {
     return a - b * trunc(a / b);
   }
@@ -70,22 +68,22 @@ fn fmod_trunc(a: f32, b: f32) -> f32 {
 {% endif %}
 @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="vec4-limit") }}
-{% if source.perAxisB is defined and source.perAxisB %}
   // B broadcasts over the innermost axis. A/C stay 128-bit packed even when
   // an odd row width makes vec4 groups cross a row boundary; derive each
   // lane's wrapped B coordinate from the flat scalar position.
-  let b0 = (i * 4u) % {{ source.axisDim }}u;
-  let b1 = (b0 + 1u) % {{ source.axisDim }}u;
-  let b2 = (b0 + 2u) % {{ source.axisDim }}u;
-  let b3 = (b0 + 3u) % {{ source.axisDim }}u;
 {% endif %}
 {% if isU32 %}
-  c[i] = a[i] % {% if source.scalarB is defined and source.scalarB %}vec4<u32>(b[0]){% elif source.perAxisB is defined and source.perAxisB %}vec4<u32>(b[b0], b[b1], b[b2], b[b3]){% else %}b[i]{% endif %};
 {% elif isI32 %}
   let av = a[i];
-  let bv = {% if source.scalarB is defined and source.scalarB %}vec4<i32>(b[0]){% elif source.perAxisB is defined and source.perAxisB %}vec4<i32>(b[b0], b[b1], b[b2], b[b3]){% else %}b[i]{% endif %};
   let rt = av % bv;
 {% if fmod %}
   c[i] = rt;
@@ -94,10 +92,10 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
 {% endif %}
 {% else %}
   let av = vec4<f32>(a[i]);
-{% if source.perAxisB is defined and source.perAxisB %}
   let bv = vec4<f32>(f32(b[b0]), f32(b[b1]), f32(b[b2]), f32(b[b3]));
 {% else %}
-  let bv = {% if source.scalarB is defined and source.scalarB %}vec4<f32>(f32(b[0])){% else %}vec4<f32>(b[i]){% endif %};
 {% endif %}
   let rt = vec4<f32>(
     fmod_trunc(av.x, bv.x),

 {% 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 }}
 {% if not isU32 and not isI32 %}
 fn fmod_trunc(a: f32, b: f32) -> f32 {
   let aa = abs(a);
   let ab = abs(b);
+  // For a zero divisor or non-finite dividend, use the WGSL expression directly
+  // so its NaN and infinity propagation is preserved.
   if (!(ab > 0.0) || bitcast<u32>(aa) >= 0x7f800000u) {
     return a - b * trunc(a / b);
   }
 {% endif %}
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 {{ flat_index_2d(note="vec4-limit") }}
+{% if perAxisB is defined and perAxisB %}
   // B broadcasts over the innermost axis. A/C stay 128-bit packed even when
   // an odd row width makes vec4 groups cross a row boundary; derive each
   // lane's wrapped B coordinate from the flat scalar position.
+  let b0 = (i * 4u) % {{ axisDimSpec }}u;
+  let b1 = (b0 + 1u) % {{ axisDimSpec }}u;
+  let b2 = (b0 + 2u) % {{ axisDimSpec }}u;
+  let b3 = (b0 + 3u) % {{ axisDimSpec }}u;
 {% endif %}
 {% if isU32 %}
+  c[i] = a[i] % {% if scalarB is defined and scalarB %}vec4<u32>(b[0]){% elif perAxisB is defined and perAxisB %}vec4<u32>(b[b0], b[b1], b[b2], b[b3]){% else %}b[i]{% endif %};
 {% elif isI32 %}
   let av = a[i];
+  let bv = {% if scalarB is defined and scalarB %}vec4<i32>(b[0]){% elif perAxisB is defined and perAxisB %}vec4<i32>(b[b0], b[b1], b[b2], b[b3]){% else %}b[i]{% endif %};
   let rt = av % bv;
 {% if fmod %}
   c[i] = rt;
 {% endif %}
 {% else %}
   let av = vec4<f32>(a[i]);
+{% if perAxisB is defined and perAxisB %}
   let bv = vec4<f32>(f32(b[b0]), f32(b[b1]), f32(b[b2]), f32(b[b3]));
 {% else %}
+  let bv = {% if scalarB is defined and scalarB %}vec4<f32>(f32(b[0])){% else %}vec4<f32>(b[i]){% endif %};
 {% endif %}
   let rt = vec4<f32>(
     fmod_trunc(av.x, bv.x),

build/webgpu/mod.wgsl.jinja CHANGED Viewed

@@ -1,30 +1,30 @@
 {% 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;
   }
@@ -92,14 +92,17 @@ fn {{ fn_name }}({% if out_numel != 0 and op_numel != 1 %}out_index: u32{% endif
 {% endif %}
 {{ offset_fn(fn_name, opShape, opRank, op_same.value, op_numel.value, outShape, outRank, out_numel.value) }}
 {%- endmacro %}{% macro binary_broadcast_offsets() %}
-{{ broadcast_offset_fn("a_offset", source.aShape, source.aRank, source.cShape, source.cRank) }}
-{{ broadcast_offset_fn("b_offset", source.bShape, source.bRank, source.cShape, source.cRank) }}
 {%- endmacro %}
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
@@ -115,7 +118,8 @@ enable f16;
 fn fmod_trunc(a: f32, b: f32) -> f32 {
   let aa = abs(a);
   let ab = abs(b);
-  // b == 0, or a non-finite: defer to the hardware formula for IEEE NaN/Inf.
   if (!(ab > 0.0) || bitcast<u32>(aa) >= 0x7f800000u) {
     return a - b * trunc(a / b);
   }
@@ -139,13 +143,13 @@ fn fmod_trunc(a: f32, b: f32) -> f32 {
 {% endif %}
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
-fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
 {{ flat_index_2d() }}
 {% if isU32 %}
-  c[i] = a[{{ broadcast_offset_call("a_offset", source.aShape, source.cShape, "i") }}] % b[{{ broadcast_offset_call("b_offset", source.bShape, source.cShape, "i") }}];
 {% elif isI32 %}
-  let av = a[{{ broadcast_offset_call("a_offset", source.aShape, source.cShape, "i") }}];
-  let bv = b[{{ broadcast_offset_call("b_offset", source.bShape, source.cShape, "i") }}];
   let rt = av % bv;
 {% if fmod %}
   c[i] = rt;
@@ -153,8 +157,8 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
   c[i] = select(rt, rt + bv, rt != 0 && (rt < 0) != (bv < 0));
 {% endif %}
 {% else %}
-  let av = f32(a[{{ broadcast_offset_call("a_offset", source.aShape, source.cShape, "i") }}]);
-  let bv = f32(b[{{ broadcast_offset_call("b_offset", source.bShape, source.cShape, "i") }}]);
   let rt = fmod_trunc(av, bv);
   c[i] = {{ scalar }}(rt);
 {% endif %}

 {% 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 %}
 {{ offset_fn(fn_name, opShape, opRank, op_same.value, op_numel.value, outShape, outRank, out_numel.value) }}
 {%- endmacro %}{% macro binary_broadcast_offsets() %}
+{% set aShape = aShape | default([]) %}
+{% set aRank = aRank | default(0) %}
+{% set bShape = bShape | default([]) %}
+{% set bRank = bRank | default(0) %}
+{% set cShape = cShape | default([]) %}
+{% set cRank = cRank | default(0) %}
+{{ broadcast_offset_fn("a_offset", aShape, aRank, cShape, cRank) }}
+{{ broadcast_offset_fn("b_offset", bShape, bRank, cShape, cRank) }}
 {%- endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 fn fmod_trunc(a: f32, b: f32) -> f32 {
   let aa = abs(a);
   let ab = abs(b);
+  // For a zero divisor or non-finite dividend, use the WGSL expression directly
+  // so its NaN and infinity propagation is preserved.
   if (!(ab > 0.0) || bitcast<u32>(aa) >= 0x7f800000u) {
     return a - b * trunc(a / b);
   }
 {% endif %}
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 {{ flat_index_2d() }}
 {% if isU32 %}
+  c[i] = a[{{ broadcast_offset_call("a_offset", aShape, cShape, "i") }}] % b[{{ broadcast_offset_call("b_offset", bShape, cShape, "i") }}];
 {% elif isI32 %}
+  let av = a[{{ broadcast_offset_call("a_offset", aShape, cShape, "i") }}];
+  let bv = b[{{ broadcast_offset_call("b_offset", bShape, cShape, "i") }}];
   let rt = av % bv;
 {% if fmod %}
   c[i] = rt;
   c[i] = select(rt, rt + bv, rt != 0 && (rt < 0) != (bv < 0));
 {% endif %}
 {% else %}
+  let av = f32(a[{{ broadcast_offset_call("a_offset", aShape, cShape, "i") }}]);
+  let bv = f32(b[{{ broadcast_offset_call("b_offset", bShape, cShape, "i") }}]);
   let rt = fmod_trunc(av, bv);
   c[i] = {{ scalar }}(rt);
 {% endif %}

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.Mod",
   "fixtureArrays": {
     "ort_int32_rank3_scalar_divisor_broadcast_input_a": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29]
   },
@@ -79,7 +78,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/math/element_wise_ops_test.cc",
         "test": "ModOpTest.Fmod_float_mixed_sign",
-        "notes": "Scalar broadcast companion: subnormal dividends smaller than the divisor are valid nonzero remainders."
       },
       "attrs": { "fmod": 1 },
       "inputs": {
@@ -457,7 +456,7 @@
       "outputs": { "c": { "dtype": "int32", "shape": [6] } },
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_mod_int64_fmod",
-        "notes": "ONNX int64 tensors use framework int32/uint32 slots where representable."
       }
     },
     {
@@ -579,7 +578,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/math/element_wise_ops_test.cc",
         "test": "ModOpTest.Int32_mod_bcast",
-        "notes": "The fmod=1 half of the same route: truncated remainder, so the sign follows the dividend and the floor-mod adjust must NOT fire. Left unpinned so the reference is exercised on the arm its pinned sibling covers only in the kernel."
       },
       "inputs": {
         "a": { "dtype": "int32", "shape": [8], "data": { "kind": "values", "values": [-9, 7, -7, 15, 0, -1, 22, -22] } },
@@ -681,7 +680,10 @@
         },
         "b": { "dtype": "float16", "shape": [4], "data": { "kind": "values", "values": [3.0, 0.5, 3.0, 0.125] } }
       },
-      "outputs": { "c": { "dtype": "float16", "shape": [4], "tolerance": 0.5 } }
     },
     {
       "name": "u32_scalar_b_non_vec4_large_values",
@@ -727,7 +729,7 @@
     {
       "name": "last_axis_broadcast_vec4_odd_row_crossing",
       "provenance": {
-        "notes": "Four vec4 groups cross odd-width row boundaries, exercising wrapped per-lane B indices. Expected values come from the TypeScript Mod reference."
       },
       "attrs": { "fmod": 1 },
       "inputs": {
@@ -743,7 +745,7 @@
     {
       "name": "last_axis_broadcast_vec4_f16_odd_row_crossing",
       "provenance": {
-        "notes": "f16 companion for the per-lane broadcast path; locks explicit widening of scalar B lanes before fmod arithmetic."
       },
       "attrs": { "fmod": 1 },
       "inputs": {

 {
   "fixtureArrays": {
     "ort_int32_rank3_scalar_divisor_broadcast_input_a": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29]
   },
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/math/element_wise_ops_test.cc",
         "test": "ModOpTest.Fmod_float_mixed_sign",
+        "notes": "On the scalar broadcast path, subnormal dividends smaller than the divisor produce valid nonzero remainders."
       },
       "attrs": { "fmod": 1 },
       "inputs": {
       "outputs": { "c": { "dtype": "int32", "shape": [6] } },
       "provenance": {
         "source": "cmake/external/onnx/onnx/backend/test/data/node/test_mod_int64_fmod",
+        "notes": "The int32 inputs exactly represent the source ONNX int64 values."
       }
     },
     {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/math/element_wise_ops_test.cc",
         "test": "ModOpTest.Int32_mod_bcast",
+        "notes": "With fmod=1, truncated remainder follows the dividend's sign and must not apply the floor-mod adjustment."
       },
       "inputs": {
         "a": { "dtype": "int32", "shape": [8], "data": { "kind": "values", "values": [-9, 7, -7, 15, 0, -1, 22, -22] } },
         },
         "b": { "dtype": "float16", "shape": [4], "data": { "kind": "values", "values": [3.0, 0.5, 3.0, 0.125] } }
       },
+      "outputs": { "c": { "dtype": "float16", "shape": [4], "tolerance": 0.01 } },
+      "provenance": {
+        "notes": "C-style fmod at the top of the float16 range needs float32 division and truncation; the remainders are exact, so the tolerance only covers the final float16 rounding."
+      }
     },
     {
       "name": "u32_scalar_b_non_vec4_large_values",
     {
       "name": "last_axis_broadcast_vec4_odd_row_crossing",
       "provenance": {
+        "notes": "Four vec4 groups cross odd-width row boundaries, exercising wrapped per-lane B indices with distinct remainders."
       },
       "attrs": { "fmod": 1 },
       "inputs": {
     {
       "name": "last_axis_broadcast_vec4_f16_odd_row_crossing",
       "provenance": {
+        "notes": "Float16 per-lane broadcasting requires scalar B lanes to be widened before fmod arithmetic."
       },
       "attrs": { "fmod": 1 },
       "inputs": {