[CUB] Add fast constant division utility

cub/cub/detail/fast_modulo_division.cuh

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: Copyright (c) 2011-2024, NVIDIA CORPORATION. All rights reserved.
+// SPDX-FileCopyrightText: Copyright (c) 2011-2026, NVIDIA CORPORATION. All rights reserved.
 // SPDX-License-Identifier: BSD-3
 
 #pragma once
@@ -18,6 +18,7 @@
 
 #include <cuda/__cmath/ceil_div.h>
 #include <cuda/__cmath/pow2.h>
+#include <cuda/std/__bit/countl.h>
 #include <cuda/std/__bit/integral.h>
 #include <cuda/std/__type_traits/conditional.h>
 #include <cuda/std/__type_traits/enable_if.h>
@@ -111,6 +112,163 @@ multiply_extract_higher_bits(T value, R multiplier)
   // clang-format on
 }
 
+/***********************************************************************************************************************
+ * Fast division by a precomputed unsigned constant
+ *
+ * A divisor of zero selects the identity operation. This provides a safe default state and lets callers use zero to
+ * represent an inactive division without introducing undefined behavior.
+ **********************************************************************************************************************/
+
+template <typename UInt>
+class fast_divide_by_constant
+{
+  static_assert(::cuda::std::is_unsigned_v<UInt>, "fast_divide_by_constant requires an unsigned integer type");
+  static_assert(sizeof(UInt) == 4 || sizeof(UInt) == 8, "fast_divide_by_constant supports 32- or 64-bit integers");
+
+  static constexpr int bits = static_cast<int>(sizeof(UInt) * CHAR_BIT);
+
+  enum class mode : unsigned char
+  {
+    identity,
+    shift,
+    multiply_shift,
+    hardware
+  };
+
+  [[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE static int ceil_log2(UInt divisor) noexcept
+  {
+    return divisor <= UInt{1} ? 0 : bits - ::cuda::std::countl_zero(divisor - UInt{1});
+  }
+
+  [[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE static UInt multiply_high(UInt lhs, UInt rhs) noexcept
+  {
+    if constexpr (sizeof(UInt) == 4)
+    {
+      return static_cast<UInt>(
+        (static_cast<::cuda::std::uint64_t>(lhs) * static_cast<::cuda::std::uint64_t>(rhs)) >> bits);
+    }
+    else
+    {
+#if _CCCL_HAS_INT128()
+      NV_IF_ELSE_TARGET(
+        NV_IS_DEVICE,
+        (return static_cast<UInt>(
+                  __umul64hi(static_cast<unsigned long long>(lhs), static_cast<unsigned long long>(rhs)));),
+        (return static_cast<UInt>((static_cast<__uint128_t>(lhs) * static_cast<__uint128_t>(rhs)) >> bits);));
+#else // ^^^ _CCCL_HAS_INT128() ^^^ / vvv !_CCCL_HAS_INT128() vvv
+      NV_IF_ELSE_TARGET(
+        NV_IS_DEVICE,
+        (return static_cast<UInt>(
+                  __umul64hi(static_cast<unsigned long long>(lhs), static_cast<unsigned long long>(rhs)));),
+        ({
+          const ::cuda::std::uint64_t lhs_low   = static_cast<::cuda::std::uint32_t>(lhs);
+          const ::cuda::std::uint64_t lhs_high  = lhs >> 32;
+          const ::cuda::std::uint64_t rhs_low   = static_cast<::cuda::std::uint32_t>(rhs);
+          const ::cuda::std::uint64_t rhs_high  = rhs >> 32;
+          const ::cuda::std::uint64_t low_low   = lhs_low * rhs_low;
+          const ::cuda::std::uint64_t low_high  = lhs_low * rhs_high;
+          const ::cuda::std::uint64_t high_low  = lhs_high * rhs_low;
+          const ::cuda::std::uint64_t high_high = lhs_high * rhs_high;
+          const ::cuda::std::uint64_t middle    = (low_low >> 32) + static_cast<::cuda::std::uint32_t>(low_high)
+                                                + static_cast<::cuda::std::uint32_t>(high_low);
+          return static_cast<UInt>(high_high + (low_high >> 32) + (high_low >> 32) + (middle >> 32));
+        }));
+#endif // !_CCCL_HAS_INT128()
+    }
+  }
+
+public:
+  _CCCL_HOST_DEVICE constexpr fast_divide_by_constant() noexcept {}
+
+  _CCCL_HOST_DEVICE explicit fast_divide_by_constant(UInt divisor) noexcept
+  {
+    init(divisor);
+  }
+
+  _CCCL_HOST_DEVICE _CCCL_FORCEINLINE void init(UInt divisor) noexcept
+  {
+    if (divisor <= UInt{1})
+    {
+      magic_ = UInt{0};
+      shift_ = 0;
+      mode_  = mode::identity;
+      return;
+    }
+    if ((divisor & (divisor - UInt{1})) == UInt{0})
+    {
+      magic_ = UInt{0};
+      shift_ = static_cast<unsigned char>(ceil_log2(divisor));
+      mode_  = mode::shift;
+      return;
+    }
+
+    const int log2_divisor = ceil_log2(divisor);
+    if (log2_divisor == bits)
+    {
+      magic_ = divisor;
+      shift_ = 0;
+      mode_  = mode::hardware;
+      return;
+    }
+    if constexpr (sizeof(UInt) == 8)
+    {
+#if _CCCL_HAS_INT128()
+      const __uint128_t numerator   = static_cast<__uint128_t>(1) << (bits + log2_divisor);
+      const __uint128_t denominator = static_cast<__uint128_t>(divisor);
+      magic_                        = static_cast<UInt>((numerator + denominator - 1) / denominator);
+#else // ^^^ _CCCL_HAS_INT128() ^^^ / vvv !_CCCL_HAS_INT128() vvv
+      UInt quotient  = 0;
+      UInt remainder = 0;
+      for (int bit = bits + log2_divisor; bit >= 0; --bit)
+      {
+        UInt next_remainder     = (remainder << 1) | (bit == bits + log2_divisor ? UInt{1} : UInt{0});
+        const bool carry        = (remainder >> (bits - 1)) != 0;
+        const UInt quotient_bit = (carry || next_remainder >= divisor) ? UInt{1} : UInt{0};
+        if (quotient_bit != 0)
+        {
+          next_remainder -= divisor;
+        }
+        remainder = next_remainder;
+        quotient  = (quotient << 1) | quotient_bit;
+      }
+      magic_ = quotient + (remainder != 0 ? UInt{1} : UInt{0});
+#endif // !_CCCL_HAS_INT128()
+    }
+    else
+    {
+      const ::cuda::std::uint64_t numerator   = ::cuda::std::uint64_t{1} << (bits + log2_divisor);
+      const ::cuda::std::uint64_t denominator = static_cast<::cuda::std::uint64_t>(divisor);
+      magic_                                  = static_cast<UInt>((numerator + denominator - 1) / denominator);
+    }
+    shift_ = static_cast<unsigned char>(log2_divisor);
+    mode_  = mode::multiply_shift;
+  }
+
+  [[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE UInt divide(UInt numerator) const noexcept
+  {
+    if (mode_ == mode::identity)
+    {
+      return numerator;
+    }
+    if (mode_ == mode::shift)
+    {
+      return numerator >> shift_;
+    }
+    if (mode_ == mode::hardware)
+    {
+      return numerator / magic_;
+    }
+
+    const UInt high = multiply_high(magic_, numerator);
+    return (((numerator - high) >> 1) + high) >> (shift_ - 1);
+  }
+
+private:
+  UInt magic_          = UInt{0};
+  unsigned char shift_ = 0;
+  mode mode_           = mode::identity;
+};
+
 /***********************************************************************************************************************
  * Fast Modulo/Division based on Precomputation
  **********************************************************************************************************************/

cub/test/internal/catch2_test_fast_divide_by_constant.cu

@@ -0,0 +1,146 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+
+#include <cub/detail/fast_modulo_division.cuh>
+
+#include <thrust/detail/raw_pointer_cast.h>
+
+#include <cuda/std/array>
+#include <cuda/std/cstddef>
+#include <cuda/std/cstdint>
+#include <cuda/std/limits>
+
+#include "c2h/utility.h"
+#include "c2h/vector.h"
+#include "cub_test_macros.h"
+
+using uint_types = c2h::type_list<::cuda::std::uint32_t, ::cuda::std::uint64_t>;
+
+template <typename UInt>
+[[nodiscard]] UInt reference_divide(UInt numerator, UInt divisor)
+{
+  return divisor == UInt{0} ? numerator : numerator / divisor;
+}
+
+template <typename UInt>
+[[nodiscard]] constexpr auto divisors()
+{
+  constexpr int bits       = static_cast<int>(sizeof(UInt) * 8);
+  constexpr UInt max_value = ::cuda::std::numeric_limits<UInt>::max();
+  constexpr UInt high_bit  = UInt{1} << (bits - 1);
+
+  return ::cuda::std::array<UInt, 15>{
+    UInt{0},
+    UInt{1},
+    UInt{2},
+    UInt{4},
+    UInt{8},
+    high_bit,
+    UInt{3},
+    UInt{5},
+    UInt{7},
+    UInt{10},
+    UInt{31},
+    high_bit - UInt{1},
+    high_bit + UInt{1},
+    max_value - UInt{1},
+    max_value};
+}
+
+template <typename UInt>
+[[nodiscard]] c2h::host_vector<UInt> make_numerators()
+{
+  constexpr int bits       = static_cast<int>(sizeof(UInt) * 8);
+  constexpr UInt max_value = ::cuda::std::numeric_limits<UInt>::max();
+  constexpr UInt high_bit  = UInt{1} << (bits - 1);
+
+  c2h::host_vector<UInt> numerators{
+    UInt{0}, UInt{1}, UInt{2}, UInt{3}, high_bit - UInt{1}, high_bit, high_bit + UInt{1}, max_value - UInt{1}, max_value};
+
+  for (const UInt divisor : divisors<UInt>())
+  {
+    if (divisor != UInt{0})
+    {
+      numerators.push_back(divisor - UInt{1});
+      numerators.push_back(divisor);
+      if (divisor != max_value)
+      {
+        numerators.push_back(divisor + UInt{1});
+      }
+      const UInt largest_multiple = max_value - (max_value % divisor);
+      numerators.push_back(largest_multiple);
+      if (largest_multiple != UInt{0})
+      {
+        numerators.push_back(largest_multiple - UInt{1});
+      }
+    }
+  }
+
+  UInt value = static_cast<UInt>(0x9e3779b9U);
+  for (int i = 0; i < 4096; ++i)
+  {
+    value = value * static_cast<UInt>(6364136223846793005ULL) + static_cast<UInt>(1442695040888963407ULL);
+    numerators.push_back(value);
+  }
+  return numerators;
+}
+
+template <typename UInt>
+__global__ void
+fast_divide_by_constant_kernel(const UInt* numerators, UInt* quotients, ::cuda::std::size_t count, UInt divisor)
+{
+  const ::cuda::std::size_t index = static_cast<::cuda::std::size_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  if (index < count)
+  {
+    const cub::detail::fast_divide_by_constant<UInt> divider{divisor};
+    quotients[index] = divider.divide(numerators[index]);
+  }
+}
+
+CUB_TEST("fast_divide_by_constant agrees with division on the host", "[util][division]", CUB_SMALL, uint_types)
+{
+  using uint_t = c2h::get<0, TestType>;
+
+  const c2h::host_vector<uint_t> numerators = make_numerators<uint_t>();
+  cub::detail::fast_divide_by_constant<uint_t> divider;
+  for (const uint_t divisor : divisors<uint_t>())
+  {
+    divider.init(divisor);
+    CAPTURE(c2h::type_name<uint_t>(), divisor);
+    for (const uint_t numerator : numerators)
+    {
+      CAPTURE(numerator);
+      REQUIRE(divider.divide(numerator) == reference_divide(numerator, divisor));
+    }
+  }
+}
+
+CUB_TEST("fast_divide_by_constant agrees with division on the device", "[util][division]", CUB_SMALL, uint_types)
+{
+  using uint_t = c2h::get<0, TestType>;
+
+  const c2h::host_vector<uint_t> numerators = make_numerators<uint_t>();
+  const c2h::device_vector<uint_t> d_numerators(numerators);
+  c2h::device_vector<uint_t> d_quotients(numerators.size());
+
+  constexpr int block_threads = 256;
+  const int blocks            = static_cast<int>((numerators.size() + block_threads - 1) / block_threads);
+  for (const uint_t divisor : divisors<uint_t>())
+  {
+    fast_divide_by_constant_kernel<<<blocks, block_threads>>>(
+      thrust::raw_pointer_cast(d_numerators.data()),
+      thrust::raw_pointer_cast(d_quotients.data()),
+      numerators.size(),
+      divisor);
+    REQUIRE(cudaSuccess == cudaPeekAtLastError());
+    REQUIRE(cudaSuccess == cudaDeviceSynchronize());
+
+    const c2h::host_vector<uint_t> quotients = d_quotients;
+    CAPTURE(c2h::type_name<uint_t>(), divisor);
+    for (::cuda::std::size_t index = 0; index < numerators.size(); ++index)
+    {
+      CAPTURE(index, numerators[index]);
+      REQUIRE(quotients[index] == reference_divide(numerators[index], divisor));
+    }
+  }
+}