fix

test/test_rangeify.py

@@ -1,7 +1,8 @@
 import unittest
 from tinygrad import Tensor, nn
-from tinygrad.helpers import RANGEIFY, Context, GlobalCounters
-from tinygrad.uop.ops import UOp
+from tinygrad.helpers import RANGEIFY, Context, GlobalCounters, getenv
+from tinygrad.uop.ops import UOp, Ops
+from tinygrad.codegen.opt import OptOps, Opt
 
 @unittest.skipIf(RANGEIFY<1, "tests only for RANGEIFY")
 class TestRangeifyAssign(unittest.TestCase):
@@ -18,7 +19,7 @@ class TestRangeifyAssign(unittest.TestCase):
     self.assertListEqual(lst, lst3)
     self.assertListEqual(lst2, B.permute(1, 0).tolist())
 
-N = 256
+N = getenv("N", 256)
 
 class TestRangeifyOpt(unittest.TestCase):
   def test_randperm(self):
@@ -93,6 +94,29 @@ class TestRangeify(unittest.TestCase):
     C = Tensor.empty(N, N)
     (A@B@C).realize()
 
+  #@unittest.skip("gemm tc")
+  def test_double_gemm_tc(self):
+    with Context(DEBUG=0):
+      A, B, C = [Tensor.randn(N, N) for _ in range(3)]
+      Tensor.realize(A, B, C)
+    args = ()
+    args += (Opt(OptOps.DEMOTE, 2, 8),)
+    # NOTE: these axes are poorly sorted
+    args += (Opt(OptOps.TC, 0, (0,0,1,1)),)
+    args += (Opt(OptOps.TC, 0, (0,0,1,0)),)
+
+    args += (Opt(OptOps.UPCAST, 0, 2),)
+    args += (Opt(OptOps.UPCAST, 1, 2),)
+    args += (Opt(OptOps.UNROLL, 0, 2),)
+    args += (Opt(OptOps.UNROLL, 1, 2),)
+
+    tst = (A@B@C).contiguous(arg=args).realize()
+    assert tst.uop.base.op is Ops.BUFFER, "buffer"
+    with Context(RANGEIFY=0, DEBUG=2):
+      GlobalCounters.reset()
+      mse = ((A@B@C)-tst).square().mean().item()
+      print(mse)
+
   def test_double_gemm_exp(self):
     A = Tensor.empty(N, N)
     B = Tensor.empty(N, N)
@@ -202,7 +226,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
+
+    # lil bigger
+    BS, HEADS, SEQLEN, EMB = 4, 32, 128, 64
 
     # bigger
     #BS, HEADS, SEQLEN, EMB = 4, 32, 1024, 64
@@ -213,15 +240,31 @@ class TestRangeify(unittest.TestCase):
     def fa():
       Tensor.manual_seed(1337)
       with Context(DEBUG=0): q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)]
-      return q.scaled_dot_product_attention(k, v).realize()
+      return q.scaled_dot_product_attention(k, v)
 
     with Context(DEBUG=4):
       GlobalCounters.reset()
-      ret = fa()
+      args = ()
+      args += (Opt(OptOps.DEMOTE, 5, 8),)
+      args += (Opt(OptOps.TC, 0, (0,0,1,3)),)
+      args += (Opt(OptOps.TC, 0, (0,0,1,0)),)
+      args += (Opt(OptOps.SWAP, 1, 4),)
+      args += (Opt(OptOps.SWAP, 2, 5),)
+      args += (Opt(OptOps.WARP, 4, 32),)
+      args += (Opt(OptOps.WARP, 5, 32),)
+      args += (Opt(OptOps.UNROLL, 1, 8),)
+      args += (Opt(OptOps.UNROLL, 2, 8),)
+      args += (Opt(OptOps.UPCAST, 0, 4),)
+      args += (Opt(OptOps.UPCAST, 1, 4),)
+      args += (Opt(OptOps.UPCAST, 2, 4),)
+      args += (Opt(OptOps.UPCAST, 3, 4),)
+      args += (Opt(OptOps.UNROLL, 0, 4),)
+      args += (Opt(OptOps.UNROLL, 3, 4),)
+      ret = fa().contiguous(arg=args).realize()
     with Context(RANGEIFY=0):
       with Context(DEBUG=2):
         GlobalCounters.reset()
