Skip to content

[cudax] Make group rank a template parameter in group dispatch #12622

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

Open
davebayer wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[cudax] Make group rank a template parameter in group dispatch #12622

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
6 changes: 3 additions & 3 deletions cudax/include/cuda/experimental/coop/__group/coalesced_group.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -254,7 +254,7 @@ public:
}

template <class _Tp, class _MappingResult, class _InLevel>
[[nodiscard]] _CCCL_DEVICE_API static _Tp
[[nodiscard]] _CCCL_DEVICE_API static constexpr _Tp
__rank_as_impl(const _MappingResult&, const _Hierarchy& __hier, const _InLevel& __in_level)
{
if constexpr (::cuda::std::is_same_v<_InLevel, level_type>)
Expand All @@ -269,14 +269,14 @@ public:

_CCCL_TEMPLATE(class _Tp, class _InLevel)
_CCCL_REQUIRES(::cuda::std::__cccl_is_integer_v<_Tp> _CCCL_AND __is_hierarchy_level_v<_InLevel>)
[[nodiscard]] _CCCL_DEVICE_API _Tp rank_as(const _InLevel& __in_level) const
[[nodiscard]] _CCCL_DEVICE_API constexpr _Tp rank_as(const _InLevel& __in_level) const
{
return __rank_as_impl<_Tp>(__mapping_result_, __hier_, __in_level);
}

_CCCL_TEMPLATE(class _InLevel)
_CCCL_REQUIRES(__is_hierarchy_level_v<_InLevel>)
[[nodiscard]] _CCCL_DEVICE_API auto rank(const _InLevel& __in_level) const
[[nodiscard]] _CCCL_DEVICE_API constexpr auto rank(const _InLevel& __in_level) const
{
return __rank_as_impl<typename _InLevel::__product_type>(__mapping_result_, __hier_, __in_level);
}
Expand Down
4 changes: 2 additions & 2 deletions cudax/include/cuda/experimental/coop/__group/generic_group.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -205,12 +205,12 @@ public:

_CCCL_TEMPLATE(class _Tp)
_CCCL_REQUIRES(::cuda::std::__cccl_is_integer_v<_Tp>)
[[nodiscard]] _CCCL_DEVICE_API _Tp rank_as(const _ParentGroup& __parent) const
[[nodiscard]] _CCCL_DEVICE_API constexpr _Tp rank_as(const _ParentGroup& __parent) const
{
return __rank_as_impl<_Tp>(__mapping_result_, __hier_, __parent);
}

[[nodiscard]] _CCCL_DEVICE_API auto rank(const _ParentGroup& __parent) const
[[nodiscard]] _CCCL_DEVICE_API constexpr auto rank(const _ParentGroup& __parent) const
{
return __rank_as_impl<typename level_type::__product_type>(__mapping_result_, __hier_, __parent);
}
Expand Down
15 changes: 4 additions & 11 deletions cudax/include/cuda/experimental/coop/__group/group_view.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -138,20 +138,13 @@ _CCCL_DEVICE_API decltype(auto) __with_static_unit_count_range(const _Group& __g
}
else
{
using _StaticMappingResult = ::cuda::experimental::coop::__mapping_result<
__uniform_static_unit_count<_MappingResult::static_group_count(), _MappingResult::static_unit_count(_First)>,
using _PerGroupMappingResult = ::cuda::experimental::coop::__mapping_result<
__per_group_count<_MappingResult::static_group_count(), _First, _MappingResult::static_unit_count(_First)>,
_MappingResult::is_always_exhaustive(),
_MappingResult::is_always_contiguous()>;

const _StaticMappingResult __static_mapping_result{
__mapping_result.group_count(),
__mapping_result.group_rank(),
__mapping_result.unit_count(),
__mapping_result.unit_rank(),
__mapping_result.lane_mask()};

return ::cuda::std::forward<_Fn>(__fn)(
group_view<_Unit, _Group, _StaticMappingResult>{__group, __static_mapping_result});
return ::cuda::std::forward<_Fn>(__fn)(group_view<_Unit, _Group, _PerGroupMappingResult>{
__group, _PerGroupMappingResult{__mapping_result.unit_rank(), __mapping_result.lane_mask()}});
}
}

Expand Down
94 changes: 91 additions & 3 deletions cudax/include/cuda/experimental/coop/__group/mapping/mapping_result.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,9 @@ struct __uniform_static_unit_count;
template <::cuda::std::size_t... _StaticUnitCounts>
struct __static_unit_counts;

template <::cuda::std::size_t _StaticGroupCount, ::cuda::std::size_t _GroupRank, ::cuda::std::size_t _StaticUnitCount>
struct __per_group_count;

template <class _CountsTag, bool _IsAlwaysExhaustive, bool _IsAlwaysContiguous>
struct __mapping_result;

