[cuco] Adds support for 1 and 2 byte key and value types in fixed_capacity_map

[ryanjspears]

Description

Makes progress on #9987

Adds support for 1 and 2 byte key and value types in fixed_capacity_map

The insert and contains tests now cover 1 and 2 byte types and dynamically assign the number of keys based on byte size. This is needed because 400 could not fit into 1 byte, thus needs to be changed accordingly.

Checklist

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

[copy-pr-bot]

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.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Changes

fixed_capacity_map now supports smaller size-compatible key/value types, derives exact-width packed CAS operands, and adds positive and expected-failure compile-time coverage for size constraints.

Packed map type support

Layer / File(s)	Summary
Packed type constraints and representation cudax/include/cuda/experimental/__cuco/detail/open_addressing/open_addressing_ref_impl.cuh, cudax/include/cuda/experimental/__cuco/fixed_capacity_map_ref.cuh	packed_cas derives its unsigned operand type from value bit width. fixed_capacity_map_ref allows mapped types up to eight bytes and requires power-of-two key and pair sizes.
Expanded fixed-capacity map coverage cudax/test/cuco/fixed_capacity_map/test_insert_and_contains.cu	Tests add 8-bit and 16-bit keys, adjust key-dependent sizing and sentinels, and validate supported slot sizes.
Constraint diagnostic tests cudax/test/CMakeLists.txt, cudax/test/cuco/fixed_capacity_map/test_*_fail.cu	Expected-failure targets verify diagnostics for non-power-of-two key and slot sizes.

Possibly related PRs

• NVIDIA/cccl#10001: Also changes the packed_cas path and its atomic operations.

Suggested labels: cuco

Suggested reviewers: caugonnet

---

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

[coderabbitai]

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

cudax/include/cuda/experimental/__cuco/detail/open_addressing/open_addressing_ref_impl.cuh (1)

701-710: 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift

important: device_buffer<__value_type> only guarantees alignof(__value_type), but packed_type can need stricter alignment. Reinterpreting __address, __expected, and __desired as packed_type* is UB for small pair-like slots such as uint8_t/uint8_t or uint16_t/uint16_t. Add an aligned packed contract or gate this path on alignof(__value_type) >= alignof(packed_type).

🧹 Nitpick comments (1)

cudax/test/cuco/fixed_capacity_map/test_insert_and_contains.cu (1)

38-39: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

suggestion: Add mixed-width key/value instantiations. map_type uses key_type for both template arguments, so these cases cover (uint8_t, uint8_t) and (uint16_t, uint16_t) but not the independently changed mapped-type contract, such as 32-bit keys with 8/16-bit values. Add a separate key/value type list or a few explicit mixed-width cases. The supplied map_type alias and new type list show that only same-width pairs are currently exercised.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: f4680b6f-be66-4b99-97c2-382fbccbe2b1

📥 Commits

Reviewing files that changed from the base of the PR and between 73a1a4b and df921f8.

📒 Files selected for processing (3)

• cudax/include/cuda/experimental/__cuco/detail/open_addressing/open_addressing_ref_impl.cuh
• cudax/include/cuda/experimental/__cuco/fixed_capacity_map_ref.cuh
• cudax/test/cuco/fixed_capacity_map/test_insert_and_contains.cu

[coderabbitai]

Actionable comments posted: 1

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: c9f14394-1f67-4604-9307-997d26e163f6

📥 Commits

Reviewing files that changed from the base of the PR and between df921f8 and ddadb4d.

📒 Files selected for processing (3)

• cudax/include/cuda/experimental/__cuco/detail/open_addressing/open_addressing_ref_impl.cuh
• cudax/include/cuda/experimental/__cuco/fixed_capacity_map_ref.cuh
• cudax/test/cuco/fixed_capacity_map/test_insert_and_contains.cu

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

• cudax/include/cuda/experimental/__cuco/detail/open_addressing/open_addressing_ref_impl.cuh

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

cudax/test/cuco/fixed_capacity_map/test_insert_and_contains.cu (1)

140-149: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

