fix rand arange folding (#10060)

* test rand range

* --amend

* fix rand arange folding

* reduce_rangeless fix

test/test_arange.py

@@ -73,6 +73,22 @@ class TestArange(unittest.TestCase):
   def test_all_opts_w_upcast_and_unroll(self):
     return self.test_all_opts([Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UNROLL, 0, 4)], [Opt(op=OptOps.GROUP, axis=0, arg=0)])
 
+class TestRand(unittest.TestCase):
+  def test_fused_rand_less_ops(self, noopt=1):
+    GlobalCounters.reset()
+    with Context(FUSE_ARANGE=0, NOOPT=noopt):
+      out = Tensor.rand(16384)
+      out.realize()
+    unfused_ops = GlobalCounters.global_ops
+
+    GlobalCounters.reset()
+    with Context(FUSE_ARANGE=1, NOOPT=noopt):
+      out = Tensor.rand(16384)
+      out.realize()
+    print(f"fused {GlobalCounters.global_ops} unfused {unfused_ops}")
+    self.assertLessEqual(GlobalCounters.global_ops, unfused_ops*2)
+  def test_fused_rand_less_ops_opt(self): self.test_fused_rand_less_ops(0)
+
 DSET, DDIM = 2048, 32
 
 class TestIndexing(unittest.TestCase):

tinygrad/codegen/devectorizer.py

@@ -332,15 +332,6 @@ def reduce_to_acc(ctx:ReduceContext, red:UOp):
   ret = functools.reduce(lambda x,y: x.alu(red.arg, y), lst)
   return acc.assign(ret) if len(reduce_range) != 0 else ret
 
-def reduce_unparented(red:UOp):
-  if red.arg not in {Ops.ADD, Ops.MAX}: return None
-  reduce_parented, reduce_unparented = partition(red.src[1:], lambda x: x in red.src[0].sparents)
-  if len(reduce_unparented) == 0: return None
-  ret = red.replace(src=(red.src[0],)+tuple(reduce_parented)) if len(reduce_parented) or red.dtype != red.src[0].dtype else red.src[0]
-  if red.arg is Ops.ADD:
-    for r in reduce_unparented: ret = ret * (r.src[1]-r.src[0]).cast(ret.dtype.scalar()).broadcast(ret.dtype.count)
-  return ret
-
 def no_vectorized_reduce(inp:UOp, red:UOp):
   if inp.dtype != red.dtype:
     # NOTE: [0 1 2 3 4 5 6 7] -> [0+4, 1+5, 2+6, 3+7]
@@ -371,6 +362,18 @@ def index_fold(buf:UOp, r:UOp, idx:UOp, r2:UOp) -> UOp|None:
   base_idx = (idx-r2.src[0])//r2.src[2]  # indexed from 0 to the length of the range
   return buf.index(base_idx.cast(r.dtype)*r.src[2] + r.src[0], (idx >= r2.src[0]) & (idx < r2.src[1]))
 
