improve fold_unrolled_divs (tinygrad#6977)

addressed tinygrad#6935
the first few terms in fold_unrolled_divs might have been folded already, so the check should first try to add those terms back. there is a case that every but one term is folded which is not an add chain anymore, so just added as a failed test case for now

test/unit/test_uop_symbolic.py

@@ -438,6 +438,21 @@ def test_arange_unrolled4(self):
     unrolled_div = (gidx+2561)//4+(gidx+2562)//4+(gidx+2560)//4+(gidx+2559)//4
     self.helper_test_variable(unrolled_div, 2559, 5118, "(gidx+2559)")
 
+  def test_arange_unrolled4_small(self):
+    gidx = Variable("gidx", 0, 3)
+    unrolled_div = (gidx)//4+(gidx+2)//4+(gidx+3)//4+(gidx+1)//4
+    self.helper_test_variable(unrolled_div, 0, 3, "gidx")
+
+    gidx = Variable("gidx", 0, 2)
+    unrolled_div = (gidx)//4+(gidx+2)//4+(gidx+3)//4+(gidx+1)//4
+    self.helper_test_variable(unrolled_div, 0, 2, "gidx")
+
+    # TODO: fix this, it has only one term and is no longer an add chain
+    with self.assertRaises(AssertionError):
+      gidx = Variable("gidx", 0, 1)
+      unrolled_div = (gidx)//4+(gidx+2)//4+(gidx+3)//4+(gidx+1)//4
+      self.helper_test_variable(unrolled_div, 0, 1, "gidx")
+
   def test_arange_unrolled2(self):
     gidx = Variable("gidx", 0, 2559)
     unrolled_div = (gidx+2559)//2+(gidx+2560)//2+3

tinygrad/ops.py

@@ -823,17 +823,23 @@ def lt_folding(x:UOp, c:int) -> Optional[UOp]:
 def fold_unrolled_divs(divs:UOp):
   # div pattern in unrolled arange
   # example: (x//4+(x+1)//4+(x+2)//4+(x+3)//4 -> x
-  add_chain, seen_const, ans = list(_get_chain(divs, BinaryOps.ADD)), [], None
+  add_chain, denominator, seen_const, ans = list(_get_chain(divs, BinaryOps.ADD)), None, [], None
   for u in add_chain:
-    if not (u.op is UOps.ALU and u.arg is BinaryOps.IDIV and u.src[1].op is UOps.CONST and u.src[1].arg==len(add_chain)): return None
+    if not (u.op is UOps.ALU and u.arg is BinaryOps.IDIV and u.src[1].op is UOps.CONST): return None
+    if denominator is None: denominator = u.src[1].arg
+    if denominator != u.src[1].arg: return None
     # assumed CONST is the last of an ADD
     if (s0:=u.src[0]).op is UOps.ALU and s0.arg is BinaryOps.ADD and s0.src[1].op is UOps.CONST and s0.src[1].op is UOps.CONST:
       seen_const.append(s0.src[1].arg)
       s0 = s0.src[0]
     else: seen_const.append(0)
     if ans is None: ans = s0
     if ans is not s0: return None
-  return ans if ans is not None and sorted(seen_const)==list(range(len(add_chain))) else None
+  if denominator is None: return None
+  # the first (denominator-len(seen_const)) terms may have been folded to 0 already
+  for i in range(denominator-len(seen_const)):
+    if ans is not None and 0 <= ans.vmin and ans.vmax + i < denominator: seen_const.append(i)
+  return ans if ans is not None and sorted(seen_const)==list(range(denominator)) else None
 
 def is_irreducible(u:UOp): return u.op in (UOps.DEFINE_VAR, UOps.SPECIAL, UOps.RANGE)