suggestion: Exercise a mixed-width pair through insert and contains, not only sizeof. The <uint8_t, uint32_t> assertion checks map instantiation and layout, but does not execute the packed CAS path where width or representation mismatches would surface.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 81d8afdd-0f78-460f-8331-697a0c6abea1

📥 Commits

Reviewing files that changed from the base of the PR and between b886248 and da4bc4a.

📒 Files selected for processing (6)

• cudax/include/cuda/experimental/__cuco/detail/open_addressing/open_addressing_ref_impl.cuh
• cudax/include/cuda/experimental/__cuco/fixed_capacity_map_ref.cuh
• cudax/test/CMakeLists.txt
• cudax/test/cuco/fixed_capacity_map/test_insert_and_contains.cu
• cudax/test/cuco/fixed_capacity_map/test_key_size_fail.cu
• cudax/test/cuco/fixed_capacity_map/test_slot_size_fail.cu

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

• cudax/include/cuda/experimental/__cuco/fixed_capacity_map_ref.cuh
• cudax/include/cuda/experimental/__cuco/detail/open_addressing/open_addressing_ref_impl.cuh

[sleeepyjack]

Thanks @ryanjspears for working on this!
Note that there is another open PR that might result in some conflicts: #10001 . CC @PointKernel

[PointKernel]

This is probably not a simple task to target right now. The original cuco::pair uses the next power-of-two alignment. Consequently, cuco::pair<int, int> is 8-byte aligned and cuco::pair<int16_t, int16_t> is 4-byte aligned, whereas the corresponding cuda::std::pair types are only 4- and 2-byte aligned.

When possible, cuco packs the key-value pair into a single integer and updates it with packed_cas. For example, a 4-byte key and 4-byte value use one 8-byte atomic CAS, which requires 8-byte alignment, not the 4-byte alignment provided by cuda::std::pair<int, int>. A power-of-two size check alone does not guarantee this alignment. The default allocation may happen to be sufficiently aligned, but external span storage and the local expected and desired objects are not guaranteed to be.

@ryanjspears Thanks for working on this. I haven’t found a clean way to preserve the original slot-alignment guarantee with cuda::std::pair yet, so we should resolve that design question before proceeding with subword support.

[PointKernel]

@ryanjspears I thought about this more, and I think we have a clean way forward without losing the packed CAS optimization.

Currently, packed_cas accesses the local expected and desired objects through packed_type*, although their declared types do not guarantee the required alignment. We can avoid these potentially misaligned accesses by bit-casting them into packed_type values.

For the storage, we only need to check the base address when constructing the storage reference. The slots are contiguous and sizeof(value_type) == sizeof(packed_type), so if the first slot satisfies atomic_ref<packed_type>::required_alignment, every slot does.

The storage reference can retain that result and use packed_cas only for suitably aligned, bitwise-compatible slot types, falling back to the existing non-packed path otherwise. This preserves the single CAS for common aligned cases such as cuda::std::pair<int, int> while safely supporting external storage.

The resulting dispatch would be approximately:

if constexpr (has_packable_representation)
{
  if (storage_ref.is_packed_cas_aligned())
  {
    return packed_cas(...);
  }
}

return non_packed_cas(...);

Could you please update the PR along these lines?

[ryanjspears]

Thanks for the feedback @PointKernel. Yes, I will submit an update shortly!

[ryanjspears]

@PointKernel I updated the dispatch to your suggestion. I also created test_misaligned_storage.cu which tests insertions over an under-aligned external storage to hit the non-packed fallback.

[PointKernel]

/ok to test 71d8664

[PointKernel]

@ryanjspears Thanks for the fast turnaround. The code looks great. I just pushed a single line change where the old code checks against __expected while we only update __expected_packed. Will ship it once CI passes.

[ryanjspears]

Thanks @PointKernel , is there another cuco issue that would be good for me to tackle? Shall I tackle #9989?

[PointKernel]

Thanks @PointKernel , is there another cuco issue that would be good for me to tackle? Shall I tackle #9989?

That will be great!

[PointKernel]

The new small-type tests expose the existing clang-tidy false positive; atomic behavior is unchanged.

[PointKernel]

/ok to test b6ce4f2