CUB - Add internal integer utils and tests (Split WarpReduce PR)

cub/cub/detail/integer_utils.cuh

@@ -0,0 +1,84 @@
+// 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/unsafe_bitcast.cuh>
+#include <cub/thread/thread_operators.cuh> // is_cuda_minimum_maximum_v
+
+#include <cuda/std/cmath> // isnan
+#include <cuda/std/limits> // numeric_limits
+#include <cuda/type_traits> // is_floating_point_v
+
+CUB_NAMESPACE_BEGIN
+
+/***********************************************************************************************************************
+ * Integer Utils
+ **********************************************************************************************************************/
+
+namespace detail
+{
+
+template <typename Input>
+_CCCL_DEVICE _CCCL_FORCEINLINE auto split_integer(Input input)
+{
+  using namespace _CUDA_VSTD;
+  static_assert(__cccl_is_integer_v<Input>);
+  constexpr auto half_bits = __num_bits_v<Input> / 2;
+  using unsigned_t         = make_unsigned_t<Input>;
+  using output_t           = __make_nbit_int_t<half_bits, is_signed_v<Input>>;
+  auto input1              = static_cast<unsigned_t>(input);
+  auto high                = static_cast<output_t>(input1 >> half_bits);
+  auto low                 = static_cast<output_t>(input1);
+  return array<output_t, 2>{high, low};
+}
+
+template <typename Input>
+_CCCL_DEVICE _CCCL_FORCEINLINE auto merge_integers(Input inputA, Input inputB)
+{
+  using namespace _CUDA_VSTD;
+  static_assert(__cccl_is_integer_v<Input>);
+  constexpr auto num_bits = __num_bits_v<Input>;
+  using unsigned_t        = __make_nbit_uint_t<num_bits>;
+  using unsigned_X2_t     = __make_nbit_uint_t<num_bits * 2>;
+  using output_t          = __make_nbit_int_t<num_bits * 2, is_signed_v<Input>>;
+  return static_cast<output_t>((static_cast<unsigned_X2_t>(inputA) << num_bits) | static_cast<unsigned_t>(inputB));
+}
+
+// When it is not possible to use native functionalities to compare floating-point values, we can convert them to
+// an integer representation that preserves the order.
+template <typename MinMaxOp, typename T>
+_CCCL_DEVICE _CCCL_FORCEINLINE auto floating_point_to_comparable_int(MinMaxOp, T value)
+{
+  using namespace _CUDA_VSTD;
+  static_assert(::cuda::is_floating_point_v<T>);
+  static_assert(is_cuda_minimum_maximum_v<MinMaxOp, T>);
+  using signed_t        = __make_nbit_int_t<__num_bits_v<T>, true>;
+  constexpr auto lowest = numeric_limits<signed_t>::lowest();
+  constexpr auto is_max = is_cuda_maximum_v<MinMaxOp, T>;
+  const auto nan        = is_max ? static_cast<T>(-numeric_limits<T>::quiet_NaN()) : numeric_limits<T>::quiet_NaN();
+  auto value1           = _CUDA_VSTD::isnan(value) ? nan : value;
+  auto value_int        = cub::detail::unsafe_bitcast<signed_t>(value1);
+  return static_cast<signed_t>(value_int < 0 ? lowest - value_int : value_int);
+}
+
+template <typename FloatingPointType, typename IntegerType>
+_CCCL_DEVICE _CCCL_FORCEINLINE auto comparable_int_to_floating_point(IntegerType value)
+{
+  static_assert(_CUDA_VSTD::__cccl_is_integer_v<IntegerType>);
+  constexpr auto lowest = _CUDA_VSTD::numeric_limits<IntegerType>::lowest();
+  auto value1           = static_cast<IntegerType>(value < 0 ? lowest - value : value);
+  return cub::detail::unsafe_bitcast<FloatingPointType>(value1);
+}
+
+} // namespace detail
+CUB_NAMESPACE_END

cub/test/internal/catch2_test_integer_utils.cu

