Add fast_divmod magic-number division helper

python/flydsl/expr/numeric.py

@@ -894,3 +894,112 @@ def __init__(self, x):
         # x is now either: Python int, or index-typed ir.Value
         # Pass directly to Numeric.__init__ (bypass Integer conversion logic)
         Numeric.__init__(self, x)
+
+
+# ---------------------------------------------------------------------------
+# Magic-number unsigned division
+# ---------------------------------------------------------------------------
+# GPUs have no hardware integer divide, so a runtime `n // d` lowers to a long
+# instruction sequence. When the divisor is known on the host (e.g. a grid
+# dimension), the divide can be replaced by a widening multiply and two shifts
+# using a precomputed multiplier, the same trick as ``cutlass::FastDivmod``.
+
+
+def fastdivmod_magic(divisor: int):
+    """Precompute the ``(magic, shift)`` pair for magic-number division by
+    ``divisor``, to be passed to :func:`fast_divmod` as kernel arguments.
+
+    Runs on the host with plain Python ints. The identity is::
+
+        n // divisor == (n * magic) >> (32 + shift)
+
+    for every ``0 <= n < 2**31``, with ``shift = max(ceil_log2(divisor) - 1, 0)``
+    and ``magic = ceil(2**(32 + shift) / divisor)``. ``magic`` is at most
+    ``2**32`` so it fits in an ``i64`` kernel argument and needs no
+    ``divisor == 1`` correction. Same non-negative, ``< 2**31`` dividend
+    contract as ``cutlass::FastDivmod``.
+    """
+    if not 0 < divisor < (1 << 31):
+        raise ValueError(f"divisor must satisfy 0 < divisor < 2**31, got {divisor}")
+    ceil_log2 = (divisor - 1).bit_length()
+    shift = max(ceil_log2 - 1, 0)
+    magic = ((1 << (32 + shift)) + divisor - 1) // divisor
+    return magic, shift
+
+
+def fast_divmod(dividend, divisor, magic, shift):
+    """Magic-number ``divmod`` returning ``(quotient, remainder)``.
+
+    ``magic`` and ``shift`` come from :func:`fastdivmod_magic`. Every argument
+    may be a DSL Numeric value (e.g. a runtime kernel argument) or a Python int.
+    Valid for ``0 <= dividend < 2**31``.
+    """
+    prod = Uint64(dividend) * Uint64(magic)
+    quotient = Int32(prod >> (Uint64(32) + Uint64(shift)))
+    remainder = Int32(dividend) - quotient * Int32(divisor)
+    return quotient, remainder
+
+
+def _ceil_log2_i32(x):
+    """``ceil(log2(x))`` for ``1 <= x < 2**31`` as ``32 - ctlz(x - 1)``.
+
+    Works on a Python int (folds) or a runtime ``Int32``. ``ctlz(0) == 32`` makes
+    ``x == 1`` yield 0.
+    """
+    from .math import ctlz
+
+    return Int32(32) - Int32(ctlz(Int32(x) - Int32(1)))
+
+
+class FastDivmod:
+    """Magic-number unsigned divmod by a fixed ``divisor``.
+
+    ``divisor`` may be a Python ``int`` (folds to constants) or a runtime
+    ``Int32`` such as a grid dimension known only at launch. The ``(magic,
+    shift)`` pair is derived once at construction via ``ceil_log2`` and a single
+    divide; the quotient is then ``(dividend * magic) >> (32 + shift)`` with no
+    runtime division. Valid for non-negative dividends below ``2**31``, the same
+    contract as ``cutlass::FastDivmod``.
+
+    It implements the DSL value protocol -- ``magic``, ``shift`` and ``divisor``
+    are the carried leaves -- so an instance can cross the host/device boundary
+    as a kernel argument or sit inside a ``@fx.struct`` / coordinate tuple.
+
+    Example::
+
+        fdm = FastDivmod(768)      # or FastDivmod(grid_n) with a runtime Int32
+        q, r = fdm.divmod(idx)
+    """
+
+    def __init__(self, divisor):
+        if isinstance(divisor, int):
+            if not 0 < divisor < (1 << 31):
+                raise ValueError(f"divisor must satisfy 0 < divisor < 2**31, got {divisor}")
+            divisor = Int32(divisor)
+        self.divisor = Int32(divisor)
+        shift = _ceil_log2_i32(self.divisor) - Int32(1)
+        shift = (shift < Int32(0)).select(Int32(0), shift)  # max(shift, 0)
+        d64 = Uint64(self.divisor)
+        numer = Uint64(0x100000000) * (Uint64(1) << Uint64(shift))
+        self.magic = (numer + d64 - Uint64(1)) // d64  # <= 2**32, exact for n < 2**31
+        self.shift = Uint32(shift)
+
+    def divmod(self, dividend):
+        return fast_divmod(dividend, self.divisor, self.magic, self.shift)
+
+    def div(self, dividend):
+        return self.divmod(dividend)[0]
+
+    def mod(self, dividend):
+        return self.divmod(dividend)[1]
+
+    def __extract_to_ir_values__(self):
+        return [self.magic.ir_value(), self.shift.ir_value(), self.divisor.ir_value()]
+
+    @classmethod
+    def __construct_from_ir_values__(cls, values, exemplar=None):
+        obj = object.__new__(cls)
+        obj.magic = Uint64(values[0])
+        obj.shift = Uint32(values[1])
+        obj.divisor = Int32(values[2])
+        return obj

