cleanups

test/test_rangeify.py

@@ -96,10 +96,10 @@ class TestRangeify(unittest.TestCase):
     out.realize()
 
   def test_flash_attention(self):
-    BS, HEADS, SEQLEN, EMB = 4, 2, 16, 8
+    #BS, HEADS, SEQLEN, EMB = 4, 2, 16, 8
 
     # bigger
-    #BS, HEADS, SEQLEN, EMB = 4, 16, 128, 64
+    BS, HEADS, SEQLEN, EMB = 4, 16, 128, 64
 
     # llama 8B
     #BS, HEADS, SEQLEN, EMB = 4, 32, 2048, 128

tinygrad/codegen/opt/postrange.py

@@ -32,7 +32,8 @@ class RKernel(Kernel):
       store_rngs = self.ast.src[0].src[2:]
 
       # filter any not in local stores
-      local_store_rngs = [x.ranges for x in self.ast.toposort() if x.op is Ops.STORE and x.src[0].dtype.addrspace == AddrSpace.LOCAL]
+      local_store_rngs = [x.ranges for x in self.ast.toposort() if (x.op is Ops.STORE and x.src[0].dtype.addrspace == AddrSpace.LOCAL) \
+                          or (x.op is Ops.BUFFERIZE and x.arg == AddrSpace.LOCAL)]
       for ls in local_store_rngs: store_rngs = [x for x in store_rngs if x in ls]
 
       store_rng = [x for x in UOp.sink(*store_rngs).toposort() if x.op is Ops.RANGE] if store_rngs else []

tinygrad/schedule/rangeify.py

@@ -356,7 +356,7 @@ pm_add_buffers_local = pm_mops+PatternMatcher([
 ])
 
 pm_add_buffers = pm_mops+PatternMatcher([
-  (UPat(Ops.BUFFERIZE, name="x"), lambda x: bufferize_to_store(x)),
+  (UPat(Ops.BUFFERIZE, name="x"), bufferize_to_store),
 
   # move RESHAPEs through MSELECT/MSTACK
   (UPat((Ops.MSELECT, Ops.MSTACK), src=UPat(Ops.RESHAPE), name="m"),

tinygrad/uop/ops.py

@@ -209,12 +209,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
       for s in self.src[:range_start[self.op]]: ret.update(s.ranges)
       for s in self.src[range_start[self.op]:]:
         if s in ret: del ret[s]
-    elif self.op in {Ops.WMMA}:
-      ret = self.src[0].ranges.copy()
-      ret.update(self.src[1].ranges)
-      ret.update(self.src[2].ranges)
-      for s in self.src[3:]:
-        if s in ret: del ret[s]
     elif self.op in {Ops.BARRIER}:
       ret = {x:None for x in self.src[0].ranges if x.arg[1] != AxisType.LOCAL}
     else: