[cudax][cuco] Port fixed_capacity_map benchmarks from cuCollections

[sleeepyjack]

Closes #9595.

Summary

Ports the relevant cuCollections static_map NVBench benchmarks to cudax fixed_capacity_map.

Adjusts the default cooperative-group size for fixed_capacity_map to align with the optimal setup used by cuCollections.

Changes

• Adds shared cudax/cuco benchmark helpers under cudax/benchmarks/bench/cuco/common/.
• Ports static_map insert benchmarks as fixed_capacity_map benchmarks:
  • fixed_capacity_map_insert_unique_capacity
  • fixed_capacity_map_insert_unique_occupancy
  • fixed_capacity_map_insert_uniform_multiplicity
• Ports static_map contains benchmarks as fixed_capacity_map benchmarks:
  • fixed_capacity_map_contains_unique_capacity
  • fixed_capacity_map_contains_unique_occupancy
  • fixed_capacity_map_contains_unique_matching_rate

[Jacobfaib]

Nit: testing/benchmarks or any non-library code shouldn't be __ugly

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 90277c0b-bcad-4a7e-92fd-a595e7e89268

📥 Commits

Reviewing files that changed from the base of the PR and between 2ea4d34 and dee3726.

📒 Files selected for processing (2)

• cudax/benchmarks/bench/cuco/common/key_generator.cuh
• cudax/test/cuco/fixed_capacity_map/test_shared_memory.cu

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

• cudax/benchmarks/bench/cuco/common/key_generator.cuh

---

📝 Walkthrough

Summary by CodeRabbit

• New Features
  • Added new CUDA benchmarks for fixed-capacity map insert and contains operations.
  • Introduced reusable benchmark presets and deterministic key-generation support for consistent, parameterized workloads.
• Bug Fixes
  • Updated the default probing configuration for fixed-capacity maps to better match common benchmark behavior.
• Tests
  • Refreshed fixed-capacity map shared-memory test setup and kernel launch/bounds logic.
• Chores
  • Added benchmark parameter presets for input size, occupancy, matching rate, and key multiplicity.

Walkthrough

Adds CUCO benchmark defaults and key generation utilities, then adds nvbench benchmarks for fixed_capacity_map::insert and fixed_capacity_map::contains. The PR also changes fixed_capacity_map’s default probing scheme.

Changes

fixed_capacity_map benchmark suite

Layer / File(s)	Summary
Benchmark defaults cudax/benchmarks/bench/cuco/common/defaults.cuh	Defines key/value type axes, default input count/occupancy/matching rate/seed, and sweep ranges for input size, occupancy, multiplicity, and matching rate.
Key generator utility cudax/benchmarks/bench/cuco/common/key_generator.cuh	Adds distribution::unique/uniform tags, device functors for uniform key generation and dropout, key_generator::generate/dropout methods, and dist_from_state conversion with nvbench type strings.
insert_async benchmark cudax/benchmarks/bench/cuco/fixed_capacity_map/insert.cu	Generates keys/pairs, builds a fixed_capacity_map sized from NumInputs/Occupancy, times insert_async and clear_async, and registers three benchmark variants.
contains_async benchmark cudax/benchmarks/bench/cuco/fixed_capacity_map/contains.cu	Builds and populates a map, applies dropout to control match rate, times contains_async, and registers three benchmark variants.
Default probing scheme change cudax/include/cuda/experimental/__cuco/fixed_capacity_map.cuh	Changes fixed_capacity_map's default _ProbingScheme from linear_probing<1, hash<_Key>> to linear_probing<4, hash<_Key>>.

Assessment against linked issues

Objective	Addressed	Explanation
Add benchmarks for fixed_capacity_map::insert and fixed_capacity_map::contains [#9595]	✅

Assessment against linked issues: Out-of-scope changes

Code Change	Explanation
fixed_capacity_map default probing scheme changed to linear_probing<4, hash<_Key>> (cudax/include/cuda/experimental/__cuco/fixed_capacity_map.cuh)	This changes library default behavior and is not required by the benchmark-addition objective in #9595.
Shared-memory test naming and kernel parameter cleanup (cudax/test/cuco/fixed_capacity_map/test_shared_memory.cu)	This test refactor is unrelated to adding the two benchmarks.
Shared-memory kernel bounds check renamed from n to num_keys (cudax/test/cuco/fixed_capacity_map/test_shared_memory.cu)	This is a local test rename, not part of the benchmark feature scope.

Suggested reviewers: shwina, jrhemstad, bernhardmgruber

---

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

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (3)

cudax/benchmarks/bench/cuco/common/key_generator.cuh (1)

207-248: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

suggestion: dropout() unconditionally shuffles [begin, end) at line 247, even when keep_prob >= 1.0 (no keys actually dropped). This side effect isn't mentioned in the docstring (lines 212-223), which only describes the drop/replace behavior. Worth documenting the shuffle's purpose (or gating it behind the keep_prob < 1.0 branch if it's only meant to intersperse dropped keys) so callers relying on this utility understand the full behavior.

As per path instructions for cudax/**/*, review should focus on "experimental API clarity" among other things.

Doc clarification example

   //! Drops keys with probability `1 - keep_prob` using the provided execution policy.
   //!
   //! Replaced keys are sampled from `[N, max_key]`, where `N` is the number of keys in the range.
+  //! The full range is always shuffled afterward, regardless of `keep_prob`.
   //!

Source: Path instructions

cudax/benchmarks/bench/cuco/fixed_capacity_map/insert.cu (1)

46-74: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

suggestion: the device/stream/memory-resource setup, key→pair transform, and map construction (lines 46-74) are duplicated almost verbatim in contains.cu (lines 46-75). Since common/ already houses shared benchmark utilities, consider extracting a shared helper (e.g. a build_map(state, keys) or similar in common/) to avoid drift between the two benchmarks as the map API evolves.

cudax/benchmarks/bench/cuco/common/defaults.cuh (1)

28-39: 🚀 Performance & Scalability | 🔵 Trivial

suggestion: trim the low-occupancy sweep The 10% point drives fixed_capacity_map_insert_unique_occupancy and fixed_capacity_map_contains_unique_occupancy to ~1e9 slots for int64_t/int64_t instances (~16 GiB of slots alone, before overhead), and that multiplies across 9 occupancy values and 2 valid type combinations. Consider raising the minimum occupancy or lowering n so the sweep stays practical on smaller GPUs.

Source: Path instructions

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: d0bd7519-6228-406b-9cb5-c72d712b117d

📥 Commits

Reviewing files that changed from the base of the PR and between 3f87270 and 2ea4d34.

📒 Files selected for processing (5)

• cudax/benchmarks/bench/cuco/common/defaults.cuh
• cudax/benchmarks/bench/cuco/common/key_generator.cuh
• cudax/benchmarks/bench/cuco/fixed_capacity_map/contains.cu
• cudax/benchmarks/bench/cuco/fixed_capacity_map/insert.cu
• cudax/include/cuda/experimental/__cuco/fixed_capacity_map.cuh

[PointKernel]

@sleeepyjack can you please share the benchmark results comparing cuco::static_map and cudax::cuco::fixed_capacity_map? Just wanted to make sure we didn't bring any surprises during the migration.

[github-actions]

😬 CI Workflow Results

🟥 Finished in 1h 08m: Pass: 98%/55 | Total: 5h 08m | Max: 1h 08m | Hits: 98%/36787

See results here.