[oleksandr-pavlyk]

Description

This PR modernizes several C2H/CUB test utilities around cuda::buffer and stream-aware test execution, then applies those utilities to the segmented-scan tests.

The main motivation is to replace direct test use of c2h::device_vector / c2h::host_vector in new code paths with buffer-based helpers backed by checked C2H memory resources. This keeps the allocation checks that C2H relies on for constrained devices, while making the segmented-scan tests use the same runtime buffer abstractions that CUB APIs increasingly exercise.

Main Changes

1. Adds checked C2H memory resources for cuda::buffer

   New helpers in <c2h/checked_memory_resource.cuh> provide buffer construction utilities such as c2h::make_device_buffer and c2h::make_host_buffer.

   c2h::make_device_buffer allocates device-accessible storage through a checked memory resource, preserving the pre-allocation device-memory checks previously available through c2h::device_vector and c2h::checked_cuda_allocator.

   c2h::make_host_buffer allocates host-accessible storage through a checked host resource. On integrated-memory systems, the host allocation path checks available device/system memory before using new/delete-style allocation, matching the intent of the previous c2h::host_vector checked allocation path.

2. Keeps the checked-memory-resource layer independent of Thrust containers

   <c2h/checked_memory_resource.cuh> is intended to provide the allocation resources and cuda::buffer factory helpers without depending on Thrust vector headers. This lets tests use checked buffer allocation without also pulling in thrust::device_vector / thrust::host_vector.

3. Adds buffer-oriented C2H generator helpers

   The C2H generator utilities are extended so tests can directly create and populate checked buffers instead of constructing a vector first and then extracting raw pointers.

   The generator-facing helpers live with the C2H generator support, primarily <c2h/generators.h> with CUDA implementation support in <c2h/detail/generators.cuh>. These generator utilities may still depend on Thrust because the existing random-data generation implementation uses Thrust/CUDA algorithm machinery. This is separate from the lower-level checked-memory-resource header, which should remain usable without Thrust vector dependencies.

   Stream-taking overloads are added for the buffer generator helpers. Those overloads enqueue their work on the caller-provided stream, so buffer population can participate in the same stream ordering as the test launch being exercised.

4. Adds stream-aware launch-wrapper overloads

   The CUB Catch2 launch wrappers now support overloads that take a CUDA stream as the first argument. This lets tests opt into stream-ordered setup and algorithm execution while still reusing the existing launch-mode matrix controlled by TEST_LAUNCH.

   The stream behavior differs by launch mode:

   • TEST_LAUNCH == 0: the stream is forwarded to the wrapped host-side CUB API calls, including both the temporary-storage sizing call and the real algorithm invocation. Synchronization is scoped to that stream.
   • TEST_LAUNCH == 1: the stream-aware overload cannot forward the host stream into the device-side CUB call because a host cudaStream_t is not valid inside the CDP-launched code. The helper synchronizes the provided stream with the host-side staging work, then uses the existing device-side launch path. The wrapped CUB call from device code receives the device-side/null stream behavior rather than the host stream.
   • TEST_LAUNCH == 2: the stream is used for CUDA graph capture and graph execution. The wrapped host-side CUB API is invoked during capture with the provided stream, and the instantiated graph is launched on that stream.

5. Migrates segmented-scan tests to the new helpers

   catch2_test_device_segmented_scan.cu is updated to use C2H checked buffers and the buffer-producing generator helpers for inputs, offsets, outputs, and host-side reference data.

   The tests that are parameterized with TEST_LAUNCH_PARAM use the stream-aware launch wrappers so the same segmented-scan coverage continues to run across direct host launch, CDP launch, and CUDA graph launch modes.

Reviewer Notes

The PR is broader than a local cleanup in catch2_test_device_segmented_scan.cu. It adds reusable C2H buffer allocation/generation infrastructure and stream-aware launch-wrapper entry points, then uses segmented scan as the first substantial migration target.

closes #10897

---

Previous edition of PR description

Segmented Scan tests to use CCCL Runtime buffers, instead of Thrust containers

Introduced c2h::make_device_buffer utility that uses added c2h::checked_device_memory_resource, a cuda::mr resource that backs cuda::device_buffer with C2H’s checked cudaMalloc path.

Uses

auto buffer = c2h::make_device_buffer<T>(stream, device, count, cuda::no_init);

auto offsets = c2h::make_device_buffer<int>(stream, device, {0, 4, 7, 9});

