FastIntDiv: fallback to normal division when values exceed 32-bit ran…

cpp/bench/prims/CMakeLists.txt

@@ -1,6 +1,6 @@
 # =============================================================================
 # cmake-format: off
-# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 # cmake-format: on
 # =============================================================================
@@ -69,7 +69,7 @@ endfunction()
 if(BUILD_PRIMS_BENCH)
   ConfigureBench(NAME CORE_BENCH PATH core/bitset.cu core/copy.cu core/memory_tracking.cu main.cpp)
 
-  ConfigureBench(NAME UTIL_BENCH PATH util/popc.cu main.cpp)
+  ConfigureBench(NAME UTIL_BENCH PATH util/fast_int_div.cu util/popc.cu main.cpp)
 
   ConfigureBench(
     NAME

cpp/bench/prims/util/fast_int_div.cu

@@ -0,0 +1,116 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <common/benchmark.hpp>
+
+#include <raft/util/fast_int_div.cuh>
+
+#include <rmm/device_buffer.hpp>
+#include <rmm/device_uvector.hpp>
+
+#include <random>
+#include <type_traits>
+#include <vector>
+
+namespace raft::bench::util {
+
+constexpr size_t kNumNumerators = 1000 * 1000;
+constexpr size_t kNumDivisors   = 100;
+
+/**
+ * A single kernel serves both variants: `DivisorT` is either `raft::util::FastIntDiv<IntT>`
+ * or a plain `IntT` (native division), and both support `operator/` and `operator%` against
+ * an `IntT` numerator.
+ */
+template <typename IntT, typename DivisorT>
+RAFT_KERNEL divmod_kernel(const IntT* numerators,
+                          const int64_t n_numerators,
+                          const DivisorT* divisors,
+                          const int64_t n_divisors,
+                          IntT* out)
+{
+  int64_t tid    = int64_t(blockIdx.x) * int64_t(blockDim.x) + int64_t(threadIdx.x);
+  int64_t stride = int64_t(gridDim.x) * int64_t(blockDim.x);
+  IntT acc       = 0;  // to prevent compiler from optimizing away the division ops
+  for (int64_t j = 0; j < n_divisors; ++j) {
+    DivisorT divisor = divisors[j];
+    for (int64_t i = tid; i < n_numerators; i += stride) {
+      IntT n = numerators[i];
+      acc ^= n / divisor;
+      acc ^= n % divisor;
+    }
+  }
+  out[tid] = acc;
+}
+
+template <typename IntT, bool UseFastIntDiv>
+struct fast_int_div_bench : public fixture {
+  using divisor_t = std::conditional_t<UseFastIntDiv, raft::util::FastIntDiv<IntT>, IntT>;
+
+  explicit fast_int_div_bench()
+    : d_numerators(kNumNumerators, stream),
+      d_divisors(size_t(kNumDivisors) * sizeof(divisor_t), stream),
+      out_d(size_t(kBlocks) * size_t(kThreads), stream)
+  {
+    std::mt19937_64 rng(42);
+    std::uniform_int_distribution<IntT> numerator_dist(std::numeric_limits<IntT>::min(),
+                                                       std::numeric_limits<IntT>::max());
+    // non-zero, non-neg divisors
+    std::uniform_int_distribution<IntT> divisor_dist(1, std::numeric_limits<IntT>::max());
+
+    std::vector<IntT> h_numerators(kNumNumerators);
+    for (auto& n : h_numerators) {
+      n = numerator_dist(rng);
+    }
+
+    std::vector<divisor_t> h_divisors;
+    h_divisors.reserve(kNumDivisors);
+    for (size_t i = 0; i < kNumDivisors; ++i) {
+      h_divisors.push_back(divisor_t(divisor_dist(rng)));
+    }
+
+    RAFT_CUDA_TRY(cudaMemcpyAsync(d_numerators.data(),
+                                  h_numerators.data(),
+                                  h_numerators.size() * sizeof(IntT),
+                                  cudaMemcpyHostToDevice,
+                                  stream));
+    RAFT_CUDA_TRY(cudaMemcpyAsync(d_divisors.data(),
+                                  h_divisors.data(),
+                                  h_divisors.size() * sizeof(divisor_t),
+                                  cudaMemcpyHostToDevice,
+                                  stream));
+    stream.synchronize();
+  }
+
+  void run_benchmark(::benchmark::State& state) override
+  {
+    const auto* divisors = static_cast<const divisor_t*>(d_divisors.data());
+    loop_on_state(state, [this, divisors]() {
+      divmod_kernel<IntT, divisor_t><<<kBlocks, kThreads, 0, stream>>>(
+        d_numerators.data(), kNumNumerators, divisors, kNumDivisors, out_d.data());
+      RAFT_CUDA_TRY(cudaPeekAtLastError());
+    });
+  }
+
+ private:
+  static constexpr int kThreads = 256;
+  static constexpr int kBlocks  = 1024;
+
+  rmm::device_uvector<IntT> d_numerators;
+  rmm::device_buffer d_divisors;
+  rmm::device_uvector<IntT> out_d;
+};
+
+using fast_int_div_i32 = fast_int_div_bench<int32_t, true>;
+using native_div_i32   = fast_int_div_bench<int32_t, false>;
+using fast_int_div_i64 = fast_int_div_bench<int64_t, true>;
+using native_div_i64   = fast_int_div_bench<int64_t, false>;
+
+RAFT_BENCH_REGISTER(fast_int_div_i32, "FastIntDiv/int32");
+RAFT_BENCH_REGISTER(native_div_i32, "NativeIntDiv/int32");
+RAFT_BENCH_REGISTER(fast_int_div_i64, "FastIntDiv/int64");
+RAFT_BENCH_REGISTER(native_div_i64, "NativeIntDiv/int64");
+
+}  // namespace raft::bench::util

cpp/include/raft/util/fast_int_div.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -10,18 +10,19 @@
 
 #include <stdint.h>
 
+#include <limits>
 #include <type_traits>
 
 namespace raft {
 namespace util {
 
+constexpr auto kInt32Min = std::numeric_limits<int32_t>::min();
+constexpr auto kInt32Max = std::numeric_limits<int32_t>::max();
+
 /**
  * @brief Perform fast integer division and modulo using a known divisor
  * From Hacker's Delight, Second Edition, Chapter 10
  *
- * @note 32b signed integer is supported.
- * @note 64b signed integers is supported for an input data up to 2^31
- * because gpu-non-native int128 is avoided for performance.
  * @todo Extend support for signed divisors
  */
 template <typename IntT>
@@ -51,12 +52,16 @@ struct FastIntDiv {
    * @brief host and device ctor's
    * @param other source object to be copied from
    */
-  HDI FastIntDiv(const FastIntDiv& other) : d(other.d), m(other.m), p(other.p) {}
+  HDI FastIntDiv(const FastIntDiv& other)
+    : d(other.d), m(other.m), p(other.p), fallback(other.fallback)
+  {
+  }
   HDI FastIntDiv& operator=(const FastIntDiv& other)
   {
-    d = other.d;
-    m = other.m;
-    p = other.p;
+    d        = other.d;
+    m        = other.m;
+    p        = other.p;
+    fallback = other.fallback;
     return *this;
   }
   /** @} */
@@ -67,6 +72,8 @@ struct FastIntDiv {
   UIntT m;
   /** the term 'p' as found in the reference chapter */
   int p;
+  /** Flag for falling back to canonical division on unsupported divisor's ranges */
+  bool fallback = false;
 
  private:
   void computeScalars()
@@ -79,6 +86,9 @@ struct FastIntDiv {
       ASSERT(false, "FastIntDiv: division by negative numbers not supported!");
     } else if (d == 0) {
       ASSERT(false, "FastIntDiv: got division by zero!");
+    } else if (int64_t(d) > kInt32Max) {
+      fallback = true;
+      return;
     }
     int64_t nc = ((1LL << 31) / d) * d - 1;
     p          = 31;
@@ -99,13 +109,16 @@ struct FastIntDiv {
  * @param divisor the denominator
  * @return the quotient
  */
-template <typename IntT>
-HDI IntT operator/(IntT n, const FastIntDiv<IntT>& divisor)
+template <typename NumIntT, typename DivIntT>
+HDI std::common_type_t<NumIntT, DivIntT> operator/(NumIntT n, const FastIntDiv<DivIntT>& divisor)
 {
-  if (divisor.d == 1) return n;
-  IntT ret = (int64_t(divisor.m) * int64_t(n)) >> divisor.p;
-  if (n < 0) ++ret;
-  return ret;
+  using CommonIntT = std::common_type_t<NumIntT, DivIntT>;
+  if (divisor.d == 1) return CommonIntT(n);
+  if (divisor.fallback || n < kInt32Min || n > kInt32Max) {
+    return CommonIntT(n) / CommonIntT(divisor.d);
+  }
+  CommonIntT ret = (int64_t(divisor.m) * int64_t(n)) >> divisor.p;
+  return ret + CommonIntT(n < 0);
 }
 
 /**
@@ -115,12 +128,14 @@ HDI IntT operator/(IntT n, const FastIntDiv<IntT>& divisor)
  * @param divisor the denominator
  * @return the remainder
  */
-template <typename IntT>
-HDI IntT operator%(IntT n, const FastIntDiv<IntT>& divisor)
+template <typename NumIntT, typename DivIntT>
+HDI std::common_type_t<NumIntT, DivIntT> operator%(NumIntT n, const FastIntDiv<DivIntT>& divisor)
 {
-  IntT quotient  = n / divisor;
-  IntT remainder = n - quotient * divisor.d;
+  using CommonIntT     = std::common_type_t<NumIntT, DivIntT>;
+  CommonIntT quotient  = n / divisor;
+  CommonIntT remainder = CommonIntT(n) - quotient * CommonIntT(divisor.d);
   return remainder;
+  // return n % divisor.d;
 }
 
 };  // namespace util

cpp/tests/util/fast_int_div.cu

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -13,42 +13,104 @@
 
 namespace raft::util {
 
-constexpr int64_t kInt32Max = std::numeric_limits<int32_t>::max();
+constexpr auto kInt64Min = std::numeric_limits<int64_t>::min();
+constexpr auto kInt64Max = std::numeric_limits<int64_t>::max();
 
-template <typename IntT>
-class FastIntDivTest : public ::testing::Test {
- protected:
-  void CompareWithNativeDivision()
-  {
-    std::vector<IntT> magnitudes{0, 1, 2, 3, 7, 13, 255, 12345, (1 << 20), kInt32Max};
-    std::vector<IntT> divisors{1, 2, 4, 7, 16, 31, 63, 128, 1000, (1 << 15), kInt32Max};
+TEST(FastIntDivTest, UnsupportedDivisors)
+{
+  ASSERT_THROW(FastIntDiv{-42}, raft::exception);
+  ASSERT_THROW(FastIntDiv{0}, raft::exception);
+}
 
-    for (IntT d : divisors) {
-      FastIntDiv fid(d);
-      for (IntT mag : magnitudes) {
-        for (IntT n : {mag, -mag}) {
-          ASSERT_EQ(n / fid, n / d) << "operator/ mismatch for numerator=" << n << " divisor=" << d;
-          ASSERT_EQ(n % fid, n % d) << "operator% mismatch for numerator=" << n << " divisor=" << d;
-        }
-      }
+TEST(FastIntDivTest, Int32Division)
+{
+  std::vector<int32_t> numerators{kInt32Min,
+                                  -(1 << 20),
+                                  -12345,
+                                  -255,
+                                  -13,
+                                  -7,
+                                  -3,
+                                  -2,
+                                  0,
+                                  1,
+                                  2,
+                                  3,
+                                  7,
+                                  13,
+                                  255,
+                                  12345,
+                                  (1 << 20),
+                                  kInt32Max};
+  std::vector<int32_t> divisors{1, 2, 4, 7, 16, 31, 63, 128, 1000, (1 << 15), kInt32Max};
+
+  for (auto d : divisors) {
+    FastIntDiv fid(d);
+    for (auto n : numerators) {
+      ASSERT_EQ(n / fid, n / d) << "operator/ mismatch for numerator=" << n << " divisor=" << d;
+      ASSERT_EQ(n % fid, n % d) << "operator% mismatch for numerator=" << n << " divisor=" << d;
     }
   }
-};
+}
 
-using FastIntDivTypes = ::testing::Types<int32_t, int64_t>;
-TYPED_TEST_CASE(FastIntDivTest, FastIntDivTypes);
+TEST(FastIntDivTest, Int64Divisors)
+{
+  std::vector<int64_t> numerators{kInt64Min,
+                                  int64_t(kInt32Min) - 12345678,
+                                  int64_t(kInt32Min),
+                                  -10007,
+                                  -1013,
+                                  0,
+                                  1013,
+                                  10007,
+                                  int64_t(kInt32Max) + 1,
+                                  int64_t(kInt32Max) + 2,
+                                  (int64_t(1) << 32),
+                                  (int64_t(1) << 33),
+                                  (int64_t(3) << 40),
+                                  kInt64Max};
 
-TYPED_TEST(FastIntDivTest, CompareWithNativeDivision) { this->CompareWithNativeDivision(); }
+  std::vector<int64_t> divisors{
+    1, 31, 129, 772, 1000, (int64_t(1) << 31), int64_t(kInt32Max) + 5, int64_t(3) << 40, kInt64Max};
 
-TEST(FastIntDiv, Int64NumeratorPastInt32Boundary)
-{
-  for (int64_t d : {129, 772, 1000}) {
+  for (auto d : divisors) {
     FastIntDiv fid(d);
-    for (int64_t n : {1LL << 31, 2147704000LL, 3LL << 30}) {  // in [2^31, 2^32)
+    for (auto n : numerators) {
       ASSERT_EQ(n / fid, n / d) << "operator/ mismatch for numerator=" << n << " divisor=" << d;
       ASSERT_EQ(n % fid, n % d) << "operator% mismatch for numerator=" << n << " divisor=" << d;
     }
   }
 }
 
+TEST(FastIntDivTest, CrossIntTypesDivision)
+{
+  // i32 numerators, i64 divisors
+  {
+    std::vector<int32_t> numerators{kInt32Min, -123, 0, 321, kInt32Max};
+    std::vector<int64_t> divisors{1, 4321, kInt64Max};
+
+    for (auto d : divisors) {
+      FastIntDiv fid(d);
+      for (auto n : numerators) {
+        ASSERT_EQ(n / fid, n / d) << "operator/ mismatch for numerator=" << n << " divisor=" << d;
+        ASSERT_EQ(n % fid, n % d) << "operator% mismatch for numerator=" << n << " divisor=" << d;
+      }
+    }
+  }
+
+  // i64 numerators, i32 divisors
+  {
+    std::vector<int64_t> numerators{kInt64Min, -1234, 0, 4321, kInt64Max};
+    std::vector<int32_t> divisors{1, 1234, kInt32Max};
+
+    for (auto d : divisors) {
+      FastIntDiv fid(d);
+      for (auto n : numerators) {
+        ASSERT_EQ(n / fid, n / d) << "operator/ mismatch for numerator=" << n << " divisor=" << d;
+        ASSERT_EQ(n % fid, n % d) << "operator% mismatch for numerator=" << n << " divisor=" << d;
+      }
+    }
+  }
+}
+
 }  // namespace raft::util