sync 91d990483a17

Browse files

Files changed (9) hide show

README.md +13 -10
build/webgpu/bench.json +0 -1
build/webgpu/manifest.json +114 -141
build/webgpu/metadata.json +21 -10
build/webgpu/resize-coord-transform.wgsl.jinja +63 -21
build/webgpu/resize-generic.wgsl.jinja +21 -27
build/webgpu/resize-linear-2x-stencil.wgsl.jinja +58 -19
build/webgpu/resize-nearest-integer-scale.wgsl.jinja +8 -11
build/webgpu/test.json +32 -1

README.md CHANGED Viewed

@@ -18,16 +18,16 @@ See the [ONNX `Upsample` spec](https://onnx.ai/onnx/operators/onnx__Upsample.htm
 ## Inputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `X` | `x` | `T` | — | — | Input tensor to upsample. | required |
-| `scales` | `scales` | `S` | `1` | — | Per-dimension scale factors, one value per input dimension. | required |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `Y` | `y` | `T` | same as `X` | — | Upsampled output tensor; each dimension is `floor(input_dimension * scale)`. | required |
 ## Attributes
@@ -46,7 +46,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
@@ -57,15 +57,18 @@ Default values (overridable per request):
 ## 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.

 ## Inputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `x` | `X` | `T` | — | — | Input tensor to upsample. | required |
+| `scales` | — | `S` | `1` | — | Per-dimension scale factors, one value per input dimension. | required |
 ## Outputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `y` | `Y` | `T` | same as `x` | — | Upsampled output tensor; each dimension is `floor(input_dimension * scale)`. | 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
+```
+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.
+This example supplies explicit metadata for:
+- `y`
 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.Upsample",
   "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
   "cases": [
     {

 {
   "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
   "cases": [
     {

build/webgpu/manifest.json CHANGED Viewed

@@ -2,222 +2,195 @@
   "domain": "ai.onnx",
   "name": "Upsample",
   "sinceVersion": 9,
-  "description": "Upsamples the input by applying a per-dimension scale factor; each output dimension equals `floor(input_dimension * scale)`. Deprecated in favor of Resize; supports `nearest` and `linear` interpolation modes.",
-  "inputs": [
-    { "role": "X", "dtype": "T", "description": "Input tensor to upsample." },
-    {
-      "role": "scales",
-      "dtype": "S",
-      "description": "Per-dimension scale factors, one value per input dimension.",
-      "rank": 1
-    }
-  ],
-  "outputs": [
-    {
-      "role": "Y",
-      "dtype": "T",
-      "description": "Upsampled output tensor; each dimension is `floor(input_dimension * scale)`.",
-      "rank": "ranks.X"
-    }
-  ],
-  "attributes": { "mode": "nearest" },
-  "attributeDescriptions": {
-    "mode": "Interpolation algorithm to use when mapping output coordinates back to input values; either `\"nearest\"` or `\"linear\"`."
-  },
   "attributeConstraints": { "mode": { "values": ["nearest", "linear"] } },
   "typeConstraints": { "T": ["float32", "float16", "int32", "int8", "uint8"], "S": ["float32"] },
-  "args": {
-    "x": { "kind": "tensor", "semantic": "X", "role": "input" },
-    "scales": { "kind": "tensor", "semantic": "scales", "role": "input" },
-    "y": { "kind": "tensor", "semantic": "Y", "role": "output" }
-  },
-  "tunables": { "WORKGROUP_SIZE": 256 },
-  "bindingSets": {
-    "static": [
-      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$outputElement" }
-    ],
-    "genericScales": [
-      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      {
-        "name": "scales",
-        "arg": "scales",
-        "semantic": "scales",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" }
-    ],
-    "nchw": [
-      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
-      {
-        "name": "scales",
-        "arg": "scales",
-        "semantic": "scales",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$outputElement" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "count", "type": "u32", "value": "numel(shapes.Y)" },
-            { "name": "C", "type": "u32", "value": "dim(shapes.X, 1)" },
-            { "name": "inH", "type": "u32", "value": "dim(shapes.X, 2)" },
-            { "name": "inW", "type": "u32", "value": "dim(shapes.X, 3)" },
-            { "name": "outH", "type": "u32", "value": "dim(shapes.Y, 2)" },
-            { "name": "outW", "type": "u32", "value": "dim(shapes.Y, 3)" }
-          ]
-        }
-      }
-    ]
   },