-        cmp = fa()
+        cmp = fa().realize()
       with Context(DEBUG=0):
         mse = ((cmp-ret)**2).sum().item()
     print(f"mse: {mse}")

tinygrad/codegen/__init__.py

@@ -75,14 +75,14 @@ def _get_rewrites_for_renderer(opts:Renderer, optimize:bool, linearizer:bool, _Q
     ret.append(RewriteStep(pm_postrange_opt, ctx=lambda _: opts, name="post optimize ast"))
 
   # ** expander (expand_rewrite) **
-  ret.append(RewriteStep(sym+migrate_indexing, name="postopt symbolic"))
+  ret.append(RewriteStep(sym+migrate_indexing+pm_group_for_reduce, name="postopt symbolic"))
+
+  # add locals
+  ret.append(RewriteStep(pm_add_buffers+rangeify_codegen+pm_flatten_range, name="add local buffers"))
 
   # expand
   ret.append(RewriteStep(sym+pm_pre_expander+pm_group_for_reduce+expander, name="expander"))
 
-  # add locals
-  ret.append(RewriteStep(pm_add_buffers+rangeify_codegen, name="add local buffers"))
-
   # ** devectorizer (full_graph_rewrite) **
   # remove reduce
   ret.append(RewriteStep(pm_reduce+gep_pushing, lambda _: ReduceContext(), name="remove_reduce"))

tinygrad/codegen/opt/__init__.py

@@ -7,6 +7,7 @@ from tinygrad.uop.ops import AxisType
 class OptOps(Enum):
   TC = auto(); UPCAST = auto(); UNROLL = auto(); LOCAL = auto(); THREAD = auto() # noqa: E702
   GROUP = auto(); GROUPTOP = auto(); NOLOCALS = auto(); PADTO = auto(); SWAP = auto() # noqa: E702
+  DEMOTE = auto(); WARP = auto() # noqa: E702
   def __lt__(self, x:OptOps): return self.value < x.value
 
 @dataclass(frozen=True, order=True)

tinygrad/codegen/opt/postrange.py

@@ -68,7 +68,7 @@ class Scheduler:
 
     # 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].ptrdtype.addrspace == AddrSpace.LOCAL) \
-                        or (x.op is Ops.BUFFERIZE and x.arg == AddrSpace.LOCAL)]
+                        or (x.op is Ops.BUFFERIZE and x.arg.addrspace == AddrSpace.LOCAL)]
     for ls in local_store_rngs: store_rngs = tuple([x for x in store_rngs if x in ls])
 
     # filter any not in reduces
@@ -172,14 +172,33 @@ class Scheduler:
         check(not self.dont_use_locals, "can't use locals")
         check(rng.arg[-1] == AxisType.REDUCE, "group is for reduce")
       ret = self.shift_to(rng, amt, opt_to_at[opt.op], top=opt.op in {OptOps.GROUPTOP, OptOps.THREAD})
+    elif opt.op is OptOps.WARP:
+      warp = UOp.range(cast(int, opt.arg), -1, AxisType.WARP)
+      ret = self.shift_to(rng, cast(int, opt.arg), AxisType.WARP, input_new_rng=warp)
+    elif opt.op is OptOps.DEMOTE:
+      _, rr = self.shift_to(rng, cast(int, opt.arg), AxisType.LOOP)
+      def do_demote(ctx, x:UOp):
+        if x.tag is not None: return None
+        mr = ctx[0]
+        nr = mr.replace(arg=ctx[0].arg[0:-2]+(mr.arg[-2]+1, mr.arg[-1]))
+        ctx[0] = nr
+        buf = x.replace(src=(x.src[0], mr)+x.src[1:], tag=1).substitute({mr:nr})
+        return UOp(Ops.APPENDINDEX, dtypes.void, (buf,mr))
+      # do the demotion
+      pm_demote = PatternMatcher([
+        (UPat(Ops.BUFFERIZE, name="x"), do_demote),
+        (UPat(Ops.INDEX, src=(UPat(Ops.APPENDINDEX, name="x"),), name="y", allow_any_len=True),
+         lambda x,y: y.replace(src=(x.src[0],)+x.src[1:]+y.src[1:])),
+      ])
+      self.ast = graph_rewrite(self.ast, pm_demote, ctx=[rr], bottom_up=True, name="demote")
     elif opt.op is OptOps.TC:
