From 30412e742da46bd4d42c361c502ed6343fe13448 Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Sun, 26 Jul 2026 20:15:11 +0200 Subject: [PATCH 1/3] Add fast_divmod magic-number division helper --- python/flydsl/expr/numeric.py | 72 +++++++++++++++++++++ python/flydsl/expr/typing.py | 6 ++ tests/unit/test_fast_divmod.py | 114 +++++++++++++++++++++++++++++++++ 3 files changed, 192 insertions(+) create mode 100644 tests/unit/test_fast_divmod.py diff --git a/python/flydsl/expr/numeric.py b/python/flydsl/expr/numeric.py index 6d1b826a3..253bbec93 100644 --- a/python/flydsl/expr/numeric.py +++ b/python/flydsl/expr/numeric.py @@ -894,3 +894,75 @@ 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 + + +class FastDivmod: + """Magic-number divmod by a fixed ``divisor``. + + Convenience wrapper for a compile-time-constant divisor: the + ``(magic, shift)`` pair is computed once at construction. For a runtime + divisor (a kernel argument) call :func:`fast_divmod` directly with a + ``(magic, shift)`` pair produced on the host by :func:`fastdivmod_magic`. + + Example:: + + fdm = FastDivmod(768) + q, r = fdm.divmod(idx) + """ + + def __init__(self, divisor: int): + self.divisor = divisor + self.magic, self.shift = fastdivmod_magic(divisor) + + 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] diff --git a/python/flydsl/expr/typing.py b/python/flydsl/expr/typing.py index 4b6ebcbd9..d3f59927c 100644 --- a/python/flydsl/expr/typing.py +++ b/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 @@ -359,6 +362,9 @@ def vec(self, n: int, elem: ir.Type) -> ir.Type: # DSL value types "Numeric", "as_numeric", + "fast_divmod", + "fastdivmod_magic", + "FastDivmod", "Boolean", "Float", "BFloat16", diff --git a/tests/unit/test_fast_divmod.py b/tests/unit/test_fast_divmod.py new file mode 100644 index 000000000..dcb3bcce7 --- /dev/null +++ b/tests/unit/test_fast_divmod.py @@ -0,0 +1,114 @@ +#!/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) From f53cb3c23f9e5ed23ea283e4cdf29e656f66878b Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Mon, 27 Jul 2026 08:59:10 +0200 Subject: [PATCH 2/3] Group fast_divmod helpers under DSL utilities in __all__ --- python/flydsl/expr/typing.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/flydsl/expr/typing.py b/python/flydsl/expr/typing.py index d3f59927c..3ed3927d8 100644 --- a/python/flydsl/expr/typing.py +++ b/python/flydsl/expr/typing.py @@ -359,12 +359,12 @@ def vec(self, n: int, elem: ir.Type) -> ir.Type: "as_dsl_value", "is_generic_address_space", "is_target_address_space", - # DSL value types - "Numeric", - "as_numeric", "fast_divmod", "fastdivmod_magic", "FastDivmod", + # DSL value types + "Numeric", + "as_numeric", "Boolean", "Float", "BFloat16", From 87e0bc31ebeb391eff5522321870ef9af2c7f681 Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Mon, 27 Jul 2026 09:35:48 +0200 Subject: [PATCH 3/3] Support dynamic divisor in FastDivmod via the value protocol --- python/flydsl/expr/numeric.py | 55 ++++++++++++++++++++++----- tests/unit/test_fast_divmod.py | 69 ++++++++++++++++++++++++++++++++++ 2 files changed, 115 insertions(+), 9 deletions(-) diff --git a/python/flydsl/expr/numeric.py b/python/flydsl/expr/numeric.py index 253bbec93..76b7f95b6 100644 --- a/python/flydsl/expr/numeric.py +++ b/python/flydsl/expr/numeric.py @@ -940,23 +940,49 @@ def fast_divmod(dividend, divisor, magic, shift): 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 divmod by a fixed ``divisor``. + """Magic-number unsigned divmod by a fixed ``divisor``. - Convenience wrapper for a compile-time-constant divisor: the - ``(magic, shift)`` pair is computed once at construction. For a runtime - divisor (a kernel argument) call :func:`fast_divmod` directly with a - ``(magic, shift)`` pair produced on the host by :func:`fastdivmod_magic`. + ``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) + fdm = FastDivmod(768) # or FastDivmod(grid_n) with a runtime Int32 q, r = fdm.divmod(idx) """ - def __init__(self, divisor: int): - self.divisor = divisor - self.magic, self.shift = fastdivmod_magic(divisor) + 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) @@ -966,3 +992,14 @@ def div(self, dividend): 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 diff --git a/tests/unit/test_fast_divmod.py b/tests/unit/test_fast_divmod.py index dcb3bcce7..87584b3f7 100644 --- a/tests/unit/test_fast_divmod.py +++ b/tests/unit/test_fast_divmod.py @@ -112,3 +112,72 @@ def launch(Q: fx.Tensor, R: fx.Tensor, stream: fx.Stream = fx.Stream(None)): 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)