Fix BlockTopK +/-0.0 handling

[pauleonix]

Description

Cause

The bug is caused by combining detail::radix::traits_t<KeyT>::bit_ordered_inversion_policy::inverse() and BFEDigitExtractor<KeyT>{}.Digit(key) as well as comparing non-digit-extracted keys with a mask built from extracted digits.

In BlockTopK we handle -0.0 manually and get rid of the corresponding bit-pattern which keeps the bug from appearing in the min case. But in the max case we invert the keys after that transformation which means that instead of both 0.0 and -0.0 being mapped to the 0.0 bit-pattern, they both get mapped to the -0.0 bit pattern.

The digit-extractor has built-in negative-zero handling for FP (it extracts the 0.0 bit pattern even from -0.0). But this handling conflicts with our use of the extracted digit to build up a prefix mask b/c we compare the mask directly with the key which still has the -0.0 bit-pattern. So keys that were registered in the histogram-pass as 0.0-bit-pattern wont be filtered correctly later.

This issue can be reproduced by having input data with tied keys being 0.0 or -0.0 and using the max select-direction. This causes the candidate bin to count all the 0.0/-0.0 keys but none of them participating in the next histogram pass since they do not compare equal with the prefix mask.

Fix

The actual fix is just 14 lines of code. All the rest is testing infrastructure I brought over from #9066 to make sure the bug is fixed and remains fixed. I found and fixed the bug there originally but the PR was de-prioritized to focus on #9077. So to still get this fix in, I extracted it into its own PR.

While one could also fix the issue by first inverting and then doing the -0.0 to 0.0 mapping, I decided to get rid of the bit pattern inversion completely. Instead, we reverse the histogram index for the max-direction (and undo the reversal when transforming the bin index into the prefix mask). This is also what other radix-ranking facilities seem to be doing (even though they don't have the same issue of building up a mask and comparing it to the keys). It feels less bug-prone, easier to understand/review and seemed to have a neutral to positive effect in benchmarks (in the context of my bigger PR #9066, I did not re-benchmark this PR).

In a follow-up PR we should look into getting a digit-extractor w/o the -0.0 handling to avoid the unnecessary conditional (dead code) introduced by it due to our manual handling of -0.0.

Checklist

• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[coderabbitai]

Note: CodeRabbit is enabled on this repository as a convenience for maintainers and contributors. Use your best judgment when considering its review comments and
suggestions — a suggested change may be inadequate, unnecessary, or safe to ignore.
Contributors are not expected to address every comment. Human reviews are what
ultimately matter for merging.

Summary

This PR fixes negative zero (-0.0) handling in CUB’s BlockTopK radix selection by making the histogram/binning and subsequent k-th prefix reconstruction explicitly SelectDirection-aware for both min and max, preventing 0.0/-0.0 from being collapsed to the wrong bit pattern in the max path.

Changes

Core algorithm fix

cub/cub/block/specializations/block_topk_air.cuh

• Updated compute_histograms to reverse bucket indexing for select::max (bucket = num_buckets - 1 - digit) instead of relying on the previous bit/digit inversion interaction.
• Updated k-th prefix reconstruction for the next pass to undo the bucket reversal when computing kth_key_prefix (mirroring the max flip using num_buckets - 1 - pass_state.bucket).
• Generalized the final scatter selection predicate to use the correct comparison operator per direction (less for min, greater for max) while keeping candidate classification (key_prefix == kth_prefix) consistent.

Test coverage

cub/test/catch2_test_block_topk.cu

• Added regression coverage for {Min,Max}Keys preserving -0.0 in output (bitwise via bit_repr and sign-bit checks), including both full-tile and partial-tile scenarios, plus ±infinity.

cub/test/catch2_test_block_topk_common.cuh

• Added shared test helpers for deterministic key generation around a boundary key (including explicit +/-0.0 tie mixing when the boundary is 0.0), direction-aware reference top-k computation, and bit-pattern diagnostics (bit_repr).

cub/test/catch2_test_device_segmented_topk_keys.cu

• Extended the device batched/segmented top-k regression test to run for both selection directions and added checks ensuring -0.0f is preserved in the output.

Walkthrough

block_topk_air is updated so SelectDirection is threaded through compute_histograms (bucket index flipped for min), the per-pass kth_key_prefix update, and the scatter predicate (direction-dependent comparator). A new shared test header and block-level test file are added, and the device-segmented regression test is generalized to cover both min and max directions.

Changes

block_topk_air SelectDirection correctness and test coverage

Layer / File(s)	Summary
block_topk_air histogram and scatter direction fixes cub/cub/block/specializations/block_topk_air.cuh	compute_histograms gains SelectDirection template parameter and flips bucket indices for min selection. Per-pass kth_key_prefix digit derivation is direction-dependent. Scatter predicate replaced with a less/greater comparator selected by SelectDirection. bit_ordered_inversion alias is relocated.
Shared test infrastructure cub/test/catch2_test_block_topk_common.cuh	New header adds RNG alias, distinct_keys, random_boundary_key, gen_keys_from_boundary_key (with float ±inf and ±0.0 special-casing), boundary_key_generator, overhang_generator, sorted_top_k, to_span, and bit_repr for bit-level mismatch diagnostics.
Block-level tests and segmented regression update cub/test/catch2_test_block_topk.cu, cub/test/catch2_test_device_segmented_topk_keys.cu	Adds topk_kernel/check_topk harness and three test cases: floating-point edge-case value preservation, boundary-key correctness on full tiles, and negative-zero sign preservation. Generalizes the segmented regression test from min-only to both min and max directions.

Suggested reviewers

• gevtushenko
• NaderAlAwar

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 39b7f3ca-5295-4397-8077-40d5d40d05f5

📥 Commits

Reviewing files that changed from the base of the PR and between f3c2b4b and c34379f.

📒 Files selected for processing (4)

• cub/cub/block/specializations/block_topk_air.cuh
• cub/test/catch2_test_block_topk.cu
• cub/test/catch2_test_block_topk_common.cuh
• cub/test/catch2_test_device_segmented_topk_keys.cu

[github-actions]

🥳 CI Workflow Results

🟩 Finished in 6h 02m: Pass: 100%/287 | Total: 2d 07h | Max: 52m 21s | Hits: 98%/199242

See results here.

[coderabbitai]

🧹 Nitpick comments (1)

cub/test/catch2_test_block_topk_common.cuh (1)

30-30: 🧹 Nitpick | 🔵 Trivial | 💤 Low value

suggestion: Include uses quote syntax; coding guidelines require angle brackets for all header inclusions.

-#include "catch2_test_device_topk_common.cuh"
+#include <catch2_test_device_topk_common.cuh>

Source: Coding guidelines

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 64ebd300-52f8-4527-8659-6045c59ee53f

📥 Commits

Reviewing files that changed from the base of the PR and between bfa14c1 and b035295.

📒 Files selected for processing (2)

• cub/test/catch2_test_block_topk.cu
• cub/test/catch2_test_block_topk_common.cuh

🚧 Files skipped from review as they are similar to previous changes (1)

• cub/test/catch2_test_block_topk.cu