+def reduce_rangeless(red:UOp):
+  # TODO: share code with reduce_unparented
+  if red.arg not in {Ops.ADD, Ops.MAX}: return None
+  if red.src[0].dtype != red.dtype: return None
+  if any(x.op in {Ops.RANGE} for x in red.src[0].toposort()): return None
+  ret = red.src[0]
+  if red.arg is Ops.ADD:
+    for r in red.src[1:]:
+      total = (r.src[1]-r.src[0]+r.src[2]-1) // r.src[2] # real count in the range
+      ret = ret * total.cast(ret.dtype.scalar()).broadcast(ret.dtype.count)
+  return ret
+
 pm_reduce_collapse = PatternMatcher([
   # put third arg in range
   (UPat(Ops.RANGE, src=(UPat.var(), UPat.var()), name="r"), lambda r: r.replace(src=r.src+(UOp.const(r.dtype, 1),))),
@@ -386,7 +389,8 @@ pm_reduce_collapse = PatternMatcher([
   # devectorize REDUCE
   (UPat(Ops.VECTORIZE, name="inp").reduce(name="red", allow_any_len=True), no_vectorized_reduce),
   # REDUCE on ADD
-  ((UPat.var("x")+UPat.var("y")).reduce(arg=Ops.ADD, allow_any_len=True), lambda x,y: x.reduce(arg=Ops.ADD) + y.reduce(arg=Ops.ADD)),
+  ((UPat.var("x")+UPat.var("y")).reduce(arg=Ops.ADD, allow_any_len=True, name="r"),
+   lambda x,y,r: x.reduce(*r.src[1:], arg=Ops.ADD) + y.reduce(*r.src[1:],arg=Ops.ADD)),
   # MUL casted bool
   ((UPat.var("x") * UPat.var("gate", dtype=dtypes.bool).cast()), lambda x,gate: gate.where(x, 0)),
   # WHERE on LOAD (works on max too)
@@ -401,8 +405,12 @@ pm_reduce_collapse = PatternMatcher([
   (UPat((Ops.INDEX, Ops.LOAD), name="alu"), no_vectorized_alu),
   # cast on RANGE (fix torch indexing)
   (UPat(Ops.RANGE, name="r").cast(name="c"), lambda r,c: r.replace(src=tuple([x.cast(c.dtype) for x in r.src]), dtype=c.dtype)),
-  # remove any ranges from a REDUCE that aren't referenced in the reduce source
-  (UPat(Ops.REDUCE, name="red"), reduce_unparented),
+  # AND on WHERE
+  ((UPat.any(UPat(Ops.DEFINE_VAR, name="x"), UPat(Ops.DEFINE_VAR).gep(name="x")) & UPat.var("y")) \
+   .where(UPat.cvar("c"), 0).reduce(arg=Ops.ADD, allow_any_len=True, name="r"),
+    lambda x,y,c,r: y.where(c, 0).reduce(*r.src[1:], arg=Ops.ADD)*x.cast(c.dtype)),
+  # remove REDUCEs that no longer have a RANGE in the src
+  (UPat(Ops.REDUCE, name="red"), reduce_rangeless),
 ])+sym
 
 def reduce_collapse(red:UOp):
@@ -415,9 +423,19 @@ def reduce_collapse(red:UOp):
         replaces[s] = UOp(Ops.DEFINE_VAR, dtype=s.dtype, arg=(f'in{len(replaces)}', s.vmin, s.vmax))
   collapse_fxn = red.substitute(replaces)
   sink = graph_rewrite(collapse_fxn, pm_reduce_collapse+devectorize, name="reduce_collapse")
+  # TODO: why is REDUCE needed here and just RANGE isn't enough?
   if any(x.op in {Ops.REDUCE, Ops.RANGE} for x in sink.toposort()): return None
   return sink.substitute({v:k for k,v in replaces.items()})
 
+def reduce_unparented(red:UOp):
+  if red.arg not in {Ops.ADD, Ops.MAX}: return None
+  reduce_parented, reduce_unparented = partition(red.src[1:], lambda x: x in red.src[0].sparents)
+  if len(reduce_unparented) == 0: return None
+  ret = red.replace(src=(red.src[0],)+tuple(reduce_parented)) if len(reduce_parented) or red.dtype != red.src[0].dtype else red.src[0]
+  if red.arg is Ops.ADD:
+    for r in reduce_unparented: ret = ret * (r.src[1]-r.src[0]).cast(ret.dtype.scalar()).broadcast(ret.dtype.count)
+  return ret
+
 pm_reduce = PatternMatcher([
   # remove any ranges from a REDUCE that aren't referenced in the reduce source
   (UPat(Ops.REDUCE, name="red"), reduce_unparented),

tinygrad/ops.py

@@ -772,7 +772,7 @@ class UPat(MathTrait):
   def view(self, st=None, **kwargs): return UPat(Ops.VIEW, self.dtype, (self,), st, **kwargs)
   def cast(self, dtype=None, **kwargs): return UPat(Ops.CAST, dtype, (self,), **kwargs)
   def bitcast(self, dtype=None): return UPat(Ops.BITCAST, dtype, (self,))
-  def gep(self, i:int): return UPat(Ops.GEP, None, (self,), (i,))
+  def gep(self, i:int|None=None, **kwargs): return UPat(Ops.GEP, None, (self,), (i,) if i is not None else None, **kwargs)
   def load(self, *src:UPat, **kwargs): return UPat(Ops.LOAD, src=(self,)+src, **kwargs)
   def store(self, *src:UPat, **kwargs): return UPat(Ops.STORE, dtypes.void, (self,)+src, **kwargs)
   def assign(self, x:UPat, **kwargs): return UPat(Ops.ASSIGN, self.dtype, (self,x), **kwargs)