[Feature Request] [Perf] LowerMagicDiv: host-precomputed magic division for launch-invariant dynamic divisors

Summary

tl.LowerMagicDiv rewrites integer division/modulo by a launch-invariant but compile-time-unknown positive int32 divisor into host-precomputed magic numbers, with each device thread executing only umulhi + shift:

q = umulhi(uint32(x), M) >> s
r = x - q * d

A plain dynamic integer division typically expands into long reciprocal-estimation sequences, corrections, or helper calls. Magic division moves all d-dependent expensive work to the host (computed once per kernel launch), so every device thread performs only an unsigned multiply-high and a right shift. This is not a floating-point approximate division — within the specified input domain the result is bit-exact with integer division.

1. Mathematical Principle

For x = q * d + r (dividend x, divisor d, quotient q, remainder r) with d ≥ 2:

k = ceil(log2(d))
p = 31 + k
M = ceil(2^p / d)
s = p - 32 = k - 1

For 0 ≤ x < 2^31:

floor(x / d) = floor(x * M / 2^p)
             = umulhi(uint32(x), M) >> s

where umulhi(a,b) returns the high 32 bits of the unsigned 32×32 product, i.e. umulhi(a,b) = floor(uint64(a) * uint64(b) / 2^32); the following right shift by s = p - 32 completes the divide-by-2^p floor.

Why it is exact. The key is that M is rounded up. Writing M = 2^p/d + ε with 0 ≤ ε < 1:

x * M / 2^p = x / d + x * ε / 2^p

Since x < 2^31 and p = 31 + k:

0 ≤ x * ε / 2^p < 1 / 2^k ≤ 1 / d

With x = q·d + r (0 ≤ r < d), x/d is at least 1/d away from the next integer; the additive error is strictly smaller than that distance, so it never crosses an integer boundary and the floor still yields the exact q = floor(x/d). This is also why the fast path requires a non-negative dividend below 2^31.

Worked example: 100 / 7

Host precomputation (d = 7):

k = ceil(log2(7)) = 3,  p = 31 + k = 34
M = ceil(2^34 / 7) = 2454267027 = 0x92492493
s = p - 32 = 2

Although M / 2^34 = 0.142857142898720… is only an upward fixed-point approximation of 1/7 = 0.142857142857143…, the error bound above guarantees the result never crosses the next integer. On device (x = 100):

umulhi(100, 2454267027) = floor(245426702700 / 2^32) = 57
q = 57 >> 2 = 14
r = x - q * d = 100 - 14 * 7 = 2

which matches 100 // 7 = 14, 100 % 7 = 2 exactly.

Special cases:

• d == 1: host sets M = 0, s = 0; device directly selects q = x, r = 0.
• d == 0: the host helper returns zero parameters without performing a division; a shape-derived zero divisor normally never launches the device work, and the domain of an executed divide-by-zero site is unchanged — magic div does not extend it.
• x < 0, d < 1, or signed int32 wrap: the fast path is not taken; an exact fallback is used.

2. Reference: CUTLASS Implementation

Primary reference: 3rdparty/cutlass/include/cutlass/fast_math.h

• find_log2: computes ceil(log2(d))
• find_divisor: computes multiplier and shift_right
• fast_divmod: device-side __umulhi(src, mul) >> shr, then rem = src - quo * div
• cutlass::FastDivmod: intended to be constructed on the host, with precomputed parameters passed through kernel Params

Core logic:

p   = 31 + ceil_log2(d);
mul = ceil(2^p / d);
shr = p - 32;
quo = d != 1 ? __umulhi(src, mul) >> shr : src;
rem = src - quo * d;

TileLang does not instantiate cutlass::FastDivmod in generated kernels. Instead it reuses the same math and the same host-precompute / device-mulhi structure, expressed as TIR intrinsics, so the optimization is applied automatically by a compiler pass to dynamic-shape indexing and can be lowered separately by the CUDA, HIP, and C codegens. Compared with using CUTLASS directly, TileLang additionally handles:

• Which TIR div/mod sites are safe to transform
• The semantic gap between TIR FloorDiv/FloorMod and C/CUDA truncating div/mod
• Fallback for invalid runtime values
• Quotient/remainder CSE for the same (x, d)
• Semantic correlation and validity-predicate reuse in conditions
• Interaction with FlattenBuffer, int64 promotion, and later Simplify on index algebra

3. Design in TileLang

Pipeline overview:

     dynamic shape / scalar divisor d
              │
              ▼
     LowerMagicDiv (device IR still uses multi-dim indexing)
       1. Normalize remainder expressions
       2. Prove divisor is launch-invariant, positive, int32
       3. Prove dividend non-negativity and condition correlation
       4. Deduplicate by equivalent divisors
       5. Rewrite div/mod into tl.magic_* intrinsics
              │
              ├────── host Bind ──────┐
              │   TileLangHostFastDivmodU32Mul(d)
              │   TileLangHostFastDivmodU32Shift(d)
              ▼                       ▼
     FlattenBuffer → ConfigIndexBitwidth → SplitHostDevice
              host computes M/s and passes them as extra
              kernel params; d still comes from the original
              dynamic shape/scalar parameter
              │
              ▼
     MagicCallHoist (after the final Simplify)
       - Merge duplicate calls by (x, d)
       - div/mod share one quotient
       - Reuse results across condition branches
       - Merge runtime validity guards
              │
              ▼
     CUDA / HIP / C codegen
       valid:   mulhi + shift, r = x - q * d
       invalid: exact floor-div / floor-mod helper

End-to-end example

Given q = x // d; r = x % d where d is a dynamic shape/scalar parameter, invariant within one launch. This launch has d = 7; one device thread has x = 100.

At compile time, LowerMagicDiv does not know d = 7, but proves it is a launch-invariant divisor and rewrites:

q = floordiv(x, d)          m = host_magic_multiplier(d)
r = floormod(x, d)    →     s = host_magic_shift(d)
                            q = tl.magic_div(x, d, m, s)
                            r = tl.magic_mod(x, d, m, s)

After MagicCallHoist, div/mod on the same (x, d) share one quotient and one validity check:

valid = (x ≥ 0) && (1 ≤ d ≤ INT32_MAX)
q = tl.magic_div_with_validity(x, d, m, s, valid)
r = tl.magic_mod_from_quotient(x, d, q, valid)

At launch time, the host sees the actual d = 7 and computes m = 0x92492493, s = 2 once; all threads reuse them:

kernel(..., d=7, tl_magic_m=0x92492493, tl_magic_s=2)

Note: m is mathematically a uint32; TileLang currently transmits its bit pattern through an int32 kernel parameter (0x92492493 reads as negative when interpreted as signed), and the device codegen re-casts to unsigned before __umulhi, so all 32 bits survive intact.

On the device thread with x = 100:

valid  = (100 ≥ 0) && (uint32(7 - 1) ≤ 0x7ffffffe) = true
high32 = __umulhi(uint32(100), uint32(0x92492493)) = 57
q      = 57 >> 2 = 14
r      = x - q * d = 100 - 14 * 7 = 2

Responsibility split:

If one thread's x violates the fast-path contract, only that thread's valid is false and it falls back to the original floor-div/mod — other threads keep the fast path.

PR

#3267 regission case in h200