@@ -0,0 +1,140 @@
+// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+#include <cub/detail/integer_utils.cuh>
+
+#include <c2h/catch2_test_helper.h>
+
+/***********************************************************************************************************************
+ * TEST CASES
+ **********************************************************************************************************************/
+
+using integral_types =
+  c2h::type_list<int16_t,
+                 uint16_t,
+                 int32_t,
+                 uint32_t,
+                 int64_t,
+                 uint64_t
+#if TEST_INT128()
+                 ,
+                 __int128_t,
+                 __uint128_t
+#endif
+                 >;
+
+using floating_point_types =
+  c2h::type_list<float,
+                 double
+#if TEST_HALF_T()
+                 ,
+                 __half
+#endif
+#if TEST_BF_T()
+                 ,
+                 __nv_bfloat16
+#endif
+                 >;
+
+using operator_types = c2h::type_list<cuda::minimum<>, cuda::maximum<>>;
+
+template <typename T>
+static __global__ void test_int_kernel(const T* input, int num_items)
+{
+  auto i = threadIdx.x + blockIdx.x * blockDim.x;
+  if (i < num_items)
+  {
+    auto value       = input[i];
+    auto [high, low] = cub::detail::split_integer(value);
+    static_assert(sizeof(high) == sizeof(T) / 2);
+    static_assert(sizeof(low) == sizeof(T) / 2);
+    static_assert(cuda::std::is_signed_v<T> == cuda::std::is_signed_v<decltype(high)>);
+    static_assert(cuda::std::is_signed_v<T> == cuda::std::is_signed_v<decltype(low)>);
+    auto result = cub::detail::merge_integers(high, low);
+    static_assert(sizeof(result) == sizeof(T));
+    static_assert(cuda::std::is_signed_v<T> == cuda::std::is_signed_v<decltype(result)>);
+    assert(value == result);
+  }
+}
+
+static __global__ void test_int_special_values_kernel()
+{
+  {
+    auto [high, low] = cub::detail::split_integer(0xAABBCCDD);
+    assert(high = 0xAABB);
+    assert(low = 0xCCDD);
+  }
+  {
+    auto [high, low] = cub::detail::split_integer(0xAABBCCDDEEFF1234);
+    assert(high = 0xAABBCCDD);
+    assert(low = 0xEEFF1234);
+  }
+}
+
+template <typename Operation, typename T>
+static __device__ void test_floating_point(Operation op, T valueA, T valueB)
+{
+  auto int_valueA = cub::detail::floating_point_to_comparable_int(op, valueA);
+  auto int_valueB = cub::detail::floating_point_to_comparable_int(op, valueB);
+  static_assert(sizeof(int_valueA) == sizeof(T));
+  static_assert(sizeof(int_valueB) == sizeof(T));
+  assert(op(valueA, valueB) == op(int_valueA, int_valueB));
+  auto result_valueA = cub::detail::comparable_int_to_floating_point<T>(int_valueA);
+  auto result_valueB = cub::detail::comparable_int_to_floating_point<T>(int_valueB);
+  static_assert(cuda::std::is_same_v<T, decltype(result_valueA)>);
+  static_assert(cuda::std::is_same_v<T, decltype(result_valueB)>);
+  assert(valueA == result_valueA);
+  assert(valueB == result_valueB);
+}
+
+template <typename Operation, typename T>
+static __global__ void test_float_kernel(Operation op, const T* input, int num_items)
+{
+  auto i = threadIdx.x + blockIdx.x * blockDim.x;
+  if (i < num_items - 1)
+  {
+    test_floating_point(op, input[i], input[i + 1]);
+  }
+}
+
+template <typename T, typename Operation>
+static __global__ void test_float_special_values_kernel(Operation op)
+{
+  test_floating_point(op, T{-0.0f}, T{0.0f});
+  test_floating_point(op, T{1.0f}, cuda::std::numeric_limits<T>::max());
+  test_floating_point(op, T{1.0f}, cuda::std::numeric_limits<T>::min());
+  test_floating_point(op, T{1.0f}, cuda::std::numeric_limits<T>::quiet_NaN());
+  test_floating_point(op, cuda::std::numeric_limits<T>::min(), cuda::std::numeric_limits<T>::max());
+  test_floating_point(op, cuda::std::numeric_limits<T>::max(), cuda::std::numeric_limits<T>::quiet_NaN());
+  test_floating_point(op, cuda::std::numeric_limits<T>::min(), cuda::std::numeric_limits<T>::quiet_NaN());
+  test_floating_point(op, cuda::std::numeric_limits<T>::quiet_NaN(), cuda::std::numeric_limits<T>::quiet_NaN());
+}
+
+C2H_TEST("Split/Merge Integers", "[Split/Merge][Random]", integral_types)
+{
+  using T              = c2h::get<0, TestType>;
+  const auto num_items = 1 << 16;
+  c2h::device_vector<T> d_in(num_items);
+  c2h::gen(C2H_SEED(1), d_in);
+  test_int_kernel<<<cuda::ceil_div(num_items, 256), 256>>>(thrust::raw_pointer_cast(d_in.data()), num_items);
+  REQUIRE(cudaSuccess == cudaPeekAtLastError());
+  REQUIRE(cudaSuccess == cudaDeviceSynchronize());
+  test_int_special_values_kernel<<<1, 1>>>();
+  REQUIRE(cudaSuccess == cudaPeekAtLastError());
+  REQUIRE(cudaSuccess == cudaDeviceSynchronize());
+}
+
+C2H_TEST(
+  "Compare Floating-point with Integers", "[CompareFloatingPointWithInt][Random]", operator_types, floating_point_types)
+{
+  using Op             = c2h::get<0, TestType>;
+  using T              = c2h::get<1, TestType>;
+  const auto num_items = 1 << 16;
+  c2h::device_vector<T> d_in(num_items);
+  c2h::gen(C2H_SEED(1), d_in);
+  test_float_kernel<<<cuda::ceil_div(num_items, 256), 256>>>(Op{}, thrust::raw_pointer_cast(d_in.data()), num_items);
+  REQUIRE(cudaSuccess == cudaPeekAtLastError());
+  REQUIRE(cudaSuccess == cudaDeviceSynchronize());
+  test_float_special_values_kernel<T><<<1, 1>>>(Op{});
+  REQUIRE(cudaSuccess == cudaPeekAtLastError());
+  REQUIRE(cudaSuccess == cudaDeviceSynchronize());
+}