Add fast_divmod magic-number division helper

python/flydsl/expr/numeric.py

@@ -27,6 +27,9 @@
     "NumericMeta",
     "Numeric",
     "as_numeric",
+    "fast_divmod",
+    "fastdivmod_magic",
+    "FastDivmod",
     "Integer",
     "Float",
     "Boolean",
@@ -961,3 +964,146 @@ 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, sit in a ``@fx.struct`` field, or travel inside a
+    plain Python tuple of kernel arguments. It is *not* a valid ``int_tuple``
+    leaf: those must be a single i32/i64 value and this carries three.
+
+    **Construct it outside the kernel when the divisor is dynamic.** Deriving
+    ``magic`` costs one 64-bit divide, which is free for a Python-int divisor
+    (it folds) but is emitted per thread if an ``Int32`` divisor is handed to
+    the constructor inside a kernel -- and a 64-bit divide is worse than the
+    32-bit one being replaced. Build it in the ``@flyc.jit`` launch wrapper and
+    pass the instance as a kernel argument; the value protocol then carries the
+    already-derived leaves in. Measured on gfx1150, one divmod by a runtime
+    divisor::
+
+        native ``//`` and ``%``                 53 instructions
+        FastDivmod(d) built in the launch wrapper   27
+        FastDivmod(d) built inside the kernel      195
+
+    Only the Python-int path range-checks the divisor; a runtime ``Int32`` that
+    is zero or >= 2**31 silently yields wrong results.
+
+    Example::
+
+        # divisor known at trace time -- folds to constants
+        fdm = FastDivmod(768)
+        q, r = fdm.divmod(idx)
+
+        # divisor known only at launch -- build it in the wrapper, not the kernel
+        @flyc.jit
+        def launch(out: fx.Tensor, d: fx.Int32, stream: fx.Stream):
+            kernel(out, FastDivmod(d)).launch(grid=..., block=..., stream=stream)
+    """
+
+    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 __get_ir_types__(cls):
+        """Leaf types, in ``__extract_to_ir_values__`` order.
+
+        Aggregates size a field by calling this unbound on the class, so it has
+        to be a classmethod; without it a ``@fx.struct`` field defaults to one
+        leaf and reconstruction runs off the end of the value list.
+        """
+        return [Uint64.ir_type, Uint32.ir_type, Int32.ir_type]
+
+    @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

tests/unit/test_fast_divmod.py

@@ -0,0 +1,227 @@
+#!/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).
+"""
+
+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 divisor *inside* the kernel.
+
+    Checks correctness only. This is deliberately not the pattern to copy: the
+    magic derivation emits a 64-bit divide per thread (195 instructions on
+    gfx1150 against 53 for a native divmod). Building it in the launch wrapper
+    instead, as test_fastdivmod_value_protocol_kernel_arg does, keeps the divide
+    out of the kernel and costs 27.
+    """
+    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)
+
+
+@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_struct_field(divisor):
+    """FastDivmod as a @fx.struct field.
+
+    Aggregates size a field by calling ``__get_ir_types__`` unbound on the type
+    and fall back to a single leaf when it is missing, so without that
+    classmethod this reconstructs from one value instead of three and raises
+    IndexError.
+    """
+    BLOCK = 256
+
+    @fx.struct
+    class Params:
+        fdm: FastDivmod
+
+    @flyc.kernel(known_block_size=[BLOCK, 1, 1])
+    def kernel(Q: fx.Tensor, R: fx.Tensor, params: Params):
+        i = fx.Int32(fx.thread_idx.x)
+        q, r = params.fdm.divmod(i)
+        fx.memref_store(q, Q, i)
+        fx.memref_store(r, R, i)
+
+    @flyc.jit
+    def launch(Q: fx.Tensor, R: fx.Tensor, d: fx.Int32, stream: fx.Stream = fx.Stream(None)):
+        kernel(Q, R, Params(fdm=FastDivmod(d))).launch(grid=(1, 1, 1), block=(BLOCK, 1, 1), stream=stream)
+
+    q = torch.zeros(BLOCK, dtype=torch.int32, device="cuda")
+    r = torch.zeros(BLOCK, dtype=torch.int32, device="cuda")
+    launch(q, r, divisor, stream=torch.cuda.Stream())
+    torch.cuda.synchronize()
+
+    n = torch.arange(BLOCK, dtype=torch.int64, device="cuda")
+    assert torch.equal(q.to(torch.int64), n // divisor)
+    assert torch.equal(r.to(torch.int64), n % divisor)