sync c928d21e6cc1

Browse files

Files changed (9) hide show

README.md +1 -0
build/webgpu/bench.json +9 -1
build/webgpu/manifest.json +160 -2
build/webgpu/metadata.json +9 -8
build/webgpu/nms-bitmask.wgsl.jinja +291 -0
build/webgpu/nms-groups-compact.wgsl.jinja +17 -20
build/webgpu/nms-groups-select.wgsl.jinja +6 -0
build/webgpu/non-max-suppression.wgsl.jinja +6 -0
build/webgpu/test.json +111 -0

README.md CHANGED Viewed

@@ -54,6 +54,7 @@ Default values (overridable per request):
 - [`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
 - [`nms-groups-compact.wgsl.jinja`](build/webgpu/nms-groups-compact.wgsl.jinja)
 - [`nms-groups-select.wgsl.jinja`](build/webgpu/nms-groups-select.wgsl.jinja)
 - [`non-max-suppression.wgsl.jinja`](build/webgpu/non-max-suppression.wgsl.jinja)

 - [`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
+- [`nms-bitmask.wgsl.jinja`](build/webgpu/nms-bitmask.wgsl.jinja)
 - [`nms-groups-compact.wgsl.jinja`](build/webgpu/nms-groups-compact.wgsl.jinja)
 - [`nms-groups-select.wgsl.jinja`](build/webgpu/nms-groups-select.wgsl.jinja)
 - [`non-max-suppression.wgsl.jinja`](build/webgpu/non-max-suppression.wgsl.jinja)

build/webgpu/bench.json CHANGED Viewed

@@ -10,7 +10,15 @@
         "iou_threshold": { "dtype": "float32", "shape": [], "data": { "kind": "constant", "value": 0.5 } },
         "score_threshold": { "dtype": "float32", "shape": [], "data": { "kind": "constant", "value": 0.1 } }
       },
-      "outputs": { "selected_indices": { "dtype": "uint32", "shape": [10, 3] } }
     },
     {
       "name": "nms-1class-2048-boxes",

         "iou_threshold": { "dtype": "float32", "shape": [], "data": { "kind": "constant", "value": 0.5 } },
         "score_threshold": { "dtype": "float32", "shape": [], "data": { "kind": "constant", "value": 0.1 } }
       },
+      "outputs": { "selected_indices": { "dtype": "uint32", "shape": [10, 3] } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "7 * dim(shapes.boxes, 0) * dim(shapes.scores, 1) * dim(shapes.boxes, 1) * dim(shapes.boxes, 1)"
+          }
+        ]
+      }
     },
     {
       "name": "nms-1class-2048-boxes",

build/webgpu/manifest.json CHANGED Viewed

@@ -67,14 +67,19 @@
     "score_threshold": { "kind": "tensor", "semantic": "score_threshold", "role": "input", "required": false },
     "selected_indices": { "kind": "tensor", "semantic": "selected_indices", "role": "output", "dtype": "uint32" }
   },