python/flydsl/expr/typing.py

@@ -18,6 +18,7 @@
 from .numeric import (
     BFloat16,
     Boolean,
+    FastDivmod,
     Float,
     Float4E2M1FN,
     Float6E2M3FN,
@@ -49,6 +50,8 @@
     _resolve_numeric_type,
     _result_numeric_type_for_op,
     as_numeric,
+    fast_divmod,
+    fastdivmod_magic,
 )
 from .primitive import *
 from .utils import lazy_classattr
@@ -356,6 +359,9 @@ def vec(self, n: int, elem: ir.Type) -> ir.Type:
     "as_dsl_value",
     "is_generic_address_space",
     "is_target_address_space",
+    "fast_divmod",
+    "fastdivmod_magic",
+    "FastDivmod",
     # DSL value types
     "Numeric",
     "as_numeric",

tests/unit/test_fast_divmod.py

@@ -0,0 +1,183 @@
+#!/usr/bin/env python3
+
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 FlyDSL Project Contributors
+
+"""Tests for magic-number division: fastdivmod_magic / fast_divmod / FastDivmod.
+
+The host magic math is checked in pure Python (L0). The device path is checked
+by running a kernel that compares fast_divmod against native ``//`` and ``%``
+for a runtime divisor (L2).
+"""
+
+from __future__ import annotations
+
+import pytest
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from flydsl.expr.numeric import FastDivmod, fast_divmod, fastdivmod_magic
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+
+DIVISORS = [1, 2, 3, 7, 8, 127, 128, 768, 1000, 1024, 12289, 32000, 65536, 128256]
+DIVIDENDS = [0, 1, 2, 5, 255, 256, 1023, 100000, 998244353, (1 << 31) - 1]
+
+
+@pytest.mark.l0_backend_agnostic
+@pytest.mark.parametrize("divisor", DIVISORS)
+def test_fastdivmod_magic_matches_floordiv(divisor):
+    magic, shift = fastdivmod_magic(divisor)
+    assert magic <= (1 << 32)
+    for n in DIVIDENDS:
+        q = (n * magic) >> (32 + shift)
+        assert q == n // divisor, f"d={divisor} n={n}: got {q}, want {n // divisor}"
+        assert n - q * divisor == n % divisor
+
+
+@pytest.mark.l0_backend_agnostic
+def test_fastdivmod_magic_rejects_out_of_range():
+    with pytest.raises(ValueError):
+        fastdivmod_magic(0)
+    with pytest.raises(ValueError):
+        fastdivmod_magic(1 << 31)
+
+
+@pytest.mark.l2_device
+@pytest.mark.rocm_lower
+@pytest.mark.skipif(torch is None or not torch.cuda.is_available(), reason="requires GPU")
+@pytest.mark.parametrize("divisor", [3, 7, 768, 32000, 128256])
+def test_fast_divmod_device_matches_native(divisor):
+    BLOCK = 256
+    NBLOCKS = 64
+    P = BLOCK * NBLOCKS
+
+    @flyc.kernel(known_block_size=[BLOCK, 1, 1])
+    def kernel(Q: fx.Tensor, R: fx.Tensor, d: fx.Int32, magic: fx.Int64, shift: fx.Int32):
+        g = fx.block_idx.x * BLOCK + fx.thread_idx.x
+        n = fx.Int32(fx.Uint32(g) * fx.Uint32(2654435761) & fx.Uint32(0x7FFFFFFF))
+        q, r = fast_divmod(n, d, magic, shift)
+        fx.memref_store(q, Q, g)
+        fx.memref_store(r, R, g)
+
+    @flyc.jit
+    def launch(
+        Q: fx.Tensor, R: fx.Tensor, d: fx.Int32, magic: fx.Int64, shift: fx.Int32, stream: fx.Stream = fx.Stream(None)
+    ):
+        kernel(Q, R, d, magic, shift).launch(grid=(NBLOCKS, 1, 1), block=(BLOCK, 1, 1), stream=stream)
+
+    magic, shift = fastdivmod_magic(divisor)
+    q = torch.zeros(P, dtype=torch.int32, device="cuda")
+    r = torch.zeros(P, dtype=torch.int32, device="cuda")
+    launch(q, r, divisor, magic, shift, stream=torch.cuda.Stream())
+    torch.cuda.synchronize()
+
+    g = torch.arange(P, dtype=torch.int64, device="cuda")
+    n = ((g * 2654435761) & 0x7FFFFFFF).to(torch.int64)
+    assert torch.equal(q.to(torch.int64), n // divisor)
+    assert torch.equal(r.to(torch.int64), n % divisor)
+
+
+@pytest.mark.l2_device
+@pytest.mark.rocm_lower
+@pytest.mark.skipif(torch is None or not torch.cuda.is_available(), reason="requires GPU")
+def test_fastdivmod_class_constant_divisor():
+    BLOCK = 128
+    P = BLOCK * 8
+    DIV = 768
+
+    @flyc.kernel(known_block_size=[BLOCK, 1, 1])
+    def kernel(Q: fx.Tensor, R: fx.Tensor):
+        g = fx.block_idx.x * BLOCK + fx.thread_idx.x
+        n = fx.Int32(fx.Uint32(g) * fx.Uint32(2654435761) & fx.Uint32(0x7FFFFFFF))
+        fdm = FastDivmod(DIV)
+        q, r = fdm.divmod(n)
+        fx.memref_store(q, Q, g)
+        fx.memref_store(r, R, g)
+
+    @flyc.jit
+    def launch(Q: fx.Tensor, R: fx.Tensor, stream: fx.Stream = fx.Stream(None)):
+        kernel(Q, R).launch(grid=(P // BLOCK, 1, 1), block=(BLOCK, 1, 1), stream=stream)
+
+    q = torch.zeros(P, dtype=torch.int32, device="cuda")
+    r = torch.zeros(P, dtype=torch.int32, device="cuda")
+    launch(q, r, stream=torch.cuda.Stream())
+    torch.cuda.synchronize()
+
+    g = torch.arange(P, dtype=torch.int64, device="cuda")
+    n = ((g * 2654435761) & 0x7FFFFFFF).to(torch.int64)
+    assert torch.equal(q.to(torch.int64), n // DIV)
+    assert torch.equal(r.to(torch.int64), n % DIV)
+
+
+@pytest.mark.l2_device
+@pytest.mark.rocm_lower
+@pytest.mark.skipif(torch is None or not torch.cuda.is_available(), reason="requires GPU")
+@pytest.mark.parametrize("divisor", [1, 3, 768, 128256])
+def test_fastdivmod_dynamic_divisor(divisor):
+    """FastDivmod built from a runtime (dynamic) divisor, magic derived on device."""
+    BLOCK = 256
+    P = BLOCK * 32
+
+    @flyc.kernel(known_block_size=[BLOCK, 1, 1])
+    def kernel(Q: fx.Tensor, R: fx.Tensor, d: fx.Int32):
+        g = fx.block_idx.x * BLOCK + fx.thread_idx.x
+        n = fx.Int32(fx.Uint32(g) * fx.Uint32(2654435761) & fx.Uint32(0x7FFFFFFF))
+        q, r = FastDivmod(d).divmod(n)
+        fx.memref_store(q, Q, g)
+        fx.memref_store(r, R, g)
+
+    @flyc.jit
+    def launch(Q: fx.Tensor, R: fx.Tensor, d: fx.Int32, stream: fx.Stream = fx.Stream(None)):
+        kernel(Q, R, d).launch(grid=(P // BLOCK, 1, 1), block=(BLOCK, 1, 1), stream=stream)
+
+    q = torch.zeros(P, dtype=torch.int32, device="cuda")
+    r = torch.zeros(P, dtype=torch.int32, device="cuda")
+    launch(q, r, divisor, stream=torch.cuda.Stream())
+    torch.cuda.synchronize()
+
+    g = torch.arange(P, dtype=torch.int64, device="cuda")
+    n = ((g * 2654435761) & 0x7FFFFFFF).to(torch.int64)
+    assert torch.equal(q.to(torch.int64), n // divisor)
+    assert torch.equal(r.to(torch.int64), n % divisor)
+
+
+@pytest.mark.l2_device
+@pytest.mark.rocm_lower
+@pytest.mark.skipif(torch is None or not torch.cuda.is_available(), reason="requires GPU")
+def test_fastdivmod_value_protocol_kernel_arg():
+    """FastDivmod crosses the host/device boundary as a kernel argument.
+
+    The launch wrapper builds it from a runtime divisor and passes the instance
+    itself; the (magic, shift, divisor) leaves are extracted/reconstructed by the
+    value protocol.
+    """
+    BLOCK = 256
+    P = BLOCK * 32
+    DIV = 768
+
+    @flyc.kernel(known_block_size=[BLOCK, 1, 1])
+    def kernel(Q: fx.Tensor, R: fx.Tensor, fdm: FastDivmod):
+        g = fx.block_idx.x * BLOCK + fx.thread_idx.x
+        n = fx.Int32(fx.Uint32(g) * fx.Uint32(2654435761) & fx.Uint32(0x7FFFFFFF))
+        q, r = fdm.divmod(n)
+        fx.memref_store(q, Q, g)
+        fx.memref_store(r, R, g)
+
+    @flyc.jit
+    def launch(Q: fx.Tensor, R: fx.Tensor, d: fx.Int32, stream: fx.Stream = fx.Stream(None)):
+        kernel(Q, R, FastDivmod(d)).launch(grid=(P // BLOCK, 1, 1), block=(BLOCK, 1, 1), stream=stream)
+
+    q = torch.zeros(P, dtype=torch.int32, device="cuda")
+    r = torch.zeros(P, dtype=torch.int32, device="cuda")
+    launch(q, r, DIV, stream=torch.cuda.Stream())
+    torch.cuda.synchronize()
+
+    g = torch.arange(P, dtype=torch.int64, device="cuda")
+    n = ((g * 2654435761) & 0x7FFFFFFF).to(torch.int64)
+    assert torch.equal(q.to(torch.int64), n // DIV)
+    assert torch.equal(r.to(torch.int64), n % DIV)