From 9c4bcf8b3bc5818955ccbcd142237711cf2badd2 Mon Sep 17 00:00:00 2001 From: Chen-Yu Yang Date: Tue, 2 Jun 2026 15:11:24 -0400 Subject: [PATCH] TODO fold symbolic arange --- test/backend/test_arange.py | 35 ++++++++++++++++++++++++++ test/null/test_attention.py | 3 ++- test/null/test_uop_symbolic.py | 20 ++++++++++++++- tinygrad/codegen/simplify.py | 3 ++- tinygrad/uop/divandmod.py | 46 +++++++++++++++++++++++++++++++++- tinygrad/uop/symbolic.py | 5 +++- 6 files changed, 107 insertions(+), 5 deletions(-) diff --git a/test/backend/test_arange.py b/test/backend/test_arange.py index 1e223f797d1e4..f7076f64d5f24 100644 --- a/test/backend/test_arange.py +++ b/test/backend/test_arange.py @@ -4,6 +4,7 @@ from tinygrad.helpers import Context, getenv, DEV from tinygrad.engine.realize import run_linear, estimate_uop, compile_linear from tinygrad.renderer.ptx import PTXRenderer +from tinygrad.uop.ops import UOp, Ops from test.helpers import needs_second_gpu class TestArange(unittest.TestCase): @@ -15,10 +16,44 @@ def _get_flops(self, tensor, desired): np.testing.assert_equal(tensor.numpy(), desired) return estimate_uop(linear.src[-1]).ops + def _assert_symbolic_equal(self, tensor, desired): + linear, var_vals = tensor.linear_with_vars() + run_linear(linear, var_vals) + np.testing.assert_equal(tensor._buffer().numpy()[:desired.size].reshape(desired.shape), desired) + return linear, var_vals + def test_arange_complexity(self): self.assertEqual(self._get_flops(Tensor.arange(256), np.arange(256)), 0) self.assertEqual(self._get_flops(Tensor.arange(2560), np.arange(2560)), 0) + def test_symbolic_arange_has_no_reduce(self): + cases = [ + (Tensor.arange(UOp.variable("n", 4, 100).bind(10)), np.arange(10), {"n": 10}), + (Tensor.arange(UOp.variable("n", 4, 100).bind(10)+4), np.arange(14), {"n": 10}), + (Tensor.arange(2, UOp.variable("n", 4, 100).bind(10)), np.arange(2, 10), {"n": 10}), + (Tensor.arange(UOp.variable("n", 4, 10).bind(6), UOp.variable("m", 20, 200).bind(25)), np.arange(6, 25), {"m": 25, "n": 6}), + (Tensor.arange(3, UOp.variable("n", 4, 100).bind(10), 2), np.arange(3, 10, 2), {"n": 10}), + ] + for t, expected, expected_vars in cases: + linear, var_vals = self._assert_symbolic_equal(t, expected) + self.assertEqual(var_vals, expected_vars) + self.assertEqual(sum(u.op is Ops.REDUCE for u in linear.toposort()), 0) + + def test_symbolic_arange_bindings(self): + for n, m in [(4, 20), (6, 25), (10, 33)]: + self._assert_symbolic_equal(Tensor.arange(UOp.variable("n", 4, 10).bind(n)+4), np.arange(n+4)) + self._assert_symbolic_equal(Tensor.arange(2, UOp.variable("n", 4, 10).bind(n)), np.arange(2, n)) + self._assert_symbolic_equal(Tensor.arange(UOp.variable("n", 4, 10).bind(n), UOp.variable("m", 20, 40).bind(m)), np.arange(n, m)) + self._assert_symbolic_equal(Tensor.arange(3, UOp.variable("n", 4, 10).bind(n), 2), np.arange(3, n, 2)) + + def test_symbolic_arange_cumalu(self): + n = UOp.variable("n", 4, 10).bind(6) + self._assert_symbolic_equal(Tensor.arange(n).cumsum(0), np.arange(6).cumsum()) + self._assert_symbolic_equal(Tensor.arange(1, n).cumprod(0), np.arange(1, 6).cumprod()) + values, indices = Tensor.arange(6).cummax(0) + np.testing.assert_equal(values.numpy(), np.arange(6)) + np.testing.assert_equal(indices.numpy(), np.arange(6)) + @unittest.skipIf(Device.DEFAULT == "CL", "flaky in CI") def test_arange_cumsum(self): np.testing.assert_equal(Tensor.arange(513).cumsum(0).numpy(), np.arange(513).cumsum()) diff --git a/test/null/test_attention.py b/test/null/test_attention.py index aa23608de913a..14e16f5ea59d3 100644 --- a/test/null/test_attention.py +++ b/test/null/test_attention.py @@ -30,7 +30,8 @@ def rope_fn(x_in, pos): return apply_rope(x_in, pos) rope_noprune(Tensor.randn(1, 2, 4, 8, dtype=dtypes.float32), v_pos.bind(1)) rope_prune(Tensor.randn(1, 2, 4, 8, dtype=dtypes.float32), v_pos.bind(1)) assert_jit_cache_len(rope_prune, 1) - assert_jit_cache_len(rope_noprune, 3) + # symbolic arange in precompute_freqs_cis no longer needs its cumsum reduction kernel + assert_jit_cache_len(rope_noprune, 2) if __name__ == '__main__': unittest.main() diff --git a/test/null/test_uop_symbolic.py b/test/null/test_uop_symbolic.py index b3fa4862dfd37..53cca09653d50 100644 --- a/test/null/test_uop_symbolic.py +++ b/test/null/test_uop_symbolic.py @@ -4,7 +4,7 @@ from tinygrad.dtype import dtypes, ConstType, DType, Invalid from test.helpers import get_uops -from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer +from tinygrad.uop.ops import UOp, Ops, AxisType, graph_rewrite, sym_infer from tinygrad.uop.symbolic import sym, commutative, pm_simplify_valid, pm_move_where_on_load from tinygrad.uop.validate import uops_to_z3 @@ -83,6 +83,24 @@ def test_lt_factors(self): def test_div_reduction(self): self.helper_test_variable(Variable("a", 2, 3)//2, 1, 1, "1") + def test_symbolic_range_divmod_reduction(self): + n = Variable("n", 1, 100) + n2 = Variable("n2", 2, 100) + r = UOp.range(n, 0, AxisType.REDUCE) + i = UOp.range(n, 1, AxisType.LOOP) + self.assertEqual((((n-1)//(1-2*n))*-1).maximum(1).simplify().render(), "1") + self.assertEqual(((n2-1)//(1-2*n2)).simplify().render(), "-1") + self.assertEqual(((2*n*r+i)%(2*n*n)).simplify().render(), "(n*r0*2+r1)") + self.assertEqual((((2*n*r+i)%(2*n*n))%(2*n-1)).simplify().render(), "(r0+r1)") + m = n+4 + r = UOp.range(m, 2, AxisType.REDUCE) + i = UOp.range(m, 3, AxisType.LOOP) + self.assertEqual((((2*m*r+i)%(2*m*m))%(2*m-1)).simplify().render(), "(r2+r3)") + m = Variable("m", 20, 200)-Variable("n3", 4, 10) + r = UOp.range(m, 4, AxisType.REDUCE) + i = UOp.range(m, 5, AxisType.LOOP) + self.assertEqual((((2*m*r+i)%(2*m*m))%(2*m-1)).simplify().render(), "(r4+r5)") + def test_equality(self): idx1 = Variable("idx1", 0, 3) idx2 = Variable("idx2", 0, 3) diff --git a/tinygrad/codegen/simplify.py b/tinygrad/codegen/simplify.py index 76b668272f88d..cf3df73a3fbde 100644 --- a/tinygrad/codegen/simplify.py +++ b/tinygrad/codegen/simplify.py @@ -132,9 +132,10 @@ def reduce_collapse(red:UOp, u:UOp, pm:PatternMatcher=pm_reduce_collapse) -> UOp included = u.toposort(gate=lambda x: r in x.ranges) if any(x.op in {Ops.STORE, Ops.REDUCE} for x in included): return None replaces: dict[UOp, UOp] = {} + range_sizes = {s for rr in red.src[1:] for s in rr.src} for u in included: for s in u.src: - if s in included or s in replaces or s.op in {Ops.CONST, Ops.PARAM, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR}: continue + if s in included or s in replaces or s in range_sizes or s.op in {Ops.CONST, Ops.PARAM, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR}: continue replaces[s] = UOp.variable(f'in{len(replaces)}', s.vmin, s.vmax, s.dtype) collapse_fxn = u.substitute(replaces).reduce(r, arg=Ops.ADD) sink = graph_rewrite(collapse_fxn, pm, name="reduce_collapse") diff --git a/tinygrad/uop/divandmod.py b/tinygrad/uop/divandmod.py index 17544173894e2..096017240f298 100644 --- a/tinygrad/uop/divandmod.py +++ b/tinygrad/uop/divandmod.py @@ -1,7 +1,34 @@ import functools, itertools, math from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp from tinygrad.dtype import dtypes -from tinygrad.helpers import floordiv, floormod, unwrap +from tinygrad.helpers import floordiv, floormod, partition, unwrap + +def _expand_mul(u:UOp) -> UOp: + def key(t:UOp): + f = t.const_factor() + return (b if f and (b:=t.divides(f)) is not None else t).tuplize + if u.op is Ops.ADD: return UOp.usum(*sorted((p for x in u.split_uop(Ops.ADD) for p in _expand_mul(x).split_uop(Ops.ADD)), key=key)).simplify() + if u.op is Ops.MUL: + a, b = (_expand_mul(x) for x in u.src) + if a.op is Ops.ADD: return UOp.usum(*sorted((p for x in a.split_uop(Ops.ADD) for p in _expand_mul(x*b).split_uop(Ops.ADD)), key=key)).simplify() + if b.op is Ops.ADD: return UOp.usum(*sorted((p for x in b.split_uop(Ops.ADD) for p in _expand_mul(a*x).split_uop(Ops.ADD)), key=key)).simplify() + return u.simplify() + +def sym_lt(x:UOp, y:UOp, inclusive=False) -> bool: + def bound(u:UOp, upper:bool) -> UOp|None: + for r in u.ranges: + u0, u1 = u.substitute({r:r.const_like(0)}), u.substitute({r:r.const_like(1)}) + coeff = _expand_mul(u1-u0) + try: + if (affine:=_expand_mul(u-(u0+coeff*r))).vmin != 0 or affine.vmax != 0: return None + inc, dec = coeff.vmin >= 0, coeff.vmax <= 0 + except ValueError: return None + if not inc and not dec: return None + u = u.substitute({r:(r.src[0]-1 if inc == upper else r.const_like(0))}) + return u + x, y = bound(x, True), bound(y, False) + if x is None or y is None: return False + return _expand_mul(x-y).vmax <= (0 if inclusive else -1) # NOTE: this cache is only on index UOps @functools.cache @@ -14,6 +41,8 @@ def fold_divmod_general(d: UOp) -> UOp|None: if y_min==y_max==0: raise ZeroDivisionError(f"{'Division' if d.op is Ops.FLOORDIV else 'Mod'} by zero trying to rewrite {x.alu(d.op, y)}") if y_min*y_max > 0 and (qv:=floordiv(x_min,y_min)) == floordiv(x_min,y_max) == floordiv(x_max,y_min) == floordiv(x_max,y_max): return x - qv*y if d.op is Ops.FLOORMOD else d.const_like(qv) + if y_min > 0 and x.vmin >= 0 and sym_lt(x, y): return x if d.op is Ops.FLOORMOD else x.const_like(0) + if d.op is Ops.FLOORDIV and x.vmin > 0 and y_max < 0 and sym_lt(x, -y): return x.const_like(-1) # split uops for the rest of the processing x_peeled, const = x.pop_const() @@ -82,6 +111,21 @@ def fold_divmod_general(d: UOp) -> UOp|None: # Reconstruct all uops including const for these checks. all_uops = list(x.split_uop(Ops.ADD)) + if d.op is Ops.FLOORMOD and y_min > 0 and x.vmin >= 0: + rems, changed = [], False + for u in all_uops: + range_parts, factor_parts = partition(u.split_uop(Ops.MUL), lambda p: p.ranges) + term, factor = math.prod(range_parts, start=u.const_like(1)), math.prod(factor_parts, start=u.const_like(1)) + if factor is y: rem_factor = factor.const_like(0) + elif factor.vmin >= 0 and sym_lt(factor, y, inclusive=True): rem_factor = factor + elif (diff:=_expand_mul(factor-y)).vmin >= 0 and sym_lt(diff, y, inclusive=True): rem_factor = diff + else: break + rems.append(term*rem_factor) + changed = changed or rem_factor is not factor + else: + rem = sum(rems, x.const_like(0)) + if changed and rem.vmin >= 0 and sym_lt(rem, y): return rem + # divide_by_gcd: x//y -> (x//gcd)//(y//gcd) gcd = UOp.gcd(*all_uops, y).simplify() if not (gcd.op is Ops.CONST and gcd.arg==1): diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 99d2469f2edd1..dd56d228de863 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -5,7 +5,7 @@ from tinygrad.dtype import ConstType, dtypes, PtrDType, can_lossless_cast, Invalid from tinygrad.helpers import partition, all_same, prod, flatten, get_single_element, unwrap, IMAGE, dedup from tinygrad.uop.decompositions import threefry2x32, xpow -from tinygrad.uop.divandmod import div_and_mod_symbolic +from tinygrad.uop.divandmod import div_and_mod_symbolic, sym_lt # ******** phase 1 of symbolic used to live in ops, it's the most generic folding rules ******** @@ -104,6 +104,7 @@ def fold_add_divmod_recombine(x:UOp) -> UOp|None: (UPat.var("x", dtype=dtypes.bool).logical_not().logical_not(), lambda x: x), (UPat.var("x", dtype=dtypes.bool).where(UPat.const(dtypes.bool, True), UPat.const(dtypes.bool, False)), lambda x: x), (UPat.var("x", dtype=dtypes.bool).where(UPat.const(dtypes.bool, False), UPat.const(dtypes.bool, True)), lambda x: x.logical_not()), + (UPat.var("x", dtype=dtypes.bool) != UPat.cvar("c"), lambda x,c: x.logical_not() if c.arg is True else None), # CAST(bool -> int) != const — CAST(True)=1, CAST(False)=0, so fold based on const value (UPat.var("x", dtype=dtypes.bool).cast(dtypes.ints+(dtypes.weakint,)) != UPat.cvar("c"), lambda x,c: x if c.arg == 0 else x.logical_not() if c.arg == 1 else x.const_like(True)), @@ -260,6 +261,8 @@ def gep_through_wmma(gep:UOp, wmma:UOp) -> UOp|None: (UPat(Ops.RANGE, src=(UPat(Ops.CONST,)), name="x"), lambda x: x.const_like(x.vmin) if x.vmin == x.vmax else None), # max folding (UPat.maximum(UPat.var("x"), UPat.var("y")), lambda x,y: x if x.vmin >= y.vmax else y if x.vmax <= y.vmin else None), + (UPat.maximum((UPat.var("x")//UPat.var("y"))*-1, UPat.cvar("one")), + lambda x,y,one: one if one.arg == 1 and x.vmin >= 0 and y.vmax < 0 and sym_lt(x, -y) else None), # TODO: why does this rule break beautiful_mnist? #((UPat.var("x")+UPat.var("z")).maximum(UPat.var("y")+UPat.var("z")), lambda x,y,z: x.maximum(y) + z), # ** two stage ALU folding **