-  "tunables": { "WORKGROUP_SIZE": 256 },
   "derive": {
     "storageBufferLimit": "min(device.limits.maxStorageBufferBindingSize, device.limits.maxBufferSize)",
     "validLayout": "ranks.boxes == 3 and ranks.scores == 3 and ranks.selected_indices == 2 and dim(shapes.boxes, 0) == dim(shapes.scores, 0) and dim(shapes.boxes, 1) == dim(shapes.scores, 2) and dim(shapes.boxes, 2) == 4 and dim(shapes.selected_indices, 1) == 3 and dim(shapes.selected_indices, 0) <= dim(shapes.boxes, 0) * dim(shapes.scores, 1) * dim(shapes.boxes, 1) and (present.max_output_boxes_per_class or dim(shapes.selected_indices, 0) == 0) and (not present.max_output_boxes_per_class or (ranks.max_output_boxes_per_class == 0 and numel(shapes.max_output_boxes_per_class) == 1)) and (not present.iou_threshold or (ranks.iou_threshold == 0 and numel(shapes.iou_threshold) == 1)) and (not present.score_threshold or (ranks.score_threshold == 0 and numel(shapes.score_threshold) == 1))",
     "numGroups": "dim(shapes.boxes, 0) * dim(shapes.scores, 1)",
     "maxKeep": "max(1, dim(shapes.boxes, 1))",
     "parallelDispatchFits": "numGroups <= device.limits.maxComputeWorkgroupsPerDimension",
-    "parallelScratchFits": "numGroups * dim(shapes.boxes, 1) * 4 <= storageBufferLimit and numGroups * maxKeep * 4 <= storageBufferLimit and numGroups * 4 <= storageBufferLimit"
   },
   "constants": { "workgroupSize": "tunables.WORKGROUP_SIZE" },
   "bindingSets": {
@@ -562,9 +567,162 @@
           ]
         }
       }
     ]
   },
   "variants": [
     {
       "id": "groups_parallel",
       "priority": 25,

     "score_threshold": { "kind": "tensor", "semantic": "score_threshold", "role": "input", "required": false },
     "selected_indices": { "kind": "tensor", "semantic": "selected_indices", "role": "output", "dtype": "uint32" }
   },
+  "tunables": { "WORKGROUP_SIZE": 256, "NMS_BITMASK_MAX_GROUPS": 4, "NMS_BITMASK_MIN_BOXES": 2048 },
   "derive": {
     "storageBufferLimit": "min(device.limits.maxStorageBufferBindingSize, device.limits.maxBufferSize)",
     "validLayout": "ranks.boxes == 3 and ranks.scores == 3 and ranks.selected_indices == 2 and dim(shapes.boxes, 0) == dim(shapes.scores, 0) and dim(shapes.boxes, 1) == dim(shapes.scores, 2) and dim(shapes.boxes, 2) == 4 and dim(shapes.selected_indices, 1) == 3 and dim(shapes.selected_indices, 0) <= dim(shapes.boxes, 0) * dim(shapes.scores, 1) * dim(shapes.boxes, 1) and (present.max_output_boxes_per_class or dim(shapes.selected_indices, 0) == 0) and (not present.max_output_boxes_per_class or (ranks.max_output_boxes_per_class == 0 and numel(shapes.max_output_boxes_per_class) == 1)) and (not present.iou_threshold or (ranks.iou_threshold == 0 and numel(shapes.iou_threshold) == 1)) and (not present.score_threshold or (ranks.score_threshold == 0 and numel(shapes.score_threshold) == 1))",
     "numGroups": "dim(shapes.boxes, 0) * dim(shapes.scores, 1)",
     "maxKeep": "max(1, dim(shapes.boxes, 1))",
     "parallelDispatchFits": "numGroups <= device.limits.maxComputeWorkgroupsPerDimension",
+    "parallelScratchFits": "numGroups * dim(shapes.boxes, 1) * 4 <= storageBufferLimit and numGroups * maxKeep * 4 <= storageBufferLimit and numGroups * 4 <= storageBufferLimit",
+    "bitmaskSortN": "pow2ceil(dim(shapes.boxes, 1))",
+    "bitmaskMaskWords": "ceilDiv(dim(shapes.boxes, 1), 32)",
+    "bitmaskPairBlocks": "ceilDiv(numGroups * dim(shapes.boxes, 1) * bitmaskMaskWords, tunables.WORKGROUP_SIZE)",
+    "bitmaskDispatchFits": "numGroups <= device.limits.maxComputeWorkgroupsPerDimension and bitmaskPairBlocks <= device.limits.maxComputeWorkgroupsPerDimension",
+    "bitmaskScratchFits": "4 * numGroups * bitmaskSortN <= storageBufferLimit and 4 * numGroups * bitmaskSortN * bitmaskMaskWords <= storageBufferLimit and numGroups * maxKeep * 4 <= storageBufferLimit and numGroups * 4 <= storageBufferLimit"
   },
   "constants": { "workgroupSize": "tunables.WORKGROUP_SIZE" },
   "bindingSets": {
           ]
         }
       }
