diff --git a/python/tilus/hidet/ir/primitives/cuda/fast_divmod.py b/python/tilus/hidet/ir/primitives/cuda/fast_divmod.py index 3815a3db..6f133912 100644 --- a/python/tilus/hidet/ir/primitives/cuda/fast_divmod.py +++ b/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. """ diff --git a/python/tilus/lang/instructions/root.py b/python/tilus/lang/instructions/root.py index e25d79ba..9b42caf1 100644 --- a/python/tilus/lang/instructions/root.py +++ b/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). diff --git a/tests/instructions/test_fast_divmod.py b/tests/instructions/test_fast_divmod.py new file mode 100644 index 00000000..15f822b2 --- /dev/null +++ b/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)