TODO fold symbolic arange

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())

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()

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)

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")

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):

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 **