+    ],
+    "bitmaskSort": [
+      {
+        "name": "scores",
+        "arg": "scores",
+        "semantic": "scores",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "f32"
+      },
+      {
+        "name": "score_threshold",
+        "arg": "score_threshold",
+        "semantic": "score_threshold",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "f32",
+        "length": 1
+      },
+      { "name": "sortKey", "semantic": "sortKey", "buffer": { "type": "storage" }, "elementType": "u32" },
+      { "name": "sortIdx", "semantic": "sortIdx", "buffer": { "type": "storage" }, "elementType": "u32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "boxesPerBatch", "type": "u32", "value": "dim(shapes.boxes, 1)" },
+            { "name": "classes", "type": "u32", "value": "dim(shapes.scores, 1)" },
+            { "name": "centerPointBox", "type": "u32", "value": "attrs.center_point_box" }
+          ]
+        }
+      }
+    ],
+    "bitmaskPairs": [
+      {
+        "name": "boxes",
+        "arg": "boxes",
+        "semantic": "boxes",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "f32"
+      },
+      {
+        "name": "iou_threshold",
+        "arg": "iou_threshold",
+        "semantic": "iou_threshold",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "f32",
+        "length": 1
+      },
+      { "name": "sortIdx", "semantic": "sortIdx", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
+      { "name": "pairMask", "semantic": "pairMask", "buffer": { "type": "storage" }, "elementType": "u32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "boxesPerBatch", "type": "u32", "value": "dim(shapes.boxes, 1)" },
+            { "name": "classes", "type": "u32", "value": "dim(shapes.scores, 1)" },
+            { "name": "centerPointBox", "type": "u32", "value": "attrs.center_point_box" }
+          ]
+        }
+      }
+    ],
+    "bitmaskWalk": [
+      {
+        "name": "max_output_boxes_per_class",
+        "arg": "max_output_boxes_per_class",
+        "semantic": "max_output_boxes_per_class",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "u32",
+        "length": 1
+      },
+      { "name": "sortIdx", "semantic": "sortIdx", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
+      { "name": "pairMask", "semantic": "pairMask", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
+      { "name": "groupSel", "semantic": "groupSel", "buffer": { "type": "storage" }, "elementType": "u32" },
+      { "name": "groupCount", "semantic": "groupCount", "buffer": { "type": "storage" }, "elementType": "u32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "boxesPerBatch", "type": "u32", "value": "dim(shapes.boxes, 1)" },
+            { "name": "classes", "type": "u32", "value": "dim(shapes.scores, 1)" },
+            { "name": "centerPointBox", "type": "u32", "value": "attrs.center_point_box" }
+          ]
+        }
+      }
     ]
   },
   "variants": [
+    {
+      "id": "bitmask_low_occupancy",
+      "priority": 30,
+      "description": "Few `(batch, class)` groups cannot occupy the device with one greedy workgroup each, so a sorted-candidate pass, an all-pairs `IoU` bitmask pass, and a mask-replay walk reproduce greedy suppression with pair parallelism; the shared compact pass lays out `selected_indices` unchanged.",
+      "when": ["validLayout", "present.max_output_boxes_per_class", "present.iou_threshold", "present.score_threshold", "dim(shapes.selected_indices, 0) > 0", "dim(shapes.boxes, 1) >= tunables.NMS_BITMASK_MIN_BOXES", "numGroups <= tunables.NMS_BITMASK_MAX_GROUPS", "bitmaskDispatchFits", "bitmaskScratchFits"],
+      "intermediates": [
+        { "id": "sortKey", "dtype": "uint32", "shape": "[numGroups * bitmaskSortN]" },
+        { "id": "sortIdx", "dtype": "uint32", "shape": "[numGroups * bitmaskSortN]" },
+        { "id": "pairMask", "dtype": "uint32", "shape": "[numGroups * bitmaskSortN * bitmaskMaskWords]" },
+        { "id": "groupSel", "dtype": "uint32", "shape": "[numGroups * maxKeep]" },
+        { "id": "groupCount", "dtype": "uint32", "shape": "[numGroups]" }
+      ],
+      "passes": [
+        {
+          "id": "sort",
+          "name": "NonMaxSuppression.BitmaskSort",
+          "source": {
+            "shader": "nms-bitmask.wgsl.jinja",
+            "inputs": { "stage": "\"sort\"", "sortN": "bitmaskSortN", "maskWords": "bitmaskMaskWords" }
+          },
+          "bindings": "bitmaskSort",
+          "dispatch": { "x": "numGroups" }
+        },
+        {
+          "id": "pairs",
+          "name": "NonMaxSuppression.BitmaskPairs",
+          "source": {
+            "shader": "nms-bitmask.wgsl.jinja",
+            "inputs": {
+              "stage": "\"pairs\"",
+              "sortN": "bitmaskSortN",
+              "maskWords": "bitmaskMaskWords",
+              "numGroups": "numGroups"
+            }
+          },
+          "bindings": "bitmaskPairs",
+          "dispatch": { "x": "bitmaskPairBlocks" }
+        },
+        {
+          "id": "walk",
+          "name": "NonMaxSuppression.BitmaskWalk",
+          "source": {
+            "shader": "nms-bitmask.wgsl.jinja",
+            "inputs": {
+              "stage": "\"walk\"",
+              "sortN": "bitmaskSortN",
+              "maskWords": "bitmaskMaskWords",
+              "maxKeep": "maxKeep"
+            }
+          },
+          "bindings": "bitmaskWalk",
+          "dispatch": { "x": "numGroups" }
+        },
+        {
+          "id": "compact",
+          "name": "NonMaxSuppression.GroupCompact",
+          "source": { "shader": "nms-groups-compact.wgsl.jinja", "inputs": { "maxKeep": "maxKeep" } },
+          "bindings": "parallelCompact",
+          "dispatch": { "x": 1 }
+        }
+      ]
+    },
     {
       "id": "groups_parallel",
       "priority": 25,

build/webgpu/metadata.json CHANGED Viewed

@@ -1,20 +1,21 @@
 {
   "name": "ai.onnx.NonMaxSuppression",
-  "id": "_ai_onnx_nonmaxsuppression_webgpu_d44f6ed",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "QfNxa5siC+dAWXaBHejTokSn3gllrgjOelQ0UNQlxEE=",
-      "manifest.json": "vZuoGvM3nYg4zijLh8n62me7YOGuyRElAeZQelssago=",
-      "nms-groups-compact.wgsl.jinja": "4z2+LMPnLoFRYfgH4bMN9D9eQw/besrivDzeM81JI5c=",
-      "nms-groups-select.wgsl.jinja": "OKea1RdhUe0+cDFuFGIrx32rOxnRbUxgDJA+0FHIHt8=",
-      "non-max-suppression.wgsl.jinja": "D4xiNb1cSvy5W1TjKi3jp4zA/H6besnr0zfsh3OAk+Q=",
-      "test.json": "ehZ2AurQHFMhpSAmDim7zkCfdJL/IdB6VaCVCakwAMY="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
   "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.NonMaxSuppression" }
 }

 {
   "name": "ai.onnx.NonMaxSuppression",
+  "id": "_ai_onnx_nonmaxsuppression_webgpu_7b023d1",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "J+izPem2HPnKTPaigtwj4dAVhN+zxO/L0S0a9ISi3Lc=",
+      "manifest.json": "xCHgSH5g6FOnX97P1CWWiuhGBA+HUT8uAlffDx4Lpv4=",
+      "nms-bitmask.wgsl.jinja": "1DpFLx6M5cSkG++rOrmF5ynIQv5cBYwBISiLy4ZnXqo=",
+      "nms-groups-compact.wgsl.jinja": "sxMP2Rsg3WNsZrnjS+sqK1DLcror9laJ056kifG1vtg=",
+      "nms-groups-select.wgsl.jinja": "mGgrizfXZ55mCm6nyWcUl7HnC2HI8dk0tsDobi6WNcs=",
+      "non-max-suppression.wgsl.jinja": "OGGmipiIE3O4HpLoXqz43q+/HFVuordD5tnL7VHEVzg=",
+      "test.json": "P5MoWNVu+EMZ9bWzzCYkFL3h1iXUidoFdd+ITSBNYiM="
     }
   },
