56fd68e7dd
CCCL (CUDA C++ Core Libraries) provides: - CUB: device/block/warp-level GPU primitives (reduce, scan, sort, topk) - Thrust: high-level parallel algorithms (transform_reduce, sort, scan) - libcudacxx: CUDA C++ standard library (atomics, barriers, memory) - cudax: experimental features (memory resources, allocators) - Tuning policies: per-SM hardware-specific algorithm parameters Competition optimization vectors mapped to CCCL: - Output TPS (83% weight): warp_reduce, block_reduce, device_topk - Input TPS (14% weight): device_scan, block_load, prefetch - Cache TPS (3% weight): prefix caching strategy patterns - Memory (0.9 util): pooled/cached/buddy allocators Source: https://github.com/NVIDIA/cccl (shallow clone, HEAD only) License: Apache-2.0
457 lines
12 KiB
Plaintext
457 lines
12 KiB
Plaintext
// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
// SPDX-License-Identifier: BSD-3-Clause
|
|
/**
|
|
* @file
|
|
* Simple binary operator functor types
|
|
*/
|
|
|
|
/******************************************************************************
|
|
* Simple functor operators
|
|
******************************************************************************/
|
|
|
|
#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/thread/thread_operators.cuh>
|
|
|
|
#include <cuda/__functional/maximum.h>
|
|
#include <cuda/__functional/minimum.h>
|
|
#include <cuda/std/__bit/bit_cast.h>
|
|
#include <cuda/std/__functional/operations.h>
|
|
#include <cuda/std/__type_traits/always_false.h>
|
|
#include <cuda/std/__type_traits/common_type.h>
|
|
#include <cuda/std/__type_traits/is_same.h>
|
|
#include <cuda/std/cstdint>
|
|
|
|
CUB_NAMESPACE_BEGIN
|
|
|
|
#ifndef _CCCL_DOXYGEN_INVOKED // Do not document
|
|
|
|
/***********************************************************************************************************************
|
|
* SIMD operators
|
|
**********************************************************************************************************************/
|
|
|
|
namespace detail
|
|
{
|
|
template <typename T>
|
|
extern _CCCL_HOST_DEVICE T simd_operation_is_not_supported_before_sm80();
|
|
|
|
template <typename T>
|
|
extern _CCCL_HOST_DEVICE T simd_operation_is_not_supported_before_sm53();
|
|
|
|
template <typename T>
|
|
struct SimdMin
|
|
{
|
|
static_assert(::cuda::std::__always_false_v<T>, "Unsupported specialization");
|
|
};
|
|
|
|
template <>
|
|
struct SimdMin<short2>
|
|
{
|
|
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE short2 operator()(short2 a, short2 b) const
|
|
{
|
|
NV_IF_ELSE_TARGET(
|
|
NV_PROVIDES_SM_90,
|
|
(auto a1 = ::cuda::std::bit_cast<uint32_t>(a); //
|
|
auto b1 = ::cuda::std::bit_cast<uint32_t>(b);
|
|
return ::cuda::std::bit_cast<short2>(::__vmins2(a1, b1));),
|
|
(return short2{::cuda::minimum<>{}(a.x, b.x), ::cuda::minimum<>{}(a.y, b.y)};))
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct SimdMin<ushort2>
|
|
{
|
|
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ushort2 operator()(ushort2 a, ushort2 b) const
|
|
{
|
|
NV_IF_ELSE_TARGET(
|
|
NV_PROVIDES_SM_90,
|
|
(auto a1 = ::cuda::std::bit_cast<uint32_t>(a); //
|
|
auto b1 = ::cuda::std::bit_cast<uint32_t>(b);
|
|
return ::cuda::std::bit_cast<ushort2>(::__vminu2(a1, b1));),
|
|
(return ushort2{::cuda::minimum<>{}(a.x, b.x), ::cuda::minimum<>{}(a.y, b.y)};))
|
|
}
|
|
};
|
|
|
|
# if _CCCL_HAS_NVFP16()
|
|
|
|
template <>
|
|
struct SimdMin<__half2>
|
|
{
|
|
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE __half2 operator()(__half2 a, __half2 b) const
|
|
{
|
|
return ::__hmin2(a, b);
|
|
}
|
|
};
|
|
|
|
# endif // _CCCL_HAS_NVFP16()
|
|
|
|
# if _CCCL_HAS_NVBF16()
|
|
|
|
template <>
|
|
struct SimdMin<__nv_bfloat162>
|
|
{
|
|
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE __nv_bfloat162 operator()(__nv_bfloat162 a, __nv_bfloat162 b) const
|
|
{
|
|
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_80,
|
|
(return ::__hmin2(a, b);),
|
|
(return simd_operation_is_not_supported_before_sm80<__nv_bfloat162>();))
|
|
}
|
|
};
|
|
|
|
# endif // _CCCL_HAS_NVBF16()
|
|
|
|
//----------------------------------------------------------------------------------------------------------------------
|
|
|
|
template <typename T>
|
|
struct SimdMax
|
|
{
|
|
static_assert(::cuda::std::__always_false_v<T>, "Unsupported specialization");
|
|
};
|
|
|
|
template <>
|
|
struct SimdMax<short2>
|
|
{
|
|
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE short2 operator()(short2 a, short2 b) const
|
|
{
|
|
NV_IF_ELSE_TARGET(
|
|
NV_PROVIDES_SM_90,
|
|
(auto a1 = ::cuda::std::bit_cast<uint32_t>(a); //
|
|
auto b1 = ::cuda::std::bit_cast<uint32_t>(b);
|
|
return ::cuda::std::bit_cast<short2>(::__vmaxs2(a1, b1));),
|
|
(return short2{::cuda::maximum<>{}(a.x, b.x), ::cuda::maximum<>{}(a.y, b.y)};))
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct SimdMax<ushort2>
|
|
{
|
|
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ushort2 operator()(ushort2 a, ushort2 b) const
|
|
{
|
|
NV_IF_ELSE_TARGET(
|
|
NV_PROVIDES_SM_90,
|
|
(auto a1 = ::cuda::std::bit_cast<uint32_t>(a); //
|
|
auto b1 = ::cuda::std::bit_cast<uint32_t>(b);
|
|
return ::cuda::std::bit_cast<ushort2>(::__vmaxu2(a1, b1));),
|
|
(return ushort2{::cuda::maximum<>{}(a.x, b.x), ::cuda::maximum<>{}(a.y, b.y)};))
|
|
}
|
|
};
|
|
|
|
# if _CCCL_HAS_NVFP16()
|
|
|
|
template <>
|
|
struct SimdMax<__half2>
|
|
{
|
|
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE __half2 operator()(__half2 a, __half2 b) const
|
|
{
|
|
return ::__hmax2(a, b);
|
|
}
|
|
};
|
|
|
|
# endif // _CCCL_HAS_NVFP16()
|
|
|
|
# if _CCCL_HAS_NVBF16()
|
|
|
|
template <>
|
|
struct SimdMax<__nv_bfloat162>
|
|
{
|
|
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE __nv_bfloat162 operator()(__nv_bfloat162 a, __nv_bfloat162 b) const
|
|
{
|
|
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_80,
|
|
(return ::__hmax2(a, b);), //
|
|
(return simd_operation_is_not_supported_before_sm80<__nv_bfloat162>();))
|
|
}
|
|
};
|
|
|
|
# endif // _CCCL_HAS_NVBF16()
|
|
|
|
//----------------------------------------------------------------------------------------------------------------------
|
|
|
|
template <typename T>
|
|
struct SimdSum
|
|
{
|
|
static_assert(::cuda::std::__always_false_v<T>, "Unsupported specialization");
|
|
};
|
|
|
|
# if _CCCL_HAS_NVFP16()
|
|
|
|
template <>
|
|
struct SimdSum<__half2>
|
|
{
|
|
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE __half2 operator()(__half2 a, __half2 b) const
|
|
{
|
|
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_53,
|
|
(return ::__hadd2(a, b);), //
|
|
(return simd_operation_is_not_supported_before_sm53<__half2>();))
|
|
}
|
|
};
|
|
|
|
# endif // _CCCL_HAS_NVFP16()
|
|
|
|
# if _CCCL_HAS_NVBF16()
|
|
|
|
template <>
|
|
struct SimdSum<__nv_bfloat162>
|
|
{
|
|
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE __nv_bfloat162 operator()(__nv_bfloat162 a, __nv_bfloat162 b) const
|
|
{
|
|
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_80,
|
|
(return ::__hadd2(a, b);), //
|
|
(return simd_operation_is_not_supported_before_sm80<__nv_bfloat162>();))
|
|
}
|
|
};
|
|
|
|
# endif // _CCCL_HAS_NVBF16()
|
|
|
|
template <>
|
|
struct SimdSum<ushort2>
|
|
{
|
|
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ushort2 operator()(ushort2 a, ushort2 b) const
|
|
{
|
|
// clang-format off
|
|
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_90,
|
|
(uint32_t ret;
|
|
asm("add.u16x2 %0, %1, %2;"
|
|
: "=r"(ret)
|
|
: "r"(::cuda::std::bit_cast<uint32_t>(a)),
|
|
"r"(::cuda::std::bit_cast<uint32_t>(b)));
|
|
return ::cuda::std::bit_cast<ushort2>(ret);),
|
|
(return ushort2{static_cast<uint16_t>(a.x + b.x), static_cast<uint16_t>(a.y + b.y)};));
|
|
// clang-format on
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct SimdSum<short2>
|
|
{
|
|
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE short2 operator()(short2 a, short2 b) const
|
|
{
|
|
// clang-format off
|
|
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_90,
|
|
(uint32_t ret;
|
|
asm("add.s16x2 %0, %1, %2;"
|
|
: "=r"(ret)
|
|
: "r"(::cuda::std::bit_cast<uint32_t>(a)),
|
|
"r"(::cuda::std::bit_cast<uint32_t>(b)));
|
|
return ::cuda::std::bit_cast<short2>(ret);),
|
|
(return short2{static_cast<int16_t>(a.x + b.x), static_cast<int16_t>(a.y + b.y)};));
|
|
// clang-format on
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct SimdSum<float2>
|
|
{
|
|
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE float2 operator()(float2 a, float2 b) const
|
|
{
|
|
// clang-format off
|
|
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_100,
|
|
(uint64_t ret;
|
|
asm("add.f32x2 %0, %1, %2;"
|
|
: "=l"(ret)
|
|
: "l"(::cuda::std::bit_cast<uint64_t>(a)),
|
|
"l"(::cuda::std::bit_cast<uint64_t>(b)));
|
|
return ::cuda::std::bit_cast<float2>(ret);),
|
|
(return float2{a.x + b.x, a.y + b.y};))
|
|
// clang-format on
|
|
}
|
|
};
|
|
|
|
//----------------------------------------------------------------------------------------------------------------------
|
|
|
|
template <typename T>
|
|
struct SimdMul
|
|
{
|
|
static_assert(::cuda::std::__always_false_v<T>, "Unsupported specialization");
|
|
};
|
|
|
|
# if _CCCL_HAS_NVFP16()
|
|
|
|
template <>
|
|
struct SimdMul<__half2>
|
|
{
|
|
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE __half2 operator()(__half2 a, __half2 b) const
|
|
{
|
|
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_53,
|
|
(return ::__hmul2(a, b);), //
|
|
(return simd_operation_is_not_supported_before_sm53<__half2>();))
|
|
}
|
|
};
|
|
|
|
# endif // _CCCL_HAS_NVFP16()
|
|
|
|
# if _CCCL_HAS_NVBF16()
|
|
|
|
template <>
|
|
struct SimdMul<__nv_bfloat162>
|
|
{
|
|
[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE __nv_bfloat162 operator()(__nv_bfloat162 a, __nv_bfloat162 b) const
|
|
{
|
|
NV_IF_ELSE_TARGET(NV_PROVIDES_SM_80,
|
|
(return ::__hmul2(a, b);),
|
|
(return simd_operation_is_not_supported_before_sm80<__nv_bfloat162>();))
|
|
}
|
|
};
|
|
|
|
# endif // _CCCL_HAS_NVBF16()
|
|
|
|
//----------------------------------------------------------------------------------------------------------------------
|
|
|
|
template <typename ReductionOp>
|
|
inline constexpr bool is_simd_operator_v = false;
|
|
|
|
template <typename T>
|
|
inline constexpr bool is_simd_operator_v<SimdSum<T>> = true;
|
|
|
|
template <typename T>
|
|
inline constexpr bool is_simd_operator_v<SimdMul<T>> = true;
|
|
|
|
template <typename T>
|
|
inline constexpr bool is_simd_operator_v<SimdMin<T>> = true;
|
|
|
|
template <typename T>
|
|
inline constexpr bool is_simd_operator_v<SimdMax<T>> = true;
|
|
|
|
//----------------------------------------------------------------------------------------------------------------------
|
|
// SIMD type
|
|
|
|
template <typename T>
|
|
struct VectorTypeX2
|
|
{
|
|
static_assert(::cuda::std::__always_false_v<T>, "Unsupported specialization");
|
|
};
|
|
|
|
template <>
|
|
struct VectorTypeX2<int16_t>
|
|
{
|
|
using type = short2;
|
|
};
|
|
|
|
template <>
|
|
struct VectorTypeX2<uint16_t>
|
|
{
|
|
using type = ushort2;
|
|
};
|
|
|
|
# if _CCCL_HAS_NVFP16()
|
|
|
|
template <>
|
|
struct VectorTypeX2<__half>
|
|
{
|
|
using type = __half2;
|
|
};
|
|
|
|
# endif // _CCCL_HAS_NVFP16()
|
|
|
|
# if _CCCL_HAS_NVBF16()
|
|
|
|
template <>
|
|
struct VectorTypeX2<__nv_bfloat16>
|
|
{
|
|
using type = __nv_bfloat162;
|
|
};
|
|
|
|
# endif // _CCCL_HAS_NVBF16()
|
|
|
|
template <typename T>
|
|
using vector_type_x2_t = typename VectorTypeX2<T>::type;
|
|
|
|
//----------------------------------------------------------------------------------------------------------------------
|
|
// Predefined CUDA operators to SIMD
|
|
|
|
template <typename ReduceOp, typename T>
|
|
struct CudaOperatorToSimdX2
|
|
{
|
|
static_assert(::cuda::std::__always_false_v<T>, "Unsupported specialization");
|
|
};
|
|
|
|
template <typename T>
|
|
struct CudaOperatorToSimdX2<::cuda::minimum<>, T>
|
|
{
|
|
using type = SimdMin<vector_type_x2_t<T>>;
|
|
};
|
|
|
|
template <typename T>
|
|
struct CudaOperatorToSimdX2<::cuda::minimum<T>, T>
|
|
{
|
|
using type = SimdMin<vector_type_x2_t<T>>;
|
|
};
|
|
|
|
template <typename T>
|
|
struct CudaOperatorToSimdX2<::cuda::maximum<>, T>
|
|
{
|
|
using type = SimdMax<vector_type_x2_t<T>>;
|
|
};
|
|
|
|
template <typename T>
|
|
struct CudaOperatorToSimdX2<::cuda::maximum<T>, T>
|
|
{
|
|
using type = SimdMax<vector_type_x2_t<T>>;
|
|
};
|
|
|
|
template <typename T>
|
|
struct CudaOperatorToSimdX2<::cuda::std::plus<>, T>
|
|
{
|
|
using type = SimdSum<vector_type_x2_t<T>>;
|
|
};
|
|
|
|
template <typename T>
|
|
struct CudaOperatorToSimdX2<::cuda::std::plus<T>, T>
|
|
{
|
|
using type = SimdSum<vector_type_x2_t<T>>;
|
|
};
|
|
|
|
template <typename T>
|
|
struct CudaOperatorToSimdX2<::cuda::std::multiplies<>, T>
|
|
{
|
|
using type = SimdMul<vector_type_x2_t<T>>;
|
|
};
|
|
|
|
template <typename T>
|
|
struct CudaOperatorToSimdX2<::cuda::std::multiplies<T>, T>
|
|
{
|
|
using type = SimdMul<vector_type_x2_t<T>>;
|
|
};
|
|
|
|
template <typename ReduceOp, typename T>
|
|
using cuda_operator_to_simd_x2_t = typename CudaOperatorToSimdX2<ReduceOp, T>::type;
|
|
|
|
//----------------------------------------------------------------------------------------------------------------------
|
|
|
|
template <typename T, typename Op>
|
|
_CCCL_DEVICE _CCCL_FORCEINLINE auto try_simd_operator(Op op)
|
|
{
|
|
using ::cuda::std::is_same_v;
|
|
constexpr bool is_supported_vector_type = is_any_short2_v<T> || is_bfloat162_v<T> || is_half2_v<T>;
|
|
if constexpr (is_cuda_std_plus_v<Op> && (is_same_v<T, float2> || is_supported_vector_type))
|
|
{
|
|
return SimdSum<T>{};
|
|
}
|
|
else if constexpr (is_cuda_minimum_v<Op> && is_supported_vector_type)
|
|
{
|
|
return SimdMin<T>{};
|
|
}
|
|
else if constexpr (is_cuda_maximum_v<Op> && is_supported_vector_type)
|
|
{
|
|
return SimdMax<T>{};
|
|
}
|
|
else
|
|
{
|
|
return op;
|
|
}
|
|
}
|
|
} // namespace detail
|
|
|
|
#endif // !_CCCL_DOXYGEN_INVOKED
|
|
|
|
CUB_NAMESPACE_END
|