diff --git a/test/null/test_uops.py b/test/null/test_uops.py index a21fae30a8..43cd44de4f 100644 --- a/test/null/test_uops.py +++ b/test/null/test_uops.py @@ -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) diff --git a/tinygrad/codegen/decomp/op.py b/tinygrad/codegen/decomp/op.py index a23142809f..4e1d5f81c4 100644 --- a/tinygrad/codegen/decomp/op.py +++ b/tinygrad/codegen/decomp/op.py @@ -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))