+  "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
   "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.NonMaxSuppression" }
 }

build/webgpu/nms-bitmask.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,291 @@

+// Bitmask NonMaxSuppression for low group counts. One workgroup per (batch,
+// class) group cannot occupy the device when there are few groups, so the
+// greedy loop's IoU work moves to a fully parallel pair pass instead:
+//
+//   sort   one workgroup per group orders candidate boxes by (score desc,
+//          index asc) in storage scratch, exactly the order the greedy
+//          argmax visits survivors;
+//   pairs  every (rank i, 32-rank word) thread stores which higher ranks
+//          box i suppresses, so IoU cost is independent of group count;
+//   walk   one workgroup per group replays the greedy scan over precomputed
+//          words: a rank is kept iff no earlier kept rank suppressed it,
+//          which reproduces greedy selection exactly.
+//
+// Scores compare by monotone IEEE key, so NaN scores are dropped up front --
+// the sweep kernel can never choose them either. The kept list lands in the
+// same `groupSel`/`groupCount` layout the shared compact pass consumes.
+{{ env.wgsl.resourceDeclarations }}
+const WG: u32 = {{ workgroupSize }}u;
+const SORT_N: u32 = {{ source.sortN }}u;
+{% if source.stage != "sort" %}
+const MASK_WORDS: u32 = {{ source.maskWords }}u;
+{% endif %}
+const INVALID: u32 = 0xffffffffu;
+{% if source.stage == "pairs" %}
+const NUM_GROUPS: u32 = {{ source.numGroups }}u;
+{% endif %}
+{% if source.stage == "walk" %}
+const MAX_KEEP: u32 = {{ source.maxKeep }}u;
+{% endif %}
+{% if source.stage != "walk" %}
+{% set needIndexedIou = false %}
+{% set needScoreAt = source.stage == "sort" %}
+{% set needBoxIou = source.stage == "pairs" %}
+{% set needScoreAt = needScoreAt is not defined or needScoreAt %}
+{% set needBoxIou = needBoxIou is not defined or needBoxIou %}
+{% if needBoxIou %}
+fn box_offset(batch: u32, box: u32) -> u32 {
+  return (batch * params.boxesPerBatch + box) * 4u;
+}
+{% endif %}{% if needScoreAt %}
+fn score_at(batch: u32, class_id: u32, box: u32) -> f32 {
+  return scores[(batch * params.classes + class_id) * params.boxesPerBatch + box];
+}
+{% endif %}{% if needBoxIou %}
+fn box_corners(batch: u32, box: u32) -> vec4<f32> {
+  let offset = box_offset(batch, box);
+  let v0 = boxes[offset];
+  let v1 = boxes[offset + 1u];
+  let v2 = boxes[offset + 2u];
+  let v3 = boxes[offset + 3u];
+  if (params.centerPointBox != 0u) {
+    let x_center = v0;
+    let y_center = v1;
+    let half_w = v2 * 0.5;
+    let half_h = v3 * 0.5;
+    return vec4<f32>(y_center - half_h, x_center - half_w, y_center + half_h, x_center + half_w);
+  }
+  return vec4<f32>(v0, v1, v2, v3);
+}
+fn iou_with_box(batch: u32, a: u32, b_box: vec4<f32>) -> f32 {
+  let a_box = box_corners(batch, a);
+  let ay1_raw = a_box.x;
+  let ax1_raw = a_box.y;
+  let ay2_raw = a_box.z;
+  let ax2_raw = a_box.w;
+  let by1_raw = b_box.x;
+  let bx1_raw = b_box.y;
+  let by2_raw = b_box.z;
+  let bx2_raw = b_box.w;
+  let ay1 = min(ay1_raw, ay2_raw);
+  let ax1 = min(ax1_raw, ax2_raw);
+  let ay2 = max(ay1_raw, ay2_raw);
+  let ax2 = max(ax1_raw, ax2_raw);
+  let by1 = min(by1_raw, by2_raw);
+  let bx1 = min(bx1_raw, bx2_raw);
+  let by2 = max(by1_raw, by2_raw);
+  let bx2 = max(bx1_raw, bx2_raw);
+  let inter_y1 = max(ay1, by1);
+  let inter_x1 = max(ax1, bx1);
+  let inter_y2 = min(ay2, by2);
+  let inter_x2 = min(ax2, bx2);
+  let inter = max(0.0, inter_y2 - inter_y1) * max(0.0, inter_x2 - inter_x1);
+  let area_a = max(0.0, ay2 - ay1) * max(0.0, ax2 - ax1);
+  let area_b = max(0.0, by2 - by1) * max(0.0, bx2 - bx1);
+  let denom = area_a + area_b - inter;
+  if (denom <= 0.0) { return 0.0; }
+  return inter / denom;
+}
+{% if needIndexedIou %}
+fn iou(batch: u32, a: u32, b: u32) -> f32 {
+  return iou_with_box(batch, a, box_corners(batch, b));
+}
+{% endif %}{% endif %}
+{% endif %}
+{% if source.stage == "sort" %}
+@compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
+  let tid = lid.x;
+  let group = wid.x;
+  let batch = group / params.classes;
+  let class_id = group % params.classes;
+  let base = group * SORT_N;
+  let score_threshold_value = score_threshold[0];
+  // Candidate keys: ~(monotone IEEE fold) sorts ascending keys as descending
+  // scores; sub-threshold, NaN, and padding slots share the maximum key and an
+  // INVALID index, so they sort strictly after every candidate.
+  for (var r = tid; r < SORT_N; r = r + WG) {
+    var key = 0xffffffffu;
+    var idx = INVALID;
+    if (r < params.boxesPerBatch) {
+      let s = score_at(batch, class_id, r);
+      let b = bitcast<u32>(s);
+      let is_nan = (b & 0x7f800000u) == 0x7f800000u && (b & 0x007fffffu) != 0u;
+      if (!is_nan && s > score_threshold_value) {
+        let folded = select(b | 0x80000000u, ~b, (b & 0x80000000u) != 0u);
+        key = ~folded;
+        idx = r;
+      }
+    }
+    sortKey[base + r] = key;
+    sortIdx[base + r] = idx;
+  }
+  storageBarrier();
+  // Bitonic network over this group's slice. Groups touch disjoint ranges, so
+  // the storage barriers only order this workgroup's own lanes.
+  var size = 2u;
+  loop {
+    if (size > SORT_N) { break; }
+    var stride = size / 2u;
+    loop {
+      if (stride == 0u) { break; }
+      for (var k = tid; k < SORT_N; k = k + WG) {
+        let partner = k ^ stride;
+        if (partner > k) {
+          let ka = sortKey[base + k];
+          let kb = sortKey[base + partner];
+          let ia = sortIdx[base + k];
+          let ib = sortIdx[base + partner];
+          let ascending = (k & size) == 0u;
+          let kBeforeP = ka < kb || (ka == kb && ia <= ib);
+          let needSwap = select(kBeforeP, !kBeforeP, ascending);
+          if (needSwap) {
+            sortKey[base + k] = kb;
+            sortKey[base + partner] = ka;
+            sortIdx[base + k] = ib;
+            sortIdx[base + partner] = ia;
+          }
+        }
+      }
+      storageBarrier();
+      stride = stride / 2u;
+    }
+    size = size * 2u;
+  }
+}
+{% elif source.stage == "pairs" %}
+@compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let rowWords = params.boxesPerBatch * MASK_WORDS;
+  let total = NUM_GROUPS * rowWords;
+  let t = gid.x;
+  if (t >= total) { return; }
+  let group = t / rowWords;
+  let rem = t % rowWords;
+  let i = rem / MASK_WORDS;
+  let word = rem % MASK_WORDS;
+  let base = group * SORT_N;
+  let batch = group / params.classes;
+  var bits = 0u;
+  let idx_i = sortIdx[base + i];
+  if (idx_i != INVALID) {
+    let corners_i = box_corners(batch, idx_i);
+    let iou_threshold_value = iou_threshold[0];
+    let j_start = word * 32u;
+    for (var b = 0u; b < 32u; b = b + 1u) {
+      let j = j_start + b;
+      if (j > i && j < params.boxesPerBatch) {
+        let idx_j = sortIdx[base + j];
+        if (idx_j != INVALID && iou_with_box(batch, idx_j, corners_i) > iou_threshold_value) {
+          bits = bits | (1u << b);
+        }
+      }
+    }
+  }
+  pairMask[(base + i) * MASK_WORDS + word] = bits;
+}
+{% else %}
+const WINDOW_WORDS: u32 = WG / 32u;
+var<workgroup> running: array<atomic<u32>, MASK_WORDS>;
+var<workgroup> admissible: array<u32, WG>;
+var<workgroup> windowSlices: array<u32, WG * WINDOW_WORDS>;
+var<workgroup> picks: array<u32, WG>;
+var<workgroup> wgPickCount: u32;
+@compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
+  let tid = lid.x;
+  let group = wid.x;
+  let base = group * SORT_N;
+  let selBase = group * MAX_KEEP;
+  let max_per_class = min(max_output_boxes_per_class[0], MAX_KEEP);
+  for (var w = tid; w < MASK_WORDS; w = w + WG) {
+    atomicStore(&running[w], 0u);
+  }
+  // Replay greedy selection one WG-rank window at a time. Lanes pre-test their
+  // rank against the running words (stable between windows), one lane resolves
+  // the window serially through the pair-mask window slice -- the only place
+  // rank order matters -- and the kept rows fold into the running words in a
+  // single barrier-free pass with each lane owning its words.
+  var kept = 0u;
+  var window_base = 0u;
+  loop {
+    if (kept >= max_per_class || window_base >= params.boxesPerBatch) { break; }
+    workgroupBarrier();
+    let r = window_base + tid;
+    var open = 0u;
+    if (r < params.boxesPerBatch) {
+      let idx_r = sortIdx[base + r];
+      if (idx_r != INVALID && (atomicLoad(&running[r / 32u]) & (1u << (r % 32u))) == 0u) {
+        open = 1u;
+      }
+    }
+    admissible[tid] = open;
+    // Stage every window rank's pair-mask window slice cooperatively, so the
+    // serial resolution below reads workgroup memory instead of chasing
+    // dependent global loads.
+    let sliceBase = (base + window_base) * MASK_WORDS + window_base / 32u;
+    for (var i = tid; i < WG * WINDOW_WORDS; i = i + WG) {
+      let t = i / WINDOW_WORDS;
+      let w = i % WINDOW_WORDS;
+      var words = 0u;
+      if (window_base + t < params.boxesPerBatch) {
+        words = pairMask[sliceBase + t * MASK_WORDS + w];
+      }
+      windowSlices[t * WINDOW_WORDS + w] = words;
+    }
+    workgroupBarrier();
+    if (tid == 0u) {
+      // Serial window resolution: a window rank is kept iff it survives the
+      // running mask (admissible) and no earlier kept rank of this same
+      // window suppresses it, read straight from its pair-mask row's slice.
+      var window_mask: array<u32, WINDOW_WORDS>;
+      for (var w = 0u; w < WINDOW_WORDS; w = w + 1u) { window_mask[w] = 0u; }
+      var found = 0u;
+      for (var t = 0u; t < WG && kept + found < max_per_class; t = t + 1u) {
+        if (admissible[t] == 0u) { continue; }
+        if ((window_mask[t / 32u] & (1u << (t % 32u))) != 0u) { continue; }
+        let pick = window_base + t;
+        picks[found] = pick;
+        groupSel[selBase + kept + found] = sortIdx[base + pick];
+        found = found + 1u;
+        for (var w = 0u; w < WINDOW_WORDS; w = w + 1u) {
+          window_mask[w] = window_mask[w] | windowSlices[t * WINDOW_WORDS + w];
+        }
+      }
+      wgPickCount = found;
+    }
+    let found = workgroupUniformLoad(&wgPickCount);
+    kept = kept + found;
+    // Fold every kept row of this window into the running words over the
+    // whole flat (row, word) grid: independent loads pipeline across lanes,
+    // and workgroup atomicOr makes concurrent same-word folds safe.
+    for (var i = tid; i < found * MASK_WORDS; i = i + WG) {
+      let k = i / MASK_WORDS;
+      let w = i % MASK_WORDS;
+      atomicOr(&running[w], pairMask[(base + picks[k]) * MASK_WORDS + w]);
+    }
+    window_base = window_base + WG;
+  }
+  if (tid == 0u) { groupCount[group] = kept; }
+}
+{% endif %}

