Use the new tuning API internally for detail::select|three_way_partition::dispatch and DevicePartition

[bernhardmgruber]

• No SASS changes for cub.bench.partition.three_way.base on SM75;80;86;90;100
• Use signed offset type for DevicePartition #8971 (required to avoid SASS changes)
• No SASS changes for cub.bench.partition.if.base on SM75;80;86;90;100
• No SASS changes for cub.bench.partition.flagged.base on SM75;80;86;90;100

Fixes: #8879
Fixes: #8380

[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: 4e509214-cb2e-4a86-96bf-61f80365fafd

📥 Commits

Reviewing files that changed from the base of the PR and between 5d6ec8e and 2e36a8b.

📒 Files selected for processing (6)

• cub/benchmarks/bench/partition/flagged.cu
• cub/benchmarks/bench/partition/if.cu
• cub/benchmarks/bench/partition/three_way.cu
• cub/cub/device/device_partition.cuh
• cub/test/catch2_test_device_partition_env.cu
• thrust/thrust/system/cuda/detail/partition.h

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

• cub/benchmarks/bench/partition/flagged.cu
• thrust/thrust/system/cuda/detail/partition.h
• cub/benchmarks/bench/partition/if.cu
• cub/test/catch2_test_device_partition_env.cu
• cub/cub/device/device_partition.cuh

---

📝 Walkthrough

Summary by CodeRabbit

• Tests

  • Added tuning verification tests to ensure partition operations (If, Flagged, three-way If) run correctly under tuned execution environments.

• Refactor

  • Streamlined partition benchmark and execution wiring to use environment-based invocation and direct partition calls, enabling execution tuning and simplifying launch/temporary-storage flow. No public API changes.

suggestion:

Walkthrough

This PR rewires DevicePartition/select dispatch to the new tuning API, replaces manual temp-size/allocation with env-based dispatch and policy_selector functors in benchmarks, updates Thrust partition dispatch, and adds tuning tests validating the tuned execution paths.

Changes

Tuning API Integration for DevicePartition and Select Dispatch

Layer / File(s)	Summary
DevicePartition environment-based overload refactoring cub/cub/device/device_partition.cuh	Removed private partition_impl; refactored env-based Flagged/If and three-way overloads to compute signed offsets inline, build default policy selectors, and call detail::select::dispatch / detail::three_way_partition::dispatch via dispatch_with_env_and_tuning.
Benchmark adoption of policy selector and tuned environment cub/benchmarks/bench/partition/flagged.cu, cub/benchmarks/bench/partition/if.cu, cub/benchmarks/bench/partition/three_way.cu	Introduce policy_selector functors; simplify type aliases; use raw device pointers; remove two-step temp-size/alloc/dispatch flow and instead use caching_allocator_t + cub_bench_env (with cuda::execution::tune(policy_selector<T>{}) when enabled) and direct cub::DevicePartition::* calls.
Thrust partition integration using CUB dispatch thrust/thrust/system/cuda/detail/partition.h	Added dispatch_partition helper performing query+execute via cub::detail::select::dispatch; unified partition helper to use single THRUST_INDEX_TYPE_DISPATCH that calls dispatch_partition.
Tuning test coverage for DevicePartition operations cub/test/catch2_test_device_partition_env.cu	Added tuning includes, defined capability-driven partition and three-way policy selectors, and added C2H_TEST cases validating tuning for DevicePartition::If, DevicePartition::Flagged, and three-way If with assertions on selected counts and tuned block sizes.

Assessment against linked issues

Objective	Addressed	Explanation
Use new tuning API for detail::select::dispatch + DevicePartition [#8879]	✅
Use new tuning API for detail::three_way_partition::dispatch [#8380]	✅

Possibly related PRs

• NVIDIA/cccl#8971: Overlaps on DevicePartition env-overload signed-offset handling and related selector construction.

Suggested reviewers

• elstehle

---

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

[coderabbitai]

Actionable comments posted: 2

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 2e262785-9664-4f8f-b738-a61ffdd14e4e

📥 Commits

Reviewing files that changed from the base of the PR and between 3f440b9 and f978ca6.

📒 Files selected for processing (6)

• cub/benchmarks/bench/partition/flagged.cu
• cub/benchmarks/bench/partition/if.cu
• cub/benchmarks/bench/partition/three_way.cu
• cub/cub/device/device_partition.cuh
• cub/test/catch2_test_device_partition_env.cu
• thrust/thrust/system/cuda/detail/partition.h

[github-actions]

🥳 CI Workflow Results

🟩 Finished in 1h 52m: Pass: 100%/340 | Total: 7d 14h | Max: 1h 51m | Hits: 62%/627542

See results here.

[Jacobfaib]

Suggested change

	cudaError_t THRUST_RUNTIME_FUNCTION dispatch_partition(
	[[nodiscard]] cudaError_t THRUST_RUNTIME_FUNCTION dispatch_partition(