better cost function

tinygrad/schedule/rangeify.py

@@ -359,10 +359,12 @@ pm_rangeify = pm_mops+PatternMatcher([
 # *****************
 # 3.5 cleanups
 
+ALWAYS_RUN_OPS = {Ops.CONTIGUOUS, Ops.COPY, Ops.ASSIGN}
+
 # you don't know in the first pass if axes are going to die, this happens if there's an EXPAND to the left
 def cleanup_dead_axes(b:UOp):
-  # if it's user contiguous or assigned to something, we don't touch it
-  if b.src[0].op in {Ops.CONTIGUOUS, Ops.COPY, Ops.ASSIGN}: return None
+  # don't optimize ALWAYS_RUN_OPS
+  if b.src[0].op in ALWAYS_RUN_OPS: return None
 
   new_rng = []
   hit = False
@@ -388,17 +390,32 @@ def remove_bufferize(src:UOp, buf:UOp, idx:UOp):
   assert all(x.op is Ops.RANGE for x in buf.src[1:])
 
   # if it's user contiguous, we never remove it
-  if src.op is Ops.CONTIGUOUS: return None
+  if src.op in ALWAYS_RUN_OPS: return None
 
-  # const reduce is okay
-  def okay_reduce(x:UOp): return all(y.op not in {Ops.BUFFER, Ops.COPY} for y in x.sparents)
-
-  # here is where we compute the cost
-  # for now just no REDUCE, COPY, or ASSIGN
-  ran = src.toposort(gate=lambda x: x.op not in {Ops.INDEX})
   # we don't want to bufferize threefry, also causes problems because not all platforms support long
-  if any(x.op in {Ops.REDUCE, Ops.COPY, Ops.BUFFER_VIEW, Ops.ASSIGN} and not okay_reduce(x) for x in ran) and src.op is not Ops.THREEFRY: return None
+  if src.op is not Ops.THREEFRY:
+    # *** here is where we compute the cost ***
+    # if we return None, the bufferize is kept
 
+    accessed_buffers = []
+    def red_gate(x):
+      if x.op is Ops.INDEX:
+        accessed_buffers.append(x)
+        return False
+      return True
+    ran = src.toposort(gate=red_gate)
+
+    # if this is generated from multiple buffers, don't remove this buffer
+    if len(dedup([x.src[0] for x in accessed_buffers])) > 2: return None
+
+    # const reduce is okay
+    # TODO: move the reduce folder to before this to prevent the need for this
+    def okay_reduce(x:UOp): return all(y.op not in {Ops.BUFFER, Ops.COPY} for y in x.sparents)
+
+    # always run this list of ops
+    if any(x.op is Ops.REDUCE and not okay_reduce(x) for x in ran): return None
+
+  # if it makes it here, the bufferize is removed
   # this is the ranges replaced
   return src.substitute(dict(zip(buf.src[1:], idx.src[1:])))