Repository navigation
feat(plugin): add kernel-sized outputs and CPU Unique - #2101
Merged
Merged
Conversation
Add an opt-in owned-output contract and route CPU Unique through it so ORT allocates data-dependent outputs after one algorithm pass. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #2101 +/- ##
==========================================
- Coverage 80.76% 80.57% -0.20%
==========================================
Files 414 428 +14
Lines 203411 209851 +6440
Branches 203411 209851 +6440
==========================================
+ Hits 164294 169086 +4792
- Misses 33605 35066 +1461
- Partials 5512 5699 +187
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
This was referenced Aug 25, 2026
justinchuby
added a commit
that referenced
this pull request
Aug 25, 2026
## Summary - extend #2101's additive kernel-sized-output API with an explicit `DeviceWorkspace` policy - add bounded CUDA `Unique` for static contiguous flattened Float32 inputs - keep the full input and all output payloads device-resident; copy only the 8-byte unique count D2H - allocate ORT outputs from returned metadata, then materialize directly into device outputs - register and claim the honest CUDA surface and decline everything else before execution ## Foundation ancestry Built from current `origin/main`; PR #2101's squash merge commit `e11051f225544e46f9ddb5757a20dd8c357e256d` is an ancestor. (`5318e9e1e` was the pre-squash PR head.) ## Device policy `KernelSizedOutputPolicy` is additive. Existing CPU `Unique` and every non-opted kernel retain the #2101 behavior unchanged: - `HostOwned` (default): kernel returns owned host bytes, as before. - `DeviceWorkspace`: metadata and reusable device state are prepared in executor-governed step workspace; after ORT allocates concrete outputs, the kernel materializes directly into those device buffers. The workspace remains owned by ORT for the complete Compute callback. No per-dispatch `cudaMalloc`/`cudaFree` is introduced. Both single-node and routed plugin paths preserve positional absent slots and device intermediates. ## CUDA algorithm and honest claim surface Claimed: - `ai.onnx::Unique`, opset 11+ - no `axis` attribute (flattened mode) - static, contiguous Float32 input - `numel <= min(1024, device max_threads_per_block)` - sorted and first-appearance (`sorted=0`) order - Y plus any positional optional subset of indices/inverse_indices/counts Declined at placement with a reason: axis mode, dynamic extents, strided layouts, non-Float32 input, larger input, malformed attributes/arity. Capture is explicitly unsupported because metadata requires an 8-byte synchronous D2H and host-side ORT allocation. Phase 1 performs one bounded one-block bitonic sort plus grouping on device. It preserves CPU semantics: all NaNs group together and sort after numbers; signed zeros group; representative Y uses the first original index; unsorted groups follow first appearance. The plan leaves first indices, counts, and inverse mapping in governed workspace. Phase 2 writes only requested outputs. This is two CUDA kernel phases, but the sort/group algorithm runs exactly once. ## Transfers and workspace The algorithm issues exactly one D2H call totaling 8 bytes for non-empty inputs and zero for empty input. Runtime transfer-count assertions fail if any additional D2H occurs; no full-input D2H is hidden behind custom telemetry. Input H2D is only the ordinary graph-boundary upload when the caller supplies host input. Workspace is `align8(3 * N * sizeof(u32)) + 8` bytes: | N | input boundary H2D | count D2H | workspace | metadata launches | materialize launches | debug host-visible total* | |---:|---:|---:|---:|---:|---:|---:| | 64 | 256 B | 8 B | 776 B | 1 | 1 | 6401 us | | 256 | 1024 B | 8 B | 3080 B | 1 | 1 | 1473 us | | 1024 | 4096 B | 8 B | 12296 B | 1 | 1 | 4061 us | `*` Windows RTX 4060, CUDA 13 dynamic-loading packages, debug test binary, warm NVRTC module. These noisy host totals include allocation/copies/synchronization and are accounting only. No speed claim or CPU/ORT A/B claim is made. ## Validation Actual GPU, serial: - direct CUDA parity: sorted/unsorted discriminating order, all four outputs, optional subset, empty/all-equal/all-distinct, NaN payload grouping, signed zero, exact shapes/dtypes/contents - telemetry: one metadata launch, one materialization launch, exactly one 8-byte D2H call, zero full-input D2H, exact governed workspace - registry/dtype/capture checks against the real CUDA registry/kernel - real ORT CUDA plugin E2E: `Unique` compiled/executed by `cuda_ep`, all four outputs and Y+inverse subset, exact values, workspace placement served - CPU Unique regression: 9/9 passed - plugin unit regression: 345/345 passed - EP API suite: 100/100 passed Mutation matrix, all killed on the physical GPU: 1. force sorted order for `sorted=0` → discriminating-order failure 2. stop grouping NaNs → indices/group-count failure 3. skip inverse/count writes → exact-content failure 4. launch the full metadata algorithm twice → launch-count failure 5. copy the full input D2H → runtime transfer-count failure Exact formatting/lint commands (all green): ```text cargo fmt -p onnx-runtime-ep-api -p onnx-runtime-ep-plugin -p onnx-runtime-ep-cuda -p onnx-runtime-ep-cuda-plugin -- --check cargo clippy -p onnx-runtime-ep-api -p onnx-runtime-ep-plugin --all-targets -- -D warnings 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 ``` ## Explicit non-verification Not implemented by design: axis mode, f16/bf16/integer dtypes, strided or dynamic input shapes, inputs above the bounded limit, or capture. No Linux/macOS GPU run and no release-mode performance A/B were performed. The broad CUDA unit suite reached 615 passed / 33 ignored with one unrelated existing failure: `runtime::tests::a_module_restored_from_cached_ptx_computes_what_a_compiled_one_does` rejects CUDA-13 PTX on this RTX driver with `CUDA_ERROR_UNSUPPORTED_PTX_VERSION`; the production runtime's CUBIN fallback and all changed GPU/plugin tests pass. ## Required Fast follow-up: reviewable fixtures **Attribution:** the required `Fast` failure was caused by this PR checking in two binary `.onnx` graph fixtures. The repository-wide `checked_in_graph_fixtures_are_not_binary_onnx` guard correctly rejected them; the guard was not weakened or allowlisted. Commit `6868c4663` replaces both binaries with human-readable `model.onnx.textproto` files and makes the CUDA E2E loader follow the established plugin convention: read TextFormat, validate the parsed model contract, convert in memory with `onnx_std::textproto::to_binary`, then call ORT `CreateSessionFromArray`. Before ORT conversion, the test now asserts for each fixture: - IR version 11 and ai.onnx opset 24; - exactly one `Unique` node and one Float32 graph input named `X`; - all four positional node-output slots; - the optional fixture's exact `[Y, absent, inverse_indices, absent]` slot pattern; - exact graph-output names and Float32/Int64 dtypes; - the `sorted` integer attribute; - no initializers. The all-output ORT E2E fixture preserves `sorted=0` and exercises first-appearance remapping; the optional-subset fixture keeps its intended `sorted=1`. Post-fix validation: ```text cargo test -p onnx-genai-metadata --test textproto_fixtures cargo test -p onnx-runtime-ep-cuda --features gpu-tests gpu_flattened_unique_matches_cpu_semantics_and_transfer_contract --lib -- --nocapture NXRT_REQUIRE_ORT_TESTS=1 cargo test -p onnx-runtime-ep-cuda-plugin --features cuda --test cuda_unique_ort_e2e -- --nocapture cargo fmt -p onnx-runtime-ep-api -p onnx-runtime-ep-plugin -p onnx-runtime-ep-cuda -p onnx-runtime-ep-cuda-plugin -- --check cargo clippy -p onnx-runtime-ep-api -p onnx-runtime-ep-plugin --all-targets -- -D warnings 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 ``` Results: fixture guard 1/1 passed; direct CUDA Unique 1/1 passed on RTX 4060; real ORT CUDA plugin E2E 2/2 passed; all exact Clippy commands passed with `-D warnings`. ## Focused re-review follow-up: unsorted fixture fidelity **Attribution:** focused Opus review correctly found that the initial binary-to-text conversion changed `unique_all_outputs` from `sorted=0` to `sorted=1`, then changed the assertions to match. That was a real fidelity and coverage regression. Commit `f9330aaca` restores the checked-in TextFormat attribute to `sorted=0`, verifies the parsed attribute is zero, and restores the original E2E oracle: - Y `[2,1,3,4]` - indices `[0,1,3,4]` - inverse `[0,1,1,2,3,2]` - counts `[1,2,2,1]` The ONNX proto3 encoder normally elides an explicit integer zero. For these in-memory ORT fixtures the loader now restores the legal wire-level `AttributeProto.i = 0` payload when `AttributeProto.type == INT`; otherwise ORT drops the valueless attribute and applies Unique's schema default `sorted=1`. The checked-in source remains reviewable textproto, and the real ORT E2E now proves the first-appearance path rather than merely inspecting source text. Post-fix results: textproto guard 1/1; direct CUDA Unique GPU test 1/1; real ORT CUDA plugin E2E 2/2; targeted CUDA and CUDA-plugin Clippy passed with `-D warnings`. No CUDA algorithm code changed. --------- Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624
justinchuby
added a commit
that referenced
this pull request
Aug 25, 2026
## Summary - migrate CPU/plugin `NonMaxSuppression` to #2101's host-only `KernelSizedOutput` contract - refactor selection into one owned result shared by native `execute` and plugin Compute - claim NMS in the plugin only after adding it to the exact kernel-sized strategy census - tighten NMS mixed-edge dtype advertisement and slot constraints ## Semantics Implements the existing ONNX NonMaxSuppression opset-10 behavior without changing selection ordering: - boxes `[batch, spatial_dimension, 4]`, Float32 - scores `[batch, classes, spatial_dimension]`, Float32 - optional positional scalars: Int64 `max_output_boxes_per_class`, Float32 `iou_threshold`, Float32 `score_threshold` - `center_point_box` 0 and 1 - output Int64 `[num_selected_indices, 3]` rows `[batch, class, box]` - strict score filter (`score > threshold`), stable lower-index tie ordering, complete per-batch/per-class ordering - empty boxes/classes/batches, zero max output, NaN scores, and threshold boundary behavior covered Absent 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_output` performs 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 owned `KernelSizedOutput`, 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 - CPU NMS targeted suite — **8 passed** - legacy overlapping NMS regression — **1 passed** - generic KernelSizedOutput plugin unit suite — **4 passed** - plugin shape/strategy census — **5 passed** - real ORT plugin NMS E2E — **1 passed**; assignment `ours=["NonMaxSuppression"]`, dynamic shape `[4,3]`, exact Int64 rows - real ORT plugin Unique regression — **1 passed** - `cargo clippy -p onnx-runtime-ep-cpu -p onnx-runtime-ep-plugin -p onnx-runtime-ep-cpu-plugin --all-targets -- -D warnings` — passed - `cargo fmt -p onnx-runtime-ep-cpu -p onnx-runtime-ep-plugin -p onnx-runtime-ep-cpu-plugin` — passed ## Mutation evidence Each mutation was applied independently, caught, and reverted: - inverted IoU keep comparison -> overlapping-box regression failed - ignored score threshold -> strict threshold/NaN test failed - swapped center-box width/height conversion -> center-point test failed - ran selection twice -> once-only counter test failed (`2 != 1`) - skipped materialization copy -> real ORT E2E returned zero rows and failed exact output - removed NMS from the census -> exact census test failed - emptied the census -> nonempty guard failed ## 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 commit `0be2d23fe6f719a5cc70fe91e53a80d26f7f2739`. Verified directly: ```text git merge-base --is-ancestor 0be2d23 HEAD # exit 0 ``` 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.rs` - `crates/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.textproto` The exact required native-backend Rust-quality command now passes with `-D warnings`: ```text cargo clippy --locked --all-targets \ -p onnx-genai-engine -p onnx-genai-server \ --features onnx-genai-engine/native-backend,onnx-genai-server/native-backend \ -- -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. 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
KernelSizedOutputcontract to the public EP APIUniquethrough the contract and remove it from the plugin declined-shape censusAPI contract
Kernel::has_kernel_sized_outputs()defaults to false. Opted kernels implementexecute_kernel_sized(inputs, requested_outputs)and return one positionalOption<KernelSizedOutput>per node output slot. Each present value owns its bytes and reports concrete shape and dtype; each omitted optional slot returnsNone.The plugin validates slot count/presence, graph-declared dtype, every extent's ORT
i64range, checked shape products, checked byte lengths, and exact returned buffer length before allocation. It allocates no absent output and performs one materialization copy per present ORT output. Routed host subgraphs move owned bytes directly into intermediate buffers and copy only final ORT sinks. Non-opted kernels retain the existing infer-shape -> allocate -> execute path.Device policy
This first slice is host-only. Device-placed EP outputs or non-host-accessible inputs are rejected before calling the algorithm, with an actionable host-EP message. The path never performs an implicit D2H payload copy.
Unique semantics
CPU
Uniquenow builds its plan once and derives Y, indices, inverse_indices, and counts from that plan. Existing flattened/axis, sorted/unsorted, NaN/signed-zero, empty, and strided semantics are preserved. Mixed output dtypes come from the returned values and must match graph value info. Optional output positions remain positional.Tests
{ ai.onnx::Unique }Uniqueassigned tocpu_ep; all four dynamic outputs andY + inverse_indicessubset matched shapes/dtypes/contentsonnx-runtime-ep-plugin: 360 passed, 3 ignoredonnx-runtime-ep-cpu-plugin: 103 passed, 1 ignoredonnx-runtime-ep-api: 100 passedExact lint/format workflow:
Mutation matrix (all killed): algorithm planning twice; empty dynamic-strategy census; skipped materialization copy; disabled malformed-byte validation.
Allocation/copy accounting
For the representative six-element f32 all-output case, assertions account for four owned result payloads totaling 128 bytes (16 B Y + 32 B indices + 48 B inverse + 32 B counts), four final ORT output allocations, one algorithm plan, and exactly four materialization copies totaling 128 bytes. This is payload/final-allocation accounting, not a claim about every algorithm-internal
Vecallocation. No speed claim is made.Reuse and scope
NMS can reuse the positional optional-output and owned-byte contract without putting its selection algorithm in shape inference. A future CUDA implementation can reuse result validation but needs an explicit device materialization policy; it must not silently copy arbitrary inputs to host.
Not implemented or verified here: NMS integration, CUDA Unique/NMS, device-resident deferred outputs, release-mode timing, or non-Windows ORT E2E. CI is asynchronous and was not waited on; required local suites and the real ORT E2E are green.