[CCCL] 瘦身 + 补全: 移除 cudax/python/libcudacxx-tests 冗余文件, 新增 c2h 测试助手 + cmake 构建系统 + 8 个 CUDA thrust examples
变更摘要:
- 删除: cudax/ (783 files, 7.2M) — 实验性组件，竞赛不需要
- 删除: python/ (226 files, 2.0M) — Python 绑定，竞赛不需要
- 删除: libcudacxx/{test,benchmarks,codegen,cmake,share} (4432 files, 31M)
保留: libcudacxx/include/ (1463 headers, cuda::std 编译依赖)
- 新增: c2h/ (27 files) — CUB Catch2 测试辅助头文件，编译 243 个测试必需
- 新增: cmake/ (29 files) — CCCL 原生 CMake 构建系统
- 新增: thrust/examples/cuda/ (7 files) + cpp_integration/ (1 file)
async_reduce, custom_temporary_allocation, explicit_cuda_stream,
global_device_vector, range_view, unwrap_pointer, wrap_pointer, device
结果: cccl_upstream 从 74M→35M (瘦身 53%), 核心内容 100% 保留:
27/27 tuning headers, 78 benchmarks, 243 tests,
60 thrust examples, 18 CUB examples, 全部编译头文件
This commit is contained in:
muh-bot
@@ -1,73 +0,0 @@
|
||||
#===----------------------------------------------------------------------===##
|
||||
#
|
||||
# Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
# under the Apache License v2.0 with LLVM Exceptions.
|
||||
# See https://llvm.org/LICENSE.txt for license information.
|
||||
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
#
|
||||
#===----------------------------------------------------------------------===##
|
||||
|
||||
cccl_get_nccl(QUIET)
|
||||
|
||||
option(
|
||||
cudax_ENABLE_NCCL
|
||||
"Enable cudax tests that link against NCCL."
|
||||
"${NCCL_FOUND}"
|
||||
)
|
||||
|
||||
# TODO(jfaibussowit):
|
||||
#
|
||||
# Remove this once NCCL supports windows and we have added NCCL to the NVHPC
|
||||
# devcontainers.
|
||||
macro(hack_cudax_handle_missing_nccl)
|
||||
if (NOT cudax_ENABLE_NCCL)
|
||||
message(
|
||||
AUTHOR_WARNING
|
||||
"Function should only be called if we have requested NCCL, this is a developer error"
|
||||
)
|
||||
endif()
|
||||
|
||||
option(
|
||||
HACK_cudax_ALLOW_MISSING_NCCL
|
||||
"Hack to work around CI issues since NCCL does not (yet) officially support Windows"
|
||||
TRUE
|
||||
)
|
||||
|
||||
if (NCCL_FOUND)
|
||||
# If we have somehow already found NCCL (e.g. they have shipped support but we have
|
||||
# forgotten to remove this function), then we have nothing to do. We still call
|
||||
# cccl_get_nccl() though because we want the found version to appear in the logs. The
|
||||
# previous call was with QUIET which won't do that.
|
||||
cccl_get_nccl(REQUIRED)
|
||||
elseif (HACK_cudax_ALLOW_MISSING_NCCL)
|
||||
set(cudax_ENABLE_NCCL FALSE)
|
||||
set(cudax_ENABLE_NCCL FALSE CACHE BOOL "" FORCE)
|
||||
else()
|
||||
# This will cause a hard error
|
||||
cccl_get_nccl(REQUIRED)
|
||||
endif()
|
||||
endmacro()
|
||||
|
||||
if (cudax_ENABLE_NCCL)
|
||||
if (WIN32 OR ("${CMAKE_CXX_COMPILER_ID}" STREQUAL "NVHPC"))
|
||||
hack_cudax_handle_missing_nccl()
|
||||
else()
|
||||
cccl_get_nccl(REQUIRED)
|
||||
endif()
|
||||
endif()
|
||||
|
||||
function(cudax_add_multi_gpu_test sub_prefix target_name_var source)
|
||||
cmake_path(GET source STEM filename)
|
||||
cudax_add_catch2_test("${target_name_var}" "multi_gpu.${sub_prefix}.${filename}" "${source}" ${ARGN})
|
||||
target_include_directories(
|
||||
"${${target_name_var}}"
|
||||
PRIVATE "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/include"
|
||||
)
|
||||
set("${target_name_var}" "${${target_name_var}}" PARENT_SCOPE)
|
||||
endfunction()
|
||||
|
||||
add_subdirectory(algorithms)
|
||||
add_subdirectory(communicators)
|
||||
add_subdirectory(concepts)
|
||||
add_subdirectory(nccl)
|
||||
@@ -1,13 +0,0 @@
|
||||
#===----------------------------------------------------------------------===##
|
||||
#
|
||||
# Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
# under the Apache License v2.0 with LLVM Exceptions.
|
||||
# See https://llvm.org/LICENSE.txt for license information.
|
||||
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
#
|
||||
#===----------------------------------------------------------------------===##
|
||||
|
||||
add_subdirectory(reduce)
|
||||
add_subdirectory(exclusive_scan)
|
||||
add_subdirectory(inclusive_scan)
|
||||
@@ -1,20 +0,0 @@
|
||||
#===----------------------------------------------------------------------===##
|
||||
#
|
||||
# Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
# under the Apache License v2.0 with LLVM Exceptions.
|
||||
# See https://llvm.org/LICENSE.txt for license information.
|
||||
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
#
|
||||
#===----------------------------------------------------------------------===##
|
||||
|
||||
if (NOT cudax_ENABLE_NCCL)
|
||||
return()
|
||||
endif()
|
||||
|
||||
file(GLOB test_srcs LIST_DIRECTORIES FALSE CONFIGURE_DEPENDS *.cu *.cpp)
|
||||
|
||||
foreach (src IN LISTS test_srcs)
|
||||
cudax_add_multi_gpu_test("algorithms.exclusive_scan" test_target "${src}")
|
||||
target_link_libraries(${test_target} PRIVATE NCCL::nccl)
|
||||
endforeach()
|
||||
@@ -1,388 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/buffer>
|
||||
#include <cuda/functional>
|
||||
#include <cuda/memory_resource>
|
||||
#include <cuda/std/array>
|
||||
#include <cuda/std/cstddef>
|
||||
#include <cuda/std/cstdint>
|
||||
#include <cuda/std/execution>
|
||||
#include <cuda/std/functional>
|
||||
#include <cuda/std/limits>
|
||||
#include <cuda/std/span>
|
||||
#include <cuda/std/type_traits>
|
||||
|
||||
#include <cuda/experimental/__multi_gpu/algorithm/scan/scan.h>
|
||||
|
||||
#include <numeric>
|
||||
#include <vector>
|
||||
|
||||
#include <algorithm_common.h>
|
||||
#include <nccl_test_common.h>
|
||||
#include <testing.cuh>
|
||||
|
||||
namespace
|
||||
{
|
||||
struct custom_plus
|
||||
{
|
||||
template <class T>
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr T operator()(const T& lhs, const T& rhs) const
|
||||
{
|
||||
return lhs + rhs;
|
||||
}
|
||||
};
|
||||
|
||||
using custom_value = c2h::custom_type_t<c2h::accumulateable_t, c2h::less_comparable_t, c2h::equal_comparable_t>;
|
||||
using value_types = c2h::type_list<cuda::std::int32_t, float, custom_value>;
|
||||
using operators = c2h::type_list<::cuda::std::plus<>, ::cuda::maximum<>, custom_plus>;
|
||||
|
||||
static_assert(cudax::nccl_transportable<custom_value>);
|
||||
|
||||
template <typename T>
|
||||
T make_value(int i)
|
||||
{
|
||||
return static_cast<T>(i);
|
||||
}
|
||||
|
||||
template <>
|
||||
custom_value make_value<>(int i)
|
||||
{
|
||||
custom_value ret{};
|
||||
|
||||
ret.key = static_cast<std::size_t>(i);
|
||||
ret.val = static_cast<std::size_t>(i);
|
||||
return ret;
|
||||
};
|
||||
|
||||
template <class T, class Op>
|
||||
[[nodiscard]] T get_identity()
|
||||
{
|
||||
if constexpr (cuda::std::is_same_v<Op, cuda::std::plus<>> || cuda::std::is_same_v<Op, custom_plus>)
|
||||
{
|
||||
return make_value<T>(0);
|
||||
}
|
||||
else if constexpr (cuda::std::is_same_v<Op, cuda::maximum<>>)
|
||||
{
|
||||
return cuda::std::numeric_limits<T>::lowest();
|
||||
}
|
||||
else
|
||||
{
|
||||
static_assert(cuda::std::__always_false_v<T, Op>, "Add handling");
|
||||
}
|
||||
}
|
||||
|
||||
template <class T, class Op>
|
||||
[[nodiscard]] std::vector<T>
|
||||
expected_for_rank(int rank, const std::vector<std::vector<T>>& inputs_by_rank, const T& init, Op op)
|
||||
{
|
||||
std::vector<T> reference;
|
||||
|
||||
for (const auto& values : inputs_by_rank)
|
||||
{
|
||||
reference.insert(reference.end(), values.begin(), values.end());
|
||||
}
|
||||
|
||||
std::vector<T> scan(reference.size());
|
||||
std::exclusive_scan(reference.begin(), reference.end(), scan.begin(), init, op);
|
||||
|
||||
cuda::std::size_t offset = 0;
|
||||
for (int r = 0; r < rank; ++r)
|
||||
{
|
||||
offset += inputs_by_rank[static_cast<cuda::std::size_t>(r)].size();
|
||||
}
|
||||
|
||||
const auto count = inputs_by_rank[static_cast<cuda::std::size_t>(rank)].size();
|
||||
return {scan.begin() + offset, scan.begin() + offset + count};
|
||||
}
|
||||
|
||||
// Run the full scan, wait for it to finish, and check that `exclusive_scan` left its argument
|
||||
// ranges untouched. This boilerplate is identical for every test regardless of how the inputs are
|
||||
// shaped.
|
||||
template <class Env, class T, class Op>
|
||||
void do_exclusive_scan(
|
||||
cuda::std::span<cudax::nccl_communicator_ref> comms,
|
||||
const std::vector<Env>& envs,
|
||||
std::vector<cuda::device_buffer<T>>& in,
|
||||
std::vector<typename cuda::device_buffer<T>::iterator>& outputs,
|
||||
const T& init,
|
||||
const T& ident,
|
||||
Op op)
|
||||
{
|
||||
const auto envs_size = envs.size();
|
||||
const auto in_copy = in;
|
||||
const auto outputs_copy = outputs;
|
||||
|
||||
INFO("init = " << init);
|
||||
INFO("ident = " << ident);
|
||||
|
||||
cudax::exclusive_scan(cudax::distributed, comms, envs, in, outputs, init, op, ident);
|
||||
|
||||
// cuda::std::execution::env has no operator==, so we can only compare the sizes.
|
||||
REQUIRE(envs.size() == envs_size);
|
||||
// Scan call should not modify the inputs in any ways
|
||||
REQUIRE(in.size() == in_copy.size());
|
||||
for (cuda::std::size_t i = 0; i < in.size(); ++i)
|
||||
{
|
||||
INFO("device = " << i);
|
||||
REQUIRE_THAT(in[i], Equals(in_copy[i]));
|
||||
}
|
||||
REQUIRE_THAT(outputs, Catch::Matchers::Equals(outputs_copy));
|
||||
}
|
||||
} // namespace
|
||||
|
||||
MULTI_GPU_TEST("exclusive_scan 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();
|
||||
// Convert to stream_ref directly, cuda::stream on their own cant be passed directly to CUB
|
||||
auto streams = std::vector<cuda::stream_ref>{streams_owned.begin(), streams_owned.end()};
|
||||
|
||||
//! [exclusive_scan]
|
||||
constexpr cuda::std::array input_values{1, 2};
|
||||
std::vector<cuda::device_buffer<int>> inputs;
|
||||
std::vector<cuda::device_buffer<int>> outputs;
|
||||
|
||||
for (cuda::std::size_t i = 0; i < comms.size(); ++i)
|
||||
{
|
||||
const auto device = comms[i].logical_device().underlying_device();
|
||||
|
||||
inputs.emplace_back(cuda::make_device_buffer<int>(streams[i], device, input_values));
|
||||
outputs.emplace_back(cuda::make_device_buffer<int>(streams[i], device, input_values.size(), cuda::no_init));
|
||||
}
|
||||
|
||||
std::vector<typename cuda::device_buffer<int>::iterator> output_iterators = make_output_iterators(outputs);
|
||||
|
||||
cudax::exclusive_scan(
|
||||
cudax::distributed,
|
||||
comms,
|
||||
// Passing streams as the environment directly
|
||||
streams,
|
||||
inputs,
|
||||
output_iterators,
|
||||
/*__init=*/0);
|
||||
|
||||
constexpr cuda::std::array expected_rank_0{0, 1};
|
||||
constexpr cuda::std::array expected_rank_1{3, 4};
|
||||
const auto expected_0 =
|
||||
cuda::make_buffer<int>(outputs[0].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_rank_0);
|
||||
const auto expected_1 =
|
||||
cuda::make_buffer<int>(outputs[1].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_rank_1);
|
||||
REQUIRE_THAT(outputs[0], Equals(expected_0));
|
||||
REQUIRE_THAT(outputs[1], Equals(expected_1));
|
||||
//! [exclusive_scan]
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("exclusive_scan, one element per rank", value_types, operators)
|
||||
{
|
||||
using T = c2h::get<0, TestType>;
|
||||
using Op = c2h::get<1, TestType>;
|
||||
|
||||
// Seed each scan with a few hardcoded initializers. The init participates in the fold the same
|
||||
// way on host and device, so any value works for every operator under test.
|
||||
const T init = make_value<T>(GENERATE(0, 1, -1, 5));
|
||||
const auto ident = get_identity<T, Op>();
|
||||
|
||||
auto comms = this->communicators();
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
// Global rank `comms[i].rank()` contributes the single value `rank`. Each local rank also gets a
|
||||
// one-element output buffer and an environment carrying its stream, so the scan is stream-ordered
|
||||
// on the correct device. `reference` mirrors the contributions of every global rank so we can
|
||||
// compute the host-side scan exactly like `exclusive_scan` does on the device.
|
||||
std::vector<cuda::device_buffer<T>> in;
|
||||
std::vector<cuda::device_buffer<T>> out;
|
||||
std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
|
||||
std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
|
||||
|
||||
in.reserve(comms.size());
|
||||
out.reserve(comms.size());
|
||||
envs.reserve(comms.size());
|
||||
for (int r = 0; r < comms.front().size(); ++r)
|
||||
{
|
||||
inputs_by_rank[static_cast<cuda::std::size_t>(r)] = std::vector<T>(1, make_value<T>(r));
|
||||
}
|
||||
for (cuda::std::size_t i = 0; i < comms.size(); ++i)
|
||||
{
|
||||
const auto& values = inputs_by_rank[static_cast<cuda::std::size_t>(comms[i].rank())];
|
||||
in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
|
||||
out.emplace_back(cuda::make_device_buffer<T>(
|
||||
streams[i], comms[i].logical_device().underlying_device(), values.size(), cuda::no_init));
|
||||
envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
|
||||
}
|
||||
|
||||
auto outputs = make_output_iterators(out);
|
||||
|
||||
do_exclusive_scan(comms, envs, in, outputs, init, ident, Op{});
|
||||
|
||||
for (cuda::std::size_t i = 0; i < out.size(); ++i)
|
||||
{
|
||||
const auto expected_values = expected_for_rank<T>(comms[i].rank(), inputs_by_rank, init, Op{});
|
||||
const auto exp = cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_values);
|
||||
|
||||
REQUIRE_THAT(out[i], Equals(exp));
|
||||
}
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("exclusive_scan, multiple elements per rank", value_types, operators)
|
||||
{
|
||||
using T = c2h::get<0, TestType>;
|
||||
using Op = c2h::get<1, TestType>;
|
||||
|
||||
// Seed each scan with a few hardcoded initializers. The init participates in the fold the same
|
||||
// way on host and device, so any value works for every operator under test.
|
||||
const T init = make_value<T>(GENERATE(0, 1, -1, 5));
|
||||
const auto ident = get_identity<T, Op>();
|
||||
|
||||
auto comms = this->communicators();
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
// Global rank `comms[i].rank()` contributes ten copies of `rank`. `exclusive_scan` first
|
||||
// computes a local prefix for each rank seeded by the prefix of all previous ranks. Each local
|
||||
// rank also gets an output buffer and an environment carrying its stream. `reference` mirrors
|
||||
// every global rank's ten contributions for the host-side scan.
|
||||
std::vector<cuda::device_buffer<T>> in;
|
||||
std::vector<cuda::device_buffer<T>> out;
|
||||
std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
|
||||
std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
|
||||
|
||||
in.reserve(comms.size());
|
||||
out.reserve(comms.size());
|
||||
envs.reserve(comms.size());
|
||||
|
||||
constexpr auto values_per_rank = 10;
|
||||
for (int r = 0; r < comms.front().size(); ++r)
|
||||
{
|
||||
inputs_by_rank[static_cast<cuda::std::size_t>(r)] = std::vector<T>(values_per_rank, make_value<T>(r));
|
||||
}
|
||||
for (cuda::std::size_t i = 0; i < comms.size(); ++i)
|
||||
{
|
||||
const auto& values = inputs_by_rank[static_cast<cuda::std::size_t>(comms[i].rank())];
|
||||
in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
|
||||
out.emplace_back(cuda::make_device_buffer<T>(
|
||||
streams[i], comms[i].logical_device().underlying_device(), values.size(), cuda::no_init));
|
||||
envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
|
||||
}
|
||||
|
||||
auto outputs = make_output_iterators(out);
|
||||
|
||||
do_exclusive_scan(comms, envs, in, outputs, init, ident, Op{});
|
||||
|
||||
for (cuda::std::size_t i = 0; i < out.size(); ++i)
|
||||
{
|
||||
const auto expected_values = expected_for_rank<T>(comms[i].rank(), inputs_by_rank, init, Op{});
|
||||
const auto exp = cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_values);
|
||||
|
||||
REQUIRE_THAT(out[i], Equals(exp));
|
||||
}
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("exclusive_scan, some ranks empty", value_types, operators)
|
||||
{
|
||||
using T = c2h::get<0, TestType>;
|
||||
using Op = c2h::get<1, TestType>;
|
||||
|
||||
const T init = make_value<T>(GENERATE(0, 1, -1, 5));
|
||||
const auto ident = get_identity<T, Op>();
|
||||
|
||||
auto comms = this->communicators();
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
// Even global ranks contribute ten copies of `rank`; odd global ranks contribute an empty input
|
||||
// range. Rank 0 is always non-empty. `exclusive_scan` must treat an empty rank as contributing
|
||||
// nothing, exactly like `std::exclusive_scan` over the surviving elements. `reference` mirrors
|
||||
// that for the host-side scan.
|
||||
std::vector<cuda::device_buffer<T>> in;
|
||||
std::vector<cuda::device_buffer<T>> out;
|
||||
std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
|
||||
std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
|
||||
|
||||
in.reserve(comms.size());
|
||||
out.reserve(comms.size());
|
||||
envs.reserve(comms.size());
|
||||
|
||||
constexpr auto values_per_rank = 10;
|
||||
for (int r = 0; r < comms.front().size(); ++r)
|
||||
{
|
||||
if (r % 2 == 0)
|
||||
{
|
||||
inputs_by_rank[static_cast<cuda::std::size_t>(r)] = std::vector<T>(values_per_rank, make_value<T>(r));
|
||||
}
|
||||
}
|
||||
for (cuda::std::size_t i = 0; i < comms.size(); ++i)
|
||||
{
|
||||
const auto& values = inputs_by_rank[static_cast<cuda::std::size_t>(comms[i].rank())];
|
||||
in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
|
||||
out.emplace_back(cuda::make_device_buffer<T>(
|
||||
streams[i], comms[i].logical_device().underlying_device(), values.size(), cuda::no_init));
|
||||
envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
|
||||
}
|
||||
|
||||
auto outputs = make_output_iterators(out);
|
||||
|
||||
do_exclusive_scan(comms, envs, in, outputs, init, ident, Op{});
|
||||
|
||||
for (cuda::std::size_t i = 0; i < out.size(); ++i)
|
||||
{
|
||||
const auto expected_values = expected_for_rank<T>(comms[i].rank(), inputs_by_rank, init, Op{});
|
||||
const auto exp = cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_values);
|
||||
|
||||
REQUIRE_THAT(out[i], Equals(exp));
|
||||
}
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("exclusive_scan, all ranks empty", value_types, operators)
|
||||
{
|
||||
using T = c2h::get<0, TestType>;
|
||||
using Op = c2h::get<1, TestType>;
|
||||
|
||||
const auto init = make_value<T>(GENERATE(0, 1, -1, 5));
|
||||
const auto ident = get_identity<T, Op>();
|
||||
|
||||
auto comms = this->communicators();
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
// No rank contributes any element. Scanning nothing produces no output values, exactly like
|
||||
// `std::exclusive_scan` over an empty range.
|
||||
std::vector<cuda::device_buffer<T>> in;
|
||||
std::vector<cuda::device_buffer<T>> out;
|
||||
std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
|
||||
std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
|
||||
|
||||
in.reserve(comms.size());
|
||||
out.reserve(comms.size());
|
||||
envs.reserve(comms.size());
|
||||
for (cuda::std::size_t i = 0; i < comms.size(); ++i)
|
||||
{
|
||||
const auto& values = inputs_by_rank[static_cast<cuda::std::size_t>(comms[i].rank())];
|
||||
in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
|
||||
out.emplace_back(cuda::make_device_buffer<T>(
|
||||
streams[i], comms[i].logical_device().underlying_device(), values.size(), cuda::no_init));
|
||||
envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
|
||||
}
|
||||
|
||||
auto outputs = make_output_iterators(out);
|
||||
|
||||
do_exclusive_scan(comms, envs, in, outputs, init, ident, Op{});
|
||||
|
||||
for (cuda::std::size_t i = 0; i < out.size(); ++i)
|
||||
{
|
||||
const auto expected_values = expected_for_rank<T>(comms[i].rank(), inputs_by_rank, init, Op{});
|
||||
const auto exp = cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_values);
|
||||
|
||||
REQUIRE_THAT(out[i], Equals(exp));
|
||||
}
|
||||
}
|
||||
@@ -1,112 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/buffer>
|
||||
#include <cuda/functional>
|
||||
#include <cuda/memory_resource>
|
||||
#include <cuda/std/cstddef>
|
||||
#include <cuda/std/cstdint>
|
||||
#include <cuda/std/execution>
|
||||
#include <cuda/std/functional>
|
||||
#include <cuda/std/iterator>
|
||||
|
||||
#include <cuda/experimental/__multi_gpu/algorithm/scan/scan.h>
|
||||
|
||||
#include <numeric>
|
||||
#include <vector>
|
||||
|
||||
#include <algorithm_common.h>
|
||||
#include <nccl_test_common.h>
|
||||
#include <testing.cuh>
|
||||
|
||||
MULTI_GPU_TEST("exclusive_scan, range overloads default values", )
|
||||
{
|
||||
using T = cuda::std::int32_t;
|
||||
using Op = ::cuda::std::plus<>;
|
||||
|
||||
constexpr auto init = T{};
|
||||
constexpr T ident = cuda::identity_element<Op, T>();
|
||||
constexpr auto op = Op{};
|
||||
|
||||
const auto comms = this->communicators();
|
||||
const auto streams = nccl_test_util::make_streams();
|
||||
|
||||
std::vector<cuda::device_buffer<T>> in;
|
||||
std::vector<cuda::device_buffer<T>> out;
|
||||
std::vector<cuda::stream_ref> envs;
|
||||
|
||||
in.reserve(comms.size());
|
||||
out.reserve(comms.size());
|
||||
envs.reserve(comms.size());
|
||||
|
||||
constexpr auto values_per_rank = 10;
|
||||
|
||||
for (cuda::std::size_t i = 0; i < comms.size(); ++i)
|
||||
{
|
||||
std::vector<T> values(values_per_rank);
|
||||
std::iota(values.begin(), values.end(), static_cast<T>(comms[i].rank() * values_per_rank + 1));
|
||||
|
||||
in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
|
||||
out.emplace_back(cuda::make_device_buffer<T>(
|
||||
streams[i], comms[i].logical_device().underlying_device(), cuda::std::size(values), cuda::no_init));
|
||||
envs.emplace_back(streams[i]);
|
||||
}
|
||||
|
||||
auto outputs = make_output_iterators(out);
|
||||
|
||||
const auto expected_values = [&] {
|
||||
std::vector<T> reference(static_cast<cuda::std::size_t>(comms.front().size()) * values_per_rank);
|
||||
|
||||
std::iota(reference.begin(), reference.end(), T{1});
|
||||
|
||||
std::vector<T> expected_values(reference.size());
|
||||
|
||||
std::exclusive_scan(reference.begin(), reference.end(), expected_values.begin(), init, op);
|
||||
|
||||
return expected_values;
|
||||
}();
|
||||
|
||||
const auto check_outputs = [&] {
|
||||
const auto exp_span = cuda::std::span{expected_values};
|
||||
|
||||
for (cuda::std::size_t i = 0; i < out.size(); ++i)
|
||||
{
|
||||
const auto expected_for_rank = exp_span.subspan(comms[i].rank() * values_per_rank, values_per_rank);
|
||||
const auto expected =
|
||||
cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_for_rank);
|
||||
|
||||
REQUIRE_THAT(out[i], Equals(expected));
|
||||
}
|
||||
};
|
||||
|
||||
SECTION("Default init, op, ident (all)")
|
||||
{
|
||||
cudax::exclusive_scan(cudax::distributed, comms, envs, in, outputs);
|
||||
check_outputs();
|
||||
}
|
||||
|
||||
SECTION("Default op, ident")
|
||||
{
|
||||
cudax::exclusive_scan(cudax::distributed, comms, envs, in, outputs, init);
|
||||
check_outputs();
|
||||
}
|
||||
|
||||
SECTION("Default ident")
|
||||
{
|
||||
cudax::exclusive_scan(cudax::distributed, comms, envs, in, outputs, init, op);
|
||||
check_outputs();
|
||||
}
|
||||
|
||||
SECTION("Default none")
|
||||
{
|
||||
cudax::exclusive_scan(cudax::distributed, comms, envs, in, outputs, init, op, ident);
|
||||
check_outputs();
|
||||
}
|
||||
}
|
||||
@@ -1,289 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/buffer>
|
||||
#include <cuda/functional>
|
||||
#include <cuda/memory_resource>
|
||||
#include <cuda/std/array>
|
||||
#include <cuda/std/cstddef>
|
||||
#include <cuda/std/cstdint>
|
||||
#include <cuda/std/execution>
|
||||
#include <cuda/std/functional>
|
||||
#include <cuda/std/limits>
|
||||
#include <cuda/std/span>
|
||||
#include <cuda/std/type_traits>
|
||||
|
||||
#include <cuda/experimental/__multi_gpu/algorithm/scan/scan.h>
|
||||
|
||||
#include <numeric>
|
||||
#include <vector>
|
||||
|
||||
#include <algorithm_common.h>
|
||||
#include <nccl_test_common.h>
|
||||
#include <testing.cuh>
|
||||
|
||||
namespace
|
||||
{
|
||||
struct custom_plus
|
||||
{
|
||||
template <class T>
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr T operator()(const T& lhs, const T& rhs) const
|
||||
{
|
||||
return lhs + rhs;
|
||||
}
|
||||
};
|
||||
|
||||
using custom_value = c2h::custom_type_t<c2h::accumulateable_t, c2h::less_comparable_t, c2h::equal_comparable_t>;
|
||||
using value_types = c2h::type_list<cuda::std::int32_t, float, custom_value>;
|
||||
using operators = c2h::type_list<::cuda::std::plus<>, ::cuda::maximum<>, custom_plus>;
|
||||
|
||||
static_assert(cudax::nccl_transportable<custom_value>);
|
||||
|
||||
template <typename T>
|
||||
T make_value(int i)
|
||||
{
|
||||
return static_cast<T>(i);
|
||||
}
|
||||
|
||||
template <>
|
||||
custom_value make_value<>(int i)
|
||||
{
|
||||
custom_value ret{};
|
||||
|
||||
ret.key = static_cast<std::size_t>(i);
|
||||
ret.val = static_cast<std::size_t>(i);
|
||||
return ret;
|
||||
};
|
||||
|
||||
template <class T, class Op>
|
||||
[[nodiscard]] T get_identity()
|
||||
{
|
||||
if constexpr (cuda::std::is_same_v<Op, cuda::std::plus<>> || cuda::std::is_same_v<Op, custom_plus>)
|
||||
{
|
||||
return make_value<T>(0);
|
||||
}
|
||||
else if constexpr (cuda::std::is_same_v<Op, cuda::maximum<>>)
|
||||
{
|
||||
return cuda::std::numeric_limits<T>::lowest();
|
||||
}
|
||||
else
|
||||
{
|
||||
static_assert(cuda::std::__always_false_v<T, Op>, "Add handling");
|
||||
}
|
||||
}
|
||||
|
||||
template <class T, class Op>
|
||||
[[nodiscard]] std::vector<T>
|
||||
expected_for_rank(int rank, const std::vector<std::vector<T>>& inputs_by_rank, const T& init, Op op)
|
||||
{
|
||||
std::vector<T> reference;
|
||||
|
||||
for (const auto& values : inputs_by_rank)
|
||||
{
|
||||
reference.insert(reference.end(), values.begin(), values.end());
|
||||
}
|
||||
|
||||
std::vector<T> scan(reference.size());
|
||||
std::exclusive_scan(reference.begin(), reference.end(), scan.begin(), init, op);
|
||||
|
||||
cuda::std::size_t offset = 0;
|
||||
for (int r = 0; r < rank; ++r)
|
||||
{
|
||||
offset += inputs_by_rank[static_cast<cuda::std::size_t>(r)].size();
|
||||
}
|
||||
|
||||
const auto count = inputs_by_rank[static_cast<cuda::std::size_t>(rank)].size();
|
||||
return {scan.begin() + offset, scan.begin() + offset + count};
|
||||
}
|
||||
|
||||
// Drive the scan through the single-communicator overload, one thread per local rank. The
|
||||
// per-rank calls must rendezvous in their collectives, so issuing them serially would deadlock.
|
||||
// Catch2 assertions remain on the main thread after all worker threads have joined.
|
||||
template <class T, class Op>
|
||||
void run_case(cuda::std::span<cudax::nccl_communicator_ref> comms,
|
||||
const std::vector<std::vector<T>>& inputs_by_rank,
|
||||
const T& init,
|
||||
const T& ident,
|
||||
Op op)
|
||||
{
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
std::vector<cuda::device_buffer<T>> in;
|
||||
std::vector<cuda::device_buffer<T>> out;
|
||||
std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
|
||||
|
||||
in.reserve(comms.size());
|
||||
out.reserve(comms.size());
|
||||
envs.reserve(comms.size());
|
||||
for (cuda::std::size_t i = 0; i < comms.size(); ++i)
|
||||
{
|
||||
const auto& values = inputs_by_rank[static_cast<cuda::std::size_t>(comms[i].rank())];
|
||||
in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
|
||||
out.emplace_back(cuda::make_device_buffer<T>(
|
||||
streams[i], comms[i].logical_device().underlying_device(), values.size(), cuda::no_init));
|
||||
envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
|
||||
}
|
||||
|
||||
const auto in_copy = in;
|
||||
auto outputs = make_output_iterators(out);
|
||||
const auto outputs_copy = outputs;
|
||||
|
||||
INFO("init = " << init);
|
||||
INFO("ident = " << ident);
|
||||
|
||||
run_threaded(comms.size(), [&](cuda::std::size_t i) {
|
||||
cudax::exclusive_scan(cudax::distributed, comms[i], envs[i], in[i], outputs[i], init, op, ident);
|
||||
});
|
||||
|
||||
REQUIRE(in.size() == in_copy.size());
|
||||
for (cuda::std::size_t i = 0; i < in.size(); ++i)
|
||||
{
|
||||
INFO("device = " << i);
|
||||
REQUIRE_THAT(in[i], Equals(in_copy[i]));
|
||||
const auto expected_values = expected_for_rank<T>(comms[i].rank(), inputs_by_rank, init, op);
|
||||
const auto expected =
|
||||
cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_values);
|
||||
|
||||
REQUIRE_THAT(out[i], Equals(expected));
|
||||
}
|
||||
REQUIRE_THAT(outputs, Catch::Matchers::Equals(outputs_copy));
|
||||
}
|
||||
} // namespace
|
||||
|
||||
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()};
|
||||
|
||||
// Must be pre-allocated since it is written to by threads
|
||||
std::vector<std::string> failed(comms.front().size());
|
||||
|
||||
// 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(cudax::distributed, 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});
|
||||
|
||||
//! [exclusive_scan_single_range]
|
||||
|
||||
// catch2 isn't thread safe by default, so we can't use the usual requires expression. So
|
||||
// we roll a hacky version of it ourselves
|
||||
if (const auto matcher = Equals(expected); !matcher.match(output))
|
||||
{
|
||||
failed[rank] = matcher.describe();
|
||||
}
|
||||
});
|
||||
|
||||
for (cuda::std::size_t i = 0; i < failed.size(); ++i)
|
||||
{
|
||||
if (const auto& err_str = failed[i]; !err_str.empty())
|
||||
{
|
||||
INFO("rank: " << i);
|
||||
REQUIRE(err_str == ""); // Should print the full error string
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("exclusive_scan single-comm, one element per rank", value_types, operators)
|
||||
{
|
||||
using T = c2h::get<0, TestType>;
|
||||
using Op = c2h::get<1, TestType>;
|
||||
|
||||
const T init = make_value<T>(GENERATE(0, 1, -1, 5));
|
||||
const auto ident = get_identity<T, Op>();
|
||||
auto comms = this->communicators();
|
||||
|
||||
std::vector<std::vector<T>> inputs_by_rank;
|
||||
|
||||
inputs_by_rank.reserve(static_cast<cuda::std::size_t>(comms.front().size()));
|
||||
for (int r = 0; r < comms.front().size(); ++r)
|
||||
{
|
||||
const auto v = {make_value<T>(r)};
|
||||
|
||||
inputs_by_rank.emplace_back(v);
|
||||
}
|
||||
|
||||
run_case(comms, inputs_by_rank, init, ident, Op{});
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("exclusive_scan single-comm, multiple elements per rank", value_types, operators)
|
||||
{
|
||||
using T = c2h::get<0, TestType>;
|
||||
using Op = c2h::get<1, TestType>;
|
||||
|
||||
const T init = make_value<T>(GENERATE(0, 1, -1, 5));
|
||||
const auto ident = get_identity<T, Op>();
|
||||
auto comms = this->communicators();
|
||||
|
||||
constexpr auto values_per_rank = 10;
|
||||
std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
|
||||
for (int r = 0; r < comms.front().size(); ++r)
|
||||
{
|
||||
inputs_by_rank[static_cast<cuda::std::size_t>(r)] = std::vector<T>(values_per_rank, make_value<T>(r));
|
||||
}
|
||||
|
||||
run_case(comms, inputs_by_rank, init, ident, Op{});
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("exclusive_scan single-comm, some ranks empty", value_types, operators)
|
||||
{
|
||||
using T = c2h::get<0, TestType>;
|
||||
using Op = c2h::get<1, TestType>;
|
||||
|
||||
const T init = make_value<T>(GENERATE(0, 1, -1, 5));
|
||||
const auto ident = get_identity<T, Op>();
|
||||
auto comms = this->communicators();
|
||||
|
||||
constexpr auto values_per_rank = 10;
|
||||
std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
|
||||
for (int r = 0; r < comms.front().size(); ++r)
|
||||
{
|
||||
if (r % 2 == 0)
|
||||
{
|
||||
inputs_by_rank[static_cast<cuda::std::size_t>(r)] = std::vector<T>(values_per_rank, make_value<T>(r));
|
||||
}
|
||||
}
|
||||
|
||||
run_case(comms, inputs_by_rank, init, ident, Op{});
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("exclusive_scan single-comm, all ranks empty", value_types, operators)
|
||||
{
|
||||
using T = c2h::get<0, TestType>;
|
||||
using Op = c2h::get<1, TestType>;
|
||||
|
||||
const T init = make_value<T>(GENERATE(0, 1, -1, 5));
|
||||
const auto ident = get_identity<T, Op>();
|
||||
auto comms = this->communicators();
|
||||
|
||||
const std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
|
||||
|
||||
run_case(comms, inputs_by_rank, init, ident, Op{});
|
||||
}
|
||||
@@ -1,120 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/buffer>
|
||||
#include <cuda/functional>
|
||||
#include <cuda/memory_resource>
|
||||
#include <cuda/std/cstddef>
|
||||
#include <cuda/std/cstdint>
|
||||
#include <cuda/std/execution>
|
||||
#include <cuda/std/functional>
|
||||
#include <cuda/std/iterator>
|
||||
|
||||
#include <cuda/experimental/__multi_gpu/algorithm/scan/scan.h>
|
||||
|
||||
#include <numeric>
|
||||
#include <vector>
|
||||
|
||||
#include <algorithm_common.h>
|
||||
#include <nccl_test_common.h>
|
||||
#include <testing.cuh>
|
||||
|
||||
MULTI_GPU_TEST("exclusive_scan single-comm, overloads default values", )
|
||||
{
|
||||
using T = cuda::std::int32_t;
|
||||
using Op = ::cuda::std::plus<>;
|
||||
|
||||
constexpr auto init = T{};
|
||||
constexpr T ident = cuda::identity_element<Op, T>();
|
||||
constexpr auto op = Op{};
|
||||
|
||||
const auto comms = this->communicators();
|
||||
const auto streams = nccl_test_util::make_streams();
|
||||
|
||||
std::vector<cuda::device_buffer<T>> in;
|
||||
std::vector<cuda::device_buffer<T>> out;
|
||||
std::vector<cuda::stream_ref> envs;
|
||||
|
||||
in.reserve(comms.size());
|
||||
out.reserve(comms.size());
|
||||
envs.reserve(comms.size());
|
||||
|
||||
constexpr auto values_per_rank = 10;
|
||||
|
||||
for (cuda::std::size_t i = 0; i < comms.size(); ++i)
|
||||
{
|
||||
std::vector<T> values(values_per_rank);
|
||||
std::iota(values.begin(), values.end(), static_cast<T>(comms[i].rank() * values_per_rank + 1));
|
||||
|
||||
in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
|
||||
out.emplace_back(cuda::make_device_buffer<T>(
|
||||
streams[i], comms[i].logical_device().underlying_device(), cuda::std::size(values), cuda::no_init));
|
||||
envs.emplace_back(streams[i]);
|
||||
}
|
||||
|
||||
auto outputs = make_output_iterators(out);
|
||||
|
||||
const auto expected_values = [&] {
|
||||
std::vector<T> reference(static_cast<cuda::std::size_t>(comms.front().size()) * values_per_rank);
|
||||
|
||||
std::iota(reference.begin(), reference.end(), T{1});
|
||||
|
||||
std::vector<T> expected_values(reference.size());
|
||||
|
||||
std::exclusive_scan(reference.begin(), reference.end(), expected_values.begin(), init, op);
|
||||
|
||||
return expected_values;
|
||||
}();
|
||||
|
||||
const auto check_outputs = [&] {
|
||||
const auto exp_span = cuda::std::span{expected_values};
|
||||
|
||||
for (cuda::std::size_t i = 0; i < out.size(); ++i)
|
||||
{
|
||||
const auto expected_for_rank = exp_span.subspan(comms[i].rank() * values_per_rank, values_per_rank);
|
||||
const auto expected =
|
||||
cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_for_rank);
|
||||
|
||||
REQUIRE_THAT(out[i], Equals(expected));
|
||||
}
|
||||
};
|
||||
|
||||
SECTION("Default init, op, ident (all)")
|
||||
{
|
||||
run_threaded(comms.size(), [&](cuda::std::size_t i) {
|
||||
cudax::exclusive_scan(cudax::distributed, comms[i], envs[i], in[i], outputs[i]);
|
||||
});
|
||||
check_outputs();
|
||||
}
|
||||
|
||||
SECTION("Default op, ident")
|
||||
{
|
||||
run_threaded(comms.size(), [&](cuda::std::size_t i) {
|
||||
cudax::exclusive_scan(cudax::distributed, comms[i], envs[i], in[i], outputs[i], init);
|
||||
});
|
||||
check_outputs();
|
||||
}
|
||||
|
||||
SECTION("Default ident")
|
||||
{
|
||||
run_threaded(comms.size(), [&](cuda::std::size_t i) {
|
||||
cudax::exclusive_scan(cudax::distributed, comms[i], envs[i], in[i], outputs[i], init, op);
|
||||
});
|
||||
check_outputs();
|
||||
}
|
||||
|
||||
SECTION("Default none")
|
||||
{
|
||||
run_threaded(comms.size(), [&](cuda::std::size_t i) {
|
||||
cudax::exclusive_scan(cudax::distributed, comms[i], envs[i], in[i], outputs[i], init, op, ident);
|
||||
});
|
||||
check_outputs();
|
||||
}
|
||||
}
|
||||
@@ -1,20 +0,0 @@
|
||||
#===----------------------------------------------------------------------===##
|
||||
#
|
||||
# Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
# under the Apache License v2.0 with LLVM Exceptions.
|
||||
# See https://llvm.org/LICENSE.txt for license information.
|
||||
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
#
|
||||
#===----------------------------------------------------------------------===##
|
||||
|
||||
if (NOT cudax_ENABLE_NCCL)
|
||||
return()
|
||||
endif()
|
||||
|
||||
file(GLOB test_srcs LIST_DIRECTORIES FALSE CONFIGURE_DEPENDS *.cu *.cpp)
|
||||
|
||||
foreach (src IN LISTS test_srcs)
|
||||
cudax_add_multi_gpu_test("algorithms.inclusive_scan" test_target "${src}")
|
||||
target_link_libraries(${test_target} PRIVATE NCCL::nccl)
|
||||
endforeach()
|
||||
@@ -1,389 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/buffer>
|
||||
#include <cuda/functional>
|
||||
#include <cuda/memory_resource>
|
||||
#include <cuda/std/array>
|
||||
#include <cuda/std/cstddef>
|
||||
#include <cuda/std/cstdint>
|
||||
#include <cuda/std/execution>
|
||||
#include <cuda/std/functional>
|
||||
#include <cuda/std/limits>
|
||||
#include <cuda/std/span>
|
||||
#include <cuda/std/type_traits>
|
||||
|
||||
#include <cuda/experimental/__multi_gpu/algorithm/scan/scan.h>
|
||||
|
||||
#include <numeric>
|
||||
#include <vector>
|
||||
|
||||
#include <algorithm_common.h>
|
||||
#include <nccl_test_common.h>
|
||||
#include <testing.cuh>
|
||||
|
||||
namespace
|
||||
{
|
||||
struct custom_plus
|
||||
{
|
||||
template <class T>
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr T operator()(const T& lhs, const T& rhs) const
|
||||
{
|
||||
return lhs + rhs;
|
||||
}
|
||||
};
|
||||
|
||||
using custom_value = c2h::custom_type_t<c2h::accumulateable_t, c2h::less_comparable_t, c2h::equal_comparable_t>;
|
||||
using value_types = c2h::type_list<cuda::std::int32_t, float, custom_value>;
|
||||
using operators = c2h::type_list<::cuda::std::plus<>, ::cuda::maximum<>, custom_plus>;
|
||||
|
||||
static_assert(cudax::nccl_transportable<custom_value>);
|
||||
|
||||
template <typename T>
|
||||
T make_value(int i)
|
||||
{
|
||||
return static_cast<T>(i);
|
||||
}
|
||||
|
||||
template <>
|
||||
custom_value make_value<>(int i)
|
||||
{
|
||||
custom_value ret{};
|
||||
|
||||
ret.key = static_cast<std::size_t>(i);
|
||||
ret.val = static_cast<std::size_t>(i);
|
||||
return ret;
|
||||
};
|
||||
|
||||
template <class T, class Op>
|
||||
[[nodiscard]] T get_identity()
|
||||
{
|
||||
if constexpr (cuda::std::is_same_v<Op, cuda::std::plus<>> || cuda::std::is_same_v<Op, custom_plus>)
|
||||
{
|
||||
return make_value<T>(0);
|
||||
}
|
||||
else if constexpr (cuda::std::is_same_v<Op, cuda::maximum<>>)
|
||||
{
|
||||
return cuda::std::numeric_limits<T>::lowest();
|
||||
}
|
||||
else
|
||||
{
|
||||
static_assert(cuda::std::__always_false_v<T, Op>, "Add handling");
|
||||
}
|
||||
}
|
||||
|
||||
template <class T, class Op>
|
||||
[[nodiscard]] std::vector<T>
|
||||
expected_for_rank(int rank, const std::vector<std::vector<T>>& inputs_by_rank, const T& init, Op op)
|
||||
{
|
||||
std::vector<T> reference;
|
||||
|
||||
for (const auto& values : inputs_by_rank)
|
||||
{
|
||||
reference.insert(reference.end(), values.begin(), values.end());
|
||||
}
|
||||
|
||||
std::vector<T> scan(reference.size());
|
||||
std::inclusive_scan(reference.begin(), reference.end(), scan.begin(), op, init);
|
||||
|
||||
cuda::std::size_t offset = 0;
|
||||
for (int r = 0; r < rank; ++r)
|
||||
{
|
||||
offset += inputs_by_rank[static_cast<cuda::std::size_t>(r)].size();
|
||||
}
|
||||
|
||||
const auto count = inputs_by_rank[static_cast<cuda::std::size_t>(rank)].size();
|
||||
return {scan.begin() + offset, scan.begin() + offset + count};
|
||||
}
|
||||
|
||||
// Run the full scan, wait for it to finish, and check that `inclusive_scan` left its argument
|
||||
// ranges untouched. This boilerplate is identical for every test regardless of how the inputs are
|
||||
// shaped.
|
||||
template <class Env, class T, class Op>
|
||||
void do_inclusive_scan(
|
||||
cuda::std::span<cudax::nccl_communicator_ref> comms,
|
||||
const std::vector<Env>& envs,
|
||||
std::vector<cuda::device_buffer<T>>& in,
|
||||
std::vector<typename cuda::device_buffer<T>::iterator>& outputs,
|
||||
const T& init,
|
||||
const T& ident,
|
||||
Op op)
|
||||
{
|
||||
const auto envs_size = envs.size();
|
||||
const auto in_copy = in;
|
||||
const auto outputs_copy = outputs;
|
||||
|
||||
INFO("init = " << init);
|
||||
INFO("ident = " << ident);
|
||||
|
||||
cudax::inclusive_scan(cudax::distributed, comms, envs, in, outputs, init, op, ident);
|
||||
|
||||
// cuda::std::execution::env has no operator==, so we can only compare the sizes.
|
||||
REQUIRE(envs.size() == envs_size);
|
||||
// Scan call should not modify the inputs in any ways
|
||||
REQUIRE(in.size() == in_copy.size());
|
||||
for (cuda::std::size_t i = 0; i < in.size(); ++i)
|
||||
{
|
||||
INFO("device = " << i);
|
||||
REQUIRE_THAT(in[i], Equals(in_copy[i]));
|
||||
}
|
||||
REQUIRE_THAT(outputs, Catch::Matchers::Equals(outputs_copy));
|
||||
}
|
||||
} // namespace
|
||||
|
||||
MULTI_GPU_TEST("inclusive_scan documentation example", c2h::type_list<int>)
|
||||
{
|
||||
auto comms = this->communicators();
|
||||
|
||||
if (comms.size() < 2)
|
||||
{
|
||||
SKIP("The inclusive_scan documentation example requires at least two local GPUs");
|
||||
}
|
||||
|
||||
auto streams_owned = nccl_test_util::make_streams();
|
||||
// Convert to stream_ref directly, cuda::stream on their own cant be passed directly to CUB
|
||||
auto streams = std::vector<cuda::stream_ref>{streams_owned.begin(), streams_owned.end()};
|
||||
|
||||
//! [inclusive_scan]
|
||||
constexpr cuda::std::array input_values{1, 2};
|
||||
std::vector<cuda::device_buffer<int>> inputs;
|
||||
std::vector<cuda::device_buffer<int>> outputs;
|
||||
|
||||
for (cuda::std::size_t i = 0; i < comms.size(); ++i)
|
||||
{
|
||||
const auto device = comms[i].logical_device().underlying_device();
|
||||
|
||||
inputs.emplace_back(cuda::make_device_buffer<int>(streams[i], device, input_values));
|
||||
outputs.emplace_back(cuda::make_device_buffer<int>(streams[i], device, input_values.size(), cuda::no_init));
|
||||
}
|
||||
|
||||
std::vector<typename cuda::device_buffer<int>::iterator> output_iterators = make_output_iterators(outputs);
|
||||
|
||||
cudax::inclusive_scan(
|
||||
cudax::distributed,
|
||||
comms,
|
||||
// Passing streams as the environment directly
|
||||
streams,
|
||||
inputs,
|
||||
output_iterators,
|
||||
/*__init=*/0);
|
||||
|
||||
constexpr cuda::std::array expected_rank_0{1, 3};
|
||||
constexpr cuda::std::array expected_rank_1{4, 6};
|
||||
const auto expected_0 =
|
||||
cuda::make_buffer<int>(outputs[0].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_rank_0);
|
||||
const auto expected_1 =
|
||||
cuda::make_buffer<int>(outputs[1].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_rank_1);
|
||||
REQUIRE_THAT(outputs[0], Equals(expected_0));
|
||||
REQUIRE_THAT(outputs[1], Equals(expected_1));
|
||||
//! [inclusive_scan]
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("inclusive_scan, one element per rank", value_types, operators)
|
||||
{
|
||||
using T = c2h::get<0, TestType>;
|
||||
using Op = c2h::get<1, TestType>;
|
||||
|
||||
// Seed each scan with a few hardcoded initializers. The init participates in the fold the same
|
||||
// way on host and device, so any value works for every operator under test.
|
||||
const T init = make_value<T>(GENERATE(0, 1, -1, 5));
|
||||
const auto ident = get_identity<T, Op>();
|
||||
|
||||
auto comms = this->communicators();
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
// Global rank `comms[i].rank()` contributes the single value `rank`. Each local rank also gets a
|
||||
// one-element output buffer and an environment carrying its stream, so the scan is stream-ordered
|
||||
// on the correct device. `reference` mirrors the contributions of every global rank so we can
|
||||
// compute the host-side scan exactly like `inclusive_scan` does on the device.
|
||||
std::vector<cuda::device_buffer<T>> in;
|
||||
std::vector<cuda::device_buffer<T>> out;
|
||||
std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
|
||||
std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
|
||||
|
||||
in.reserve(comms.size());
|
||||
out.reserve(comms.size());
|
||||
envs.reserve(comms.size());
|
||||
for (int r = 0; r < comms.front().size(); ++r)
|
||||
{
|
||||
inputs_by_rank[static_cast<cuda::std::size_t>(r)] = std::vector<T>(1, make_value<T>(r));
|
||||
}
|
||||
for (cuda::std::size_t i = 0; i < comms.size(); ++i)
|
||||
{
|
||||
const auto& values = inputs_by_rank[static_cast<cuda::std::size_t>(comms[i].rank())];
|
||||
in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
|
||||
out.emplace_back(cuda::make_device_buffer<T>(
|
||||
streams[i], comms[i].logical_device().underlying_device(), values.size(), cuda::no_init));
|
||||
envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
|
||||
}
|
||||
|
||||
auto outputs = make_output_iterators(out);
|
||||
|
||||
do_inclusive_scan(comms, envs, in, outputs, init, ident, Op{});
|
||||
|
||||
for (cuda::std::size_t i = 0; i < out.size(); ++i)
|
||||
{
|
||||
const auto expected_values = expected_for_rank<T>(comms[i].rank(), inputs_by_rank, init, Op{});
|
||||
const auto exp = cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_values);
|
||||
|
||||
REQUIRE_THAT(out[i], Equals(exp));
|
||||
}
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("inclusive_scan, multiple elements per rank", value_types, operators)
|
||||
{
|
||||
using T = c2h::get<0, TestType>;
|
||||
using Op = c2h::get<1, TestType>;
|
||||
|
||||
// Seed each scan with a few hardcoded initializers. The init participates in the fold the same
|
||||
// way on host and device, so any value works for every operator under test.
|
||||
const T init = make_value<T>(GENERATE(0, 1, -1, 5));
|
||||
const auto ident = get_identity<T, Op>();
|
||||
|
||||
auto comms = this->communicators();
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
// Global rank `comms[i].rank()` contributes ten copies of `rank`. `inclusive_scan` first
|
||||
// computes a local prefix for each rank seeded by the prefix of all previous ranks. Each local
|
||||
// rank also gets an output buffer and an environment carrying its stream. `reference` mirrors
|
||||
// every global rank's ten contributions for the host-side scan.
|
||||
std::vector<cuda::device_buffer<T>> in;
|
||||
std::vector<cuda::device_buffer<T>> out;
|
||||
std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
|
||||
std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
|
||||
|
||||
in.reserve(comms.size());
|
||||
out.reserve(comms.size());
|
||||
envs.reserve(comms.size());
|
||||
|
||||
constexpr auto values_per_rank = 10;
|
||||
for (int r = 0; r < comms.front().size(); ++r)
|
||||
{
|
||||
const auto v = make_value<T>(r);
|
||||
inputs_by_rank[static_cast<cuda::std::size_t>(r)] = std::vector<T>(values_per_rank, v);
|
||||
}
|
||||
for (cuda::std::size_t i = 0; i < comms.size(); ++i)
|
||||
{
|
||||
const auto& values = inputs_by_rank[static_cast<cuda::std::size_t>(comms[i].rank())];
|
||||
in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
|
||||
out.emplace_back(cuda::make_device_buffer<T>(
|
||||
streams[i], comms[i].logical_device().underlying_device(), values.size(), cuda::no_init));
|
||||
envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
|
||||
}
|
||||
|
||||
auto outputs = make_output_iterators(out);
|
||||
|
||||
do_inclusive_scan(comms, envs, in, outputs, init, ident, Op{});
|
||||
|
||||
for (cuda::std::size_t i = 0; i < out.size(); ++i)
|
||||
{
|
||||
const auto expected_values = expected_for_rank<T>(comms[i].rank(), inputs_by_rank, init, Op{});
|
||||
const auto exp = cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_values);
|
||||
|
||||
REQUIRE_THAT(out[i], Equals(exp));
|
||||
}
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("inclusive_scan, some ranks empty", value_types, operators)
|
||||
{
|
||||
using T = c2h::get<0, TestType>;
|
||||
using Op = c2h::get<1, TestType>;
|
||||
|
||||
const T init = make_value<T>(GENERATE(0, 1, -1, 5));
|
||||
const auto ident = get_identity<T, Op>();
|
||||
|
||||
auto comms = this->communicators();
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
// Even global ranks contribute ten copies of `rank`; odd global ranks contribute an empty input
|
||||
// range. Rank 0 is always non-empty. `inclusive_scan` must treat an empty rank as contributing
|
||||
// nothing, exactly like `std::inclusive_scan` over the surviving elements. `reference` mirrors
|
||||
// that for the host-side scan.
|
||||
std::vector<cuda::device_buffer<T>> in;
|
||||
std::vector<cuda::device_buffer<T>> out;
|
||||
std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
|
||||
std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
|
||||
|
||||
in.reserve(comms.size());
|
||||
out.reserve(comms.size());
|
||||
envs.reserve(comms.size());
|
||||
|
||||
constexpr auto values_per_rank = 10;
|
||||
for (int r = 0; r < comms.front().size(); ++r)
|
||||
{
|
||||
if (r % 2 == 0)
|
||||
{
|
||||
inputs_by_rank[static_cast<cuda::std::size_t>(r)] = std::vector<T>(values_per_rank, make_value<T>(r));
|
||||
}
|
||||
}
|
||||
for (cuda::std::size_t i = 0; i < comms.size(); ++i)
|
||||
{
|
||||
const auto& values = inputs_by_rank[static_cast<cuda::std::size_t>(comms[i].rank())];
|
||||
in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
|
||||
out.emplace_back(cuda::make_device_buffer<T>(
|
||||
streams[i], comms[i].logical_device().underlying_device(), values.size(), cuda::no_init));
|
||||
envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
|
||||
}
|
||||
|
||||
auto outputs = make_output_iterators(out);
|
||||
|
||||
do_inclusive_scan(comms, envs, in, outputs, init, ident, Op{});
|
||||
|
||||
for (cuda::std::size_t i = 0; i < out.size(); ++i)
|
||||
{
|
||||
const auto expected_values = expected_for_rank<T>(comms[i].rank(), inputs_by_rank, init, Op{});
|
||||
const auto exp = cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_values);
|
||||
|
||||
REQUIRE_THAT(out[i], Equals(exp));
|
||||
}
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("inclusive_scan, all ranks empty", value_types, operators)
|
||||
{
|
||||
using T = c2h::get<0, TestType>;
|
||||
using Op = c2h::get<1, TestType>;
|
||||
|
||||
const auto init = make_value<T>(GENERATE(0, 1, -1, 5));
|
||||
const auto ident = get_identity<T, Op>();
|
||||
|
||||
auto comms = this->communicators();
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
// No rank contributes any element. Scanning nothing produces no output values, exactly like
|
||||
// `std::inclusive_scan` over an empty range.
|
||||
std::vector<cuda::device_buffer<T>> in;
|
||||
std::vector<cuda::device_buffer<T>> out;
|
||||
std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
|
||||
std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
|
||||
|
||||
in.reserve(comms.size());
|
||||
out.reserve(comms.size());
|
||||
envs.reserve(comms.size());
|
||||
for (cuda::std::size_t i = 0; i < comms.size(); ++i)
|
||||
{
|
||||
const auto& values = inputs_by_rank[static_cast<cuda::std::size_t>(comms[i].rank())];
|
||||
in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
|
||||
out.emplace_back(cuda::make_device_buffer<T>(
|
||||
streams[i], comms[i].logical_device().underlying_device(), values.size(), cuda::no_init));
|
||||
envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
|
||||
}
|
||||
|
||||
auto outputs = make_output_iterators(out);
|
||||
|
||||
do_inclusive_scan(comms, envs, in, outputs, init, ident, Op{});
|
||||
|
||||
for (cuda::std::size_t i = 0; i < out.size(); ++i)
|
||||
{
|
||||
const auto expected_values = expected_for_rank<T>(comms[i].rank(), inputs_by_rank, init, Op{});
|
||||
const auto exp = cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_values);
|
||||
|
||||
REQUIRE_THAT(out[i], Equals(exp));
|
||||
}
|
||||
}
|
||||
@@ -1,112 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/buffer>
|
||||
#include <cuda/functional>
|
||||
#include <cuda/memory_resource>
|
||||
#include <cuda/std/cstddef>
|
||||
#include <cuda/std/cstdint>
|
||||
#include <cuda/std/execution>
|
||||
#include <cuda/std/functional>
|
||||
#include <cuda/std/iterator>
|
||||
|
||||
#include <cuda/experimental/__multi_gpu/algorithm/scan/scan.h>
|
||||
|
||||
#include <numeric>
|
||||
#include <vector>
|
||||
|
||||
#include <algorithm_common.h>
|
||||
#include <nccl_test_common.h>
|
||||
#include <testing.cuh>
|
||||
|
||||
MULTI_GPU_TEST("inclusive_scan, range overloads default values", )
|
||||
{
|
||||
using T = cuda::std::int32_t;
|
||||
using Op = ::cuda::std::plus<>;
|
||||
|
||||
constexpr auto init = T{};
|
||||
constexpr T ident = cuda::identity_element<Op, T>();
|
||||
constexpr auto op = Op{};
|
||||
|
||||
const auto comms = this->communicators();
|
||||
const auto streams = nccl_test_util::make_streams();
|
||||
|
||||
std::vector<cuda::device_buffer<T>> in;
|
||||
std::vector<cuda::device_buffer<T>> out;
|
||||
std::vector<cuda::stream_ref> envs;
|
||||
|
||||
in.reserve(comms.size());
|
||||
out.reserve(comms.size());
|
||||
envs.reserve(comms.size());
|
||||
|
||||
constexpr auto values_per_rank = 10;
|
||||
|
||||
for (cuda::std::size_t i = 0; i < comms.size(); ++i)
|
||||
{
|
||||
std::vector<T> values(values_per_rank);
|
||||
std::iota(values.begin(), values.end(), static_cast<T>(comms[i].rank() * values_per_rank + 1));
|
||||
|
||||
in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
|
||||
out.emplace_back(cuda::make_device_buffer<T>(
|
||||
streams[i], comms[i].logical_device().underlying_device(), cuda::std::size(values), cuda::no_init));
|
||||
envs.emplace_back(streams[i]);
|
||||
}
|
||||
|
||||
auto outputs = make_output_iterators(out);
|
||||
|
||||
const auto expected_values = [&] {
|
||||
std::vector<T> reference(static_cast<cuda::std::size_t>(comms.front().size()) * values_per_rank);
|
||||
|
||||
std::iota(reference.begin(), reference.end(), T{1});
|
||||
|
||||
std::vector<T> expected_values(reference.size());
|
||||
|
||||
std::inclusive_scan(reference.begin(), reference.end(), expected_values.begin(), op, init);
|
||||
|
||||
return expected_values;
|
||||
}();
|
||||
|
||||
const auto check_outputs = [&] {
|
||||
const auto exp_span = cuda::std::span{expected_values};
|
||||
|
||||
for (cuda::std::size_t i = 0; i < out.size(); ++i)
|
||||
{
|
||||
const auto expected_for_rank = exp_span.subspan(comms[i].rank() * values_per_rank, values_per_rank);
|
||||
const auto expected =
|
||||
cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_for_rank);
|
||||
|
||||
REQUIRE_THAT(out[i], Equals(expected));
|
||||
}
|
||||
};
|
||||
|
||||
SECTION("Default init, op, ident (all)")
|
||||
{
|
||||
cudax::inclusive_scan(cudax::distributed, comms, envs, in, outputs);
|
||||
check_outputs();
|
||||
}
|
||||
|
||||
SECTION("Default op, ident")
|
||||
{
|
||||
cudax::inclusive_scan(cudax::distributed, comms, envs, in, outputs, init);
|
||||
check_outputs();
|
||||
}
|
||||
|
||||
SECTION("Default ident")
|
||||
{
|
||||
cudax::inclusive_scan(cudax::distributed, comms, envs, in, outputs, init, op);
|
||||
check_outputs();
|
||||
}
|
||||
|
||||
SECTION("Default none")
|
||||
{
|
||||
cudax::inclusive_scan(cudax::distributed, comms, envs, in, outputs, init, op, ident);
|
||||
check_outputs();
|
||||
}
|
||||
}
|
||||
@@ -1,290 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/buffer>
|
||||
#include <cuda/functional>
|
||||
#include <cuda/memory_resource>
|
||||
#include <cuda/std/array>
|
||||
#include <cuda/std/cstddef>
|
||||
#include <cuda/std/cstdint>
|
||||
#include <cuda/std/execution>
|
||||
#include <cuda/std/functional>
|
||||
#include <cuda/std/limits>
|
||||
#include <cuda/std/span>
|
||||
#include <cuda/std/type_traits>
|
||||
|
||||
#include <cuda/experimental/__multi_gpu/algorithm/scan/scan.h>
|
||||
|
||||
#include <numeric>
|
||||
#include <vector>
|
||||
|
||||
#include <algorithm_common.h>
|
||||
#include <nccl_test_common.h>
|
||||
#include <testing.cuh>
|
||||
|
||||
namespace
|
||||
{
|
||||
struct custom_plus
|
||||
{
|
||||
template <class T>
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr T operator()(const T& lhs, const T& rhs) const
|
||||
{
|
||||
return lhs + rhs;
|
||||
}
|
||||
};
|
||||
|
||||
using custom_value = c2h::custom_type_t<c2h::accumulateable_t, c2h::less_comparable_t, c2h::equal_comparable_t>;
|
||||
using value_types = c2h::type_list<cuda::std::int32_t, float, custom_value>;
|
||||
using operators = c2h::type_list<::cuda::std::plus<>, ::cuda::maximum<>, custom_plus>;
|
||||
|
||||
static_assert(cudax::nccl_transportable<custom_value>);
|
||||
|
||||
template <typename T>
|
||||
T make_value(int i)
|
||||
{
|
||||
return static_cast<T>(i);
|
||||
}
|
||||
|
||||
template <>
|
||||
custom_value make_value<>(int i)
|
||||
{
|
||||
custom_value ret{};
|
||||
|
||||
ret.key = static_cast<std::size_t>(i);
|
||||
ret.val = static_cast<std::size_t>(i);
|
||||
return ret;
|
||||
};
|
||||
|
||||
template <class T, class Op>
|
||||
[[nodiscard]] T get_identity()
|
||||
{
|
||||
if constexpr (cuda::std::is_same_v<Op, cuda::std::plus<>> || cuda::std::is_same_v<Op, custom_plus>)
|
||||
{
|
||||
return make_value<T>(0);
|
||||
}
|
||||
else if constexpr (cuda::std::is_same_v<Op, cuda::maximum<>>)
|
||||
{
|
||||
return cuda::std::numeric_limits<T>::lowest();
|
||||
}
|
||||
else
|
||||
{
|
||||
static_assert(cuda::std::__always_false_v<T, Op>, "Add handling");
|
||||
}
|
||||
}
|
||||
|
||||
template <class T, class Op>
|
||||
[[nodiscard]] std::vector<T>
|
||||
expected_for_rank(int rank, const std::vector<std::vector<T>>& inputs_by_rank, const T& init, Op op)
|
||||
{
|
||||
std::vector<T> reference;
|
||||
|
||||
for (const auto& values : inputs_by_rank)
|
||||
{
|
||||
reference.insert(reference.end(), values.begin(), values.end());
|
||||
}
|
||||
|
||||
std::vector<T> scan(reference.size());
|
||||
std::inclusive_scan(reference.begin(), reference.end(), scan.begin(), op, init);
|
||||
|
||||
cuda::std::size_t offset = 0;
|
||||
for (int r = 0; r < rank; ++r)
|
||||
{
|
||||
offset += inputs_by_rank[static_cast<cuda::std::size_t>(r)].size();
|
||||
}
|
||||
|
||||
const auto count = inputs_by_rank[static_cast<cuda::std::size_t>(rank)].size();
|
||||
return {scan.begin() + offset, scan.begin() + offset + count};
|
||||
}
|
||||
|
||||
// Drive the scan through the single-communicator overload, one thread per local rank. The
|
||||
// per-rank calls must rendezvous in their collectives, so issuing them serially would deadlock.
|
||||
// Catch2 assertions remain on the main thread after all worker threads have joined.
|
||||
template <class T, class Op>
|
||||
void run_case(cuda::std::span<cudax::nccl_communicator_ref> comms,
|
||||
const std::vector<std::vector<T>>& inputs_by_rank,
|
||||
const T& init,
|
||||
const T& ident,
|
||||
Op op)
|
||||
{
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
std::vector<cuda::device_buffer<T>> in;
|
||||
std::vector<cuda::device_buffer<T>> out;
|
||||
std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
|
||||
|
||||
in.reserve(comms.size());
|
||||
out.reserve(comms.size());
|
||||
envs.reserve(comms.size());
|
||||
for (cuda::std::size_t i = 0; i < comms.size(); ++i)
|
||||
{
|
||||
const auto& values = inputs_by_rank[static_cast<cuda::std::size_t>(comms[i].rank())];
|
||||
in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
|
||||
out.emplace_back(cuda::make_device_buffer<T>(
|
||||
streams[i], comms[i].logical_device().underlying_device(), values.size(), cuda::no_init));
|
||||
envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
|
||||
}
|
||||
|
||||
const auto in_copy = in;
|
||||
auto outputs = make_output_iterators(out);
|
||||
const auto outputs_copy = outputs;
|
||||
|
||||
INFO("init = " << init);
|
||||
INFO("ident = " << ident);
|
||||
|
||||
run_threaded(comms.size(), [&](cuda::std::size_t i) {
|
||||
cudax::inclusive_scan(cudax::distributed, comms[i], envs[i], in[i], outputs[i], init, op, ident);
|
||||
});
|
||||
|
||||
REQUIRE(in.size() == in_copy.size());
|
||||
for (cuda::std::size_t i = 0; i < in.size(); ++i)
|
||||
{
|
||||
INFO("device = " << i);
|
||||
REQUIRE_THAT(in[i], Equals(in_copy[i]));
|
||||
const auto expected_values = expected_for_rank<T>(comms[i].rank(), inputs_by_rank, init, op);
|
||||
const auto expected =
|
||||
cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_values);
|
||||
|
||||
REQUIRE_THAT(out[i], Equals(expected));
|
||||
}
|
||||
REQUIRE_THAT(outputs, Catch::Matchers::Equals(outputs_copy));
|
||||
}
|
||||
} // namespace
|
||||
|
||||
MULTI_GPU_TEST("inclusive_scan single-comm documentation example", c2h::type_list<int>)
|
||||
{
|
||||
auto comms = this->communicators();
|
||||
|
||||
if (comms.size() < 2)
|
||||
{
|
||||
SKIP("The inclusive_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()};
|
||||
|
||||
// Must be pre-allocated since it is written to by threads
|
||||
std::vector<std::string> failed(comms.front().size());
|
||||
|
||||
// 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();
|
||||
|
||||
//! [inclusive_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::inclusive_scan(cudax::distributed, 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 + 1, 3 * rank + 3});
|
||||
|
||||
//! [inclusive_scan_single_range]
|
||||
|
||||
// catch2 isn't thread safe by default, so we can't use the usual requires expression. So
|
||||
// we roll a hacky version of it ourselves
|
||||
if (const auto matcher = Equals(expected); !matcher.match(output))
|
||||
{
|
||||
failed[rank] = matcher.describe();
|
||||
}
|
||||
});
|
||||
|
||||
for (cuda::std::size_t i = 0; i < failed.size(); ++i)
|
||||
{
|
||||
if (const auto& err_str = failed[i]; !err_str.empty())
|
||||
{
|
||||
INFO("rank: " << i);
|
||||
REQUIRE(err_str == ""); // Should print the full error string
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("inclusive_scan single-comm, one element per rank", value_types, operators)
|
||||
{
|
||||
using T = c2h::get<0, TestType>;
|
||||
using Op = c2h::get<1, TestType>;
|
||||
|
||||
const T init = make_value<T>(GENERATE(0, 1, -1, 5));
|
||||
const auto ident = get_identity<T, Op>();
|
||||
auto comms = this->communicators();
|
||||
|
||||
std::vector<std::vector<T>> inputs_by_rank;
|
||||
|
||||
inputs_by_rank.reserve(static_cast<cuda::std::size_t>(comms.front().size()));
|
||||
for (int r = 0; r < comms.front().size(); ++r)
|
||||
{
|
||||
const auto v = {make_value<T>(r)};
|
||||
|
||||
inputs_by_rank.emplace_back(v);
|
||||
}
|
||||
|
||||
run_case(comms, inputs_by_rank, init, ident, Op{});
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("inclusive_scan single-comm, multiple elements per rank", value_types, operators)
|
||||
{
|
||||
using T = c2h::get<0, TestType>;
|
||||
using Op = c2h::get<1, TestType>;
|
||||
|
||||
const T init = make_value<T>(GENERATE(0, 1, -1, 5));
|
||||
const auto ident = get_identity<T, Op>();
|
||||
auto comms = this->communicators();
|
||||
|
||||
constexpr auto values_per_rank = 10;
|
||||
std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
|
||||
for (int r = 0; r < comms.front().size(); ++r)
|
||||
{
|
||||
const auto value = make_value<T>(r);
|
||||
inputs_by_rank[static_cast<cuda::std::size_t>(r)] = std::vector<T>(values_per_rank, value);
|
||||
}
|
||||
|
||||
run_case(comms, inputs_by_rank, init, ident, Op{});
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("inclusive_scan single-comm, some ranks empty", value_types, operators)
|
||||
{
|
||||
using T = c2h::get<0, TestType>;
|
||||
using Op = c2h::get<1, TestType>;
|
||||
|
||||
const T init = make_value<T>(GENERATE(0, 1, -1, 5));
|
||||
const auto ident = get_identity<T, Op>();
|
||||
auto comms = this->communicators();
|
||||
|
||||
constexpr auto values_per_rank = 10;
|
||||
std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
|
||||
for (int r = 0; r < comms.front().size(); ++r)
|
||||
{
|
||||
if (r % 2 == 0)
|
||||
{
|
||||
inputs_by_rank[static_cast<cuda::std::size_t>(r)] = std::vector<T>(values_per_rank, make_value<T>(r));
|
||||
}
|
||||
}
|
||||
|
||||
run_case(comms, inputs_by_rank, init, ident, Op{});
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("inclusive_scan single-comm, all ranks empty", value_types, operators)
|
||||
{
|
||||
using T = c2h::get<0, TestType>;
|
||||
using Op = c2h::get<1, TestType>;
|
||||
|
||||
const T init = make_value<T>(GENERATE(0, 1, -1, 5));
|
||||
const auto ident = get_identity<T, Op>();
|
||||
auto comms = this->communicators();
|
||||
|
||||
const std::vector<std::vector<T>> inputs_by_rank(static_cast<cuda::std::size_t>(comms.front().size()));
|
||||
|
||||
run_case(comms, inputs_by_rank, init, ident, Op{});
|
||||
}
|
||||
@@ -1,120 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/buffer>
|
||||
#include <cuda/functional>
|
||||
#include <cuda/memory_resource>
|
||||
#include <cuda/std/cstddef>
|
||||
#include <cuda/std/cstdint>
|
||||
#include <cuda/std/execution>
|
||||
#include <cuda/std/functional>
|
||||
#include <cuda/std/iterator>
|
||||
|
||||
#include <cuda/experimental/__multi_gpu/algorithm/scan/scan.h>
|
||||
|
||||
#include <numeric>
|
||||
#include <vector>
|
||||
|
||||
#include <algorithm_common.h>
|
||||
#include <nccl_test_common.h>
|
||||
#include <testing.cuh>
|
||||
|
||||
MULTI_GPU_TEST("inclusive_scan single-comm, overloads default values", )
|
||||
{
|
||||
using T = cuda::std::int32_t;
|
||||
using Op = ::cuda::std::plus<>;
|
||||
|
||||
constexpr auto init = T{};
|
||||
constexpr T ident = cuda::identity_element<Op, T>();
|
||||
constexpr auto op = Op{};
|
||||
|
||||
const auto comms = this->communicators();
|
||||
const auto streams = nccl_test_util::make_streams();
|
||||
|
||||
std::vector<cuda::device_buffer<T>> in;
|
||||
std::vector<cuda::device_buffer<T>> out;
|
||||
std::vector<cuda::stream_ref> envs;
|
||||
|
||||
in.reserve(comms.size());
|
||||
out.reserve(comms.size());
|
||||
envs.reserve(comms.size());
|
||||
|
||||
constexpr auto values_per_rank = 10;
|
||||
|
||||
for (cuda::std::size_t i = 0; i < comms.size(); ++i)
|
||||
{
|
||||
std::vector<T> values(values_per_rank);
|
||||
std::iota(values.begin(), values.end(), static_cast<T>(comms[i].rank() * values_per_rank + 1));
|
||||
|
||||
in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
|
||||
out.emplace_back(cuda::make_device_buffer<T>(
|
||||
streams[i], comms[i].logical_device().underlying_device(), cuda::std::size(values), cuda::no_init));
|
||||
envs.emplace_back(streams[i]);
|
||||
}
|
||||
|
||||
auto outputs = make_output_iterators(out);
|
||||
|
||||
const auto expected_values = [&] {
|
||||
std::vector<T> reference(static_cast<cuda::std::size_t>(comms.front().size()) * values_per_rank);
|
||||
|
||||
std::iota(reference.begin(), reference.end(), T{1});
|
||||
|
||||
std::vector<T> expected_values(reference.size());
|
||||
|
||||
std::inclusive_scan(reference.begin(), reference.end(), expected_values.begin(), op, init);
|
||||
|
||||
return expected_values;
|
||||
}();
|
||||
|
||||
const auto check_outputs = [&] {
|
||||
const auto exp_span = cuda::std::span{expected_values};
|
||||
|
||||
for (cuda::std::size_t i = 0; i < out.size(); ++i)
|
||||
{
|
||||
const auto expected_for_rank = exp_span.subspan(comms[i].rank() * values_per_rank, values_per_rank);
|
||||
const auto expected =
|
||||
cuda::make_buffer<T>(out[i].stream(), cuda::mr::legacy_pinned_memory_resource{}, expected_for_rank);
|
||||
|
||||
REQUIRE_THAT(out[i], Equals(expected));
|
||||
}
|
||||
};
|
||||
|
||||
SECTION("Default init, op, ident (all)")
|
||||
{
|
||||
run_threaded(comms.size(), [&](cuda::std::size_t i) {
|
||||
cudax::inclusive_scan(cudax::distributed, comms[i], envs[i], in[i], outputs[i]);
|
||||
});
|
||||
check_outputs();
|
||||
}
|
||||
|
||||
SECTION("Default op, ident")
|
||||
{
|
||||
run_threaded(comms.size(), [&](cuda::std::size_t i) {
|
||||
cudax::inclusive_scan(cudax::distributed, comms[i], envs[i], in[i], outputs[i], init);
|
||||
});
|
||||
check_outputs();
|
||||
}
|
||||
|
||||
SECTION("Default ident")
|
||||
{
|
||||
run_threaded(comms.size(), [&](cuda::std::size_t i) {
|
||||
cudax::inclusive_scan(cudax::distributed, comms[i], envs[i], in[i], outputs[i], init, op);
|
||||
});
|
||||
check_outputs();
|
||||
}
|
||||
|
||||
SECTION("Default none")
|
||||
{
|
||||
run_threaded(comms.size(), [&](cuda::std::size_t i) {
|
||||
cudax::inclusive_scan(cudax::distributed, comms[i], envs[i], in[i], outputs[i], init, op, ident);
|
||||
});
|
||||
check_outputs();
|
||||
}
|
||||
}
|
||||
@@ -1,20 +0,0 @@
|
||||
#===----------------------------------------------------------------------===##
|
||||
#
|
||||
# Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
# under the Apache License v2.0 with LLVM Exceptions.
|
||||
# See https://llvm.org/LICENSE.txt for license information.
|
||||
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
#
|
||||
#===----------------------------------------------------------------------===##
|
||||
|
||||
if (NOT cudax_ENABLE_NCCL)
|
||||
return()
|
||||
endif()
|
||||
|
||||
file(GLOB test_srcs LIST_DIRECTORIES FALSE CONFIGURE_DEPENDS *.cu *.cpp)
|
||||
|
||||
foreach (src IN LISTS test_srcs)
|
||||
cudax_add_multi_gpu_test("algorithms.reduce" test_target "${src}")
|
||||
target_link_libraries(${test_target} PRIVATE NCCL::nccl)
|
||||
endforeach()
|
||||
@@ -1,341 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/buffer>
|
||||
#include <cuda/functional>
|
||||
#include <cuda/memory_resource>
|
||||
#include <cuda/std/cstdint>
|
||||
#include <cuda/std/execution>
|
||||
#include <cuda/std/functional>
|
||||
#include <cuda/std/type_traits>
|
||||
|
||||
#include <cuda/experimental/__multi_gpu/algorithm/reduce/reduce.h>
|
||||
|
||||
#include <numeric>
|
||||
#include <vector>
|
||||
|
||||
#include <algorithm_common.h>
|
||||
#include <nccl_test_common.h>
|
||||
#include <testing.cuh>
|
||||
|
||||
namespace
|
||||
{
|
||||
struct custom_plus
|
||||
{
|
||||
template <class T>
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr T operator()(const T& lhs, const T& rhs) const
|
||||
{
|
||||
return lhs + rhs;
|
||||
}
|
||||
};
|
||||
|
||||
using custom_value = c2h::custom_type_t<c2h::accumulateable_t, c2h::less_comparable_t, c2h::equal_comparable_t>;
|
||||
using value_types = c2h::type_list<cuda::std::int32_t, float, custom_value>;
|
||||
using operators = c2h::type_list<::cuda::std::plus<>, ::cuda::maximum<>, custom_plus>;
|
||||
|
||||
static_assert(cudax::nccl_transportable<custom_value>);
|
||||
|
||||
template <typename T>
|
||||
T make_value(int i)
|
||||
{
|
||||
return static_cast<T>(i);
|
||||
}
|
||||
|
||||
template <>
|
||||
custom_value make_value<>(int i)
|
||||
{
|
||||
custom_value ret{};
|
||||
|
||||
ret.key = static_cast<std::size_t>(i);
|
||||
ret.val = static_cast<std::size_t>(i);
|
||||
return ret;
|
||||
};
|
||||
|
||||
template <class T, class Op>
|
||||
[[nodiscard]] T get_identity()
|
||||
{
|
||||
if constexpr (cuda::std::is_same_v<Op, cuda::std::plus<>> || cuda::std::is_same_v<Op, custom_plus>)
|
||||
{
|
||||
return make_value<T>(0);
|
||||
}
|
||||
else if constexpr (cuda::std::is_same_v<Op, cuda::maximum<>>)
|
||||
{
|
||||
return cuda::std::numeric_limits<T>::lowest();
|
||||
}
|
||||
else
|
||||
{
|
||||
static_assert(cuda::std::__always_false_v<T, Op>, "Add handling");
|
||||
}
|
||||
}
|
||||
|
||||
// Run the full reduction, wait for it to finish, and check that `reduce` left its argument ranges
|
||||
// untouched. This boilerplate is identical for every test regardless of how the inputs are shaped.
|
||||
template <class Env, class T, class Op>
|
||||
void do_reduce(cuda::std::span<cudax::nccl_communicator_ref> comms,
|
||||
const std::vector<Env>& envs,
|
||||
std::vector<cuda::device_buffer<T>>& in,
|
||||
std::vector<typename cuda::device_buffer<T>::iterator>& outputs,
|
||||
const T& init,
|
||||
const T& ident,
|
||||
Op op)
|
||||
{
|
||||
const auto envs_size = envs.size();
|
||||
const auto in_copy = in;
|
||||
const auto outputs_copy = outputs;
|
||||
|
||||
INFO("init = " << init);
|
||||
INFO("ident = " << ident);
|
||||
|
||||
cudax::reduce(cudax::broadcasted, comms, envs, in, outputs, init, op, ident);
|
||||
|
||||
// cuda::std::execution::env has no operator==, so we can only compare the sizes.
|
||||
REQUIRE(envs.size() == envs_size);
|
||||
// Reduction call should not modify the inputs in any ways
|
||||
REQUIRE(in.size() == in_copy.size());
|
||||
for (cuda::std::size_t i = 0; i < in.size(); ++i)
|
||||
{
|
||||
INFO("device = " << i);
|
||||
REQUIRE_THAT(in[i], Equals(in_copy[i]));
|
||||
}
|
||||
REQUIRE_THAT(outputs, Catch::Matchers::Equals(outputs_copy));
|
||||
}
|
||||
} // namespace
|
||||
|
||||
MULTI_GPU_TEST("reduce, one element per rank", value_types, operators)
|
||||
{
|
||||
using T = c2h::get<0, TestType>;
|
||||
using Op = c2h::get<1, TestType>;
|
||||
|
||||
// Seed each reduction with a few hardcoded initializers. The init participates in the fold the
|
||||
// same way on host and device, so any value works for every operator under test.
|
||||
const T init = make_value<T>(GENERATE(0, 1, -1, 5));
|
||||
const auto ident = get_identity<T, Op>();
|
||||
|
||||
auto comms = this->communicators();
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
// Global rank `comms[i].rank()` contributes the single value `rank`. Each local rank also gets a
|
||||
// one-element output buffer and an environment carrying its stream, so the reduction is
|
||||
// stream-ordered on the correct device. `reference` mirrors the contributions of every global
|
||||
// rank so we can fold them on the host exactly like `reduce` does on the device.
|
||||
std::vector<cuda::device_buffer<T>> in;
|
||||
std::vector<cuda::device_buffer<T>> out;
|
||||
std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
|
||||
|
||||
in.reserve(comms.size());
|
||||
out.reserve(comms.size());
|
||||
envs.reserve(comms.size());
|
||||
for (cuda::std::size_t i = 0; i < comms.size(); ++i)
|
||||
{
|
||||
const auto values = {make_value<T>(comms[i].rank())};
|
||||
in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
|
||||
out.emplace_back(
|
||||
cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), 1, cuda::no_init));
|
||||
envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
|
||||
}
|
||||
|
||||
auto outputs = make_output_iterators(out);
|
||||
|
||||
do_reduce(comms, envs, in, outputs, init, ident, Op{});
|
||||
|
||||
const T expected = [&] {
|
||||
std::vector<T> reference;
|
||||
|
||||
reference.reserve(comms.front().size());
|
||||
for (int r = 0; r < comms.front().size(); ++r)
|
||||
{
|
||||
reference.push_back(make_value<T>(r));
|
||||
}
|
||||
|
||||
return std::accumulate(reference.begin(), reference.end(), init, Op{});
|
||||
}();
|
||||
|
||||
for (const auto& buf : out)
|
||||
{
|
||||
const auto exp = cuda::make_buffer(buf.stream(), cuda::mr::legacy_pinned_memory_resource{}, 1, expected);
|
||||
|
||||
REQUIRE_THAT(buf, Equals(exp));
|
||||
}
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("reduce, multiple elements per rank", value_types, operators)
|
||||
{
|
||||
using T = c2h::get<0, TestType>;
|
||||
using Op = c2h::get<1, TestType>;
|
||||
|
||||
// Seed each reduction with a few hardcoded initializers. The init participates in the fold the
|
||||
// same way on host and device, so any value works for every operator under test.
|
||||
const T init = make_value<T>(GENERATE(0, 1, -1, 5));
|
||||
const auto ident = get_identity<T, Op>();
|
||||
|
||||
auto comms = this->communicators();
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
// Global rank `comms[i].rank()` contributes ten copies of `rank`. `reduce` first does a local
|
||||
// CUB reduction of each rank's range, then combines the partials across ranks. Each local rank
|
||||
// also gets a one-element output buffer and an environment carrying its stream. `reference`
|
||||
// mirrors every global rank's ten contributions for the host-side fold.
|
||||
std::vector<cuda::device_buffer<T>> in;
|
||||
std::vector<cuda::device_buffer<T>> out;
|
||||
std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
|
||||
|
||||
in.reserve(comms.size());
|
||||
out.reserve(comms.size());
|
||||
envs.reserve(comms.size());
|
||||
|
||||
constexpr auto values_per_rank = 10;
|
||||
for (cuda::std::size_t i = 0; i < comms.size(); ++i)
|
||||
{
|
||||
const auto v = make_value<T>(comms[i].rank());
|
||||
const std::vector<T> values(values_per_rank, v);
|
||||
|
||||
in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
|
||||
out.emplace_back(
|
||||
cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), 1, cuda::no_init));
|
||||
envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
|
||||
}
|
||||
|
||||
auto outputs = make_output_iterators(out);
|
||||
|
||||
do_reduce(comms, envs, in, outputs, init, ident, Op{});
|
||||
|
||||
const T expected = [&] {
|
||||
std::vector<T> reference;
|
||||
|
||||
reference.reserve(comms.front().size() * values_per_rank);
|
||||
for (int r = 0; r < comms.front().size(); ++r)
|
||||
{
|
||||
const auto v = make_value<T>(r);
|
||||
|
||||
reference.insert(reference.end(), values_per_rank, v);
|
||||
}
|
||||
|
||||
return std::accumulate(reference.begin(), reference.end(), init, Op{});
|
||||
}();
|
||||
|
||||
for (const auto& buf : out)
|
||||
{
|
||||
const auto exp = cuda::make_buffer(buf.stream(), cuda::mr::legacy_pinned_memory_resource{}, 1, expected);
|
||||
|
||||
REQUIRE_THAT(buf, Equals(exp));
|
||||
}
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("reduce, some ranks empty", value_types, operators)
|
||||
{
|
||||
using T = c2h::get<0, TestType>;
|
||||
using Op = c2h::get<1, TestType>;
|
||||
|
||||
const T init = make_value<T>(GENERATE(0, 1, -1, 5));
|
||||
const auto ident = get_identity<T, Op>();
|
||||
|
||||
auto comms = this->communicators();
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
// Even global ranks contribute ten copies of `rank`; odd global ranks contribute an empty input
|
||||
// range. Rank 0 (the reduction root) is always non-empty. `reduce` must treat an empty rank as
|
||||
// contributing nothing, exactly like `std::accumulate` over the surviving elements. `reference`
|
||||
// mirrors that for the host-side fold.
|
||||
std::vector<cuda::device_buffer<T>> in;
|
||||
std::vector<cuda::device_buffer<T>> out;
|
||||
std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
|
||||
|
||||
in.reserve(comms.size());
|
||||
out.reserve(comms.size());
|
||||
envs.reserve(comms.size());
|
||||
|
||||
constexpr auto values_per_rank = 10;
|
||||
for (cuda::std::size_t i = 0; i < comms.size(); ++i)
|
||||
{
|
||||
const auto rank = comms[i].rank();
|
||||
if (rank % 2 == 0)
|
||||
{
|
||||
const std::vector<T> values(values_per_rank, make_value<T>(rank));
|
||||
in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
|
||||
}
|
||||
else
|
||||
{
|
||||
in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device()));
|
||||
}
|
||||
out.emplace_back(
|
||||
cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), 1, cuda::no_init));
|
||||
envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
|
||||
}
|
||||
|
||||
auto outputs = make_output_iterators(out);
|
||||
|
||||
do_reduce(comms, envs, in, outputs, init, ident, Op{});
|
||||
|
||||
const T expected = [&] {
|
||||
std::vector<T> reference;
|
||||
|
||||
reference.reserve(comms.front().size() * values_per_rank);
|
||||
for (int r = 0; r < comms.front().size(); ++r)
|
||||
{
|
||||
if (r % 2 == 0)
|
||||
{
|
||||
reference.insert(reference.end(), values_per_rank, make_value<T>(r));
|
||||
}
|
||||
}
|
||||
|
||||
return std::accumulate(reference.begin(), reference.end(), init, Op{});
|
||||
}();
|
||||
|
||||
for (const auto& buf : out)
|
||||
{
|
||||
const auto exp = cuda::make_buffer(buf.stream(), cuda::mr::legacy_pinned_memory_resource{}, 1, expected);
|
||||
|
||||
REQUIRE_THAT(buf, Equals(exp));
|
||||
}
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("reduce, all ranks empty", value_types, operators)
|
||||
{
|
||||
using T = c2h::get<0, TestType>;
|
||||
using Op = c2h::get<1, TestType>;
|
||||
|
||||
const auto init = make_value<T>(GENERATE(0, 1, -1, 5));
|
||||
const auto ident = get_identity<T, Op>();
|
||||
|
||||
auto comms = this->communicators();
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
// No rank contributes any element. Reducing nothing seeded by `init` is just `init`, so every
|
||||
// output must equal `init` regardless of the operator.
|
||||
std::vector<cuda::device_buffer<T>> in;
|
||||
std::vector<cuda::device_buffer<T>> out;
|
||||
std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
|
||||
|
||||
in.reserve(comms.size());
|
||||
out.reserve(comms.size());
|
||||
envs.reserve(comms.size());
|
||||
for (cuda::std::size_t i = 0; i < comms.size(); ++i)
|
||||
{
|
||||
in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device()));
|
||||
out.emplace_back(
|
||||
cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), 1, cuda::no_init));
|
||||
envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
|
||||
}
|
||||
|
||||
auto outputs = make_output_iterators(out);
|
||||
|
||||
do_reduce(comms, envs, in, outputs, init, ident, Op{});
|
||||
|
||||
// Reducing nothing seeded by `init` yields `init`, exactly like `std::accumulate` over an empty
|
||||
// range.
|
||||
const T expected = init;
|
||||
|
||||
for (const auto& buf : out)
|
||||
{
|
||||
const auto exp = cuda::make_buffer(buf.stream(), cuda::mr::legacy_pinned_memory_resource{}, 1, expected);
|
||||
|
||||
REQUIRE_THAT(buf, Equals(exp));
|
||||
}
|
||||
}
|
||||
@@ -1,112 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/buffer>
|
||||
#include <cuda/memory_resource>
|
||||
#include <cuda/std/cstdint>
|
||||
#include <cuda/std/execution>
|
||||
#include <cuda/std/type_traits>
|
||||
|
||||
#include <cuda/experimental/__multi_gpu/algorithm/reduce/reduce.h>
|
||||
|
||||
#include <numeric>
|
||||
#include <vector>
|
||||
|
||||
#include <algorithm_common.h>
|
||||
#include <nccl_test_common.h>
|
||||
#include <testing.cuh>
|
||||
|
||||
MULTI_GPU_TEST("reduce, range overloads default values", )
|
||||
{
|
||||
using T = cuda::std::int32_t;
|
||||
using Op = ::cuda::std::plus<>;
|
||||
|
||||
constexpr auto init = T{};
|
||||
constexpr T ident = cuda::identity_element<Op, T>();
|
||||
constexpr auto op = Op{};
|
||||
|
||||
auto comms = this->communicators();
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
std::vector<cuda::device_buffer<T>> in;
|
||||
std::vector<cuda::device_buffer<T>> out;
|
||||
std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
|
||||
|
||||
in.reserve(comms.size());
|
||||
out.reserve(comms.size());
|
||||
envs.reserve(comms.size());
|
||||
|
||||
constexpr auto values_per_rank = 10;
|
||||
for (cuda::std::size_t i = 0; i < comms.size(); ++i)
|
||||
{
|
||||
const std::vector<T> values(values_per_rank, static_cast<T>(comms[i].rank()));
|
||||
|
||||
in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
|
||||
out.emplace_back(
|
||||
cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), 1, cuda::no_init));
|
||||
envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
|
||||
}
|
||||
|
||||
auto outputs = make_output_iterators(out);
|
||||
|
||||
const auto expected = [&] {
|
||||
std::vector<T> reference;
|
||||
|
||||
reference.reserve(comms.front().size() * values_per_rank);
|
||||
for (int r = 0; r < comms.front().size(); ++r)
|
||||
{
|
||||
reference.insert(reference.end(), values_per_rank, static_cast<T>(r));
|
||||
}
|
||||
|
||||
const auto val = std::accumulate(reference.begin(), reference.end(), init, op);
|
||||
|
||||
return cuda::make_buffer(cuda::stream_ref{::CUstream{}}, cuda::mr::legacy_pinned_memory_resource{}, 1, val);
|
||||
}();
|
||||
|
||||
SECTION("Default init, op, ident (all)")
|
||||
{
|
||||
cudax::reduce(cudax::broadcasted, comms, envs, in, outputs);
|
||||
|
||||
for (const auto& buf : out)
|
||||
{
|
||||
REQUIRE_THAT(buf, Equals(expected));
|
||||
}
|
||||
}
|
||||
|
||||
SECTION("Default op, ident")
|
||||
{
|
||||
cudax::reduce(cudax::broadcasted, comms, envs, in, outputs, init);
|
||||
|
||||
for (const auto& buf : out)
|
||||
{
|
||||
REQUIRE_THAT(buf, Equals(expected));
|
||||
}
|
||||
}
|
||||
|
||||
SECTION("Default ident")
|
||||
{
|
||||
cudax::reduce(cudax::broadcasted, comms, envs, in, outputs, init, op);
|
||||
|
||||
for (const auto& buf : out)
|
||||
{
|
||||
REQUIRE_THAT(buf, Equals(expected));
|
||||
}
|
||||
}
|
||||
|
||||
SECTION("Default none")
|
||||
{
|
||||
cudax::reduce(cudax::broadcasted, comms, envs, in, outputs, init, op, ident);
|
||||
|
||||
for (const auto& buf : out)
|
||||
{
|
||||
REQUIRE_THAT(buf, Equals(expected));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,347 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/buffer>
|
||||
#include <cuda/functional>
|
||||
#include <cuda/memory_resource>
|
||||
#include <cuda/std/cstdint>
|
||||
#include <cuda/std/execution>
|
||||
#include <cuda/std/functional>
|
||||
#include <cuda/std/type_traits>
|
||||
|
||||
#include <cuda/experimental/__multi_gpu/algorithm/reduce/reduce.h>
|
||||
|
||||
#include <exception>
|
||||
#include <future>
|
||||
#include <numeric>
|
||||
#include <vector>
|
||||
|
||||
#include <algorithm_common.h>
|
||||
#include <nccl_test_common.h>
|
||||
#include <testing.cuh>
|
||||
|
||||
namespace
|
||||
{
|
||||
struct custom_plus
|
||||
{
|
||||
template <class T>
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr T operator()(const T& lhs, const T& rhs) const
|
||||
{
|
||||
return lhs + rhs;
|
||||
}
|
||||
};
|
||||
|
||||
using custom_value = c2h::custom_type_t<c2h::accumulateable_t, c2h::less_comparable_t, c2h::equal_comparable_t>;
|
||||
using value_types = c2h::type_list<cuda::std::int32_t, float, custom_value>;
|
||||
using operators = c2h::type_list<::cuda::std::plus<>, ::cuda::maximum<>, custom_plus>;
|
||||
|
||||
static_assert(cudax::nccl_transportable<custom_value>);
|
||||
|
||||
template <typename T>
|
||||
T make_value(int i)
|
||||
{
|
||||
return static_cast<T>(i);
|
||||
}
|
||||
|
||||
template <>
|
||||
custom_value make_value<>(int i)
|
||||
{
|
||||
custom_value ret{};
|
||||
|
||||
ret.key = static_cast<std::size_t>(i);
|
||||
ret.val = static_cast<std::size_t>(i);
|
||||
return ret;
|
||||
};
|
||||
|
||||
template <class T, class Op>
|
||||
[[nodiscard]] T get_identity()
|
||||
{
|
||||
if constexpr (cuda::std::is_same_v<Op, cuda::std::plus<>> || cuda::std::is_same_v<Op, custom_plus>)
|
||||
{
|
||||
return make_value<T>(0);
|
||||
}
|
||||
else if constexpr (cuda::std::is_same_v<Op, cuda::maximum<>>)
|
||||
{
|
||||
return cuda::std::numeric_limits<T>::lowest();
|
||||
}
|
||||
else
|
||||
{
|
||||
static_assert(cuda::std::__always_false_v<T, Op>, "Add handling");
|
||||
}
|
||||
}
|
||||
|
||||
// Drive the reduction through the single-communicator overload of `reduce`, one thread per
|
||||
// rank. That overload opens its own NCCL group on a single communicator, so issuing the
|
||||
// per-rank calls serially on one thread would deadlock at `ncclGroupEnd`. Running each rank on
|
||||
// its own thread lets the per-thread groups rendezvous across ranks. Only the `reduce` call
|
||||
// happens on the worker threads; every Catch2 assertion runs on the main thread after the
|
||||
// join, since the assertion macros are not safe to fire concurrently.
|
||||
template <class Env, class T, class Op>
|
||||
void do_reduce_threaded(
|
||||
cuda::std::span<cudax::nccl_communicator_ref> comms,
|
||||
std::vector<Env>& envs,
|
||||
std::vector<cuda::device_buffer<T>>& in,
|
||||
std::vector<typename cuda::device_buffer<T>::iterator>& outputs,
|
||||
const T& init,
|
||||
const T& ident,
|
||||
Op op)
|
||||
{
|
||||
const auto in_copy = in;
|
||||
const auto outputs_copy = outputs;
|
||||
|
||||
INFO("init = " << init);
|
||||
INFO("ident = " << ident);
|
||||
|
||||
run_threaded(comms.size(), [&](cuda::std::size_t i) {
|
||||
cudax::reduce(cudax::broadcasted, comms[i], envs[i], in[i], outputs[i], init, op, ident);
|
||||
});
|
||||
|
||||
// Reduction call should not modify the inputs in any ways
|
||||
REQUIRE(in.size() == in_copy.size());
|
||||
for (cuda::std::size_t i = 0; i < in.size(); ++i)
|
||||
{
|
||||
INFO("device = " << i);
|
||||
REQUIRE_THAT(in[i], Equals(in_copy[i]));
|
||||
}
|
||||
REQUIRE_THAT(outputs, Catch::Matchers::Equals(outputs_copy));
|
||||
}
|
||||
} // namespace
|
||||
|
||||
MULTI_GPU_TEST("reduce single-comm, one element per rank", value_types, operators)
|
||||
{
|
||||
using T = c2h::get<0, TestType>;
|
||||
using Op = c2h::get<1, TestType>;
|
||||
|
||||
// Seed each reduction with a few hardcoded initializers. The init participates in the fold the
|
||||
// same way on host and device, so any value works for every operator under test.
|
||||
const T init = make_value<T>(GENERATE(0, 1, -1, 5));
|
||||
const auto ident = get_identity<T, Op>();
|
||||
|
||||
auto comms = this->communicators();
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
// Global rank `comms[i].rank()` contributes the single value `rank`. Each local rank also gets a
|
||||
// one-element output buffer and an environment carrying its stream, so the reduction is
|
||||
// stream-ordered on the correct device. `reference` mirrors the contributions of every global
|
||||
// rank so we can fold them on the host exactly like `reduce` does on the device.
|
||||
std::vector<cuda::device_buffer<T>> in;
|
||||
std::vector<cuda::device_buffer<T>> out;
|
||||
std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
|
||||
|
||||
in.reserve(comms.size());
|
||||
out.reserve(comms.size());
|
||||
envs.reserve(comms.size());
|
||||
for (cuda::std::size_t i = 0; i < comms.size(); ++i)
|
||||
{
|
||||
const auto values = {make_value<T>(comms[i].rank())};
|
||||
in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
|
||||
out.emplace_back(
|
||||
cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), 1, cuda::no_init));
|
||||
envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
|
||||
}
|
||||
|
||||
auto outputs = make_output_iterators(out);
|
||||
|
||||
do_reduce_threaded(comms, envs, in, outputs, init, ident, Op{});
|
||||
|
||||
const T expected = [&] {
|
||||
std::vector<T> reference;
|
||||
|
||||
reference.reserve(comms.front().size());
|
||||
for (int r = 0; r < comms.front().size(); ++r)
|
||||
{
|
||||
reference.push_back(make_value<T>(r));
|
||||
}
|
||||
|
||||
return std::accumulate(reference.begin(), reference.end(), init, Op{});
|
||||
}();
|
||||
|
||||
for (const auto& buf : out)
|
||||
{
|
||||
const auto exp = cuda::make_buffer(buf.stream(), cuda::mr::legacy_pinned_memory_resource{}, 1, expected);
|
||||
|
||||
REQUIRE_THAT(buf, Equals(exp));
|
||||
}
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("reduce single-comm, multiple elements per rank", value_types, operators)
|
||||
{
|
||||
using T = c2h::get<0, TestType>;
|
||||
using Op = c2h::get<1, TestType>;
|
||||
|
||||
// Seed each reduction with a few hardcoded initializers. The init participates in the fold the
|
||||
// same way on host and device, so any value works for every operator under test.
|
||||
const T init = make_value<T>(GENERATE(0, 1, -1, 5));
|
||||
const auto ident = get_identity<T, Op>();
|
||||
|
||||
auto comms = this->communicators();
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
// Global rank `comms[i].rank()` contributes ten copies of `rank`. `reduce` first does a local
|
||||
// CUB reduction of each rank's range, then combines the partials across ranks. Each local rank
|
||||
// also gets a one-element output buffer and an environment carrying its stream. `reference`
|
||||
// mirrors every global rank's ten contributions for the host-side fold.
|
||||
std::vector<cuda::device_buffer<T>> in;
|
||||
std::vector<cuda::device_buffer<T>> out;
|
||||
std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
|
||||
|
||||
in.reserve(comms.size());
|
||||
out.reserve(comms.size());
|
||||
envs.reserve(comms.size());
|
||||
|
||||
constexpr auto values_per_rank = 10;
|
||||
for (cuda::std::size_t i = 0; i < comms.size(); ++i)
|
||||
{
|
||||
const auto v = make_value<T>(comms[i].rank());
|
||||
const std::vector<T> values(values_per_rank, v);
|
||||
|
||||
in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
|
||||
out.emplace_back(
|
||||
cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), 1, cuda::no_init));
|
||||
envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
|
||||
}
|
||||
|
||||
auto outputs = make_output_iterators(out);
|
||||
|
||||
do_reduce_threaded(comms, envs, in, outputs, init, ident, Op{});
|
||||
|
||||
const T expected = [&] {
|
||||
std::vector<T> reference;
|
||||
|
||||
reference.reserve(comms.front().size() * values_per_rank);
|
||||
for (int r = 0; r < comms.front().size(); ++r)
|
||||
{
|
||||
const auto v = make_value<T>(r);
|
||||
|
||||
reference.insert(reference.end(), values_per_rank, v);
|
||||
}
|
||||
|
||||
return std::accumulate(reference.begin(), reference.end(), init, Op{});
|
||||
}();
|
||||
|
||||
for (const auto& buf : out)
|
||||
{
|
||||
const auto exp = cuda::make_buffer(buf.stream(), cuda::mr::legacy_pinned_memory_resource{}, 1, expected);
|
||||
|
||||
REQUIRE_THAT(buf, Equals(exp));
|
||||
}
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("reduce single-comm, some ranks empty", value_types, operators)
|
||||
{
|
||||
using T = c2h::get<0, TestType>;
|
||||
using Op = c2h::get<1, TestType>;
|
||||
|
||||
const T init = make_value<T>(GENERATE(0, 1, -1, 5));
|
||||
const auto ident = get_identity<T, Op>();
|
||||
|
||||
auto comms = this->communicators();
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
// Even global ranks contribute ten copies of `rank`; odd global ranks contribute an empty input
|
||||
// range. Rank 0 (the reduction root) is always non-empty. `reduce` must treat an empty rank as
|
||||
// contributing nothing, exactly like `std::accumulate` over the surviving elements. `reference`
|
||||
// mirrors that for the host-side fold.
|
||||
std::vector<cuda::device_buffer<T>> in;
|
||||
std::vector<cuda::device_buffer<T>> out;
|
||||
std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
|
||||
|
||||
in.reserve(comms.size());
|
||||
out.reserve(comms.size());
|
||||
envs.reserve(comms.size());
|
||||
|
||||
constexpr auto values_per_rank = 10;
|
||||
for (cuda::std::size_t i = 0; i < comms.size(); ++i)
|
||||
{
|
||||
const auto rank = comms[i].rank();
|
||||
if (rank % 2 == 0)
|
||||
{
|
||||
const std::vector<T> values(values_per_rank, make_value<T>(rank));
|
||||
in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
|
||||
}
|
||||
else
|
||||
{
|
||||
in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device()));
|
||||
}
|
||||
out.emplace_back(
|
||||
cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), 1, cuda::no_init));
|
||||
envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
|
||||
}
|
||||
|
||||
auto outputs = make_output_iterators(out);
|
||||
|
||||
do_reduce_threaded(comms, envs, in, outputs, init, ident, Op{});
|
||||
|
||||
const T expected = [&] {
|
||||
std::vector<T> reference;
|
||||
|
||||
reference.reserve(comms.front().size() * values_per_rank);
|
||||
for (int r = 0; r < comms.front().size(); ++r)
|
||||
{
|
||||
if (r % 2 == 0)
|
||||
{
|
||||
reference.insert(reference.end(), values_per_rank, make_value<T>(r));
|
||||
}
|
||||
}
|
||||
|
||||
return std::accumulate(reference.begin(), reference.end(), init, Op{});
|
||||
}();
|
||||
|
||||
for (const auto& buf : out)
|
||||
{
|
||||
const auto exp = cuda::make_buffer(buf.stream(), cuda::mr::legacy_pinned_memory_resource{}, 1, expected);
|
||||
|
||||
REQUIRE_THAT(buf, Equals(exp));
|
||||
}
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("reduce single-comm, all ranks empty", value_types, operators)
|
||||
{
|
||||
using T = c2h::get<0, TestType>;
|
||||
using Op = c2h::get<1, TestType>;
|
||||
|
||||
const auto init = make_value<T>(GENERATE(0, 1, -1, 5));
|
||||
const auto ident = get_identity<T, Op>();
|
||||
|
||||
auto comms = this->communicators();
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
// No rank contributes any element. Reducing nothing seeded by `init` is just `init`, so every
|
||||
// output must equal `init` regardless of the operator.
|
||||
std::vector<cuda::device_buffer<T>> in;
|
||||
std::vector<cuda::device_buffer<T>> out;
|
||||
std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
|
||||
|
||||
in.reserve(comms.size());
|
||||
out.reserve(comms.size());
|
||||
envs.reserve(comms.size());
|
||||
for (cuda::std::size_t i = 0; i < comms.size(); ++i)
|
||||
{
|
||||
in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device()));
|
||||
out.emplace_back(
|
||||
cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), 1, cuda::no_init));
|
||||
envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
|
||||
}
|
||||
|
||||
auto outputs = make_output_iterators(out);
|
||||
|
||||
do_reduce_threaded(comms, envs, in, outputs, init, ident, Op{});
|
||||
|
||||
// Reducing nothing seeded by `init` yields `init`, exactly like `std::accumulate` over an empty
|
||||
// range.
|
||||
const T expected = init;
|
||||
|
||||
for (const auto& buf : out)
|
||||
{
|
||||
const auto exp = cuda::make_buffer(buf.stream(), cuda::mr::legacy_pinned_memory_resource{}, 1, expected);
|
||||
|
||||
REQUIRE_THAT(buf, Equals(exp));
|
||||
}
|
||||
}
|
||||
@@ -1,122 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/buffer>
|
||||
#include <cuda/memory_resource>
|
||||
#include <cuda/std/cstdint>
|
||||
#include <cuda/std/execution>
|
||||
#include <cuda/std/type_traits>
|
||||
|
||||
#include <cuda/experimental/__multi_gpu/algorithm/reduce/reduce.h>
|
||||
|
||||
#include <exception>
|
||||
#include <future>
|
||||
#include <numeric>
|
||||
#include <vector>
|
||||
|
||||
#include <algorithm_common.h>
|
||||
#include <nccl_test_common.h>
|
||||
#include <testing.cuh>
|
||||
|
||||
MULTI_GPU_TEST("reduce single-comm, overloads default values", )
|
||||
{
|
||||
using T = cuda::std::int32_t;
|
||||
using Op = ::cuda::std::plus<>;
|
||||
|
||||
constexpr auto init = T{};
|
||||
constexpr T ident = cuda::identity_element<Op, T>();
|
||||
constexpr auto op = Op{};
|
||||
|
||||
auto comms = this->communicators();
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
std::vector<cuda::device_buffer<T>> in;
|
||||
std::vector<cuda::device_buffer<T>> out;
|
||||
std::vector<decltype(::cuda::std::execution::env{::cuda::stream_ref{streams[0]}})> envs;
|
||||
|
||||
in.reserve(comms.size());
|
||||
out.reserve(comms.size());
|
||||
envs.reserve(comms.size());
|
||||
|
||||
constexpr auto values_per_rank = 10;
|
||||
for (cuda::std::size_t i = 0; i < comms.size(); ++i)
|
||||
{
|
||||
const std::vector<T> values(values_per_rank, static_cast<T>(comms[i].rank()));
|
||||
|
||||
in.emplace_back(cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), values));
|
||||
out.emplace_back(
|
||||
cuda::make_device_buffer<T>(streams[i], comms[i].logical_device().underlying_device(), 1, cuda::no_init));
|
||||
envs.emplace_back(::cuda::std::execution::env{::cuda::stream_ref{streams[i]}});
|
||||
}
|
||||
|
||||
auto outputs = make_output_iterators(out);
|
||||
|
||||
const auto expected = [&] {
|
||||
std::vector<T> reference;
|
||||
|
||||
reference.reserve(comms.front().size() * values_per_rank);
|
||||
for (int r = 0; r < comms.front().size(); ++r)
|
||||
{
|
||||
reference.insert(reference.end(), values_per_rank, static_cast<T>(r));
|
||||
}
|
||||
|
||||
const auto val = std::accumulate(reference.begin(), reference.end(), init, op);
|
||||
|
||||
return cuda::make_buffer(cuda::stream_ref{::CUstream{}}, cuda::mr::legacy_pinned_memory_resource{}, 1, val);
|
||||
}();
|
||||
|
||||
SECTION("Default init, op, ident (all)")
|
||||
{
|
||||
run_threaded(comms.size(), [&](cuda::std::size_t i) {
|
||||
cudax::reduce(cudax::broadcasted, comms[i], envs[i], in[i], outputs[i]);
|
||||
});
|
||||
|
||||
for (const auto& buf : out)
|
||||
{
|
||||
REQUIRE_THAT(buf, Equals(expected));
|
||||
}
|
||||
}
|
||||
|
||||
SECTION("Default op, ident")
|
||||
{
|
||||
run_threaded(comms.size(), [&](cuda::std::size_t i) {
|
||||
cudax::reduce(cudax::broadcasted, comms[i], envs[i], in[i], outputs[i], init);
|
||||
});
|
||||
|
||||
for (const auto& buf : out)
|
||||
{
|
||||
REQUIRE_THAT(buf, Equals(expected));
|
||||
}
|
||||
}
|
||||
|
||||
SECTION("Default ident")
|
||||
{
|
||||
run_threaded(comms.size(), [&](cuda::std::size_t i) {
|
||||
cudax::reduce(cudax::broadcasted, comms[i], envs[i], in[i], outputs[i], init, op);
|
||||
});
|
||||
|
||||
for (const auto& buf : out)
|
||||
{
|
||||
REQUIRE_THAT(buf, Equals(expected));
|
||||
}
|
||||
}
|
||||
|
||||
SECTION("Default none")
|
||||
{
|
||||
run_threaded(comms.size(), [&](cuda::std::size_t i) {
|
||||
cudax::reduce(cudax::broadcasted, comms[i], envs[i], in[i], outputs[i], init, op, ident);
|
||||
});
|
||||
|
||||
for (const auto& buf : out)
|
||||
{
|
||||
REQUIRE_THAT(buf, Equals(expected));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,11 +0,0 @@
|
||||
#===----------------------------------------------------------------------===##
|
||||
#
|
||||
# Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
# under the Apache License v2.0 with LLVM Exceptions.
|
||||
# See https://llvm.org/LICENSE.txt for license information.
|
||||
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
#
|
||||
#===----------------------------------------------------------------------===##
|
||||
|
||||
add_subdirectory(nccl)
|
||||
@@ -1,20 +0,0 @@
|
||||
#===----------------------------------------------------------------------===##
|
||||
#
|
||||
# Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
# under the Apache License v2.0 with LLVM Exceptions.
|
||||
# See https://llvm.org/LICENSE.txt for license information.
|
||||
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
#
|
||||
#===----------------------------------------------------------------------===##
|
||||
|
||||
if (NOT cudax_ENABLE_NCCL)
|
||||
return()
|
||||
endif()
|
||||
|
||||
file(GLOB test_srcs LIST_DIRECTORIES FALSE CONFIGURE_DEPENDS *.cu *.cpp)
|
||||
|
||||
foreach (src IN LISTS test_srcs)
|
||||
cudax_add_multi_gpu_test("communicators.nccl" test_target "${src}")
|
||||
target_link_libraries(${test_target} PRIVATE NCCL::nccl)
|
||||
endforeach()
|
||||
@@ -1,223 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/devices>
|
||||
#include <cuda/std/type_traits>
|
||||
#include <cuda/std/utility>
|
||||
|
||||
#include <cuda/experimental/__multi_gpu/nccl_communicator.h>
|
||||
#include <cuda/experimental/__multi_gpu/nccl_communicator_ref.h>
|
||||
|
||||
#include <nccl.h>
|
||||
#include <nccl_test_common.h>
|
||||
|
||||
namespace
|
||||
{
|
||||
[[nodiscard]] ncclComm_t make_nccl_communicator_handle()
|
||||
{
|
||||
if (cuda::devices.size() == 0)
|
||||
{
|
||||
SKIP("No CUDA devices visible");
|
||||
}
|
||||
|
||||
const int device = cuda::devices[0].get();
|
||||
ncclComm_t handle{};
|
||||
|
||||
const ncclResult_t result = ncclCommInitAll(&handle, 1, &device);
|
||||
|
||||
INFO("NCCL: " << ncclGetErrorString(result));
|
||||
REQUIRE(result == ncclSuccess);
|
||||
|
||||
return handle;
|
||||
}
|
||||
} // namespace
|
||||
|
||||
C2H_TEST("nccl_communicator_ref typedefs", "[multi_gpu]")
|
||||
{
|
||||
STATIC_REQUIRE(::cuda::std::is_same_v<cudax::nccl_communicator_ref::native_handle_type, ncclComm_t>);
|
||||
STATIC_REQUIRE(
|
||||
::cuda::std::is_same_v<cudax::nccl_communicator_ref::group_guard_type,
|
||||
decltype(::cuda::std::declval<const cudax::nccl_communicator_ref&>().group_guard())>);
|
||||
}
|
||||
|
||||
C2H_TEST("nccl_communicator(s) not constructible from NCCL_COMM_NULL", "[multi_gpu]")
|
||||
{
|
||||
SECTION("ref")
|
||||
{
|
||||
STATIC_REQUIRE(!::cuda::std::is_constructible_v<cudax::nccl_communicator_ref, decltype(NCCL_COMM_NULL)>);
|
||||
STATIC_REQUIRE(!::cuda::std::is_constructible_v<cudax::nccl_communicator_ref, cuda::std::nullptr_t>);
|
||||
}
|
||||
|
||||
SECTION("owning")
|
||||
{
|
||||
STATIC_REQUIRE(!::cuda::std::is_constructible_v<cudax::nccl_communicator, decltype(NCCL_COMM_NULL)>);
|
||||
STATIC_REQUIRE(!::cuda::std::is_constructible_v<cudax::nccl_communicator, cuda::std::nullptr_t>);
|
||||
STATIC_REQUIRE(!::cuda::std::is_constructible_v<cudax::nccl_communicator, ncclComm_t>);
|
||||
}
|
||||
}
|
||||
|
||||
C2H_TEST("nccl_communicator basic", "[multi_gpu][nccl]")
|
||||
{
|
||||
SECTION("is move-only")
|
||||
{
|
||||
STATIC_REQUIRE(!cuda::std::is_copy_constructible_v<cudax::nccl_communicator>);
|
||||
STATIC_REQUIRE(!cuda::std::is_copy_assignable_v<cudax::nccl_communicator>);
|
||||
STATIC_REQUIRE(cuda::std::is_move_constructible_v<cudax::nccl_communicator>);
|
||||
STATIC_REQUIRE(cuda::std::is_nothrow_move_constructible_v<cudax::nccl_communicator>);
|
||||
STATIC_REQUIRE(cuda::std::is_move_assignable_v<cudax::nccl_communicator>);
|
||||
STATIC_REQUIRE(cuda::std::is_nothrow_move_assignable_v<cudax::nccl_communicator>);
|
||||
}
|
||||
|
||||
SECTION("factory construction")
|
||||
{
|
||||
STATIC_REQUIRE(
|
||||
cuda::std::is_same_v<decltype(cudax::nccl_communicator::from_native_handle(cuda::std::declval<ncclComm_t>())),
|
||||
cudax::nccl_communicator>);
|
||||
|
||||
//! [nccl_communicator_construction]
|
||||
const ncclComm_t handle = make_nccl_communicator_handle();
|
||||
|
||||
auto comm = cuda::experimental::nccl_communicator::from_native_handle(handle);
|
||||
|
||||
// comm owns the handle now
|
||||
REQUIRE(comm.native_handle() == handle);
|
||||
//! [nccl_communicator_construction]
|
||||
}
|
||||
|
||||
SECTION("factory construction with logical device")
|
||||
{
|
||||
STATIC_REQUIRE(
|
||||
cuda::std::is_same_v<decltype(cudax::nccl_communicator::from_native_handle(
|
||||
cuda::std::declval<ncclComm_t>(), cuda::std::declval<cudax::logical_device>())),
|
||||
cudax::nccl_communicator>);
|
||||
|
||||
//! [nccl_communicator_construction_with_logical_device]
|
||||
const ncclComm_t handle = make_nccl_communicator_handle();
|
||||
const auto device = cudax::logical_device{cuda::devices[0]};
|
||||
|
||||
auto comm = cudax::nccl_communicator::from_native_handle(handle, device);
|
||||
|
||||
REQUIRE(comm.native_handle() == handle);
|
||||
REQUIRE(comm.logical_device() == device);
|
||||
//! [nccl_communicator_construction_with_logical_device]
|
||||
}
|
||||
|
||||
SECTION("no_init construction")
|
||||
{
|
||||
STATIC_REQUIRE(cuda::std::is_nothrow_constructible_v<cudax::nccl_communicator, cuda::no_init_t>);
|
||||
|
||||
//! [nccl_communicator_no_init_construction]
|
||||
const auto comm = cudax::nccl_communicator{cuda::no_init};
|
||||
|
||||
REQUIRE(comm.native_handle() == ncclComm_t{NCCL_COMM_NULL});
|
||||
//! [nccl_communicator_no_init_construction]
|
||||
|
||||
REQUIRE(comm.rank() == 0);
|
||||
REQUIRE(comm.size() == 0);
|
||||
}
|
||||
|
||||
SECTION("release")
|
||||
{
|
||||
//! [nccl_communicator_release]
|
||||
const ncclComm_t handle = make_nccl_communicator_handle();
|
||||
|
||||
auto comm = cuda::experimental::nccl_communicator::from_native_handle(handle);
|
||||
|
||||
const auto released_handle = comm.release();
|
||||
|
||||
// comm contains the null handle after release
|
||||
REQUIRE(comm.native_handle() == ncclComm_t{NCCL_COMM_NULL});
|
||||
REQUIRE(released_handle == handle);
|
||||
//! [nccl_communicator_release]
|
||||
|
||||
// so that we clean up properly
|
||||
[[maybe_unused]] const auto _ = cudax::nccl_communicator::from_native_handle(handle);
|
||||
}
|
||||
|
||||
SECTION("move construction")
|
||||
{
|
||||
//! [nccl_communicator_move_construction]
|
||||
const ncclComm_t handle = make_nccl_communicator_handle();
|
||||
|
||||
auto source = cudax::nccl_communicator::from_native_handle(handle);
|
||||
auto destination = cudax::nccl_communicator{cuda::std::move(source)};
|
||||
|
||||
// moved-from communicator is now invalid
|
||||
REQUIRE(source.native_handle() == ncclComm_t{NCCL_COMM_NULL});
|
||||
REQUIRE(destination.native_handle() == handle);
|
||||
//! [nccl_communicator_move_construction]
|
||||
}
|
||||
|
||||
SECTION("move assignment")
|
||||
{
|
||||
//! [nccl_communicator_move_assignment]
|
||||
auto source = cuda::experimental::nccl_communicator::from_native_handle(make_nccl_communicator_handle());
|
||||
auto destination = cuda::experimental::nccl_communicator::from_native_handle(make_nccl_communicator_handle());
|
||||
|
||||
// Save the native handle to verify that ownership is transferred.
|
||||
const auto handle = source.native_handle();
|
||||
|
||||
destination = cuda::std::move(source);
|
||||
|
||||
REQUIRE(source.native_handle() == ncclComm_t{NCCL_COMM_NULL});
|
||||
REQUIRE(destination.native_handle() == handle);
|
||||
//! [nccl_communicator_move_assignment]
|
||||
}
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("nccl_communicator_ref basic", )
|
||||
{
|
||||
SECTION("rank and size")
|
||||
{
|
||||
int i = 0;
|
||||
|
||||
for (auto& comm : this->communicators())
|
||||
{
|
||||
REQUIRE(comm.rank() == i);
|
||||
REQUIRE(comm.size() == static_cast<int>(cuda::devices.size()));
|
||||
++i;
|
||||
}
|
||||
}
|
||||
|
||||
SECTION("native handle")
|
||||
{
|
||||
for (auto& comm : this->communicators())
|
||||
{
|
||||
REQUIRE(comm.native_handle() != NCCL_COMM_NULL);
|
||||
}
|
||||
}
|
||||
|
||||
SECTION("logical device")
|
||||
{
|
||||
int i = 0;
|
||||
|
||||
for (auto& comm : this->communicators())
|
||||
{
|
||||
REQUIRE(comm.logical_device().underlying_device() == cuda::devices[i]);
|
||||
++i;
|
||||
}
|
||||
}
|
||||
|
||||
SECTION("group_guard round trip")
|
||||
{
|
||||
// Opening and closing a guard with no enqueued ops must not throw.
|
||||
[[maybe_unused]] auto g = this->communicators().front().group_guard();
|
||||
}
|
||||
|
||||
SECTION("device mismatch throws")
|
||||
{
|
||||
if (cuda::devices.size() > 1)
|
||||
{
|
||||
REQUIRE_THROWS_WITH(
|
||||
cudax::nccl_communicator_ref(this->communicators()[0].native_handle(), cudax::logical_device{cuda::devices[1]}),
|
||||
"Inconsistent devices, NCCL communicator device and provided logical device do not match");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,487 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/buffer>
|
||||
#include <cuda/devices>
|
||||
#include <cuda/functional>
|
||||
#include <cuda/memory_pool>
|
||||
#include <cuda/memory_resource>
|
||||
#include <cuda/std/cstddef>
|
||||
#include <cuda/std/cstdint>
|
||||
#include <cuda/std/functional>
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include <nccl_test_common.h>
|
||||
|
||||
namespace
|
||||
{
|
||||
constexpr cuda::std::int32_t ROOT_RANK = 0;
|
||||
} // namespace
|
||||
|
||||
MULTI_GPU_TEST("nccl_communicator_ref all_reduce sum", )
|
||||
{
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
// Rank r contributes {r+1, r+1, r+1, r+1}; the element-wise sum is the n-th triangular number.
|
||||
std::vector<cuda::device_buffer<cuda::std::int32_t>> send;
|
||||
std::vector<cuda::device_buffer<cuda::std::int32_t>> recv;
|
||||
|
||||
for (int i = 0; i < static_cast<int>(cuda::devices.size()); ++i)
|
||||
{
|
||||
auto pool = cuda::device_default_memory_pool(cuda::devices[i]);
|
||||
const auto values = {i + 1, i + 1, i + 1, i + 1};
|
||||
auto& s = send.emplace_back(streams[i], pool, values);
|
||||
recv.emplace_back(cuda::make_buffer(streams[i], pool, s.size(), cuda::std::int32_t{-1}));
|
||||
}
|
||||
|
||||
{
|
||||
auto g = this->communicators().front().group_guard();
|
||||
|
||||
for (cuda::std::size_t i = 0; i < cuda::devices.size(); ++i)
|
||||
{
|
||||
this->communicators()[i].all_reduce(
|
||||
g, send[i].data(), recv[i].data(), send[i].size(), ::cuda::std::plus<>{}, streams[i]);
|
||||
}
|
||||
}
|
||||
|
||||
for (auto& stream : streams)
|
||||
{
|
||||
stream.sync();
|
||||
}
|
||||
|
||||
const auto sum = static_cast<cuda::std::int32_t>(cuda::devices.size() * (cuda::devices.size() + 1) / 2);
|
||||
|
||||
for (auto& buf : recv)
|
||||
{
|
||||
auto pool = cuda::mr::legacy_pinned_memory_resource{};
|
||||
const auto expected = cuda::make_buffer(buf.stream(), pool, buf.size(), sum);
|
||||
const auto actual = cuda::make_buffer(buf.stream(), pool, buf);
|
||||
|
||||
REQUIRE_THAT(actual, Equals(expected));
|
||||
}
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("nccl_communicator_ref all_reduce maximum", )
|
||||
{
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
std::vector<cuda::device_buffer<cuda::std::int32_t>> send;
|
||||
std::vector<cuda::device_buffer<cuda::std::int32_t>> recv;
|
||||
|
||||
for (int i = 0; i < static_cast<int>(cuda::devices.size()); ++i)
|
||||
{
|
||||
auto pool = cuda::device_default_memory_pool(cuda::devices[i]);
|
||||
const auto values = {i, 100 - i, 2 * i};
|
||||
auto& s = send.emplace_back(streams[i], pool, values);
|
||||
recv.emplace_back(cuda::make_buffer(streams[i], pool, s.size(), cuda::std::int32_t{-1}));
|
||||
}
|
||||
|
||||
{
|
||||
auto g = this->communicators().front().group_guard();
|
||||
|
||||
for (cuda::std::size_t i = 0; i < cuda::devices.size(); ++i)
|
||||
{
|
||||
this->communicators()[i].all_reduce(
|
||||
g, send[i].data(), recv[i].data(), send[i].size(), ::cuda::maximum<>{}, streams[i]);
|
||||
}
|
||||
}
|
||||
|
||||
for (auto& stream : streams)
|
||||
{
|
||||
stream.sync();
|
||||
}
|
||||
|
||||
const auto n = static_cast<cuda::std::int32_t>(cuda::devices.size());
|
||||
|
||||
for (auto& buf : recv)
|
||||
{
|
||||
auto pool = cuda::mr::legacy_pinned_memory_resource{};
|
||||
const auto expected_values = {n - 1, cuda::std::int32_t{100}, 2 * (n - 1)};
|
||||
const auto expected = cuda::make_buffer(buf.stream(), pool, expected_values);
|
||||
const auto actual = cuda::make_buffer(buf.stream(), pool, buf);
|
||||
|
||||
REQUIRE_THAT(actual, Equals(expected));
|
||||
}
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("nccl_communicator_ref reduce sum to root 0", )
|
||||
{
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
std::vector<cuda::device_buffer<cuda::std::int32_t>> send;
|
||||
std::vector<cuda::device_buffer<cuda::std::int32_t>> recv;
|
||||
|
||||
for (int i = 0; i < static_cast<int>(cuda::devices.size()); ++i)
|
||||
{
|
||||
auto pool = cuda::device_default_memory_pool(cuda::devices[i]);
|
||||
const auto values = {i + 1, i + 1, i + 1, i + 1};
|
||||
auto& s = send.emplace_back(streams[i], pool, values);
|
||||
recv.emplace_back(cuda::make_buffer(streams[i], pool, s.size(), cuda::std::int32_t{-1}));
|
||||
}
|
||||
|
||||
{
|
||||
auto g = this->communicators().front().group_guard();
|
||||
|
||||
for (cuda::std::size_t i = 0; i < cuda::devices.size(); ++i)
|
||||
{
|
||||
this->communicators()[i].reduce(
|
||||
g, send[i].data(), recv[i].data(), send[i].size(), ::cuda::std::plus<>{}, ROOT_RANK, streams[i]);
|
||||
}
|
||||
}
|
||||
|
||||
for (auto& stream : streams)
|
||||
{
|
||||
stream.sync();
|
||||
}
|
||||
|
||||
const auto sum = static_cast<cuda::std::int32_t>(cuda::devices.size() * (cuda::devices.size() + 1) / 2);
|
||||
auto pool = cuda::mr::legacy_pinned_memory_resource{};
|
||||
const auto actual = cuda::make_buffer(recv[ROOT_RANK].stream(), pool, recv[ROOT_RANK]);
|
||||
const auto expected = cuda::make_buffer(actual.stream(), pool, actual.size(), sum);
|
||||
|
||||
REQUIRE_THAT(actual, Equals(expected));
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("nccl_communicator_ref broadcast from root 0", )
|
||||
{
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
// Only root's send buffer is read; give every rank the same literal so the source is obvious.
|
||||
std::vector<cuda::device_buffer<cuda::std::int32_t>> send;
|
||||
std::vector<cuda::device_buffer<cuda::std::int32_t>> recv;
|
||||
|
||||
for (int i = 0; i < static_cast<int>(cuda::devices.size()); ++i)
|
||||
{
|
||||
auto pool = cuda::device_default_memory_pool(cuda::devices[i]);
|
||||
const auto values = {10, 20, 30, 40};
|
||||
auto& s = send.emplace_back(streams[i], pool, values);
|
||||
recv.emplace_back(cuda::make_buffer(streams[i], pool, s.size(), cuda::std::int32_t{-1}));
|
||||
}
|
||||
|
||||
{
|
||||
auto g = this->communicators().front().group_guard();
|
||||
|
||||
for (cuda::std::size_t i = 0; i < cuda::devices.size(); ++i)
|
||||
{
|
||||
this->communicators()[i].broadcast(g, send[i].data(), recv[i].data(), send[i].size(), ROOT_RANK, streams[i]);
|
||||
}
|
||||
}
|
||||
|
||||
for (auto& stream : streams)
|
||||
{
|
||||
stream.sync();
|
||||
}
|
||||
|
||||
for (auto& buf : recv)
|
||||
{
|
||||
auto pool = cuda::mr::legacy_pinned_memory_resource{};
|
||||
const auto expected_values = {
|
||||
cuda::std::int32_t{10}, cuda::std::int32_t{20}, cuda::std::int32_t{30}, cuda::std::int32_t{40}};
|
||||
const auto expected = cuda::make_buffer(buf.stream(), pool, expected_values);
|
||||
const auto actual = cuda::make_buffer(buf.stream(), pool, buf);
|
||||
|
||||
REQUIRE_THAT(actual, Equals(expected));
|
||||
}
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("nccl_communicator_ref all_gather", )
|
||||
{
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
// Rank r contributes {10*r, 10*r+1}; every rank ends up with the concatenation in rank order.
|
||||
std::vector<cuda::device_buffer<cuda::std::int32_t>> send;
|
||||
std::vector<cuda::device_buffer<cuda::std::int32_t>> recv;
|
||||
|
||||
for (int i = 0; i < static_cast<int>(cuda::devices.size()); ++i)
|
||||
{
|
||||
auto pool = cuda::device_default_memory_pool(cuda::devices[i]);
|
||||
const auto values = {10 * i, 10 * i + 1};
|
||||
auto& s = send.emplace_back(streams[i], pool, values);
|
||||
recv.emplace_back(cuda::make_buffer(streams[i], pool, s.size() * cuda::devices.size(), cuda::std::int32_t{-1}));
|
||||
}
|
||||
|
||||
{
|
||||
auto g = this->communicators().front().group_guard();
|
||||
|
||||
for (cuda::std::size_t i = 0; i < cuda::devices.size(); ++i)
|
||||
{
|
||||
this->communicators()[i].all_gather(g, send[i].data(), recv[i].data(), send[i].size(), streams[i]);
|
||||
}
|
||||
}
|
||||
|
||||
for (auto& stream : streams)
|
||||
{
|
||||
stream.sync();
|
||||
}
|
||||
|
||||
const cuda::std::size_t per_rank = send.front().size();
|
||||
std::vector<cuda::std::int32_t> expected_values(per_rank * cuda::devices.size());
|
||||
|
||||
for (cuda::std::size_t r = 0; r < cuda::devices.size(); ++r)
|
||||
{
|
||||
expected_values[r * per_rank] = static_cast<cuda::std::int32_t>(10 * r);
|
||||
expected_values[(r * per_rank) + 1] = static_cast<cuda::std::int32_t>((10 * r) + 1);
|
||||
}
|
||||
|
||||
for (auto& buf : recv)
|
||||
{
|
||||
auto pool = cuda::mr::legacy_pinned_memory_resource{};
|
||||
const auto expected = cuda::make_buffer<cuda::std::int32_t>(buf.stream(), pool, expected_values);
|
||||
const auto actual = cuda::make_buffer(buf.stream(), pool, buf);
|
||||
|
||||
REQUIRE_THAT(actual, Equals(expected));
|
||||
}
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("nccl_communicator_ref gather_v to root 0", )
|
||||
{
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
// Rank r contributes (2 + r) elements: {10*r, 10*r+1, ...}. Root concatenates them in rank order.
|
||||
std::vector<cuda::device_buffer<cuda::std::int32_t>> send;
|
||||
|
||||
for (int i = 0; i < static_cast<int>(cuda::devices.size()); ++i)
|
||||
{
|
||||
auto pool = cuda::device_default_memory_pool(cuda::devices[i]);
|
||||
|
||||
std::vector<cuda::std::int32_t> h(2 + static_cast<cuda::std::size_t>(i));
|
||||
|
||||
for (cuda::std::size_t k = 0; k < h.size(); ++k)
|
||||
{
|
||||
h[k] = (10 * i) + static_cast<int>(k);
|
||||
}
|
||||
|
||||
send.emplace_back(streams[i], pool, h);
|
||||
}
|
||||
|
||||
std::vector<cuda::std::size_t> recv_counts(cuda::devices.size());
|
||||
std::vector<cuda::std::size_t> displs(cuda::devices.size());
|
||||
cuda::std::size_t total = 0;
|
||||
|
||||
for (cuda::std::size_t r = 0; r < cuda::devices.size(); ++r)
|
||||
{
|
||||
recv_counts[r] = send[r].size();
|
||||
displs[r] = total;
|
||||
total += recv_counts[r];
|
||||
}
|
||||
|
||||
auto root_pool = cuda::device_default_memory_pool(cuda::devices[ROOT_RANK]);
|
||||
auto recv = cuda::make_buffer(streams[ROOT_RANK], root_pool, total, cuda::std::int32_t{-1});
|
||||
|
||||
{
|
||||
auto g = this->communicators().front().group_guard();
|
||||
|
||||
for (cuda::std::size_t i = 0; i < cuda::devices.size(); ++i)
|
||||
{
|
||||
this->communicators()[i].gather_v(
|
||||
g, send[i].data(), send[i].size(), recv.data(), recv_counts.data(), displs.data(), ROOT_RANK, streams[i]);
|
||||
}
|
||||
}
|
||||
|
||||
for (auto& stream : streams)
|
||||
{
|
||||
stream.sync();
|
||||
}
|
||||
|
||||
std::vector<cuda::std::int32_t> expected_values(total);
|
||||
|
||||
for (cuda::std::size_t r = 0; r < cuda::devices.size(); ++r)
|
||||
{
|
||||
for (cuda::std::size_t k = 0; k < recv_counts[r]; ++k)
|
||||
{
|
||||
expected_values[displs[r] + k] = static_cast<cuda::std::int32_t>((10 * r) + k);
|
||||
}
|
||||
}
|
||||
|
||||
auto pool = cuda::mr::legacy_pinned_memory_resource{};
|
||||
const auto expected = cuda::make_buffer<cuda::std::int32_t>(recv.stream(), pool, expected_values);
|
||||
const auto actual = cuda::make_buffer(recv.stream(), pool, recv);
|
||||
|
||||
REQUIRE_THAT(actual, Equals(expected));
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("nccl_communicator_ref all_to_all_v", )
|
||||
{
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
// Two elements exchanged with each peer; the block for peer j sits at offset block*j.
|
||||
constexpr cuda::std::size_t block = 2;
|
||||
|
||||
std::vector<cuda::std::size_t> counts(cuda::devices.size(), block);
|
||||
std::vector<cuda::std::size_t> displs(cuda::devices.size());
|
||||
|
||||
for (cuda::std::size_t j = 0; j < cuda::devices.size(); ++j)
|
||||
{
|
||||
displs[j] = block * j;
|
||||
}
|
||||
|
||||
// Rank r block destined for peer j encodes 100*r + 10*j + k.
|
||||
std::vector<cuda::device_buffer<cuda::std::int32_t>> send;
|
||||
std::vector<cuda::device_buffer<cuda::std::int32_t>> recv;
|
||||
|
||||
for (int i = 0; i < static_cast<int>(cuda::devices.size()); ++i)
|
||||
{
|
||||
auto pool = cuda::device_default_memory_pool(cuda::devices[i]);
|
||||
|
||||
std::vector<cuda::std::int32_t> h(block * cuda::devices.size());
|
||||
|
||||
for (cuda::std::size_t j = 0; j < cuda::devices.size(); ++j)
|
||||
{
|
||||
h[block * j] = static_cast<cuda::std::int32_t>((100 * i) + (10 * j));
|
||||
h[(block * j) + 1] = static_cast<cuda::std::int32_t>((100 * i) + (10 * j) + 1);
|
||||
}
|
||||
|
||||
auto& s = send.emplace_back(streams[i], pool, h);
|
||||
recv.emplace_back(cuda::make_buffer(streams[i], pool, s.size(), cuda::std::int32_t{-1}));
|
||||
}
|
||||
|
||||
{
|
||||
auto g = this->communicators().front().group_guard();
|
||||
|
||||
for (cuda::std::size_t i = 0; i < cuda::devices.size(); ++i)
|
||||
{
|
||||
this->communicators()[i].all_to_all_v(
|
||||
g, send[i].data(), counts.data(), displs.data(), recv[i].data(), counts.data(), displs.data(), streams[i]);
|
||||
}
|
||||
}
|
||||
|
||||
for (auto& stream : streams)
|
||||
{
|
||||
stream.sync();
|
||||
}
|
||||
|
||||
// Rank r receives from peer i the block i sent to r, placed at block*i: 100*i + 10*r + k.
|
||||
for (cuda::std::size_t r = 0; r < cuda::devices.size(); ++r)
|
||||
{
|
||||
std::vector<cuda::std::int32_t> expected_values(block * cuda::devices.size());
|
||||
|
||||
for (cuda::std::size_t i = 0; i < cuda::devices.size(); ++i)
|
||||
{
|
||||
expected_values[block * i] = static_cast<cuda::std::int32_t>((100 * i) + (10 * r));
|
||||
expected_values[(block * i) + 1] = static_cast<cuda::std::int32_t>((100 * i) + (10 * r) + 1);
|
||||
}
|
||||
|
||||
auto pool = cuda::mr::legacy_pinned_memory_resource{};
|
||||
const auto expected = cuda::make_buffer<cuda::std::int32_t>(recv[r].stream(), pool, expected_values);
|
||||
const auto actual = cuda::make_buffer(recv[r].stream(), pool, recv[r]);
|
||||
|
||||
REQUIRE_THAT(actual, Equals(expected));
|
||||
}
|
||||
}
|
||||
|
||||
#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 28, 0)
|
||||
|
||||
MULTI_GPU_TEST("nccl_communicator_ref gather to root", )
|
||||
{
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
std::vector<cuda::device_buffer<cuda::std::int32_t>> send;
|
||||
std::vector<cuda::device_buffer<cuda::std::int32_t>> recv;
|
||||
|
||||
for (int i = 0; i < static_cast<int>(cuda::devices.size()); ++i)
|
||||
{
|
||||
auto pool = cuda::device_default_memory_pool(cuda::devices[i]);
|
||||
const auto values = {10 * i, 10 * i + 1};
|
||||
auto& s = send.emplace_back(streams[i], pool, values);
|
||||
recv.emplace_back(cuda::make_buffer(streams[i], pool, s.size() * cuda::devices.size(), cuda::std::int32_t{-1}));
|
||||
}
|
||||
|
||||
{
|
||||
auto g = this->communicators().front().group_guard();
|
||||
|
||||
for (cuda::std::size_t i = 0; i < cuda::devices.size(); ++i)
|
||||
{
|
||||
this->communicators()[i].gather(g, send[i].data(), recv[i].data(), send[i].size(), ROOT_RANK, streams[i]);
|
||||
}
|
||||
}
|
||||
|
||||
for (auto& stream : streams)
|
||||
{
|
||||
stream.sync();
|
||||
}
|
||||
|
||||
const cuda::std::size_t per_rank = send.front().size();
|
||||
std::vector<cuda::std::int32_t> expected_values(per_rank * cuda::devices.size());
|
||||
|
||||
for (cuda::std::size_t r = 0; r < cuda::devices.size(); ++r)
|
||||
{
|
||||
expected_values[r * per_rank] = static_cast<cuda::std::int32_t>(10 * r);
|
||||
expected_values[r * per_rank + 1] = static_cast<cuda::std::int32_t>(10 * r + 1);
|
||||
}
|
||||
|
||||
auto pool = cuda::mr::legacy_pinned_memory_resource{};
|
||||
const auto expected = cuda::make_buffer<cuda::std::int32_t>(recv[ROOT_RANK].stream(), pool, expected_values);
|
||||
const auto actual = cuda::make_buffer(recv[ROOT_RANK].stream(), pool, recv[ROOT_RANK]);
|
||||
|
||||
REQUIRE_THAT(actual, Equals(expected));
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("nccl_communicator_ref all_to_all", )
|
||||
{
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
// Two elements exchanged with each peer; the block for peer j sits at offset block*j.
|
||||
constexpr cuda::std::size_t block = 2;
|
||||
|
||||
// Rank r block destined for peer j encodes 100*r + 10*j + k.
|
||||
std::vector<cuda::device_buffer<cuda::std::int32_t>> send;
|
||||
std::vector<cuda::device_buffer<cuda::std::int32_t>> recv;
|
||||
|
||||
for (int i = 0; i < static_cast<int>(cuda::devices.size()); ++i)
|
||||
{
|
||||
auto pool = cuda::device_default_memory_pool(cuda::devices[i]);
|
||||
|
||||
std::vector<cuda::std::int32_t> h(block * cuda::devices.size());
|
||||
|
||||
for (cuda::std::size_t j = 0; j < cuda::devices.size(); ++j)
|
||||
{
|
||||
h[block * j] = static_cast<cuda::std::int32_t>(100 * i + 10 * j);
|
||||
h[block * j + 1] = static_cast<cuda::std::int32_t>(100 * i + 10 * j + 1);
|
||||
}
|
||||
|
||||
auto& s = send.emplace_back(streams[i], pool, h);
|
||||
recv.emplace_back(cuda::make_buffer(streams[i], pool, s.size(), cuda::std::int32_t{-1}));
|
||||
}
|
||||
|
||||
{
|
||||
auto g = this->communicators().front().group_guard();
|
||||
|
||||
for (cuda::std::size_t i = 0; i < cuda::devices.size(); ++i)
|
||||
{
|
||||
this->communicators()[i].all_to_all(g, send[i].data(), recv[i].data(), block, streams[i]);
|
||||
}
|
||||
}
|
||||
|
||||
for (auto& stream : streams)
|
||||
{
|
||||
stream.sync();
|
||||
}
|
||||
|
||||
// Rank r receives from peer i the block i sent to r, placed at block*i: 100*i + 10*r + k.
|
||||
for (cuda::std::size_t r = 0; r < cuda::devices.size(); ++r)
|
||||
{
|
||||
std::vector<cuda::std::int32_t> expected_values(block * cuda::devices.size());
|
||||
|
||||
for (cuda::std::size_t i = 0; i < cuda::devices.size(); ++i)
|
||||
{
|
||||
expected_values[block * i] = static_cast<cuda::std::int32_t>(100 * i + 10 * r);
|
||||
expected_values[block * i + 1] = static_cast<cuda::std::int32_t>(100 * i + 10 * r + 1);
|
||||
}
|
||||
|
||||
auto pool = cuda::mr::legacy_pinned_memory_resource{};
|
||||
const auto expected = cuda::make_buffer<cuda::std::int32_t>(recv[r].stream(), pool, expected_values);
|
||||
const auto actual = cuda::make_buffer(recv[r].stream(), pool, recv[r]);
|
||||
|
||||
REQUIRE_THAT(actual, Equals(expected));
|
||||
}
|
||||
}
|
||||
|
||||
#endif // NCCL_VERSION_CODE >= NCCL_VERSION(2, 28, 0)
|
||||
@@ -1,92 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/functional>
|
||||
#include <cuda/std/functional>
|
||||
|
||||
#include <cuda/experimental/__multi_gpu/concepts.h>
|
||||
#include <cuda/experimental/__multi_gpu/nccl_communicator_ref.h>
|
||||
|
||||
#include <functional>
|
||||
|
||||
#include <testing.cuh>
|
||||
|
||||
namespace
|
||||
{
|
||||
struct payload
|
||||
{
|
||||
int from;
|
||||
int value;
|
||||
};
|
||||
|
||||
struct non_trivial
|
||||
{
|
||||
non_trivial(const non_trivial&) {} // NOLINT(modernize-use-equals-default)
|
||||
|
||||
int value;
|
||||
};
|
||||
|
||||
struct unsupported_op
|
||||
{};
|
||||
} // namespace
|
||||
|
||||
C2H_TEST("nccl_communicator_ref concept conformance", "[multi_gpu][nccl]")
|
||||
{
|
||||
STATIC_REQUIRE(cudax::__communicator<cudax::nccl_communicator_ref>);
|
||||
STATIC_REQUIRE(cudax::__has_send<cudax::nccl_communicator_ref, int*>);
|
||||
STATIC_REQUIRE(cudax::__has_send<cudax::nccl_communicator_ref, payload*>);
|
||||
STATIC_REQUIRE(!cudax::__has_send<cudax::nccl_communicator_ref, non_trivial*>);
|
||||
STATIC_REQUIRE(cudax::__has_recv<cudax::nccl_communicator_ref, int*>);
|
||||
STATIC_REQUIRE(cudax::__has_recv<cudax::nccl_communicator_ref, payload*>);
|
||||
STATIC_REQUIRE(!cudax::__has_recv<cudax::nccl_communicator_ref, non_trivial*>);
|
||||
STATIC_REQUIRE(cudax::__has_reduce<cudax::nccl_communicator_ref, int*>);
|
||||
STATIC_REQUIRE(!cudax::__has_reduce<cudax::nccl_communicator_ref, payload*>);
|
||||
STATIC_REQUIRE(cudax::__has_all_reduce<cudax::nccl_communicator_ref>);
|
||||
STATIC_REQUIRE(!cudax::__has_all_reduce<cudax::nccl_communicator_ref, payload*>);
|
||||
STATIC_REQUIRE(cudax::__has_gather<cudax::nccl_communicator_ref>);
|
||||
STATIC_REQUIRE(cudax::__has_gather<cudax::nccl_communicator_ref, payload*>);
|
||||
STATIC_REQUIRE(!cudax::__has_gather<cudax::nccl_communicator_ref, non_trivial*>);
|
||||
STATIC_REQUIRE(cudax::__has_gather_v<cudax::nccl_communicator_ref>);
|
||||
STATIC_REQUIRE(cudax::__has_gather_v<cudax::nccl_communicator_ref, payload*>);
|
||||
STATIC_REQUIRE(!cudax::__has_gather_v<cudax::nccl_communicator_ref, non_trivial*>);
|
||||
STATIC_REQUIRE(cudax::__has_all_gather<cudax::nccl_communicator_ref>);
|
||||
STATIC_REQUIRE(cudax::__has_all_gather<cudax::nccl_communicator_ref, payload*>);
|
||||
STATIC_REQUIRE(!cudax::__has_all_gather<cudax::nccl_communicator_ref, non_trivial*>);
|
||||
STATIC_REQUIRE(cudax::__has_broadcast<cudax::nccl_communicator_ref>);
|
||||
STATIC_REQUIRE(cudax::__has_broadcast<cudax::nccl_communicator_ref, payload*>);
|
||||
STATIC_REQUIRE(!cudax::__has_broadcast<cudax::nccl_communicator_ref, non_trivial*>);
|
||||
STATIC_REQUIRE(cudax::__has_all_to_all<cudax::nccl_communicator_ref>);
|
||||
STATIC_REQUIRE(cudax::__has_all_to_all<cudax::nccl_communicator_ref, payload*>);
|
||||
STATIC_REQUIRE(!cudax::__has_all_to_all<cudax::nccl_communicator_ref, non_trivial*>);
|
||||
STATIC_REQUIRE(cudax::__has_all_to_all_v<cudax::nccl_communicator_ref>);
|
||||
STATIC_REQUIRE(cudax::__has_all_to_all_v<cudax::nccl_communicator_ref, payload*>);
|
||||
STATIC_REQUIRE(!cudax::__has_all_to_all_v<cudax::nccl_communicator_ref, non_trivial*>);
|
||||
|
||||
STATIC_REQUIRE(cudax::nccl_transportable<int>);
|
||||
STATIC_REQUIRE(cudax::nccl_transportable<int*>);
|
||||
STATIC_REQUIRE(cudax::nccl_transportable<const int*>);
|
||||
STATIC_REQUIRE(cudax::nccl_transportable<const volatile int* const>);
|
||||
STATIC_REQUIRE(cudax::nccl_transportable<float>);
|
||||
STATIC_REQUIRE(cudax::nccl_transportable<::cuda::std::int32_t>);
|
||||
STATIC_REQUIRE(cudax::nccl_transportable<void>);
|
||||
STATIC_REQUIRE(cudax::nccl_transportable<payload>);
|
||||
STATIC_REQUIRE(!cudax::nccl_transportable<non_trivial>);
|
||||
|
||||
STATIC_REQUIRE(cudax::nccl_reducible<int, cuda::std::plus<>>);
|
||||
STATIC_REQUIRE(cudax::nccl_reducible<int, cuda::std::multiplies<>>);
|
||||
STATIC_REQUIRE(cudax::nccl_reducible<int, cuda::maximum<>>);
|
||||
STATIC_REQUIRE(cudax::nccl_reducible<int, cuda::minimum<>>);
|
||||
STATIC_REQUIRE(!cudax::nccl_reducible<payload, cuda::std::plus<>>);
|
||||
STATIC_REQUIRE(!cudax::nccl_reducible<int, unsupported_op>);
|
||||
|
||||
STATIC_REQUIRE(cudax::nccl_reducible<int, std::plus<>>);
|
||||
STATIC_REQUIRE(cudax::nccl_reducible<int, std::multiplies<>>);
|
||||
STATIC_REQUIRE(!cudax::nccl_reducible<payload, std::plus<>>);
|
||||
}
|
||||
@@ -1,150 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/buffer>
|
||||
#include <cuda/devices>
|
||||
#include <cuda/memory_pool>
|
||||
#include <cuda/memory_resource>
|
||||
#include <cuda/std/cstdint>
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include <nccl_test_common.h>
|
||||
|
||||
namespace
|
||||
{
|
||||
struct payload
|
||||
{
|
||||
cuda::std::int32_t from;
|
||||
cuda::std::int32_t index;
|
||||
};
|
||||
} // namespace
|
||||
|
||||
// Ring exchange via send/recv. Rank r contributes {r, r, r}.
|
||||
MULTI_GPU_TEST("nccl_communicator_ref send/recv ring", )
|
||||
{
|
||||
if (cuda::devices.size() == 1)
|
||||
{
|
||||
// NCCL disallows self send/recv on a 1-rank comm.
|
||||
REQUIRE(this->communicators().front().rank() == 0);
|
||||
REQUIRE(this->communicators().front().size() == 1);
|
||||
return;
|
||||
}
|
||||
|
||||
const int n = static_cast<int>(cuda::devices.size());
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
std::vector<cuda::device_buffer<cuda::std::int32_t>> send;
|
||||
std::vector<cuda::device_buffer<cuda::std::int32_t>> recv;
|
||||
|
||||
for (int i = 0; i < n; ++i)
|
||||
{
|
||||
auto pool = cuda::device_default_memory_pool(cuda::devices[i]);
|
||||
|
||||
auto& s = send.emplace_back(cuda::make_buffer(streams[i], pool, 3, i));
|
||||
recv.emplace_back(cuda::make_buffer<cuda::std::int32_t>(streams[i], pool, s.size(), -1));
|
||||
}
|
||||
|
||||
{
|
||||
auto g = this->communicators().front().group_guard();
|
||||
|
||||
for (int i = 0; i < n; ++i)
|
||||
{
|
||||
const int prev = (i + n - 1) % n;
|
||||
const int next = (i + 1) % n;
|
||||
|
||||
this->communicators()[i].recv(g, recv[i].data(), recv[i].size(), prev, streams[i]);
|
||||
this->communicators()[i].send(g, send[i].data(), send[i].size(), next, streams[i]);
|
||||
}
|
||||
}
|
||||
|
||||
for (auto& stream : streams)
|
||||
{
|
||||
stream.sync();
|
||||
}
|
||||
|
||||
// Rank r received from its predecessor (r-1): {r-1, r-1, r-1}.
|
||||
for (int r = 0; r < n; ++r)
|
||||
{
|
||||
const cuda::std::int32_t prev = (r + n - 1) % n;
|
||||
|
||||
auto pool = cuda::mr::legacy_pinned_memory_resource{};
|
||||
const cuda::host_buffer<cuda::std::int32_t> expected =
|
||||
cuda::make_buffer(recv[r].stream(), pool, recv[r].size(), prev);
|
||||
const cuda::host_buffer<cuda::std::int32_t> actual = cuda::make_buffer(recv[r].stream(), pool, recv[r]);
|
||||
|
||||
REQUIRE_THAT(actual, Equals(expected));
|
||||
}
|
||||
}
|
||||
|
||||
MULTI_GPU_TEST("nccl_communicator_ref send/recv transports trivially copyable payload", )
|
||||
{
|
||||
if (cuda::devices.size() == 1)
|
||||
{
|
||||
// NCCL disallows self send/recv on a 1-rank comm.
|
||||
REQUIRE(this->communicators().front().rank() == 0);
|
||||
REQUIRE(this->communicators().front().size() == 1);
|
||||
return;
|
||||
}
|
||||
|
||||
const int n = static_cast<int>(cuda::devices.size());
|
||||
auto streams = nccl_test_util::make_streams();
|
||||
|
||||
std::vector<cuda::device_buffer<payload>> send;
|
||||
std::vector<cuda::device_buffer<payload>> recv;
|
||||
|
||||
for (int i = 0; i < n; ++i)
|
||||
{
|
||||
auto pool = cuda::device_default_memory_pool(cuda::devices[i]);
|
||||
|
||||
std::vector<payload> h(3);
|
||||
|
||||
for (cuda::std::size_t k = 0; k < h.size(); ++k)
|
||||
{
|
||||
h[k] = payload{static_cast<cuda::std::int32_t>(i), static_cast<cuda::std::int32_t>(k)};
|
||||
}
|
||||
|
||||
auto& s = send.emplace_back(streams[i], pool, h);
|
||||
recv.emplace_back(cuda::make_buffer<payload>(streams[i], pool, s.size(), cuda::no_init));
|
||||
}
|
||||
|
||||
{
|
||||
auto g = this->communicators().front().group_guard();
|
||||
|
||||
for (int i = 0; i < n; ++i)
|
||||
{
|
||||
const int prev = (i + n - 1) % n;
|
||||
const int next = (i + 1) % n;
|
||||
|
||||
this->communicators()[i].recv(g, recv[i].data(), recv[i].size(), prev, streams[i]);
|
||||
this->communicators()[i].send(g, send[i].data(), send[i].size(), next, streams[i]);
|
||||
}
|
||||
}
|
||||
|
||||
for (auto& stream : streams)
|
||||
{
|
||||
stream.sync();
|
||||
}
|
||||
|
||||
for (int r = 0; r < n; ++r)
|
||||
{
|
||||
const int prev = (r + n - 1) % n;
|
||||
|
||||
auto pool = cuda::mr::legacy_pinned_memory_resource{};
|
||||
const cuda::host_buffer<payload> actual = cuda::make_buffer(recv[r].stream(), pool, recv[r]);
|
||||
|
||||
actual.stream().sync();
|
||||
for (cuda::std::size_t k = 0; k < actual.size(); ++k)
|
||||
{
|
||||
REQUIRE(actual[k].from == static_cast<cuda::std::int32_t>(prev));
|
||||
REQUIRE(actual[k].index == static_cast<cuda::std::int32_t>(k));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,15 +0,0 @@
|
||||
#===----------------------------------------------------------------------===##
|
||||
#
|
||||
# Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
# under the Apache License v2.0 with LLVM Exceptions.
|
||||
# See https://llvm.org/LICENSE.txt for license information.
|
||||
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
#
|
||||
#===----------------------------------------------------------------------===##
|
||||
|
||||
file(GLOB test_srcs LIST_DIRECTORIES FALSE CONFIGURE_DEPENDS *.cu *.cpp)
|
||||
|
||||
foreach (src IN LISTS test_srcs)
|
||||
cudax_add_multi_gpu_test("concepts" test_target "${src}")
|
||||
endforeach()
|
||||
@@ -1,66 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#ifndef _CUDAX_TEST_MULTI_GPU_COLLECTIVE_CONCEPTS_COMMON_CUH
|
||||
#define _CUDAX_TEST_MULTI_GPU_COLLECTIVE_CONCEPTS_COMMON_CUH
|
||||
|
||||
#include <cuda/__stream/stream_ref.h>
|
||||
#include <cuda/std/__cstddef/types.h>
|
||||
#include <cuda/std/cstdint>
|
||||
|
||||
#include "concepts_common.cuh"
|
||||
|
||||
namespace cudax_multi_gpu_concepts
|
||||
{
|
||||
struct collective_communicator_model : communicator_model
|
||||
{
|
||||
template <class Tp, class Op>
|
||||
void reduce(group_guard_type&, Tp*, Tp*, ::cuda::std::size_t, Op, ::cuda::std::int32_t, ::cuda::stream_ref);
|
||||
|
||||
template <class Tp, class Op>
|
||||
void all_reduce(group_guard_type&, Tp*, Tp*, ::cuda::std::size_t, Op, ::cuda::stream_ref);
|
||||
|
||||
template <class Tp>
|
||||
void gather(group_guard_type&, Tp*, Tp*, ::cuda::std::size_t, ::cuda::std::int32_t, ::cuda::stream_ref);
|
||||
|
||||
template <class Tp>
|
||||
void gather_v(
|
||||
group_guard_type&,
|
||||
Tp*,
|
||||
::cuda::std::size_t,
|
||||
Tp*,
|
||||
const ::cuda::std::size_t*,
|
||||
const ::cuda::std::size_t*,
|
||||
::cuda::std::int32_t,
|
||||
::cuda::stream_ref);
|
||||
|
||||
template <class Tp>
|
||||
void all_gather(group_guard_type&, Tp*, Tp*, ::cuda::std::size_t, ::cuda::stream_ref);
|
||||
|
||||
template <class Tp>
|
||||
void broadcast(group_guard_type&, Tp*, Tp*, ::cuda::std::size_t, ::cuda::std::int32_t, ::cuda::stream_ref);
|
||||
|
||||
template <class Tp>
|
||||
void all_to_all(group_guard_type&, Tp*, Tp*, ::cuda::std::size_t, ::cuda::stream_ref);
|
||||
|
||||
template <class Tp>
|
||||
void all_to_all_v(
|
||||
group_guard_type&,
|
||||
Tp*,
|
||||
const ::cuda::std::size_t*,
|
||||
const ::cuda::std::size_t*,
|
||||
Tp*,
|
||||
const ::cuda::std::size_t*,
|
||||
const ::cuda::std::size_t*,
|
||||
::cuda::stream_ref);
|
||||
};
|
||||
} // namespace cudax_multi_gpu_concepts
|
||||
|
||||
#endif // _CUDAX_TEST_MULTI_GPU_COLLECTIVE_CONCEPTS_COMMON_CUH
|
||||
@@ -1,44 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/experimental/__multi_gpu/concepts.h>
|
||||
|
||||
#include <testing.cuh>
|
||||
|
||||
#include "concepts_common.cuh"
|
||||
|
||||
namespace
|
||||
{
|
||||
namespace types = cudax_multi_gpu_concepts;
|
||||
|
||||
// nvcc ignores [[maybe_unused]] entirely
|
||||
_CCCL_BEGIN_NV_DIAG_SUPPRESS(177)
|
||||
|
||||
struct no_send : types::basic_communicator_model
|
||||
{
|
||||
template <class Tp>
|
||||
void recv(group_guard_type&, Tp*, ::cuda::std::size_t, ::cuda::std::int32_t, ::cuda::stream_ref);
|
||||
};
|
||||
|
||||
struct no_recv : types::basic_communicator_model
|
||||
{
|
||||
template <class Tp>
|
||||
void send(group_guard_type&, Tp*, ::cuda::std::size_t, ::cuda::std::int32_t, ::cuda::stream_ref);
|
||||
};
|
||||
|
||||
_CCCL_END_NV_DIAG_SUPPRESS()
|
||||
} // namespace
|
||||
|
||||
C2H_TEST("communicator concept", "[multi_gpu][concepts]")
|
||||
{
|
||||
STATIC_REQUIRE(cudax::__communicator<types::communicator_model>);
|
||||
STATIC_REQUIRE(!cudax::__communicator<no_send>);
|
||||
STATIC_REQUIRE(!cudax::__communicator<no_recv>);
|
||||
}
|
||||
@@ -1,43 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#ifndef _CUDAX_TEST_MULTI_GPU_CONCEPTS_COMMON_CUH
|
||||
#define _CUDAX_TEST_MULTI_GPU_CONCEPTS_COMMON_CUH
|
||||
|
||||
#include <cuda/__stream/stream_ref.h>
|
||||
#include <cuda/std/__cstddef/types.h>
|
||||
#include <cuda/std/cstdint>
|
||||
|
||||
namespace cudax_multi_gpu_concepts
|
||||
{
|
||||
struct group_guard
|
||||
{};
|
||||
|
||||
struct basic_communicator_model
|
||||
{
|
||||
using native_handle_type = int;
|
||||
using group_guard_type = group_guard;
|
||||
|
||||
native_handle_type native_handle() noexcept;
|
||||
::cuda::std::int32_t rank() noexcept;
|
||||
::cuda::std::int32_t size() noexcept;
|
||||
group_guard_type group_guard();
|
||||
};
|
||||
|
||||
struct communicator_model : basic_communicator_model
|
||||
{
|
||||
template <class Tp>
|
||||
void send(group_guard_type&, Tp*, ::cuda::std::size_t, ::cuda::std::int32_t, ::cuda::stream_ref);
|
||||
template <class Tp>
|
||||
void recv(group_guard_type&, Tp*, ::cuda::std::size_t, ::cuda::std::int32_t, ::cuda::stream_ref);
|
||||
};
|
||||
} // namespace cudax_multi_gpu_concepts
|
||||
|
||||
#endif // _CUDAX_TEST_MULTI_GPU_CONCEPTS_COMMON_CUH
|
||||
@@ -1,43 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/__stream/stream_ref.h>
|
||||
#include <cuda/std/__cstddef/types.h>
|
||||
|
||||
#include <cuda/experimental/__multi_gpu/concepts.h>
|
||||
|
||||
#include <testing.cuh>
|
||||
|
||||
#include "collective_concepts_common.cuh"
|
||||
|
||||
namespace
|
||||
{
|
||||
namespace types = cudax_multi_gpu_concepts;
|
||||
|
||||
// nvcc ignores [[maybe_unused]] entirely
|
||||
_CCCL_BEGIN_NV_DIAG_SUPPRESS(177)
|
||||
|
||||
struct all_gather_returns_int : types::communicator_model
|
||||
{
|
||||
template <class Tp>
|
||||
int all_gather(group_guard_type&, Tp*, Tp*, ::cuda::std::size_t, ::cuda::stream_ref);
|
||||
};
|
||||
|
||||
_CCCL_END_NV_DIAG_SUPPRESS()
|
||||
} // namespace
|
||||
|
||||
C2H_TEST("__has_all_gather concept", "[multi_gpu][concepts]")
|
||||
{
|
||||
STATIC_REQUIRE(cudax::__has_all_gather<types::collective_communicator_model>);
|
||||
STATIC_REQUIRE(cudax::__has_all_gather<types::collective_communicator_model, long*>);
|
||||
STATIC_REQUIRE(!cudax::__has_all_gather<types::communicator_model>);
|
||||
|
||||
STATIC_REQUIRE(!cudax::__has_all_gather<all_gather_returns_int>);
|
||||
}
|
||||
@@ -1,53 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/__stream/stream_ref.h>
|
||||
#include <cuda/std/__cstddef/types.h>
|
||||
#include <cuda/std/__type_traits/enable_if.h>
|
||||
#include <cuda/std/__type_traits/is_void.h>
|
||||
|
||||
#include <cuda/experimental/__multi_gpu/concepts.h>
|
||||
|
||||
#include <testing.cuh>
|
||||
|
||||
#include "collective_concepts_common.cuh"
|
||||
|
||||
namespace
|
||||
{
|
||||
namespace types = cudax_multi_gpu_concepts;
|
||||
|
||||
// nvcc ignores [[maybe_unused]] entirely
|
||||
_CCCL_BEGIN_NV_DIAG_SUPPRESS(177)
|
||||
|
||||
struct all_reduce_rejects_void : types::communicator_model
|
||||
{
|
||||
template <class Tp, class Op, ::cuda::std::enable_if_t<!::cuda::std::is_void_v<Tp>, int> = 0>
|
||||
void all_reduce(group_guard_type&, Tp*, Tp*, ::cuda::std::size_t, Op, ::cuda::stream_ref);
|
||||
};
|
||||
|
||||
struct all_reduce_returns_int : types::communicator_model
|
||||
{
|
||||
template <class Tp, class Op>
|
||||
int all_reduce(group_guard_type&, Tp*, Tp*, ::cuda::std::size_t, Op, ::cuda::stream_ref);
|
||||
};
|
||||
|
||||
_CCCL_END_NV_DIAG_SUPPRESS()
|
||||
} // namespace
|
||||
|
||||
C2H_TEST("__has_all_reduce concept", "[multi_gpu][concepts]")
|
||||
{
|
||||
STATIC_REQUIRE(cudax::__has_all_reduce<types::collective_communicator_model>);
|
||||
STATIC_REQUIRE(cudax::__has_all_reduce<types::collective_communicator_model, long*>);
|
||||
STATIC_REQUIRE(!cudax::__has_all_reduce<types::communicator_model>);
|
||||
|
||||
STATIC_REQUIRE(!cudax::__has_all_reduce<all_reduce_returns_int>);
|
||||
STATIC_REQUIRE(cudax::__has_all_reduce<all_reduce_rejects_void, int*>);
|
||||
STATIC_REQUIRE(!cudax::__has_all_reduce<all_reduce_rejects_void, void*>);
|
||||
}
|
||||
@@ -1,43 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/__stream/stream_ref.h>
|
||||
#include <cuda/std/__cstddef/types.h>
|
||||
|
||||
#include <cuda/experimental/__multi_gpu/concepts.h>
|
||||
|
||||
#include <testing.cuh>
|
||||
|
||||
#include "collective_concepts_common.cuh"
|
||||
|
||||
namespace
|
||||
{
|
||||
namespace types = cudax_multi_gpu_concepts;
|
||||
|
||||
// nvcc ignores [[maybe_unused]] entirely
|
||||
_CCCL_BEGIN_NV_DIAG_SUPPRESS(177)
|
||||
|
||||
struct all_to_all_returns_int : types::communicator_model
|
||||
{
|
||||
template <class Tp>
|
||||
int all_to_all(group_guard_type&, Tp*, Tp*, ::cuda::std::size_t, ::cuda::stream_ref);
|
||||
};
|
||||
|
||||
_CCCL_END_NV_DIAG_SUPPRESS()
|
||||
} // namespace
|
||||
|
||||
C2H_TEST("__has_all_to_all concept", "[multi_gpu][concepts]")
|
||||
{
|
||||
STATIC_REQUIRE(cudax::__has_all_to_all<types::collective_communicator_model>);
|
||||
STATIC_REQUIRE(cudax::__has_all_to_all<types::collective_communicator_model, long*>);
|
||||
STATIC_REQUIRE(!cudax::__has_all_to_all<types::communicator_model>);
|
||||
|
||||
STATIC_REQUIRE(!cudax::__has_all_to_all<all_to_all_returns_int>);
|
||||
}
|
||||
@@ -1,51 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/__stream/stream_ref.h>
|
||||
#include <cuda/std/__cstddef/types.h>
|
||||
|
||||
#include <cuda/experimental/__multi_gpu/concepts.h>
|
||||
|
||||
#include <testing.cuh>
|
||||
|
||||
#include "collective_concepts_common.cuh"
|
||||
|
||||
namespace
|
||||
{
|
||||
namespace types = cudax_multi_gpu_concepts;
|
||||
|
||||
// nvcc ignores [[maybe_unused]] entirely
|
||||
_CCCL_BEGIN_NV_DIAG_SUPPRESS(177)
|
||||
|
||||
struct all_to_all_v_returns_int : types::communicator_model
|
||||
{
|
||||
template <class Tp>
|
||||
int all_to_all_v(
|
||||
group_guard_type&,
|
||||
Tp*,
|
||||
const ::cuda::std::size_t*,
|
||||
const ::cuda::std::size_t*,
|
||||
Tp*,
|
||||
const ::cuda::std::size_t*,
|
||||
const ::cuda::std::size_t*,
|
||||
::cuda::stream_ref);
|
||||
};
|
||||
|
||||
_CCCL_END_NV_DIAG_SUPPRESS()
|
||||
} // namespace
|
||||
|
||||
C2H_TEST("__has_all_to_all_v concept", "[multi_gpu][concepts]")
|
||||
{
|
||||
STATIC_REQUIRE(cudax::__has_all_to_all_v<types::collective_communicator_model>);
|
||||
STATIC_REQUIRE(cudax::__has_all_to_all_v<types::collective_communicator_model, long*>);
|
||||
STATIC_REQUIRE(!cudax::__has_all_to_all_v<types::communicator_model>);
|
||||
|
||||
STATIC_REQUIRE(!cudax::__has_all_to_all_v<all_to_all_v_returns_int>);
|
||||
}
|
||||
@@ -1,44 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/__stream/stream_ref.h>
|
||||
#include <cuda/std/__cstddef/types.h>
|
||||
#include <cuda/std/cstdint>
|
||||
|
||||
#include <cuda/experimental/__multi_gpu/concepts.h>
|
||||
|
||||
#include <testing.cuh>
|
||||
|
||||
#include "collective_concepts_common.cuh"
|
||||
|
||||
namespace
|
||||
{
|
||||
namespace types = cudax_multi_gpu_concepts;
|
||||
|
||||
// nvcc ignores [[maybe_unused]] entirely
|
||||
_CCCL_BEGIN_NV_DIAG_SUPPRESS(177)
|
||||
|
||||
struct broadcast_returns_int : types::communicator_model
|
||||
{
|
||||
template <class Tp>
|
||||
int broadcast(group_guard_type&, Tp*, Tp*, ::cuda::std::size_t, ::cuda::std::int32_t, ::cuda::stream_ref);
|
||||
};
|
||||
|
||||
_CCCL_END_NV_DIAG_SUPPRESS()
|
||||
} // namespace
|
||||
|
||||
C2H_TEST("__has_broadcast concept", "[multi_gpu][concepts]")
|
||||
{
|
||||
STATIC_REQUIRE(cudax::__has_broadcast<types::collective_communicator_model>);
|
||||
STATIC_REQUIRE(cudax::__has_broadcast<types::collective_communicator_model, long*>);
|
||||
STATIC_REQUIRE(!cudax::__has_broadcast<types::communicator_model>);
|
||||
|
||||
STATIC_REQUIRE(!cudax::__has_broadcast<broadcast_returns_int>);
|
||||
}
|
||||
@@ -1,44 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/__stream/stream_ref.h>
|
||||
#include <cuda/std/__cstddef/types.h>
|
||||
#include <cuda/std/cstdint>
|
||||
|
||||
#include <cuda/experimental/__multi_gpu/concepts.h>
|
||||
|
||||
#include <testing.cuh>
|
||||
|
||||
#include "collective_concepts_common.cuh"
|
||||
|
||||
namespace
|
||||
{
|
||||
namespace types = cudax_multi_gpu_concepts;
|
||||
|
||||
// nvcc ignores [[maybe_unused]] entirely
|
||||
_CCCL_BEGIN_NV_DIAG_SUPPRESS(177)
|
||||
|
||||
struct gather_returns_int : types::communicator_model
|
||||
{
|
||||
template <class Tp>
|
||||
int gather(group_guard_type&, Tp*, Tp*, ::cuda::std::size_t, ::cuda::std::int32_t, ::cuda::stream_ref);
|
||||
};
|
||||
|
||||
_CCCL_END_NV_DIAG_SUPPRESS()
|
||||
} // namespace
|
||||
|
||||
C2H_TEST("__has_gather concept", "[multi_gpu][concepts]")
|
||||
{
|
||||
STATIC_REQUIRE(cudax::__has_gather<types::collective_communicator_model>);
|
||||
STATIC_REQUIRE(cudax::__has_gather<types::collective_communicator_model, long*>);
|
||||
STATIC_REQUIRE(!cudax::__has_gather<types::communicator_model>);
|
||||
|
||||
STATIC_REQUIRE(!cudax::__has_gather<gather_returns_int>);
|
||||
}
|
||||
@@ -1,51 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/__stream/stream_ref.h>
|
||||
#include <cuda/std/__cstddef/types.h>
|
||||
#include <cuda/std/cstdint>
|
||||
|
||||
#include <cuda/experimental/__multi_gpu/concepts.h>
|
||||
|
||||
#include <testing.cuh>
|
||||
|
||||
#include "collective_concepts_common.cuh"
|
||||
|
||||
namespace
|
||||
{
|
||||
namespace types = cudax_multi_gpu_concepts;
|
||||
|
||||
// nvcc ignores [[maybe_unused]] entirely
|
||||
_CCCL_BEGIN_NV_DIAG_SUPPRESS(177)
|
||||
|
||||
struct gather_v_returns_int : types::communicator_model
|
||||
{
|
||||
template <class Tp>
|
||||
int gather_v(group_guard_type&,
|
||||
Tp*,
|
||||
::cuda::std::size_t,
|
||||
Tp*,
|
||||
const ::cuda::std::size_t*,
|
||||
const ::cuda::std::size_t*,
|
||||
::cuda::std::int32_t,
|
||||
::cuda::stream_ref);
|
||||
};
|
||||
|
||||
_CCCL_END_NV_DIAG_SUPPRESS()
|
||||
} // namespace
|
||||
|
||||
C2H_TEST("__has_gather_v concept", "[multi_gpu][concepts]")
|
||||
{
|
||||
STATIC_REQUIRE(cudax::__has_gather_v<types::collective_communicator_model>);
|
||||
STATIC_REQUIRE(cudax::__has_gather_v<types::collective_communicator_model, long*>);
|
||||
STATIC_REQUIRE(!cudax::__has_gather_v<types::communicator_model>);
|
||||
|
||||
STATIC_REQUIRE(!cudax::__has_gather_v<gather_v_returns_int>);
|
||||
}
|
||||
@@ -1,50 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/__stream/stream_ref.h>
|
||||
#include <cuda/std/__cstddef/types.h>
|
||||
#include <cuda/std/cstdint>
|
||||
|
||||
#include <cuda/experimental/__multi_gpu/concepts.h>
|
||||
|
||||
#include <testing.cuh>
|
||||
|
||||
#include "concepts_common.cuh"
|
||||
|
||||
namespace
|
||||
{
|
||||
namespace types = cudax_multi_gpu_concepts;
|
||||
|
||||
// nvcc ignores [[maybe_unused]] entirely
|
||||
_CCCL_BEGIN_NV_DIAG_SUPPRESS(177)
|
||||
|
||||
struct recv_returns_int : types::basic_communicator_model
|
||||
{
|
||||
template <class Tp>
|
||||
int recv(group_guard_type&, Tp*, ::cuda::std::size_t, ::cuda::std::int32_t, ::cuda::stream_ref);
|
||||
};
|
||||
|
||||
struct no_recv : types::basic_communicator_model
|
||||
{
|
||||
template <class Tp>
|
||||
void send(group_guard_type&, Tp*, ::cuda::std::size_t, ::cuda::std::int32_t, ::cuda::stream_ref);
|
||||
};
|
||||
|
||||
_CCCL_END_NV_DIAG_SUPPRESS()
|
||||
} // namespace
|
||||
|
||||
C2H_TEST("__has_recv concept", "[multi_gpu][concepts]")
|
||||
{
|
||||
STATIC_REQUIRE(cudax::__has_recv<types::communicator_model>);
|
||||
STATIC_REQUIRE(cudax::__has_recv<types::communicator_model, long*>);
|
||||
STATIC_REQUIRE(!cudax::__has_recv<no_recv>);
|
||||
|
||||
STATIC_REQUIRE(!cudax::__has_recv<recv_returns_int>);
|
||||
}
|
||||
@@ -1,44 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/__stream/stream_ref.h>
|
||||
#include <cuda/std/__cstddef/types.h>
|
||||
#include <cuda/std/cstdint>
|
||||
|
||||
#include <cuda/experimental/__multi_gpu/concepts.h>
|
||||
|
||||
#include <testing.cuh>
|
||||
|
||||
#include "collective_concepts_common.cuh"
|
||||
|
||||
namespace
|
||||
{
|
||||
namespace types = cudax_multi_gpu_concepts;
|
||||
|
||||
// nvcc ignores [[maybe_unused]] entirely
|
||||
_CCCL_BEGIN_NV_DIAG_SUPPRESS(177)
|
||||
|
||||
struct reduce_returns_int : types::communicator_model
|
||||
{
|
||||
template <class Tp, class Op>
|
||||
int reduce(group_guard_type&, Tp*, Tp*, ::cuda::std::size_t, Op, ::cuda::std::int32_t, ::cuda::stream_ref);
|
||||
};
|
||||
|
||||
_CCCL_END_NV_DIAG_SUPPRESS()
|
||||
} // namespace
|
||||
|
||||
C2H_TEST("__has_reduce concept", "[multi_gpu][concepts]")
|
||||
{
|
||||
STATIC_REQUIRE(cudax::__has_reduce<types::collective_communicator_model>);
|
||||
STATIC_REQUIRE(cudax::__has_reduce<types::collective_communicator_model, long*>);
|
||||
STATIC_REQUIRE(!cudax::__has_reduce<types::communicator_model>);
|
||||
|
||||
STATIC_REQUIRE(!cudax::__has_reduce<reduce_returns_int>);
|
||||
}
|
||||
@@ -1,50 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/__stream/stream_ref.h>
|
||||
#include <cuda/std/__cstddef/types.h>
|
||||
#include <cuda/std/cstdint>
|
||||
|
||||
#include <cuda/experimental/__multi_gpu/concepts.h>
|
||||
|
||||
#include <testing.cuh>
|
||||
|
||||
#include "concepts_common.cuh"
|
||||
|
||||
namespace
|
||||
{
|
||||
namespace types = cudax_multi_gpu_concepts;
|
||||
|
||||
// nvcc ignores [[maybe_unused]] entirely
|
||||
_CCCL_BEGIN_NV_DIAG_SUPPRESS(177)
|
||||
|
||||
struct no_send : types::basic_communicator_model
|
||||
{
|
||||
template <class Tp>
|
||||
void send_sync(group_guard_type&, Tp*, ::cuda::std::size_t, ::cuda::std::int32_t);
|
||||
};
|
||||
|
||||
struct send_returns_int : no_send
|
||||
{
|
||||
template <class Tp>
|
||||
int send(group_guard_type&, Tp*, ::cuda::std::size_t, ::cuda::std::int32_t, ::cuda::stream_ref);
|
||||
};
|
||||
|
||||
_CCCL_END_NV_DIAG_SUPPRESS()
|
||||
} // namespace
|
||||
|
||||
C2H_TEST("__has_send concept", "[multi_gpu][concepts]")
|
||||
{
|
||||
STATIC_REQUIRE(cudax::__has_send<types::communicator_model>);
|
||||
STATIC_REQUIRE(cudax::__has_send<types::communicator_model, long*>);
|
||||
|
||||
STATIC_REQUIRE(!cudax::__has_send<no_send>);
|
||||
STATIC_REQUIRE(!cudax::__has_send<send_returns_int>);
|
||||
}
|
||||
@@ -1,76 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cuda/buffer>
|
||||
#include <cuda/std/cstddef>
|
||||
|
||||
#include <exception>
|
||||
#include <future>
|
||||
#include <vector>
|
||||
|
||||
// One output iterator per local output buffer. Collected after `out` is fully built so the
|
||||
// iterators do not dangle across reallocations.
|
||||
template <class T>
|
||||
[[nodiscard]] std::vector<typename cuda::device_buffer<T>::iterator>
|
||||
make_output_iterators(std::vector<cuda::device_buffer<T>>& out)
|
||||
{
|
||||
std::vector<typename cuda::device_buffer<T>::iterator> outputs;
|
||||
|
||||
outputs.reserve(out.size());
|
||||
for (auto& buf : out)
|
||||
{
|
||||
outputs.push_back(buf.begin());
|
||||
}
|
||||
return outputs;
|
||||
}
|
||||
|
||||
template <class Fn>
|
||||
void run_threaded(cuda::std::size_t num_ranks, Fn fn)
|
||||
{
|
||||
// Every rank must be launched before any is waited on: the single-communicator `reduce`
|
||||
// blocks on a collective, so calling `get()` on rank 0's future before rank 1 is even
|
||||
// started would deadlock. Launch all futures into the vector first, then drain them.
|
||||
std::vector<std::future<void>> futures;
|
||||
|
||||
futures.reserve(num_ranks);
|
||||
for (cuda::std::size_t i = 0; i < num_ranks; ++i)
|
||||
{
|
||||
futures.push_back(std::async(std::launch::async, fn, i));
|
||||
}
|
||||
|
||||
// `std::async` stashes any exception thrown by `fn` in the future and `get()` rethrows it on
|
||||
// the main thread, where Catch2 can report it as a normal failure. Any not-yet-drained
|
||||
// future still joins its thread in its destructor, so a throw here never leaves a peer
|
||||
// waiting on an unposted collective. Drain every future so a failure on rank 0 does not mask
|
||||
// one on a peer.
|
||||
std::exception_ptr error = nullptr;
|
||||
|
||||
for (auto& f : futures)
|
||||
{
|
||||
try
|
||||
{
|
||||
f.get();
|
||||
}
|
||||
catch (...)
|
||||
{
|
||||
if (!error)
|
||||
{
|
||||
error = std::current_exception();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (error)
|
||||
{
|
||||
std::rethrow_exception(error);
|
||||
}
|
||||
}
|
||||
@@ -1,94 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#ifndef CUDAX_TEST_MULTI_NCCL_TEST_COMMON_H
|
||||
#define CUDAX_TEST_MULTI_NCCL_TEST_COMMON_H
|
||||
|
||||
#include <cuda/devices>
|
||||
#include <cuda/std/cstddef>
|
||||
#include <cuda/std/span>
|
||||
|
||||
#include <cuda/experimental/__multi_gpu/nccl_communicator.h>
|
||||
#include <cuda/experimental/__multi_gpu/nccl_communicator_ref.h>
|
||||
#include <cuda/experimental/stream.cuh>
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include <nccl.h>
|
||||
|
||||
#include <c2h/catch2_test_helper.h>
|
||||
|
||||
namespace cudax = ::cuda::experimental;
|
||||
|
||||
namespace nccl_test_util
|
||||
{
|
||||
// One stream per rank, each current on its own device.
|
||||
[[nodiscard]] inline std::vector<cudax::stream> make_streams()
|
||||
{
|
||||
return {cuda::devices.begin(), cuda::devices.end()};
|
||||
}
|
||||
|
||||
[[nodiscard]] inline const std::vector<cudax::nccl_communicator>& nccl_comms()
|
||||
{
|
||||
static const auto comms = []() -> std::vector<cudax::nccl_communicator> {
|
||||
if (cuda::devices.size() == 0)
|
||||
{
|
||||
SKIP("No CUDA devices visible");
|
||||
}
|
||||
|
||||
std::vector<int> devs;
|
||||
|
||||
devs.reserve(cuda::devices.size());
|
||||
for (auto d : cuda::devices)
|
||||
{
|
||||
devs.emplace_back(d.get());
|
||||
}
|
||||
|
||||
std::vector<ncclComm_t> raw_comms(devs.size());
|
||||
|
||||
const ncclResult_t result = ncclCommInitAll(raw_comms.data(), static_cast<int>(devs.size()), devs.data());
|
||||
|
||||
INFO("NCCL: " << ncclGetErrorString(result));
|
||||
REQUIRE(result == ncclSuccess);
|
||||
|
||||
std::vector<cudax::nccl_communicator> comms;
|
||||
comms.reserve(raw_comms.size());
|
||||
|
||||
for (const auto comm : raw_comms)
|
||||
{
|
||||
comms.emplace_back(cudax::nccl_communicator::from_native_handle(comm));
|
||||
}
|
||||
|
||||
return comms;
|
||||
}();
|
||||
|
||||
return comms;
|
||||
}
|
||||
|
||||
// Caches a single-process, multi-GPU NCCL communicator world for the life of the entire test
|
||||
// suite.
|
||||
template <class = void>
|
||||
class nccl_comm_fixture
|
||||
{
|
||||
public:
|
||||
[[nodiscard]] cuda::std::span<cudax::nccl_communicator_ref> communicators()
|
||||
{
|
||||
return wrappers_;
|
||||
}
|
||||
|
||||
private:
|
||||
std::vector<cudax::nccl_communicator_ref> wrappers_{nccl_comms().begin(), nccl_comms().end()};
|
||||
};
|
||||
|
||||
#define MULTI_GPU_TEST(NAME, ...) \
|
||||
C2H_TEST_WITH_FIXTURE(::nccl_test_util::nccl_comm_fixture, NAME, "[multi_gpu][nccl]", __VA_ARGS__)
|
||||
} // namespace nccl_test_util
|
||||
|
||||
#endif // CUDAX_TEST_MULTI_GPU_NCCL_TEST_COMMON_H
|
||||
@@ -1,15 +0,0 @@
|
||||
#===----------------------------------------------------------------------===##
|
||||
#
|
||||
# Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
# under the Apache License v2.0 with LLVM Exceptions.
|
||||
# See https://llvm.org/LICENSE.txt for license information.
|
||||
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
#
|
||||
#===----------------------------------------------------------------------===##
|
||||
|
||||
file(GLOB test_srcs LIST_DIRECTORIES FALSE CONFIGURE_DEPENDS *.cu *.cpp)
|
||||
|
||||
foreach (src IN LISTS test_srcs)
|
||||
cudax_add_multi_gpu_test("nccl" test_target "${src}")
|
||||
endforeach()
|
||||
@@ -1,129 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/std/__type_traits/is_same.h>
|
||||
#include <cuda/std/__type_traits/underlying_type.h>
|
||||
|
||||
#include <cuda/experimental/__nccl/abi_compatible.h>
|
||||
|
||||
#include <c2h/catch2_test_helper.h>
|
||||
|
||||
namespace
|
||||
{
|
||||
namespace abi_detail = ::cuda::experimental::__nccl::__abi_detail;
|
||||
|
||||
// NOLINTBEGIN(bugprone-reserved-identifier)
|
||||
enum Foo_enum
|
||||
{
|
||||
};
|
||||
|
||||
#if _CCCL_OS(WINDOWS)
|
||||
using FooEnum = int;
|
||||
#else
|
||||
using FooEnum = unsigned int;
|
||||
#endif
|
||||
|
||||
struct Foo_st;
|
||||
using FooStruct = Foo_st*;
|
||||
|
||||
// An enum whose underlying type is fixed, so that mismatch tests are deterministic
|
||||
// regardless of how the implementation picks the underlying type of an unfixed enum.
|
||||
enum class CharEnum : char
|
||||
{
|
||||
};
|
||||
|
||||
enum class IntEnum : int
|
||||
{
|
||||
};
|
||||
|
||||
struct Bar_st;
|
||||
// NOLINTEND(bugprone-reserved-identifier)
|
||||
} // namespace
|
||||
|
||||
C2H_TEST("nccl __abi_compatible type comparisons", "[multi_gpu][nccl]")
|
||||
{
|
||||
// --- Identical / scalar types -------------------------------------------------
|
||||
STATIC_REQUIRE(abi_detail::__abi_compatible<int, int>());
|
||||
STATIC_REQUIRE(!abi_detail::__abi_compatible<int, float>());
|
||||
|
||||
// remove_cv only strips the top level, so a top-level cv difference is still compatible.
|
||||
STATIC_REQUIRE(abi_detail::__abi_compatible<const int, int>());
|
||||
STATIC_REQUIRE(abi_detail::__abi_compatible<volatile int, int>());
|
||||
STATIC_REQUIRE(abi_detail::__abi_compatible<const volatile int, int>());
|
||||
|
||||
// Different-width / signedness integers are distinct types and must not be compatible.
|
||||
STATIC_REQUIRE(!abi_detail::__abi_compatible<int, long>());
|
||||
STATIC_REQUIRE(!abi_detail::__abi_compatible<int, short>());
|
||||
STATIC_REQUIRE(!abi_detail::__abi_compatible<int, unsigned int>());
|
||||
STATIC_REQUIRE(!abi_detail::__abi_compatible<char, signed char>());
|
||||
|
||||
// --- Pointers -----------------------------------------------------------------
|
||||
STATIC_REQUIRE(abi_detail::__abi_compatible<const char**, const char* const*>());
|
||||
STATIC_REQUIRE(!abi_detail::__abi_compatible<const char*, const int*>());
|
||||
|
||||
// Deep cv-qualifications on the pointee are stripped at each level of recursion.
|
||||
STATIC_REQUIRE(abi_detail::__abi_compatible<int*, const int*>());
|
||||
STATIC_REQUIRE(abi_detail::__abi_compatible<int**, const int* const*>());
|
||||
|
||||
// Multi-level pointers must agree in depth.
|
||||
STATIC_REQUIRE(abi_detail::__abi_compatible<int**, int**>());
|
||||
STATIC_REQUIRE(!abi_detail::__abi_compatible<int**, int*>());
|
||||
|
||||
// A pointer is never compatible with a non-pointer (the `&&` in the pointer branch).
|
||||
STATIC_REQUIRE(!abi_detail::__abi_compatible<int*, int>());
|
||||
STATIC_REQUIRE(!abi_detail::__abi_compatible<int, int*>());
|
||||
|
||||
// void* is not compatible with a typed pointer: the pointees differ.
|
||||
STATIC_REQUIRE(!abi_detail::__abi_compatible<void*, int*>());
|
||||
STATIC_REQUIRE(abi_detail::__abi_compatible<void*, void*>());
|
||||
|
||||
// Opaque struct pointers: identical handle types and matching opaque pointees.
|
||||
STATIC_REQUIRE(abi_detail::__abi_compatible<Foo_st*, FooStruct>());
|
||||
STATIC_REQUIRE(!abi_detail::__abi_compatible<Foo_st*, Bar_st*>());
|
||||
|
||||
// --- Enums --------------------------------------------------------------------
|
||||
STATIC_REQUIRE(::cuda::std::is_same_v<FooEnum, ::cuda::std::underlying_type_t<Foo_enum>>);
|
||||
|
||||
// enum vs its underlying type (either side), and enum-vs-enum through the pointee.
|
||||
STATIC_REQUIRE(abi_detail::__abi_compatible<FooEnum, Foo_enum>());
|
||||
STATIC_REQUIRE(abi_detail::__abi_compatible<Foo_enum, FooEnum>());
|
||||
STATIC_REQUIRE(abi_detail::__abi_compatible<FooEnum*, Foo_enum*>());
|
||||
STATIC_REQUIRE(!abi_detail::__abi_compatible<FooEnum*, Foo_enum>());
|
||||
|
||||
// Both sides enums: compatible iff their underlying types match exactly.
|
||||
STATIC_REQUIRE(abi_detail::__abi_compatible<CharEnum, CharEnum>());
|
||||
STATIC_REQUIRE(!abi_detail::__abi_compatible<CharEnum, IntEnum>());
|
||||
|
||||
// An enum is compatible with its exact underlying type but not a mismatched one.
|
||||
STATIC_REQUIRE(abi_detail::__abi_compatible<CharEnum, char>());
|
||||
STATIC_REQUIRE(!abi_detail::__abi_compatible<CharEnum, int>());
|
||||
STATIC_REQUIRE(abi_detail::__abi_compatible<IntEnum, int>());
|
||||
STATIC_REQUIRE(!abi_detail::__abi_compatible<IntEnum, char>());
|
||||
|
||||
// --- Function pointers --------------------------------------------------------
|
||||
STATIC_REQUIRE(abi_detail::__abi_compatible<int (*)(Foo_st*), int (*)(FooStruct)>());
|
||||
|
||||
// Return type, arity, and argument types must all match.
|
||||
STATIC_REQUIRE(!abi_detail::__abi_compatible<int (*)(int), float (*)(int)>());
|
||||
STATIC_REQUIRE(!abi_detail::__abi_compatible<int (*)(int), int (*)(int, int)>());
|
||||
STATIC_REQUIRE(!abi_detail::__abi_compatible<int (*)(int), int (*)(long)>());
|
||||
|
||||
// void return and no-argument functions.
|
||||
STATIC_REQUIRE(abi_detail::__abi_compatible<void (*)(), void (*)()>());
|
||||
STATIC_REQUIRE(!abi_detail::__abi_compatible<void (*)(), int (*)()>());
|
||||
|
||||
// Per-argument ABI compatibility recurses (enum / opaque-pointer arguments).
|
||||
STATIC_REQUIRE(abi_detail::__abi_compatible<void (*)(FooEnum), void (*)(Foo_enum)>());
|
||||
STATIC_REQUIRE(abi_detail::__abi_compatible<int (*)(Foo_st*, FooEnum), int (*)(FooStruct, Foo_enum)>());
|
||||
STATIC_REQUIRE(!abi_detail::__abi_compatible<void (*)(CharEnum), void (*)(IntEnum)>());
|
||||
|
||||
// A function type is not compatible with a non-function type.
|
||||
STATIC_REQUIRE(!abi_detail::__abi_compatible<int (*)(int), int*>());
|
||||
}
|
||||
@@ -1,100 +0,0 @@
|
||||
//===----------------------------------------------------------------------===//
|
||||
//
|
||||
// Part of CUDA Experimental in CUDA C++ Core Libraries,
|
||||
// under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://llvm.org/LICENSE.txt for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include <cuda/functional> // cuda::maximum, cuda::minimum
|
||||
#include <cuda/std/cstdint>
|
||||
#include <cuda/std/functional> // cuda::std::plus, cuda::std::multiplies
|
||||
|
||||
#include <cuda/experimental/__nccl/nccl_api.h>
|
||||
|
||||
#include <c2h/catch2_test_helper.h>
|
||||
|
||||
namespace
|
||||
{
|
||||
// A trivially-copyable aggregate with no corresponding NCCL data type.
|
||||
struct trivial_aggregate
|
||||
{
|
||||
int a;
|
||||
double b;
|
||||
};
|
||||
|
||||
// A non-trivially-copyable type (user-provided copy ctor).
|
||||
struct non_trivial
|
||||
{
|
||||
non_trivial(const non_trivial&) {} // NOLINT(modernize-use-equals-default)
|
||||
int a;
|
||||
};
|
||||
} // namespace
|
||||
|
||||
C2H_TEST("nccl type and reduction-op traits", "[multi_gpu][nccl]")
|
||||
{
|
||||
namespace nccl = ::cuda::experimental::__nccl;
|
||||
|
||||
// --- __has_nccl_type_of -------------------------------------------------------
|
||||
//
|
||||
// True only for the fixed set of types __nccl_type_of() recognizes.
|
||||
STATIC_REQUIRE(nccl::__has_nccl_type_of<::cuda::std::int8_t>);
|
||||
STATIC_REQUIRE(nccl::__has_nccl_type_of<::cuda::std::uint8_t>);
|
||||
STATIC_REQUIRE(nccl::__has_nccl_type_of<::cuda::std::int32_t>);
|
||||
STATIC_REQUIRE(nccl::__has_nccl_type_of<::cuda::std::uint32_t>);
|
||||
STATIC_REQUIRE(nccl::__has_nccl_type_of<::cuda::std::int64_t>);
|
||||
STATIC_REQUIRE(nccl::__has_nccl_type_of<::cuda::std::uint64_t>);
|
||||
STATIC_REQUIRE(nccl::__has_nccl_type_of<float>);
|
||||
STATIC_REQUIRE(nccl::__has_nccl_type_of<double>);
|
||||
STATIC_REQUIRE(nccl::__has_nccl_type_of<bool>);
|
||||
|
||||
// cv-qualifications are stripped via remove_cvref_t before the lookup.
|
||||
STATIC_REQUIRE(nccl::__has_nccl_type_of<const float>);
|
||||
STATIC_REQUIRE(nccl::__has_nccl_type_of<volatile int>);
|
||||
STATIC_REQUIRE(nccl::__has_nccl_type_of<float&>);
|
||||
|
||||
// Unsupported types have no NCCL data type.
|
||||
STATIC_REQUIRE(!nccl::__has_nccl_type_of<long double>);
|
||||
STATIC_REQUIRE(!nccl::__has_nccl_type_of<void>);
|
||||
STATIC_REQUIRE(!nccl::__has_nccl_type_of<int*>);
|
||||
STATIC_REQUIRE(!nccl::__has_nccl_type_of<trivial_aggregate>);
|
||||
|
||||
// --- __nccl_type_of_v ---------------------------------------------------------
|
||||
STATIC_REQUIRE(nccl::__nccl_type_of_v<bool> == nccl::__ncclChar);
|
||||
STATIC_REQUIRE(nccl::__nccl_type_of_v<::cuda::std::int8_t> == nccl::__ncclInt8);
|
||||
STATIC_REQUIRE(nccl::__nccl_type_of_v<::cuda::std::uint8_t> == nccl::__ncclUint8);
|
||||
STATIC_REQUIRE(nccl::__nccl_type_of_v<::cuda::std::int32_t> == nccl::__ncclInt32);
|
||||
STATIC_REQUIRE(nccl::__nccl_type_of_v<::cuda::std::uint32_t> == nccl::__ncclUint32);
|
||||
STATIC_REQUIRE(nccl::__nccl_type_of_v<::cuda::std::int64_t> == nccl::__ncclInt64);
|
||||
STATIC_REQUIRE(nccl::__nccl_type_of_v<::cuda::std::uint64_t> == nccl::__ncclUint64);
|
||||
STATIC_REQUIRE(nccl::__nccl_type_of_v<float> == nccl::__ncclFloat);
|
||||
STATIC_REQUIRE(nccl::__nccl_type_of_v<double> == nccl::__ncclDouble);
|
||||
|
||||
// cv-ref qualifiers are removed before mapping.
|
||||
STATIC_REQUIRE(nccl::__nccl_type_of_v<const double&> == nccl::__ncclDouble);
|
||||
|
||||
// --- __nccl_redop_of_v --------------------------------------------------------
|
||||
STATIC_REQUIRE(nccl::__nccl_redop_of_v<::cuda::std::plus<>> == nccl::__ncclSum);
|
||||
STATIC_REQUIRE(nccl::__nccl_redop_of_v<::cuda::std::multiplies<>> == nccl::__ncclProd);
|
||||
STATIC_REQUIRE(nccl::__nccl_redop_of_v<::cuda::maximum<>> == nccl::__ncclMax);
|
||||
STATIC_REQUIRE(nccl::__nccl_redop_of_v<::cuda::minimum<>> == nccl::__ncclMin);
|
||||
STATIC_REQUIRE(nccl::__nccl_redop_of_v<const ::cuda::std::plus<int>&> == nccl::__ncclSum);
|
||||
|
||||
// --- __has_nccl_redop ---------------------------------------------------------
|
||||
STATIC_REQUIRE(nccl::__has_nccl_redop_of<::cuda::std::plus<>>);
|
||||
STATIC_REQUIRE(nccl::__has_nccl_redop_of<::cuda::std::multiplies<>>);
|
||||
STATIC_REQUIRE(nccl::__has_nccl_redop_of<::cuda::maximum<>>);
|
||||
STATIC_REQUIRE(nccl::__has_nccl_redop_of<::cuda::minimum<>>);
|
||||
|
||||
// Typed specializations are recognized too.
|
||||
STATIC_REQUIRE(nccl::__has_nccl_redop_of<::cuda::std::plus<int>>);
|
||||
|
||||
// cv-ref qualifiers are stripped before the lookup.
|
||||
STATIC_REQUIRE(nccl::__has_nccl_redop_of<const ::cuda::std::plus<>&>);
|
||||
|
||||
// Unsupported functors and non-functor types are rejected.
|
||||
STATIC_REQUIRE(!nccl::__has_nccl_redop_of<::cuda::std::minus<>>);
|
||||
STATIC_REQUIRE(!nccl::__has_nccl_redop_of<int>);
|
||||
}
|
||||
Reference in New Issue
Block a user