[Fix] Keep the fast_divmod multiplier within 32 bits

python/tilus/hidet/ir/primitives/cuda/fast_divmod.py

@@ -17,7 +17,7 @@
 
 Two primitive functions are registered:
 
-1. ``fastdiv(a, b)`` — semantic placeholder for ``floor(a / b)`` where a >= 0 and b > 0.
+1. ``fastdiv(a, b)`` — semantic placeholder for ``floor(a / b)`` where 0 <= a < 2^31 and b > 0.
    A Hidet IR pass (LowerFastDivPass) replaces this with the runtime computation using
    precomputed magic multiplier and shift.
 
@@ -43,7 +43,7 @@ def register_fast_divmod_primitives():
     @no_type_check
     @script
     def cuda_fastdiv(a: i32, b: i32) -> i32:
-        """Semantic placeholder for floor(a / b) where a >= 0 and b > 0.
+        """Semantic placeholder for floor(a / b) where 0 <= a < 2^31 and b > 0.
 
         This function will be lowered by LowerFastDivPass to use precomputed
         magic multiplier and shift. It should never survive to codegen.
@@ -78,20 +78,21 @@ def cuda_fastdiv_runtime(a: i32, multiplier: i32, shift: i32) -> i32:
     def cuda_fastdiv_precompute_m(b: uint32) -> uint32:
         """Host-side: compute magic multiplier for divisor b.
 
-        Uses the Hacker's Delight algorithm for unsigned division by constant.
         For power-of-2 divisors, returns 0 (sentinel) and the runtime uses a simple shift.
-        For non-power-of-2 divisors, finds the smallest p >= 32 such that
-        2^p > nc * (d - 2^p % d), then m = (2^p + d - 2^p % d) / d.
+        For non-power-of-2 divisors, returns m = ceil(2^p / b) with p = 31 + ceil(log2(b)), as
+        CUTLASS FastDivmod does. This is exact for dividends in [0, 2^31) and always fits in 32 bits,
+        while a multiplier that is exact for the whole uint32 range needs 33 bits for some divisors.
         """
         attrs.func_kind = "cpu_internal"
         # Check power of 2: b & (b - 1) == 0
         if (b & (b - uint32(1))) == uint32(0):
             return uint32(0)
-        nc: uint64 = (uint64(1) << uint64(32)) - uint64(1) - uint64((uint64(1) << uint64(32)) % uint64(b))
-        p: uint32 = 32
-        while (uint64(1) << uint64(p)) <= nc * uint64(uint64(b) - (uint64(1) << uint64(p)) % uint64(b)):
+        p: uint32 = 31
+        tmp: uint32 = b - uint32(1)
+        while tmp > uint32(0):
+            tmp = tmp >> uint32(1)
             p = p + uint32(1)
-        m: uint64 = ((uint64(1) << uint64(p)) + uint64(b) - (uint64(1) << uint64(p)) % uint64(b)) / uint64(b)
+        m: uint64 = ((uint64(1) << uint64(p)) + uint64(b) - uint64(1)) / uint64(b)
         return uint32(m)
 
     @no_type_check
@@ -100,7 +101,7 @@ def cuda_fastdiv_precompute_s(b: uint32) -> uint32:
         """Host-side: compute shift for divisor b.
 
         For power-of-2 divisors, returns log2(b).
-        For non-power-of-2 divisors, returns p - 32 where p is from the Hacker's Delight algorithm.
+        For non-power-of-2 divisors, returns p - 32 with p = 31 + ceil(log2(b)), see cuda_fastdiv_precompute_m.
         """
         attrs.func_kind = "cpu_internal"
         if (b & (b - uint32(1))) == uint32(0):
@@ -111,9 +112,10 @@ def cuda_fastdiv_precompute_s(b: uint32) -> uint32:
                 tmp = tmp >> uint32(1)
                 s = s + uint32(1)
             return s
-        nc: uint64 = (uint64(1) << uint64(32)) - uint64(1) - uint64((uint64(1) << uint64(32)) % uint64(b))
-        p: uint32 = 32
-        while (uint64(1) << uint64(p)) <= nc * uint64(uint64(b) - (uint64(1) << uint64(p)) % uint64(b)):
+        p: uint32 = 31
+        rest: uint32 = b - uint32(1)
+        while rest > uint32(0):
+            rest = rest >> uint32(1)
             p = p + uint32(1)
         return p - uint32(32)
 
@@ -122,7 +124,7 @@ def cuda_fastdiv_precompute_s(b: uint32) -> uint32:
 
 
 def fastdiv(a: Expr, b: Expr) -> Expr:
-    """Semantic floor division: floor(a / b) for a >= 0, b > 0.
+    """Semantic floor division: floor(a / b) for 0 <= a < 2^31, b > 0.
 
     Will be lowered by LowerFastDivPass to use precomputed multiplier and shift.
     """

python/tilus/lang/instructions/root.py

@@ -2032,7 +2032,7 @@ def assign(self, dst: RegisterTensor, src: RegisterTensor) -> None:
     def fast_divmod(self, a: Union[Expr, int], b: Union[Expr, int]) -> tuple:
         """Fast integer division and modulo using precomputed magic multiplier.
 
