Skip to content

[Rust] Add Block TopK keys and pairs - #12246

Draft
tpn wants to merge 321 commits into
NVIDIA:mainfrom
tpn:codex/cuda-coop-rust-topk
Draft

tpn wants to merge 321 commits into
NVIDIA:mainfrom
tpn:codex/cuda-coop-rust-topk

Conversation

@tpn

@tpn tpn commented Oct 24, 2026 •

Copy link
Copy Markdown
Contributor

Add Rust Block TopK min/max selection for keys and key/value pairs. Results are unsorted, preserve their inputs and return fresh payloads with exactly the selected prefix valid. Runtime or constant counts select full or partial inputs; every input slot must remain readable even for empty selections. Key/value types and alignments are independent, and equal-key order is unspecified.

This follows MergeSortPairs #12245 and uses the pinned Oxide compiler. Review the 21-file owning diff; Files changed against main also includes preceding Rust API layers.

Local Rust 1.99.1 validation passed host tests, strict checks, six inference probes, 17 native requests, the exact-capacity regression and thirteen compiler controls. Both release consumers passed artifact inspection. The capacity regression retains one native call and 2,112 bytes of scratch. Numerical checks, three sanitizers and four invalid-count/input traps with immediate successful health runs passed.

Earlier source reviews and MSRV evidence retain their historical attribution. Timing was not repeated, and no performance or historical code-equality claim is made. Current-head host CI passes (110 tests/doctests; 18 native tests excluded), along with the three upstream checks.

tpn added 30 commits October 20, 2026 21:32
Describe the implemented physical and logical warp algorithms and their
scratch requirements. Link payload readers to Load/Store and storage.

Signed-off-by: Trent Nelson <[email protected]>
Expose inferred fresh Scan results with common operators, typed seeds
and checked temporary storage. Add a generic Load-Scan-Store example
and focused native, compiler and runtime coverage.

Signed-off-by: Trent Nelson <[email protected]>
Preserve the recognized prefix query through generic group helpers and
pin the Rust workspace and examples to the supporting compiler revision.

Signed-off-by: Trent Nelson <[email protected]>
Preserve input values and expose results only at the group root.
Support typed operators, counts, algorithms and scratch policies through
the native provider, with generic examples and checked safe outputs.

Signed-off-by: Trent Nelson <[email protected]>
Copy intermediate values only for lanes that own a batch. Preserve
collective execution and both output layouts, and check copy ownership
without evaluating uninitialized values.

Signed-off-by: Trent Nelson <[email protected]>
Reduce payload columns across physical and logical warps into fresh
payloads with calculated capacity and checked result validity.

Provide native CUB specialization, generic examples and documentation
through the existing checked payload and managed storage interfaces.
Distinct kernel shapes can request identical Warp source. Include the
validated XYZ dimensions in each identity so one batch retains unique
symbols and ordered responses without rejecting valid consumers.

Cover different sizes and equal-thread-count shapes in Warp families.
Preserve input while converting striped and blocked per-thread layouts
into a fresh payload with the same type, extent and alignment. Use the
existing native CUB provider and compiler-planned scratch storage.

Add a generic standalone example, a guide and focused behavior and
participation coverage.
Shift blocked-order values by one item while preserving input and
returning a fresh payload. Keep the exposed edge invalid until the
caller initializes it, and guard the native output read at that
boundary.

Include inferred generic helpers, a runnable example and guide, and
behavioral coverage for boundary validity and collective
participation.
Preserve the input and return fresh differences with typed neighbors
and counts. Copy inactive suffix items and validate the complete input
before native reads, including calls with a zero count.

Reuse the existing payload and scratch interfaces for the seven CUB
forms. Add the generic example, guide and focused compiler and runtime
coverage. Correct the synchronization function name in Shuffle
documentation.
Return fresh head and tail flags while preserving numeric inputs.
Support typed outside neighbors and independent paired results with
existing payload, validity and shared-storage interfaces.

Add a generic XYZ example, guide and focused compiler and runtime
coverage for the eight supported CUB forms.

Signed-off-by: Trent Nelson <[email protected]>
Return fresh striped counters with a selectable type and extent while
preserving the input payload. Validate sample bounds before native CUB
execution and plan the complete counter and algorithm scratch region.

Include a generic example, behavioral checks and the compiler
requirement in the Rust guide.

Signed-off-by: Trent Nelson <[email protected]>
Sort selected key bits in either direction into a fresh blocked payload.
Preserve input values through the checked collective and storage APIs,
with inferred types and independent static or runtime bounds.
Return fresh key and value payloads with stable selected-digit ordering
and preserve both inputs. Use the native CUB pair specialization with
independent payload types, alignments and checked validity.

Signed-off-by: Trent Nelson <[email protected]>
Return stable selected-digit ranks as fresh i32 payloads and preserve
integer keys. Support static digit widths and checked uniform bounds on
complete XYZ blocks through the native CUB rank specialization.

Signed-off-by: Trent Nelson <[email protected]>
Sort numeric keys into fresh results with checked prefix validity and
configurable order and temporary storage.