-      check(len(self.applied_opts) == 0, "tensor core opts must be first") # TODO: remove the need for this by having warps
+      #check(len(self.applied_opts) == 0, "tensor core opts must be first") # TODO: remove the need for this by having warps
       check(opt.axis is not None, "tensor core opts must have an axis")
-      check(opt.arg is not None and isinstance(opt.arg, tuple) and len(opt.arg) == 3, "tensor core opts must have valid arg")
+      check(opt.arg is not None and isinstance(opt.arg, tuple) and len(opt.arg) >= 3, "tensor core opts must have valid arg")
       check(-1 <= (tc_select:=cast(tuple, opt.arg)[0]) < len(self.opts.tensor_cores), "tensor core opts must have valid tc_select")
       check(0 <= (tc_opt:=cast(tuple, opt.arg)[1]) <= 2, "tensor core opts must have valid tc_opt")
       check(0 < (use_tensor_cores:=cast(tuple, opt.arg)[2]) <= 2, "use_tensor_cores value is not valid")
-      try: ret = self._apply_tc_opt(use_tensor_cores, cast(int, opt.axis), tc_select, tc_opt)
+      try: ret = self._apply_tc_opt(use_tensor_cores, cast(int, opt.axis), tc_select, tc_opt, opt.arg[3] if len(opt.arg) > 3 else 0)
       except ValueError as e: raise KernelOptError(str(e))
       check(ret is not None, "no tensor core available")
     elif opt.op is OptOps.PADTO:
@@ -203,7 +222,7 @@ class Scheduler:
         altrng = self.rngs[opt.arg]
       except IndexError:
         raise KernelOptError
-      check(rng.arg[-1] == AxisType.GLOBAL and altrng.arg[-1] == AxisType.GLOBAL, "swap only for globals")
+      check(rng.arg[-1] == AxisType.LOOP and altrng.arg[-1] == AxisType.LOOP, "swap only for globals")
       self.ast = self.ast.substitute({rng:rng.replace(arg=(*altrng.arg[0:-1], rng.arg[-1]), tag=1),
                                       altrng:altrng.replace(arg=(*rng.arg[0:-1], altrng.arg[-1]), tag=1)})
       self.ast = graph_rewrite(self.ast, remove_tags)
@@ -213,10 +232,10 @@ class Scheduler:
     if append_opt: self.applied_opts.append(opt)
     return ret
 
-  def _apply_tc_opt(self, use_tensor_cores:int, axis:int, tc_select:int, opt_level:int) -> None|list[UOp]:
+  def _apply_tc_opt(self, use_tensor_cores:int, axis:int, tc_select:int, opt_level:int, reduce_choice:int) -> None|list[UOp]:
     reduceops = [x for x in self.ast.toposort() if x.op is Ops.REDUCE]
     if not len(reduceops): raise KernelOptError("no reduce ops for TensorCore")
-    reduceop = reduceops[0]
+    reduceop = reduceops[reduce_choice]
     if use_tensor_cores and reduceop is not None and reduceop.arg is Ops.ADD:
       mul = reduceop.src[0] if reduceop.src[0].op is not Ops.CAST else reduceop.src[0].src[0]
       if mul.op is not Ops.MUL: return None
@@ -261,7 +280,7 @@ class Scheduler:
           ne: list[UOp] = []
           for opt in tc.opts:
             if opt[0] == "l":
-              axes[int(opt[1])], new_range = self.shift_to(axes[int(opt[1])], 2, AxisType.LOCAL, input_new_rng=warp%2)
+              axes[int(opt[1])], new_range = self.shift_to(axes[int(opt[1])], 2, AxisType.WARP, input_new_rng=warp%2)
               warp //= 2
             elif opt[0] == "u":
               axes[int(opt[1])], new_range = self.shift_to(axes[int(opt[1])], 2, AxisType.UPCAST)
