Skip to content

Use result policies for MGMN algorithms #9979

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Jacobfaib merged 1 commit into from
Jul 17, 2026
Merged

Use result policies for MGMN algorithms #9979

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 15 additions & 2 deletions cudax/include/cuda/experimental/__multi_gpu/algorithm/reduce/reduce.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@

#include <cuda/experimental/__multi_gpu/algorithm/common.h>
#include <cuda/experimental/__multi_gpu/concepts.h>
#include <cuda/experimental/__utility/result_policy.cuh>

#include <vector>

Expand Down Expand Up @@ -231,6 +232,7 @@ _CCCL_CONCEPT __indirectly_binary_reducible = _CCCL_REQUIRES_EXPR((_Fn, _Tp, _It
//! element is 0. For maximum and minimum, the identity values are INT_MIN, and INT_MAX
//! respectively.
//!
//! @tparam _Policy The return value policy. Currently only `broadcasted_t` is supported.
//! @tparam _CommRange The range of communicators. Each element must model the communicator
//! concept.
//! @tparam _EnvRange The range of execution environments. Each environment supplies the
Expand All @@ -242,6 +244,7 @@ _CCCL_CONCEPT __indirectly_binary_reducible = _CCCL_REQUIRES_EXPR((_Fn, _Tp, _It
//! element type.
//! @tparam _BinaryOp The binary reduction operator type. Defaults to `::cuda::std::plus<>`.
//!
//! @param[in] __policy The result policy object.
//! @param[in] __comms The range of communicators.
//! @param[in] __envs The range of execution environments. The execution environment must
//! contain a stream.
Expand All @@ -250,7 +253,8 @@ _CCCL_CONCEPT __indirectly_binary_reducible = _CCCL_REQUIRES_EXPR((_Fn, _Tp, _It
//! @param[in] __init The initial value seeding each reduction.
//! @param[in] __op The binary reduction operator.
//! @param[in] __ident The identity element to be used in case of empty ranges.
_CCCL_TEMPLATE(class _CommRange,
_CCCL_TEMPLATE(class _Policy,
class _CommRange,
class _EnvRange,
class _InputRangeRange,
class _OutputItRange,
Expand All @@ -260,6 +264,7 @@ _CCCL_REQUIRES(__range_of_communicators<_CommRange> _CCCL_AND ::cuda::std::range
__detail::__range_of_sized_random_access_ranges<_InputRangeRange> _CCCL_AND
__detail::__range_of_output_iters<_OutputItRange, _Tp>)
_CCCL_HOST_API void reduce(
[[maybe_unused]] const __result_policy_base<_Policy>& __policy,
_CommRange&& __comms,
_EnvRange&& __envs,
_InputRangeRange&& __range_of_inputs,
Expand All @@ -269,6 +274,9 @@ _CCCL_HOST_API void reduce(
_Tp __ident = ::cuda::identity_element<_BinaryOp, _Tp>())
{
static_assert(::cuda::std::ranges::sized_range<_CommRange>);
static_assert(::cuda::std::same_as<_Policy, broadcasted_t>,
"Only broadcasted results are currently supported. Please open an issue at "
"github.com/NVIDIA/cccl/issue requesting support for your specified policy.");

using __properties =
::cuda::experimental::__detail::__in_range_out_it_properties<_InputRangeRange, _OutputItRange, _EnvRange>;
Expand Down Expand Up @@ -324,6 +332,7 @@ _CCCL_HOST_API void reduce(
//! output iterator)` to the range-based overload. See the range overload for further
//! discussion.
//!
//! @tparam _Policy The return value policy. Currently only `broadcasted_t` is supported.
//! @tparam _Comm The communicator type. Must model the communicator concept.
//! @tparam _Env The execution environment type. Supplies the stream and optional memory
//! resource.
Expand All @@ -332,14 +341,16 @@ _CCCL_HOST_API void reduce(
//! @tparam _Tp The reduction and result value type.
//! @tparam _BinaryOp The binary reduction operator type. Defaults to `::cuda::std::plus<>`.
//!
//! @param[in] __policy The result policy object.
//! @param[in] __comm The communicator.
//! @param[in] __env The execution environment. Must contain a stream.
//! @param[in] __input The input range to reduce.
//! @param[out] __output The output iterator receiving the result.
//! @param[in] __init The initial value seeding the reduction.
//! @param[in] __op The binary reduction operator.
//! @param[in] __ident The identity element to be used in case of empty ranges.
_CCCL_TEMPLATE(class _Comm,
_CCCL_TEMPLATE(class _Policy,
class _Comm,
class _Env,
class _InputRange,
class _OutputIt,
Expand All @@ -348,6 +359,7 @@ _CCCL_TEMPLATE(class _Comm,
_CCCL_REQUIRES(__communicator<_Comm> _CCCL_AND ::cuda::std::ranges::random_access_range<_InputRange>
_CCCL_AND ::cuda::std::output_iterator<_OutputIt, _Tp>)
_CCCL_HOST_API void reduce(
const __result_policy_base<_Policy>& __policy,
_Comm&& __comm,
_Env&& __env,
_InputRange&& __input,
Expand All @@ -357,6 +369,7 @@ _CCCL_HOST_API void reduce(
_Tp __ident = ::cuda::identity_element<_BinaryOp, _Tp>())
{
::cuda::experimental::reduce(
__policy,
::cuda::std::span<::cuda::std::remove_reference_t<_Comm>, 1>{::cuda::std::addressof(__comm), 1},
::cuda::std::span<::cuda::std::remove_reference_t<_Env>, 1>{::cuda::std::addressof(__env), 1},
::cuda::std::span<::cuda::std::remove_reference_t<_InputRange>, 1>{::cuda::std::addressof(__input), 1},
Expand Down
34 changes: 30 additions & 4 deletions cudax/include/cuda/experimental/__multi_gpu/algorithm/scan/scan.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@
#include <cuda/experimental/__multi_gpu/algorithm/common.h>
#include <cuda/experimental/__multi_gpu/algorithm/reduce/reduce.h>
#include <cuda/experimental/__multi_gpu/concepts.h>
#include <cuda/experimental/__utility/result_policy.cuh>

#include <vector>

Expand All @@ -65,13 +66,15 @@ enum class __kind : ::cuda::std::uint8_t
};

template <__kind _Kind,
class _Policy,
class _CommRange,
class _EnvRange,
class _InputRangeRange,
class _OutputItRange,
class _Tp,
class _BinaryOp>
_CCCL_HOST_API void __scan(
const __result_policy_base<_Policy>&,
_CommRange&& __comms,
_EnvRange&& __envs,
_InputRangeRange&& __range_of_inputs,
Expand All @@ -81,6 +84,9 @@ _CCCL_HOST_API void __scan(
_Tp __ident)
{
static_assert(::cuda::std::ranges::sized_range<_CommRange>);
static_assert(::cuda::std::same_as<_Policy, distributed_t>,
"Only distributed results are currently supported. Please open an issue at "
"github.com/NVIDIA/cccl/issue requesting support for your specified policy.");

using __properties =
::cuda::experimental::__detail::__in_range_out_it_properties<_InputRangeRange, _OutputItRange, _EnvRange>;
Expand Down Expand Up @@ -251,6 +257,7 @@ _CCCL_HOST_API void __scan(
//! Each environment supplies the stream and optional memory resource for its local rank. The
//! results are ready after the work enqueued on the corresponding streams completes.
//!
//! @tparam _Policy The return value policy. Currently only `distributed_t` is supported.
//! @tparam _CommRange The range of communicators. Each element must model the communicator
//! concept.
//! @tparam _EnvRange The range of execution environments.
Expand All @@ -260,6 +267,7 @@ _CCCL_HOST_API void __scan(
//! @tparam _Tp The scan value type. Deduced by default from the input element type.
//! @tparam _BinaryOp The binary scan operator type. Defaults to `::cuda::std::plus<>`.
//!
//! @param[in] __policy The result policy object.
//! @param[in] __comms The range of communicators.
//! @param[in] __envs The range of execution environments. Each environment must contain a
//! stream.
Expand All @@ -268,7 +276,8 @@ _CCCL_HOST_API void __scan(
//! @param[in] __init The initial value for the scan.
//! @param[in] __op The associative binary scan operator.
//! @param[in] __ident The identity element for `__op`.
_CCCL_TEMPLATE(class _CommRange,
_CCCL_TEMPLATE(class _Policy,
class _CommRange,
class _EnvRange,
class _InputRangeRange,
class _OutputItRange,
Expand All @@ -278,6 +287,7 @@ _CCCL_REQUIRES(__range_of_communicators<_CommRange> _CCCL_AND ::cuda::std::range
__detail::__range_of_sized_random_access_ranges<_InputRangeRange> _CCCL_AND
__detail::__range_of_output_iters<_OutputItRange, _Tp>)
_CCCL_HOST_API void exclusive_scan(
const __result_policy_base<_Policy>& __policy,
_CommRange&& __comms,
_EnvRange&& __envs,
_InputRangeRange&& __range_of_inputs,
Expand All @@ -287,6 +297,7 @@ _CCCL_HOST_API void exclusive_scan(
_Tp __ident = ::cuda::identity_element<_BinaryOp, _Tp>())
{
::cuda::experimental::__detail::__scan::__scan<::cuda::experimental::__detail::__scan::__kind::__exclusive>(
__policy,
::cuda::std::forward<_CommRange>(__comms),
::cuda::std::forward<_EnvRange>(__envs),
::cuda::std::forward<_InputRangeRange>(__range_of_inputs),
Expand Down Expand Up @@ -318,6 +329,7 @@ _CCCL_HOST_API void exclusive_scan(
//! The environment supplies the stream and optional memory resource for the local rank. The
//! results are ready after the work enqueued on that stream completes.
//!
//! @tparam _Policy The return value policy. Currently only `distributed_t` is supported.
//! @tparam _Comm The communicator type. Must model the communicator concept.
//! @tparam _Env The execution environment type. Supplies the stream and optional memory
//! resource.
Expand All @@ -326,14 +338,16 @@ _CCCL_HOST_API void exclusive_scan(
//! @tparam _Tp The scan value type. Deduced by default from the input element type.
//! @tparam _BinaryOp The binary scan operator type. Defaults to `::cuda::std::plus<>`.
//!
//! @param[in] __policy The result policy object.
//! @param[in] __comm The communicator.
//! @param[in] __env The execution environment. Must contain a stream.
//! @param[in] __input The input range to scan.
//! @param[out] __output The output iterator receiving the local scan results.
//! @param[in] __init The initial value for the scan.
//! @param[in] __op The associative binary scan operator.
//! @param[in] __ident The identity element for `__op`.
_CCCL_TEMPLATE(class _Comm,
_CCCL_TEMPLATE(class _Policy,
class _Comm,
class _Env,
class _InputRange,
class _OutputIt,
Expand All @@ -342,6 +356,7 @@ _CCCL_TEMPLATE(class _Comm,
_CCCL_REQUIRES(__communicator<_Comm> _CCCL_AND ::cuda::std::ranges::random_access_range<_InputRange>
_CCCL_AND ::cuda::std::output_iterator<_OutputIt, _Tp>)
_CCCL_HOST_API void exclusive_scan(
const __result_policy_base<_Policy>& __policy,
_Comm&& __comm,
_Env&& __env,
_InputRange&& __input,
Expand All @@ -351,6 +366,7 @@ _CCCL_HOST_API void exclusive_scan(
_Tp __ident = ::cuda::identity_element<_BinaryOp, _Tp>())
{
::cuda::experimental::exclusive_scan(
__policy,
::cuda::std::span<::cuda::std::remove_reference_t<_Comm>, 1>{::cuda::std::addressof(__comm), 1},
::cuda::std::span<::cuda::std::remove_reference_t<_Env>, 1>{::cuda::std::addressof(__env), 1},
::cuda::std::span<::cuda::std::remove_reference_t<_InputRange>, 1>{::cuda::std::addressof(__input), 1},
Expand Down Expand Up @@ -388,6 +404,7 @@ _CCCL_HOST_API void exclusive_scan(
//! Each environment supplies the stream and optional memory resource for its local rank. The
//! results are ready after the work enqueued on the corresponding streams completes.
//!
//! @tparam _Policy The return value policy. Currently only `distributed_t` is supported.
//! @tparam _CommRange The range of communicators. Each element must model the communicator
//! concept.
//! @tparam _EnvRange The range of execution environments.
Expand All @@ -397,6 +414,7 @@ _CCCL_HOST_API void exclusive_scan(
//! @tparam _Tp The scan value type. Deduced by default from the input element type.
//! @tparam _BinaryOp The binary scan operator type. Defaults to `::cuda::std::plus<>`.
//!
//! @param[in] __policy The result policy object.
//! @param[in] __comms The range of communicators.
//! @param[in] __envs The range of execution environments. Each environment must contain a
//! stream.
Expand All @@ -405,7 +423,8 @@ _CCCL_HOST_API void exclusive_scan(
//! @param[in] __init The initial value for the scan.
//! @param[in] __op The associative binary scan operator.
//! @param[in] __ident The identity element for `__op`.
_CCCL_TEMPLATE(class _CommRange,
_CCCL_TEMPLATE(class _Policy,
class _CommRange,
class _EnvRange,
class _InputRangeRange,
class _OutputItRange,
Expand All @@ -415,6 +434,7 @@ _CCCL_REQUIRES(__range_of_communicators<_CommRange> _CCCL_AND ::cuda::std::range
__detail::__range_of_sized_random_access_ranges<_InputRangeRange> _CCCL_AND
__detail::__range_of_output_iters<_OutputItRange, _Tp>)
_CCCL_HOST_API void inclusive_scan(
const __result_policy_base<_Policy>& __policy,
_CommRange&& __comms,
_EnvRange&& __envs,
_InputRangeRange&& __range_of_inputs,
Expand All @@ -424,6 +444,7 @@ _CCCL_HOST_API void inclusive_scan(
_Tp __ident = ::cuda::identity_element<_BinaryOp, _Tp>())
{
::cuda::experimental::__detail::__scan::__scan<::cuda::experimental::__detail::__scan::__kind::__inclusive>(
__policy,
::cuda::std::forward<_CommRange>(__comms),
::cuda::std::forward<_EnvRange>(__envs),
::cuda::std::forward<_InputRangeRange>(__range_of_inputs),
Expand Down Expand Up @@ -455,6 +476,7 @@ _CCCL_HOST_API void inclusive_scan(
//! The environment supplies the stream and optional memory resource for the local rank. The
//! results are ready after the work enqueued on that stream completes.
//!
//! @tparam _Policy The return value policy. Currently only `distributed_t` is supported.
//! @tparam _Comm The communicator type. Must model the communicator concept.
//! @tparam _Env The execution environment type. Supplies the stream and optional memory
//! resource.
Expand All @@ -463,14 +485,16 @@ _CCCL_HOST_API void inclusive_scan(
//! @tparam _Tp The scan value type. Deduced by default from the input element type.
//! @tparam _BinaryOp The binary scan operator type. Defaults to `::cuda::std::plus<>`.
//!
//! @param[in] __policy The result policy object.
//! @param[in] __comm The communicator.
//! @param[in] __env The execution environment. Must contain a stream.
//! @param[in] __input The input range to scan.
//! @param[out] __output The output iterator receiving the local scan results.
//! @param[in] __init The initial value for the scan.
//! @param[in] __op The associative binary scan operator.
//! @param[in] __ident The identity element for `__op`.
_CCCL_TEMPLATE(class _Comm,
_CCCL_TEMPLATE(class _Policy,
class _Comm,
class _Env,
class _InputRange,
class _OutputIt,
Expand All @@ -479,6 +503,7 @@ _CCCL_TEMPLATE(class _Comm,
_CCCL_REQUIRES(__communicator<_Comm> _CCCL_AND ::cuda::std::ranges::random_access_range<_InputRange>
_CCCL_AND ::cuda::std::output_iterator<_OutputIt, _Tp>)
_CCCL_HOST_API void inclusive_scan(
const __result_policy_base<_Policy>& __policy,
_Comm&& __comm,
_Env&& __env,
_InputRange&& __input,
Expand All @@ -488,6 +513,7 @@ _CCCL_HOST_API void inclusive_scan(
_Tp __ident = ::cuda::identity_element<_BinaryOp, _Tp>())
{
::cuda::experimental::inclusive_scan(
__policy,
::cuda::std::span<::cuda::std::remove_reference_t<_Comm>, 1>{::cuda::std::addressof(__comm), 1},
::cuda::std::span<::cuda::std::remove_reference_t<_Env>, 1>{::cuda::std::addressof(__env), 1},
::cuda::std::span<::cuda::std::remove_reference_t<_InputRange>, 1>{::cuda::std::addressof(__input), 1},
Expand Down
3 changes: 2 additions & 1 deletion cudax/test/multi_gpu/algorithms/exclusive_scan/range_basic.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -123,7 +123,7 @@ void do_exclusive_scan(
INFO("init = " << init);
INFO("ident = " << ident);

cudax::exclusive_scan(comms, envs, in, outputs, init, op, 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);
Expand Down Expand Up @@ -167,6 +167,7 @@ MULTI_GPU_TEST("exclusive_scan documentation example", c2h::type_list<int>)
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,
Expand Down
8 changes: 4 additions & 4 deletions cudax/test/multi_gpu/algorithms/exclusive_scan/range_defaults.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,25 +88,25 @@ MULTI_GPU_TEST("exclusive_scan, range overloads default values", )

SECTION("Default init, op, ident (all)")
{
cudax::exclusive_scan(comms, envs, in, outputs);
cudax::exclusive_scan(cudax::distributed, comms, envs, in, outputs);
check_outputs();
}

SECTION("Default op, ident")
{
cudax::exclusive_scan(comms, envs, in, outputs, init);
cudax::exclusive_scan(cudax::distributed, comms, envs, in, outputs, init);
check_outputs();
}

SECTION("Default ident")
{
cudax::exclusive_scan(comms, envs, in, outputs, init, op);
cudax::exclusive_scan(cudax::distributed, comms, envs, in, outputs, init, op);
check_outputs();
}

SECTION("Default none")
{
cudax::exclusive_scan(comms, envs, in, outputs, init, op, ident);
cudax::exclusive_scan(cudax::distributed, comms, envs, in, outputs, init, op, ident);
check_outputs();
}
}
4 changes: 2 additions & 2 deletions cudax/test/multi_gpu/algorithms/exclusive_scan/single_comm_basic.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,7 @@ void run_case(cuda::std::span<cudax::nccl_communicator_ref> comms,
INFO("ident = " << ident);

run_threaded(comms.size(), [&](cuda::std::size_t i) {
cudax::exclusive_scan(comms[i], envs[i], in[i], outputs[i], init, op, ident);
cudax::exclusive_scan(cudax::distributed, comms[i], envs[i], in[i], outputs[i], init, op, ident);
});

REQUIRE(in.size() == in_copy.size());
Expand Down Expand Up @@ -184,7 +184,7 @@ MULTI_GPU_TEST("exclusive_scan single-comm documentation example", c2h::type_lis
auto input = cuda::make_device_buffer<int>(environment, device, input_values);
auto output = cuda::make_device_buffer<int>(environment, device, input_values.size(), cuda::no_init);

cudax::exclusive_scan(communicator, environment, input, output.begin(), /*__init=*/0);
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();
Expand Down
8 changes: 4 additions & 4 deletions cudax/test/multi_gpu/algorithms/exclusive_scan/single_comm_defaults.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -89,31 +89,31 @@ MULTI_GPU_TEST("exclusive_scan single-comm, overloads default values", )
SECTION("Default init, op, ident (all)")
{
run_threaded(comms.size(), [&](cuda::std::size_t i) {
cudax::exclusive_scan(comms[i], envs[i], in[i], outputs[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(comms[i], envs[i], in[i], outputs[i], init);
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(comms[i], envs[i], in[i], outputs[i], init, op);
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(comms[i], envs[i], in[i], outputs[i], init, op, ident);
cudax::exclusive_scan(cudax::distributed, comms[i], envs[i], in[i], outputs[i], init, op, ident);
});
check_outputs();
}
Expand Down
Loading
Loading