Signed-off-by: Trent Nelson <[email protected]>
Carry numeric values with their keys through Block and Warp merge
sorts, returning fresh payloads with checked prefix validity.
Add a generic example and document storage and comparator behavior.

Signed-off-by: Trent Nelson <[email protected]>
Provide min and max selection for keys and key-value pairs with typed
counts, partial input limits, fresh results and checked validity.
Select one native call per operation so exclusive storage charges one
scratch region. Include generic examples and focused diagnostics.

Signed-off-by: Trent Nelson <[email protected]>
Keep the owning implementation and tests unchanged while applying
its scoped compiler, lock and documentation compatibility updates.

Signed-off-by: Trent Nelson <[email protected]>
Keep the owning implementation and tests unchanged while applying
its scoped compiler, lock and documentation compatibility updates.

Signed-off-by: Trent Nelson <[email protected]>
Keep the owning implementation and tests unchanged while applying
its scoped compiler, lock and documentation compatibility updates.

Signed-off-by: Trent Nelson <[email protected]>
Keep the owning implementation and tests unchanged while applying
its scoped compiler, lock and documentation compatibility updates.

Signed-off-by: Trent Nelson <[email protected]>
Keep the owning implementation and tests unchanged while applying
its scoped compiler, lock and documentation compatibility updates.

Signed-off-by: Trent Nelson <[email protected]>
Keep the owning implementation and tests unchanged while applying
its scoped compiler, lock and documentation compatibility updates.

Signed-off-by: Trent Nelson <[email protected]>
Keep the owning implementation and tests unchanged while applying
its scoped compiler, lock and documentation compatibility updates.

Signed-off-by: Trent Nelson <[email protected]>
Keep the owning implementation and tests unchanged while applying
its scoped compiler, lock and documentation compatibility updates.

Signed-off-by: Trent Nelson <[email protected]>
tpn added 25 commits October 23, 2026 14:08
Merge the corrected prerequisite while preserving this owner.

Signed-off-by: Trent Nelson <[email protected]>
Merge the corrected prerequisite while preserving this owner.

Signed-off-by: Trent Nelson <[email protected]>
Merge the corrected prerequisite while preserving this owner.

Signed-off-by: Trent Nelson <[email protected]>
Merge the corrected prerequisite while preserving this owner.

Signed-off-by: Trent Nelson <[email protected]>
Merge the corrected prerequisite while preserving this owner.

Signed-off-by: Trent Nelson <[email protected]>
Use the scoped barrier diagnostic for conflicting Warp accesses.
Accept the two finite sibling-warp read loops after reconvergence,
matching the migrated checker and existing logical-warp policy.

Signed-off-by: Trent Nelson <[email protected]>
Merge the corrected prerequisite while preserving this owner.

Signed-off-by: Trent Nelson <[email protected]>
Merge the corrected prerequisite while preserving this owner.

Signed-off-by: Trent Nelson <[email protected]>
Merge the corrected prerequisite while preserving this owner.

Signed-off-by: Trent Nelson <[email protected]>
Merge the corrected prerequisite while preserving this owner.

Signed-off-by: Trent Nelson <[email protected]>
Merge the corrected prerequisite while preserving this owner.

Signed-off-by: Trent Nelson <[email protected]>
Merge the corrected prerequisite while preserving this owner.

Signed-off-by: Trent Nelson <[email protected]>
Merge the corrected prerequisite while preserving this owner.

Signed-off-by: Trent Nelson <[email protected]>
Merge the corrected prerequisite while preserving this owner.

Signed-off-by: Trent Nelson <[email protected]>
Merge the corrected prerequisite while preserving this owner.

Signed-off-by: Trent Nelson <[email protected]>
Merge the corrected prerequisite while preserving this owner.

Signed-off-by: Trent Nelson <[email protected]>
Merge the corrected prerequisite while preserving this owner.

Signed-off-by: Trent Nelson <[email protected]>
Merge the corrected prerequisite while preserving this owner.

Signed-off-by: Trent Nelson <[email protected]>
Merge the corrected prerequisite while preserving this owner.

Signed-off-by: Trent Nelson <[email protected]>
Merge the corrected prerequisite while preserving this owner.

Signed-off-by: Trent Nelson <[email protected]>
Merge the corrected prerequisite while preserving this owner.

Signed-off-by: Trent Nelson <[email protected]>
Merge the corrected prerequisite while preserving this owner.

Signed-off-by: Trent Nelson <[email protected]>
Merge the corrected prerequisite while preserving this owner.

Signed-off-by: Trent Nelson <[email protected]>
Merge the corrected prerequisite while preserving this owner.

Signed-off-by: Trent Nelson <[email protected]>
Merge the corrected prerequisite while preserving this owner.

Signed-off-by: Trent Nelson <[email protected]>
@copy-pr-bot

copy-pr-bot Bot commented Oct 24, 2026

Copy link
Copy Markdown
Contributor

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@cccl-authenticator-app cccl-authenticator-app Bot moved this from Todo to In Progress in CCCL Oct 24, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Status: In Progress

Development

Successfully merging this pull request may close these issues.

1 participant