This replaces uses of c2h::device_vector in tests.

Similarly, c2h::host_vector is replaced with c2h::host_buffer that uses operator new/operator delete-based device-bound synchronous memory resource.

---

Checklist

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

[coderabbitai]

📝 Summary

Summary by CodeRabbit

• New Features

  • Added checked host/device memory resources with allocation validation, memory diagnostics, and convenient buffer creation helpers.
  • Added device, host, and paired buffer generators for random data and uniform segment offsets.
  • Added stream-aware generator operations and support for custom CUDA stream-like objects.
  • Added cross-platform environment-variable parsing utilities.

• Bug Fixes

  • Improved allocation failure handling, size-overflow checks, alignment validation, and integrated-device memory checks.

• Documentation

  • Expanded launch-wrapper documentation for stream-ordered setup, teardown, graph, host, and device-side execution.

Walkthrough

This change adds checked CUDA memory resources, stream-aware C2H generators, shared CUDA test utilities, and environment-aware launch wrappers. Segmented scan tests migrate from Thrust vectors to CUDA buffers with explicit streams, asynchronous transfers, host-backed expectations, and shared matchers.

Changes

CUDA buffer migration

Layer / File(s)	Summary
Memory and generator foundation c2h/include/c2h/detail/checked_memory.cuh, c2h/include/c2h/checked_memory_resource.cuh, c2h/include/c2h/buffer_generators.cuh, c2h/include/c2h/generator_types.h, c2h/include/c2h/vector_generators.h, c2h/generators*.cu	Checked allocation utilities, CUDA buffer generators, validated uniform-offset sizing, and stream-aware generator overloads are added.
Shared test support c2h/include/c2h/catch2_test_helper.h, c2h/include/c2h/detail/*, cub/test/catch2_test_cuda_utils.cuh, cub/test/catch2_test_device_segmented_scan_utils.cuh, cub/test/catch2_test_memory_resources.h	Shared device and stream setup, asynchronous copy helpers, host buffers, matchers, seed/NVTX helpers, bounded mismatch reporting, and allocation handling are added or reorganized.
Stream-aware launch support cub/test/catch2_test_launch_helper.h, cub/test/catch2_test_launch_wrapper.cu, cub/test/catch2_test_device_for_env.cu, cub/test/catch2_test_device_segmented_scan_env.cu, docs/cub/developer/test_overview.rst	Launch wrappers support explicit streams, stream-like objects, environment-based APIs, graph launches, device launches, and allocation verification.
Core segmented scan test migration cub/test/catch2_test_device_segmented_scan.cu, cub/test/catch2_test_device_segmented_scan_noncommutative.cu	Core segmented scan tests use CUDA buffers, owning streams, asynchronous copies, host references, and shared validation.
Segmented scan API and variant migrations cub/test/catch2_test_device_segmented_scan_api.cu, cub/test/catch2_test_device_segmented_scan_env_api.cu, cub/test/catch2_test_device_segmented_scan_multi_segment.cu	API, environment, multi-segment, empty-segment, invalid-offset, fancy-iterator, and custom-policy tests use explicit streams and CUDA buffers instead of Thrust vectors.

Assessment against linked issues

Objective	Addressed	Explanation
Replace thrust::device_vector in segmented scan tests with CUDA buffers [#10897]	✅

Out-of-scope changes

Code Change	Explanation
Checked memory resources and allocation utilities (c2h/include/c2h/detail/checked_memory.cuh:1-352, c2h/include/c2h/checked_memory_resource.cuh:1-148)	These provide general C2H memory infrastructure and are not part of the linked segmented scan test migration objective.
Stream-aware generator APIs (c2h/generators.cu:76-188, c2h/generators_uniform_offsets.cu:34-117, c2h/generators_vector.cu:56-70)	These change C2H generator behavior and public APIs beyond replacing segmented scan test containers.
Launch wrapper environment support (cub/test/catch2_test_launch_helper.h:503-540)	This adds general environment launch handling rather than only migrating segmented scan test storage.

Suggested reviewers: bernhardmgruber

Merge Risk: 🔵 Low · up to fd0a5

This test-infrastructure change moves segmented scans to checked CUDA buffers and explicit streams. Remaining concerns are bounded validation and edge-case robustness issues, so the change is mergeable with owner follow-up rather than showing a demonstrated production failure.

---

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

[github-actions]

⚠️ CCCL compile-time benchmark comparison: Public headers compile-time bench

Compile-time benchmark reporting did not produce a summary manifest.

Artifacts: reports and traces

[oleksandr-pavlyk]

@coderabbitai Regarding classification of this change as Medium Merge Risk, I think the stream-ordering risk is overstated here. The PR does use Runtime streams for buffer/copy operations, but the validation paths are not reading segmented scan outputs from an unordered stream.

The common launch helper synchronizes before returning for host, device-side, and graph launch modes. The env launch helper does the same. For direct default-stream validation, the shared segmented-scan test helper explicitly fences with cudaDeviceSynchronize() before routing cuda::copy_bytes through a real Runtime stream, then synchronizes that copy stream. Non-default paths copy and validate on the same stream.

For generated inputs, C2H uses c2h::device_policy, which is thrust::cuda::par, not par_nosync.

So I do not think “some scans and validation operations still use independently synchronized streams” is enough to justify moderate merge risk unless there is a specific call site where a scan launch can race a later validation copy. Could you point to that call site?

On annotations: C2H header APIs added for the checked device memory resource are annotated. Host-side test utilities are intentionally not annotated with _CCCL_HOST_API per CCCL coding policy, so those should not be counted as missing API annotations.

[coderabbitai]

Tip

For best results, initiate chat on the files or code changes.

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

files=(
  cub/test/catch2_test_device_segmented_scan_utils.cuh
  cub/test/catch2_test_device_segmented_scan.cu
  cub/test/catch2_test_device_segmented_scan_api.cu
  cub/test/catch2_test_device_segmented_scan_env.cu
  cub/test/catch2_test_device_segmented_scan_env_api.cu
  cub/test/catch2_test_device_segmented_scan_multi_segment.cu
  cub/test/catch2_test_device_segmented_scan_noncommutative.cu
  c2h/include/c2h/checked_allocator.cuh
  c2h/include/c2h/generators.h
)

for file in "${files[@]}"; do
  if [[ -f "$file" ]]; then
    echo "===== OUTLINE: $file ====="
    ast-grep outline "$file" --items all 2>/dev/null || true
    echo "===== ORDERING OPERATIONS: $file ====="
    rg -n -C 4 \
      'cuda(DeviceSynchronize|StreamSynchronize)|\.sync\(|copy_(to_host|to_device)|copy_bytes|DeviceSegmented(Scan|Reduce)|device_policy|thrust::cuda::par|_CCCL_HOST_API' \
      "$file" || true
  fi
done

Length of output: 63108

---

🏁 Script executed:

#!/bin/bash
set -euo pipefail

for name in catch2_test_launch_helper.h catch2_test_env_launch_helper.h; do
  echo "===== CANDIDATES: $name ====="
  fd -a -t f "^${name}$" . -x sh -c '
    echo "----- $1 -----"
    ast-grep outline "$1" --items all 2>/dev/null || true
    rg -n -C 8 "sync|Synchronize|stream|DECLARE_LAUNCH_WRAPPER|launch" "$1" || true
  ' sh {}
done

echo "===== STREAM SETUP AND WRAPPER INVOCATIONS ====="
sed -n '140,210p' cub/test/catch2_test_device_segmented_scan.cu
sed -n '120,210p' cub/test/catch2_test_device_segmented_scan_multi_segment.cu
sed -n '90,165p' cub/test/catch2_test_device_segmented_scan_noncommutative.cu

echo "===== C2H DEVICE POLICY ====="
rg -n -C 6 'device_policy|par_nosync|thrust::cuda::par' c2h/include cub/test -g '*.{h,hpp,cuh,cu}' || true

echo "===== ENV VALIDATION AFTER CUSTOM ENV LAUNCH ====="
sed -n '315,365p' cub/test/catch2_test_device_segmented_scan_env_api.cu

Length of output: 194760

---

suggestion: @oleksandr-pavlyk I found no scan-to-validation-copy race in the changed paths.

catch2_test_launch_helper.h synchronizes before return for host, device-side, and graph launch modes. catch2_test_env_launch_helper.h does the same.

In cub/test/catch2_test_device_segmented_scan.cu, scans use the default stream. copy_to_host calls cudaDeviceSynchronize() before it uses copy_stream. Non-default paths use the same stream for the scan and validation copy.

c2h::device_policy uses thrust::cuda::par, not par_nosync. The new public APIs in c2h/include/c2h/checked_allocator.cuh have _CCCL_HOST_API. The test utility functions are not public C2H APIs.

I withdraw the stream-ordering and annotation concerns. They do not justify Medium Merge Risk for this PR.

---

✏️ Learnings added

Learnt from: oleksandr-pavlyk
URL: https://github.com/NVIDIA/cccl/pull/10898

Timestamp: 2026-08-21T16:36:45.231Z
Learning: In CCCL C2H tests, `c2h::device_policy` is defined as `thrust::cuda::par` with a checked allocator. `c2h::nosync_device_policy` is the separate `thrust::cuda::par_nosync` policy.

---

Learnt from: oleksandr-pavlyk
URL: https://github.com/NVIDIA/cccl/pull/10898

Timestamp: 2026-08-21T16:36:45.231Z
Learning: In the CCCL segmented-scan test utilities, host-side functions in `cub/test/catch2_test_device_segmented_scan_utils.cuh` intentionally do not use `_CCCL_HOST_API`. Public C2H APIs, including `c2h::checked_device_memory_resource` and `c2h::make_device_buffer` in `c2h/include/c2h/checked_allocator.cuh`, require the annotation.

Note: Learnings are effective only in the context of similar code segments. To apply general coding guidelines, please configure review instructions. You can manage existing learnings in the app.

_{You are interacting with an AI system.}

[Jacobfaib]

Could we also extend some of this work to cudax? I ended up reinventing a lot of this under cudax/test/multi_gpu/, it would be great to replace it with these!

[coderabbitai]

Actionable comments posted: 1

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 2c30052e-6501-4bea-b4fe-e6ab7c193076

📥 Commits

Reviewing files that changed from the base of the PR and between c397427 and d7ecdd7.

📒 Files selected for processing (7)

• c2h/include/c2h/buffer_generators.cuh
• c2h/include/c2h/catch2_seed.h
• c2h/include/c2h/detail/checked_memory.cuh
• c2h/include/c2h/detail/generators.cuh
• cub/test/catch2_test_device_segmented_scan.cu
• cub/test/catch2_test_env_launch_helper.h
• cub/test/catch2_test_launch_helper.h

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

• cub/test/catch2_test_launch_helper.h
• c2h/include/c2h/buffer_generators.cuh
• cub/test/catch2_test_device_segmented_scan.cu
• c2h/include/c2h/detail/checked_memory.cuh

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

[oleksandr-pavlyk]

_{Commented in CodeRabbit Change Stack}

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

cub/test/catch2_test_launch_helper.h (1)

622-626: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

suggestion: this check is conditional here and at line 454, but line 497 makes expected_allocation_size mandatory for TEST_LAUNCH == 1. A test that omits the property therefore still passes for launcher ids 0 and 2, and only fails to compile for id 1. Make the property mandatory in all three paths so the allocation-size check always runs.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: aed4e980-0af5-4022-9cc1-2d2558cc512e

📥 Commits

Reviewing files that changed from the base of the PR and between df2ecac and fd0a5e8.

📒 Files selected for processing (4)

• cub/test/catch2_test_device_for_env.cu
• cub/test/catch2_test_device_segmented_scan_env.cu
• cub/test/catch2_test_launch_helper.h
• docs/cub/developer/test_overview.rst

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

[coderabbitai]

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

important: the prose and the example disagree. Line 386 tells the reader to pass a type supported by cuda::get_stream, but line 390 passes cuda::get_stream(stream), which is the already-extracted stream_ref. Show the stream-like object directly so readers do not copy the extraction call.

-    cub_reduce_sum(cuda::get_stream(stream), d_in, d_out, n);
+    cub_reduce_sum(stream, d_in, d_out, n);

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	pass any type supported by ``cuda::get_stream`` as the wrapper's first argument:
	.. code-block:: c++
	cub_reduce_sum(cuda::get_stream(stream), d_in, d_out, n);
	pass any type supported by ``cuda::get_stream`` as the wrapper's first argument:
	.. code-block:: c++
	cub_reduce_sum(stream, d_in, d_out, n);

[github-actions]

🥳 CI Workflow Results

🟩 Finished in 2h 45m: Pass: 100%/163 | Total: 5d 06h | Max: 2h 45m | Hits: 60%/241704

See results here.

[oleksandr-pavlyk]

This PR has been split into smaller PRs, most of which have been merged. Closing this PR now