Expand Down Expand Up @@ -227,7 +230,7 @@ struct __mapping_result<__static_unit_counts<_StaticUnitCounts...>, _IsAlwaysExh
return static_cast<::cuda::std::uint32_t>(_StaticGroupCount);
}

[[nodiscard]] _CCCL_DEVICE_API ::cuda::std::uint32_t group_rank() const
[[nodiscard]] _CCCL_DEVICE_API constexpr ::cuda::std::uint32_t group_rank() const
{
if constexpr (!_IsAlwaysExhaustive)
{
Expand Down Expand Up @@ -256,8 +259,12 @@ struct __mapping_result<__static_unit_counts<_StaticUnitCounts...>, _IsAlwaysExh
[[nodiscard]] _CCCL_DEVICE_API static constexpr ::cuda::std::size_t static_unit_count(::cuda::std::size_t __i)
{
_CCCL_ASSERT(__i < _StaticGroupCount, "__i out of range");
constexpr ::cuda::std::size_t __counts[]{_StaticUnitCounts...};
return __counts[__i];

// Use fold expression rather than constexpr array to keep runtime queries in registers.
::cuda::std::size_t __index = 0;
::cuda::std::size_t __result = 0;
((__result = (__i == __index++) ? _StaticUnitCounts : __result), ...);

@Jacobfaib Jacobfaib Nov 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is so cursed...

@davebayer davebayer Nov 13, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just standard library things :D

return __result;
}

[[nodiscard]] _CCCL_DEVICE_API ::cuda::std::uint32_t unit_count() const
Expand Down Expand Up @@ -310,6 +317,87 @@ struct __mapping_result<__static_unit_counts<_StaticUnitCounts...>, _IsAlwaysExh
}
};

// This mapping result is only ever used in group.with_static_unit_count(...) where it's constructed from the group's
// mapping result. That means that the mapping result is always valid. We still keep the is_always_exhaustive() property
// from the original mapping result.
template <::cuda::std::size_t _StaticGroupCount,
::cuda::std::size_t _GroupRank,
::cuda::std::size_t _StaticUnitCount,
bool _IsAlwaysExhaustive,
bool _IsAlwaysContiguous>
struct __mapping_result<__per_group_count<_StaticGroupCount, _GroupRank, _StaticUnitCount>,
_IsAlwaysExhaustive,
_IsAlwaysContiguous>
{
static_assert(_StaticGroupCount != ::cuda::std::dynamic_extent, "group count must be statically known");
static_assert(_GroupRank < _StaticGroupCount, "group rank out of range");
static_assert(_StaticUnitCount != ::cuda::std::dynamic_extent, "unit count must be statically known");

::cuda::std::uint32_t __unit_rank_;
::cuda::device::lane_mask __lane_mask_;

[[nodiscard]] _CCCL_DEVICE_API static constexpr ::cuda::std::size_t static_group_count()
{
return _StaticGroupCount;
}

[[nodiscard]] _CCCL_DEVICE_API constexpr ::cuda::std::uint32_t group_count() const
{
return static_cast<::cuda::std::uint32_t>(_StaticGroupCount);
}

[[nodiscard]] _CCCL_DEVICE_API constexpr ::cuda::std::uint32_t group_rank() const
{
return static_cast<::cuda::std::uint32_t>(_GroupRank);
}

[[nodiscard]] _CCCL_DEVICE_API static constexpr bool has_uniform_static_unit_count()
{
return true;
}

[[nodiscard]] _CCCL_DEVICE_API static constexpr ::cuda::std::size_t static_unit_count()
{
return _StaticUnitCount;
}

[[nodiscard]] _CCCL_DEVICE_API static constexpr ::cuda::std::size_t static_unit_count(::cuda::std::size_t __i)
{
_CCCL_ASSERT(__i < _StaticGroupCount, "__i out of range");
return _StaticUnitCount;
}

[[nodiscard]] _CCCL_DEVICE_API constexpr ::cuda::std::uint32_t unit_count() const
{
return static_cast<::cuda::std::uint32_t>(_StaticUnitCount);

@Jacobfaib Jacobfaib Nov 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
return static_cast<::cuda::std::uint32_t>(_StaticUnitCount);
return ::cuda::std::uint32_t{_StaticUnitCount};

This form will error at compile time if _StaticUnitCount cannot fit in a uint32_t. Should do the same for group_rank() as well

@davebayer davebayer Nov 13, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It would already static_asserted by this point if the static unit count wouldn't be representable. But you are probably right, I'll fix it

}

[[nodiscard]] _CCCL_DEVICE_API ::cuda::std::uint32_t unit_rank() const
{
return __unit_rank_;
}

[[nodiscard]] _CCCL_DEVICE_API ::cuda::device::lane_mask lane_mask() const
{
return __lane_mask_;
}

[[nodiscard]] _CCCL_DEVICE_API constexpr bool is_valid() const
{
return true;
}

[[nodiscard]] _CCCL_DEVICE_API static constexpr bool is_always_exhaustive()
{
return _IsAlwaysExhaustive;
}

[[nodiscard]] _CCCL_DEVICE_API static constexpr bool is_always_contiguous()
{
return _IsAlwaysContiguous;
}
};

