sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

FLOORDIV to SHR for powers of 2 [pr] (#17669)

Browse Source
This commit is contained in:
chenyu
2026-08-21 16:28:04 -04:00
committed by GitHub
parent 12f889aaad
commit daa154aa22
2 changed files with 9 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/null/test_uops.py
+4 -3
View File
@@ -301,9 +301,9 @@ class TestFastIdiv(unittest.TestCase):
self.assertNotIn(Ops.CMOD, ops, f"For dtype={dt} FLOORMOD by pow2 left a MOD")
self.assertNotIn(Ops.FLOORMOD, ops, f"For dtype={dt} FLOORMOD survived past late rewrite")
def test_floordiv_power_of_two_uint(self):
# uint FLOORDIV by a power of two lowers to a shift, leaving no IDIV/FLOORDIV in the kernel
for dt in (dtypes.uint32, dtypes.uint64):
def test_floordiv_power_of_two(self):
# FLOORDIV by a power of two lowers to a shift, with no round toward zero correction (a shift is exactly floor division)
for dt in (dtypes.int32, dtypes.uint32, dtypes.int64, dtypes.uint64):
g = UOp.param(0, dt, (3,))
c = UOp.const(2).cast(dt)
a = UOp(Ops.FLOORDIV, dt, (g.index(c), c))
@@ -311,6 +311,7 @@ class TestFastIdiv(unittest.TestCase):
ops = [x.op for x in uops]
self.assertIn(Ops.SHR, ops, f"For dtype={dt} FLOORDIV by power of two did not simplify to shift")
self.assertNotIn(Ops.CDIV, ops, f"For dtype={dt} FLOORDIV by power of two did not simplify to shift")
self.assertNotIn(Ops.CMOD, ops, f"For dtype={dt} FLOORDIV by pow2 kept the round toward zero correction")
self.assertNotIn(Ops.FLOORDIV, ops, f"For dtype={dt} FLOORDIV survived past late rewrite")
@Context(DISABLE_FAST_IDIV=0)
tinygrad/codegen/decomp/op.py
+5 -1
View File
@@ -75,7 +75,11 @@ powers_of_two: dict[int, int] = {2**i:i for i in range(64)}
@functools.cache
def get_simplifying_rewrite_patterns(ops:tuple[Ops, ...]) -> PatternMatcher:
# these are rewrites that make things simpler
pat: list[tuple[UPat, Callable]] = [(UPat.var("a")//UPat.var("b"), floordiv_to_idiv)]
pat: list[tuple[UPat, Callable]] = []
# FLOORDIV by 2**y -> x >> y (an arithmetic shift is exactly floor division for any sign); fires before floordiv_to_idiv
if Ops.SHR in ops: pat.append((UPat.var("x", dtypes.ints)//UPat.cvar("c"),
lambda x,c: x >> v if (v:=powers_of_two.get(c.val, 0)) else None))
pat.append((UPat.var("a")//UPat.var("b"), floordiv_to_idiv))
# FLOORMOD by 2**y -> x & (2**y-1) (correct floor mod for any sign in two's complement); fires before floormod_to_mod
if Ops.AND in ops: pat.append((UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.val-1) if c.val in powers_of_two else None))
pat.append((UPat.var("a")%UPat.var("b"), floormod_to_mod))