Repository navigation
feat(cpu-plugin): execute NMS once for dynamic output - #2112
Merged
Merged
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2112 +/- ##
==========================================
+ Coverage 80.54% 81.19% +0.65%
==========================================
Files 413 429 +16
Lines 194861 215388 +20527
Branches 194861 215388 +20527
==========================================
+ Hits 156944 174878 +17934
- Misses 32390 34737 +2347
- Partials 5527 5773 +246
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
justinchuby
force-pushed
the
feat/plugin-nms-kernel-sized
branch
from
August 25, 2026 14:28
848ef7b to
d06a290
Compare
Reuse KernelSizedOutput so NonMaxSuppression selects once, returns owned Int64 rows, and materializes one dynamically sized plugin output. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624
justinchuby
force-pushed
the
feat/plugin-nms-kernel-sized
branch
from
August 25, 2026 15:01
d06a290 to
17a3c5b
Compare
justinchuby
added a commit
that referenced
this pull request
Aug 25, 2026
## Summary - add bounded f32 CUDA `NonMaxSuppression` with exact CPU row/order semantics from #2112 - reuse #2113's generic `DeviceWorkspace` prepare → 8-byte count → ORT device allocation → materialize policy - keep boxes, scores, optional scalars, suppression state, and final selected rows device-resident - claim only static contiguous bounded geometry and fail closed outside it ## Claimed surface and semantics The CUDA EP claims ONNX NMS opset 10 for: - contiguous static Float32 boxes `[batch, boxes, 4]` - contiguous static Float32 scores `[batch, classes, boxes]` - optional device-resident Int64 `max_output_boxes_per_class` - optional device-resident Float32 `iou_threshold` / `score_threshold` - `center_point_box` 0 or 1 - at most **256 boxes** and **256 batch×class groups** Output is exact contiguous Int64 `[selected,3]` rows `[batch,class,box]`. Score filtering is strict (`score > threshold`); ties keep lower box indices; output groups are batch-major then class-major. Defaults, empty inputs, zero max output, no selections, both center modes, negative-max rejection, and threshold semantics follow #2112's CPU implementation. Float16/BFloat16/Float64, symbolic geometry, strided inputs, malformed ranks/dtypes/scalars, more than 256 boxes, or more than 256 groups decline at claim time where metadata permits. ## Algorithm and complexity The prepare pipeline launches one block per batch/class group. Threads filter scores and perform a deterministic shared-memory bitonic sort (`O(N log²N)` parallel work). Thread 0 then performs bounded greedy suppression (`O(N²)` worst case), explicitly capped at `N <= 256`; groups run independently in parallel. A second one-thread metadata launch sums per-group selected counts. Materialization launches one block per group and writes the exact rows using the persisted selected indices/counts. Selection never runs twice. ## DeviceWorkspace, memory, and capture This uses #2113's existing `KernelSizedOutputPolicy::DeviceWorkspace`; no new deferred-output policy exists. Workspace contains fixed-capacity selected indices, per-group counts, and the 8-byte total count. It is overflow-checked, 8-byte aligned, `StepScoped`, and allocated through the EP governor. There is no NMS-side `cudaMalloc`. Exactly **8 bytes** (the total selected count) cross D2H. Full boxes/scores/scalars and output remain on device; instrumentation reports `full_input_d2h_bytes == 0`. Workspace remains live through materialization under the generic DeviceWorkspace contract. CUDA graph capture fails closed because the 8-byte D2H synchronization and host-driven ORT dynamic allocation are not capture-safe. ## Actual GPU validation RTX 4060 Laptop GPU, driver 591.55, pinned CUDA 13.1 wheels, serialized: - `cargo test -p onnx-runtime-ep-cuda --features gpu-tests --test nms_gpu -- --test-threads=1` — **6 passed, 0 failed, 0 ignored** - real ORT CUDA plugin NMS E2E — **1 passed**; assigned to our EP, dynamic `[4,3]`, exact rows, prepare/count/materialize each once, D2H=8, full-input D2H=0 - CUDA Unique DeviceWorkspace regression E2E — **2 passed** - unchanged CPU NMS targeted oracle — **8 passed** - CUDA covered-op duplicate guard — **1 passed** - source-derived CPU/CUDA gap guard — **1 passed** - conformance-profile duplicate guard — **1 passed** - `cargo clippy -p onnx-runtime-ep-cuda --all-targets --features cuda,gpu-tests -- -D warnings` — passed - `cargo clippy -p onnx-runtime-ep-cuda-plugin --all-targets --features cuda -- -D warnings` — passed - `cargo fmt -p onnx-runtime-ep-cuda -p onnx-runtime-ep-cuda-plugin` — passed The broad pre-existing `every_covered_op_has_a_conformance_entry` guard remains red for `PagedAttention`, `Mish`, `Celu`, `TensorScatter`, and newly merged `Unique`; NMS has its dedicated profile entry and is absent from that missing set. ## Mutation evidence Each mutation was applied independently, caught, and reverted: - inverted IoU comparison -> multi-batch/class CPU parity failed - ignored score threshold -> threshold test retained a below-threshold row - swapped center-box width/height conversion -> center-mode parity failed - reversed equal-score index tie order -> tie-order test failed - reported full-input D2H -> `full_input_d2h_bytes == 0` guard failed - ran the prepare/selection kernel twice -> exactly-once launch guard failed (`2 != 1`) - skipped device materialization -> real ORT E2E returned zero rows and failed exact output ## Representative structural measurement For **batch=1, classes=2, boxes=32, max_output=8**: - prepare/sort/suppress launches: **1** (grid has 2 group blocks) - count-reduction launches: **1** - materialization launches: **1** - selected rows: **16** - governed workspace: **272 bytes** - D2H: **8 bytes** - full-input D2H: **0 bytes** No latency or throughput claim is made; no A/B benchmark was run. ## Not verified - Linux, H100, or H200 execution - geometry beyond the explicit 256-box / 256-group bound (declined) - half/f64 or strided CUDA NMS (declined) - CUDA graph capture (explicitly unsupported) - throughput versus ORT/PyTorch CUDA NMS ## Independent post-rejection revision (`52fd41566`) **Ownership:** Luv rejected the prior revision. Batty independently authored this revision under reviewer lockout; Leon did not advise, pair, or contribute. Blocking findings fixed: 1. **Default score threshold parity:** omitted `score_threshold` now uses IEEE `-infinity`, matching CPU/ONNX rather than `f32::MIN`. A discriminating GPU test uses scores `[f32::MIN, -INF]`: omitted and explicit `-INF` both select only the `f32::MIN` box. 2. **Signed-zero ordering:** the device comparator now implements Rust `f32::total_cmp` bit ordering before descending comparison. Identical overlapping boxes with lower-index `-0` and higher-index `+0` select the higher index, matching CPU. 3. **Optional scalar claim gate:** every present optional input must have a known static rank-0 shape at `supports_op`. `[1]`, `[0]`, symbolic rank-1, and rank-2 shapes decline before kernel creation; valid `[]` remains claimed. 4. **Parallel telemetry race:** all NMS GPU tests use one poison-recovering target-wide mutex. The default-parallel and explicit serial runs are both 9/9 green. Removing serialization reproduces telemetry corruption (`prepare_launches` > 1) and concurrent VMM reservation failures. Validation on the physical CUDA host: ```text cargo test -p onnx-runtime-ep-cuda --features gpu-tests --test nms_gpu # 9 passed, 0 failed, 0 ignored (default threads) cargo test -p onnx-runtime-ep-cuda --features gpu-tests --test nms_gpu -- --test-threads=1 # 9 passed, 0 failed, 0 ignored cargo test -p onnx-runtime-ep-cpu nms_ --lib # 8 passed NXRT_REQUIRE_ORT_TESTS=1 cargo test -p onnx-runtime-ep-cuda-plugin --features cuda --test cuda_unique_ort_e2e # 3 passed: NMS E2E 1 + Unique DeviceWorkspace regressions 2 cargo clippy -p onnx-runtime-ep-cuda --all-targets --features gpu-tests -- -D warnings cargo clippy -p onnx-runtime-ep-cuda-plugin --all-targets --features cuda -- -D warnings # both passed ``` Mutation matrix, all killed: - restore the `f32::MIN` default -> omitted-threshold test returns empty instead of box 0; - restore numeric-equality/index zero comparator -> signed-zero test selects box 0 instead of box 1; - remove optional scalar shape gate -> `[1]` is claimed and the claim test fails; - replace the shared mutex with per-call mutexes -> default-parallel suite exposes cross-test telemetry/VMM races. No broad NMS behavior or plugin contract changed beyond the reviewed findings. --------- Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
NonMaxSuppressionto feat(plugin): add kernel-sized outputs and CPU Unique #2101's host-onlyKernelSizedOutputcontractexecuteand plugin ComputeSemantics
Implements the existing ONNX NonMaxSuppression opset-10 behavior without changing selection ordering:
[batch, spatial_dimension, 4], Float32[batch, classes, spatial_dimension], Float32max_output_boxes_per_class, Float32iou_threshold, Float32score_thresholdcenter_point_box0 and 1[num_selected_indices, 3]rows[batch, class, box]score > threshold), stable lower-index tie ordering, complete per-batch/per-class orderingAbsent and trailing optional inputs remain positional. Strided host boxes/scores are materialized by the existing CPU dense helpers. Device inputs fail at the host-accessibility gate before selection; there is no implicit payload D2H.
Single execution and materialization
compute_owned_outputperforms validation, dense materialization, selection, shape construction, and byte encoding once. Native execution validates/copies that owned result into its preallocated tensor. Plugin Compute receives the same ownedKernelSizedOutput, asks ORT for the exact[selected,3]allocation, and performs one final byte copy.Representative test geometry: 1 batch × 2 classes × 3 boxes selected 4 rows. Selection counter: exactly 1. Materialized output: 4 × 3 × 8 = 96 bytes, one copy.
Validation
ours=["NonMaxSuppression"], dynamic shape[4,3], exact Int64 rowscargo clippy -p onnx-runtime-ep-cpu -p onnx-runtime-ep-plugin -p onnx-runtime-ep-cpu-plugin --all-targets -- -D warnings— passedcargo fmt -p onnx-runtime-ep-cpu -p onnx-runtime-ep-plugin -p onnx-runtime-ep-cpu-plugin— passedMutation evidence
Each mutation was applied independently, caught, and reverted:
2 != 1)Host-only limitation / CUDA follow-up
This PR is deliberately CPU/plugin single-execution only. Kernel-sized outputs are host-only today; CUDA NMS must wait for the CUDA Unique device-sized-output policy to prove device allocation/materialization ownership. No CUDA NMS implementation or device payload copy is introduced here.
Performance
No throughput claim is made; no A/B benchmark was run. The structural improvement is removal of the formerly required double-algorithm path: the plugin now selects once and copies only the final owned output bytes once.
Latest rebase / Rust quality attribution (2026-08-25)
Rebased again onto latest
origin/main, which now contains #2115 merge commit0be2d23fe6f719a5cc70fe91e53a80d26f7f2739. Verified directly:The PR diff remains exactly these six NMS/plugin files and no server/engine files:
crates/onnx-runtime-ep-cpu/src/kernels/{mod.rs,selection.rs}crates/onnx-runtime-ep-plugin/src/compute.rscrates/onnx-runtime-ep-cpu-plugin/tests/{plugin_ort_e2e.rs,shape_inference_coverage.rs}crates/onnx-runtime-ep-cpu-plugin/tests/fixtures/non_max_suppression_kernel_sized/model.onnx.textprotoThe exact required native-backend Rust-quality command now passes with
-D warnings:Post-rebase targeted evidence is also green: CPU NMS 8 passed; generic KernelSizedOutput 4 passed; plugin census 5 passed; real ORT NMS E2E 1 passed; real ORT Unique regression 1 passed.