-  "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
   "variants": [
     {
       "id": "nearest_integer_scale_vec4",
       "priority": 21,
-      "when": ["ranks.X == 4", "ranks.Y == 4", "dim(shapes.X, 0) == dim(shapes.Y, 0)", "dim(shapes.X, 1) == dim(shapes.Y, 1)", "numel(shapes.Y) > 0", "attrs.mode == \"nearest\"", "(dim(shapes.Y, 2) == dim(shapes.X, 2) * 2 or dim(shapes.Y, 2) == dim(shapes.X, 2) * 3 or dim(shapes.Y, 2) == dim(shapes.X, 2) * 4)", "(dim(shapes.Y, 3) == dim(shapes.X, 3) * 2 or dim(shapes.Y, 3) == dim(shapes.X, 3) * 3 or dim(shapes.Y, 3) == dim(shapes.X, 3) * 4)", "dim(shapes.Y, 3) % 4 == 0", "f16Ok(dtypes.T)"],
-      "constants": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"", "outputElement": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Upsample.nearestIntegerScaleVec4",
-          "source": {
-            "shader": "resize-nearest-integer-scale.wgsl.jinja",
-            "inputs": { "xShape": "shapes.X", "yShape": "shapes.Y", "vec4": true }
-          },
-          "bindings": "static",
-          "dispatch": { "gridStride": "numel(shapes.Y) / 4", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "nearest_integer_scale",
       "priority": 20,
-      "when": ["ranks.X == 4", "ranks.Y == 4", "dim(shapes.X, 0) == dim(shapes.Y, 0)", "dim(shapes.X, 1) == dim(shapes.Y, 1)", "numel(shapes.Y) > 0", "attrs.mode == \"nearest\"", "(dim(shapes.Y, 2) == dim(shapes.X, 2) * 2 or dim(shapes.Y, 2) == dim(shapes.X, 2) * 3 or dim(shapes.Y, 2) == dim(shapes.X, 2) * 4)", "(dim(shapes.Y, 3) == dim(shapes.X, 3) * 2 or dim(shapes.Y, 3) == dim(shapes.X, 3) * 3 or dim(shapes.Y, 3) == dim(shapes.X, 3) * 4)", "f16Ok(dtypes.T)"],
-      "constants": { "outputElement": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "Upsample.nearestIntegerScale",
-          "source": {
-            "shader": "resize-nearest-integer-scale.wgsl.jinja",
-            "inputs": { "xShape": "shapes.X", "yShape": "shapes.Y", "vec4": false }
-          },
-          "bindings": "static",
-          "dispatch": { "gridStride": "numel(shapes.Y)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "linear_2x_stencil_vec4",
       "priority": 20,
-      "when": ["ranks.X == 4", "ranks.Y == 4", "dim(shapes.X, 0) == dim(shapes.Y, 0)", "dim(shapes.X, 1) == dim(shapes.Y, 1)", "numel(shapes.Y) > 0", "attrs.mode == \"linear\"", "dim(shapes.Y, 2) == dim(shapes.X, 2) * 2", "dim(shapes.Y, 3) == dim(shapes.X, 3) * 2", "dim(shapes.Y, 3) % 4 == 0", "f16Ok(dtypes.T)"],
-      "constants": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"", "outputElement": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Upsample.linear2xStencilVec4",
-          "source": {
-            "shader": "resize-linear-2x-stencil.wgsl.jinja",
-            "inputs": { "xShape": "shapes.X", "yShape": "shapes.Y", "coordinateMode": "\"asymmetric\"", "vec4": true }
           },
-          "bindings": "static",
-          "dispatch": { "gridStride": "numel(shapes.Y) / 4", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "linear_2x_stencil",
       "priority": 19,
-      "when": ["ranks.X == 4", "ranks.Y == 4", "dim(shapes.X, 0) == dim(shapes.Y, 0)", "dim(shapes.X, 1) == dim(shapes.Y, 1)", "numel(shapes.Y) > 0", "attrs.mode == \"linear\"", "dim(shapes.Y, 2) == dim(shapes.X, 2) * 2", "dim(shapes.Y, 3) == dim(shapes.X, 3) * 2", "f16Ok(dtypes.T)"],
-      "constants": { "outputElement": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "Upsample.linear2xStencil",
-          "source": {
-            "shader": "resize-linear-2x-stencil.wgsl.jinja",
-            "inputs": { "xShape": "shapes.X", "yShape": "shapes.Y", "coordinateMode": "\"asymmetric\"" }
           },
-          "bindings": "static",
-          "dispatch": { "gridStride": "numel(shapes.Y)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "nchw_rank4_vec4",
       "priority": 1,
-      "when": ["ranks.X == 4", "ranks.Y == 4", "dim(shapes.X, 0) == dim(shapes.Y, 0)", "dim(shapes.X, 1) == dim(shapes.Y, 1)", "dim(shapes.Y, 3) % 4 == 0", "f16Ok(dtypes.T)"],
-      "constants": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"", "outputElement": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Upsample.Vec4",
-          "source": {
-            "shader": "resize-coord-transform.wgsl.jinja",
-            "inputs": {
-              "mode": "attrs.mode",
-              "coordinateMode": "\"asymmetric\"",
-              "nearestMode": "\"floor\"",
-              "vec4": true,
-              "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
-              "scalesBuffer": true
-            }
           },
-          "bindings": "nchw",
-          "dispatch": { "threads": "numel(shapes.Y) / 4", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "nchw_rank4",
-      "when": ["ranks.X == 4", "ranks.Y == 4", "dim(shapes.X, 0) == dim(shapes.Y, 0)", "dim(shapes.X, 1) == dim(shapes.Y, 1)", "f16Ok(dtypes.T)"],
-      "constants": { "outputElement": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "Upsample",
-          "source": {
-            "shader": "resize-coord-transform.wgsl.jinja",
-            "inputs": {
-              "mode": "attrs.mode",
-              "coordinateMode": "\"asymmetric\"",
-              "nearestMode": "\"floor\"",
-              "scalesBuffer": true
-            }
           },
-          "bindings": "nchw",
-          "dispatch": { "threads": "numel(shapes.Y)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "generic_rank_n",
       "priority": 0,
-      "when": ["ranks.X == ranks.Y", "ranks.X >= 1", "numel(shapes.Y) > 0", "(ranks.X != 4 or dim(shapes.X, 0) != dim(shapes.Y, 0) or dim(shapes.X, 1) != dim(shapes.Y, 1))", "f16Ok(dtypes.T)"],
       "passes": [
         {
           "id": "main",
           "name": "Upsample.Generic",
-          "source": {
-            "shader": "resize-generic.wgsl.jinja",
-            "inputs": {
-              "xShape": "shapes.X",
-              "yShape": "shapes.Y",
-              "mode": "attrs.mode",
-              "coordinateMode": "\"asymmetric\"",
-              "nearestMode": "\"floor\"",
-              "scalesBuffer": true
-            }
           },
-          "bindings": "genericScales",
-          "dispatch": { "gridStride": "numel(shapes.Y)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     }

   "domain": "ai.onnx",
   "name": "Upsample",
   "sinceVersion": 9,
+  "inputs": { "x": { "onnx": "X", "dtype": "T" }, "scales": { "dtype": "S", "rank": 1 } },
+  "outputs": { "y": { "onnx": "Y", "dtype": "T", "rank": "ranks.x" } },
+  "attributes": { "mode": { "default": "nearest" } },
   "attributeConstraints": { "mode": { "values": ["nearest", "linear"] } },
   "typeConstraints": { "T": ["float32", "float16", "int32", "int8", "uint8"], "S": ["float32"] },
+  "tunables": { "WORKGROUP_SIZE": { "default": 256 } },
+  "derive": { "scalar": "dtypes.T" },
+  "when": ["f16Ok(dtypes.T)"],
+  "bindings": {
+    "y": { "buffer": "storage", "elementType": "$outputElement" },
+    "scales": { "buffer": "read-only-storage", "elementType": "f32" },
+    "params": {
+      "buffer": "uniform",
+      "struct": [
+        { "name": "count", "type": "u32", "value": "numel(shapes.y)" },
+        { "name": "C", "type": "u32", "value": "dim(shapes.x, 1)" },
+        { "name": "inH", "type": "u32", "value": "dim(shapes.x, 2)" },
+        { "name": "inW", "type": "u32", "value": "dim(shapes.x, 3)" },
+        { "name": "outH", "type": "u32", "value": "dim(shapes.y, 2)" },
+        { "name": "outW", "type": "u32", "value": "dim(shapes.y, 3)" }
+      ]
+    }
   },
   "variants": [
     {
       "id": "nearest_integer_scale_vec4",
       "priority": 21,
+      "when": ["ranks.x == 4", "ranks.y == 4", "dim(shapes.x, 0) == dim(shapes.y, 0)", "dim(shapes.x, 1) == dim(shapes.y, 1)", "numel(shapes.y) > 0", "attrs.mode == \"nearest\"", "(dim(shapes.y, 2) == dim(shapes.x, 2) * 2 or dim(shapes.y, 2) == dim(shapes.x, 2) * 3 or dim(shapes.y, 2) == dim(shapes.x, 2) * 4)", "(dim(shapes.y, 3) == dim(shapes.x, 3) * 2 or dim(shapes.y, 3) == dim(shapes.x, 3) * 3 or dim(shapes.y, 3) == dim(shapes.x, 3) * 4)", "dim(shapes.y, 3) % 4 == 0"],
+      "derive": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"", "outputElement": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Upsample.nearestIntegerScaleVec4",
+          "shader": "resize-nearest-integer-scale.wgsl.jinja",
+          "derive": { "xShape": "shapes.x", "yShape": "shapes.y", "vec4": true },
+          "bindings": ["x", "y"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
+            "y": 1,
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "nearest_integer_scale",
       "priority": 20,
+      "when": ["ranks.x == 4", "ranks.y == 4", "dim(shapes.x, 0) == dim(shapes.y, 0)", "dim(shapes.x, 1) == dim(shapes.y, 1)", "numel(shapes.y) > 0", "attrs.mode == \"nearest\"", "(dim(shapes.y, 2) == dim(shapes.x, 2) * 2 or dim(shapes.y, 2) == dim(shapes.x, 2) * 3 or dim(shapes.y, 2) == dim(shapes.x, 2) * 4)", "(dim(shapes.y, 3) == dim(shapes.x, 3) * 2 or dim(shapes.y, 3) == dim(shapes.x, 3) * 3 or dim(shapes.y, 3) == dim(shapes.x, 3) * 4)"],
+      "derive": { "outputElement": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "Upsample.nearestIntegerScale",
+          "shader": "resize-nearest-integer-scale.wgsl.jinja",
+          "derive": { "xShape": "shapes.x", "yShape": "shapes.y", "vec4": false },
+          "bindings": ["x", "y"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
+            "y": 1,
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "linear_2x_stencil_vec4",
       "priority": 20,
+      "when": ["ranks.x == 4", "ranks.y == 4", "dim(shapes.x, 0) == dim(shapes.y, 0)", "dim(shapes.x, 1) == dim(shapes.y, 1)", "numel(shapes.y) > 0", "attrs.mode == \"linear\"", "dim(shapes.y, 2) == dim(shapes.x, 2) * 2", "dim(shapes.y, 3) == dim(shapes.x, 3) * 2", "dim(shapes.y, 3) % 4 == 0"],
+      "derive": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"", "outputElement": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Upsample.linear2xStencilVec4",
+          "shader": "resize-linear-2x-stencil.wgsl.jinja",
+          "derive": {
+            "xShape": "shapes.x",
+            "yShape": "shapes.y",
+            "coordinateMode": "\"asymmetric\"",
+            "vec4": true,
+            "scalesBuffer": true
           },
+          "bindings": ["x", "scales", "y"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
+            "y": 1,
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "linear_2x_stencil",
       "priority": 19,
+      "when": ["ranks.x == 4", "ranks.y == 4", "dim(shapes.x, 0) == dim(shapes.y, 0)", "dim(shapes.x, 1) == dim(shapes.y, 1)", "numel(shapes.y) > 0", "attrs.mode == \"linear\"", "dim(shapes.y, 2) == dim(shapes.x, 2) * 2", "dim(shapes.y, 3) == dim(shapes.x, 3) * 2"],
+      "derive": { "outputElement": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "Upsample.linear2xStencil",
+          "shader": "resize-linear-2x-stencil.wgsl.jinja",
+          "derive": {
+            "xShape": "shapes.x",
+            "yShape": "shapes.y",
+            "coordinateMode": "\"asymmetric\"",
+            "scalesBuffer": true
           },
+          "bindings": ["x", "scales", "y"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
+            "y": 1,
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "nchw_rank4_vec4",
       "priority": 1,
+      "when": ["ranks.x == 4", "ranks.y == 4", "dim(shapes.x, 0) == dim(shapes.y, 0)", "dim(shapes.x, 1) == dim(shapes.y, 1)", "dim(shapes.y, 3) % 4 == 0"],
+      "derive": { "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"", "outputElement": "\"vec4<\" ~ dtypes.T ~ \">\"" },
       "passes": [
         {
           "id": "main",
           "name": "Upsample.Vec4",
+          "shader": "resize-coord-transform.wgsl.jinja",
+          "derive": {
+            "modeSpec": "attrs.mode",
+            "coordinateMode": "\"asymmetric\"",
+            "nearestMode": "\"floor\"",
+            "vec4": true,
+            "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
+            "scalesBuffer": true
           },
+          "bindings": ["x", "scales", "y", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "nchw_rank4",
+      "when": ["ranks.x == 4", "ranks.y == 4", "dim(shapes.x, 0) == dim(shapes.y, 0)", "dim(shapes.x, 1) == dim(shapes.y, 1)"],
+      "derive": { "outputElement": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "Upsample",
+          "shader": "resize-coord-transform.wgsl.jinja",
+          "derive": {
+            "modeSpec": "attrs.mode",
+            "coordinateMode": "\"asymmetric\"",
+            "nearestMode": "\"floor\"",
+            "scalesBuffer": true
           },
+          "bindings": ["x", "scales", "y", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "generic_rank_n",
       "priority": 0,
+      "when": ["ranks.x == ranks.y", "ranks.x >= 1", "numel(shapes.y) > 0", "(ranks.x != 4 or dim(shapes.x, 0) != dim(shapes.y, 0) or dim(shapes.x, 1) != dim(shapes.y, 1))"],
       "passes": [
         {
           "id": "main",
           "name": "Upsample.Generic",
+          "shader": "resize-generic.wgsl.jinja",
+          "derive": {
+            "xShape": "shapes.x",
+            "yShape": "shapes.y",
+            "modeSpec": "attrs.mode",
+            "coordinateMode": "\"asymmetric\"",
+            "nearestMode": "\"floor\"",
+            "scalesBuffer": true
           },
+          "bindings": ["x", "scales", { "arg": "y" }],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
+            "y": 1,
+            "z": 1
+          }
         }
       ]
     }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,21 +1,32 @@
 {
   "name": "ai.onnx.Upsample",
-  "id": "_ai_onnx_upsample_webgpu_a452be7",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "XMODCq/PrRSKwq5XRxOvlQoOhjCxnml+3+c9dAA0vv8=",
-      "manifest.json": "RXmICJOLSLjz7QW9zMGE5gJgKMAeqkER3fq+e7O64K0=",
-      "resize-coord-transform.wgsl.jinja": "+HABv5jMRtHp8VOXFY1Wi/ZjTIFxOJe6dExMJFAZ7H0=",
-      "resize-generic.wgsl.jinja": "JYlXh0pPD1wW2uCSm9KhS+rJqNimezHaq4sKb+iqv0M=",
-      "resize-linear-2x-stencil.wgsl.jinja": "cyI43wkZQdrjNUiaXRP2bfvLBuy04VvQQ5rFTuYXMd0=",
-      "resize-nearest-integer-scale.wgsl.jinja": "fgry1tCB+VAmoiWgqHWsy3EL2G47mtvtZhVPP27Z0Fg=",
-      "test.json": "FWRXvZ7UwTNUb3+n9G8xf4jp53QFrImqWPe3z0yUcqo="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.Upsample" }
 }

 {
   "name": "ai.onnx.Upsample",
+  "id": "_ai_onnx_upsample_webgpu_226c6a3",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "F6Vym+nah84vMiKJZvh+qKarMEzp84955NrGH18TWBs=",
+      "manifest.json": "ZZF7ZLeywO+hHftIgE5IeP4cufl+b3GKDYU9AV/inS4=",
+      "resize-coord-transform.wgsl.jinja": "WzhlhFxNz877wX29wcns2y+IJTC5g/IzYX5NEH27n6s=",
+      "resize-generic.wgsl.jinja": "bAUivK8i1ccgtOFXGB7Y9kTjyRfql5TUb+sikbvOMPk=",
+      "resize-linear-2x-stencil.wgsl.jinja": "r98ldVQUx6+bQnoqw9Fjt4JtygJqj5yEKXAv2DOJ5ic=",
+      "resize-nearest-integer-scale.wgsl.jinja": "GnmGPu4MyJrir+dzp3ucQ/ADVW9mJDcU1czn6L/RPF8=",
+      "test.json": "Cmo87nuUyMyjpettO/Nks1mu7Z00DCEnYyaHf5fQNsI="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "nearest_integer_scale_vec4": ["resize-nearest-integer-scale.wgsl.jinja"],
+      "nearest_integer_scale": ["resize-nearest-integer-scale.wgsl.jinja"],
+      "linear_2x_stencil_vec4": ["resize-linear-2x-stencil.wgsl.jinja"],
+      "linear_2x_stencil": ["resize-linear-2x-stencil.wgsl.jinja"],
+      "nchw_rank4_vec4": ["resize-coord-transform.wgsl.jinja"],
+      "nchw_rank4": ["resize-coord-transform.wgsl.jinja"],
+      "generic_rank_n": ["resize-generic.wgsl.jinja"]
+    }
+  }
 }

build/webgpu/resize-coord-transform.wgsl.jinja CHANGED Viewed

@@ -3,9 +3,8 @@
 // Nearest-index rounding is load-bearing for ONNX compatibility:
 // round_prefer_ceil ties toward positive infinity, while round_prefer_floor
 // uses its specified epsilon.
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 {% set wg = workgroupSize if workgroupSize is defined else tunables.WORKGROUP_SIZE %}
@@ -19,7 +18,7 @@ fn source_coord(out_coord: u32, in_size: u32, out_size: u32, explicit_scale: f32
 /* Backend-stable nearest-neighbor coordinate selection. Inferred align-corners,
  * asymmetric, and half-pixel transforms use exact integer ratios. This matters
  * at nearest-mode ties: reassociating an f32 division into a reciprocal
- * multiply can move an exact 0.5 to either side of the boundary. */{% if source.mode == "nearest" %}
 fn nearest_index(v: f32) -> i32 {
   return i32(floor(v));
 }
@@ -67,7 +66,7 @@ fn input_at(base: u32, h: i32, w: i32) -> f32 {
   return f32(x[base + hh * params.inW + ww]);
 }
-{% if source.mode == "linear" %}
 // Clamp before choosing neighbors. At an edge both neighbors coincide and
 // receive half weight; the four-term order is observable for integer outputs.
 fn bilinear_sample(base: u32, source_y: f32, source_x: f32) -> f32 {
@@ -92,20 +91,24 @@ fn bilinear_sample(base: u32, source_y: f32, source_x: f32) -> f32 {
 {% endif %}
 @compute @workgroup_size({{ wg }})
-fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
   // 2D-folded flat index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension limit.
-  let flat = gid.x + gid.y * nwg.x * {{ wg }}u;
-  // Upsample: the per-axis scales are a runtime input tensor, unknown at plan
-  // build time, so read the true spatial scales (NCHW axes 2/3) directly instead
-  // of inferring out/in (floor(in*scale) loses the fractional part when the scale
-  // is non-integer, e.g. in=2 scale=2.333 -> out=4 -> inferred 2.0 != 2.333).
   let scaleH_val = f32(scales[2]);
   let scaleW_val = f32(scales[3]);
-{% if source.vec4 %}
-  // One thread owns 4 contiguous output columns (outW % 4 == 0, so a 4-group
-  // never crosses a row): all 4 lanes share oh/c/n and the source row coord sy;
-  // only sx differs per lane. One 128-bit vec4 store. Bit-identical to scalar.
   let count4 = params.count / 4u;
   if (flat >= count4) { return; }
   let i = flat * 4u;
@@ -115,14 +118,49 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
   let c = (row / params.outH) % params.C;
   let n = row / (params.outH * params.C);
   let base = (n * params.C + c) * params.inH * params.inW;
-{% if source.mode != "nearest" %}
   let sy = source_coord(oh, params.inH, params.outH, scaleH_val);
   let sx0 = source_coord(ow0, params.inW, params.outW, scaleW_val);
   let sx1 = source_coord(ow0 + 1u, params.inW, params.outW, scaleW_val);
   let sx2 = source_coord(ow0 + 2u, params.inW, params.outW, scaleW_val);
   let sx3 = source_coord(ow0 + 3u, params.inW, params.outW, scaleW_val);
 {% endif %}
-  {% if source.mode == "nearest" %}
   let ih = nearest_coord_index(oh, params.inH, params.outH, scaleH_val);
   let r = vec4<f32>(
     input_at(base, ih, nearest_coord_index(ow0, params.inW, params.outW, scaleW_val)),
@@ -136,7 +174,11 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
     bilinear_sample(base, sy, sx2),
     bilinear_sample(base, sy, sx3));
   {% endif %}
-  y[flat] = {{ source.vecType }}(r);
 {% else %}
   let i = flat;
   if (i >= params.count) { return; }
@@ -145,11 +187,11 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
   let c = (i / (params.outW * params.outH)) % params.C;
   let n = i / (params.outW * params.outH * params.C);
   let base = (n * params.C + c) * params.inH * params.inW;
-{% if source.mode != "nearest" %}
   let sy = source_coord(oh, params.inH, params.outH, scaleH_val);
   let sx = source_coord(ow, params.inW, params.outW, scaleW_val);
 {% endif %}
-  {% if source.mode == "nearest" %}
   let ih = nearest_coord_index(oh, params.inH, params.outH, scaleH_val);
   let iw = nearest_coord_index(ow, params.inW, params.outW, scaleW_val);
   let result = input_at(base, ih, iw);

 // Nearest-index rounding is load-bearing for ONNX compatibility:
 // round_prefer_ceil ties toward positive infinity, while round_prefer_floor
 // uses its specified epsilon.
+{% set scalesBuffer = scalesBuffer is defined and scalesBuffer %}
+{% set scalarX4 = scalarX4 is defined and scalarX4 %}
 {{ env.wgsl.resourceDeclarations }}
 {% set wg = workgroupSize if workgroupSize is defined else tunables.WORKGROUP_SIZE %}
 /* Backend-stable nearest-neighbor coordinate selection. Inferred align-corners,
  * asymmetric, and half-pixel transforms use exact integer ratios. This matters
  * at nearest-mode ties: reassociating an f32 division into a reciprocal
+ * multiply can move an exact 0.5 to either side of the boundary. */{% if modeSpec == "nearest" %}
 fn nearest_index(v: f32) -> i32 {
   return i32(floor(v));
 }
   return f32(x[base + hh * params.inW + ww]);
 }
+{% if modeSpec == "linear" %}
 // Clamp before choosing neighbors. At an edge both neighbors coincide and
 // receive half weight; the four-term order is observable for integer outputs.
 fn bilinear_sample(base: u32, source_y: f32, source_x: f32) -> f32 {
 {% endif %}
 @compute @workgroup_size({{ wg }})
+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 flat = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wg }}u;
+{% if scalesBuffer %}
+  // Read spatial scales from NCHW axes 2 and 3. Inferring a scale from the
+  // rounded output extent would discard its fractional part; for example,
+  // input 2 scaled by 2.333 has output extent 4, whose extent ratio is 2.0.
   let scaleH_val = f32(scales[2]);
   let scaleW_val = f32(scales[3]);
+{% else %}
+  let scaleH_val = params.scaleH;
+  let scaleW_val = params.scaleW;
+{% endif %}
+{% if vec4 %}
+  // One invocation owns four contiguous output columns. Because outW is
+  // divisible by four, the group never crosses a row; all lanes share n, c,
+  // oh, and source row coordinate sy, while sx differs per lane.
   let count4 = params.count / 4u;
   if (flat >= count4) { return; }
   let i = flat * 4u;
   let c = (row / params.outH) % params.C;
   let n = row / (params.outH * params.C);
   let base = (n * params.C + c) * params.inH * params.inW;
+{% if modeSpec != "nearest" %}
+  let sy = source_coord(oh, params.inH, params.outH, scaleH_val);
+  let sx0 = source_coord(ow0, params.inW, params.outW, scaleW_val);
+  let sx1 = source_coord(ow0 + 1u, params.inW, params.outW, scaleW_val);
+  let sx2 = source_coord(ow0 + 2u, params.inW, params.outW, scaleW_val);
+  let sx3 = source_coord(ow0 + 3u, params.inW, params.outW, scaleW_val);
+{% endif %}
+  {% if modeSpec == "nearest" %}
+  let ih = nearest_coord_index(oh, params.inH, params.outH, scaleH_val);
+  let r = vec4<f32>(
+    input_at(base, ih, nearest_coord_index(ow0, params.inW, params.outW, scaleW_val)),
+    input_at(base, ih, nearest_coord_index(ow0 + 1u, params.inW, params.outW, scaleW_val)),
+    input_at(base, ih, nearest_coord_index(ow0 + 2u, params.inW, params.outW, scaleW_val)),
+    input_at(base, ih, nearest_coord_index(ow0 + 3u, params.inW, params.outW, scaleW_val)));
+  {% else %}
+  let r = vec4<f32>(
+    bilinear_sample(base, sy, sx0),
+    bilinear_sample(base, sy, sx1),
+    bilinear_sample(base, sy, sx2),
+    bilinear_sample(base, sy, sx3));
+  {% endif %}
+  y[flat] = {{ vecType }}(r);
+{% elif scalarX4 %}
+  // Each invocation computes four row-local columns without a vec4 storage
+  // binding. The final one to three columns use guarded scalar stores.
+  let groups_per_row = (params.outW + 3u) / 4u;
+  let group_count = (params.count / params.outW) * groups_per_row;
+  if (flat >= group_count) { return; }
+  let group_in_row = flat % groups_per_row;
+  let row = flat / groups_per_row;
+  let ow0 = group_in_row * 4u;
+  let oh = row % params.outH;
+  let c = (row / params.outH) % params.C;
+  let n = row / (params.outH * params.C);
+  let base = (n * params.C + c) * params.inH * params.inW;
+{% if modeSpec != "nearest" %}
   let sy = source_coord(oh, params.inH, params.outH, scaleH_val);
   let sx0 = source_coord(ow0, params.inW, params.outW, scaleW_val);
   let sx1 = source_coord(ow0 + 1u, params.inW, params.outW, scaleW_val);
   let sx2 = source_coord(ow0 + 2u, params.inW, params.outW, scaleW_val);
   let sx3 = source_coord(ow0 + 3u, params.inW, params.outW, scaleW_val);
 {% endif %}
+  {% if modeSpec == "nearest" %}
   let ih = nearest_coord_index(oh, params.inH, params.outH, scaleH_val);
   let r = vec4<f32>(
     input_at(base, ih, nearest_coord_index(ow0, params.inW, params.outW, scaleW_val)),
     bilinear_sample(base, sy, sx2),
     bilinear_sample(base, sy, sx3));
   {% endif %}
+  let i = row * params.outW + ow0;
+  y[i] = {{ scalar }}(r.x);
+  if (ow0 + 1u < params.outW) { y[i + 1u] = {{ scalar }}(r.y); }
+  if (ow0 + 2u < params.outW) { y[i + 2u] = {{ scalar }}(r.z); }
+  if (ow0 + 3u < params.outW) { y[i + 3u] = {{ scalar }}(r.w); }
 {% else %}
   let i = flat;
   if (i >= params.count) { return; }
   let c = (i / (params.outW * params.outH)) % params.C;
   let n = i / (params.outW * params.outH * params.C);
   let base = (n * params.C + c) * params.inH * params.inW;
+{% if modeSpec != "nearest" %}
   let sy = source_coord(oh, params.inH, params.outH, scaleH_val);
   let sx = source_coord(ow, params.inW, params.outW, scaleW_val);
 {% endif %}
+  {% if modeSpec == "nearest" %}
   let ih = nearest_coord_index(oh, params.inH, params.outH, scaleH_val);
   let iw = nearest_coord_index(ow, params.inW, params.outW, scaleW_val);
   let result = input_at(base, ih, iw);

build/webgpu/resize-generic.wgsl.jinja CHANGED Viewed

@@ -1,18 +1,13 @@
-// Rank-generic Resize for arbitrary ranks and resized axes, including
-// tf_crop_and_resize with an ROI and extrapolation value. The per-axis
-// coordinate transform follows the rank-specialized paths so half_pixel,
-// pytorch_half_pixel, asymmetric, and align_corners stay
-// consistent. Linear interpolation runs only over resized axes; passthrough
-// axes contribute one tap.
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 {% set wg = workgroupSize if workgroupSize is defined else tunables.WORKGROUP_SIZE %}
-{%- set xs = source.xShape %}
-{% set ys = source.yShape %}
 {% set r = xs | length -%}
-{% set usesExplicitScale = source.coordinateMode not in ["align_corners", "tf_crop_and_resize"] %}
 {% set resolvedScales = namespace(values = []) %}
 {% for a in range(r) %}
@@ -36,10 +31,9 @@ enable f16;
 const WG: u32 = {{ wg }}u;
-// Runtime form used by the shared exact-nearest helper. The generic kernel's
-// interpolation paths keep their baked expressions below; nearest uses this
-// function only when an exact inferred-ratio path is unavailable.
-{% if source.mode == "nearest" and k > 0 and (source.roi is not defined or (source.roi | length) == 0) and source.coordinateMode != "align_corners" %}
 fn source_coord(out_coord: u32, in_size: u32, out_size: u32, explicit_scale: f32) -> f32 {
   return select(f32(out_coord) * f32(in_size) / f32(out_size), f32(out_coord) / explicit_scale, explicit_scale > 0.0);
@@ -57,10 +51,10 @@ fn source_coord(out_coord: u32, in_size: u32, out_size: u32, explicit_scale: f32
   (f32({{ oc }}) / scales[{{ a }}])
 {% endmacro %}
-{% if source.mode == "cubic" %}
 // Keys cubic convolution kernel with coefficient coeffA.
 fn cubic_weight(s_in: f32) -> f32 {
-  let a = f32({{ source.coeffA }});
   let s = abs(s_in);
   let s2 = s * s;
   let s3 = s2 * s;
@@ -70,15 +64,15 @@ fn cubic_weight(s_in: f32) -> f32 {
 }
 {% endif %}
-{% if source.mode == "nearest" and k > 0 %}
 /* Backend-stable nearest-neighbor coordinate selection. Inferred align-corners,
  * asymmetric, and half-pixel transforms use exact integer ratios. This matters
  * at nearest-mode ties: reassociating an f32 division into a reciprocal
- * multiply can move an exact 0.5 to either side of the boundary. */{% set ratioMode = source.coordinateMode in ["align_corners", "asymmetric"] %}
-{% set inferredHalfPixelMode = source.coordinateMode in ["pytorch_half_pixel", "half_pixel", "half_pixel_symmetric"] %}
-{% set noRoi = source.roi is not defined or (source.roi | length) == 0 %}
-{% set usesExplicitScale = source.coordinateMode not in ["align_corners", "tf_crop_and_resize"] %}
-{% if source.mode == "nearest" %}
 fn nearest_index(v: f32) -> i32 {
   return i32(floor(v));
 }
@@ -177,11 +171,11 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
 {% endfor %}
     // Per-axis source coordinates.
 {% for a in range(r) %}
-{% if a in resized.list and (source.mode != "nearest" or (source.roi is defined and source.roi)) %}
     let s{{ a }} = {{ src_coord(a, out_coord(a)) }};
 {% endif %}
 {% endfor %}
-{% if source.mode == "nearest" %}
 {% for a in range(r) %}
 {% if a in resized.list %}
 {% set explicitScale = "f32(scales[" ~ a ~ "])" %}
@@ -191,7 +185,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
 {% endif %}
 {% endfor %}
     y[i] = {{ scalar }}(f32(x[0u{% for a in range(r) %} + idx{{ a }} * {{ istride(a) }}u{% endfor %}]));
-{% elif source.mode == "cubic" %}
     // Separable Keys cubic: 4 taps per resized axis (offsets -1..2), passthrough
     // axes contribute a single tap. 4^k taps total.
 {% for a in resized.list %}

+{% set coeffA = coeffA | default(0) %}
+// Rank-generic interpolation for arbitrary ranks and resized axes. Scales are
+// read from the input buffer, and linear interpolation runs only over resized
+// axes; passthrough axes contribute one tap.
 {{ env.wgsl.resourceDeclarations }}
 {% set wg = workgroupSize if workgroupSize is defined else tunables.WORKGROUP_SIZE %}
+{%- set xs = xShape %}
+{% set ys = yShape %}
 {% set r = xs | length -%}
+{% set usesExplicitScale = coordinateMode not in ["align_corners", "tf_crop_and_resize"] %}
 {% set resolvedScales = namespace(values = []) %}
 {% for a in range(r) %}
 const WG: u32 = {{ wg }}u;
+// Runtime coordinate form for nearest interpolation when an exact
+// inferred-ratio expression is unavailable.
+{% if modeSpec == "nearest" and k > 0 and (roi is not defined or (roi | length) == 0) and coordinateMode != "align_corners" %}
 fn source_coord(out_coord: u32, in_size: u32, out_size: u32, explicit_scale: f32) -> f32 {
   return select(f32(out_coord) * f32(in_size) / f32(out_size), f32(out_coord) / explicit_scale, explicit_scale > 0.0);
   (f32({{ oc }}) / scales[{{ a }}])
 {% endmacro %}
+{% if modeSpec == "cubic" %}
 // Keys cubic convolution kernel with coefficient coeffA.
 fn cubic_weight(s_in: f32) -> f32 {
+  let a = f32({{ coeffA }});
   let s = abs(s_in);
   let s2 = s * s;
   let s3 = s2 * s;
 }
 {% endif %}
+{% if modeSpec == "nearest" and k > 0 %}
 /* Backend-stable nearest-neighbor coordinate selection. Inferred align-corners,
  * asymmetric, and half-pixel transforms use exact integer ratios. This matters
  * at nearest-mode ties: reassociating an f32 division into a reciprocal
+ * multiply can move an exact 0.5 to either side of the boundary. */{% set ratioMode = coordinateMode in ["align_corners", "asymmetric"] %}
+{% set inferredHalfPixelMode = coordinateMode in ["pytorch_half_pixel", "half_pixel", "half_pixel_symmetric"] %}
+{% set noRoi = roi is not defined or (roi | length) == 0 %}
+{% set usesExplicitScale = coordinateMode not in ["align_corners", "tf_crop_and_resize"] %}
+{% if modeSpec == "nearest" %}
 fn nearest_index(v: f32) -> i32 {
   return i32(floor(v));
 }
 {% endfor %}
     // Per-axis source coordinates.
 {% for a in range(r) %}
+{% if a in resized.list and (modeSpec != "nearest" or (roi is defined and roi)) %}
     let s{{ a }} = {{ src_coord(a, out_coord(a)) }};
 {% endif %}
 {% endfor %}
+{% if modeSpec == "nearest" %}
 {% for a in range(r) %}
 {% if a in resized.list %}
 {% set explicitScale = "f32(scales[" ~ a ~ "])" %}
 {% endif %}
 {% endfor %}
     y[i] = {{ scalar }}(f32(x[0u{% for a in range(r) %} + idx{{ a }} * {{ istride(a) }}u{% endfor %}]));
+{% elif modeSpec == "cubic" %}
     // Separable Keys cubic: 4 taps per resized axis (offsets -1..2), passthrough
     // axes contribute a single tap. 4^k taps total.
 {% for a in resized.list %}

build/webgpu/resize-linear-2x-stencil.wgsl.jinja CHANGED Viewed

@@ -1,35 +1,55 @@
-// NCHW bilinear Resize with an exact 2x upscale on both spatial axes. The
 // interpolation weights collapse to fixed per-parity stencils, replacing
 // per-element floating-point coordinate divisions with integer parity tests
 // and constants.
 //
-// For half_pixel and pytorch_half_pixel, src = o / 2 - 0.25: even outputs use
-// floor q - 1 with fraction 0.75, and odd outputs use floor q with fraction
-// 0.25. For asymmetric mode, src = o / 2: even outputs use fraction 0.0 and odd
-// outputs use 0.5. align_corners remains on the generic fallback because its
-// position-dependent scale does not reduce to fixed stencils.
 //
 // Input coordinates clamp at the edges, including the q - 1 and q + 1 border
-// taps. The f32 weight construction is bit-identical to the generic fallback.
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 {% set wg = workgroupSize if workgroupSize is defined else tunables.WORKGROUP_SIZE %}
 const WG: u32 = {{ wg }}u;
-const IN_H: u32 = {{ source.xShape[2] }}u;
-const IN_W: u32 = {{ source.xShape[3] }}u;
-const IN_HW: u32 = {{ source.xShape[2] * source.xShape[3] }}u;
-const OUT_H: u32 = {{ source.yShape[2] }}u;
-const OUT_W: u32 = {{ source.yShape[3] }}u;
-const COUNT: u32 = {{ source.yShape[0] * source.yShape[1] * source.yShape[2] * source.yShape[3] }}u;
 fn input_at(base: u32, h: i32, w: i32) -> f32 {
   let hh = u32(clamp(h, 0, i32(IN_H) - 1));
   let ww = u32(clamp(w, 0, i32(IN_W) - 1));
   return f32(x[base + hh * IN_W + ww]);
 }
 @compute @workgroup_size({{ wg }})
 fn main(
@@ -37,13 +57,16 @@ fn main(
   @builtin(num_workgroups) nwg: vec3<u32>
 ) {
   let stride = nwg.x * WG;
-{% if source.vec4 %}
   // Vec4 path: one thread owns 4 contiguous output columns (OUT_W % 4 == 0, so
   // a 4-group never crosses a row). All 4 lanes share the same oh -> same source
   // rows (y0,y1) and the same 4-column source window (qx0-1 .. qx0+2); load that
   // window once, compute the 4 bilinear results with the baked per-lane parity
-  // weights, and emit a single 128-bit vec4 store. Bit-identical to the scalar
-  // path (same input_at, same (1-w) split, same f32 operation order).
   let OUT_W4 = OUT_W / 4u;
   let COUNT4 = COUNT / 4u;
   for (var v = gid.x; v < COUNT4; v += stride) {
@@ -53,6 +76,16 @@ fn main(
     let nc = tt / OUT_H;
     let base = nc * IN_HW;
     let ow0 = ow4 * 4u;
     let qx0 = i32(ow0 / 2u);
     let y0 = i32(oh / 2u);
     let wy = select(0.0, 0.5, (oh & 1u) == 1u);
@@ -77,6 +110,12 @@ fn main(
     let oh = t % OUT_H;
     let nc = t / OUT_H;
     let base = nc * IN_HW;
     let qy = i32(oh / 2u);
     let qx = i32(ow / 2u);
     let oddY = (oh & 1u) == 1u;

+{% set scalesBuffer = scalesBuffer if scalesBuffer is defined else false %}
+// NCHW bilinear 2x upscale on both spatial axes. The
 // interpolation weights collapse to fixed per-parity stencils, replacing
 // per-element floating-point coordinate divisions with integer parity tests
 // and constants.
 //
+// With asymmetric coordinates, src = o / 2: even outputs use fraction 0.0 and
+// odd outputs use fraction 0.5.
 //
 // Input coordinates clamp at the edges, including the q - 1 and q + 1 border
+// taps. Interpolation weights are constructed in f32 before the output cast.
 {{ env.wgsl.resourceDeclarations }}
 {% set wg = workgroupSize if workgroupSize is defined else tunables.WORKGROUP_SIZE %}
 const WG: u32 = {{ wg }}u;
+const IN_H: u32 = {{ xShape[2] }}u;
+const IN_W: u32 = {{ xShape[3] }}u;
+const IN_HW: u32 = {{ xShape[2] * xShape[3] }}u;
+const OUT_H: u32 = {{ yShape[2] }}u;
+const OUT_W: u32 = {{ yShape[3] }}u;
+const COUNT: u32 = {{ yShape[0] * yShape[1] * yShape[2] * yShape[3] }}u;
 fn input_at(base: u32, h: i32, w: i32) -> f32 {
   let hh = u32(clamp(h, 0, i32(IN_H) - 1));
   let ww = u32(clamp(w, 0, i32(IN_W) - 1));
   return f32(x[base + hh * IN_W + ww]);
 }
+{% if scalesBuffer %}
+// The output shape condition floor(in * scale) == 2 * in also admits scales
+// slightly above 2. Read the spatial scales: exactly 2 uses the stencils;
+// otherwise sample asymmetric coordinates with source = output / scale.
+fn general_sample(base: u32, oh: u32, ow: u32) -> f32 {
+  let sy = clamp(f32(oh) / f32(scales[2]), 0.0, f32(IN_H - 1u));
+  let sx = clamp(f32(ow) / f32(scales[3]), 0.0, f32(IN_W - 1u));
+  let y0 = i32(floor(sy));
+  let x0 = i32(floor(sx));
+  let y1 = min(y0 + 1, i32(IN_H) - 1);
+  let x1 = min(x0 + 1, i32(IN_W) - 1);
+  let same_y = y0 == y1;
+  let same_x = x0 == x1;
+  let wy0 = select(abs(sy - f32(y1)), 0.5, same_y);
+  let wy1 = select(abs(sy - f32(y0)), 0.5, same_y);
+  let wx0 = select(abs(sx - f32(x1)), 0.5, same_x);
+  let wx1 = select(abs(sx - f32(x0)), 0.5, same_x);
+  let v00 = input_at(base, y0, x0);
+  let v01 = input_at(base, y0, x1);
+  let v10 = input_at(base, y1, x0);
+  let v11 = input_at(base, y1, x1);
+  return wx0 * wy0 * v00 + wx1 * wy0 * v01 + wx0 * wy1 * v10 + wx1 * wy1 * v11;
+}
+{% endif %}
 @compute @workgroup_size({{ wg }})
 fn main(
   @builtin(num_workgroups) nwg: vec3<u32>
 ) {
   let stride = nwg.x * WG;
+{% if scalesBuffer %}
+  let exact2x = f32(scales[2]) == 2.0 && f32(scales[3]) == 2.0;
+{% endif %}
+{% if vec4 %}
   // Vec4 path: one thread owns 4 contiguous output columns (OUT_W % 4 == 0, so
   // a 4-group never crosses a row). All 4 lanes share the same oh -> same source
   // rows (y0,y1) and the same 4-column source window (qx0-1 .. qx0+2); load that
   // window once, compute the 4 bilinear results with the baked per-lane parity
+  // weights, and emit a single vec4 store. Each component uses the same
+  // input_at calls, weight split, and f32 operation order.
   let OUT_W4 = OUT_W / 4u;
   let COUNT4 = COUNT / 4u;
   for (var v = gid.x; v < COUNT4; v += stride) {
     let nc = tt / OUT_H;
     let base = nc * IN_HW;
     let ow0 = ow4 * 4u;
+{% if scalesBuffer %}
+    if (!exact2x) {
+      y[v] = vec4<{{ scalar }}>(vec4<f32>(
+        general_sample(base, oh, ow0),
+        general_sample(base, oh, ow0 + 1u),
+        general_sample(base, oh, ow0 + 2u),
+        general_sample(base, oh, ow0 + 3u)));
+      continue;
+    }
+{% endif %}
     let qx0 = i32(ow0 / 2u);
     let y0 = i32(oh / 2u);
     let wy = select(0.0, 0.5, (oh & 1u) == 1u);
     let oh = t % OUT_H;
     let nc = t / OUT_H;
     let base = nc * IN_HW;
+{% if scalesBuffer %}
+    if (!exact2x) {
+      y[i] = {{ scalar }}(general_sample(base, oh, ow));
+      continue;
+    }
+{% endif %}
     let qy = i32(oh / 2u);
     let qx = i32(ow / 2u);
     let oddY = (oh & 1u) == 1u;

build/webgpu/resize-nearest-integer-scale.wgsl.jinja CHANGED Viewed

@@ -7,25 +7,22 @@
 // ow = q*s + r gives an offset whose magnitude is strictly below one half, so
 // ties cannot occur and every nearest rounding rule returns q. Downsampling and
 // other coordinate-mode combinations require the generic coordinate transform.
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 {% set wg = workgroupSize if workgroupSize is defined else tunables.WORKGROUP_SIZE %}
-{%- set inH = source.xShape[2] %}
-{% set inW = source.xShape[3] %}
-{% set outH = source.yShape[2] %}
-{% set outW = source.yShape[3] %}
 {% set scaleW = ((outW / inW) | int) %}
-{% set rows = source.yShape[0] * source.yShape[1] * outH %}
 const WG: u32 = {{ wg }}u;
 const IN_W: u32 = {{ inW }}u;
 const IN_HW: u32 = {{ inH * inW }}u;
 const OUT_H: u32 = {{ outH }}u;
 const SCALE_H: u32 = {{ ((outH / inH) | int) }}u;
-{% if not source.vec4 or scaleW == 3 %}const SCALE_W: u32 = {{ scaleW }}u;
 {% endif %}
-{% if source.vec4 %}
 const OUT_W_VECS: u32 = {{ ((outW / 4) | int) }}u;
 const COUNT: u32 = {{ (((rows * outW) / 4) | int) }}u;
 {% else %}
@@ -39,7 +36,7 @@ fn main(
   @builtin(num_workgroups) nwg: vec3<u32>
 ) {
   let stride = nwg.x * WG;
-{% if source.vec4 %}
   for (var v = gid.x; v < COUNT; v += stride) {
     let ow0 = (v % OUT_W_VECS) * 4u;
     let t = v / OUT_W_VECS;

 // ow = q*s + r gives an offset whose magnitude is strictly below one half, so
 // ties cannot occur and every nearest rounding rule returns q. Downsampling and
 // other coordinate-mode combinations require the generic coordinate transform.
 {{ env.wgsl.resourceDeclarations }}
 {% set wg = workgroupSize if workgroupSize is defined else tunables.WORKGROUP_SIZE %}
+{%- set inH = xShape[2] %}
+{% set inW = xShape[3] %}
+{% set outH = yShape[2] %}
+{% set outW = yShape[3] %}
 {% set scaleW = ((outW / inW) | int) %}
+{% set rows = yShape[0] * yShape[1] * outH %}
 const WG: u32 = {{ wg }}u;
 const IN_W: u32 = {{ inW }}u;
 const IN_HW: u32 = {{ inH * inW }}u;
 const OUT_H: u32 = {{ outH }}u;
 const SCALE_H: u32 = {{ ((outH / inH) | int) }}u;
+{% if not vec4 or scaleW == 3 %}const SCALE_W: u32 = {{ scaleW }}u;
 {% endif %}
+{% if vec4 %}
 const OUT_W_VECS: u32 = {{ ((outW / 4) | int) }}u;
 const COUNT: u32 = {{ (((rows * outW) / 4) | int) }}u;
 {% else %}
   @builtin(num_workgroups) nwg: vec3<u32>
 ) {
   let stride = nwg.x * WG;
+{% if vec4 %}
   for (var v = gid.x; v < COUNT; v += stride) {
     let ow0 = (v % OUT_W_VECS) * 4u;
     let t = v / OUT_W_VECS;

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.Upsample",
   "cases": [
     {
       "name": "nearest_2x_f32",
@@ -472,6 +471,38 @@
         }
       },
       "outputs": { "y": { "dtype": "float32", "shape": [1, 1, 2, 1, 2, 1, 2, 6], "tolerance": 0 } }
     }
   ]
 }

 {
   "cases": [
     {
       "name": "nearest_2x_f32",
         }
       },
       "outputs": { "y": { "dtype": "float32", "shape": [1, 1, 2, 1, 2, 1, 2, 6], "tolerance": 0 } }
+    },
+    {
+      "name": "linear_scale_2p1_not_exact_2x_f32",
+      "provenance": {
+        "notes": "scales 2.1 on a 3x3 input gives floor(3 * 2.1) = 6 outputs per axis, the same shape as an exact 2x upsample, but the asymmetric source coordinate is o / 2.1, not o / 2. A 2x stencil chosen from the shapes alone would compute the wrong interpolation weights."
+      },
+      "attrs": { "mode": "linear" },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 1, 3, 3],
+          "data": { "kind": "values", "values": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0] }
+        },
+        "scales": { "dtype": "float32", "shape": [4], "data": { "kind": "values", "values": [1.0, 1.0, 2.1, 2.1] } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 1, 6, 6], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "linear_scale_2p1_not_exact_2x_vec4_f32",
+      "provenance": {
+        "notes": "A 3-by-2 input with scales 2.1 produces a 6-by-4 vector-aligned output while retaining source coordinates based on 2.1 rather than 2."
+      },
+      "attrs": { "mode": "linear" },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 1, 3, 2],
+          "data": { "kind": "values", "values": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] }
+        },
+        "scales": { "dtype": "float32", "shape": [4], "data": { "kind": "values", "values": [1.0, 1.0, 2.1, 2.1] } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 1, 6, 4], "tolerance": 0.000001 } }
     }
   ]
 }