template <bool _IsAlwaysContiguous>
[[nodiscard]] _CCCL_DEVICE_API ::cuda::device::lane_mask __make_lane_mask_for_n(
::cuda::device::lane_mask __prev_lane_mask, ::cuda::std::uint32_t __n, ::cuda::std::uint32_t __rank)
Expand Down
6 changes: 3 additions & 3 deletions cudax/include/cuda/experimental/coop/__group/this_group.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -254,7 +254,7 @@ public:
}

template <class _Tp, class _MappingResult, class _InLevel>
[[nodiscard]] _CCCL_DEVICE_API static _Tp
[[nodiscard]] _CCCL_DEVICE_API static constexpr _Tp
__rank_as_impl(const _MappingResult&, const _Hierarchy& __hier, const _InLevel& __in_level)
{
if constexpr (::cuda::std::is_same_v<_InLevel, _Level>)
Expand All @@ -269,14 +269,14 @@ public:

_CCCL_TEMPLATE(class _Tp, class _InLevel)
_CCCL_REQUIRES(::cuda::std::__cccl_is_integer_v<_Tp> _CCCL_AND __is_hierarchy_level_v<_InLevel>)
[[nodiscard]] _CCCL_DEVICE_API _Tp rank_as(const _InLevel& __in_level) const
[[nodiscard]] _CCCL_DEVICE_API constexpr _Tp rank_as(const _InLevel& __in_level) const
{
return __rank_as_impl<_Tp>(__mapping_result_, __hier_, __in_level);
}

_CCCL_TEMPLATE(class _InLevel)
_CCCL_REQUIRES(__is_hierarchy_level_v<_InLevel>)
[[nodiscard]] _CCCL_DEVICE_API auto rank(const _InLevel& __in_level) const
[[nodiscard]] _CCCL_DEVICE_API constexpr auto rank(const _InLevel& __in_level) const
{
return __rank_as_impl<typename _InLevel::__product_type>(__mapping_result_, __hier_, __in_level);
}
Expand Down
4 changes: 2 additions & 2 deletions cudax/include/cuda/experimental/coop/__group/virtual_group.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -250,12 +250,12 @@ public:

_CCCL_TEMPLATE(class _Tp)
_CCCL_REQUIRES(::cuda::std::__cccl_is_integer_v<_Tp>)
[[nodiscard]] _CCCL_DEVICE_API _Tp rank_as(const _ParentGroup& __parent) const
[[nodiscard]] _CCCL_DEVICE_API constexpr _Tp rank_as(const _ParentGroup& __parent) const
{
return __rank_as_impl<_Tp>(__mapping_result_, __hier_, __parent);
}

[[nodiscard]] _CCCL_DEVICE_API auto rank(const _ParentGroup& __parent) const
[[nodiscard]] _CCCL_DEVICE_API constexpr auto rank(const _ParentGroup& __parent) const
{
return __rank_as_impl<typename level_type::__product_type>(__mapping_result_, __hier_, __parent);
}
Expand Down
39 changes: 39 additions & 0 deletions cudax/test/coop/group/mapping/non_uniform_mapping.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,40 @@ __device__ void test_result_storage()
CHECK(uniform.unit_count() == 8);
}

template <unsigned GroupRank, bool Exhaustive, bool Contiguous>
__device__ void test_per_group_result()
{
using Result =
cudax::coop::__mapping_result<cudax::coop::__per_group_count<3, GroupRank, 12>, Exhaustive, Contiguous>;

static_assert(cudax::coop::__group_mapping_result<Result>);
static_assert(Result::static_group_count() == 3);
static_assert(Result::has_uniform_static_unit_count());
static_assert(Result::static_unit_count() == 12);
static_assert(Result::static_unit_count(GroupRank) == 12);
static_assert(Result::is_always_exhaustive() == Exhaustive);
static_assert(Result::is_always_contiguous() == Contiguous);

for (unsigned rank = 0; rank < 12; ++rank)
{
const auto mask = cuda::device::lane_mask{0x55u};
const Result result{rank, mask};

CHECK(result.is_valid());
CHECK(result.group_count() == 3);
CHECK(result.group_rank() == GroupRank);
CHECK(result.unit_count() == 12);
CHECK(result.unit_rank() == rank);
CHECK(result.lane_mask() == mask);
}

if constexpr (Exhaustive)
{
constexpr Result result{11, cuda::device::lane_mask::all()};
static_assert(result.group_rank() == GroupRank);
}
}