build/webgpu/nms-groups-compact.wgsl.jinja CHANGED Viewed

@@ -9,34 +9,31 @@
 const WG: u32 = {{ workgroupSize }}u;
 const MAX_KEEP: u32 = {{ source.maxKeep }}u;
-var<workgroup> wgTotal: u32;
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
   let tid = lid.x;
   let numGroups = params.batches * params.classes;
-  if (tid == 0u) {
-    var total = 0u;
-    for (var group = 0u; group < numGroups; group = group + 1u) {
-      let batch = group / params.classes;
-      let class_id = group % params.classes;
-      let selBase = group * MAX_KEEP;
-      let n = groupCount[group];
-      for (var k = 0u; k < n; k = k + 1u) {
-        if (total < params.capacity) {
-          let box = groupSel[selBase + k];
-          let row = total * 3u;
-          selected_indices[row] = batch;
-          selected_indices[row + 1u] = class_id;
-          selected_indices[row + 2u] = box;
-        }
-        total = total + 1u;
       }
     }
-    wgTotal = total;
   }
-  let total = workgroupUniformLoad(&wgTotal);
   // Initialize any defensive tail rows to (0,0,0).
   let written = min(total, params.capacity);

 const WG: u32 = {{ workgroupSize }}u;
 const MAX_KEEP: u32 = {{ source.maxKeep }}u;
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(local_invocation_id) lid: vec3<u32>) {
   let tid = lid.x;
   let numGroups = params.batches * params.classes;
+  // Every lane walks the same group-count prefix (uniform, cache-served), so
+  // the running offset needs no shared state while each group's rows copy in
+  // parallel WG-wide strides.
+  var total = 0u;
+  for (var group = 0u; group < numGroups; group = group + 1u) {
+    let batch = group / params.classes;
+    let class_id = group % params.classes;
+    let selBase = group * MAX_KEEP;
+    let n = groupCount[group];
+    for (var k = tid; k < n; k = k + WG) {
+      if (total + k < params.capacity) {
+        let box = groupSel[selBase + k];
+        let row = (total + k) * 3u;
+        selected_indices[row] = batch;
+        selected_indices[row + 1u] = class_id;
+        selected_indices[row + 2u] = box;
       }
     }
+    total = total + n;
   }
   // Initialize any defensive tail rows to (0,0,0).
   let written = min(total, params.capacity);

build/webgpu/nms-groups-select.wgsl.jinja CHANGED Viewed

@@ -16,13 +16,18 @@ const WG: u32 = {{ workgroupSize }}u;
 const MAX_KEEP: u32 = {{ source.maxKeep }}u;
 const INVALID: u32 = 0xffffffffu;
 fn box_offset(batch: u32, box: u32) -> u32 {
   return (batch * params.boxesPerBatch + box) * 4u;
 }
 fn score_at(batch: u32, class_id: u32, box: u32) -> f32 {
   return scores[(batch * params.classes + class_id) * params.boxesPerBatch + box];
 }
 fn box_corners(batch: u32, box: u32) -> vec4<f32> {
   let offset = box_offset(batch, box);
@@ -69,6 +74,7 @@ fn iou_with_box(batch: u32, a: u32, b_box: vec4<f32>) -> f32 {
   if (denom <= 0.0) { return 0.0; }
   return inter / denom;
 }
 var<workgroup> wgScore: array<f32, WG>;

 const MAX_KEEP: u32 = {{ source.maxKeep }}u;
 const INVALID: u32 = 0xffffffffu;
+{% set needScoreAt = needScoreAt is not defined or needScoreAt %}
+{% set needBoxIou = needBoxIou is not defined or needBoxIou %}
+{% if needBoxIou %}
 fn box_offset(batch: u32, box: u32) -> u32 {
   return (batch * params.boxesPerBatch + box) * 4u;
 }
+{% endif %}{% if needScoreAt %}
 fn score_at(batch: u32, class_id: u32, box: u32) -> f32 {
   return scores[(batch * params.classes + class_id) * params.boxesPerBatch + box];
 }
+{% endif %}{% if needBoxIou %}
 fn box_corners(batch: u32, box: u32) -> vec4<f32> {
   let offset = box_offset(batch, box);
   if (denom <= 0.0) { return 0.0; }
   return inter / denom;
 }
+{% endif %}
 var<workgroup> wgScore: array<f32, WG>;

build/webgpu/non-max-suppression.wgsl.jinja CHANGED Viewed

@@ -1,12 +1,17 @@
 {{ env.wgsl.resourceDeclarations }}
 fn box_offset(batch: u32, box: u32) -> u32 {
   return (batch * params.boxesPerBatch + box) * 4u;
 }
 fn score_at(batch: u32, class_id: u32, box: u32) -> f32 {
   return scores[(batch * params.classes + class_id) * params.boxesPerBatch + box];
 }
 fn box_corners(batch: u32, box: u32) -> vec4<f32> {
   let offset = box_offset(batch, box);
@@ -57,6 +62,7 @@ fn iou_with_box(batch: u32, a: u32, b_box: vec4<f32>) -> f32 {
 fn iou(batch: u32, a: u32, b: u32) -> f32 {
   return iou_with_box(batch, a, box_corners(batch, b));
 }
 fn selected_box(row: u32) -> u32 {

 {{ env.wgsl.resourceDeclarations }}
+{% set needScoreAt = needScoreAt is not defined or needScoreAt %}
+{% set needBoxIou = needBoxIou is not defined or needBoxIou %}
+{% if needBoxIou %}
 fn box_offset(batch: u32, box: u32) -> u32 {
   return (batch * params.boxesPerBatch + box) * 4u;
 }
+{% endif %}{% if needScoreAt %}
 fn score_at(batch: u32, class_id: u32, box: u32) -> f32 {
   return scores[(batch * params.classes + class_id) * params.boxesPerBatch + box];
 }
+{% endif %}{% if needBoxIou %}
 fn box_corners(batch: u32, box: u32) -> vec4<f32> {
   let offset = box_offset(batch, box);
 fn iou(batch: u32, a: u32, b: u32) -> f32 {
   return iou_with_box(batch, a, box_corners(batch, b));
 }
+{% endif %}
 fn selected_box(row: u32) -> u32 {

build/webgpu/test.json CHANGED Viewed

@@ -1174,6 +1174,117 @@
         "score_threshold": { "dtype": "float32", "shape": [], "data": { "kind": "values", "values": [0.5] } }
       },
       "outputs": { "selected_indices": { "dtype": "uint32", "shape": [0, 3], "tolerance": 0 } }
     }
   ]
 }

         "score_threshold": { "dtype": "float32", "shape": [], "data": { "kind": "values", "values": [0.5] } }
       },
       "outputs": { "selected_indices": { "dtype": "uint32", "shape": [0, 3], "tolerance": 0 } }
+    },
+    {
+      "name": "bitmask_1class_2048_grid_sorted_keep128",
+      "provenance": {
+        "notes": "2048 disjoint linspace boxes with ascending scores route to the bitmask variant; zero overlaps make selection the top-128 scores, so the sorted walk must emit indices in descending index order."
+      },
+      "attrs": { "center_point_box": 0 },
+      "inputs": {
+        "boxes": {
+          "dtype": "float32",
+          "shape": [1, 2048, 4],
+          "data": { "kind": "linspace", "start": 0.0, "end": 8191.0 }
+        },
+        "scores": {
+          "dtype": "float32",
+          "shape": [1, 1, 2048],
+          "data": { "kind": "linspace", "start": 0.1, "end": 1.0 }
+        },
+        "max_output_boxes_per_class": { "dtype": "uint32", "shape": [], "data": { "kind": "values", "values": [128] } },
+        "iou_threshold": { "dtype": "float32", "shape": [], "data": { "kind": "values", "values": [0.5] } },
+        "score_threshold": { "dtype": "float32", "shape": [], "data": { "kind": "values", "values": [0.05] } }
+      },
+      "outputs": { "selected_indices": { "dtype": "uint32", "shape": [128, 3], "tolerance": 0 } }
+    },
+    {
+      "name": "bitmask_1class_4096_two_positions_collapse",
+      "provenance": {
+        "notes": "4096 boxes cycling two identical positions collapse under the pair mask to the best-scored box of each position; descending scores make those boxes 0 and 1."
+      },
+      "attrs": { "center_point_box": 0 },
+      "inputs": {
+        "boxes": {
+          "dtype": "float32",
+          "shape": [1, 4096, 4],
+          "data": { "kind": "cycle", "values": [0.0, 0.0, 10.0, 10.0, 100.0, 100.0, 110.0, 110.0] }
+        },
+        "scores": {
+          "dtype": "float32",
+          "shape": [1, 1, 4096],
+          "data": { "kind": "linspace", "start": 1.0, "end": 0.001 }
+        },
+        "max_output_boxes_per_class": { "dtype": "uint32", "shape": [], "data": { "kind": "values", "values": [4096] } },
+        "iou_threshold": { "dtype": "float32", "shape": [], "data": { "kind": "values", "values": [0.5] } },
+        "score_threshold": { "dtype": "float32", "shape": [], "data": { "kind": "values", "values": [0.0] } }
+      },
+      "outputs": { "selected_indices": { "dtype": "uint32", "shape": [2, 3], "tolerance": 0 } }
+    },
+    {
+      "name": "bitmask_2batch_2048_two_positions",
+      "provenance": {
+        "notes": "Two batches of 2048 two-position boxes give two groups on the bitmask route; each batch keeps its own pair of survivors, exercising per-group sort, mask, and walk offsets."
+      },
+      "attrs": { "center_point_box": 0 },
+      "inputs": {
+        "boxes": {
+          "dtype": "float32",
+          "shape": [2, 2048, 4],
+          "data": { "kind": "cycle", "values": [0.0, 0.0, 10.0, 10.0, 100.0, 100.0, 110.0, 110.0] }
+        },
+        "scores": {
+          "dtype": "float32",
+          "shape": [2, 1, 2048],
+          "data": { "kind": "linspace", "start": 1.0, "end": 0.001 }
+        },
+        "max_output_boxes_per_class": { "dtype": "uint32", "shape": [], "data": { "kind": "values", "values": [2048] } },
+        "iou_threshold": { "dtype": "float32", "shape": [], "data": { "kind": "values", "values": [0.5] } },
+        "score_threshold": { "dtype": "float32", "shape": [], "data": { "kind": "values", "values": [0.0] } }
+      },
+      "outputs": { "selected_indices": { "dtype": "uint32", "shape": [4, 3], "tolerance": 0 } }
+    },
+    {
+      "name": "bitmask_1class_2048_all_tied_scores_keep64",
+      "provenance": {
+        "notes": "All 2048 disjoint boxes share one score, so ordering is decided purely by the index tie-break; the walk must keep boxes 0..63 exactly."
+      },
+      "attrs": { "center_point_box": 0 },
+      "inputs": {
+        "boxes": {
+          "dtype": "float32",
+          "shape": [1, 2048, 4],
+          "data": { "kind": "linspace", "start": 0.0, "end": 8191.0 }
+        },
+        "scores": { "dtype": "float32", "shape": [1, 1, 2048], "data": { "kind": "cycle", "values": [0.75] } },
+        "max_output_boxes_per_class": { "dtype": "uint32", "shape": [], "data": { "kind": "values", "values": [64] } },
+        "iou_threshold": { "dtype": "float32", "shape": [], "data": { "kind": "values", "values": [0.5] } },
+        "score_threshold": { "dtype": "float32", "shape": [], "data": { "kind": "values", "values": [0.0] } }
+      },
+      "outputs": { "selected_indices": { "dtype": "uint32", "shape": [64, 3], "tolerance": 0 } }
+    },
+    {
+      "name": "bitmask_1class_2048_threshold_half_keep32",
+      "provenance": {
+        "notes": "A 0.55 score threshold admits just under half of 2048 disjoint linspace-scored boxes and `max_output_boxes_per_class` clamps the kept prefix to 32 on the bitmask route."
+      },
+      "attrs": { "center_point_box": 0 },
+      "inputs": {
+        "boxes": {
+          "dtype": "float32",
+          "shape": [1, 2048, 4],
+          "data": { "kind": "linspace", "start": 0.0, "end": 8191.0 }
+        },
+        "scores": {
+          "dtype": "float32",
+          "shape": [1, 1, 2048],
+          "data": { "kind": "linspace", "start": 0.1, "end": 1.0 }
+        },
+        "max_output_boxes_per_class": { "dtype": "uint32", "shape": [], "data": { "kind": "values", "values": [32] } },
+        "iou_threshold": { "dtype": "float32", "shape": [], "data": { "kind": "values", "values": [0.5] } },
+        "score_threshold": { "dtype": "float32", "shape": [], "data": { "kind": "values", "values": [0.55] } }
+      },
+      "outputs": { "selected_indices": { "dtype": "uint32", "shape": [32, 3], "tolerance": 0 } }
     }
   ]
 }