Add multi-GPU exclusive_scan#9796
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
Jacobfaib
force-pushed
the
jacobf/2026-07-09/mgmn-scan
branch
from
e7a8bdf to
562b112
Compare
Jacobfaib
force-pushed
the
jacobf/2026-07-09/mgmn-scan
branch
from
562b112 to
1d78bb9
Compare
|
/ok to test |
📝 WalkthroughSummary by CodeRabbit
WalkthroughChangesAdds multi-GPU Multi-GPU algorithms
Suggested labels: Suggested reviewers: Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: a892325c-662d-4f16-b33a-6b4568ca2cd7
📒 Files selected for processing (14)
cudax/include/cuda/experimental/__multi_gpu/algorithm/reduce/reduce.hcudax/include/cuda/experimental/__multi_gpu/algorithm/scan/scan.hcudax/test/multi_gpu/algorithms/CMakeLists.txtcudax/test/multi_gpu/algorithms/reduce/range_basic.cucudax/test/multi_gpu/algorithms/reduce/range_defaults.cucudax/test/multi_gpu/algorithms/reduce/single_comm_basic.cucudax/test/multi_gpu/algorithms/reduce/single_comm_defaults.cucudax/test/multi_gpu/algorithms/scan/CMakeLists.txtcudax/test/multi_gpu/algorithms/scan/range_basic.cucudax/test/multi_gpu/algorithms/scan/range_defaults.cucudax/test/multi_gpu/algorithms/scan/single_comm_basic.cucudax/test/multi_gpu/algorithms/scan/single_comm_defaults.cucudax/test/multi_gpu/include/algorithm_common.hdocs/cudax/Doxyfile
| for (auto&& [__comm, __env, __part, __inputs, __out] : | ||
| ::cuda::std::ranges::views::zip(__comms, __envs, __partials, __range_of_inputs, __outputs)) | ||
| { | ||
| auto __prefix = ::cuda::experimental::__detail::__make_safe_uninitialized_buffer<_Tp>( | ||
| __part.stream(), __part.memory_resource(), /*__size=*/1, __env); | ||
|
|
||
| { | ||
| const auto __num_items = __comm.rank(); | ||
| // Root rank has no preceding partials and therefore starts directly with init. Other ranks | ||
| // include root rank partial, which already contains init, so their prefix reduction starts | ||
| // with ident. | ||
| const auto& __prefix_init = __comm.rank() == __ROOT_RANK ? __init : __ident; | ||
|
|
||
| // TODO(jfaibussowit): | ||
| // | ||
| // We could potentially just fold this into the exclusive scan kernel directly with fancy | ||
| // iterators. Basically, every rank needs to compute its local prefix sum on the init | ||
| // dereference. If that only happens once, we can probably reasonably fuse these | ||
| // kernels. If it happens multiple times, we need to cache the result. | ||
| __CUDAX_MULTI_GPU_DISPATCH( | ||
| __comm.logical_device(), | ||
| __num_items, | ||
| CUB_NS_QUALIFIER::DeviceReduce::Reduce, | ||
| (__part.begin(), __prefix.begin(), __num_items_fixed, __op, __prefix_init, __env)); | ||
| } | ||
| { | ||
| const auto __num_items = ::cuda::std::ranges::size(__inputs); | ||
|
|
||
| // Don't rely on CTAD for FutureValue, we want to be absolutely sure the iterator does not | ||
| // convert to something we dont expect | ||
| using __future_type = CUB_NS_QUALIFIER::FutureValue<typename __properties::__buffer_type::value_type, | ||
| typename __properties::__buffer_type::iterator>; | ||
|
|
||
| __CUDAX_MULTI_GPU_DISPATCH( | ||
| __comm.logical_device(), | ||
| __num_items, | ||
| CUB_NS_QUALIFIER::DeviceScan::ExclusiveScan, | ||
| (::cuda::std::ranges::begin(__inputs), | ||
| __out, | ||
| __op, | ||
| /*__init=*/__future_type{__prefix.begin()}, | ||
| __num_items_fixed, | ||
| __env)); | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Locate the macro definition and confirm whether it defines __num_items_fixed internally.
rg -n '__CUDAX_MULTI_GPU_DISPATCH' -C5 cudax/include
rg -n '__num_items_fixed' cudax/includeRepository: NVIDIA/cccl
Length of output: 7317
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== common.h macro =="
sed -n '50,90p' cudax/include/cuda/experimental/__multi_gpu/algorithm/common.h
echo
echo "== scan.h relevant block =="
sed -n '198,245p' cudax/include/cuda/experimental/__multi_gpu/algorithm/scan/scan.h
echo
echo "== reduce.h relevant block =="
sed -n '96,170p' cudax/include/cuda/experimental/__multi_gpu/algorithm/reduce/reduce.hRepository: NVIDIA/cccl
Length of output: 7188
important: __num_items_fixed is not declared here, and __CUDAX_MULTI_GPU_DISPATCH just forwards its arguments. These dispatch calls will fail to compile unless this name is added or replaced with __num_items.
| MULTI_GPU_TEST("exclusive_scan single-comm documentation example", c2h::type_list<int>) | ||
| { | ||
| auto comms = this->communicators(); | ||
|
|
||
| if (comms.size() < 2) | ||
| { | ||
| SKIP("The exclusive_scan documentation example requires at least two local GPUs"); | ||
| } | ||
|
|
||
| auto streams_owned = nccl_test_util::make_streams(); | ||
| auto streams = std::vector<cuda::stream_ref>{streams_owned.begin(), streams_owned.end()}; | ||
|
|
||
| // Every communicator rank must invoke the collective concurrently. | ||
| run_threaded(comms.size(), [&](cuda::std::size_t i) { | ||
| auto& communicator = comms[i]; | ||
| auto environment = streams[i]; | ||
| const auto device = communicator.logical_device().underlying_device(); | ||
|
|
||
| //! [exclusive_scan_single_range] | ||
| constexpr cuda::std::array input_values{1, 2}; | ||
|
|
||
| auto input = cuda::make_device_buffer<int>(environment, device, input_values); | ||
| auto output = cuda::make_device_buffer<int>(environment, device, input_values.size(), cuda::no_init); | ||
|
|
||
| cudax::exclusive_scan(communicator, environment, input, output.begin(), /*__init=*/0); | ||
|
|
||
| // Every rank contributes {1, 2}, so rank r starts with a prefix of 3 * r. | ||
| const auto rank = communicator.rank(); | ||
| const auto expected = | ||
| cuda::make_buffer<int>(output.stream(), cuda::mr::legacy_pinned_memory_resource{}, {3 * rank, 3 * rank + 1}); | ||
| REQUIRE_THAT(output, Equals(expected)); | ||
| //! [exclusive_scan_single_range] | ||
| }); | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the test file and the threaded helper used by this example.
git ls-files cudax/test/multi_gpu/algorithms/scan/single_comm_basic.cu
ast-grep outline cudax/test/multi_gpu/algorithms/scan/single_comm_basic.cu --view expanded || true
rg -n "run_threaded|REQUIRE_THAT|exclusive_scan" cudax/test/multi_gpu/algorithms/scan/single_comm_basic.cu cudax/test/multi_gpu/algorithms/scan -S
# Find the helper implementation and any thread-safety notes nearby.
rg -n "std::async|launch::async|future|try/catch|REQUIRE.*thread|user-spawned thread|terminate the process" cudax/test -SRepository: NVIDIA/cccl
Length of output: 18088
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Read the relevant helper implementation in small slices if found.
helper="$(rg -n --no-heading "run_threaded" cudax/test -S | head -n 20 | cut -d: -f1 | sort -u | head -n 1 || true)"
if [ -n "${helper}" ]; then
echo "HELPER_FILE=${helper}"
wc -l "$helper"
sed -n '1,220p' "$helper"
fiRepository: NVIDIA/cccl
Length of output: 7995
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Show the exact lines around the test and helper implementation.
sed -n '120,220p' cudax/test/multi_gpu/algorithms/scan/single_comm_basic.cu
echo '---'
rg -n "run_threaded" cudax/test/multi_gpu/algorithms/scan -SRepository: NVIDIA/cccl
Length of output: 4657
🌐 Web query:
Catch2 assertion inside user-spawned thread terminate the process REQUIRE thread safety documentation
💡 Result:
By default, Catch2 assertion macros (including REQUIRE) are not thread-safe, and using them from multiple threads concurrently can lead to undefined behavior, including process termination or crashes [1][2][3]. However, Catch2 does support thread-safe assertions as an optional, opt-in feature [3]. To enable this, you must compile your project with the macro CATCH_CONFIG_THREAD_SAFE_ASSERTIONS defined [3]. This support was introduced in Catch2 version 3.9.0 and became a non-experimental, officially supported feature in version 3.12.0 [3]. When enabled, Catch2 introduces internal locking (using a std::recursive_mutex) for assertion and logging macros, allowing multiple user-spawned threads to use these macros simultaneously [4][3]. Note the following implications: 1. Performance Impact: Enabling this feature imposes a performance cost on all assertion macros, even in single-threaded parts of your test suite, due to the global locking mechanism [4][3]. 2. Scope: Thread safety is generally limited to assertion macros (e.g., REQUIRE, CHECK) and logging/message macros (e.g., INFO, WARN) [1][3]. Other macros, such as those related to benchmarks, sections, generators, or top-level test cases, are not thread-safe [1]. 3. Alternative Workaround: If you cannot enable CATCH_CONFIG_THREAD_SAFE_ASSERTIONS or are using an older version of Catch2, a common workaround is to synchronize access to assertion macros manually using a std::mutex and a guard block [5]. Example of manual synchronization: std::mutex g_catch2_mutex; #define REQUIRE_TS(expr) { std::lock_guard guard(g_catch2_mutex); REQUIRE(expr); } Always verify your specific Catch2 version's documentation, as features and configuration macros may evolve [3].
Citations:
- 1: https://github.com/catchorg/Catch2/blob/devel/docs/thread-safety.md
- 2: https://ryankert01.github.io/Catch2/limitations.html
- 3: https://github.com/catchorg/Catch2/blob/devel/docs/configuration.md
- 4: Make assertions thread-safe catchorg/Catch2#2948
- 5: Regression due to Large output redirect refactor catchorg/Catch2#2935
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Check whether this repository enables Catch2 thread-safe assertions anywhere.
rg -n "CATCH_CONFIG_THREAD_SAFE_ASSERTIONS|thread-safe assertions|thread_safe_assertions" -S .Repository: NVIDIA/cccl
Length of output: 149
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the helper and surrounding comments directly.
sed -n '1,120p' cudax/test/multi_gpu/include/algorithm_common.h
echo '---'
sed -n '150,205p' cudax/test/multi_gpu/algorithms/scan/single_comm_basic.cuRepository: NVIDIA/cccl
Length of output: 4577
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Find other run_threaded call sites and where assertions are placed relative to them.
rg -n "run_threaded\\(" cudax/test/multi_gpu -S
echo '---'
rg -n "REQUIRE_THAT\\(|REQUIRE\\(" cudax/test/multi_gpu/algorithms -S | sed -n '1,120p'Repository: NVIDIA/cccl
Length of output: 5745
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Look for any repo-local Catch2 configuration or version hints.
rg -n "catch2|CATCH2|CATCH_CONFIG" cudax/test . -g '!**/build/**' -SRepository: NVIDIA/cccl
Length of output: 50368
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Search the whole repository for Catch2 thread-safe assertion configuration.
rg -n "CATCH_CONFIG_THREAD_SAFE_ASSERTIONS" -S .Repository: NVIDIA/cccl
Length of output: 149
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the full helper comments and code for exception propagation.
sed -n '1,90p' cudax/test/multi_gpu/include/algorithm_common.hRepository: NVIDIA/cccl
Length of output: 2438
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Check whether Catch2 thread-safe assertions are enabled anywhere in the repo's test setup.
rg -n "CATCH_CONFIG_THREAD_SAFE_ASSERTIONS|thread-safe assertions|catch2.*thread" -S .github ci cudax thrust cub libcudacxx python cRepository: NVIDIA/cccl
Length of output: 1444
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Look for compile definitions that may affect Catch2 thread safety in test targets.
rg -n "add_compile_definitions|target_compile_definitions|CATCH_CONFIG_" -S cudax ci .github CMakeLists.txt CMakePresets.jsonRepository: NVIDIA/cccl
Length of output: 348
critical: move the REQUIRE_THAT out of the run_threaded lambda, or enable Catch2 thread-safe assertions for this target. Concurrent Catch2 assertions from worker threads are undefined behavior here.
This comment has been minimized.
This comment has been minimized.
😬 CI Workflow Results🟥 Finished in 3h 19m: Pass: 98%/57 | Total: 1d 15h | Max: 1h 27m | Hits: 5%/351263See results here. |
1 participant
Description
closes
Checklist