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,33 @@
 
 #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.
+ * **Usage**
+ *
+ * Construct the divisor once and call `/` and `%` operators repeatedly with
+ * different numerators.
+
+ * @code{.cpp}
+ * raft::util::FastIntDiv<int32_t> div(stride);
+ * int32_t quotient  = flat_index / div;
+ * int32_t remainder = flat_index % div;
+ * @endcode
+ *
+ * @note It will auto-fallback to the plain division when the divisor or the numerator
+ * is beyond the 32-bit signed integer range.
+ *
  * @todo Extend support for signed divisors
  */
 template <typename IntT>
@@ -51,12 +66,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 +86,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 +100,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;
@@ -94,33 +118,44 @@ struct FastIntDiv {
 
 /**
  * @brief Division overload, so that FastIntDiv can be transparently switched
- *        to even on device
+ *
+ * @note Not meant to be called directly, but via `n / div` where `div` is a
+ *       `FastIntDiv` instance
+ *
  * @param n numerator
- * @param divisor the denominator
+ * @param divisor the precomputed divisor
  * @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);
 }
 
 /**
- * @brief Modulo overload, so that FastIntDiv can be transparently switched
- *        to even on device
+ * @brief Modulo overload enabling transparent use of `FastIntDiv` with `%`.
+ *
+ * @note Not meant to be called directly, but via `n % div` where `div` is a
+ *       `FastIntDiv` instance
+ *
  * @param n numerator
- * @param divisor the denominator
+ * @param divisor the precomputed divisor
  * @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