fix(build): remove 1106 extra CCCL headers — keep only 288 needed by allocator
commit 4c365b8c added 1106 CCCL device-level headers (294K lines) to
qwen3_6_scripts/cccl_preload/include/. These are NOT used by the allocator
preload (which only needs cub/util_allocator.cuh + 288 transitive deps)
and cannot compile on corex CUDA 10.2 anyway.
The extra headers doubled docker context from 15MB to 31MB, likely
causing platform build timeout or size limit failure.
Restoring to the original 288-header set that is proven to compile and run.
This commit is contained in:
@@ -1,54 +0,0 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
|
||||
// SPDX-License-Identifier: BSD-3
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cub/detail/type_traits.cuh> // static_size_v
|
||||
#include <cub/util_namespace.cuh>
|
||||
|
||||
#include <cuda/std/__iterator/iterator_traits.h>
|
||||
#include <cuda/std/__type_traits/conditional.h>
|
||||
#include <cuda/std/__type_traits/is_same.h>
|
||||
#include <cuda/std/__utility/integer_sequence.h>
|
||||
#include <cuda/std/array>
|
||||
#include <cuda/std/cstddef>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
namespace detail
|
||||
{
|
||||
#ifndef _CCCL_DOXYGEN_INVOKED // Do not document
|
||||
|
||||
/***********************************************************************************************************************
|
||||
* Generic Array-like to Array Conversion
|
||||
**********************************************************************************************************************/
|
||||
|
||||
template <typename CastType, typename Input, ::cuda::std::size_t... i>
|
||||
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::array<CastType, static_size_v<Input>>
|
||||
to_array_impl(const Input& input, ::cuda::std::index_sequence<i...>)
|
||||
{
|
||||
using ArrayType = ::cuda::std::array<CastType, static_size_v<Input>>;
|
||||
return ArrayType{static_cast<CastType>(input[i])...};
|
||||
}
|
||||
|
||||
template <typename CastType = void, typename Input>
|
||||
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::array<CastType, static_size_v<Input>>
|
||||
to_array(const Input& input)
|
||||
{
|
||||
using InputType = ::cuda::std::iter_value_t<Input>;
|
||||
using CastType1 = ::cuda::std::_If<::cuda::std::is_same_v<CastType, void>, InputType, CastType>;
|
||||
return to_array_impl<CastType1>(input, ::cuda::std::make_index_sequence<static_size_v<Input>>{});
|
||||
}
|
||||
|
||||
#endif // !_CCCL_DOXYGEN_INVOKED
|
||||
} // namespace detail
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,126 +0,0 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cuda/std/__algorithm/lower_bound.h>
|
||||
#include <cuda/std/__algorithm/upper_bound.h>
|
||||
#include <cuda/std/__iterator/iterator_traits.h>
|
||||
#include <cuda/std/cstddef>
|
||||
#include <cuda/std/tuple>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail::find
|
||||
{
|
||||
constexpr ::cuda::std::ptrdiff_t linear_lower_bound_threshold = 8;
|
||||
|
||||
template <typename RangeIteratorT, typename RangeNumItemsT, typename CompareOpT, typename Mode>
|
||||
struct comp_wrapper_t
|
||||
{
|
||||
RangeIteratorT first;
|
||||
RangeNumItemsT num_items;
|
||||
CompareOpT op;
|
||||
|
||||
template <typename Value, typename Output>
|
||||
_CCCL_DEVICE _CCCL_FORCEINLINE void operator()(::cuda::std::tuple<Value, Output> args) const
|
||||
{
|
||||
using DifferenceT = ::cuda::std::iter_difference_t<RangeIteratorT>;
|
||||
const auto last = first + static_cast<DifferenceT>(num_items);
|
||||
|
||||
::cuda::std::get<1>(args) = Mode::Invoke(first, last, ::cuda::std::get<0>(args), op);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename Mode, typename RangeIteratorT, typename RangeNumItemsT, typename CompareOpT>
|
||||
_CCCL_HOST_DEVICE auto make_comp_wrapper(RangeIteratorT first, RangeNumItemsT num_items, CompareOpT comp)
|
||||
{
|
||||
return comp_wrapper_t<RangeIteratorT, RangeNumItemsT, CompareOpT, Mode>{first, num_items, comp};
|
||||
}
|
||||
|
||||
struct lower_bound
|
||||
{
|
||||
template <typename RangeIteratorT, typename DifferenceT, typename T, typename CompareOpT>
|
||||
_CCCL_DEVICE _CCCL_FORCEINLINE static DifferenceT
|
||||
Linear(RangeIteratorT first, DifferenceT num_items, const T& value, CompareOpT comp)
|
||||
{
|
||||
DifferenceT retval = 0;
|
||||
for (DifferenceT i = 0; i < num_items; ++i)
|
||||
{
|
||||
retval += static_cast<DifferenceT>(comp(first[i], value));
|
||||
}
|
||||
|
||||
return retval;
|
||||
}
|
||||
|
||||
template <typename RangeIteratorT, typename T, typename CompareOpT>
|
||||
_CCCL_DEVICE _CCCL_FORCEINLINE static ::cuda::std::ptrdiff_t
|
||||
Invoke(RangeIteratorT first, RangeIteratorT last, const T& value, CompareOpT comp)
|
||||
{
|
||||
return ::cuda::std::lower_bound(first, last, value, comp) - first;
|
||||
}
|
||||
};
|
||||
|
||||
struct upper_bound
|
||||
{
|
||||
template <typename RangeIteratorT, typename DifferenceT, typename T, typename CompareOpT>
|
||||
_CCCL_DEVICE _CCCL_FORCEINLINE static DifferenceT
|
||||
Linear(RangeIteratorT first, DifferenceT num_items, const T& value, CompareOpT comp)
|
||||
{
|
||||
DifferenceT retval = 0;
|
||||
for (DifferenceT i = 0; i < num_items; ++i)
|
||||
{
|
||||
retval += static_cast<DifferenceT>(!comp(value, first[i]));
|
||||
}
|
||||
|
||||
return retval;
|
||||
}
|
||||
|
||||
template <typename RangeIteratorT, typename T, typename CompareOpT>
|
||||
_CCCL_DEVICE _CCCL_FORCEINLINE static ::cuda::std::ptrdiff_t
|
||||
Invoke(RangeIteratorT first, RangeIteratorT last, const T& value, CompareOpT comp)
|
||||
{
|
||||
return ::cuda::std::upper_bound(first, last, value, comp) - first;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename RangeIteratorT, typename RangeNumItemsT, typename CompareOpT, typename Mode>
|
||||
struct binary_search_transform_op_t
|
||||
{
|
||||
RangeIteratorT first;
|
||||
RangeNumItemsT num_items;
|
||||
CompareOpT op;
|
||||
|
||||
template <typename Value>
|
||||
_CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::ptrdiff_t operator()(const Value& value) const
|
||||
{
|
||||
using DifferenceT = ::cuda::std::iter_difference_t<RangeIteratorT>;
|
||||
const auto count = static_cast<DifferenceT>(num_items);
|
||||
|
||||
if (num_items <= static_cast<RangeNumItemsT>(linear_lower_bound_threshold))
|
||||
{
|
||||
return Mode::Linear(first, count, value, op);
|
||||
}
|
||||
|
||||
return Mode::Invoke(first, first + count, value, op);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename Mode, typename RangeIteratorT, typename RangeNumItemsT, typename CompareOpT>
|
||||
_CCCL_HOST_DEVICE auto make_binary_search_transform_op(RangeIteratorT first, RangeNumItemsT num_items, CompareOpT comp)
|
||||
{
|
||||
return binary_search_transform_op_t<RangeIteratorT, RangeNumItemsT, CompareOpT, Mode>{first, num_items, comp};
|
||||
}
|
||||
} // namespace detail::find
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,160 +0,0 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cuda/__device/compute_capability.h>
|
||||
#include <cuda/std/__type_traits/is_empty.h>
|
||||
#include <cuda/std/__utility/forward.h>
|
||||
#include <cuda/std/__utility/integer_sequence.h>
|
||||
#include <cuda/std/array>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail
|
||||
{
|
||||
// makes a functor that gets the policy for CC from PolicySelector when called
|
||||
template <typename PolicySelector, int CC>
|
||||
struct policy_getter : PolicySelector
|
||||
{
|
||||
_CCCL_HOST_DEVICE_API _CCCL_FORCEINLINE constexpr auto operator()() const
|
||||
{
|
||||
return PolicySelector::operator()(::cuda::compute_capability{CC});
|
||||
}
|
||||
};
|
||||
|
||||
// Device-only variant for kernel-side compile-time policy queries.
|
||||
template <typename PolicySelector, int CC>
|
||||
struct device_policy_getter : PolicySelector
|
||||
{
|
||||
_CCCL_DEVICE_API _CCCL_FORCEINLINE constexpr auto operator()() const
|
||||
{
|
||||
return PolicySelector::operator()(::cuda::compute_capability{CC});
|
||||
}
|
||||
};
|
||||
|
||||
#if !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC)
|
||||
# if _CCCL_STD_VER < 2020 && !_CCCL_COMPILER(GCC, <, 8)
|
||||
template <typename CudaCcSeq, typename PolicySelector, size_t... Is>
|
||||
struct lowest_cc_resolver;
|
||||
|
||||
// we keep the compile-time build up of the mapping table outside a template parameterized by a user-provided callable
|
||||
template <int... CudaCcs, typename PolicySelector, size_t... Is>
|
||||
struct lowest_cc_resolver<::cuda::std::integer_sequence<int, CudaCcs...>, PolicySelector, Is...>
|
||||
{
|
||||
static_assert(sizeof...(CudaCcs) == sizeof...(Is));
|
||||
|
||||
using policy_t = decltype(PolicySelector{}(::cuda::compute_capability{}));
|
||||
|
||||
static constexpr ::cuda::compute_capability all_ccs[sizeof...(Is)]{::cuda::compute_capability{CudaCcs}...};
|
||||
static constexpr policy_t all_policies[sizeof...(Is)]{PolicySelector{}(all_ccs[Is])...};
|
||||
|
||||
_CCCL_HOST_DEVICE_API static constexpr auto find_lowest(size_t i) -> ::cuda::compute_capability
|
||||
{
|
||||
const auto& policy = all_policies[i];
|
||||
while (i > 0 && policy == all_policies[i - 1])
|
||||
{
|
||||
--i;
|
||||
}
|
||||
return all_ccs[i];
|
||||
}
|
||||
|
||||
static constexpr ::cuda::compute_capability lowest_cc_with_same_policy[sizeof...(Is)]{find_lowest(Is)...};
|
||||
};
|
||||
# endif // if _CCCL_STD_VER < 2020 && !_CCCL_COMPILER(GCC, <, 8)
|
||||
|
||||
// GCC below 12 ICEs in some cases when creating an integral_constant holding a policy
|
||||
# if _CCCL_STD_VER >= 2020 && _CCCL_COMPILER(GCC, <, 12)
|
||||
template <typename Tp, Tp P>
|
||||
struct policy_constant
|
||||
{
|
||||
_CCCL_API constexpr auto operator()() const noexcept
|
||||
{
|
||||
return P;
|
||||
}
|
||||
};
|
||||
# else // _CCCL_STD_VER >= 2020 && _CCCL_COMPILER(GCC, <, 12)
|
||||
template <typename Tp, Tp P> // using <auto P> will miscompile on GCC 12
|
||||
using policy_constant = ::cuda::std::integral_constant<Tp, P>;
|
||||
# endif // _CCCL_STD_VER >= 2020 && _CCCL_COMPILER(GCC, <, 12)
|
||||
|
||||
template <typename PolicySelector, typename FunctorT, size_t... Is>
|
||||
CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch_to_cc_list(
|
||||
PolicySelector policy_selector, ::cuda::compute_capability device_cc, FunctorT&& f, ::cuda::std::index_sequence<Is...>)
|
||||
{
|
||||
constexpr auto all_ccs = ::cuda::__target_compute_capabilities();
|
||||
|
||||
_CCCL_ASSERT(((device_cc == all_ccs[Is]) || ...),
|
||||
"device_cc must appear in the list of compute capabilities compiled for");
|
||||
|
||||
cudaError_t e = cudaErrorInvalidDeviceFunction;
|
||||
# if _CCCL_STD_VER >= 2020
|
||||
// In C++20, we just create an integral_constant holding the policy, because policies are structural types in C++20.
|
||||
// This causes f to be only instantiated for each distinct policy, since the same policy for different arches results
|
||||
// in the same integral_constant type passed to f
|
||||
using policy_t = decltype(policy_selector(::cuda::compute_capability{}));
|
||||
(..., (device_cc == all_ccs[Is] ? (e = f(policy_constant<policy_t, policy_selector(all_ccs[Is])>{})) : cudaSuccess));
|
||||
# else // _CCCL_STD_VER >= 2020
|
||||
# if _CCCL_COMPILER(GCC, <, 8)
|
||||
// GCC 7 ICEs on constexpr evaluation of policy comparisons, so we skip the lowest-CC-with-same-policy optimization
|
||||
// and instantiate f for each CC directly. This may increase compile time and binary size.
|
||||
(...,
|
||||
(device_cc == all_ccs[Is] ? (e = f(policy_getter<PolicySelector, all_ccs[Is].get()>{policy_selector}))
|
||||
: cudaSuccess));
|
||||
# else // _CCCL_COMPILER(GCC, <, 8)
|
||||
// In C++17, we have to collapse architectures with the same policies ourselves, so we instantiate call_for_cc once
|
||||
// per policy on the lowest CC which produces the same policy
|
||||
using resolver_t =
|
||||
lowest_cc_resolver<::cuda::std::integer_sequence<int, all_ccs[Is].get()...>, PolicySelector, Is...>;
|
||||
(...,
|
||||
(device_cc == all_ccs[Is]
|
||||
? (e = f(policy_getter<PolicySelector, resolver_t::lowest_cc_with_same_policy[Is].get()>{policy_selector}))
|
||||
: cudaSuccess));
|
||||
# endif // _CCCL_COMPILER(GCC, <, 8)
|
||||
# endif // _CCCL_STD_VER >= 2020
|
||||
return e;
|
||||
}
|
||||
|
||||
//! Takes a policy hub and instantiates f with the minimum possible number of nullary functor types that return a policy
|
||||
//! at compile-time (if possible), and then calls the appropriate instantiation based on a runtime GPU architecture.
|
||||
//! Depending on the used compiler, C++ standard, and available macros, a different number of instantiations may be
|
||||
//! produced.
|
||||
template <typename PolicySelector, typename F>
|
||||
CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t
|
||||
dispatch_compute_cap(PolicySelector policy_selector, ::cuda::compute_capability device_cc, F&& f)
|
||||
{
|
||||
// when not using CCCL.C, policy_selector is empty since all information is contained in its type
|
||||
static_assert(::cuda::std::is_empty_v<PolicySelector>);
|
||||
return dispatch_to_cc_list(
|
||||
policy_selector,
|
||||
device_cc,
|
||||
::cuda::std::forward<F>(f),
|
||||
::cuda::std::make_index_sequence<::cuda::__target_compute_capabilities().size()>{});
|
||||
}
|
||||
|
||||
#else // !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC)
|
||||
|
||||
// if we are compiling CCCL.C with runtime policies, we cannot query the policy hub at compile time
|
||||
_CCCL_EXEC_CHECK_DISABLE
|
||||
template <typename PolicySelector, typename F>
|
||||
_CCCL_HOST_DEVICE_API _CCCL_FORCEINLINE cudaError_t
|
||||
dispatch_compute_cap(PolicySelector policy_selector, ::cuda::compute_capability device_cc, F&& f)
|
||||
{
|
||||
return f([&] {
|
||||
return policy_selector(device_cc);
|
||||
});
|
||||
}
|
||||
#endif // !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC)
|
||||
} // namespace detail
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,135 +0,0 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2011-2024, NVIDIA CORPORATION. All rights reserved.
|
||||
// SPDX-License-Identifier: BSD-3
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cuda/std/__iterator/iterator_traits.h>
|
||||
#include <cuda/std/__type_traits/common_type.h>
|
||||
#include <cuda/std/__type_traits/conditional.h>
|
||||
#include <cuda/std/__type_traits/is_integral.h>
|
||||
#include <cuda/std/__type_traits/is_same.h>
|
||||
#include <cuda/std/__type_traits/is_unsigned.h>
|
||||
#include <cuda/std/__type_traits/remove_cv.h>
|
||||
#include <cuda/std/cstdint>
|
||||
#include <cuda/std/limits>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail
|
||||
{
|
||||
/**
|
||||
* choose_offset checks NumItemsT, the type of the num_items parameter, and
|
||||
* selects the offset type based on it.
|
||||
*/
|
||||
template <typename NumItemsT>
|
||||
struct choose_offset
|
||||
{
|
||||
// NumItemsT must be an integral type (but not bool).
|
||||
static_assert(::cuda::std::is_integral_v<NumItemsT>
|
||||
&& !::cuda::std::is_same_v<::cuda::std::remove_cv_t<NumItemsT>, bool>,
|
||||
"NumItemsT must be an integral type, but not bool");
|
||||
|
||||
// Unsigned integer type for global offsets.
|
||||
using type = ::cuda::std::_If<(sizeof(NumItemsT) <= 4), uint32_t, unsigned long long>;
|
||||
};
|
||||
|
||||
/**
|
||||
* choose_offset_t is an alias template that checks NumItemsT, the type of the num_items parameter, and
|
||||
* selects the offset type based on it.
|
||||
*/
|
||||
template <typename NumItemsT>
|
||||
using choose_offset_t = typename choose_offset<NumItemsT>::type;
|
||||
|
||||
/**
|
||||
* promote_small_offset checks NumItemsT, the type of the num_items parameter, and
|
||||
* promotes any integral type smaller than 32 bits to a signed 32-bit integer type.
|
||||
*/
|
||||
template <typename NumItemsT>
|
||||
struct promote_small_offset
|
||||
{
|
||||
// NumItemsT must be an integral type (but not bool).
|
||||
static_assert(::cuda::std::is_integral_v<NumItemsT>
|
||||
&& !::cuda::std::is_same_v<::cuda::std::remove_cv_t<NumItemsT>, bool>,
|
||||
"NumItemsT must be an integral type, but not bool");
|
||||
|
||||
// Unsigned integer type for global offsets.
|
||||
using type = ::cuda::std::_If<(sizeof(NumItemsT) < 4), int32_t, NumItemsT>;
|
||||
};
|
||||
|
||||
/**
|
||||
* promote_small_offset_t is an alias template that checks NumItemsT, the type of the num_items parameter, and
|
||||
* promotes any integral type smaller than 32 bits to a signed 32-bit integer type.
|
||||
*/
|
||||
template <typename NumItemsT>
|
||||
using promote_small_offset_t = typename promote_small_offset<NumItemsT>::type;
|
||||
|
||||
/**
|
||||
* choose_signed_offset checks NumItemsT, the type of the num_items parameter, and
|
||||
* selects the offset type to be either int32 or int64, such that the selected offset type covers the range of NumItemsT
|
||||
* unless it was uint64, in which case int64 will be used.
|
||||
*/
|
||||
template <typename NumItemsT>
|
||||
struct choose_signed_offset
|
||||
{
|
||||
// NumItemsT must be an integral type (but not bool).
|
||||
static_assert(::cuda::std::is_integral_v<NumItemsT>
|
||||
&& !::cuda::std::is_same_v<::cuda::std::remove_cv_t<NumItemsT>, bool>,
|
||||
"NumItemsT must be an integral type, but not bool");
|
||||
|
||||
// Signed integer type for global offsets.
|
||||
// uint32 -> int64, else
|
||||
// LEQ 4B -> int32, else
|
||||
// int64
|
||||
using type = ::cuda::std::_If<(::cuda::std::is_integral_v<NumItemsT> && ::cuda::std::is_unsigned_v<NumItemsT>),
|
||||
::cuda::std::int64_t,
|
||||
::cuda::std::_If<(sizeof(NumItemsT) <= 4), ::cuda::std::int32_t, ::cuda::std::int64_t>>;
|
||||
|
||||
/**
|
||||
* Checks if the given num_items can be covered by the selected offset type. If not, returns cudaErrorInvalidValue,
|
||||
* otherwise returns cudaSuccess.
|
||||
*/
|
||||
static _CCCL_HOST_DEVICE _CCCL_FORCEINLINE cudaError_t is_exceeding_offset_type(NumItemsT num_items)
|
||||
{
|
||||
_CCCL_DIAG_PUSH
|
||||
_CCCL_DIAG_SUPPRESS_MSVC(4127) /* conditional expression is constant */
|
||||
if (sizeof(NumItemsT) >= 8 && num_items > static_cast<NumItemsT>(::cuda::std::numeric_limits<type>::max()))
|
||||
{
|
||||
return cudaErrorInvalidValue;
|
||||
}
|
||||
_CCCL_DIAG_POP
|
||||
return cudaSuccess;
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* choose_signed_offset_t is an alias template that checks NumItemsT, the type of the num_items parameter, and
|
||||
* selects the corresponding signed offset type based on it.
|
||||
*/
|
||||
template <typename NumItemsT>
|
||||
using choose_signed_offset_t = typename choose_signed_offset<NumItemsT>::type;
|
||||
|
||||
/**
|
||||
* common_iterator_value sets member type to the common_type of
|
||||
* value_type for all argument types. used to get OffsetT in
|
||||
* DeviceSegmentedReduce.
|
||||
*/
|
||||
template <typename... Iter>
|
||||
struct common_iterator_value
|
||||
{
|
||||
using type = ::cuda::std::common_type_t<::cuda::std::__iter_value_type<Iter>...>;
|
||||
};
|
||||
template <typename... Iter>
|
||||
using common_iterator_value_t = typename common_iterator_value<Iter...>::type;
|
||||
} // namespace detail
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,82 +0,0 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cub/detail/choose_offset.cuh>
|
||||
|
||||
#include <cuda/__argument/argument.h>
|
||||
#include <cuda/std/__type_traits/is_same.h>
|
||||
#include <cuda/std/__utility/declval.h>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail
|
||||
{
|
||||
#if !_CCCL_COMPILER(NVRTC)
|
||||
// Preserve deferred problem sizes for dispatch and canonicalize immediate values to CUB's offset type.
|
||||
template <typename NumItemsT>
|
||||
[[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto make_num_items_dispatch_arg(NumItemsT num_items) noexcept
|
||||
{
|
||||
using args_traits_t = ::cuda::args::__traits<NumItemsT>;
|
||||
|
||||
if constexpr (args_traits_t::is_deferred)
|
||||
{
|
||||
return num_items;
|
||||
}
|
||||
else
|
||||
{
|
||||
using offset_t = choose_offset_t<typename args_traits_t::element_type>;
|
||||
return static_cast<offset_t>(::cuda::args::__unwrap(num_items));
|
||||
}
|
||||
}
|
||||
|
||||
// Forms a kernel parameter from a single-value argument without reading a deferred source.
|
||||
// Immediate values are converted to TargetT. Deferred arguments are unwrapped to their source, erasing bounds from
|
||||
// the kernel type and payload.
|
||||
template <typename TargetT, typename ParameterT>
|
||||
[[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE constexpr auto parameter_from_host(ParameterT parameter) noexcept
|
||||
{
|
||||
using args_traits_t = ::cuda::args::__traits<ParameterT>;
|
||||
static_assert(args_traits_t::is_single_value, "parameter must contain a single value");
|
||||
|
||||
if constexpr (args_traits_t::is_deferred)
|
||||
{
|
||||
return ::cuda::args::__unwrap(parameter);
|
||||
}
|
||||
else
|
||||
{
|
||||
return static_cast<TargetT>(::cuda::args::__unwrap(parameter));
|
||||
}
|
||||
}
|
||||
|
||||
template <typename TargetT, typename ParameterT>
|
||||
using parameter_from_host_t = decltype(parameter_from_host<TargetT>(::cuda::std::declval<ParameterT>()));
|
||||
#endif // !_CCCL_COMPILER(NVRTC)
|
||||
|
||||
// Forms a value from a kernel parameter, reading element zero when the parameter is a deferred source.
|
||||
template <typename TargetT, typename ParameterT>
|
||||
[[nodiscard]] _CCCL_DEVICE_API _CCCL_FORCEINLINE TargetT parameter_from_device(ParameterT parameter) noexcept
|
||||
{
|
||||
if constexpr (::cuda::std::is_same_v<ParameterT, TargetT>)
|
||||
{
|
||||
return parameter;
|
||||
}
|
||||
else
|
||||
{
|
||||
return static_cast<TargetT>(parameter[0]);
|
||||
}
|
||||
}
|
||||
} // namespace detail
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,150 +0,0 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved.
|
||||
// SPDX-License-Identifier: BSD-3
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cub/agent/single_pass_scan_operators.cuh>
|
||||
#include <cub/device/dispatch/tuning/common.cuh>
|
||||
|
||||
#include <cuda/std/__concepts/same_as.h>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail
|
||||
{
|
||||
template <typename DelayConstructor>
|
||||
inline constexpr auto lookback_delay_policy_from_type = 0;
|
||||
|
||||
template <unsigned int L2WriteLatency>
|
||||
inline constexpr auto lookback_delay_policy_from_type<no_delay_constructor_t<L2WriteLatency>> =
|
||||
LookbackDelayPolicy{LookbackDelayAlgorithm::no_delay, 0, L2WriteLatency};
|
||||
|
||||
template <unsigned int Delay, unsigned int L2WriteLatency>
|
||||
inline constexpr auto lookback_delay_policy_from_type<fixed_delay_constructor_t<Delay, L2WriteLatency>> =
|
||||
LookbackDelayPolicy{LookbackDelayAlgorithm::fixed_delay, Delay, L2WriteLatency};
|
||||
|
||||
template <unsigned int Delay, unsigned int L2WriteLatency>
|
||||
inline constexpr auto lookback_delay_policy_from_type<exponential_backoff_constructor_t<Delay, L2WriteLatency>> =
|
||||
LookbackDelayPolicy{LookbackDelayAlgorithm::exponential_backoff, Delay, L2WriteLatency};
|
||||
|
||||
template <unsigned int Delay, unsigned int L2WriteLatency>
|
||||
inline constexpr auto lookback_delay_policy_from_type<exponential_backoff_jitter_constructor_t<Delay, L2WriteLatency>> =
|
||||
LookbackDelayPolicy{LookbackDelayAlgorithm::exponential_backoff_jitter, Delay, L2WriteLatency};
|
||||
|
||||
template <unsigned int Delay, unsigned int L2WriteLatency>
|
||||
inline constexpr auto
|
||||
lookback_delay_policy_from_type<exponential_backoff_jitter_window_constructor_t<Delay, L2WriteLatency>> =
|
||||
LookbackDelayPolicy{LookbackDelayAlgorithm::exponential_backoff_jitter_window, Delay, L2WriteLatency};
|
||||
|
||||
template <unsigned int Delay, unsigned int L2WriteLatency>
|
||||
inline constexpr auto
|
||||
lookback_delay_policy_from_type<exponential_backon_jitter_window_constructor_t<Delay, L2WriteLatency>> =
|
||||
LookbackDelayPolicy{LookbackDelayAlgorithm::exponential_backon_jitter_window, Delay, L2WriteLatency};
|
||||
|
||||
template <unsigned int Delay, unsigned int L2WriteLatency>
|
||||
inline constexpr auto lookback_delay_policy_from_type<exponential_backon_jitter_constructor_t<Delay, L2WriteLatency>> =
|
||||
LookbackDelayPolicy{LookbackDelayAlgorithm::exponential_backon_jitter, Delay, L2WriteLatency};
|
||||
|
||||
template <unsigned int Delay, unsigned int L2WriteLatency>
|
||||
inline constexpr auto lookback_delay_policy_from_type<exponential_backon_constructor_t<Delay, L2WriteLatency>> =
|
||||
LookbackDelayPolicy{LookbackDelayAlgorithm::exponential_backon, Delay, L2WriteLatency};
|
||||
|
||||
template <unsigned int Delay, unsigned int L2WriteLatency, unsigned int GridThreshold>
|
||||
inline constexpr auto
|
||||
lookback_delay_policy_from_type<reduce_by_key_delay_constructor_t<Delay, L2WriteLatency, GridThreshold>> =
|
||||
LookbackDelayPolicy{LookbackDelayAlgorithm::__reduce_by_key, Delay, L2WriteLatency};
|
||||
|
||||
template <LookbackDelayAlgorithm Kind, unsigned int Delay, unsigned int L2WriteLatency>
|
||||
struct delay_constructor_for;
|
||||
|
||||
template <unsigned int Delay, unsigned int L2WriteLatency>
|
||||
struct delay_constructor_for<LookbackDelayAlgorithm::no_delay, Delay, L2WriteLatency>
|
||||
{
|
||||
using type = no_delay_constructor_t<L2WriteLatency>;
|
||||
};
|
||||
|
||||
template <unsigned int Delay, unsigned int L2WriteLatency>
|
||||
struct delay_constructor_for<LookbackDelayAlgorithm::fixed_delay, Delay, L2WriteLatency>
|
||||
{
|
||||
using type = fixed_delay_constructor_t<Delay, L2WriteLatency>;
|
||||
};
|
||||
|
||||
template <unsigned int Delay, unsigned int L2WriteLatency>
|
||||
struct delay_constructor_for<LookbackDelayAlgorithm::exponential_backoff, Delay, L2WriteLatency>
|
||||
{
|
||||
using type = exponential_backoff_constructor_t<Delay, L2WriteLatency>;
|
||||
};
|
||||
|
||||
template <unsigned int Delay, unsigned int L2WriteLatency>
|
||||
struct delay_constructor_for<LookbackDelayAlgorithm::exponential_backoff_jitter, Delay, L2WriteLatency>
|
||||
{
|
||||
using type = exponential_backoff_jitter_constructor_t<Delay, L2WriteLatency>;
|
||||
};
|
||||
|
||||
template <unsigned int Delay, unsigned int L2WriteLatency>
|
||||
struct delay_constructor_for<LookbackDelayAlgorithm::exponential_backoff_jitter_window, Delay, L2WriteLatency>
|
||||
{
|
||||
using type = exponential_backoff_jitter_window_constructor_t<Delay, L2WriteLatency>;
|
||||
};
|
||||
|
||||
template <unsigned int Delay, unsigned int L2WriteLatency>
|
||||
struct delay_constructor_for<LookbackDelayAlgorithm::exponential_backon_jitter_window, Delay, L2WriteLatency>
|
||||
{
|
||||
using type = exponential_backon_jitter_window_constructor_t<Delay, L2WriteLatency>;
|
||||
};
|
||||
|
||||
template <unsigned int Delay, unsigned int L2WriteLatency>
|
||||
struct delay_constructor_for<LookbackDelayAlgorithm::exponential_backon_jitter, Delay, L2WriteLatency>
|
||||
{
|
||||
using type = exponential_backon_jitter_constructor_t<Delay, L2WriteLatency>;
|
||||
};
|
||||
|
||||
template <unsigned int Delay, unsigned int L2WriteLatency>
|
||||
struct delay_constructor_for<LookbackDelayAlgorithm::exponential_backon, Delay, L2WriteLatency>
|
||||
{
|
||||
using type = exponential_backon_constructor_t<Delay, L2WriteLatency>;
|
||||
};
|
||||
|
||||
template <unsigned int Delay, unsigned int L2WriteLatency>
|
||||
struct delay_constructor_for<LookbackDelayAlgorithm::__reduce_by_key, Delay, L2WriteLatency>
|
||||
{
|
||||
using type = reduce_by_key_delay_constructor_t<Delay, L2WriteLatency>;
|
||||
};
|
||||
|
||||
template <LookbackDelayAlgorithm Kind, unsigned int Delay, unsigned int L2WriteLatency>
|
||||
using delay_constructor_t = typename delay_constructor_for<Kind, Delay, L2WriteLatency>::type;
|
||||
|
||||
_CCCL_HOST_DEVICE_API constexpr auto default_delay_constructor_policy(bool is_primitive_or_trivially_copyable)
|
||||
{
|
||||
if (is_primitive_or_trivially_copyable)
|
||||
{
|
||||
return LookbackDelayPolicy{LookbackDelayAlgorithm::fixed_delay, 350, 450};
|
||||
}
|
||||
return LookbackDelayPolicy{LookbackDelayAlgorithm::no_delay, 0, 450};
|
||||
}
|
||||
|
||||
_CCCL_HOST_DEVICE_API constexpr auto default_reduce_by_key_delay_constructor_policy(
|
||||
int key_size,
|
||||
int value_size,
|
||||
bool key_is_primitive_or_trivially_copyable,
|
||||
bool value_is_primitive_or_trivially_copyable)
|
||||
{
|
||||
if (value_is_primitive_or_trivially_copyable && (value_size + key_size < 16))
|
||||
{
|
||||
return LookbackDelayPolicy{LookbackDelayAlgorithm::__reduce_by_key, 350, 450};
|
||||
}
|
||||
return default_delay_constructor_policy(key_is_primitive_or_trivially_copyable);
|
||||
}
|
||||
} // namespace detail
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,94 +0,0 @@
|
||||
/*
|
||||
* Copyright 2021 NVIDIA Corporation
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cub/util_namespace.cuh>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail
|
||||
{
|
||||
/**
|
||||
* @brief It's a double-buffer storage wrapper for multi-pass stream
|
||||
* transformations that require more than one storage array for
|
||||
* streaming intermediate results back and forth.
|
||||
*
|
||||
* Many multi-pass computations require a pair of "ping-pong" storage buffers
|
||||
* (e.g., one for reading from and the other for writing to, and then
|
||||
* vice-versa for the subsequent pass). This structure wraps a set of device
|
||||
* buffers.
|
||||
*
|
||||
* Unlike `cub::DoubleBuffer` this class doesn't provide a "selector" member
|
||||
* to track which buffer is "current". The main reason for this class existence
|
||||
* is the performance difference. Since `cub::DoubleBuffer` relies on the
|
||||
* runtime variable to index pointers arrays, they are placed in the local
|
||||
* memory instead of registers. Local memory accesses significantly affect
|
||||
* performance. On the contrary, this class swaps pointer, so all operations
|
||||
* can be performed in registers.
|
||||
*/
|
||||
template <typename T>
|
||||
class device_double_buffer
|
||||
{
|
||||
/// Pair of device buffer pointers
|
||||
T* m_current_buffer{};
|
||||
T* m_alternate_buffer{};
|
||||
|
||||
public:
|
||||
/**
|
||||
* @param d_current
|
||||
* The currently valid buffer
|
||||
*
|
||||
* @param d_alternate
|
||||
* Alternate storage buffer of the same size as @p d_current
|
||||
*/
|
||||
_CCCL_HOST_DEVICE _CCCL_FORCEINLINE device_double_buffer(T* current, T* alternate)
|
||||
: m_current_buffer(current)
|
||||
, m_alternate_buffer(alternate)
|
||||
{}
|
||||
|
||||
/// \brief Return pointer to the currently valid buffer
|
||||
_CCCL_HOST_DEVICE _CCCL_FORCEINLINE T* current() const
|
||||
{
|
||||
return m_current_buffer;
|
||||
}
|
||||
|
||||
/// \brief Return pointer to the currently invalid buffer
|
||||
_CCCL_HOST_DEVICE _CCCL_FORCEINLINE T* alternate() const
|
||||
{
|
||||
return m_alternate_buffer;
|
||||
}
|
||||
|
||||
_CCCL_HOST_DEVICE void swap()
|
||||
{
|
||||
T* tmp = m_current_buffer;
|
||||
m_current_buffer = m_alternate_buffer;
|
||||
m_alternate_buffer = tmp;
|
||||
}
|
||||
};
|
||||
} // namespace detail
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,97 +0,0 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
|
||||
// SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cub/util_device.cuh>
|
||||
|
||||
#include <cuda/__runtime/api_wrapper.h>
|
||||
#include <cuda/__stream/stream_ref.h>
|
||||
#include <cuda/std/__exception/terminate.h>
|
||||
#include <cuda/std/cstdint>
|
||||
|
||||
#include <nv/target>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail
|
||||
{
|
||||
// TODO(gevtushenko/srinivasyadav18): move cudax `device_memory_resource` to `cuda::__device_memory_resource` and remove
|
||||
// this implementation
|
||||
struct device_memory_resource
|
||||
{
|
||||
CUB_RUNTIME_FUNCTION void* allocate(size_t bytes, size_t /* alignment */)
|
||||
{
|
||||
void* ptr{nullptr};
|
||||
NV_IF_ELSE_TARGET(
|
||||
NV_IS_HOST,
|
||||
(_CCCL_TRY_CUDA_API(::cudaMallocAsync, "allocate failed to allocate with cudaMallocAsync", &ptr, bytes, NULL);),
|
||||
({
|
||||
_CubLog("%s\n", "cub::detail::device_memory_resource::allocate not supported from device code.");
|
||||
::cuda::std::terminate();
|
||||
}));
|
||||
_CCCL_ASSERT(ptr != nullptr, "allocate failed to allocate with cudaMallocAsync");
|
||||
return ptr;
|
||||
}
|
||||
|
||||
CUB_RUNTIME_FUNCTION void deallocate(void* ptr, size_t /* bytes */)
|
||||
{
|
||||
NV_IF_ELSE_TARGET( //
|
||||
NV_IS_HOST,
|
||||
(_CCCL_TRY_CUDA_API(::cudaFree, "deallocate failed", ptr);),
|
||||
({
|
||||
_CubLog("%s\n", "cub::detail::device_memory_resource::deallocate not supported from device code.");
|
||||
::cuda::std::terminate();
|
||||
}));
|
||||
}
|
||||
|
||||
CUB_RUNTIME_FUNCTION void* allocate(::cuda::stream_ref stream, size_t bytes, size_t /* alignment */)
|
||||
{
|
||||
return allocate(stream, bytes);
|
||||
}
|
||||
|
||||
CUB_RUNTIME_FUNCTION void* allocate(::cuda::stream_ref stream, size_t bytes)
|
||||
{
|
||||
void* ptr{nullptr};
|
||||
NV_IF_ELSE_TARGET( //
|
||||
NV_IS_HOST,
|
||||
({
|
||||
_CCCL_TRY_CUDA_API(
|
||||
::cudaMallocAsync, "allocate failed to allocate with cudaMallocAsync", &ptr, bytes, stream.get());
|
||||
}),
|
||||
({
|
||||
_CubLog("%s\n", "cub::detail::device_memory_resource::allocate not supported from device code.");
|
||||
::cuda::std::terminate();
|
||||
}));
|
||||
return ptr;
|
||||
}
|
||||
|
||||
CUB_RUNTIME_FUNCTION void deallocate(::cuda::stream_ref stream, void* ptr, size_t bytes, size_t /* alignment */)
|
||||
{
|
||||
deallocate(stream, ptr, bytes);
|
||||
}
|
||||
|
||||
CUB_RUNTIME_FUNCTION void deallocate(::cuda::stream_ref stream, void* ptr, size_t /* bytes */)
|
||||
{
|
||||
NV_IF_ELSE_TARGET( //
|
||||
NV_IS_HOST,
|
||||
(_CCCL_TRY_CUDA_API(::cudaFreeAsync, "deallocate failed", ptr, stream.get());),
|
||||
({
|
||||
_CubLog("%s\n", "cub::detail::device_memory_resource::deallocate not supported from device code.");
|
||||
::cuda::std::terminate();
|
||||
}));
|
||||
}
|
||||
};
|
||||
} // namespace detail
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,137 +0,0 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
|
||||
// SPDX-License-Identifier: BSD-3-Clause
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cub/detail/device_memory_resource.cuh>
|
||||
#include <cub/detail/temporary_storage.cuh>
|
||||
|
||||
#include <cuda/__execution/tune.h>
|
||||
#include <cuda/__functional/call_or.h>
|
||||
#include <cuda/__memory_resource/get_memory_resource.h>
|
||||
#include <cuda/__stream/get_stream.h>
|
||||
#include <cuda/std/__execution/env.h>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail
|
||||
{
|
||||
//! @cond
|
||||
//! Generic environment-based algorithm dispatch wrapper
|
||||
//!
|
||||
//! Handles common boilerplate for all env-based algorithms:
|
||||
//! - Query stream, memory resource, and tuning from environment
|
||||
//! - Two-phase call (query temp storage size, then execute)
|
||||
//! - Temporary storage allocation/deallocation
|
||||
//! - Memory resource querying from environment
|
||||
//!
|
||||
//! @param env The execution environment
|
||||
//! @param algorithm_callable Callable that invokes the algorithm implementation with determinism specified
|
||||
template <typename EnvT, typename AlgorithmCallable>
|
||||
CUB_RUNTIME_FUNCTION static cudaError_t dispatch_with_env(const EnvT& env, AlgorithmCallable&& algorithm_callable)
|
||||
{
|
||||
// Query stream from environment
|
||||
auto stream = ::cuda::__call_or(::cuda::get_stream, ::cuda::stream_ref{cudaStream_t{}}, env);
|
||||
|
||||
// Query memory resource from environment
|
||||
auto mr = ::cuda::__call_or(::cuda::mr::__get_memory_resource, detail::device_memory_resource{}, env);
|
||||
|
||||
// Query tuning from environment
|
||||
const auto tuning = ::cuda::__call_or(::cuda::execution::__get_tuning, ::cuda::std::execution::env<>{}, env);
|
||||
|
||||
void* d_temp_storage = nullptr;
|
||||
size_t temp_storage_bytes = 0;
|
||||
|
||||
// Phase 1: Query temporary storage size
|
||||
if (const auto error = algorithm_callable(tuning, d_temp_storage, temp_storage_bytes, stream.get()))
|
||||
{
|
||||
return error;
|
||||
}
|
||||
|
||||
// Allocate temporary storage
|
||||
if (const auto error = CubDebug(detail::temporary_storage::allocate(stream, d_temp_storage, temp_storage_bytes, mr)))
|
||||
{
|
||||
return error;
|
||||
}
|
||||
|
||||
// Phase 2: Execute algorithm
|
||||
const auto error = algorithm_callable(tuning, d_temp_storage, temp_storage_bytes, stream.get());
|
||||
|
||||
// Deallocate temporary storage (always attempt, even on error)
|
||||
const auto deallocate_error =
|
||||
CubDebug(detail::temporary_storage::deallocate(stream, d_temp_storage, temp_storage_bytes, mr));
|
||||
|
||||
// Algorithm error takes precedence over deallocation error
|
||||
return (error != cudaSuccess) ? error : deallocate_error;
|
||||
}
|
||||
//! @endcond
|
||||
|
||||
template <typename DefaultPolicySelector, typename EnvT, typename AlgorithmCallable>
|
||||
CUB_RUNTIME_FUNCTION static cudaError_t
|
||||
dispatch_with_env_and_tuning(const EnvT& env, AlgorithmCallable&& algorithm_callable)
|
||||
{
|
||||
return detail::dispatch_with_env(
|
||||
env,
|
||||
[&algorithm_callable](
|
||||
[[maybe_unused]] auto tuning_env, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) {
|
||||
using policy_t = decltype(DefaultPolicySelector{}(::cuda::compute_capability{}));
|
||||
using policy_selector =
|
||||
::cuda::std::execution::__query_result_or_t<decltype(tuning_env), policy_t, DefaultPolicySelector>;
|
||||
return algorithm_callable(policy_selector{}, d_temp_storage, temp_storage_bytes, stream);
|
||||
});
|
||||
}
|
||||
|
||||
//! @cond
|
||||
//! Generic environment-based algorithm dispatch wrapper
|
||||
//!
|
||||
//! Handles common boilerplate for env-based algorithms with user provided memory:
|
||||
//! - Query stream, and tuning from environment
|
||||
//! - Single-phase call passing user provided memory and size
|
||||
//!
|
||||
//! @param env The execution environment
|
||||
//! @param[in] d_temp_storage @devicestorage
|
||||
//! @param[in,out] temp_storage_bytes Reference to size in bytes of `d_temp_storage` allocation
|
||||
//! @param algorithm_callable Callable that invokes the algorithm implementation with determinism specified
|
||||
template <typename EnvT, typename AlgorithmCallable>
|
||||
CUB_RUNTIME_FUNCTION static cudaError_t dispatch_with_env(
|
||||
void* d_temp_storage, size_t& temp_storage_bytes, const EnvT& env, AlgorithmCallable&& algorithm_callable)
|
||||
{
|
||||
// Query stream from environment
|
||||
auto stream = ::cuda::__call_or(::cuda::get_stream, ::cuda::stream_ref{cudaStream_t{}}, env);
|
||||
|
||||
// Query tuning from environment
|
||||
const auto tuning = ::cuda::__call_or(::cuda::execution::__get_tuning, ::cuda::std::execution::env<>{}, env);
|
||||
|
||||
return algorithm_callable(tuning, d_temp_storage, temp_storage_bytes, stream.get());
|
||||
}
|
||||
//! @endcond
|
||||
|
||||
template <typename DefaultPolicySelector, typename EnvT, typename AlgorithmCallable>
|
||||
CUB_RUNTIME_FUNCTION static cudaError_t dispatch_with_env_and_tuning(
|
||||
void* d_temp_storage, size_t& temp_storage_bytes, const EnvT& env, AlgorithmCallable&& algorithm_callable)
|
||||
{
|
||||
return detail::dispatch_with_env(
|
||||
d_temp_storage,
|
||||
temp_storage_bytes,
|
||||
env,
|
||||
[&algorithm_callable](
|
||||
[[maybe_unused]] auto tuning_env, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) {
|
||||
using policy_t = decltype(DefaultPolicySelector{}(::cuda::compute_capability{}));
|
||||
using policy_selector =
|
||||
::cuda::std::execution::__query_result_or_t<decltype(tuning_env), policy_t, DefaultPolicySelector>;
|
||||
return algorithm_callable(policy_selector{}, d_temp_storage, temp_storage_bytes, stream);
|
||||
});
|
||||
}
|
||||
//! @endcond
|
||||
} // namespace detail
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,226 +0,0 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2011-2024, NVIDIA CORPORATION. All rights reserved.
|
||||
// SPDX-License-Identifier: BSD-3
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cub/detail/type_traits.cuh> // implicit_prom_t
|
||||
#include <cub/util_type.cuh> // _CCCL_HAS_INT128()
|
||||
|
||||
#include <cuda/__cmath/ceil_div.h>
|
||||
#include <cuda/__cmath/pow2.h>
|
||||
#include <cuda/std/__bit/integral.h>
|
||||
#include <cuda/std/__type_traits/conditional.h>
|
||||
#include <cuda/std/__type_traits/enable_if.h>
|
||||
#include <cuda/std/__type_traits/integral_constant.h>
|
||||
#include <cuda/std/__type_traits/is_integral.h>
|
||||
#include <cuda/std/__type_traits/is_same.h>
|
||||
#include <cuda/std/__type_traits/is_signed.h>
|
||||
#include <cuda/std/__type_traits/make_unsigned.h>
|
||||
#include <cuda/std/climits> // CHAR_BIT
|
||||
#include <cuda/std/cstdint> // uint64_t
|
||||
#include <cuda/std/limits>
|
||||
|
||||
#if defined(CCCL_ENABLE_DEVICE_ASSERTIONS)
|
||||
_CCCL_BEGIN_NV_DIAG_SUPPRESS(186) // pointless comparison of unsigned integer with zero
|
||||
#endif // CCCL_ENABLE_DEVICE_ASSERTIONS
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail
|
||||
{
|
||||
/***********************************************************************************************************************
|
||||
* larger_unsigned_type
|
||||
**********************************************************************************************************************/
|
||||
|
||||
template <typename T, typename = void>
|
||||
struct larger_unsigned_type
|
||||
{
|
||||
using type = void;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct larger_unsigned_type<T, ::cuda::std::enable_if_t<(sizeof(T) < 4)>>
|
||||
{
|
||||
using type = ::cuda::std::uint32_t;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct larger_unsigned_type<T, ::cuda::std::enable_if_t<(sizeof(T) == 4)>>
|
||||
{
|
||||
using type = ::cuda::std::uint64_t;
|
||||
};
|
||||
|
||||
#if _CCCL_HAS_INT128()
|
||||
|
||||
template <typename T>
|
||||
struct larger_unsigned_type<T, ::cuda::std::enable_if_t<(sizeof(T) == 8)>>
|
||||
{
|
||||
using type = __uint128_t;
|
||||
};
|
||||
|
||||
#endif // _CCCL_HAS_INT128()
|
||||
|
||||
template <typename T>
|
||||
using larger_unsigned_type_t = typename larger_unsigned_type<T>::type;
|
||||
|
||||
template <typename T>
|
||||
using unsigned_implicit_prom_t = ::cuda::std::make_unsigned_t<implicit_prom_t<T>>;
|
||||
|
||||
template <typename T>
|
||||
using supported_integral =
|
||||
::cuda::std::bool_constant<::cuda::std::is_integral_v<T> && !::cuda::std::is_same_v<T, bool> && (sizeof(T) <= 8)>;
|
||||
|
||||
/***********************************************************************************************************************
|
||||
* Extract higher bits after multiplication
|
||||
**********************************************************************************************************************/
|
||||
|
||||
template <typename DivisorType, typename T, typename R>
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE unsigned_implicit_prom_t<DivisorType>
|
||||
multiply_extract_higher_bits(T value, R multiplier)
|
||||
{
|
||||
static_assert(supported_integral<T>::value, "unsupported type");
|
||||
static_assert(supported_integral<R>::value, "unsupported type");
|
||||
if constexpr (::cuda::std::is_signed_v<T>)
|
||||
{
|
||||
_CCCL_ASSERT(value >= 0, "value must be non-negative");
|
||||
}
|
||||
if constexpr (::cuda::std::is_signed_v<R>)
|
||||
{
|
||||
_CCCL_ASSERT(multiplier >= 0, "multiplier must be non-negative");
|
||||
}
|
||||
static constexpr int NumBits = sizeof(DivisorType) * CHAR_BIT;
|
||||
using unsigned_t = unsigned_implicit_prom_t<DivisorType>;
|
||||
using larger_t = larger_unsigned_type_t<DivisorType>;
|
||||
// clang-format off
|
||||
NV_IF_ELSE_TARGET(
|
||||
NV_IS_HOST,
|
||||
(return static_cast<unsigned_t>((static_cast<larger_t>(value) * multiplier) >> NumBits);),
|
||||
({return (sizeof(T) == 8)
|
||||
? static_cast<unsigned_t>(__umul64hi(value, multiplier))
|
||||
: static_cast<unsigned_t>((static_cast<larger_t>(value) * multiplier) >> NumBits);}));
|
||||
// clang-format on
|
||||
}
|
||||
|
||||
/***********************************************************************************************************************
|
||||
* Fast Modulo/Division based on Precomputation
|
||||
**********************************************************************************************************************/
|
||||
|
||||
_CCCL_DIAG_PUSH
|
||||
_CCCL_DIAG_SUPPRESS_MSVC(4127) /* conditional expression is constant */
|
||||
|
||||
template <typename T1>
|
||||
class fast_div_mod
|
||||
{
|
||||
static_assert(supported_integral<T1>::value, "unsupported type");
|
||||
|
||||
// uint16_t is a special case that would requires complex logic. Workaround: convert to int
|
||||
using T = ::cuda::std::conditional_t<::cuda::std::is_same_v<T1, ::cuda::std::uint16_t>, int, T1>;
|
||||
using unsigned_t = unsigned_implicit_prom_t<T>;
|
||||
|
||||
public:
|
||||
template <typename R>
|
||||
struct result
|
||||
{
|
||||
using common_t = decltype(R{} / T{});
|
||||
common_t quotient;
|
||||
common_t remainder;
|
||||
};
|
||||
|
||||
fast_div_mod() = delete;
|
||||
|
||||
_CCCL_HOST_DEVICE explicit fast_div_mod(T divisor) noexcept
|
||||
: _divisor{static_cast<unsigned_t>(divisor)}
|
||||
{
|
||||
using larger_t = larger_unsigned_type_t<T>;
|
||||
_CCCL_ASSERT(divisor > 0, "divisor must be positive");
|
||||
auto udivisor = static_cast<unsigned_t>(divisor);
|
||||
// the following branches are needed to avoid negative shift
|
||||
if (::cuda::is_power_of_two(udivisor))
|
||||
{
|
||||
_shift_right = ::cuda::std::bit_width(udivisor) - 1;
|
||||
return;
|
||||
}
|
||||
else if (sizeof(T) == 8 && divisor == 3)
|
||||
{
|
||||
return;
|
||||
}
|
||||
constexpr int BitSize = sizeof(T) * CHAR_BIT; // 32
|
||||
constexpr int BitOffset = BitSize / 16; // 2
|
||||
int num_bits = ::cuda::std::bit_width(udivisor) + 1;
|
||||
_CCCL_ASSERT(static_cast<size_t>(num_bits + BitSize - BitOffset) < sizeof(larger_t) * CHAR_BIT, "overflow error");
|
||||
// without explicit power-of-two check, num_bits needs to replace +1 with !::cuda::is_power_of_two(udivisor)
|
||||
_multiplier = static_cast<unsigned_t>(::cuda::ceil_div(larger_t{1} << (num_bits + BitSize - BitOffset), //
|
||||
static_cast<larger_t>(divisor)));
|
||||
_shift_right = num_bits - BitOffset;
|
||||
_CCCL_ASSERT(_multiplier != 0, "overflow error");
|
||||
}
|
||||
|
||||
fast_div_mod(const fast_div_mod&) noexcept = default;
|
||||
|
||||
fast_div_mod(fast_div_mod&&) noexcept = default;
|
||||
|
||||
template <typename R>
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE result<R> operator()(R dividend) const noexcept
|
||||
{
|
||||
static_assert(supported_integral<R>::value, "unsupported type");
|
||||
using common_t = decltype(R{} / T{});
|
||||
using ucommon_t = ::cuda::std::make_unsigned_t<common_t>;
|
||||
using result_t = result<R>;
|
||||
_CCCL_ASSERT(dividend >= 0, "divisor must be non-negative");
|
||||
auto udividend = static_cast<ucommon_t>(dividend);
|
||||
if (_divisor == 1)
|
||||
{
|
||||
return result_t{static_cast<common_t>(dividend), common_t{}};
|
||||
}
|
||||
else if (_divisor > unsigned_t{::cuda::std::numeric_limits<T>::max() / 2})
|
||||
{
|
||||
auto quotient = udividend >= static_cast<ucommon_t>(_divisor);
|
||||
return result_t{static_cast<common_t>(quotient), static_cast<common_t>(udividend - (quotient * _divisor))};
|
||||
}
|
||||
else if (sizeof(T) == 8 && _divisor == 3)
|
||||
{
|
||||
return result_t{static_cast<common_t>(udividend / 3), static_cast<common_t>(udividend % 3)};
|
||||
}
|
||||
auto higher_bits = (_multiplier == 0) ? udividend : multiply_extract_higher_bits<T>(dividend, _multiplier);
|
||||
auto quotient = higher_bits >> _shift_right;
|
||||
auto remainder = udividend - (quotient * _divisor);
|
||||
_CCCL_ASSERT(quotient == udividend / _divisor, "wrong quotient");
|
||||
_CCCL_ASSERT(remainder < (ucommon_t) _divisor, "remainder out of range");
|
||||
return result_t{static_cast<common_t>(quotient), static_cast<common_t>(remainder)};
|
||||
}
|
||||
|
||||
template <typename R>
|
||||
_CCCL_HOST_DEVICE _CCCL_FORCEINLINE friend implicit_prom_t<T> operator/(R dividend, fast_div_mod div) noexcept
|
||||
{
|
||||
return div(dividend).quotient;
|
||||
}
|
||||
|
||||
template <typename R>
|
||||
_CCCL_HOST_DEVICE _CCCL_FORCEINLINE friend implicit_prom_t<T> operator%(R dividend, fast_div_mod div) noexcept
|
||||
{
|
||||
return div(dividend).remainder;
|
||||
}
|
||||
|
||||
private:
|
||||
unsigned_t _divisor = 1;
|
||||
unsigned_t _multiplier = 0;
|
||||
unsigned _shift_right = 0;
|
||||
};
|
||||
_CCCL_DIAG_POP
|
||||
} // namespace detail
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
|
||||
#if defined(CCCL_ENABLE_DEVICE_ASSERTIONS)
|
||||
_CCCL_END_NV_DIAG_SUPPRESS()
|
||||
#endif // CCCL_ENABLE_DEVICE_ASSERTIONS
|
||||
@@ -1,147 +0,0 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cub/util_device.cuh>
|
||||
|
||||
#include <thrust/system/cuda/detail/core/triple_chevron_launch.h>
|
||||
|
||||
#include <cuda/__device/compute_capability.h>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail
|
||||
{
|
||||
struct TripleChevronFactory
|
||||
{
|
||||
CUB_RUNTIME_FUNCTION void __assert_pdl_allowed(bool dependent_launch) const
|
||||
{
|
||||
if (dependent_launch)
|
||||
{
|
||||
[[maybe_unused]] int sm_version = 0;
|
||||
_CCCL_ASSERT(SmVersion(sm_version) == cudaSuccess, "Failed to query SM compute capability");
|
||||
if (sm_version >= 900)
|
||||
{
|
||||
[[maybe_unused]] ::cuda::compute_capability cc;
|
||||
_CCCL_ASSERT(PtxComputeCap(cc) == cudaSuccess, "Failed to query PTX compute capability");
|
||||
_CCCL_ASSERT((cc >= ::cuda::compute_capability{9, 0}),
|
||||
"Enabling PDL for a kernel launch requires CC 9.0+ PTX/SASS when running on SM90+");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
CUB_RUNTIME_FUNCTION THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron operator()(
|
||||
dim3 grid, dim3 block, ::cuda::std::size_t shared_mem, ::cudaStream_t stream, bool dependent_launch = false) const
|
||||
{
|
||||
__assert_pdl_allowed(dependent_launch);
|
||||
return THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(grid, block, shared_mem, stream, dependent_launch);
|
||||
}
|
||||
|
||||
template <class T = void>
|
||||
CUB_RUNTIME_FUNCTION ::cudaError_t PtxVersion(int& version)
|
||||
{
|
||||
return cub::PtxVersion<T>(version);
|
||||
}
|
||||
|
||||
template <class T = void>
|
||||
CUB_RUNTIME_FUNCTION ::cudaError_t PtxComputeCap(::cuda::compute_capability& cc) const
|
||||
{
|
||||
return ptx_compute_cap<T>(cc);
|
||||
}
|
||||
|
||||
_CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t MultiProcessorCount(int& sm_count) const
|
||||
{
|
||||
int device_ordinal;
|
||||
::cudaError_t error = CubDebug(::cudaGetDevice(&device_ordinal));
|
||||
if (::cudaSuccess != error)
|
||||
{
|
||||
return error;
|
||||
}
|
||||
|
||||
// Get SM count
|
||||
return ::cudaDeviceGetAttribute(&sm_count, ::cudaDevAttrMultiProcessorCount, device_ordinal);
|
||||
}
|
||||
|
||||
template <typename Kernel>
|
||||
_CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t
|
||||
MaxSmOccupancy(int& sm_occupancy, Kernel kernel_ptr, int block_size, int dynamic_smem_bytes = 0)
|
||||
{
|
||||
return ::cudaOccupancyMaxActiveBlocksPerMultiprocessor(&sm_occupancy, kernel_ptr, block_size, dynamic_smem_bytes);
|
||||
}
|
||||
|
||||
_CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t MaxGridDimX(int& max_grid_dim_x) const
|
||||
{
|
||||
int device_ordinal;
|
||||
::cudaError_t error = CubDebug(::cudaGetDevice(&device_ordinal));
|
||||
if (::cudaSuccess != error)
|
||||
{
|
||||
return error;
|
||||
}
|
||||
|
||||
// Get max grid dimension
|
||||
return ::cudaDeviceGetAttribute(&max_grid_dim_x, ::cudaDevAttrMaxGridDimX, device_ordinal);
|
||||
}
|
||||
|
||||
_CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t
|
||||
MemsetAsync(void* dst, unsigned char value, size_t num_bytes, ::cudaStream_t stream) const
|
||||
{
|
||||
return ::cudaMemsetAsync(dst, value, num_bytes, stream);
|
||||
}
|
||||
|
||||
_CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t
|
||||
MemcpyAsync(void* dst, const void* src, size_t num_bytes, ::cudaMemcpyKind kind, ::cudaStream_t stream) const
|
||||
{
|
||||
return ::cudaMemcpyAsync(dst, src, num_bytes, kind, stream);
|
||||
}
|
||||
|
||||
// TODO(bgruber): this is very similar to thrust::cuda_cub::core::get_max_shared_memory_per_block. We should unify
|
||||
// this.
|
||||
_CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION cudaError_t MaxSharedMemory(int& max_shared_memory) const
|
||||
{
|
||||
int device = 0;
|
||||
auto error = CubDebug(cudaGetDevice(&device));
|
||||
if (error != cudaSuccess)
|
||||
{
|
||||
return error;
|
||||
}
|
||||
|
||||
return cudaDeviceGetAttribute(&max_shared_memory, cudaDevAttrMaxSharedMemoryPerBlock, device);
|
||||
}
|
||||
|
||||
template <typename Kernel>
|
||||
_CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t
|
||||
max_dynamic_smem_size_for(int& max_dynamic_smem_size, [[maybe_unused]] Kernel kernel_ptr)
|
||||
{
|
||||
NV_IF_ELSE_TARGET(NV_IS_HOST, //
|
||||
({ return MaxPotentialDynamicSmemBytes(max_dynamic_smem_size, kernel_ptr); }),
|
||||
({
|
||||
::cudaFuncAttributes func_attrs{};
|
||||
if (const auto error = CubDebug(::cudaFuncGetAttributes(&func_attrs, kernel_ptr)))
|
||||
{
|
||||
return error;
|
||||
}
|
||||
max_dynamic_smem_size = func_attrs.maxDynamicSharedSizeBytes;
|
||||
return cudaSuccess;
|
||||
}))
|
||||
}
|
||||
|
||||
template <typename Kernel>
|
||||
_CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t set_max_dynamic_smem_size_for(Kernel kernel_ptr, int smem_size)
|
||||
{
|
||||
return CubDebug(::cudaFuncSetAttribute(kernel_ptr, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
|
||||
}
|
||||
};
|
||||
} // namespace detail
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,113 +0,0 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
// SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cub/detail/fast_modulo_division.cuh> // fast_div_mod
|
||||
|
||||
#include <cuda/std/__mdspan/extents.h>
|
||||
#include <cuda/std/__type_traits/make_unsigned.h>
|
||||
#include <cuda/std/__utility/integer_sequence.h>
|
||||
#include <cuda/std/array>
|
||||
#include <cuda/std/cstddef>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
namespace detail
|
||||
{
|
||||
_CCCL_DIAG_PUSH
|
||||
_CCCL_DIAG_SUPPRESS_MSVC(4702) // unreachable code (even if there are no branches!)
|
||||
|
||||
// Compute the submdspan size of a given rank
|
||||
template <typename IndexType, size_t... Extents>
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr ::cuda::std::make_unsigned_t<IndexType>
|
||||
size_range(const ::cuda::std::extents<IndexType, Extents...>& ext, int start, int end)
|
||||
{
|
||||
_CCCL_ASSERT(start >= 0 && end <= static_cast<int>(ext.rank()), "invalid start or end");
|
||||
::cuda::std::make_unsigned_t<IndexType> s = 1;
|
||||
for (auto i = start; i < end; i++)
|
||||
{
|
||||
s *= ext.extent(i);
|
||||
}
|
||||
return s;
|
||||
}
|
||||
|
||||
_CCCL_DIAG_POP // MSVC(4702)
|
||||
|
||||
template <typename IndexType, size_t... Extents>
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr ::cuda::std::make_unsigned_t<IndexType>
|
||||
size(const ::cuda::std::extents<IndexType, Extents...>& ext)
|
||||
{
|
||||
return cub::detail::size_range(ext, 0, static_cast<int>(ext.rank()));
|
||||
}
|
||||
|
||||
template <bool IsLayoutRight, int Position, typename IndexType, size_t... E>
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE_API auto sub_size_fast_div_mod_impl(const ::cuda::std::extents<IndexType, E...>& ext)
|
||||
{
|
||||
using fast_mod_div_t = fast_div_mod<IndexType>;
|
||||
constexpr auto start = IsLayoutRight ? Position + 1 : 0;
|
||||
constexpr auto end = IsLayoutRight ? sizeof...(E) : Position;
|
||||
return fast_mod_div_t(cub::detail::size_range(ext, start, end));
|
||||
}
|
||||
|
||||
// precompute modulo/division for each submdspan size (by rank)
|
||||
template <bool IsLayoutRight, typename IndexType, size_t... E, size_t... Positions>
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE_API auto
|
||||
sub_sizes_fast_div_mod(const ::cuda::std::extents<IndexType, E...>& ext, ::cuda::std::index_sequence<Positions...> = {})
|
||||
{
|
||||
using fast_mod_div_t = fast_div_mod<IndexType>;
|
||||
using array_t = ::cuda::std::array<fast_mod_div_t, sizeof...(Positions)>;
|
||||
return array_t{cub::detail::sub_size_fast_div_mod_impl<IsLayoutRight, Positions>(ext)...};
|
||||
}
|
||||
|
||||
// precompute modulo/division for each mdspan extent
|
||||
template <typename IndexType, size_t... E, size_t... Positions>
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE_API auto
|
||||
extents_fast_div_mod(const ::cuda::std::extents<IndexType, E...>& ext, ::cuda::std::index_sequence<Positions...> = {})
|
||||
{
|
||||
using fast_mod_div_t = fast_div_mod<IndexType>;
|
||||
using array_t = ::cuda::std::array<fast_mod_div_t, sizeof...(Positions)>;
|
||||
return array_t{fast_mod_div_t(ext.extent(Positions))...};
|
||||
}
|
||||
|
||||
// GCC <= 9 constexpr workaround: Extent must be passed as type only, even const Extent& doesn't work
|
||||
template <typename Extents>
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr bool are_extents_in_range_static(int start, int end)
|
||||
{
|
||||
for (auto i = start; i < end; i++)
|
||||
{
|
||||
if (Extents::static_extent(i) == ::cuda::std::dynamic_extent)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
template <typename MappingTypeLhs, typename MappingTypeRhs>
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE_API bool
|
||||
have_same_strides(const MappingTypeLhs& mapping_lhs, const MappingTypeRhs& mapping_rhs)
|
||||
{
|
||||
auto extents_lhs = mapping_lhs.extents();
|
||||
auto extents_rhs = mapping_rhs.extents();
|
||||
_CCCL_ASSERT(extents_lhs.rank() == extents_rhs.rank(), "extents must have the same rank");
|
||||
for (size_t i = 0; i < extents_lhs.rank(); i++)
|
||||
{
|
||||
if (mapping_lhs.stride(i) != mapping_rhs.stride(i))
|
||||
{
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
} // namespace detail
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,691 +0,0 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
|
||||
// SPDX-License-Identifier: BSD-3
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cuda/std/__algorithm/max.h>
|
||||
#include <cuda/std/__algorithm/min.h>
|
||||
#include <cuda/std/__bit/bit_cast.h>
|
||||
#include <cuda/std/__cmath/exponential_functions.h>
|
||||
#include <cuda/std/__cmath/isinf.h>
|
||||
#include <cuda/std/__type_traits/enable_if.h>
|
||||
#include <cuda/std/__type_traits/is_arithmetic.h>
|
||||
#include <cuda/std/__type_traits/is_floating_point.h>
|
||||
#include <cuda/std/__type_traits/is_same.h>
|
||||
#include <cuda/std/array>
|
||||
#include <cuda/std/climits>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail::rfa
|
||||
{
|
||||
// jump table for indexing into data
|
||||
inline constexpr int cub_rfa_max_jump = 5;
|
||||
static_assert(cub_rfa_max_jump <= 5, "cub_rfa_max_jump must be less than or equal to 5");
|
||||
|
||||
template <typename FType, int Len>
|
||||
static _CCCL_DEVICE FType* get_shared_bin_array()
|
||||
{
|
||||
static __shared__ FType bin_computed_array[Len];
|
||||
return bin_computed_array;
|
||||
}
|
||||
|
||||
//! Class to hold a reproducible summation of the numbers passed to it
|
||||
//!
|
||||
//! @param FType Floating-point data type; either `float` or `double
|
||||
//! @param Fold Number of collectors in the binned number (K-fold), used for reproducible summation. Defaults to 3.
|
||||
template <class FType, int Fold = 3, ::cuda::std::enable_if_t<::cuda::std::is_floating_point_v<FType>>* = nullptr>
|
||||
class alignas(2 * sizeof(FType)) ReproducibleFloatingAccumulator
|
||||
{
|
||||
public:
|
||||
using ftype = FType;
|
||||
|
||||
private:
|
||||
::cuda::std::array<ftype, 2 * Fold> data{};
|
||||
|
||||
/// Floating-point precision bin width
|
||||
static constexpr int bin_width = ::cuda::std::is_same_v<ftype, double> ? 40 : 13;
|
||||
static constexpr int min_exp = ::cuda::std::numeric_limits<ftype>::min_exponent;
|
||||
static constexpr int max_exp = ::cuda::std::numeric_limits<ftype>::max_exponent;
|
||||
static constexpr int mant_dig = ::cuda::std::numeric_limits<ftype>::digits;
|
||||
|
||||
public:
|
||||
/// Binned floating-point maximum index
|
||||
static constexpr int max_index = ((max_exp - min_exp + mant_dig - 1) / bin_width) - 1;
|
||||
|
||||
// The maximum floating-point fold supported by the library
|
||||
static constexpr auto max_fold = max_index + 1;
|
||||
|
||||
_CCCL_DEVICE static ftype initialize_bin(int index) noexcept
|
||||
{
|
||||
if (index == 0)
|
||||
{
|
||||
if constexpr (::cuda::std::is_same_v<ftype, float>)
|
||||
{
|
||||
return ::cuda::std::ldexp(0.75, max_exp);
|
||||
}
|
||||
else
|
||||
{
|
||||
return 2.0 * ::cuda::std::ldexp(0.75, max_exp - 1);
|
||||
}
|
||||
}
|
||||
|
||||
if (index > 0 && index <= max_index)
|
||||
{
|
||||
return ::cuda::std::ldexp(0.75, max_exp + mant_dig - bin_width + 1 - index * bin_width);
|
||||
}
|
||||
else
|
||||
{
|
||||
return ::cuda::std::ldexp(0.75, max_exp + mant_dig - bin_width + 1 - max_index * bin_width);
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
/// Binned floating-point compression factor
|
||||
/// This factor is used to scale down inputs before deposition into the bin of
|
||||
/// highest index
|
||||
static constexpr auto compression = 1.0 / (1 << (mant_dig - bin_width + 1));
|
||||
/// Binned double precision expansion factor
|
||||
/// This factor is used to scale up inputs after deposition into the bin of
|
||||
/// highest index
|
||||
static constexpr auto expansion = 1.0 * (1 << (mant_dig - bin_width + 1));
|
||||
static constexpr auto exp_bias = max_exp - 2;
|
||||
|
||||
/// Return a binned floating-point bin
|
||||
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE static ftype binned_bins(int index)
|
||||
{
|
||||
ftype* bins = get_shared_bin_array<ftype, max_index + max_fold>();
|
||||
return bins[index];
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE static uint32_t& get_bit_representation(float& x) noexcept
|
||||
{
|
||||
return *reinterpret_cast<uint32_t*>(&x);
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE static uint64_t& get_bit_representation(double& x) noexcept
|
||||
{
|
||||
return *reinterpret_cast<uint64_t*>(&x);
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE static uint32_t get_bit_representation(const float& x) noexcept
|
||||
{
|
||||
return ::cuda::std::bit_cast<uint32_t>(x);
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE static uint64_t get_bit_representation(const double& x) noexcept
|
||||
{
|
||||
return ::cuda::std::bit_cast<uint64_t>(x);
|
||||
}
|
||||
|
||||
/// Return primary vector value const ref
|
||||
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE const ftype& primary(int i) const noexcept
|
||||
{
|
||||
if constexpr (Fold <= cub_rfa_max_jump)
|
||||
{
|
||||
switch (i)
|
||||
{
|
||||
case 0:
|
||||
if constexpr (Fold >= 1)
|
||||
{
|
||||
return data[0];
|
||||
}
|
||||
[[fallthrough]];
|
||||
case 1:
|
||||
if constexpr (Fold >= 2)
|
||||
{
|
||||
return data[1];
|
||||
}
|
||||
[[fallthrough]];
|
||||
case 2:
|
||||
if constexpr (Fold >= 3)
|
||||
{
|
||||
return data[2];
|
||||
}
|
||||
[[fallthrough]];
|
||||
case 3:
|
||||
if constexpr (Fold >= 4)
|
||||
{
|
||||
return data[3];
|
||||
}
|
||||
[[fallthrough]];
|
||||
case 4:
|
||||
if constexpr (Fold >= 5)
|
||||
{
|
||||
return data[4];
|
||||
}
|
||||
[[fallthrough]];
|
||||
default:
|
||||
return data[Fold - 1];
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
return data[i];
|
||||
}
|
||||
}
|
||||
|
||||
/// Return carry vector value const ref
|
||||
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE const ftype& carry(int i) const noexcept
|
||||
{
|
||||
if (Fold <= cub_rfa_max_jump)
|
||||
{
|
||||
switch (i)
|
||||
{
|
||||
case 0:
|
||||
if (Fold >= 1)
|
||||
{
|
||||
return data[Fold + 0];
|
||||
}
|
||||
[[fallthrough]];
|
||||
case 1:
|
||||
if (Fold >= 2)
|
||||
{
|
||||
return data[Fold + 1];
|
||||
}
|
||||
[[fallthrough]];
|
||||
case 2:
|
||||
if (Fold >= 3)
|
||||
{
|
||||
return data[Fold + 2];
|
||||
}
|
||||
[[fallthrough]];
|
||||
case 3:
|
||||
if (Fold >= 4)
|
||||
{
|
||||
return data[Fold + 3];
|
||||
}
|
||||
[[fallthrough]];
|
||||
case 4:
|
||||
if (Fold >= 5)
|
||||
{
|
||||
return data[Fold + 4];
|
||||
}
|
||||
[[fallthrough]];
|
||||
default:
|
||||
return data[2 * Fold - 1];
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
return data[Fold + i]; // NOLINT(bugprone-misplaced-widening-cast)
|
||||
}
|
||||
}
|
||||
|
||||
/// Return primary vector value ref
|
||||
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ftype& primary(int i) noexcept
|
||||
{
|
||||
const auto& c = *this;
|
||||
return const_cast<ftype&>(c.primary(i));
|
||||
}
|
||||
|
||||
/// Return carry vector value ref
|
||||
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ftype& carry(int i) noexcept
|
||||
{
|
||||
const auto& c = *this;
|
||||
return const_cast<ftype&>(c.carry(i));
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE static int exp_val(const ftype x) noexcept
|
||||
{
|
||||
const auto bits = get_bit_representation(x);
|
||||
return (bits >> (mant_dig - 1)) & (2 * max_exp - 1);
|
||||
}
|
||||
|
||||
/// Get index of float-point precision
|
||||
/// The index of a non-binned type is the smallest index a binned type would
|
||||
/// need to have to sum it reproducibly. Higher indices correspond to smaller
|
||||
/// bins.
|
||||
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE static int binned_dindex(const ftype x)
|
||||
{
|
||||
int exp = exp_val(x);
|
||||
|
||||
if (exp != 0)
|
||||
{
|
||||
return ((max_exp + exp_bias) - exp) / bin_width;
|
||||
}
|
||||
if (x == 0.0)
|
||||
{
|
||||
return max_index;
|
||||
}
|
||||
else
|
||||
{
|
||||
(void) ::cuda::std::frexpf(x, &exp);
|
||||
return (::cuda::std::min) ((max_exp - exp) / bin_width, +max_index);
|
||||
}
|
||||
}
|
||||
|
||||
/// Get index of manually specified binned double precision
|
||||
/// The index of a binned type is the bin that it corresponds to. Higher
|
||||
/// indices correspond to smaller bins.
|
||||
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE int binned_index() const
|
||||
{
|
||||
return ((max_exp + mant_dig - bin_width + 1 + exp_bias) - exp_val(primary(0))) / bin_width;
|
||||
}
|
||||
|
||||
/// Check if index of manually specified binned floating-point is 0
|
||||
/// A quick check to determine if the index is 0
|
||||
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE bool is_binned_index_zero() const
|
||||
{
|
||||
return exp_val(primary(0)) == max_exp + exp_bias;
|
||||
}
|
||||
|
||||
//! Update manually specified binned fp with a scalar (X -> Y)
|
||||
//!
|
||||
//! This method updates the binned fp to an index suitable for adding numbers
|
||||
//! with absolute value less than @p max_abs_val
|
||||
_CCCL_DEVICE void binned_update(const ftype max_abs_val)
|
||||
{
|
||||
int X_index = binned_dindex(max_abs_val);
|
||||
int shift = binned_index() - X_index;
|
||||
if (shift > 0)
|
||||
{
|
||||
_CCCL_PRAGMA_UNROLL_FULL()
|
||||
for (int i = Fold - 1; i >= 1; i--)
|
||||
{
|
||||
if (i < shift)
|
||||
{
|
||||
break;
|
||||
}
|
||||
primary(i) = primary((i - shift));
|
||||
carry(i) = carry((i - shift));
|
||||
}
|
||||
_CCCL_PRAGMA_UNROLL_FULL()
|
||||
for (int j = 0; j < Fold; j++)
|
||||
{
|
||||
if (j >= shift)
|
||||
{
|
||||
break;
|
||||
}
|
||||
primary(j) = binned_bins(j + X_index);
|
||||
carry(j) = 0.0;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
//! Add scalar @p X to suitably binned manually specified binned fp (Y += X)
|
||||
//!
|
||||
//! Performs the operation Y += X on an binned type Y where the index of Y is
|
||||
//! larger than the index of @p X
|
||||
_CCCL_DEVICE void binned_deposit(const ftype X)
|
||||
{
|
||||
ftype M;
|
||||
ftype x = X;
|
||||
|
||||
if (is_binned_index_zero())
|
||||
{
|
||||
M = primary(0);
|
||||
ftype qd = x * compression;
|
||||
auto& ql = get_bit_representation(qd);
|
||||
ql |= 1;
|
||||
qd += M;
|
||||
primary(0) = qd;
|
||||
M -= qd;
|
||||
M *= expansion * 0.5;
|
||||
x += M;
|
||||
x += M;
|
||||
_CCCL_PRAGMA_UNROLL_FULL()
|
||||
for (int i = 1; i < Fold - 1; i++)
|
||||
{
|
||||
M = primary(i);
|
||||
qd = x;
|
||||
ql |= 1;
|
||||
qd += M;
|
||||
primary(i) = qd;
|
||||
M -= qd;
|
||||
x += M;
|
||||
}
|
||||
qd = x;
|
||||
ql |= 1;
|
||||
primary((Fold - 1)) += qd;
|
||||
}
|
||||
else
|
||||
{
|
||||
ftype qd = x;
|
||||
auto& ql = get_bit_representation(qd);
|
||||
_CCCL_PRAGMA_UNROLL_FULL()
|
||||
for (int i = 0; i < Fold - 1; i++)
|
||||
{
|
||||
M = primary(i);
|
||||
qd = x;
|
||||
ql |= 1;
|
||||
qd += M;
|
||||
primary(i) = qd;
|
||||
M -= qd;
|
||||
x += M;
|
||||
}
|
||||
qd = x;
|
||||
ql |= 1;
|
||||
primary((Fold - 1)) += qd;
|
||||
}
|
||||
}
|
||||
|
||||
//! Renormalize manually specified binned double precision
|
||||
//!
|
||||
//! Renormalization keeps the primary vector within the necessary bins by
|
||||
//! shifting over to the carry vector
|
||||
_CCCL_DEVICE _CCCL_FORCEINLINE void binned_renorm()
|
||||
{
|
||||
_CCCL_PRAGMA_UNROLL_FULL()
|
||||
for (int i = 0; i < Fold; i++)
|
||||
{
|
||||
auto tmp_renormd = primary(i);
|
||||
auto& tmp_renorml = get_bit_representation(tmp_renormd);
|
||||
|
||||
carry(i) += static_cast<int>((tmp_renorml >> (mant_dig - 3)) & 3) - 2;
|
||||
|
||||
tmp_renorml &= ~(1ull << (mant_dig - 3));
|
||||
tmp_renorml |= 1ull << (mant_dig - 2);
|
||||
primary(i) = tmp_renormd;
|
||||
}
|
||||
}
|
||||
|
||||
//! Add scalar to manually specified binned fp (Y += X)
|
||||
//!
|
||||
//! Performs the operation Y += X on an binned type Y
|
||||
_CCCL_DEVICE _CCCL_FORCEINLINE void binned_add(const ftype x)
|
||||
{
|
||||
binned_update(x);
|
||||
binned_deposit(x);
|
||||
binned_renorm();
|
||||
}
|
||||
|
||||
//! Add two manually specified binned fp (Y += X)
|
||||
//! Performs the operation Y += X
|
||||
//!
|
||||
//! @param x Another binned fp of the same type
|
||||
_CCCL_DEVICE void binned_add(const ReproducibleFloatingAccumulator& x)
|
||||
{
|
||||
const auto X_index = x.binned_index();
|
||||
const auto Y_index = this->binned_index();
|
||||
const auto shift = Y_index - X_index;
|
||||
if (shift > 0)
|
||||
{
|
||||
// shift Y upwards and add X to Y
|
||||
_CCCL_PRAGMA_UNROLL_FULL()
|
||||
for (int i = Fold - 1; i >= 1; i--)
|
||||
{
|
||||
if (i < shift)
|
||||
{
|
||||
break;
|
||||
}
|
||||
primary(i) = x.primary(i) + (primary((i - shift)) - binned_bins(i - shift + Y_index));
|
||||
carry(i) = x.carry(i) + carry((i - shift));
|
||||
}
|
||||
_CCCL_PRAGMA_UNROLL_FULL()
|
||||
for (int i = 0; i < Fold; i++)
|
||||
{
|
||||
if (i == shift)
|
||||
{
|
||||
break;
|
||||
}
|
||||
primary(i) = x.primary(i);
|
||||
carry(i) = x.carry(i);
|
||||
}
|
||||
}
|
||||
else if (shift < 0)
|
||||
{
|
||||
// shift X upwards and add X to Y
|
||||
_CCCL_PRAGMA_UNROLL_FULL()
|
||||
for (int i = 0; i < Fold; i++)
|
||||
{
|
||||
if (i < -shift)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
primary(i) += x.primary((i + shift)) - binned_bins(X_index + i + shift);
|
||||
carry(i) += x.carry((i + shift));
|
||||
}
|
||||
}
|
||||
else if (shift == 0)
|
||||
{
|
||||
// add X to Y
|
||||
_CCCL_PRAGMA_UNROLL_FULL()
|
||||
for (int i = 0; i < Fold; i++)
|
||||
{
|
||||
primary(i) += x.primary(i) - binned_bins(i + X_index);
|
||||
carry(i) += x.carry(i);
|
||||
}
|
||||
}
|
||||
|
||||
binned_renorm();
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_DEVICE double conv_binned_to_double() const
|
||||
{
|
||||
int i = 0;
|
||||
double Y = 0.0;
|
||||
const auto X_index = binned_index();
|
||||
if (X_index <= (3 * mant_dig) / bin_width)
|
||||
{
|
||||
double scale_down = ::cuda::std::ldexpf(0.5f, 1 - (2 * mant_dig - bin_width));
|
||||
double scale_up = ::cuda::std::ldexpf(0.5f, 1 - (2 * mant_dig - bin_width));
|
||||
int scaled = ::cuda::std::max(::cuda::std::min(Fold, (3 * mant_dig) / bin_width - X_index), 0);
|
||||
if (X_index == 0)
|
||||
{
|
||||
Y += carry(0) * ((binned_bins(0 + X_index) / 6.0) * scale_down * expansion);
|
||||
Y += carry(1) * ((binned_bins(1 + X_index) / 6.0) * scale_down);
|
||||
Y += (primary(0) - binned_bins(0 + X_index)) * scale_down * expansion;
|
||||
i = 2;
|
||||
}
|
||||
else
|
||||
{
|
||||
Y += carry(0) * ((binned_bins(0 + X_index) / 6.0) * scale_down);
|
||||
i = 1;
|
||||
}
|
||||
for (; i < scaled; i++)
|
||||
{
|
||||
Y += carry(i) * ((binned_bins(i + X_index) / 6.0) * scale_down);
|
||||
Y += (primary((i - 1)) - binned_bins(i - 1 + X_index)) * scale_down;
|
||||
}
|
||||
if (i == Fold)
|
||||
{
|
||||
Y += (primary((Fold - 1)) - binned_bins(Fold - 1 + X_index)) * scale_down;
|
||||
return Y * scale_up;
|
||||
}
|
||||
if (::cuda::std::isinf(Y * scale_up))
|
||||
{
|
||||
return Y * scale_up;
|
||||
}
|
||||
Y *= scale_up;
|
||||
for (; i < Fold; i++)
|
||||
{
|
||||
Y += carry(i) * (binned_bins(i + X_index) / 6.0);
|
||||
Y += primary((i - 1)) - binned_bins(i - 1 + X_index);
|
||||
}
|
||||
Y += primary((Fold - 1)) - binned_bins(Fold - 1 + X_index);
|
||||
}
|
||||
else
|
||||
{
|
||||
Y += carry(0) * (binned_bins(0 + X_index) / 6.0);
|
||||
for (i = 1; i < Fold; i++)
|
||||
{
|
||||
Y += carry(i) * (binned_bins(i + X_index) / 6.0);
|
||||
Y += (primary((i - 1)) - binned_bins(i - 1 + X_index));
|
||||
}
|
||||
Y += (primary((Fold - 1)) - binned_bins(Fold - 1 + X_index));
|
||||
}
|
||||
return Y;
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_DEVICE float conv_binned_to_float() const
|
||||
{
|
||||
int i = 0;
|
||||
double Y = 0.0;
|
||||
|
||||
// Note that the following order of summation is in order of decreasing
|
||||
// exponent. The following code is specific to SBWIDTH=13, FLT_MANT_DIG=24, and
|
||||
// the number of carries equal to 1.
|
||||
const auto X_index = binned_index();
|
||||
if (X_index == 0)
|
||||
{
|
||||
Y += static_cast<double>(carry(0)) * static_cast<double>(binned_bins(0 + X_index) / 6.0)
|
||||
* static_cast<double>(expansion);
|
||||
Y += static_cast<double>(carry(1)) * static_cast<double>(binned_bins(1 + X_index) / 6.0);
|
||||
Y += static_cast<double>(primary(0) - binned_bins(0 + X_index)) * static_cast<double>(expansion);
|
||||
i = 2;
|
||||
}
|
||||
else
|
||||
{
|
||||
Y += static_cast<double>(carry(0)) * static_cast<double>((binned_bins(0 + X_index) / 6.0));
|
||||
i = 1;
|
||||
}
|
||||
for (; i < Fold; i++)
|
||||
{
|
||||
Y += static_cast<double>(carry(i)) * static_cast<double>(binned_bins(i + X_index) / 6.0);
|
||||
Y += static_cast<double>(primary(i - 1) - binned_bins(i - 1 + X_index));
|
||||
}
|
||||
Y += static_cast<double>(primary(Fold - 1) - binned_bins(Fold - 1 + X_index));
|
||||
return static_cast<float>(Y);
|
||||
}
|
||||
|
||||
public:
|
||||
ReproducibleFloatingAccumulator() = default;
|
||||
|
||||
/// Set the binned fp to zero
|
||||
_CCCL_DEVICE void zero() noexcept
|
||||
{
|
||||
data = {};
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE constexpr int endurance() const noexcept
|
||||
{
|
||||
return 1 << (mant_dig - bin_width - 2);
|
||||
}
|
||||
|
||||
//! Accumulate an arithmetic @p x into the binned fp.
|
||||
//! NOTE: Casts @p x to the type of the binned fp
|
||||
_CCCL_TEMPLATE(typename U)
|
||||
_CCCL_REQUIRES(::cuda::std::is_arithmetic_v<U>)
|
||||
_CCCL_DEVICE ReproducibleFloatingAccumulator& operator+=(const U x)
|
||||
{
|
||||
binned_add(static_cast<ftype>(x));
|
||||
return *this;
|
||||
}
|
||||
|
||||
//! Accumulate-subtract an arithmetic @p x into the binned fp.
|
||||
//! NOTE: Casts @p x to the type of the binned fp
|
||||
_CCCL_TEMPLATE(typename U)
|
||||
_CCCL_REQUIRES(::cuda::std::is_arithmetic_v<U>)
|
||||
_CCCL_DEVICE ReproducibleFloatingAccumulator& operator-=(const U x)
|
||||
{
|
||||
binned_add(-static_cast<ftype>(x));
|
||||
return *this;
|
||||
}
|
||||
|
||||
/// Accumulate a binned fp @p x into the binned fp.
|
||||
_CCCL_DEVICE ReproducibleFloatingAccumulator& operator+=(const ReproducibleFloatingAccumulator& other)
|
||||
{
|
||||
binned_add(other);
|
||||
return *this;
|
||||
}
|
||||
|
||||
//! Accumulate-subtract a binned fp @p other into the binned fp.
|
||||
//! NOTE: Makes a copy and performs arithmetic; slow.
|
||||
_CCCL_DEVICE ReproducibleFloatingAccumulator& operator-=(const ReproducibleFloatingAccumulator& other)
|
||||
{
|
||||
const auto temp = -other;
|
||||
binned_add(temp);
|
||||
}
|
||||
|
||||
_CCCL_DEVICE friend bool operator==(const ReproducibleFloatingAccumulator& a, const ReproducibleFloatingAccumulator& b)
|
||||
{
|
||||
return a.data == b.data;
|
||||
}
|
||||
|
||||
_CCCL_DEVICE friend bool operator!=(const ReproducibleFloatingAccumulator& a, const ReproducibleFloatingAccumulator& b)
|
||||
{
|
||||
return !(a == b);
|
||||
}
|
||||
|
||||
//! Sets this binned fp equal to the arithmetic value @p x
|
||||
//! NOTE: Casts @p x to the type of the binned fp
|
||||
_CCCL_TEMPLATE(typename U)
|
||||
_CCCL_REQUIRES(::cuda::std::is_arithmetic_v<U>)
|
||||
_CCCL_DEVICE ReproducibleFloatingAccumulator& operator=(const U x)
|
||||
{
|
||||
zero();
|
||||
binned_add(static_cast<ftype>(x));
|
||||
return *this;
|
||||
}
|
||||
|
||||
//! Returns the negative of this binned fp
|
||||
//! NOTE: Makes a copy and performs arithmetic; slow.
|
||||
[[nodiscard]] _CCCL_DEVICE ReproducibleFloatingAccumulator operator-() const
|
||||
{
|
||||
ReproducibleFloatingAccumulator temp = *this;
|
||||
if (primary(0) != 0.0)
|
||||
{
|
||||
_CCCL_PRAGMA_UNROLL_FULL()
|
||||
for (int i = 0; i < Fold; i++)
|
||||
{
|
||||
temp.primary(i) = binned_bins(i + binned_index()) - (primary(i) - binned_bins(i + binned_index()));
|
||||
temp.carry(i) = -carry(i);
|
||||
}
|
||||
}
|
||||
return temp;
|
||||
}
|
||||
|
||||
/// Convert this binned fp into its native floating-point representation
|
||||
[[nodiscard]] _CCCL_DEVICE ftype conv_to_fp() const
|
||||
{
|
||||
if (::cuda::std::is_same_v<ftype, float>)
|
||||
{
|
||||
return conv_binned_to_float();
|
||||
}
|
||||
else
|
||||
{
|
||||
return conv_binned_to_double();
|
||||
}
|
||||
}
|
||||
|
||||
/// Add @p x to the binned fp
|
||||
_CCCL_DEVICE void add(const ftype x)
|
||||
{
|
||||
binned_add(x);
|
||||
}
|
||||
|
||||
//////////////////////////////////////
|
||||
// MANUAL OPERATIONS; USE WISELY
|
||||
//////////////////////////////////////
|
||||
|
||||
//! Rebins for repeated accumulation of scalars with magnitude <= @p mav
|
||||
//!
|
||||
//! Once rebinned, `endurance` values <= @p mav can be added to the accumulator
|
||||
//! with `unsafe_add` after which `renorm()` must be called. See the source of
|
||||
//!`add()` for an example
|
||||
_CCCL_DEVICE void set_max_val(const ftype mav)
|
||||
{
|
||||
binned_update(mav);
|
||||
}
|
||||
|
||||
//! Add @p x to the binned fp
|
||||
//!
|
||||
//! This is intended to be used after a call to `set_max_abs_val()`
|
||||
_CCCL_DEVICE void unsafe_add(const ftype x)
|
||||
{
|
||||
binned_deposit(x);
|
||||
}
|
||||
|
||||
//! Renormalizes the binned fp
|
||||
//!
|
||||
//! This is intended to be used after a call to `set_max_abs_val()` and one or
|
||||
//! more calls to `unsafe_add()`
|
||||
_CCCL_DEVICE void renorm()
|
||||
{
|
||||
binned_renorm();
|
||||
}
|
||||
};
|
||||
} // namespace detail::rfa
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,145 +0,0 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cuda/argument>
|
||||
#include <cuda/std/__type_traits/integral_constant.h>
|
||||
#include <cuda/std/__type_traits/remove_cvref.h>
|
||||
#include <cuda/std/__utility/forward.h>
|
||||
#include <cuda/std/cstddef>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail::params
|
||||
{
|
||||
// =====================================================================
|
||||
// get_param — unified segment parameter access
|
||||
// =====================================================================
|
||||
|
||||
//! @brief Returns the value of an argument for a given segment index.
|
||||
//!
|
||||
//! @param[in] __arg Argument or argument wrapper to read.
|
||||
//! @param[in] __index Segment index to read for sequence arguments.
|
||||
//! @return The single argument value, or the sequence element at the given index.
|
||||
_CCCL_TEMPLATE(class _Tp, class _SegmentIndexT)
|
||||
_CCCL_REQUIRES((!::cuda::args::__is_wrapper_v<::cuda::std::remove_cvref_t<_Tp>>) )
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE constexpr auto get_param(_Tp&& __arg, [[maybe_unused]] _SegmentIndexT __index) noexcept
|
||||
{
|
||||
if constexpr (::cuda::args::__traits<::cuda::std::remove_cvref_t<_Tp>>::is_single_value)
|
||||
{
|
||||
return __arg;
|
||||
}
|
||||
else
|
||||
{
|
||||
return __arg[__index];
|
||||
}
|
||||
}
|
||||
|
||||
template <auto _Value, class _Tp, class _SegmentIndexT>
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE constexpr auto
|
||||
get_param(const ::cuda::args::constant<_Value, _Tp>& __arg, [[maybe_unused]] _SegmentIndexT __index) noexcept
|
||||
{
|
||||
return ::cuda::args::__unwrap(__arg);
|
||||
}
|
||||
|
||||
template <class _Arg, class _StaticBounds, class _SegmentIndexT>
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE constexpr auto
|
||||
get_param(const ::cuda::args::immediate<_Arg, _StaticBounds>& __arg, [[maybe_unused]] _SegmentIndexT __index) noexcept
|
||||
{
|
||||
return ::cuda::args::__unwrap(__arg);
|
||||
}
|
||||
|
||||
template <class _Arg, class _StaticBounds, class _SegmentIndexT>
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE constexpr auto
|
||||
get_param(const ::cuda::args::deferred<_Arg, _StaticBounds>& __arg, [[maybe_unused]] _SegmentIndexT __index) noexcept
|
||||
{
|
||||
return ::cuda::args::__unwrap(__arg);
|
||||
}
|
||||
|
||||
template <class _Arg, class _StaticBounds, class _SegmentIndexT>
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE constexpr auto
|
||||
get_param(const ::cuda::args::deferred_sequence<_Arg, _StaticBounds>& __arg, _SegmentIndexT __index) noexcept
|
||||
{
|
||||
return ::cuda::args::__unwrap(__arg)[__index];
|
||||
}
|
||||
|
||||
// =====================================================================
|
||||
// Discrete parameter support
|
||||
// =====================================================================
|
||||
|
||||
//! @brief Specifies a list of supported options for a parameter.
|
||||
template <typename T, T... Options>
|
||||
struct supported_options
|
||||
{
|
||||
static constexpr ::cuda::std::size_t count = sizeof...(Options);
|
||||
};
|
||||
|
||||
//! @brief Static discrete parameter — a single compile-time value that is also its only supported option.
|
||||
//!
|
||||
//! Holds no runtime value, so it cannot be put into a state that disagrees with its supported option, and
|
||||
//! @c dispatch_impl therefore always matches it. This is the safe representation for a compile-time-fixed discrete
|
||||
//! parameter (e.g. a statically known top-k selection direction): modeling such a parameter with a runtime value
|
||||
//! instead would risk that value silently disagreeing with the supported option (a no-op dispatch unless
|
||||
//! @c CCCL_ENABLE_ASSERTIONS is set).
|
||||
template <typename T, T Value>
|
||||
struct static_discrete_param
|
||||
{
|
||||
using value_type = T;
|
||||
using supported_options_t = supported_options<T, Value>;
|
||||
|
||||
template <typename SegmentIndexT>
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE constexpr T get_param(SegmentIndexT) const noexcept
|
||||
{
|
||||
return Value;
|
||||
}
|
||||
};
|
||||
|
||||
// =====================================================================
|
||||
// Discrete dispatch
|
||||
// =====================================================================
|
||||
|
||||
//! @brief Translates a runtime parameter value into a compile-time constant by matching
|
||||
//! against a list of supported options.
|
||||
//!
|
||||
//! @param[in] val Runtime value to match.
|
||||
//! @param[in] __supported_options Supported values for the parameter.
|
||||
//! @param[in] f Functor invoked with the matched compile-time constant.
|
||||
//! @return `true` if the value matches one of the supported options.
|
||||
template <typename T, T... Opts, typename Functor>
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE bool
|
||||
dispatch_impl(T val, [[maybe_unused]] supported_options<T, Opts...> __supported_options, Functor&& f)
|
||||
{
|
||||
const bool match_found = ((val == Opts ? (f(::cuda::std::integral_constant<T, Opts>{}), true) : false) || ...);
|
||||
_CCCL_ASSERT(match_found, "The given runtime parameter value is not in the supported list");
|
||||
return match_found;
|
||||
}
|
||||
|
||||
//! @brief Dispatcher that resolves a discrete parameter to a compile-time constant
|
||||
//! and invokes a functor with the matched option.
|
||||
//!
|
||||
//! @param[in] param Discrete parameter to resolve.
|
||||
//! @param[in] segment_id Segment index to read from `param`.
|
||||
//! @param[in] f Functor invoked with the matched compile-time constant.
|
||||
//! @return `true` if the parameter value matches one of its supported options.
|
||||
template <typename ParamT, typename SegmentIndexT, typename Functor>
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE bool dispatch_discrete(ParamT param, SegmentIndexT segment_id, Functor&& f)
|
||||
{
|
||||
using supported_list = typename ParamT::supported_options_t;
|
||||
auto param_value = param.get_param(segment_id);
|
||||
return CUB_NS_QUALIFIER::detail::params::dispatch_impl(
|
||||
param_value, supported_list{}, ::cuda::std::forward<Functor>(f));
|
||||
}
|
||||
} // namespace detail::params
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,167 +0,0 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved.
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.
|
||||
// SPDX-License-Identifier: BSD-3
|
||||
|
||||
/**
|
||||
* @file Utilities for strong memory operations.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cub/util_ptx.cuh>
|
||||
#include <cub/util_type.cuh>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
#ifndef _CCCL_DOXYGEN_INVOKED // Do not document
|
||||
|
||||
namespace detail
|
||||
{
|
||||
static _CCCL_DEVICE _CCCL_FORCEINLINE uint4 load_relaxed(uint4 const* ptr)
|
||||
{
|
||||
uint4 retval;
|
||||
NV_IF_ELSE_TARGET(
|
||||
NV_PROVIDES_SM_70,
|
||||
(asm volatile("ld.relaxed.gpu.v4.u32 {%0, %1, %2, %3}, [%4];" : "=r"(retval.x),
|
||||
"=r"(retval.y),
|
||||
"=r"(retval.z),
|
||||
"=r"(retval.w) : "l"(ptr) : "memory");),
|
||||
(asm volatile("ld.cg.v4.u32 {%0, %1, %2, %3}, [%4];" : "=r"(retval.x),
|
||||
"=r"(retval.y),
|
||||
"=r"(retval.z),
|
||||
"=r"(retval.w) : "l"(ptr) : "memory");));
|
||||
return retval;
|
||||
}
|
||||
|
||||
static _CCCL_DEVICE _CCCL_FORCEINLINE ulonglong2 load_relaxed(ulonglong2 const* ptr)
|
||||
{
|
||||
ulonglong2 retval;
|
||||
NV_IF_ELSE_TARGET(
|
||||
NV_PROVIDES_SM_70,
|
||||
(asm volatile("ld.relaxed.gpu.v2.u64 {%0, %1}, [%2];" : "=l"(retval.x), "=l"(retval.y) : "l"(ptr) : "memory");),
|
||||
(asm volatile("ld.cg.v2.u64 {%0, %1}, [%2];" : "=l"(retval.x), "=l"(retval.y) : "l"(ptr) : "memory");));
|
||||
return retval;
|
||||
}
|
||||
|
||||
static _CCCL_DEVICE _CCCL_FORCEINLINE ushort4 load_relaxed(ushort4 const* ptr)
|
||||
{
|
||||
ushort4 retval;
|
||||
NV_IF_ELSE_TARGET(
|
||||
NV_PROVIDES_SM_70,
|
||||
(asm volatile("ld.relaxed.gpu.v4.u16 {%0, %1, %2, %3}, [%4];" : "=h"(retval.x),
|
||||
"=h"(retval.y),
|
||||
"=h"(retval.z),
|
||||
"=h"(retval.w) : "l"(ptr) : "memory");),
|
||||
(asm volatile("ld.cg.v4.u16 {%0, %1, %2, %3}, [%4];" : "=h"(retval.x),
|
||||
"=h"(retval.y),
|
||||
"=h"(retval.z),
|
||||
"=h"(retval.w) : "l"(ptr) : "memory");));
|
||||
return retval;
|
||||
}
|
||||
|
||||
static _CCCL_DEVICE _CCCL_FORCEINLINE uint2 load_relaxed(uint2 const* ptr)
|
||||
{
|
||||
uint2 retval;
|
||||
NV_IF_ELSE_TARGET(
|
||||
NV_PROVIDES_SM_70,
|
||||
(asm volatile("ld.relaxed.gpu.v2.u32 {%0, %1}, [%2];" : "=r"(retval.x), "=r"(retval.y) : "l"(ptr) : "memory");),
|
||||
(asm volatile("ld.cg.v2.u32 {%0, %1}, [%2];" : "=r"(retval.x), "=r"(retval.y) : "l"(ptr) : "memory");));
|
||||
return retval;
|
||||
}
|
||||
|
||||
static _CCCL_DEVICE _CCCL_FORCEINLINE unsigned long long load_relaxed(unsigned long long const* ptr)
|
||||
{
|
||||
unsigned long long retval;
|
||||
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70,
|
||||
(asm volatile("ld.relaxed.gpu.u64 %0, [%1];" : "=l"(retval) : "l"(ptr) : "memory");),
|
||||
(asm volatile("ld.cg.u64 %0, [%1];" : "=l"(retval) : "l"(ptr) : "memory");));
|
||||
return retval;
|
||||
}
|
||||
|
||||
static _CCCL_DEVICE _CCCL_FORCEINLINE unsigned int load_relaxed(unsigned int const* ptr)
|
||||
{
|
||||
unsigned int retval;
|
||||
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70,
|
||||
(asm volatile("ld.relaxed.gpu.u32 %0, [%1];" : "=r"(retval) : "l"(ptr) : "memory");),
|
||||
(asm volatile("ld.cg.u32 %0, [%1];" : "=r"(retval) : "l"(ptr) : "memory");));
|
||||
|
||||
return retval;
|
||||
}
|
||||
|
||||
static _CCCL_DEVICE _CCCL_FORCEINLINE unsigned short load_relaxed(unsigned short const* ptr)
|
||||
{
|
||||
unsigned short retval;
|
||||
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70,
|
||||
(asm volatile("ld.relaxed.gpu.u16 %0, [%1];" : "=h"(retval) : "l"(ptr) : "memory");),
|
||||
(asm volatile("ld.cg.u16 %0, [%1];" : "=h"(retval) : "l"(ptr) : "memory");));
|
||||
return retval;
|
||||
}
|
||||
|
||||
static _CCCL_DEVICE _CCCL_FORCEINLINE unsigned char load_relaxed(unsigned char const* ptr)
|
||||
{
|
||||
unsigned short retval;
|
||||
NV_IF_ELSE_TARGET(
|
||||
NV_PROVIDES_SM_70,
|
||||
(asm volatile("{"
|
||||
" .reg .u8 datum;"
|
||||
" ld.relaxed.gpu.u8 datum, [%1];"
|
||||
" cvt.u16.u8 %0, datum;"
|
||||
"}" : "=h"(retval) : "l"(ptr) : "memory");),
|
||||
(asm volatile("{"
|
||||
" .reg .u8 datum;"
|
||||
" ld.cg.u8 datum, [%1];"
|
||||
" cvt.u16.u8 %0, datum;"
|
||||
"}" : "=h"(retval) : "l"(ptr) : "memory");));
|
||||
return (unsigned char) retval;
|
||||
}
|
||||
|
||||
static _CCCL_DEVICE _CCCL_FORCEINLINE ulonglong2 load_acquire(ulonglong2 const* ptr)
|
||||
{
|
||||
ulonglong2 retval;
|
||||
NV_IF_ELSE_TARGET(
|
||||
NV_PROVIDES_SM_70,
|
||||
(asm volatile("ld.acquire.gpu.v2.u64 {%0, %1}, [%2];" : "=l"(retval.x), "=l"(retval.y) : "l"(ptr) : "memory");),
|
||||
({
|
||||
asm volatile("ld.cg.v2.u64 {%0, %1}, [%2];" : "=l"(retval.x), "=l"(retval.y) : "l"(ptr) : "memory");
|
||||
__threadfence();
|
||||
}));
|
||||
return retval;
|
||||
}
|
||||
|
||||
static _CCCL_DEVICE _CCCL_FORCEINLINE uint2 load_acquire(uint2 const* ptr)
|
||||
{
|
||||
uint2 retval;
|
||||
NV_IF_ELSE_TARGET(
|
||||
NV_PROVIDES_SM_70,
|
||||
(asm volatile("ld.acquire.gpu.v2.u32 {%0, %1}, [%2];" : "=r"(retval.x), "=r"(retval.y) : "l"(ptr) : "memory");),
|
||||
({
|
||||
asm volatile("ld.cg.v2.u32 {%0, %1}, [%2];" : "=r"(retval.x), "=r"(retval.y) : "l"(ptr) : "memory");
|
||||
__threadfence();
|
||||
}));
|
||||
return retval;
|
||||
}
|
||||
|
||||
static _CCCL_DEVICE _CCCL_FORCEINLINE unsigned int load_acquire(unsigned int const* ptr)
|
||||
{
|
||||
unsigned int retval;
|
||||
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70,
|
||||
(asm volatile("ld.acquire.gpu.u32 %0, [%1];" : "=r"(retval) : "l"(ptr) : "memory");),
|
||||
(asm volatile("ld.cg.u32 %0, [%1];" : "=r"(retval) : "l"(ptr) : "memory"); __threadfence();));
|
||||
|
||||
return retval;
|
||||
}
|
||||
} // namespace detail
|
||||
|
||||
#endif // _CCCL_DOXYGEN_INVOKED
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,223 +0,0 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved.
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.
|
||||
// SPDX-License-Identifier: BSD-3
|
||||
|
||||
/**
|
||||
* @file Utilities for strong memory operations.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cub/util_ptx.cuh>
|
||||
#include <cub/util_type.cuh>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
#ifndef _CCCL_DOXYGEN_INVOKED // Do not document
|
||||
|
||||
namespace detail
|
||||
{
|
||||
static _CCCL_DEVICE _CCCL_FORCEINLINE void store_relaxed(uint4* ptr, uint4 val)
|
||||
{
|
||||
NV_IF_ELSE_TARGET(
|
||||
NV_PROVIDES_SM_70,
|
||||
(asm volatile("st.relaxed.gpu.v4.u32 [%0], {%1, %2, %3, %4};" : : "l"(ptr),
|
||||
"r"(val.x),
|
||||
"r"(val.y),
|
||||
"r"(val.z),
|
||||
"r"(val.w) : "memory");),
|
||||
(asm volatile(
|
||||
"st.cg.v4.u32 [%0], {%1, %2, %3, %4};" : : "l"(ptr), "r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w) : "memory");));
|
||||
}
|
||||
|
||||
static _CCCL_DEVICE _CCCL_FORCEINLINE void store_relaxed(ulonglong2* ptr, ulonglong2 val)
|
||||
{
|
||||
NV_IF_ELSE_TARGET(
|
||||
NV_PROVIDES_SM_70,
|
||||
(asm volatile("st.relaxed.gpu.v2.u64 [%0], {%1, %2};" : : "l"(ptr), "l"(val.x), "l"(val.y) : "memory");),
|
||||
(asm volatile("st.cg.v2.u64 [%0], {%1, %2};" : : "l"(ptr), "l"(val.x), "l"(val.y) : "memory");));
|
||||
}
|
||||
|
||||
static _CCCL_DEVICE _CCCL_FORCEINLINE void store_relaxed(ushort4* ptr, ushort4 val)
|
||||
{
|
||||
NV_IF_ELSE_TARGET(
|
||||
NV_PROVIDES_SM_70,
|
||||
(asm volatile("st.relaxed.gpu.v4.u16 [%0], {%1, %2, %3, %4};" : : "l"(ptr),
|
||||
"h"(val.x),
|
||||
"h"(val.y),
|
||||
"h"(val.z),
|
||||
"h"(val.w) : "memory");),
|
||||
(asm volatile(
|
||||
"st.cg.v4.u16 [%0], {%1, %2, %3, %4};" : : "l"(ptr), "h"(val.x), "h"(val.y), "h"(val.z), "h"(val.w) : "memory");));
|
||||
}
|
||||
|
||||
static _CCCL_DEVICE _CCCL_FORCEINLINE void store_relaxed(uint2* ptr, uint2 val)
|
||||
{
|
||||
NV_IF_ELSE_TARGET(
|
||||
NV_PROVIDES_SM_70,
|
||||
(asm volatile("st.relaxed.gpu.v2.u32 [%0], {%1, %2};" : : "l"(ptr), "r"(val.x), "r"(val.y) : "memory");),
|
||||
(asm volatile("st.cg.v2.u32 [%0], {%1, %2};" : : "l"(ptr), "r"(val.x), "r"(val.y) : "memory");));
|
||||
}
|
||||
|
||||
static _CCCL_DEVICE _CCCL_FORCEINLINE void store_relaxed(unsigned long long* ptr, unsigned long long val)
|
||||
{
|
||||
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70,
|
||||
(asm volatile("st.relaxed.gpu.u64 [%0], %1;" : : "l"(ptr), "l"(val) : "memory");),
|
||||
(asm volatile("st.cg.u64 [%0], %1;" : : "l"(ptr), "l"(val) : "memory");));
|
||||
}
|
||||
|
||||
static _CCCL_DEVICE _CCCL_FORCEINLINE void store_relaxed(unsigned int* ptr, unsigned int val)
|
||||
{
|
||||
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70,
|
||||
(asm volatile("st.relaxed.gpu.u32 [%0], %1;" : : "l"(ptr), "r"(val) : "memory");),
|
||||
(asm volatile("st.cg.u32 [%0], %1;" : : "l"(ptr), "r"(val) : "memory");));
|
||||
}
|
||||
|
||||
static _CCCL_DEVICE _CCCL_FORCEINLINE void store_relaxed(unsigned short* ptr, unsigned short val)
|
||||
{
|
||||
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70,
|
||||
(asm volatile("st.relaxed.gpu.u16 [%0], %1;" : : "l"(ptr), "h"(val) : "memory");),
|
||||
(asm volatile("st.cg.u16 [%0], %1;" : : "l"(ptr), "h"(val) : "memory");));
|
||||
}
|
||||
|
||||
static _CCCL_DEVICE _CCCL_FORCEINLINE void store_relaxed(unsigned char* ptr, unsigned char val)
|
||||
{
|
||||
NV_IF_ELSE_TARGET(
|
||||
NV_PROVIDES_SM_70,
|
||||
(asm volatile("{"
|
||||
" .reg .u8 datum;"
|
||||
" cvt.u8.u16 datum, %1;"
|
||||
" st.relaxed.gpu.u8 [%0], datum;"
|
||||
"}" : : "l"(ptr),
|
||||
"h"((unsigned short) val) : "memory");),
|
||||
(asm volatile("{"
|
||||
" .reg .u8 datum;"
|
||||
" cvt.u8.u16 datum, %1;"
|
||||
" st.cg.u8 [%0], datum;"
|
||||
"}" : : "l"(ptr),
|
||||
"h"((unsigned short) val) : "memory");));
|
||||
}
|
||||
|
||||
_CCCL_DEVICE _CCCL_FORCEINLINE void store_release(uint4* ptr, uint4 val)
|
||||
{
|
||||
NV_IF_ELSE_TARGET(
|
||||
NV_PROVIDES_SM_70,
|
||||
(asm volatile("st.release.gpu.v4.u32 [%0], {%1, %2, %3, %4};" : : "l"(ptr),
|
||||
"r"(val.x),
|
||||
"r"(val.y),
|
||||
"r"(val.z),
|
||||
"r"(val.w) : "memory");),
|
||||
({
|
||||
__threadfence();
|
||||
asm volatile("st.cg.v4.u32 [%0], {%1, %2, %3, %4};"
|
||||
:
|
||||
: "l"(ptr), "r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w)
|
||||
: "memory");
|
||||
}));
|
||||
}
|
||||
|
||||
_CCCL_DEVICE _CCCL_FORCEINLINE void store_release(ulonglong2* ptr, ulonglong2 val)
|
||||
{
|
||||
NV_IF_ELSE_TARGET(
|
||||
NV_PROVIDES_SM_70,
|
||||
(asm volatile("st.release.gpu.v2.u64 [%0], {%1, %2};" : : "l"(ptr), "l"(val.x), "l"(val.y) : "memory");),
|
||||
({
|
||||
__threadfence();
|
||||
asm volatile("st.cg.v2.u64 [%0], {%1, %2};" : : "l"(ptr), "l"(val.x), "l"(val.y) : "memory");
|
||||
}));
|
||||
}
|
||||
|
||||
_CCCL_DEVICE _CCCL_FORCEINLINE void store_release(ushort4* ptr, ushort4 val)
|
||||
{
|
||||
NV_IF_ELSE_TARGET(
|
||||
NV_PROVIDES_SM_70,
|
||||
(asm volatile("st.release.gpu.v4.u16 [%0], {%1, %2, %3, %4};" : : "l"(ptr),
|
||||
"h"(val.x),
|
||||
"h"(val.y),
|
||||
"h"(val.z),
|
||||
"h"(val.w) : "memory");),
|
||||
({
|
||||
__threadfence();
|
||||
asm volatile("st.cg.v4.u16 [%0], {%1, %2, %3, %4};"
|
||||
:
|
||||
: "l"(ptr), "h"(val.x), "h"(val.y), "h"(val.z), "h"(val.w)
|
||||
: "memory");
|
||||
}));
|
||||
}
|
||||
|
||||
_CCCL_DEVICE _CCCL_FORCEINLINE void store_release(uint2* ptr, uint2 val)
|
||||
{
|
||||
NV_IF_ELSE_TARGET(
|
||||
NV_PROVIDES_SM_70,
|
||||
(asm volatile("st.release.gpu.v2.u32 [%0], {%1, %2};" : : "l"(ptr), "r"(val.x), "r"(val.y) : "memory");),
|
||||
({
|
||||
__threadfence();
|
||||
asm volatile("st.cg.v2.u32 [%0], {%1, %2};" : : "l"(ptr), "r"(val.x), "r"(val.y) : "memory");
|
||||
}));
|
||||
}
|
||||
|
||||
_CCCL_DEVICE _CCCL_FORCEINLINE void store_release(unsigned long long* ptr, unsigned long long val)
|
||||
{
|
||||
NV_IF_ELSE_TARGET(
|
||||
NV_PROVIDES_SM_70, (asm volatile("st.release.gpu.u64 [%0], %1;" : : "l"(ptr), "l"(val) : "memory");), ({
|
||||
__threadfence();
|
||||
asm volatile("st.cg.u64 [%0], %1;" : : "l"(ptr), "l"(val) : "memory");
|
||||
}));
|
||||
}
|
||||
|
||||
_CCCL_DEVICE _CCCL_FORCEINLINE void store_release(unsigned int* ptr, unsigned int val)
|
||||
{
|
||||
NV_IF_ELSE_TARGET(
|
||||
NV_PROVIDES_SM_70, (asm volatile("st.release.gpu.u32 [%0], %1;" : : "l"(ptr), "r"(val) : "memory");), ({
|
||||
__threadfence();
|
||||
asm volatile("st.cg.u32 [%0], %1;" : : "l"(ptr), "r"(val) : "memory");
|
||||
}));
|
||||
}
|
||||
|
||||
_CCCL_DEVICE _CCCL_FORCEINLINE void store_release(unsigned short* ptr, unsigned short val)
|
||||
{
|
||||
NV_IF_ELSE_TARGET(
|
||||
NV_PROVIDES_SM_70, (asm volatile("st.release.gpu.u16 [%0], %1;" : : "l"(ptr), "h"(val) : "memory");), ({
|
||||
__threadfence();
|
||||
asm volatile("st.cg.u16 [%0], %1;" : : "l"(ptr), "h"(val) : "memory");
|
||||
}));
|
||||
}
|
||||
|
||||
_CCCL_DEVICE _CCCL_FORCEINLINE void store_release(unsigned char* ptr, unsigned char val)
|
||||
{
|
||||
NV_IF_ELSE_TARGET(
|
||||
NV_PROVIDES_SM_70,
|
||||
(asm volatile("{"
|
||||
" .reg .u8 datum;"
|
||||
" cvt.u8.u16 datum, %1;"
|
||||
" st.release.gpu.u8 [%0], datum;"
|
||||
"}" : : "l"(ptr),
|
||||
"h"((unsigned short) val) : "memory");),
|
||||
({
|
||||
__threadfence();
|
||||
asm volatile(
|
||||
"{"
|
||||
" .reg .u8 datum;"
|
||||
" cvt.u8.u16 datum, %1;"
|
||||
" st.cg.u8 [%0], datum;"
|
||||
"}"
|
||||
:
|
||||
: "l"(ptr), "h"((unsigned short) val)
|
||||
: "memory");
|
||||
}));
|
||||
}
|
||||
} // namespace detail
|
||||
|
||||
#endif // _CCCL_DOXYGEN_INVOKED
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,380 +0,0 @@
|
||||
/*
|
||||
* Copyright 2021 NVIDIA Corporation
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cub/util_namespace.cuh>
|
||||
#include <cub/util_temporary_storage.cuh>
|
||||
|
||||
#include <cuda/__stream/stream_ref.h>
|
||||
#include <cuda/std/__algorithm/max.h>
|
||||
#include <cuda/std/__exception/exception_macros.h>
|
||||
#include <cuda/std/cstddef>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail::temporary_storage
|
||||
{
|
||||
class slot;
|
||||
|
||||
template <typename T>
|
||||
class alias;
|
||||
|
||||
template <int SlotsCount>
|
||||
class layout;
|
||||
|
||||
/**
|
||||
* @brief Temporary storage slot that can be considered a C++ union with an
|
||||
* arbitrary fields count.
|
||||
*
|
||||
* @warning slot lifetime is defined by the lifetime of the associated layout.
|
||||
* It's impossible to request new array if layout is already mapped.
|
||||
*
|
||||
* @par A Simple Example
|
||||
* @code
|
||||
* auto slot = temporary_storage.get_slot(0);
|
||||
*
|
||||
* // Add fields into the slot
|
||||
* // Create an int alias with 0 elements:
|
||||
* auto int_array = slot->create_alias<int>();
|
||||
* // Create a double alias with 2 elements:
|
||||
* auto double_array = slot->create_alias<double>(2);
|
||||
* // Create a char alias with 0 elements:
|
||||
* auto empty_array = slot->create_alias<char>();
|
||||
* // Slot size is defined by double_array size (2 * sizeof(double))
|
||||
*
|
||||
* if (condition)
|
||||
* {
|
||||
* int_array.grow(42);
|
||||
* // Now slot size is defined by int_array size (42 * sizeof(int))
|
||||
* }
|
||||
*
|
||||
* // Temporary storage mapping
|
||||
* // ...
|
||||
|
||||
* int *d_int_array = int_array.get();
|
||||
* double *d_double_array = double_array.get();
|
||||
* char *d_empty_array = empty_array.get(); // Guaranteed to return nullptr
|
||||
* @endcode
|
||||
*/
|
||||
class slot
|
||||
{
|
||||
size_t m_size{};
|
||||
void* m_pointer{};
|
||||
|
||||
public:
|
||||
slot() = default;
|
||||
|
||||
/**
|
||||
* @brief Returns an array of type @p T and length @p elements
|
||||
*/
|
||||
template <typename T>
|
||||
_CCCL_HOST_DEVICE alias<T> create_alias(size_t elements = 0);
|
||||
|
||||
private:
|
||||
_CCCL_HOST_DEVICE void set_bytes_required(size_t new_size)
|
||||
{
|
||||
m_size = (::cuda::std::max) (m_size, new_size);
|
||||
}
|
||||
|
||||
_CCCL_HOST_DEVICE size_t get_bytes_required() const
|
||||
{
|
||||
return m_size;
|
||||
}
|
||||
|
||||
_CCCL_HOST_DEVICE void set_storage(void* ptr)
|
||||
{
|
||||
m_pointer = ptr;
|
||||
}
|
||||
_CCCL_HOST_DEVICE void* get_storage() const
|
||||
{
|
||||
return m_pointer;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
friend class alias;
|
||||
|
||||
template <int>
|
||||
friend class layout;
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Named memory region of a temporary storage slot
|
||||
*
|
||||
* @par Overview
|
||||
* This class provides a typed wrapper of a temporary slot memory region.
|
||||
* It can be considered as a field in the C++ union. It's only possible to
|
||||
* increase the array size.
|
||||
*
|
||||
* @warning alias lifetime is defined by the lifetime of the associated slot
|
||||
* It's impossible to grow the array if the layout is already mapped.
|
||||
*/
|
||||
template <typename T>
|
||||
class alias
|
||||
{
|
||||
slot& m_slot;
|
||||
size_t m_elements{};
|
||||
|
||||
_CCCL_HOST_DEVICE explicit alias(slot& slot, size_t elements = 0)
|
||||
: m_slot(slot)
|
||||
, m_elements(elements)
|
||||
{
|
||||
this->update_slot();
|
||||
}
|
||||
|
||||
_CCCL_HOST_DEVICE void update_slot()
|
||||
{
|
||||
m_slot.set_bytes_required(m_elements * sizeof(T));
|
||||
}
|
||||
|
||||
public:
|
||||
alias() = delete;
|
||||
|
||||
/**
|
||||
* @brief Increases the number of elements
|
||||
*
|
||||
* @warning
|
||||
* This method should be called before temporary storage mapping stage.
|
||||
*
|
||||
* @param[in] new_elements Increases the memory region occupied in the
|
||||
* temporary slot to fit up to @p new_elements items
|
||||
* of type @p T.
|
||||
*/
|
||||
_CCCL_HOST_DEVICE void grow(size_t new_elements)
|
||||
{
|
||||
m_elements = new_elements;
|
||||
this->update_slot();
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Returns pointer to array
|
||||
*
|
||||
* If the @p elements number is equal to zero, or storage layout isn't mapped,
|
||||
* @p nullptr is returned.
|
||||
*/
|
||||
_CCCL_HOST_DEVICE T* get() const
|
||||
{
|
||||
if (m_elements == 0)
|
||||
{
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
return reinterpret_cast<T*>(m_slot.get_storage());
|
||||
}
|
||||
|
||||
friend class slot;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
_CCCL_HOST_DEVICE alias<T> slot::create_alias(size_t elements)
|
||||
{
|
||||
return alias<T>(*this, elements);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Temporary storage layout represents a structure with
|
||||
* @p SlotsCount union-like fields
|
||||
*
|
||||
* The layout can be mapped to a temporary buffer only once.
|
||||
*
|
||||
* @par A Simple Example
|
||||
* @code
|
||||
* cub::detail::temporary_storage::layout<3> temporary_storage;
|
||||
*
|
||||
* auto slot_1 = temporary_storage.get_slot(0);
|
||||
* auto slot_2 = temporary_storage.get_slot(1);
|
||||
*
|
||||
* // Add fields into the first slot
|
||||
* auto int_array = slot_1->create_alias<int>(1);
|
||||
* auto double_array = slot_1->create_alias<double>(2);
|
||||
*
|
||||
* // Add fields into the second slot
|
||||
* auto char_array = slot_2->create_alias<char>();
|
||||
*
|
||||
* // The equivalent C++ structure could look like
|
||||
* // struct StorageLayout
|
||||
* // {
|
||||
* // union {
|
||||
* // } slot_0;
|
||||
* // std::byte padding_0[256 - sizeof (slot_0)];
|
||||
* //
|
||||
* // union {
|
||||
* // int alias_0[1];
|
||||
* // double alias_1[2];
|
||||
* // } slot_1;
|
||||
* // std::byte padding_1[256 - sizeof (slot_1)];
|
||||
* //
|
||||
* // union {
|
||||
* // char alias_0[0];
|
||||
* // } slot_2;
|
||||
* // std::byte padding_2[256 - sizeof (slot_2)];
|
||||
* // };
|
||||
*
|
||||
* // The third slot is empty
|
||||
*
|
||||
* // Temporary storage mapping
|
||||
* if (d_temp_storage == nullptr)
|
||||
* {
|
||||
* temp_storage_bytes = temporary_storage.get_size();
|
||||
* return;
|
||||
* }
|
||||
* else
|
||||
* {
|
||||
* temporary_storage.map_to_buffer(d_temp_storage, temp_storage_bytes);
|
||||
* }
|
||||
*
|
||||
* // Use pointers
|
||||
* int *d_int_array = int_array.get();
|
||||
* double *d_double_array = double_array.get();
|
||||
* char *d_char_array = char_array.get();
|
||||
* @endcode
|
||||
*/
|
||||
template <int SlotsCount>
|
||||
class layout
|
||||
{
|
||||
slot m_slots[SlotsCount];
|
||||
size_t m_sizes[SlotsCount];
|
||||
void* m_pointers[SlotsCount];
|
||||
bool m_layout_was_mapped{};
|
||||
|
||||
public:
|
||||
layout() = default;
|
||||
|
||||
_CCCL_HOST_DEVICE slot* get_slot(int slot_id)
|
||||
{
|
||||
if (slot_id < SlotsCount)
|
||||
{
|
||||
return &m_slots[slot_id];
|
||||
}
|
||||
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Returns required temporary storage size in bytes
|
||||
*/
|
||||
_CCCL_HOST_DEVICE size_t get_size()
|
||||
{
|
||||
this->prepare_interface();
|
||||
|
||||
// alias_temporaries can return error only in mapping stage, so it's safe to ignore it here.
|
||||
size_t temp_storage_bytes{};
|
||||
[[maybe_unused]] const auto error = detail::alias_temporaries(nullptr, temp_storage_bytes, m_pointers, m_sizes);
|
||||
_CCCL_ASSERT(error == cudaSuccess, "");
|
||||
|
||||
if (temp_storage_bytes == 0)
|
||||
{
|
||||
// The current CUB convention implies that there are two stages for each
|
||||
// device-scope function call. The first one returns the required storage
|
||||
// size. The second stage consumes temporary storage to perform some work.
|
||||
// The only way to distinguish between the two stages is by checking the
|
||||
// value of the temporary storage pointer. If zero bytes are requested,
|
||||
// `cudaMalloc` will return `nullptr`. This fact makes it impossible to
|
||||
// distinguish between the two stages, so we request some fixed amount of
|
||||
// bytes (even if we don't need it) to have a non-null temporary storage
|
||||
// pointer.
|
||||
return 1;
|
||||
}
|
||||
|
||||
return temp_storage_bytes;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Maps the layout to the temporary storage buffer.
|
||||
*/
|
||||
_CCCL_HOST_DEVICE cudaError_t map_to_buffer(void* d_temp_storage, size_t temp_storage_bytes)
|
||||
{
|
||||
if (m_layout_was_mapped)
|
||||
{
|
||||
return cudaErrorAlreadyMapped;
|
||||
}
|
||||
|
||||
this->prepare_interface();
|
||||
|
||||
if (cudaError_t error = detail::alias_temporaries(d_temp_storage, temp_storage_bytes, m_pointers, m_sizes))
|
||||
{
|
||||
return error;
|
||||
}
|
||||
|
||||
for (size_t slot_id = 0; slot_id < SlotsCount; slot_id++)
|
||||
{
|
||||
m_slots[slot_id].set_storage(m_pointers[slot_id]);
|
||||
}
|
||||
|
||||
m_layout_was_mapped = true;
|
||||
return cudaSuccess;
|
||||
}
|
||||
|
||||
private:
|
||||
_CCCL_HOST_DEVICE void prepare_interface()
|
||||
{
|
||||
if (m_layout_was_mapped)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
for (size_t slot_id = 0; slot_id < SlotsCount; slot_id++)
|
||||
{
|
||||
const size_t slot_size = m_slots[slot_id].get_bytes_required();
|
||||
|
||||
m_sizes[slot_id] = slot_size;
|
||||
m_pointers[slot_id] = nullptr;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
template <typename MRT>
|
||||
CUB_RUNTIME_FUNCTION cudaError_t
|
||||
allocate(::cuda::stream_ref stream, void*& d_temp_storage, size_t temp_storage_bytes, MRT& mr)
|
||||
{
|
||||
_CCCL_TRY
|
||||
{
|
||||
d_temp_storage = mr.allocate(stream, temp_storage_bytes, alignof(::cuda::std::max_align_t));
|
||||
return cudaSuccess;
|
||||
}
|
||||
_CCCL_CATCH_ALL
|
||||
{
|
||||
return cudaErrorMemoryAllocation;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename MRT>
|
||||
CUB_RUNTIME_FUNCTION cudaError_t
|
||||
deallocate(::cuda::stream_ref stream, void* d_temp_storage, size_t temp_storage_bytes, MRT& mr)
|
||||
{
|
||||
_CCCL_TRY
|
||||
{
|
||||
mr.deallocate(stream, d_temp_storage, temp_storage_bytes, alignof(::cuda::std::max_align_t));
|
||||
return cudaSuccess;
|
||||
}
|
||||
_CCCL_CATCH_ALL
|
||||
{
|
||||
return cudaErrorMemoryAllocation;
|
||||
}
|
||||
}
|
||||
} // namespace detail::temporary_storage
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,185 +0,0 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
// SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
/**
|
||||
* \file
|
||||
* Wrappers and extensions around <type_traits> utilities.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cub/util_cpp_dialect.cuh>
|
||||
#include <cub/util_namespace.cuh>
|
||||
|
||||
#include <cuda/std/__concepts/concept_macros.h> // IWYU pragma: keep
|
||||
#include <cuda/std/__fwd/array.h>
|
||||
#include <cuda/std/__fwd/mdspan.h>
|
||||
#include <cuda/std/__fwd/span.h>
|
||||
#include <cuda/std/__type_traits/always_false.h>
|
||||
#include <cuda/std/__type_traits/conditional.h>
|
||||
#include <cuda/std/__type_traits/enable_if.h>
|
||||
#include <cuda/std/__type_traits/integral_constant.h>
|
||||
#include <cuda/std/__type_traits/is_same.h>
|
||||
#include <cuda/std/__type_traits/is_signed_integer.h>
|
||||
#include <cuda/std/__type_traits/is_unsigned_integer.h>
|
||||
#include <cuda/std/__type_traits/remove_cv.h>
|
||||
#include <cuda/std/__type_traits/void_t.h>
|
||||
#include <cuda/std/__utility/declval.h>
|
||||
#include <cuda/std/cstddef>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
namespace detail
|
||||
{
|
||||
template <typename T, typename... TArgs>
|
||||
inline constexpr bool is_one_of_v = (::cuda::std::is_same_v<T, TArgs> || ...);
|
||||
|
||||
template <typename T, typename V, typename = void>
|
||||
struct has_binary_call_operator : ::cuda::std::false_type
|
||||
{};
|
||||
|
||||
template <typename T, typename V>
|
||||
struct has_binary_call_operator<
|
||||
T,
|
||||
V,
|
||||
::cuda::std::void_t<decltype(::cuda::std::declval<T>()(::cuda::std::declval<V>(), ::cuda::std::declval<V>()))>>
|
||||
: ::cuda::std::true_type
|
||||
{};
|
||||
|
||||
/***********************************************************************************************************************
|
||||
* Array-like type traits
|
||||
**********************************************************************************************************************/
|
||||
|
||||
template <typename T>
|
||||
inline constexpr bool is_fixed_size_random_access_range_v = false;
|
||||
|
||||
template <typename T, size_t N>
|
||||
inline constexpr bool is_fixed_size_random_access_range_v<T[N]> = true;
|
||||
|
||||
template <typename T, size_t N>
|
||||
inline constexpr bool is_fixed_size_random_access_range_v<::cuda::std::array<T, N>> = true;
|
||||
|
||||
template <typename T, size_t N>
|
||||
inline constexpr bool is_fixed_size_random_access_range_v<::cuda::std::span<T, N>> = N != ::cuda::std::dynamic_extent;
|
||||
|
||||
template <typename T, typename E, typename L, typename A>
|
||||
inline constexpr bool is_fixed_size_random_access_range_v<::cuda::std::mdspan<T, E, L, A>> =
|
||||
E::rank() == 1 && E::rank_dynamic() == 0;
|
||||
|
||||
/***********************************************************************************************************************
|
||||
* static_size: a type trait that returns the number of elements in an Array-like type
|
||||
**********************************************************************************************************************/
|
||||
|
||||
template <typename T>
|
||||
inline constexpr int static_size_v = ::cuda::std::enable_if_t<::cuda::std::__always_false_v<T>>{};
|
||||
|
||||
template <typename T, size_t N>
|
||||
inline constexpr int static_size_v<T[N]> = N;
|
||||
|
||||
template <typename T, size_t N>
|
||||
inline constexpr int static_size_v<::cuda::std::array<T, N>> = N;
|
||||
|
||||
template <typename T, size_t N>
|
||||
inline constexpr int static_size_v<::cuda::std::span<T, N>> =
|
||||
::cuda::std::enable_if_t<N != ::cuda::std::dynamic_extent, int>{N};
|
||||
|
||||
template <typename T, typename E, typename L, typename A>
|
||||
inline constexpr int static_size_v<::cuda::std::mdspan<T, E, L, A>> =
|
||||
::cuda::std::enable_if_t<E::rank() == 1 && E::rank_dynamic() == 0, int>{E::static_extent(0)};
|
||||
|
||||
template <typename T>
|
||||
using implicit_prom_t = decltype(+T{});
|
||||
|
||||
/***********************************************************************************************************************
|
||||
* Extended floating point traits
|
||||
**********************************************************************************************************************/
|
||||
// half
|
||||
|
||||
template <typename>
|
||||
inline constexpr bool is_half_impl_v = false;
|
||||
|
||||
template <typename>
|
||||
inline constexpr bool is_half2_impl_v = false;
|
||||
|
||||
#if _CCCL_HAS_NVFP16()
|
||||
|
||||
template <>
|
||||
inline constexpr bool is_half_impl_v<__half> = true;
|
||||
|
||||
template <>
|
||||
inline constexpr bool is_half2_impl_v<__half2> = true;
|
||||
|
||||
#endif // _CCCL_HAS_NVFP16
|
||||
|
||||
template <typename T>
|
||||
inline constexpr bool is_half_v = is_half_impl_v<::cuda::std::remove_cv_t<T>>;
|
||||
|
||||
template <typename T>
|
||||
inline constexpr bool is_half2_v = is_half2_impl_v<::cuda::std::remove_cv_t<T>>;
|
||||
|
||||
template <typename T>
|
||||
inline constexpr bool is_any_half_v = is_half_impl_v<T> || is_half2_impl_v<T>;
|
||||
|
||||
//----------------------------------------------------------------------------------------------------------------------
|
||||
// bfloat16
|
||||
|
||||
template <typename>
|
||||
inline constexpr bool is_bfloat16_impl_v = false;
|
||||
|
||||
template <typename>
|
||||
inline constexpr bool is_bfloat162_impl_v = false;
|
||||
|
||||
#if _CCCL_HAS_NVBF16()
|
||||
|
||||
template <>
|
||||
inline constexpr bool is_bfloat16_impl_v<__nv_bfloat16> = true;
|
||||
|
||||
template <>
|
||||
inline constexpr bool is_bfloat162_impl_v<__nv_bfloat162> = true;
|
||||
|
||||
#endif // _CCCL_HAS_NVBF16
|
||||
|
||||
template <typename T>
|
||||
inline constexpr bool is_bfloat16_v = is_bfloat16_impl_v<::cuda::std::remove_cv_t<T>>;
|
||||
|
||||
template <typename T>
|
||||
inline constexpr bool is_bfloat162_v = is_bfloat162_impl_v<::cuda::std::remove_cv_t<T>>;
|
||||
|
||||
template <typename T>
|
||||
inline constexpr bool is_any_bfloat16_v = is_bfloat16_v<T> || is_bfloat162_v<T>;
|
||||
|
||||
//----------------------------------------------------------------------------------------------------------------------
|
||||
// short2/ushort2
|
||||
|
||||
template <typename T>
|
||||
inline constexpr bool is_any_short2_impl_v = false;
|
||||
|
||||
template <>
|
||||
inline constexpr bool is_any_short2_impl_v<short2> = true;
|
||||
|
||||
template <>
|
||||
inline constexpr bool is_any_short2_impl_v<ushort2> = true;
|
||||
|
||||
template <typename T>
|
||||
inline constexpr bool is_any_short2_v = is_any_short2_impl_v<::cuda::std::remove_cv_t<T>>;
|
||||
|
||||
//----------------------------------------------------------------------------------------------------------------------
|
||||
|
||||
// - promote small integer types to their corresponding 32-bit promotion type
|
||||
// - address the incompatibility between linux/windows for int/long
|
||||
template <typename T>
|
||||
using signed_promotion_t = ::cuda::std::conditional_t<
|
||||
::cuda::std::__cccl_is_signed_integer_v<T> && sizeof(T) <= sizeof(int),
|
||||
int,
|
||||
::cuda::std::conditional_t<::cuda::std::__cccl_is_unsigned_integer_v<T> && sizeof(T) <= sizeof(uint32_t), uint32_t, T>>;
|
||||
} // namespace detail
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,48 +0,0 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.
|
||||
// SPDX-License-Identifier: BSD-3
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cuda/__type_traits/is_trivially_copyable.h>
|
||||
#include <cuda/std/__new/device_new.h>
|
||||
#include <cuda/std/__type_traits/enable_if.h>
|
||||
#include <cuda/std/__utility/forward.h>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail
|
||||
{
|
||||
#if _CCCL_CUDA_COMPILER(NVHPC)
|
||||
template <typename T, typename U>
|
||||
_CCCL_HOST_DEVICE void uninitialized_copy_single(T* ptr, U&& val)
|
||||
{
|
||||
// NVBug 3384810
|
||||
new (ptr) T(::cuda::std::forward<U>(val));
|
||||
}
|
||||
#else // ^^^ _CCCL_CUDA_COMPILER(NVHPC) ^^^ / vvv !_CCCL_CUDA_COMPILER(NVHPC) vvv
|
||||
template <typename T, typename U, ::cuda::std::enable_if_t<::cuda::is_trivially_copyable_v<T>, int> = 0>
|
||||
_CCCL_HOST_DEVICE void uninitialized_copy_single(T* ptr, U&& val)
|
||||
{
|
||||
// gevtushenko: placement new should work here as well, but the code generated for copy assignment is sometimes better
|
||||
*ptr = ::cuda::std::forward<U>(val);
|
||||
}
|
||||
|
||||
template <typename T, typename U, ::cuda::std::enable_if_t<!::cuda::is_trivially_copyable_v<T>, int> = 0>
|
||||
_CCCL_HOST_DEVICE void uninitialized_copy_single(T* ptr, U&& val)
|
||||
{
|
||||
new (ptr) T(::cuda::std::forward<U>(val));
|
||||
}
|
||||
#endif // !_CCCL_CUDA_COMPILER(NVHPC)
|
||||
} // namespace detail
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,34 +0,0 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2011-2025, NVIDIA CORPORATION. All rights reserved.
|
||||
// SPDX-License-Identifier: BSD-3
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
namespace detail
|
||||
{
|
||||
#ifndef _CCCL_DOXYGEN_INVOKED // Do not document
|
||||
|
||||
// NOTE: bit_cast cannot be always used because __half, __nv_bfloat16, etc. are not trivially copyable
|
||||
template <typename Output, typename Input>
|
||||
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE Output unsafe_bitcast(const Input& input)
|
||||
{
|
||||
Output output;
|
||||
static_assert(sizeof(input) == sizeof(output), "wrong size");
|
||||
// NOLINTNEXTLINE(bugprone-undefined-memory-manipulation)
|
||||
::memcpy(&output, &input, sizeof(input));
|
||||
return output;
|
||||
}
|
||||
|
||||
#endif // !_CCCL_DOXYGEN_INVOKED
|
||||
} // namespace detail
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,83 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cub/detail/warpspeed/optimize_smem_ptr.cuh>
|
||||
|
||||
#include <cuda/std/__type_traits/is_constant_evaluated.h>
|
||||
#include <cuda/std/cstdint>
|
||||
|
||||
#include <nv/target>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail::warpspeed
|
||||
{
|
||||
struct SmemAllocator
|
||||
{
|
||||
::cuda::std::uint32_t mPtrSmem32 = 0;
|
||||
int mAllocatedSize = 0;
|
||||
|
||||
_CCCL_HOST_DEVICE_API constexpr SmemAllocator() noexcept
|
||||
{
|
||||
// we only need the real pointer at runtime in device code
|
||||
if (!::cuda::std::is_constant_evaluated())
|
||||
{
|
||||
NV_IF_TARGET(NV_IS_DEVICE, mPtrSmem32 = dynamic_smem_base();)
|
||||
}
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_DEVICE_API static ::cuda::std::uint32_t dynamic_smem_base() noexcept
|
||||
{
|
||||
extern __shared__ char warpSpeedDynamicSmemBase[];
|
||||
return __cvta_generic_to_shared(warpSpeedDynamicSmemBase);
|
||||
}
|
||||
|
||||
// SmemAllocator is a non-copyable, non-movable type. It must be passed by
|
||||
// (mutable) reference to be useful.
|
||||
SmemAllocator(const SmemAllocator&) = delete; // Delete copy constructor
|
||||
SmemAllocator(SmemAllocator&&) = delete; // Delete move constructor
|
||||
SmemAllocator& operator=(const SmemAllocator&) = delete; // Delete copy assignment
|
||||
SmemAllocator& operator=(SmemAllocator&&) = delete; // Delete move assignment
|
||||
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr void* alloc(::cuda::std::uint32_t size, ::cuda::std::uint32_t align = 0)
|
||||
{
|
||||
// Align mPtrSmem32 to requested alignment (round-up)
|
||||
::cuda::std::uint32_t ptrAllocation32 = (mPtrSmem32 + (align - 1)) & ~(align - 1);
|
||||
|
||||
// Move base pointer and update allocated size
|
||||
mAllocatedSize += static_cast<int>(size + ptrAllocation32 - mPtrSmem32);
|
||||
mPtrSmem32 = ptrAllocation32 + size;
|
||||
|
||||
// we only need the pointer at runtime in device code
|
||||
if (!::cuda::std::is_constant_evaluated())
|
||||
{
|
||||
NV_IF_TARGET(
|
||||
NV_IS_DEVICE,
|
||||
(
|
||||
// Convert allocated smem address to generic pointer
|
||||
void* mPtrAllocation = __cvta_shared_to_generic(ptrAllocation32);
|
||||
// Ensure alignment calculation does not move down into rest of kernel code.
|
||||
return optimizeSmemPtr(mPtrAllocation);))
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr ::cuda::std::uint32_t sizeBytes() const
|
||||
{
|
||||
return mAllocatedSize;
|
||||
}
|
||||
};
|
||||
} // namespace detail::warpspeed
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,52 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cuda/std/cassert>
|
||||
|
||||
/*
|
||||
* _WS_CONSTANT_ASSERT: an assertion that is intended to be verified at compile time.
|
||||
*
|
||||
* A _WS_CONSTANT_ASSERT asserts something that the compiler (optimizer) can verify
|
||||
* at compile time. Therefore, it does not result in an actual call to assert in
|
||||
* the compiled binary. This allows checking various properties that cannot be
|
||||
* verified using static_assert.
|
||||
*
|
||||
* To ensure that all _WS_CONSTANT_ASSERTs are in fact eliminated, compile with
|
||||
* -D_WARPSPEED_FORCE_ASSERT_AT_COMPILE_TIME. With this macro defined, any
|
||||
* _WS_CONSTANT_ASSERT failure will output illegal PTX containing the error message.
|
||||
* As a result, compilation will fail.
|
||||
*
|
||||
* Compiling with -D_WARPSPEED_FORCE_ASSERT_AT_COMPILE_TIME has the additional
|
||||
* advantage that violating any of the assertions can be detected at compile
|
||||
* time and before even running the code.
|
||||
*
|
||||
*/
|
||||
|
||||
#if defined(_WARPSPEED_FORCE_ASSERT_AT_COMPILE_TIME) && defined(__CUDA_ARCH__)
|
||||
// When _WARPSPEED_FORCE_ASSERT_AT_COMPILE_TIME is defined and compiling for device, output illegal PTX.
|
||||
// This causes the compilation to fail.
|
||||
# define _WS_CONSTANT_ASSERT(expr, msg) \
|
||||
do \
|
||||
{ \
|
||||
if (!(expr)) \
|
||||
{ \
|
||||
asm volatile(".pragma \"\n" __FILE__ "(" _CCCL_TO_STRING( \
|
||||
__LINE__) "): %0" \
|
||||
": error: constant assertion failed with '" msg "'\n\";" ::"C"(__func__)); \
|
||||
} \
|
||||
} while (0)
|
||||
#else
|
||||
// Host or !_WARPSPEED_FORCE_ASSERT_AT_COMPILE_TIME
|
||||
# define _WS_CONSTANT_ASSERT(expr, msg) _CCCL_ASSERT((expr), msg)
|
||||
#endif
|
||||
@@ -1,338 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cub/detail/strong_load.cuh>
|
||||
#include <cub/detail/strong_store.cuh>
|
||||
#include <cub/detail/warpspeed/special_registers.cuh>
|
||||
#include <cub/thread/thread_store.cuh>
|
||||
#include <cub/warp/specializations/warp_redux.cuh>
|
||||
#include <cub/warp/warp_reduce.cuh>
|
||||
|
||||
#include <cuda/__cmath/pow2.h>
|
||||
#include <cuda/__functional/operator_properties.h>
|
||||
#include <cuda/__memory/is_aligned.h>
|
||||
#include <cuda/__ptx/instructions/get_sreg.h>
|
||||
#include <cuda/__type_traits/is_trivially_copyable.h>
|
||||
#include <cuda/std/__algorithm/min.h>
|
||||
#include <cuda/std/__bit/popcount.h>
|
||||
#include <cuda/std/__type_traits/is_same.h>
|
||||
#include <cuda/std/__type_traits/underlying_type.h>
|
||||
|
||||
#if !_CCCL_HAS_NV_ATOMIC_BUILTINS()
|
||||
# include <cuda/atomic>
|
||||
#endif // !_CCCL_HAS_NV_ATOMIC_BUILTINS()
|
||||
|
||||
#include <nv/target>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail::warpspeed
|
||||
{
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE_API _CCCL_CONSTEVAL ::cuda::std::size_t max_native_atomic_size() noexcept
|
||||
{
|
||||
#if _CCCL_CUDA_COMPILER(NVHPC)
|
||||
return 8;
|
||||
#else // ^^^ _CCCL_CUDA_COMPILER(NVHPC) ^^^ / vvv !_CCCL_CUDA_COMPILER(NVHPC) vvv
|
||||
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_90, (return 16;), (return 8;))
|
||||
#endif // !_CCCL_CUDA_COMPILER(NVHPC)
|
||||
}
|
||||
|
||||
enum scan_state : ::cuda::std::uint32_t
|
||||
{
|
||||
empty = 0,
|
||||
tile_aggregate = 1,
|
||||
};
|
||||
|
||||
template <typename AccumT>
|
||||
struct tile_state_unaligned_t
|
||||
{
|
||||
scan_state state;
|
||||
AccumT value;
|
||||
};
|
||||
|
||||
// some older nvcc versions do not evaluate next_power_of_two() at compile time when called inside an attribute, so we
|
||||
// have to force constant evaluation by assigning the result to a template parameter
|
||||
template <typename AccumT,
|
||||
::cuda::std::size_t _Alignment = ::cuda::next_power_of_two(sizeof(tile_state_unaligned_t<AccumT>))>
|
||||
struct alignas(_Alignment) tile_state_t : tile_state_unaligned_t<AccumT>
|
||||
{};
|
||||
|
||||
#if __cccl_ptx_isa >= 860
|
||||
|
||||
template <typename AccumT>
|
||||
_CCCL_DEVICE_API void
|
||||
storeTileAggregate(tile_state_t<AccumT>* ptrTileStates, scan_state scanState, AccumT aggr, int index, int num_tiles)
|
||||
{
|
||||
_CCCL_ASSERT(::cuda::is_aligned(ptrTileStates, alignof(tile_state_t<AccumT>)), "");
|
||||
_CCCL_ASSERT(index >= 0 && index < num_tiles, "Reading out of bounds tile state");
|
||||
|
||||
if constexpr (sizeof(tile_state_t<AccumT>) <= cub::detail::warpspeed::max_native_atomic_size()
|
||||
&& ::cuda::is_trivially_copyable_v<tile_state_t<AccumT>>)
|
||||
{
|
||||
static_assert(::cuda::is_power_of_two(sizeof(tile_state_t<AccumT>)));
|
||||
tile_state_t<AccumT> tmp{scanState, aggr};
|
||||
|
||||
# if _CCCL_HAS_NV_ATOMIC_BUILTINS()
|
||||
__nv_atomic_store(ptrTileStates + index, &tmp, __NV_ATOMIC_RELAXED, __NV_THREAD_SCOPE_DEVICE);
|
||||
# else // ^^^ _CCCL_HAS_NV_ATOMIC_BUILTINS() ^^^ / vvv !_CCCL_HAS_NV_ATOMIC_BUILTINS() vvv
|
||||
::cuda::atomic_ref<tile_state_t<AccumT>, ::cuda::std::thread_scope_device>{ptrTileStates[index]}.store(
|
||||
tmp, ::cuda::std::memory_order_relaxed);
|
||||
# endif // !_CCCL_HAS_NV_ATOMIC_BUILTINS()
|
||||
}
|
||||
else
|
||||
{
|
||||
ThreadStore<STORE_CG>(&ptrTileStates[index].value, aggr);
|
||||
using state_int = ::cuda::std::underlying_type_t<scan_state>;
|
||||
store_release(reinterpret_cast<state_int*>(&ptrTileStates[index].state), scanState);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename AccumT>
|
||||
_CCCL_DEVICE_API tile_state_t<AccumT> loadTileAggregate(tile_state_t<AccumT>* ptrTileStates, int index, int num_tiles)
|
||||
{
|
||||
_CCCL_ASSERT(::cuda::is_aligned(ptrTileStates, alignof(tile_state_t<AccumT>)), "");
|
||||
_CCCL_ASSERT(index >= 0 && index < num_tiles, "Reading out of bounds tile state");
|
||||
|
||||
tile_state_t<AccumT> res;
|
||||
if constexpr (sizeof(tile_state_t<AccumT>) <= cub::detail::warpspeed::max_native_atomic_size()
|
||||
&& ::cuda::is_trivially_copyable_v<tile_state_t<AccumT>>)
|
||||
{
|
||||
static_assert(::cuda::is_power_of_two(sizeof(tile_state_t<AccumT>)));
|
||||
# if _CCCL_HAS_NV_ATOMIC_BUILTINS()
|
||||
__nv_atomic_load(ptrTileStates + index, &res, __NV_ATOMIC_RELAXED, __NV_THREAD_SCOPE_DEVICE);
|
||||
# else // ^^^ _CCCL_HAS_NV_ATOMIC_BUILTINS() ^^^ / vvv !_CCCL_HAS_NV_ATOMIC_BUILTINS() vvv
|
||||
res = ::cuda::atomic_ref<tile_state_t<AccumT>, ::cuda::std::thread_scope_device>{ptrTileStates[index]}.load(
|
||||
::cuda::std::memory_order_relaxed);
|
||||
# endif // !_CCCL_HAS_NV_ATOMIC_BUILTINS()
|
||||
}
|
||||
else
|
||||
{
|
||||
using state_int = ::cuda::std::underlying_type_t<scan_state>;
|
||||
res.state = static_cast<scan_state>(load_acquire(reinterpret_cast<const state_int*>(&ptrTileStates[index].state)));
|
||||
res.value = ThreadLoad<LOAD_CG>(&ptrTileStates[index].value);
|
||||
}
|
||||
return res;
|
||||
}
|
||||
|
||||
// warpLoadLookahead loads tmp states:
|
||||
// idxTileCur + [0; 32 * numTileStatesPerThread[
|
||||
//
|
||||
// The states are loaded in laneId order and warp-strided:
|
||||
//
|
||||
// outTmpStates[0] contains:
|
||||
// Lane 0: idxTileCur + 0
|
||||
// Lane 1: idxTileCur + 1
|
||||
// ...
|
||||
// Lane 31: idxTileCur + 31
|
||||
//
|
||||
// outTmpStates[1] contains:
|
||||
// Lane 0: idxTileCur + 32
|
||||
// ...
|
||||
// Lane 31 idxTileCur + 63
|
||||
//
|
||||
// If the index idxTileCur + ii of the loaded state is equal to or exceeds idxTileNext, i.e., idxTileCur + ii >=
|
||||
// idxTileNext, then the state is not loaded from memory and set to empty.
|
||||
template <int numTileStatesPerThread, typename AccumT>
|
||||
_CCCL_DEVICE_API void warpLoadLookahead(
|
||||
int laneIdx,
|
||||
tile_state_t<AccumT> (&outTileStates)[numTileStatesPerThread],
|
||||
tile_state_t<AccumT>* ptrTileStates,
|
||||
int idxTileCur,
|
||||
int idxTileNext,
|
||||
int num_tiles)
|
||||
{
|
||||
for (int i = 0; i < numTileStatesPerThread; ++i)
|
||||
{
|
||||
const int idxTileLookahead = idxTileCur + 32 * i + laneIdx;
|
||||
if (idxTileLookahead < idxTileNext)
|
||||
{
|
||||
outTileStates[i] = loadTileAggregate(ptrTileStates, idxTileLookahead, num_tiles);
|
||||
}
|
||||
else
|
||||
{
|
||||
// If we are looking ahead of idxTileNext, then set state to empty
|
||||
outTileStates[i].state = scan_state::empty;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// warpIncrementalLookahead takes the latest known aggrExclusiveCtaPrev and its tile index, idxTilePrev (which's
|
||||
// aggregate is NOT included in aggrExclusiveCtaPrev), and computes the aggrExclusiveCta for the next tile of interest,
|
||||
// idxTileNext (where the returned value will NOT include the aggregate of idxTileNext).
|
||||
//
|
||||
// It does so by loading states in chunks of 32 * numTileStatesPerThread elements, starting from idxTilePrev + 1. From
|
||||
// the chunk of states, it tries to advance its knowledge of aggrExclusiveCta as much as possible. It loops until it can
|
||||
// calculate the value of aggrExclusiveCta from the preceding states.
|
||||
//
|
||||
// The function must be called from a single warp. All passed arguments must be warp-uniform.
|
||||
template <int numTileStatesPerThread, typename AccumT, typename ScanOpT>
|
||||
[[nodiscard]] _CCCL_DEVICE_API _CCCL_FORCEINLINE AccumT warpIncrementalLookahead(
|
||||
SpecialRegisters specialRegisters,
|
||||
tile_state_t<AccumT>* ptrTileStates,
|
||||
const int idxTilePrev,
|
||||
const AccumT aggrExclusiveCtaPrev,
|
||||
const int idxTileNext,
|
||||
ScanOpT& scan_op,
|
||||
const int num_tiles)
|
||||
{
|
||||
const int laneIdx = static_cast<int>(specialRegisters.laneIdx);
|
||||
[[maybe_unused]] const ::cuda::std::uint32_t lanemaskEq = ::cuda::ptx::get_sreg_lanemask_eq();
|
||||
|
||||
int idxTileCur = idxTilePrev;
|
||||
AccumT aggrExclusiveCtaCur = aggrExclusiveCtaPrev;
|
||||
|
||||
using warp_reduce_t = WarpReduce<AccumT>;
|
||||
static_assert(::cuda::std::is_same_v<typename warp_reduce_t::TempStorage, Uninitialized<NullType>>,
|
||||
"WarpReduce for a full warp must not require temporary storage");
|
||||
[[maybe_unused]] typename warp_reduce_t::TempStorage temp_storage;
|
||||
|
||||
while (idxTileCur < idxTileNext)
|
||||
{
|
||||
tile_state_t<AccumT> regTmpStates[numTileStatesPerThread];
|
||||
warpLoadLookahead(laneIdx, regTmpStates, ptrTileStates, idxTileCur, idxTileNext, num_tiles);
|
||||
|
||||
for (int idx = 0; idx < numTileStatesPerThread; ++idx)
|
||||
{
|
||||
// Bitmask with 1 bits indicating which lane has a tile aggregate
|
||||
const ::cuda::std::uint32_t warp_has_aggregate_mask =
|
||||
__ballot_sync(0xffffffffu, regTmpStates[idx].state == scan_state::tile_aggregate);
|
||||
|
||||
// Bitmask with 1 bits for all rightmost lanes having a tile aggregate
|
||||
const ::cuda::std::uint32_t warp_right_aggregates_mask = warp_has_aggregate_mask & (~warp_has_aggregate_mask - 1);
|
||||
|
||||
// Cannot reduce if no rightmost tile aggregates
|
||||
if (warp_right_aggregates_mask == 0)
|
||||
{
|
||||
break;
|
||||
}
|
||||
|
||||
const ::cuda::std::uint32_t warp_right_aggregates_count = ::cuda::std::popcount(warp_right_aggregates_mask);
|
||||
|
||||
// Accumulate the rightmost tile aggregates
|
||||
AccumT local_aggr;
|
||||
NV_IF_ELSE_TARGET(
|
||||
NV_PROVIDES_SM_80,
|
||||
({ // NOTE: Inlined from warp_reduce_shfl
|
||||
if constexpr (is_warp_redux_op_supported_sm80<ScanOpT, AccumT>)
|
||||
{
|
||||
const bool use_value = lanemaskEq & warp_right_aggregates_mask;
|
||||
const AccumT value = use_value ? regTmpStates[idx].value : cuda::identity_element<ScanOpT, AccumT>();
|
||||
local_aggr = cub::detail::warp_redux_sm80(value, ~0, scan_op);
|
||||
}
|
||||
else
|
||||
{
|
||||
// TODO(bgruber): this generates a LOT of SASS. I think it can do better.
|
||||
local_aggr =
|
||||
warp_reduce_t{temp_storage}.Reduce(regTmpStates[idx].value, scan_op, warp_right_aggregates_count);
|
||||
}
|
||||
}),
|
||||
(local_aggr =
|
||||
warp_reduce_t{temp_storage}.Reduce(regTmpStates[idx].value, scan_op, warp_right_aggregates_count);))
|
||||
|
||||
// We never initialized aggrExclusiveCtaCur when starting look ahead at tile 0
|
||||
aggrExclusiveCtaCur = idxTileCur == 0 ? local_aggr : scan_op(aggrExclusiveCtaCur, local_aggr);
|
||||
idxTileCur += warp_right_aggregates_count;
|
||||
|
||||
// we can only continue on the next 32 tile states, if we consumed all 32 of this iteration
|
||||
if (warp_right_aggregates_count < 32)
|
||||
{
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return aggrExclusiveCtaCur; // must only be valid in lane_0
|
||||
}
|
||||
|
||||
// Deterministic version of warpIncrementalLookahead that returns the same aggrExclusiveCta. The difference is that it
|
||||
// always starts the lookahead from a tile index that is a multiple of 32. The left pointer (idxTilePrev) is itself
|
||||
// always a multiple of 32, as it starts at 0 and is only ever advanced by whole batches of 32, so the lookahead resumes
|
||||
// from there directly. Because every reduction begins at the same fixed tiles, no matter which tiles happened to finish
|
||||
// first, the order in which values are summed is always the same and the result is identical on every run.
|
||||
// idxTilePrev/aggrExclusiveCtaPrev are updated by reference to the last multiple of 32.
|
||||
template <int numTileStatesPerThread, typename AccumT, typename ScanOpT>
|
||||
[[nodiscard]] _CCCL_DEVICE_API _CCCL_FORCEINLINE AccumT warpIncrementalLookaheadStable(
|
||||
SpecialRegisters specialRegisters,
|
||||
tile_state_t<AccumT>* ptrTileStates,
|
||||
int& idxTilePrev,
|
||||
AccumT& aggrExclusiveCtaPrev,
|
||||
const int idxTileNext,
|
||||
ScanOpT& scan_op,
|
||||
const int num_tiles)
|
||||
{
|
||||
const int laneIdx = static_cast<int>(specialRegisters.laneIdx);
|
||||
const ::cuda::std::uint32_t lanemaskEq = ::cuda::ptx::get_sreg_lanemask_eq();
|
||||
|
||||
int idxTileCur = idxTilePrev;
|
||||
AccumT aggrExclusiveCtaCur = aggrExclusiveCtaPrev;
|
||||
|
||||
using warp_reduce_t = WarpReduce<AccumT>;
|
||||
static_assert(::cuda::std::is_same_v<typename warp_reduce_t::TempStorage, Uninitialized<NullType>>,
|
||||
"WarpReduce for a full warp must not require temporary storage");
|
||||
[[maybe_unused]] typename warp_reduce_t::TempStorage temp_storage;
|
||||
|
||||
while (idxTileCur < idxTileNext)
|
||||
{
|
||||
tile_state_t<AccumT> regTmpStates[numTileStatesPerThread];
|
||||
warpLoadLookahead(laneIdx, regTmpStates, ptrTileStates, idxTileCur, idxTileNext, num_tiles);
|
||||
|
||||
for (int idx = 0; idx < numTileStatesPerThread; ++idx)
|
||||
{
|
||||
// Bitmask with 1 bits indicating which lane has a tile aggregate
|
||||
const ::cuda::std::uint32_t warp_has_aggregate_mask =
|
||||
__ballot_sync(0xffffffffu, regTmpStates[idx].state == scan_state::tile_aggregate);
|
||||
|
||||
// Bitmask with 1 bits for the contiguous run of lanes having a tile aggregate starting from LSB
|
||||
const ::cuda::std::uint32_t warp_right_aggregates_mask = warp_has_aggregate_mask & (~warp_has_aggregate_mask - 1);
|
||||
|
||||
const ::cuda::std::uint32_t warp_right_aggregates_count = ::cuda::std::popcount(warp_right_aggregates_mask);
|
||||
|
||||
// Only reduce once 32 contiguous tile aggregates are available, so the reduction order is fixed.
|
||||
const ::cuda::std::uint32_t expected_count =
|
||||
static_cast<::cuda::std::uint32_t>(::cuda::std::min(32, idxTileNext - idxTileCur));
|
||||
if (warp_right_aggregates_count < expected_count)
|
||||
{
|
||||
break;
|
||||
}
|
||||
|
||||
const bool use_value = lanemaskEq & warp_right_aggregates_mask;
|
||||
const AccumT value = use_value ? regTmpStates[idx].value : cuda::identity_element<ScanOpT, AccumT>();
|
||||
const AccumT local_aggr = warp_reduce_t{temp_storage}.Reduce(value, scan_op);
|
||||
|
||||
if (expected_count == 32)
|
||||
{
|
||||
aggrExclusiveCtaCur = idxTileCur == 0 ? local_aggr : scan_op(aggrExclusiveCtaCur, local_aggr);
|
||||
idxTileCur += 32;
|
||||
}
|
||||
else
|
||||
{
|
||||
const AccumT full_aggr = idxTileCur == 0 ? local_aggr : scan_op(aggrExclusiveCtaCur, local_aggr);
|
||||
idxTilePrev = idxTileCur;
|
||||
aggrExclusiveCtaPrev = aggrExclusiveCtaCur;
|
||||
return full_aggr;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Only reached when idxTileNext is a multiple of 32; otherwise the final partial batch full aggregate returns inside
|
||||
// the loop above.
|
||||
idxTilePrev = idxTileNext;
|
||||
aggrExclusiveCtaPrev = aggrExclusiveCtaCur;
|
||||
return aggrExclusiveCtaCur; // must only be valid in lane_0
|
||||
}
|
||||
|
||||
#endif // __cccl_ptx_isa >= 860
|
||||
} // namespace detail::warpspeed
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,42 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cuda/std/cstdint>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail::warpspeed
|
||||
{
|
||||
// Move register to uniform register
|
||||
|
||||
// For int32_t and uint32_t, we can use the CREDUX instruction, which is coupled and has a constant latency.
|
||||
// For 64-bit types, we still use __shfl_sync
|
||||
|
||||
[[nodiscard]] _CCCL_DEVICE_API inline int makeWarpUniform(int x)
|
||||
{
|
||||
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_90, (return __reduce_min_sync(~0, x);), (return x;));
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_DEVICE_API inline ::cuda::std::uint32_t makeWarpUniform(::cuda::std::uint32_t x)
|
||||
{
|
||||
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_90, (return __reduce_min_sync(~0, x);), (return x;));
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_DEVICE_API inline ::cuda::std::uint64_t makeWarpUniform(::cuda::std::uint64_t x)
|
||||
{
|
||||
return __shfl_sync(~0, x, 0);
|
||||
}
|
||||
} // namespace detail::warpspeed
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,38 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail::warpspeed
|
||||
{
|
||||
template <typename _Tp>
|
||||
[[nodiscard]] _CCCL_DEVICE_API _Tp* optimizeSmemPtr(const _Tp* smemGeneric)
|
||||
{
|
||||
// See https://nvbugspro.nvidia.com/bug/4907996
|
||||
|
||||
// 1. Convert to 32-bit shared memory pointer
|
||||
::cuda::std::uint32_t smem32 = __cvta_generic_to_shared(smemGeneric);
|
||||
// 2. Pretend to NVVM that the 32-bit pointer is modified. This is required to avoid NVVM constant
|
||||
// propagation from pulling the smem32 definition into loops and branches in subsequent code.
|
||||
asm("" : "+r"(smem32));
|
||||
// 3. Make a generic pointer to smem that is constructed using `__cvta_shared_to_generic`. This
|
||||
// benefits from an
|
||||
// optimization pass in NVVM that performs the following simplification:
|
||||
// __cvta_generic_to_shared(__cvta_shared_to_generic(x)) => x.
|
||||
// In our case, `x` is smem32, which is exactly what we want.
|
||||
return reinterpret_cast<_Tp*>(__cvta_shared_to_generic(smem32));
|
||||
}
|
||||
} // namespace detail::warpspeed
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,45 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cub/detail/warpspeed/resource/smem_ref.cuh>
|
||||
#include <cub/detail/warpspeed/resource/smem_resource_raw.cuh>
|
||||
|
||||
#include <cuda/std/cstdint>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail::warpspeed
|
||||
{
|
||||
template <typename _Tp>
|
||||
struct SmemPhase
|
||||
{
|
||||
SmemResourceRaw& mSmemResourceRaw;
|
||||
int mCurPhase;
|
||||
|
||||
_CCCL_DEVICE_API SmemPhase(SmemResourceRaw& smemResourceRaw, int phase) noexcept
|
||||
: mSmemResourceRaw(smemResourceRaw)
|
||||
, mCurPhase(phase)
|
||||
{}
|
||||
|
||||
[[nodiscard]] _CCCL_DEVICE_API SmemRef<_Tp> acquireRef()
|
||||
{
|
||||
// Wait on barrier
|
||||
mSmemResourceRaw.acquire(mCurPhase);
|
||||
// Return ref
|
||||
return SmemRef<_Tp>(mSmemResourceRaw, mCurPhase);
|
||||
}
|
||||
};
|
||||
} // namespace detail::warpspeed
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,92 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cub/detail/warpspeed/resource/smem_resource_raw.cuh>
|
||||
#include <cub/detail/warpspeed/squad/squad.cuh>
|
||||
|
||||
#include <cuda/std/cstdint>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail::warpspeed
|
||||
{
|
||||
template <typename _Tp>
|
||||
struct SmemRef
|
||||
{
|
||||
SmemResourceRaw& mSmemResourceRaw;
|
||||
int mCurPhase;
|
||||
bool mTxCountIsSet = false;
|
||||
int mTxCount = 0;
|
||||
bool mDoFenceLdsToAsyncProxy = false;
|
||||
|
||||
_CCCL_DEVICE_API SmemRef(SmemResourceRaw& smemResourceRaw, int phase) noexcept
|
||||
: mSmemResourceRaw(smemResourceRaw)
|
||||
, mCurPhase(phase)
|
||||
{}
|
||||
// SmemRef is a non-copyable, non-movable type. It must be passed by (mutable)
|
||||
// reference to be useful. The reason is that it in case of an accidental copy
|
||||
// or move the destructor is called twice. This leads to double-arrivals on
|
||||
// barriers and results in deadlock or a hardware fault.
|
||||
SmemRef(const SmemRef&) = delete; // Delete copy constructor
|
||||
SmemRef(SmemRef&&) = delete; // Delete move constructor
|
||||
SmemRef& operator=(const SmemRef&) = delete; // Delete copy assignment
|
||||
SmemRef& operator=(SmemRef&&) = delete; // Delete move assignment
|
||||
|
||||
_CCCL_DEVICE_API ~SmemRef()
|
||||
{
|
||||
if (mDoFenceLdsToAsyncProxy)
|
||||
{
|
||||
mSmemResourceRaw.fenceLdsToAsyncProxy();
|
||||
}
|
||||
if (mTxCountIsSet)
|
||||
{
|
||||
mSmemResourceRaw.releaseTx(mCurPhase, mTxCount);
|
||||
}
|
||||
else
|
||||
{
|
||||
mSmemResourceRaw.release(mCurPhase);
|
||||
}
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_DEVICE_API _Tp& data() noexcept
|
||||
{
|
||||
return *static_cast<_Tp*>(mSmemResourceRaw.data());
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_DEVICE_API int sizeBytes() const noexcept
|
||||
{
|
||||
return mSmemResourceRaw.mSizeBytes;
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_DEVICE_API uint64_t* ptrCurBarrierRelease()
|
||||
{
|
||||
return mSmemResourceRaw.ptrCurBarrierRelease(mCurPhase);
|
||||
}
|
||||
|
||||
_CCCL_DEVICE_API void squadIncreaseTxCount(const Squad& squad, int txCount)
|
||||
{
|
||||
mTxCountIsSet = true;
|
||||
// Only leader thread increments txCount
|
||||
txCount = squad.isLeaderThread() ? txCount : 0;
|
||||
mTxCount += txCount;
|
||||
}
|
||||
|
||||
_CCCL_DEVICE_API void setFenceLdsToAsyncProxy() noexcept
|
||||
{
|
||||
mDoFenceLdsToAsyncProxy = true;
|
||||
}
|
||||
};
|
||||
} // namespace detail::warpspeed
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,60 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cub/detail/warpspeed/allocators/smem_allocator.cuh>
|
||||
#include <cub/detail/warpspeed/resource/smem_resource_raw.cuh>
|
||||
#include <cub/detail/warpspeed/resource/smem_stage.cuh>
|
||||
#include <cub/detail/warpspeed/sync_handler.cuh>
|
||||
#include <cub/detail/warpspeed/values.cuh>
|
||||
|
||||
#include <cuda/std/__utility/to_underlying.h>
|
||||
#include <cuda/std/cstdint>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail::warpspeed
|
||||
{
|
||||
template <typename _Tp>
|
||||
struct SmemResource : SmemResourceRaw
|
||||
{
|
||||
template <int stageCount>
|
||||
_CCCL_HOST_DEVICE_API SmemResource(SyncHandler& syncHandler, _Tp (&smemBuffer)[stageCount])
|
||||
: SmemResourceRaw(syncHandler, smemBuffer, sizeof(smemBuffer[0]), sizeof(smemBuffer[0]), stageCount)
|
||||
{}
|
||||
|
||||
_CCCL_HOST_DEVICE_API constexpr SmemResource(
|
||||
SyncHandler& syncHandler, SmemAllocator& smemAllocator, Stages stages, Elems elems = Elems{1})
|
||||
: SmemResourceRaw(makeSmemResourceRaw(syncHandler, smemAllocator, stages, elems))
|
||||
{}
|
||||
|
||||
[[nodiscard]] _CCCL_DEVICE_API SmemStage<_Tp> nextStage() noexcept
|
||||
{
|
||||
return SmemStage<_Tp>(*this);
|
||||
}
|
||||
|
||||
private:
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE_API static constexpr SmemResourceRaw
|
||||
makeSmemResourceRaw(SyncHandler& syncHandler, SmemAllocator& smemAllocator, Stages stages, Elems elems = Elems{1})
|
||||
{
|
||||
int align = alignof(_Tp);
|
||||
int sizeBytes = ::cuda::std::to_underlying(elems) * sizeof(_Tp);
|
||||
int strideBytes = sizeBytes;
|
||||
|
||||
void* ptrBase = smemAllocator.alloc(::cuda::std::to_underlying(stages) * strideBytes, align);
|
||||
return {syncHandler, ptrBase, sizeBytes, strideBytes, ::cuda::std::to_underlying(stages)};
|
||||
}
|
||||
};
|
||||
} // namespace detail::warpspeed
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,182 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cub/detail/warpspeed/allocators/smem_allocator.cuh>
|
||||
#include <cub/detail/warpspeed/constant_assert.cuh>
|
||||
#include <cub/detail/warpspeed/squad/squad_desc.cuh>
|
||||
#include <cub/detail/warpspeed/sync_handler.cuh>
|
||||
|
||||
#include <cuda/__ptx/instructions/fence.h>
|
||||
#include <cuda/__ptx/instructions/mbarrier_arrive.h>
|
||||
#include <cuda/__ptx/instructions/mbarrier_wait.h>
|
||||
#include <cuda/__ptx/ptx_dot_variants.h>
|
||||
#include <cuda/std/__type_traits/is_constant_evaluated.h>
|
||||
#include <cuda/std/cstdint>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail::warpspeed
|
||||
{
|
||||
struct SmemResourceRaw
|
||||
{
|
||||
static constexpr int mMaxNumPhases = 4;
|
||||
|
||||
int mStageCurrent = 0;
|
||||
|
||||
int mResourceHandle;
|
||||
::cuda::std::uint8_t* mPtrBase{};
|
||||
int mSizeBytes;
|
||||
int mStride;
|
||||
int mStageCount;
|
||||
int mNumPhases = 0;
|
||||
|
||||
::cuda::std::uint64_t* mPtrBar[mMaxNumPhases]{};
|
||||
int mParity[mMaxNumPhases]{};
|
||||
|
||||
_CCCL_HOST_DEVICE_API constexpr SmemResourceRaw(
|
||||
SyncHandler& syncHandler, void* ptrBase, int sizeBytes, int strideBytes, int stageCount) noexcept
|
||||
: mResourceHandle(syncHandler.registerResource(stageCount))
|
||||
, mSizeBytes(sizeBytes)
|
||||
, mStride(strideBytes)
|
||||
, mStageCount(stageCount)
|
||||
{
|
||||
// we don't need the pointer during constant evaluation (and casting is not allowed)
|
||||
if (!::cuda::std::is_constant_evaluated())
|
||||
{
|
||||
mPtrBase = static_cast<::cuda::std::uint8_t*>(ptrBase);
|
||||
}
|
||||
|
||||
for (int pi = 0; pi < mMaxNumPhases; ++pi)
|
||||
{
|
||||
mParity[pi] = pi == 0 ? 1 : 0;
|
||||
}
|
||||
}
|
||||
|
||||
template <int numSquads>
|
||||
_CCCL_HOST_DEVICE_API constexpr void
|
||||
addPhase(SyncHandler& syncHandler, ::cuda::std::uint64_t* ptrBarrier, const SquadDesc (&squads)[numSquads])
|
||||
{
|
||||
int numOwningThreads = squadCountThreads(squads);
|
||||
|
||||
int curPhase = mNumPhases;
|
||||
mNumPhases++;
|
||||
|
||||
syncHandler.registerPhase(mResourceHandle, numOwningThreads, ptrBarrier);
|
||||
mPtrBar[curPhase] = ptrBarrier;
|
||||
}
|
||||
|
||||
template <int numSquads>
|
||||
_CCCL_HOST_DEVICE_API constexpr void
|
||||
addPhase(SyncHandler& syncHandler, SmemAllocator& smemAllocator, const SquadDesc (&squads)[numSquads])
|
||||
{
|
||||
void* ptrBar_raw = smemAllocator.alloc(mStageCount * sizeof(::cuda::std::uint64_t), alignof(::cuda::std::uint64_t));
|
||||
// we don't need the pointer during constant evaluation (and casting is not allowed)
|
||||
::cuda::std::uint64_t* ptrBar = nullptr;
|
||||
if (!::cuda::std::is_constant_evaluated())
|
||||
{
|
||||
ptrBar = static_cast<::cuda::std::uint64_t*>(ptrBar_raw);
|
||||
}
|
||||
addPhase(syncHandler, ptrBar, squads);
|
||||
}
|
||||
|
||||
_CCCL_HOST_DEVICE_API void
|
||||
addPhase(SyncHandler& syncHandler, ::cuda::std::uint64_t* ptrBarrier, const SquadDesc& squad)
|
||||
{
|
||||
const SquadDesc squads[] = {squad};
|
||||
addPhase(syncHandler, ptrBarrier, squads);
|
||||
}
|
||||
|
||||
_CCCL_HOST_DEVICE_API constexpr void
|
||||
addPhase(SyncHandler& syncHandler, SmemAllocator& smemAllocator, const SquadDesc& squad)
|
||||
{
|
||||
const SquadDesc squads[] = {squad};
|
||||
addPhase(syncHandler, smemAllocator, squads);
|
||||
}
|
||||
|
||||
_CCCL_DEVICE_API void incrementStage()
|
||||
{
|
||||
if (mStageCurrent == mStageCount - 1)
|
||||
{
|
||||
mStageCurrent = 0;
|
||||
// We loop over all phases with a conditional on resNumPhases. If we
|
||||
// directly loop over only resNumPhases, then the SROA optimization does
|
||||
// not kick in and the mParity array is spilled to the stack.
|
||||
for (int pi = 0; pi < mMaxNumPhases; ++pi)
|
||||
{
|
||||
if (pi < mNumPhases)
|
||||
{
|
||||
mParity[pi] ^= 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
mStageCurrent++;
|
||||
}
|
||||
}
|
||||
|
||||
_CCCL_DEVICE_API void* data()
|
||||
{
|
||||
return (void*) (mPtrBase + mStageCurrent * mStride);
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_DEVICE_API ::cuda::std::uint64_t* ptrCurBarrierRelease(int phase)
|
||||
{
|
||||
::cuda::std::uint64_t* ptrBarPhase = mPtrBar[phase];
|
||||
_WS_CONSTANT_ASSERT(phase < mNumPhases, "Phase exceeds limit.");
|
||||
return &ptrBarPhase[mStageCurrent];
|
||||
}
|
||||
_CCCL_DEVICE_API void release(int phase)
|
||||
{
|
||||
_WS_CONSTANT_ASSERT(phase < mNumPhases, "Phase exceeds limit.");
|
||||
::cuda::ptx::mbarrier_arrive(ptrCurBarrierRelease(phase));
|
||||
}
|
||||
|
||||
_CCCL_DEVICE_API void releaseTx(int phase, int txCount)
|
||||
{
|
||||
_WS_CONSTANT_ASSERT(phase < mNumPhases, "Phase exceeds limit.");
|
||||
::cuda::ptx::mbarrier_arrive_expect_tx(
|
||||
::cuda::ptx::sem_release, ::cuda::ptx::scope_cta, ::cuda::ptx::space_shared, ptrCurBarrierRelease(phase), txCount);
|
||||
}
|
||||
|
||||
_CCCL_DEVICE_API void fenceLdsToAsyncProxy()
|
||||
{
|
||||
::cuda::ptx::fence_proxy_async(::cuda::ptx::space_shared);
|
||||
}
|
||||
|
||||
_CCCL_DEVICE_API void releaseLdsToAsyncProxy(int phase)
|
||||
{
|
||||
// First fence
|
||||
fenceLdsToAsyncProxy();
|
||||
// Then perform a normal release
|
||||
release(phase);
|
||||
}
|
||||
|
||||
_CCCL_DEVICE_API void acquire(int phase)
|
||||
{
|
||||
_WS_CONSTANT_ASSERT(phase < mNumPhases, "Phase exceeds limit.");
|
||||
|
||||
// The release of the previous phase occurs on the `phase - 1` barrier. So
|
||||
// that is what we wait on.
|
||||
int phaseAcq = (mNumPhases + phase - 1) % mNumPhases;
|
||||
::cuda::std::uint64_t* ptrBarPhase = mPtrBar[phaseAcq];
|
||||
|
||||
while (!::cuda::ptx::mbarrier_try_wait_parity(&ptrBarPhase[mStageCurrent], mParity[phase]))
|
||||
{
|
||||
}
|
||||
}
|
||||
};
|
||||
} // namespace detail::warpspeed
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,91 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cub/detail/warpspeed/constant_assert.cuh>
|
||||
#include <cub/detail/warpspeed/resource/smem_phase.cuh>
|
||||
#include <cub/detail/warpspeed/resource/smem_resource_raw.cuh>
|
||||
|
||||
#include <cuda/std/__tuple_dir/tuple_element.h>
|
||||
#include <cuda/std/__tuple_dir/tuple_size.h>
|
||||
#include <cuda/std/cstdint>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail::warpspeed
|
||||
{
|
||||
template <typename _Tp>
|
||||
struct SmemStage
|
||||
{
|
||||
SmemResourceRaw& mSmemResourceRaw;
|
||||
|
||||
_CCCL_DEVICE_API SmemStage(SmemResourceRaw& smemResourceRaw) noexcept
|
||||
: mSmemResourceRaw(smemResourceRaw)
|
||||
{}
|
||||
|
||||
_CCCL_DEVICE_API ~SmemStage()
|
||||
{
|
||||
mSmemResourceRaw.incrementStage();
|
||||
}
|
||||
|
||||
// SmemStage is a non-copyable, non-movable type. It must be passed by (mutable)
|
||||
// reference to be useful. The reason is that it in case of an accidental copy
|
||||
// or move the destructor is called twice. This leads to double-increment of
|
||||
// the stage index and results in deadlock or a hardware fault.
|
||||
SmemStage(const SmemStage&) = delete; // Delete copy constructor
|
||||
SmemStage(SmemStage&&) = delete; // Delete move constructor
|
||||
SmemStage& operator=(const SmemStage&) = delete; // Delete copy assignment
|
||||
SmemStage& operator=(const SmemStage&&) = delete; // Delete move assignment
|
||||
};
|
||||
|
||||
// Helper: Container to expose SmemPhase for structured binding
|
||||
template <typename _Tp, ::cuda::std::size_t numPhases>
|
||||
struct SmemPhaseStructuredBinding
|
||||
{
|
||||
SmemResourceRaw& mSmemResourceRaw;
|
||||
|
||||
template <::cuda::std::size_t _Index>
|
||||
[[nodiscard]] _CCCL_DEVICE_API SmemPhase<_Tp> get() const
|
||||
{
|
||||
return SmemPhase<_Tp>(mSmemResourceRaw, _Index);
|
||||
}
|
||||
};
|
||||
|
||||
// The binding function
|
||||
template <::cuda::std::size_t numPhases, typename _Tp>
|
||||
[[nodiscard]] _CCCL_DEVICE_API SmemPhaseStructuredBinding<_Tp, numPhases> bindPhases(SmemStage<_Tp>& smemStage)
|
||||
{
|
||||
_WS_CONSTANT_ASSERT(smemStage.mSmemResourceRaw.mNumPhases == numPhases,
|
||||
"Number of bound phases must match resource phases.");
|
||||
|
||||
return SmemPhaseStructuredBinding<_Tp, numPhases>{smemStage.mSmemResourceRaw};
|
||||
}
|
||||
} // namespace detail::warpspeed
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
|
||||
// Tuple protocol specializations
|
||||
namespace std
|
||||
{
|
||||
template <typename _Tp, size_t numPhases>
|
||||
struct tuple_size<CUB_NS_QUALIFIER::detail::warpspeed::SmemPhaseStructuredBinding<_Tp, numPhases>>
|
||||
{
|
||||
static constexpr size_t value = numPhases;
|
||||
};
|
||||
|
||||
template <typename _Tp, size_t _Index, ::cuda::std::size_t numPhases>
|
||||
struct tuple_element<_Index, CUB_NS_QUALIFIER::detail::warpspeed::SmemPhaseStructuredBinding<_Tp, numPhases>>
|
||||
{
|
||||
using type = CUB_NS_QUALIFIER::detail::warpspeed::SmemPhase<_Tp>;
|
||||
};
|
||||
} // namespace std
|
||||
@@ -1,44 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cub/detail/warpspeed/make_warp_uniform.cuh>
|
||||
|
||||
#include <cuda/__ptx/instructions/get_sreg.h>
|
||||
#include <cuda/std/cstdint>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail::warpspeed
|
||||
{
|
||||
// Commonly used special registers that we should cache in registers or uniform
|
||||
// registers.
|
||||
struct SpecialRegisters
|
||||
{
|
||||
const ::cuda::std::uint32_t clusterCtaRank;
|
||||
const ::cuda::std::uint32_t blockIdxX;
|
||||
const ::cuda::std::uint32_t threadIdxX;
|
||||
const ::cuda::std::uint32_t warpIdx;
|
||||
const ::cuda::std::uint32_t laneIdx;
|
||||
};
|
||||
|
||||
[[nodiscard]] _CCCL_DEVICE_API inline SpecialRegisters getSpecialRegisters()
|
||||
{
|
||||
::cuda::std::uint32_t clusterCtaRank = ::cuda::ptx::get_sreg_cluster_ctarank();
|
||||
::cuda::std::uint32_t threadIdxX = threadIdx.x;
|
||||
::cuda::std::uint32_t warpIdx = makeWarpUniform(threadIdxX / 32);
|
||||
return {clusterCtaRank, blockIdx.x, threadIdxX, warpIdx, ::cuda::ptx::get_sreg_laneid()};
|
||||
}
|
||||
} // namespace detail::warpspeed
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,388 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cub/detail/warpspeed/resource/smem_ref.cuh>
|
||||
#include <cub/detail/warpspeed/squad/squad.cuh>
|
||||
|
||||
#include <cuda/__memory/align_down.h>
|
||||
#include <cuda/__memory/align_up.h>
|
||||
#include <cuda/__ptx/instructions/cp_async_bulk.h>
|
||||
#include <cuda/__ptx/instructions/cp_async_bulk_commit_group.h>
|
||||
#include <cuda/__ptx/instructions/cp_async_bulk_wait_group.h>
|
||||
#include <cuda/__ptx/instructions/elect_sync.h>
|
||||
#include <cuda/__ptx/instructions/fence.h>
|
||||
#include <cuda/std/__type_traits/make_nbit_int.h>
|
||||
#include <cuda/std/cstdint>
|
||||
|
||||
#include <nv/target>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail::warpspeed
|
||||
{
|
||||
#if __cccl_ptx_isa >= 860
|
||||
|
||||
template <typename Tp>
|
||||
struct CpAsyncOobInfo
|
||||
{
|
||||
// The aligned up and down pointers below must be ::cuda::std::byte*, since the nearest aligned up/down ptr may not
|
||||
// point to a multiple of sizeof(Tp). E.g. a uchar3* pointing to address 0x...5 will be aligned down to 0x...0, but
|
||||
// that's not a valid start for an uchar3 in that array. So we must express all aligned pointers in bytes here.
|
||||
|
||||
::cuda::std::byte* ptrGmem;
|
||||
::cuda::std::byte* ptrGmemStartAlignDown;
|
||||
::cuda::std::byte* ptrGmemStartAlignUp;
|
||||
::cuda::std::byte* ptrGmemEnd;
|
||||
::cuda::std::byte* ptrGmemEndAlignDown;
|
||||
::cuda::std::byte* ptrGmemEndAlignUp;
|
||||
::cuda::std::uint32_t overCopySizeBytes;
|
||||
::cuda::std::uint32_t underCopySizeBytes;
|
||||
::cuda::std::uint32_t origCopySizeBytes;
|
||||
::cuda::std::uint32_t smemStartSkipBytes; // ptrSmem + smemStartSkipBytes will point to the first valid element copied
|
||||
// from ptrGmem
|
||||
::cuda::std::uint32_t smemEndBytesAfter16BBoundary; // number of bytes after the last 16B boundary in GMEM/SMEM that
|
||||
// still contains valid (partial) elements
|
||||
};
|
||||
|
||||
template <typename Tp>
|
||||
_CCCL_DEVICE_API _CCCL_FORCEINLINE CpAsyncOobInfo<Tp> prepareCpAsyncOob(Tp* ptrGmem, ::cuda::std::uint32_t sizeElem)
|
||||
{
|
||||
auto ptrGmemBytes = reinterpret_cast<::cuda::std::byte*>(ptrGmem);
|
||||
auto ptrGmemEnd = reinterpret_cast<::cuda::std::byte*>(ptrGmem + sizeElem);
|
||||
|
||||
// We will copy from [ptrGmemBase, ptrGmemEnd). Both pointers have to be 16B aligned.
|
||||
::cuda::std::byte* ptrGmemStartAlignDown = ::cuda::align_down(ptrGmemBytes, 16);
|
||||
::cuda::std::byte* ptrGmemStartAlignUp = ::cuda::align_up(ptrGmemBytes, 16);
|
||||
::cuda::std::byte* ptrGmemEndAlignUp = ::cuda::align_up(ptrGmemEnd, 16);
|
||||
::cuda::std::byte* ptrGmemEndAlignDown = ::cuda::align_down(ptrGmemEnd, 16);
|
||||
|
||||
// Compute the final copy size in bytes. It can be either sizeElem or sizeElem + 16 / sizeof(T).
|
||||
const auto origCopySizeBytes = static_cast<::cuda::std::uint32_t>(sizeof(Tp) * sizeElem);
|
||||
const auto overCopySizeBytes = static_cast<::cuda::std::uint32_t>(ptrGmemEndAlignUp - ptrGmemStartAlignDown);
|
||||
auto underCopySizeBytes = static_cast<::cuda::std::uint32_t>(ptrGmemEndAlignDown - ptrGmemStartAlignUp);
|
||||
if (origCopySizeBytes < underCopySizeBytes)
|
||||
{
|
||||
// If ptrGmemStart and ptrGmemEnd are aligned to [1, .., 15] bytes, then
|
||||
// when we align the one up and the other down we get overflow. We check for
|
||||
// that here. In that case, the undercopy size is zero.
|
||||
underCopySizeBytes = 0;
|
||||
}
|
||||
|
||||
_CCCL_ASSERT(overCopySizeBytes % 16 == 0, "");
|
||||
_CCCL_ASSERT(underCopySizeBytes % 16 == 0, "");
|
||||
|
||||
return {
|
||||
ptrGmemBytes,
|
||||
ptrGmemStartAlignDown,
|
||||
ptrGmemStartAlignUp,
|
||||
ptrGmemEnd,
|
||||
ptrGmemEndAlignDown,
|
||||
ptrGmemEndAlignUp,
|
||||
overCopySizeBytes,
|
||||
underCopySizeBytes,
|
||||
origCopySizeBytes,
|
||||
static_cast<::cuda::std::uint32_t>(ptrGmemBytes - ptrGmemStartAlignDown),
|
||||
static_cast<::cuda::std::uint32_t>(ptrGmemEnd - ptrGmemEndAlignDown),
|
||||
};
|
||||
}
|
||||
|
||||
template <typename ResourceTp, typename Tp>
|
||||
_CCCL_DEVICE_API void squadLoadBulk(Squad squad, SmemRef<ResourceTp>& refDestSmem, CpAsyncOobInfo<Tp> cpAsyncOobInfo)
|
||||
{
|
||||
::cuda::std::byte* ptrSmem = refDestSmem.data().inout;
|
||||
_CCCL_ASSERT(::cuda::is_aligned(ptrSmem, 16), "");
|
||||
::cuda::std::uint64_t* ptrBar = refDestSmem.ptrCurBarrierRelease();
|
||||
|
||||
if constexpr (alignof(Tp) >= 16)
|
||||
{
|
||||
// for alignments larger than 16, we can just bulk copy, even just a single element
|
||||
if (squad.isLeaderThread())
|
||||
{
|
||||
::cuda::ptx::cp_async_bulk(
|
||||
::cuda::std::conditional_t<__cccl_ptx_isa >= 860, ::cuda::ptx::space_shared_t, ::cuda::ptx::space_cluster_t>{},
|
||||
::cuda::ptx::space_global,
|
||||
ptrSmem,
|
||||
cpAsyncOobInfo.ptrGmem,
|
||||
cpAsyncOobInfo.origCopySizeBytes,
|
||||
ptrBar);
|
||||
}
|
||||
refDestSmem.squadIncreaseTxCount(squad, cpAsyncOobInfo.underCopySizeBytes);
|
||||
}
|
||||
else
|
||||
{
|
||||
// for alignments smaller than 16, we can overcopy but need to declare the ignored bytes left and right
|
||||
# if __cccl_ptx_isa >= 920
|
||||
if (squad.isLeaderThread())
|
||||
{
|
||||
::cuda::ptx::cp_async_bulk_ignore_oob(
|
||||
::cuda::ptx::space_shared,
|
||||
::cuda::ptx::space_global,
|
||||
ptrSmem,
|
||||
cpAsyncOobInfo.ptrGmemStartAlignDown,
|
||||
cpAsyncOobInfo.overCopySizeBytes,
|
||||
/* ignore left */ cpAsyncOobInfo.smemStartSkipBytes,
|
||||
/* ignore right */ cpAsyncOobInfo.ptrGmemEndAlignUp - cpAsyncOobInfo.ptrGmemEnd,
|
||||
ptrBar);
|
||||
}
|
||||
refDestSmem.squadIncreaseTxCount(squad, cpAsyncOobInfo.overCopySizeBytes);
|
||||
# else // __cccl_ptx_isa >= 920
|
||||
// if we don't have cp_async_bulk_ignore_oob, we have to undercopy and copy head and tail elements manually
|
||||
|
||||
// handle small copies first. If we have less than 16 bytes we may not straddle a 16B boundary
|
||||
if (cpAsyncOobInfo.origCopySizeBytes < 16)
|
||||
{
|
||||
const auto elemCount = cpAsyncOobInfo.origCopySizeBytes / sizeof(Tp);
|
||||
_CCCL_ASSERT(elemCount <= squad.threadCount(), "");
|
||||
if (squad.threadRank() < elemCount)
|
||||
{
|
||||
reinterpret_cast<Tp*>(ptrSmem + cpAsyncOobInfo.smemStartSkipBytes)[squad.threadRank()] =
|
||||
reinterpret_cast<const Tp*>(cpAsyncOobInfo.ptrGmem)[squad.threadRank()];
|
||||
}
|
||||
return; // no bulk copy has been performed so we don't need to update the tx count of any barrier
|
||||
}
|
||||
|
||||
// copies larger than 16 byte which straddle at least one 16B boundary, so we have dedicated start and end copies
|
||||
|
||||
const bool doStartCopy = cpAsyncOobInfo.smemStartSkipBytes > 0;
|
||||
|
||||
::cuda::std::byte* ptrSmemMiddle = ptrSmem;
|
||||
if (doStartCopy)
|
||||
{
|
||||
ptrSmemMiddle += 16;
|
||||
}
|
||||
|
||||
// TODO(bgruber): we could skip the middle if underCopySizeBytes is zero
|
||||
if (squad.isLeaderThread())
|
||||
{
|
||||
::cuda::ptx::cp_async_bulk(
|
||||
::cuda::std::conditional_t<__cccl_ptx_isa >= 860, ::cuda::ptx::space_shared_t, ::cuda::ptx::space_cluster_t>{},
|
||||
::cuda::ptx::space_global,
|
||||
ptrSmemMiddle,
|
||||
cpAsyncOobInfo.ptrGmemStartAlignUp,
|
||||
cpAsyncOobInfo.underCopySizeBytes,
|
||||
ptrBar);
|
||||
}
|
||||
refDestSmem.squadIncreaseTxCount(squad, cpAsyncOobInfo.underCopySizeBytes);
|
||||
|
||||
// we cannot use Tp to load the head and tail elements, because sizeof(Tp) may be larger than alignof(Tp)
|
||||
using load_word_t = ::cuda::std::__make_nbit_uint_t<alignof(Tp) * CHAR_BIT>;
|
||||
|
||||
const int head_elements = (cpAsyncOobInfo.ptrGmemStartAlignUp - cpAsyncOobInfo.ptrGmem) / sizeof(load_word_t);
|
||||
const int tail_elements = (cpAsyncOobInfo.ptrGmemEnd - cpAsyncOobInfo.ptrGmemEndAlignDown) / sizeof(load_word_t);
|
||||
_CCCL_ASSERT(head_elements <= squad.threadCount(), "");
|
||||
_CCCL_ASSERT(tail_elements <= squad.threadCount(), "");
|
||||
load_word_t head_value, tail_value;
|
||||
if (squad.threadRank() < head_elements)
|
||||
{
|
||||
head_value = reinterpret_cast<const load_word_t*>(cpAsyncOobInfo.ptrGmem)[squad.threadRank()];
|
||||
}
|
||||
if (squad.threadRank() < tail_elements)
|
||||
{
|
||||
tail_value = reinterpret_cast<const load_word_t*>(cpAsyncOobInfo.ptrGmemEndAlignDown)[squad.threadRank()];
|
||||
}
|
||||
|
||||
if (squad.threadRank() < head_elements)
|
||||
{
|
||||
reinterpret_cast<load_word_t*>(ptrSmem + cpAsyncOobInfo.smemStartSkipBytes)[squad.threadRank()] = head_value;
|
||||
}
|
||||
if (squad.threadRank() < tail_elements)
|
||||
{
|
||||
reinterpret_cast<load_word_t*>(ptrSmemMiddle + cpAsyncOobInfo.underCopySizeBytes)[squad.threadRank()] =
|
||||
tail_value;
|
||||
}
|
||||
# endif // __cccl_ptx_isa >= 920
|
||||
}
|
||||
}
|
||||
|
||||
_CCCL_DEVICE_API _CCCL_FORCEINLINE void squadStoreMasked16B(
|
||||
Squad squad,
|
||||
::cuda::std::byte* dstGmem,
|
||||
const ::cuda::std::byte* srcSmem,
|
||||
::cuda::std::uint16_t byteMask,
|
||||
int firstByte,
|
||||
int lastByte)
|
||||
{
|
||||
NV_IF_ELSE_TARGET(
|
||||
NV_PROVIDES_SM_100,
|
||||
(if (::cuda::ptx::elect_sync(~0)) {
|
||||
::cuda::ptx::cp_async_bulk_cp_mask(
|
||||
::cuda::ptx::space_global, ::cuda::ptx::space_shared, dstGmem, srcSmem, /*size*/ 16, byteMask);
|
||||
}),
|
||||
({
|
||||
const int rank = squad.threadRank();
|
||||
if (firstByte <= rank && rank < lastByte)
|
||||
{
|
||||
dstGmem[rank] = srcSmem[rank];
|
||||
}
|
||||
}));
|
||||
}
|
||||
|
||||
template <typename OutputT>
|
||||
_CCCL_DEVICE_API void
|
||||
squadStoreBulkSync(Squad squad, CpAsyncOobInfo<OutputT> cpAsyncOobInfo, const ::cuda::std::byte* srcSmem)
|
||||
{
|
||||
// This function performs either 1 copy, or three copies, depending on the
|
||||
// size and alignment of the output tile in global memory.
|
||||
//
|
||||
// If the output tile is contained in a single 16-byte aligned and sized region, then we
|
||||
// only perform a single masked copy.
|
||||
//
|
||||
// If the output tile is larger than 16 bytes or straddles two 16-byte aligned and sized regions, then
|
||||
// we perform up to three copies:
|
||||
// - One copy for the first up to 15 bytes at the start of the region.
|
||||
// - One copy that starts at a 16-byte aligned address and ends at the latest 16-byte aligned address.
|
||||
// - One copy that cleans up the last up to 15 bytes.
|
||||
if (squad.isLeaderWarp())
|
||||
{
|
||||
// Acquire shared memory in async proxy
|
||||
// Perform fence.proxy.async with full warp to avoid BSSY+BSYNC
|
||||
::cuda::ptx::fence_proxy_async(::cuda::ptx::space_shared);
|
||||
|
||||
# if _CCCL_CUDA_COMPILER(NVCC, <, 13, 3)
|
||||
// for some reason the optimizer propagates some information from the computation of
|
||||
// overCopySizeBytes to the masked bulk copy below and generates an unaligned access error.
|
||||
// The artificial read modification of overCopySizeBytes prevents the propagation here works around this.
|
||||
// It also solves the issue described in nvbug 5848313 by accident on nvcc 13.2+
|
||||
asm volatile("" : "+r"(cpAsyncOobInfo.overCopySizeBytes));
|
||||
# endif // _CCCL_CUDA_COMPILER(NVCC, <, 13, 3)
|
||||
|
||||
const bool doStartCopy = cpAsyncOobInfo.smemStartSkipBytes > 0;
|
||||
const bool doEndCopy = cpAsyncOobInfo.smemEndBytesAfter16BBoundary > 0;
|
||||
const bool doMiddleCopy = cpAsyncOobInfo.ptrGmemStartAlignUp != cpAsyncOobInfo.ptrGmemEndAlignUp;
|
||||
|
||||
constexpr ::cuda::std::uint16_t byteMask = 0xFFFF;
|
||||
const ::cuda::std::uint16_t byteMaskStart = byteMask << cpAsyncOobInfo.smemStartSkipBytes;
|
||||
const ::cuda::std::uint16_t byteMaskEnd = byteMask >> (16 - cpAsyncOobInfo.smemEndBytesAfter16BBoundary) % 16;
|
||||
// byteMaskStart contains zeroes at the left
|
||||
# if _CCCL_CUDA_COMPILER(NVCC, >=, 13, 2)
|
||||
const ::cuda::std::uint16_t byteMaskSmall = byteMaskStart & byteMaskEnd;
|
||||
# else // _CCCL_CUDA_COMPILER(NVCC, >=, 13, 2)
|
||||
// `ptxas fatal : (C7907) Internal compiler error`, see nvbug 5848313
|
||||
const ::cuda::std::uint16_t byteMaskSmall =
|
||||
byteMaskStart & (byteMask >> (16 - (cpAsyncOobInfo.ptrGmemEnd - cpAsyncOobInfo.ptrGmemStartAlignDown)));
|
||||
# endif // _CCCL_CUDA_COMPILER(NVCC, >=, 13, 2)
|
||||
|
||||
const ::cuda::std::byte* ptrSmemMiddle = srcSmem;
|
||||
if (doStartCopy)
|
||||
{
|
||||
ptrSmemMiddle += 16;
|
||||
}
|
||||
|
||||
if (doMiddleCopy)
|
||||
{
|
||||
// Copy the middle part. Starting at byte 0 or 16 in shared memory. This
|
||||
// is the large copy. We perform this one first, so that the compiler can
|
||||
// (hopefully) hide all the arithmetic behind this instruction.
|
||||
if (::cuda::ptx::elect_sync(~0))
|
||||
{
|
||||
// need to work around another optimizer bug, see: https://github.com/NVIDIA/cccl/issues/8644
|
||||
# if _CCCL_CUDA_COMPILER(NVCC, <, 13, 3)
|
||||
asm volatile("" : "+l"(cpAsyncOobInfo.ptrGmemStartAlignUp));
|
||||
# endif // _CCCL_CUDA_COMPILER(NVCC, <, 13, 3)
|
||||
::cuda::ptx::cp_async_bulk(
|
||||
::cuda::ptx::space_global,
|
||||
::cuda::ptx::space_shared,
|
||||
cpAsyncOobInfo.ptrGmemStartAlignUp,
|
||||
ptrSmemMiddle,
|
||||
cpAsyncOobInfo.underCopySizeBytes);
|
||||
}
|
||||
if (doStartCopy)
|
||||
{
|
||||
// need to work around yet another optimizer bug, see: https://github.com/NVIDIA/cccl/issues/8838
|
||||
# if _CCCL_CUDA_COMPILER(NVCC, <, 13, 3)
|
||||
asm volatile("" : "+l"(cpAsyncOobInfo.ptrGmemStartAlignDown));
|
||||
asm volatile("" : "+l"(srcSmem));
|
||||
# endif // _CCCL_CUDA_COMPILER(NVCC, <, 13, 3)
|
||||
// Copy a subset of the first 16 bytes
|
||||
squadStoreMasked16B(
|
||||
squad,
|
||||
cpAsyncOobInfo.ptrGmemStartAlignDown,
|
||||
srcSmem,
|
||||
byteMaskStart,
|
||||
static_cast<int>(cpAsyncOobInfo.smemStartSkipBytes),
|
||||
16);
|
||||
}
|
||||
if (doEndCopy)
|
||||
{
|
||||
# if _CCCL_CUDA_COMPILER(NVHPC)
|
||||
// nvc++ seems to have an optimizer bug, crashing with an unaligned access error below. The addresses are fine
|
||||
// when printed, so let's shake the optimizer a bit.
|
||||
asm volatile("" : "+l"(cpAsyncOobInfo.ptrGmemEndAlignDown));
|
||||
# endif // _CCCL_CUDA_COMPILER(NVHPC)
|
||||
|
||||
// Copy a subset of the last 16 bytes
|
||||
squadStoreMasked16B(
|
||||
squad,
|
||||
cpAsyncOobInfo.ptrGmemEndAlignDown,
|
||||
ptrSmemMiddle + cpAsyncOobInfo.underCopySizeBytes,
|
||||
byteMaskEnd,
|
||||
0,
|
||||
static_cast<int>(cpAsyncOobInfo.smemEndBytesAfter16BBoundary));
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
// Copy a subset of the first 16 bytes
|
||||
squadStoreMasked16B(
|
||||
squad,
|
||||
cpAsyncOobInfo.ptrGmemStartAlignDown,
|
||||
srcSmem,
|
||||
byteMaskSmall,
|
||||
static_cast<int>(cpAsyncOobInfo.smemStartSkipBytes),
|
||||
static_cast<int>(cpAsyncOobInfo.ptrGmemEnd - cpAsyncOobInfo.ptrGmemStartAlignDown));
|
||||
}
|
||||
// Commit and wait for store to have completed reading from shared memory
|
||||
::cuda::ptx::cp_async_bulk_commit_group();
|
||||
::cuda::ptx::cp_async_bulk_wait_group_read(::cuda::ptx::n32_t<0>{});
|
||||
}
|
||||
}
|
||||
|
||||
#endif // __cccl_ptx_isa >= 860
|
||||
|
||||
template <typename InputT, typename AccumT, int ElemPerThread>
|
||||
_CCCL_DEVICE_API void squadLoadSmem(Squad squad, AccumT (&outReg)[ElemPerThread], const InputT* smemBuf)
|
||||
{
|
||||
for (int i = 0; i < ElemPerThread; ++i)
|
||||
{
|
||||
const int elem_idx = squad.threadRank() * ElemPerThread + i;
|
||||
outReg[i] = smemBuf[elem_idx];
|
||||
}
|
||||
}
|
||||
|
||||
template <typename OutputT, typename AccumT, int ElemPerThread>
|
||||
_CCCL_DEVICE_API void squadStoreSmem(Squad squad, OutputT* smemBuf, const AccumT (&inReg)[ElemPerThread])
|
||||
{
|
||||
for (int i = 0; i < ElemPerThread; ++i)
|
||||
{
|
||||
const int elem_idx = squad.threadRank() * ElemPerThread + i;
|
||||
smemBuf[elem_idx] = inReg[i];
|
||||
}
|
||||
}
|
||||
|
||||
template <typename OutputT, typename AccumT, int ElemPerThread>
|
||||
_CCCL_DEVICE_API void
|
||||
squadStoreSmemPartial(Squad squad, OutputT* smemBuf, const AccumT (&inReg)[ElemPerThread], int beginIndex, int endIndex)
|
||||
{
|
||||
for (int i = 0; i < ElemPerThread; ++i)
|
||||
{
|
||||
const int elem_idx = squad.threadRank() * ElemPerThread + i;
|
||||
if (beginIndex <= elem_idx && elem_idx < endIndex)
|
||||
{
|
||||
smemBuf[elem_idx - beginIndex] = inReg[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
} // namespace detail::warpspeed
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,159 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cub/detail/warpspeed/special_registers.cuh>
|
||||
#include <cub/detail/warpspeed/squad/squad_desc.cuh>
|
||||
|
||||
#include <cuda/__ptx/instructions/elect_sync.h>
|
||||
#include <cuda/std/array>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail::warpspeed
|
||||
{
|
||||
// Squad - device squad instance
|
||||
//
|
||||
// A squad is a collection of warps that work together in a warp-specialized
|
||||
// kernel. A warp-specialized kernel has multiple squads that perform part of
|
||||
// the computation.
|
||||
//
|
||||
// The Squad class is a device runtime instance of a squad. It provides
|
||||
// functionality to determine the rank of the current thread or warp in the
|
||||
// squad, and to sync all threads in the squad.
|
||||
struct Squad : SquadDesc
|
||||
{
|
||||
SpecialRegisters mSpecialRegisters;
|
||||
bool mIsWarpLeader = false;
|
||||
bool mIsLeaderWarp = false;
|
||||
|
||||
_CCCL_DEVICE_API Squad(SquadDesc squadStatic, SpecialRegisters specialRegisters)
|
||||
: SquadDesc(squadStatic)
|
||||
, mSpecialRegisters(specialRegisters)
|
||||
{
|
||||
mIsWarpLeader = ::cuda::ptx::elect_sync(~0);
|
||||
mIsLeaderWarp = warpRank() == 0;
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_DEVICE_API int warpRank() const
|
||||
{
|
||||
return static_cast<int>(mSpecialRegisters.warpIdx % this->warpCount());
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_DEVICE_API int threadRank() const
|
||||
{
|
||||
return static_cast<int>(mSpecialRegisters.threadIdxX % this->threadCount());
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_DEVICE_API bool isLeaderThread() const
|
||||
{
|
||||
return mIsWarpLeader && mIsLeaderWarp;
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_DEVICE_API bool isLeaderWarp() const
|
||||
{
|
||||
return mIsLeaderWarp;
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_DEVICE_API bool isLeaderThreadOfWarp() const
|
||||
{
|
||||
return mIsWarpLeader;
|
||||
}
|
||||
|
||||
_CCCL_DEVICE_API void syncThreads() const
|
||||
{
|
||||
// barrier 0 is reserved for __syncthreads(). We use barrier ids 1, ...
|
||||
const int barrierIdx = this->mSquadIdx + 1;
|
||||
|
||||
__barrier_sync_count(barrierIdx, this->threadCount());
|
||||
}
|
||||
};
|
||||
// squadDispatch
|
||||
//
|
||||
// squadDispatch is used at the start of the kernel. It takes an array of squad
|
||||
// descriptors and determines which squad the current thread belongs to. The
|
||||
// lambda `f: (Squad) -> void` is called with the squad currently active on this
|
||||
// thread.
|
||||
//
|
||||
// Typically, the user will call the kernel body with the active squad.
|
||||
//
|
||||
// Implementation notes:
|
||||
//
|
||||
// Dispatch to squad based on warp index using a binary search. This balances
|
||||
// the number of BRA instructions per squad and avoids NVVM inserting BRX
|
||||
// instructions. BRX instructions require a jump table that is loaded from GCC,
|
||||
// which incurs latency.
|
||||
//
|
||||
// The benefit of this function for fastScan is that adding a new squad doesn't
|
||||
// require code changes in the dispatch. For low-latency inference, I hope that
|
||||
// the avoidance of linear search and BRX instructions translates into latency
|
||||
// reductions.
|
||||
//
|
||||
template <int numSquads, typename F>
|
||||
_CCCL_DEVICE_API _CCCL_FORCEINLINE void
|
||||
squadDispatch(SpecialRegisters sr, const SquadDesc (&squads)[numSquads], F f, int warpIdxStart = 0)
|
||||
{
|
||||
static_assert(numSquads > 0);
|
||||
if (numSquads == 1)
|
||||
{
|
||||
// Leaf
|
||||
SquadDesc squad = squads[0];
|
||||
|
||||
if (static_cast<unsigned>(warpIdxStart) <= sr.warpIdx
|
||||
&& sr.warpIdx < static_cast<unsigned>(warpIdxStart + squad.warpCount()))
|
||||
{
|
||||
f(Squad(squad, sr));
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
constexpr int mid = numSquads / 2;
|
||||
// Left
|
||||
int warpIdxStartMid = warpIdxStart;
|
||||
for (int gi = 0; gi < mid; ++gi)
|
||||
{
|
||||
warpIdxStartMid += squads[gi].warpCount();
|
||||
}
|
||||
if (sr.warpIdx < static_cast<unsigned>(warpIdxStartMid))
|
||||
{
|
||||
if constexpr (0 < mid)
|
||||
{
|
||||
SquadDesc squadsLeft[mid];
|
||||
for (int gi = 0; gi < mid; ++gi)
|
||||
{
|
||||
squadsLeft[gi] = squads[gi];
|
||||
}
|
||||
squadDispatch(sr, squadsLeft, f, warpIdxStart);
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
SquadDesc squadsRight[numSquads - mid]{};
|
||||
for (int gi = 0; gi < numSquads - mid; ++gi)
|
||||
{
|
||||
squadsRight[gi] = squads[mid + gi];
|
||||
}
|
||||
squadDispatch(sr, squadsRight, f, warpIdxStartMid);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <::cuda::std::size_t numSquads, typename F>
|
||||
_CCCL_DEVICE_API _CCCL_FORCEINLINE void
|
||||
squadDispatch(SpecialRegisters sr, ::cuda::std::array<SquadDesc, numSquads> squads, F f, int warpIdxStart = 0)
|
||||
{
|
||||
squadDispatch<numSquads>(sr, squads.__elems_, f, warpIdxStart);
|
||||
}
|
||||
} // namespace detail::warpspeed
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,77 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail::warpspeed
|
||||
{
|
||||
// SquadDesc - squad descriptor
|
||||
//
|
||||
// A squad is a collection of warps that work together in a warp-specialized
|
||||
// kernel. A warp-specialized kernel has multiple squads that perform part of
|
||||
// the computation.
|
||||
//
|
||||
// SquadDesc is a host+device constexpr-compatible class that allows describing
|
||||
// the warp-specialized layout of a kernel.
|
||||
//
|
||||
// SquadDesc is constexpr-compatible and can be created on host and device.
|
||||
struct SquadDesc
|
||||
{
|
||||
int mSquadIdx = -1;
|
||||
int mWarpCount = -1;
|
||||
|
||||
_CCCL_HIDE_FROM_ABI constexpr SquadDesc() = default;
|
||||
_CCCL_HOST_DEVICE_API constexpr SquadDesc(int squadIdx, int warpCount) noexcept
|
||||
: mSquadIdx(squadIdx)
|
||||
, mWarpCount(warpCount)
|
||||
{}
|
||||
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr int warpCount() const noexcept
|
||||
{
|
||||
return mWarpCount;
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr int threadCount() const noexcept
|
||||
{
|
||||
return 32 * warpCount();
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator==(SquadDesc lhs, SquadDesc rhs) noexcept
|
||||
{
|
||||
return lhs.mSquadIdx == rhs.mSquadIdx;
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator!=(SquadDesc lhs, SquadDesc rhs) noexcept
|
||||
{
|
||||
return lhs.mSquadIdx != rhs.mSquadIdx;
|
||||
}
|
||||
};
|
||||
// squadCountThreads
|
||||
//
|
||||
// Utility function to count the number of threads in an array of squad
|
||||
// descriptors. It is used to launch a kernel with the correct number of
|
||||
// threads.
|
||||
template <int numSquads>
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr int squadCountThreads(const SquadDesc (&squads)[numSquads]) noexcept
|
||||
{
|
||||
int sumThreads = 0;
|
||||
for (int gi = 0; gi < numSquads; ++gi)
|
||||
{
|
||||
sumThreads += squads[gi].threadCount();
|
||||
}
|
||||
return sumThreads;
|
||||
}
|
||||
} // namespace detail::warpspeed
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,144 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cub/detail/warpspeed/constant_assert.cuh>
|
||||
#include <cub/detail/warpspeed/special_registers.cuh>
|
||||
|
||||
#include <cuda/__ptx/instructions/mbarrier_init.h>
|
||||
#include <cuda/std/cstdint>
|
||||
|
||||
#include <nv/target>
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail::warpspeed
|
||||
{
|
||||
// SkipSync is a tag type that is used to indicate that a SyncHandler.blockInit
|
||||
// should forgo syncing.
|
||||
struct SkipSync
|
||||
{};
|
||||
|
||||
struct SyncHandler
|
||||
{
|
||||
// reducing these values to the actually used number of resources and phases does not improve performance
|
||||
static constexpr int mMaxNumResources = 10;
|
||||
static constexpr int mMaxNumPhases = 4;
|
||||
|
||||
// Whether barriers have been initialized.
|
||||
bool mHasInitialized = false;
|
||||
|
||||
// Arrays of barrier locations, number of stages, number of owning threads.
|
||||
int mNextResourceHandle = 0;
|
||||
int mNumStages[mMaxNumResources]{};
|
||||
int mNumPhases[mMaxNumResources]{};
|
||||
int mNumOwningThreads[mMaxNumResources][mMaxNumPhases]{};
|
||||
::cuda::std::uint64_t* mPtrBar[mMaxNumResources][mMaxNumPhases]{};
|
||||
|
||||
constexpr SyncHandler() = default;
|
||||
|
||||
// we need constant destruction for the host side single stage SMEM amount, which is only possible in C++20
|
||||
#if _CCCL_STD_VER >= 2020
|
||||
_CCCL_HOST_DEVICE_API constexpr ~SyncHandler()
|
||||
{
|
||||
_WS_CONSTANT_ASSERT(mHasInitialized, "SyncHandler must have been initialized at end of kernel.");
|
||||
}
|
||||
#endif // _CCCL_STD_VER >= 2020
|
||||
|
||||
// SyncHandler is a non-copyable, non-movable type. It must be passed by
|
||||
// (mutable) reference to be useful.
|
||||
SyncHandler(const SyncHandler&) = delete; // Delete copy constructor
|
||||
SyncHandler(SyncHandler&&) = delete; // Delete move constructor
|
||||
SyncHandler& operator=(const SyncHandler&) = delete; // Delete copy assignment
|
||||
SyncHandler& operator=(const SyncHandler&&) = delete; // Delete move assignment
|
||||
|
||||
// registerResource and registerPhase can be called on host and device.
|
||||
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr int registerResource(int numStages)
|
||||
{
|
||||
_WS_CONSTANT_ASSERT(!mHasInitialized, "Cannot register resource after SyncHandler has been initialized.");
|
||||
// Avoid exceeding the max number of stages
|
||||
_WS_CONSTANT_ASSERT(mNextResourceHandle < mMaxNumResources, "Cannot register more than 10 resources.");
|
||||
|
||||
// Get a handle
|
||||
int handle = mNextResourceHandle;
|
||||
mNextResourceHandle++;
|
||||
// Set the number of stages
|
||||
mNumStages[handle] = numStages;
|
||||
|
||||
return handle;
|
||||
}
|
||||
|
||||
_CCCL_HOST_DEVICE_API void constexpr registerPhase(int resourceHandle, int numOwningThreads, uint64_t* ptrBar)
|
||||
{
|
||||
_WS_CONSTANT_ASSERT(!mHasInitialized, "Cannot register phase after SyncHandler has been initialized.");
|
||||
_WS_CONSTANT_ASSERT(resourceHandle < mNextResourceHandle, "Invalid resource handle.");
|
||||
|
||||
// Get phase index:
|
||||
int curPhase = mNumPhases[resourceHandle];
|
||||
_WS_CONSTANT_ASSERT(curPhase < mMaxNumPhases, "Cannot register more phases than maximum.");
|
||||
|
||||
mNumOwningThreads[resourceHandle][curPhase] = numOwningThreads;
|
||||
mPtrBar[resourceHandle][curPhase] = ptrBar;
|
||||
|
||||
mNumPhases[resourceHandle]++;
|
||||
}
|
||||
|
||||
// clusterInitSync can only be called on device.
|
||||
template <int NumThreads>
|
||||
_CCCL_DEVICE_API void clusterInitSync(SpecialRegisters sr, SkipSync)
|
||||
{
|
||||
_WS_CONSTANT_ASSERT(!mHasInitialized, "Cannot initialize SyncHandler twice.");
|
||||
mHasInitialized = true;
|
||||
|
||||
// All warps iterate through all resources and phases. Since all array indices have to be statically resolved by the
|
||||
// SROA optimization to avoid spilling to local memory, we cannot split the iteration among warps etc.
|
||||
for (int ri = 0; ri < mMaxNumResources; ri++)
|
||||
{
|
||||
if (ri >= mNextResourceHandle)
|
||||
{
|
||||
break;
|
||||
}
|
||||
const int resNumPhases = mNumPhases[ri];
|
||||
const int numStages = mNumStages[ri];
|
||||
|
||||
for (int pi = 0; pi < mMaxNumPhases; pi++)
|
||||
{
|
||||
if (pi >= resNumPhases)
|
||||
{
|
||||
break;
|
||||
}
|
||||
|
||||
uint64_t* ptrBar = mPtrBar[ri][pi];
|
||||
int numOwningThreads = mNumOwningThreads[ri][pi];
|
||||
// use block strided iteration to vectorize setup of barriers
|
||||
for (int si = static_cast<int>(sr.threadIdxX); si < numStages; si += NumThreads)
|
||||
{
|
||||
::cuda::ptx::mbarrier_init(&ptrBar[si], numOwningThreads);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <int NumThreads>
|
||||
_CCCL_DEVICE_API void clusterInitSync(SpecialRegisters sr)
|
||||
{
|
||||
NV_IF_TARGET(NV_PROVIDES_SM_90, ({
|
||||
clusterInitSync<NumThreads>(sr, SkipSync{});
|
||||
__cluster_barrier_arrive_relaxed();
|
||||
__cluster_barrier_wait();
|
||||
}))
|
||||
}
|
||||
};
|
||||
} // namespace detail::warpspeed
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,33 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
CUB_NAMESPACE_BEGIN
|
||||
|
||||
namespace detail::warpspeed
|
||||
{
|
||||
enum class Stages : int
|
||||
{
|
||||
};
|
||||
enum class Elems : int
|
||||
{
|
||||
};
|
||||
enum class Warps : int
|
||||
{
|
||||
};
|
||||
enum class Align : int
|
||||
{
|
||||
};
|
||||
} // namespace detail::warpspeed
|
||||
|
||||
CUB_NAMESPACE_END
|
||||
@@ -1,27 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
|
||||
#pragma once
|
||||
|
||||
#include <cub/config.cuh>
|
||||
|
||||
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
|
||||
# pragma GCC system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
|
||||
# pragma clang system_header
|
||||
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
|
||||
# pragma system_header
|
||||
#endif // no system header
|
||||
|
||||
#include <cub/detail/warpspeed/allocators/smem_allocator.cuh>
|
||||
#include <cub/detail/warpspeed/constant_assert.cuh>
|
||||
#include <cub/detail/warpspeed/make_warp_uniform.cuh>
|
||||
#include <cub/detail/warpspeed/resource/smem_phase.cuh>
|
||||
#include <cub/detail/warpspeed/resource/smem_ref.cuh>
|
||||
#include <cub/detail/warpspeed/resource/smem_resource.cuh>
|
||||
#include <cub/detail/warpspeed/resource/smem_resource_raw.cuh>
|
||||
#include <cub/detail/warpspeed/resource/smem_stage.cuh>
|
||||
#include <cub/detail/warpspeed/special_registers.cuh>
|
||||
#include <cub/detail/warpspeed/squad/squad.cuh>
|
||||
#include <cub/detail/warpspeed/squad/squad_desc.cuh>
|
||||
#include <cub/detail/warpspeed/sync_handler.cuh>
|
||||
#include <cub/detail/warpspeed/values.cuh>
|
||||
Reference in New Issue
Block a user