@@ -288,6 +307,11 @@ class Scheduler:
             # axes to range number (was done in lowerer)
             tc_upcast_axes = tuple([tuple([(self.rngs[a].arg[0], sz) for a,sz in v]) for v in tc_upcast_axes])
             tc_reduce_axes = tuple([self.rngs[a].arg[0] for a in tc_reduce_axes])
+            #print("ne", [x.arg for x in ne])
+            #print(tc_upcast_axes)
+            #print(tc_reduce_axes)
+            tc_upcast_axes = (((ne[0].arg[0], 2),), ((ne[0].arg[0], 2),), ((ne[0].arg[0], 2),))
+            #print("hack", tc_upcast_axes)
 
             # construct the op
             # TODO: remove tc_upcast_axes from the arg

tinygrad/schedule/rangeify.py

@@ -519,7 +519,7 @@ pm_limit_bufs = PatternMatcher([(UPat(set.union(GroupOp.Binary, GroupOp.Ternary)
 # BUFFERIZE returns the BUFFER ready for INDEXing (doing this will make splitting a lot easier)
 # NOTE: this has been fixed up a bit
 
-def bufferize_to_store(x:UOp):
+def bufferize_to_store(x:UOp, allow_locals=True):
   rngs = x.src[1:]
   shape = tuple([int(r.vmax+1) for r in rngs])
   size = prod(shape)
@@ -552,6 +552,7 @@ def bufferize_to_store(x:UOp):
     return ret.replace(tag=x.tag)
 
   # handle locals
+  if not allow_locals: return None
   tag = x.arg.device
   if tag is None: tag = UOp.unique().arg # TODO: hack
   buf = UOp(Ops.DEFINE_LOCAL, sdtype, arg=tag)
@@ -559,14 +560,22 @@ def bufferize_to_store(x:UOp):
   # TODO: how is this unified?
   return buf.reshape(shape).index(*rngs, dtype=sdtype).store(x.src[0], *rngs, dtype=sdtype).forced_reshape(shape, dtype=x.dtype)
 
-pm_add_buffers = pm_mops+to_bufferview+PatternMatcher([
-  (UPat(Ops.BUFFERIZE, name="x"), bufferize_to_store),
 
+_pm_add_buffers = pm_mops+to_bufferview+PatternMatcher([
   # move RESHAPEs through MSELECT/MSTACK
   (UPat((Ops.MSELECT, Ops.MSTACK), src=UPat(Ops.RESHAPE), name="m"),
    lambda m: m.replace(src=tuple([x.src[0] for x in m.src]), tag=None).reshape(m.src[0].arg).rtag(m.tag)),
 ])
 
+pm_add_buffers_nolocals = PatternMatcher([
+  (UPat(Ops.BUFFERIZE, name="x"), lambda x: bufferize_to_store(x, allow_locals=False)),
+])+_pm_add_buffers
+
+# all local bufferization should happen later
+pm_add_buffers = PatternMatcher([
+  (UPat(Ops.BUFFERIZE, name="x"), bufferize_to_store),
+])+_pm_add_buffers
+
 # *****************
 # 5. split into kernels
 
@@ -741,7 +750,7 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
   if getenv("VIZ"): graph_rewrite(tsink, PatternMatcher([]), name="View Tagged Rangeify")
 
   # bufferize -> store
-  tsink = graph_rewrite(tsink, pm_add_buffers, bottom_up=True, name="bufferize to store")
+  tsink = graph_rewrite(tsink, pm_add_buffers_nolocals, bottom_up=True, name="bufferize to store")
   tsink = graph_rewrite(tsink, split_kernels, ctx=uop_list, name="split kernels")
 
   # if a kernel depends on a buffer, and that buffer is later assigned to, make the assign depend on the kernel's assign

tinygrad/uop/__init__.py

@@ -18,7 +18,7 @@ class Ops(FastEnum):
   COPY = auto(); BUFFER = auto(); BUFFER_VIEW = auto(); MSELECT = auto(); MSTACK = auto() # noqa: E702
 
   # create buffer
-  BUFFERIZE = auto()
+  BUFFERIZE = auto(); APPENDINDEX = auto() # noqa: E702
 
   # ops that adjust the behavior of the scheduler
   CONTIGUOUS = auto(); CONTIGUOUS_BACKWARD = auto(); DETACH = auto(); FUSE = auto() # noqa: E702