-        Computes ``(a // b, a % b)`` for ``a >= 0`` and ``b > 0`` using the
+        Computes ``(a // b, a % b)`` for ``0 <= a < 2^31`` and ``b > 0`` using the
         FastDivmod algorithm. The divisor ``b`` must be a grid-constant expression
         (known at kernel launch time). A Hidet IR pass will lower this to
         ``__umulhi(a, magic) >> shift`` with host-side precomputation of
@@ -2041,7 +2041,7 @@ def fast_divmod(self, a: Union[Expr, int], b: Union[Expr, int]) -> tuple:
         Parameters
         ----------
         a : Expr
-            The dividend (int32, must be >= 0).
+            The dividend (int32, must satisfy 0 <= a < 2^31).
         b : Expr
             The divisor (int32, must be > 0 and grid-constant).
 

tests/instructions/test_fast_divmod.py

@@ -0,0 +1,122 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import typing
+
+import pytest
+import tilus
+import torch
+from tilus import int32
+from tilus.hidet.ir.dtypes import uint32
+from tilus.hidet.ir.primitives.cuda.fast_divmod import fastdiv_precompute_m, fastdiv_precompute_s
+from tilus.hidet.lang import attrs, script
+from tilus.hidet.lang.script import script_module
+from tilus.target import nvgpu_sm80, scope
+from tilus.testing import requires
+
+MAX_DIVIDEND = 2**31 - 1
+# Every small divisor, plus the divisors around each power of two up to the int32 limit.
+DIVISORS = sorted({*range(1, 1025), *(2**k + 1 for k in range(11, 31)), *(2**k - 1 for k in range(11, 32))})
+
+
+def build_precompute_kernel():
+    with script_module() as module:
+
+        @typing.no_type_check
+        @script
+        def _precompute(m: ~uint32, s: ~uint32, b: ~uint32, n: int32):
+            attrs.func_kind = "cpu_kernel"
+            for i in range(n):
+                m[i] = fastdiv_precompute_m(b[i])
+                s[i] = fastdiv_precompute_s(b[i])
+
+    with scope(nvgpu_sm80):
+        return module.build()
+
+
+def fastdiv_runtime(a: torch.Tensor, multiplier: int, shift: int) -> torch.Tensor:
+    """Evaluate cuda_fastdiv_runtime on the host for non-negative dividends."""
+    if multiplier == 0:
+        return a >> shift
+    # a * multiplier fits in int64 because a < 2^31 and multiplier < 2^32.
+    return ((a * multiplier) >> 32) >> shift
+
+
+def boundary_dividends(b: int) -> torch.Tensor:
+    """Multiples of b and their predecessors at both ends of the dividend range, where the error would show up first."""
+    top = MAX_DIVIDEND // b
+    factors = torch.cat([torch.arange(1, 17), torch.arange(max(top - 15, 1), top + 1)])
+    multiples = factors[factors <= top] * b
+    return torch.cat([torch.tensor([0, MAX_DIVIDEND]), multiples - 1, multiples])
+
+
+def test_precomputed_multiplier_is_exact(monkeypatch):
+    # The kernel runs on the host, so building it needs nvcc but no GPU.
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    kernel = build_precompute_kernel()
+    monkeypatch.undo()
+
+    n = len(DIVISORS)
+    m = torch.zeros(n, dtype=torch.uint32)
+    s = torch.zeros(n, dtype=torch.uint32)
+    kernel(m, s, torch.tensor(DIVISORS, dtype=torch.uint32), n)
+
+    wrong = []
+    for b, multiplier, shift in zip(DIVISORS, m.tolist(), s.tolist()):
+        a = boundary_dividends(b)
+        if not torch.equal(fastdiv_runtime(a, multiplier, shift), a // b):
+            wrong.append(b)
+    assert not wrong, "{} of {} divisors give a wrong quotient, the first are {}".format(
+        len(wrong), len(DIVISORS), wrong[:8]
+    )
+
+
+class FastDivmodExample(tilus.Script):
+    def __call__(self, base: int32, b: int32, quotient: ~int32, remainder: ~int32):
+        self.attrs.blocks = (256, 1, 1)
+        self.attrs.warps = 1
+
+        q, r = self.fast_divmod(base + self.blockIdx.x, b)
+        g_q = self.global_view(quotient, dtype=int32, shape=[256])
+        g_r = self.global_view(remainder, dtype=int32, shape=[256])
+        with self.thread_group(0, num_threads=1):
+            self.store_global(
+                g_q, self.register_tensor(dtype=int32, shape=[1], init=lambda i: q), offsets=[self.blockIdx.x]
+            )
+            self.store_global(
+                g_r, self.register_tensor(dtype=int32, shape=[1], init=lambda i: r), offsets=[self.blockIdx.x]
+            )
+
+
+@requires.nvgpu_sm80
+@pytest.mark.parametrize(
+    "base, b",
+    [
+        (0, 3),
+        (0, 7),
+        (0, 56),
+        (0, 64),
+        (MAX_DIVIDEND - 255, 7),
+    ],
+)
+def test_fast_divmod(base, b):
+    device = "cuda"
+    quotient = torch.full((256,), -1, dtype=torch.int32, device=device)
+    remainder = torch.full((256,), -1, dtype=torch.int32, device=device)
+
+    FastDivmodExample()(base, b, quotient, remainder)
+
+    a = torch.arange(base, base + 256, dtype=torch.int64, device=device)
+    torch.testing.assert_close(quotient.to(torch.int64), a // b)
+    torch.testing.assert_close(remainder.to(torch.int64), a % b)