template <unsigned N, bool Static, bool Exhaustive, bool Contiguous = true, class Config>
__device__ void test_group_by_counts(Config config)
{
Expand Down Expand Up @@ -286,6 +320,11 @@ struct TestKernel
{
test_other_mappings(config);

test_per_group_result<0, true, true>();
test_per_group_result<1, true, false>();
test_per_group_result<2, false, true>();
test_per_group_result<2, false, false>();

// Storage must preserve both exhaustive and contiguous guarantees independently.
test_result_storage<true, true>();
test_result_storage<true, false>();
Expand Down
44 changes: 44 additions & 0 deletions cudax/test/coop/group/with_static_unit_count.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,8 @@ template <class Config, class Level>
__device__ void test_hardware_group(Config config, Level level)
{
const auto group = cudax::coop::make_this_group(level, config);
static_assert(group.rank(level) == 0);
static_assert(group.template rank_as<unsigned>(level) == 0);
test_uniform(group);

// Changing to thread units must preserve any available static count.
Expand Down Expand Up @@ -194,6 +196,19 @@ __device__ void test_nonuniform(const Group& group, unsigned units_per_warp)
static_assert(count == Original::static_unit_count(0) || count == Original::static_unit_count(1));
CHECK(count == expected);

// The selected rank is part of the result type, even when two groups share a count.
using Mapping = typename decltype(view)::__mapping_result_type;
static_assert(cudax::coop::__group_mapping_result<Mapping>);
constexpr auto group_rank = Mapping{0, cuda::device::lane_mask::all()}.group_rank();
static_assert(count == Original::static_unit_count(group_rank));
static_assert(
cuda::std::is_same_v<
Mapping,
cudax::coop::__mapping_result<cudax::coop::__per_group_count<Original::static_group_count(), group_rank, count>,
Original::is_always_exhaustive(),
Original::is_always_contiguous()>>);
CHECK(group_rank == group.__mapping_result().group_rank());

check_mapping(group, view);

// Copy deduction must retain the specialized view type.
Expand Down Expand Up @@ -234,6 +249,22 @@ __device__ void test_nonuniform_groups(Config config, Counts counts)
// Check nonuniform mappings while retaining the parent block synchronizer.
const cudax::coop::virtual_group virtual_group{cuda::warp, parent, mapping};
test_nonuniform(virtual_group, 1);
virtual_group.with_static_unit_count([&](auto view) {
using Mapping = typename decltype(view)::__mapping_result_type;
constexpr auto rank = Mapping{0, cuda::device::lane_mask::all()}.group_rank();

// NVCC 12.0 cannot constant-evaluate this query through a lambda capture.
#if _CCCL_CUDA_COMPILER(NVCC, ==, 12, 0)
const cudax::coop::this_block rank_parent{config};
static_assert(view.rank(rank_parent) == rank);
static_assert(view.template rank_as<unsigned>(rank_parent) == rank);
#else // _CCCL_CUDA_COMPILER(NVCC, ==, 12, 0)
static_assert(view.rank(parent) == rank);
static_assert(view.template rank_as<unsigned>(parent) == rank);
#endif // _CCCL_CUDA_COMPILER(NVCC, ==, 12, 0)

CHECK(view.rank(parent) == virtual_group.rank(parent));
});
CHECK(virtual_group.count(parent) == 3);

const cudax::coop::group_view threads{cuda::gpu_thread, virtual_group};
Expand All @@ -251,6 +282,19 @@ __device__ void test_nonuniform_groups(Config config, Counts counts)
// Super-group queries and synchronization must survive specialization too.
group.with_static_unit_count([&](auto view) {
CHECK(view.count(parent) == group.count(parent));
using Mapping = typename decltype(view)::__mapping_result_type;
constexpr auto rank = Mapping{0, cuda::device::lane_mask::all()}.group_rank();

// NVCC 12.0 cannot constant-evaluate this query through a lambda capture.
#if _CCCL_CUDA_COMPILER(NVCC, ==, 12, 0)
const cudax::coop::this_block rank_parent{config};
static_assert(view.rank(rank_parent) == rank);
static_assert(view.template rank_as<unsigned>(rank_parent) == rank);
#else // _CCCL_CUDA_COMPILER(NVCC, ==, 12, 0)
static_assert(view.rank(parent) == rank);
static_assert(view.template rank_as<unsigned>(parent) == rank);
#endif // _CCCL_CUDA_COMPILER(NVCC, ==, 12, 0)

CHECK(view.rank(parent) == group.rank(parent));
view.sync();
});
Expand Down
Loading