Merge branch 'master' into simpler_postrange

test/test_rangeify.py

@@ -11,6 +11,26 @@ class TestRangeify(unittest.TestCase):
     ba = A.expand(N, N)
     ((ba+1).sum(axis=1) + (ba+2).sum(axis=0)).realize()
 
+  def test_partial_contig(self):
+    A = Tensor.empty(64, 64, 64)
+    ret = A.sum(axis=2).contiguous(arg=(1,)).sum(axis=1)
+    ret.realize()
+
+  def test_double_gemm_real(self):
+    def go():
+      with Context(DEBUG=0):
+        Tensor.manual_seed(1337)
+        A,B,C = [Tensor.randn(N, N) for _ in range(3)]
+        Tensor.realize(A, B, C)
+      GlobalCounters.reset()
+      return (A@B@C).realize()
+    rng = go()
+    with Context(RANGEIFY=0, DEBUG=2):
+      ref = go()
+      mse = ((rng-ref)**2).sum().item()
+    print(f"mse: {mse}")
+    self.assertLessEqual(mse, 1e-2)
+
   def test_double_gemm(self):
     A = Tensor.empty(N, N)
     B = Tensor.empty(N, N)
@@ -96,10 +116,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, 32, 1024, 64
 
     # llama 8B
     #BS, HEADS, SEQLEN, EMB = 4, 32, 2048, 128

tinygrad/codegen/__init__.py

@@ -9,17 +9,18 @@ from tinygrad.renderer import Renderer
 # import all pattern matchers here
 from tinygrad.codegen.lowerer import pm_lowerer, get_index
 from tinygrad.codegen.quantize import pm_quant
-from tinygrad.codegen.gpudims import pm_add_gpudims
+from tinygrad.codegen.gpudims import pm_add_gpudims, pm_tensor_cores, pm_group_for_reduce, pm_fix_locals, pm_bufferize_loop
 from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing
 from tinygrad.uop.decompositions import get_late_rewrite_patterns
 from tinygrad.codegen.late.expander import migrate_indexing, expander, pm_pre_expander
 from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
   ReduceContext, correct_load_store, pm_render
 from tinygrad.codegen.late.linearize import block_create, pm_blockend_merge, block_merge, pm_finalize, BlockContext
-from tinygrad.codegen.opt import pm_get_optimization, pm_do_optimize
+from tinygrad.codegen.opt import pm_get_optimization, pm_do_optimize, pm_postrange_opt
 from tinygrad.codegen.opt.swizzler import view_left, view_right, fix_kernel_ops
 from tinygrad.codegen.opt.postrange import pm_postrange_opt
 from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen
+from tinygrad.codegen.opt.postrange import pm_flatten_range
 
 @dataclass
 class RewriteStep:
@@ -58,21 +59,32 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC
 
   # this is kernel.py
   if not _RANGEIFY: ret.append(RewriteStep(pm_get_optimization, ctx=lambda _: opts, name="get optimization"))
+
   if not _POSTOPT and not _RANGEIFY: ret.append(RewriteStep(pm_do_optimize, ctx=lambda _: opts, name="optimize ast"))
 
   if _QUANTIZE and opts.device in {"CPU", "DSP"}: ret.append(RewriteStep(pm_quant, name="quantize"))
   ret.append(RewriteStep(pm_lowerer, get_index, name="lowerer", bottom_up=True))
+  
+  # add tensor cores
+  if _RANGEIFY:
+    ret.append(RewriteStep(pm_bufferize_loop, name="bufferize loop"))
+    ret.append(RewriteStep(pm_tensor_cores, lambda _: ({}, opts), name="tensor cores", bottom_up=True))
 
   if _POSTOPT or _RANGEIFY: ret.append(RewriteStep(pm_postrange_opt, ctx=lambda _: opts, name="post optimize ast"))
 
   # ** expander (expand_rewrite) **
   ret.append(RewriteStep(sym+migrate_indexing, name="initial symbolic"))
 
-  # expand
-  ret.append(RewriteStep(sym+pm_pre_expander+expander, name="expander"))
+  ret.append(RewriteStep(pm_fix_locals, name="fix locals"))
 
   # add locals
-  ret.append(RewriteStep(pm_add_buffers_local+rangeify_codegen, name="add local buffers"))
+  ret.append(RewriteStep(pm_flatten_range+pm_add_buffers_local+rangeify_codegen+pm_group_for_reduce, name="add local buffers"))
+
+  # add gpu dims (late). this also handles UNROLL range
+  ret.append(RewriteStep(pm_add_gpudims, lambda _: opts, name="add gpudims"))
+
+  # expand
+  ret.append(RewriteStep(sym+pm_pre_expander+expander, name="expander"))
 
   # ** devectorizer (full_graph_rewrite) **
   # remove reduce

tinygrad/codegen/gpudims.py

@@ -1,9 +1,10 @@
-import math
-from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType
-from tinygrad.helpers import all_int, dedup
-from tinygrad.dtype import dtypes
+import math, functools, operator
+from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType, graph_rewrite
+from tinygrad.helpers import all_int, partition, flatten, prod, dedup, USE_TC, DEBUG
+from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.shape.view import get_contraction
 from tinygrad.renderer import Renderer
+from tinygrad.codegen.opt.tc import TensorCore
 
 def _group_dims(dims:tuple[sint, ...], max_sizes:tuple[int, ...]):
   # TODO: symbolic shape
@@ -88,7 +89,181 @@ def add_gpudims(ctx:Renderer, s:UOp):
     except ValueError: continue
   return s.substitute(subs)
 
+def fix_reduce_unroll(x:UOp):
+  reduce_range, reduce_expand = partition(x.src[1:], lambda y: y.op is Ops.RANGE)
+  if len(reduce_expand) == 0: return None
+  reduce_expand = [x for x in reduce_expand if x.op is not Ops.CONST]
+  assert all(x.op is Ops.UNROLL for x in reduce_expand), f"not all UNROLLS in {reduce_expand}"
+  ret = x.src[0]
+  if len(contract_axis:=flatten(x.arg for x in reduce_expand)):
+    ret = UOp(Ops.CONTRACT, x.dtype.vec(prod(x[1] for x in contract_axis)), (ret,), tuple(contract_axis), tag=1)
+  # REDUCE supports both "horizontal" reduction and range reduction. the horizontal elements are taken in the nearest group
+  return x.replace(src=(ret,)+tuple(reduce_range))
+
+def fix_store_unroll(x:UOp):
+  store_expand, store_range = partition(x.src[2:], lambda y: y.op is Ops.UNROLL)
+  if len(store_expand) == 0: return None
+  return UOp(Ops.CONTRACT, dtypes.void, (x.replace(src=x.src[:2]+tuple(store_range)),), tuple(flatten(x.arg for x in store_expand)), tag=1)
+
+def fix_group_for_reduce(x:UOp):
+  reduce_gfr, reduce_r = partition(x.src[1:], lambda u: u.op is Ops.RANGE and u.arg[1] == AxisType.GROUP_REDUCE)
+  if len(reduce_gfr) == 0: return None
+
+  # NOTE: if there's other locals here, we need them in the buffer too
+  upstream_locals = [u for u in x.toposort() if u.op is Ops.RANGE and u.arg[1] == AxisType.LOCAL]
+
+  # do only the non grouped reduces early
+  ret = x.replace(src=(x.src[0],)+tuple(reduce_r))
+  reduce_loop = [x.replace(arg=(x.arg[0]+100, AxisType.REDUCE)) for x in reduce_gfr]
+  buf = ret.bufferize(*upstream_locals, *reduce_gfr, arg=AddrSpace.LOCAL).index(*upstream_locals, *reduce_loop)
+
+  # gate with an if on the store + do the final reduce
+  buf = UOp(Ops.IF, dtype=buf.dtype, src=(functools.reduce(operator.and_, [x.eq(0) for x in reduce_gfr]), buf))
+  return buf.reduce(*reduce_loop, arg=x.arg)
+
+pm_group_for_reduce = PatternMatcher([
+  # fix group for reduce
+  (UPat(Ops.REDUCE, name="x"), fix_group_for_reduce),
+])
+
 pm_add_gpudims = PatternMatcher([
   # add gpudims must be last
   (UPat(Ops.SINK, name="s"), add_gpudims),
+  # rewrite UPCAST/UNROLL range to something to be expanded
+  (UPat(Ops.RANGE, name="r"),
+   lambda r: UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(s:=r.vmax+1), tuple(range(s))),), ((r.arg[0],s),)) \
+    if r.arg[1] in {AxisType.UNROLL, AxisType.UPCAST} else None),
+  # fix REDUCEs with UNROLLs
+  (UPat(Ops.REDUCE, name="x"), fix_reduce_unroll),
+  (UPat(Ops.STORE, name="x"), fix_store_unroll),
 ])
+
+def apply_tensor_cores(ctx:tuple[dict, Renderer], in0:UOp, in1:UOp, r_range:UOp, reduceop:UOp):
+  if not USE_TC: return None
+  # tensor cores have three ranges. X, Y, and REDUCE
+  in0_ranges = sorted([u for u in in0.ranges if u not in in1.ranges], key=lambda x: x.arg[0])
+  in1_ranges = sorted([u for u in in1.ranges if u not in in0.ranges], key=lambda x: x.arg[0])
+  #print(len(in0_ranges), len(in1_ranges))
+  if not len(in0_ranges) or not len(in1_ranges): return None
+  in0_range, in1_range = in0_ranges[0], in1_ranges[0]
+  if DEBUG >= 2: print('TC', in0_range.arg, in1_range.arg, r_range.arg)
+
+  # confirm the dtype and size is good
+  tc_opts: list[TensorCore] = []
+  for tc in ctx[1].tensor_cores:
+    if reduceop.dtype == tc.dtype_out and in0.dtype == tc.dtype_in and in1.dtype == tc.dtype_in:
+      if all(i <= j for i,j in zip(tc.dims, [in0_range.vmax+1, in1_range.vmax+1, r_range.vmax+1])):
+        tc_opts.append(tc)
+  if len(tc_opts) == 0: return None
+  tc = tc_opts[0]
+
+  # create the new ranges as speced by the tensor core
+  old_range = [in0_range, in1_range, r_range]
+  new_range = [r.replace(src=(r.src[0]//tc.dims[i],)) for i,r in enumerate(old_range)]
+  new_reduce_range = new_range[2]
+  tc_range = 9050 #+ r_range.arg[0]*100
+  red_ranges = []
+
+  ne: list[UOp] = []
+  for o in tc.opts:
+    lrange = UOp.range(dtypes.int, 2, tc_range, AxisType.UPCAST if o[0] == "u" else AxisType.LOCAL)
+    ne.append(lrange)
+    tc_range += 1
+    new_range[1-int(o[1])] = (2 * new_range[1-int(o[1])]) + lrange
+  for _, amt in tc.get_reduce_axes():
+    lrange = UOp.range(dtypes.int, amt, tc_range, AxisType.UNROLL)
+    ne.append(lrange)
+    red_ranges.append(lrange)
+    tc_range += 1
+    new_range[2] = (amt * new_range[2]) + lrange
+  tne = [x.replace(tag=1) for x in ne]
+
+  # replace ranges in other parts of the graph
+  for x,y in zip(old_range, new_range): ctx[0][x] = y
+
+  # apply the swizzled ranges to the srcs
+  srcs = [s.substitute(dict(zip(old_range, new_range))).substitute(dict(zip(ne, tne))) for s in (in0, in1)]
+  srcs = [x.substitute(dict(zip(tne, [ne[i] for i in p]))) for x,p in zip(srcs, tc.permutes_for_shape_str(tc.base_shape_str()))]
+
+  ned = dict(zip(tc.base_shape_str(), ne))
+  tc_reduce_axes = tuple([ned[f"r{i}"].arg[0] for i in range(len(tc.get_reduce_axes()))])
+  base_upcast_axes = tuple([(ned[s].arg[0], 2) for s in tc.base_upcast_axes()])
+  tc_upcast_axes = tuple([base_upcast_axes[:int(math.log2(tc.elements_per_thread[i]))] for i in range(3)])
+
+  # construct the op
+  # TODO: remove tc_upcast_axes from the arg
+  wmma_arg = (str(tc), tc.dims, tc.dtype_in, tc.dtype_out, ctx[1].device, tc.threads, tc_upcast_axes, tc_reduce_axes)
+  wmma = UOp(Ops.WMMA, dtype=tc.dtype_out.vec(tc.elements_per_thread[2]), src=(
+    UOp(Ops.CONTRACT, dtype=srcs[0].dtype.vec(tc.elements_per_thread[0]), src=(srcs[0],), arg=tc_upcast_axes[0]),
+    UOp(Ops.CONTRACT, dtype=srcs[1].dtype.vec(tc.elements_per_thread[1]), src=(srcs[1],), arg=tc_upcast_axes[1]),
+    UOp.const(tc.dtype_out.vec(tc.elements_per_thread[2]), 0.0)), arg=wmma_arg)
+  tc_uop = UOp(Ops.UNROLL, tc.dtype_out, (wmma,), arg=tc_upcast_axes[2])
+  ret = tc_uop.reduce(new_reduce_range, arg=Ops.ADD)
+  # confirm the UNROLLs aren't actually used, these need to be broadcast MUL
+  assert all(u not in red_ranges for u in ret.toposort()), "UNROLLs in TC"
+  return ret
+
+from tinygrad.codegen.opt.postrange import pm_flatten_range
+
+pm_tensor_cores = PatternMatcher([
+  ((UPat.var("in0")*UPat.var("in1")).reduce(UPat(Ops.RANGE, name="r_range"), name="reduceop", arg=Ops.ADD), apply_tensor_cores),
+
+  # replace range
+  #(UPat(Ops.RANGE, name="r"), lambda ctx,r: ctx[0].get(r, None)),
+  (UPat(Ops.SINK, name="s"), lambda ctx,s: graph_rewrite(s.substitute(ctx[0]), pm_flatten_range, name="flatten")),
+])
+
+def fix_bufferize(x:UOp):
+  if x.arg != AddrSpace.LOCAL: return None
+  locals_left = [r for r in x.ranges if r.arg[1] == AxisType.LOCAL]
+  if not len(locals_left): return None
+  acc = x.size
+  st = []
+  for l in locals_left:
+    st.append(l*acc)
+    acc *= l.vmax+1
+  return x.replace(src=(x.src[0],) + tuple(locals_left[::-1]) + x.src[1:]).index(sum(st))
+
+
+pm_double_index = PatternMatcher([
+  (UPat(Ops.INDEX, src=(UPat(Ops.INDEX, src=(UPat.var("b"), UPat.var("x"))), UPat.var("y"))), lambda b,x,y: b.index(x+y)),
+])
+
+pm_fix_locals = pm_double_index+PatternMatcher([
+  (UPat(Ops.BUFFERIZE, name="x"), fix_bufferize),
+])
+
+B = 8
+
+def loop_store(x:UOp):
+  r_maybe = [r for r in x.src[0].ranges if r.arg[0] == 2 and r.tag is None]
+  ur_maybe = [r for r in x.ranges if r.arg[0] < 0]
+  #print("store", len(r_maybe), len(ur_maybe))
+  if not len(r_maybe) or not len(ur_maybe): return None
+
+  r = r_maybe[0]
+  ur = ur_maybe[0]
+  rr = r.replace(src=(r.src[0]//B,), tag=1)
+  return x.substitute({r:rr*B+ur})
+
+def loop_bufferize(x:UOp):
+  if x.arg != AddrSpace.LOCAL: return None
+  r_maybe = [r for r in x.ranges if r.arg[0] == 2 and r.tag is None]
+  ur_maybe = [r for r in x.ranges if r.arg[0] < 0]
+
+  if len(ur_maybe):
+    ur = ur_maybe[0]
+  else:
+    if len(r_maybe) == 0: return None
+    r = r_maybe[0]
+    ur = UOp.range(dtypes.int, B, -2010)
+    rr = r.replace(src=(r.src[0]//B,), tag=1)
+    x = x.substitute({r:rr*B + ur})
+
+  ur1 = UOp.range(dtypes.int, B, ur.arg[0]+1)
+  return x.replace(src=(x.src[0],ur)+x.src[1:]).index(x.size*ur1)
+
+pm_bufferize_loop = pm_double_index+PatternMatcher([
+  (UPat(Ops.BUFFERIZE, name="x"), loop_bufferize),
+  (UPat(Ops.STORE, name="x"), loop_store),
+])

tinygrad/codegen/opt/__init__.py

@@ -1,6 +1,7 @@
 # opt opinionatedly transforms an ast into an optimized ast using either heuristics or beam search
 
 from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.codegen.opt.postrange import RKernel, pm_flatten_range
 from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, KernelInfo
 from tinygrad.helpers import NOOPT, BEAM, USE_TC, getenv
@@ -44,3 +45,15 @@ def apply_opt(ast:UOp, renderer:Renderer):
 pm_do_optimize = PatternMatcher([
   (UPat(Ops.SINK, name="ast"), lambda ctx,ast: apply_opt(ast, ctx) if ast.arg is not None and ast.arg.opts_to_apply is not None else None),
 ])
+
+# ** postrange **
+
+def apply_ropt(ast:UOp, renderer:Renderer):
+  k = RKernel(ast, opts=renderer)
+  if ast.arg is not None: k.apply_opts(ast.arg.opts_to_apply)
+  return k.get_optimized_ast()
+
+pm_postrange_opt = pm_flatten_range+PatternMatcher([
+  (UPat(Ops.SINK, name="ast"), lambda ctx,ast: apply_ropt(ast, ctx) if ast.arg is None or \
+   (ast.arg is not None and ast.arg.opts_to_apply is not None) else None),
+])

tinygrad/codegen/opt/kernel.py

@@ -92,10 +92,6 @@ class Kernel:
     global_loops = AxisType.GLOBAL if self.opts.has_local else AxisType.LOOP
     self.axis_types: list[AxisType] = [AxisType.REDUCE if resolve(x!=y) else global_loops for x,y in zip(self.output_shape, self.full_shape)]
 
-    # confirm all reduce axes are at the end
-    if (final_reduces := [x for x in self.axis_types if x == AxisType.REDUCE]) and final_reduces != self.axis_types[-len(final_reduces):]:
-      raise RuntimeError(f"reduces are not at the end of the shape {self.full_shape} -> {self.output_shape}")
-
   def copy(self):
     ret = type(self).__new__(type(self))
 

tinygrad/codegen/opt/postrange.py

@@ -1,3 +1,168 @@
+import math
+from tinygrad.dtype import AddrSpace
+from tinygrad.uop.ops import UOp, Ops, sint, ssimplify, AxisType, KernelInfo, PatternMatcher, UPat, graph_rewrite, _substitute
+from tinygrad.uop.symbolic import symbolic
+from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
+from tinygrad.renderer import Renderer
+from tinygrad.dtype import dtypes
+
+def flatten_range(r:UOp):
+  off = 2 if r.op is Ops.STORE else 1
+  rngs = r.src[off:]
+  if not len(rngs): return None
+  new_rngs = [x for x in UOp.sink(*rngs).toposort() if x.op is Ops.RANGE]
+  return r.replace(src=r.src[:off]+tuple(new_rngs))
+
+pm_flatten_range = PatternMatcher([
+  # real ranges only
+  (UPat((Ops.REDUCE, Ops.STORE), name="r"), flatten_range),
+])
+
+def count_divmod(x:UOp): return len([u for u in x.toposort() if u.op in {Ops.IDIV, Ops.MOD}])
+
+class RKernel(Kernel):
+  def __init__(self, ast:UOp, opts:Renderer|None=None):
+    self.rng = sorted([u for u in ast.toposort() if u.op is Ops.RANGE and u.vmax > 0], key=lambda x: x.arg)
+    super().__init__(ast, opts)
+    self.sts.clear()
+
+    # convert LOOP to GLOBAL
+    self.replaces = {}
+    if self.opts.has_local:
+      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) \
+                          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 []
+      rng = [x.replace(arg=(x.arg[0], AxisType.GLOBAL)) if x.arg[1] == AxisType.LOOP and x in store_rng else x for x in self.rng]
+      self.replaces.update(dict(zip(self.rng, rng)))
+      self.rng = rng
+
+  def simplify_merge_adjacent(self):
+    return
+    # NOTE: this one is better than the one in kernel.py, which is kind of a problem
+    terminators = [u for u in self.ast.toposort() if u.op in {Ops.REDUCE, Ops.STORE}]
+    termination = {}
+    for t in terminators:
+      for u in t.src[1 if t.op is Ops.REDUCE else 2:]: termination[u] = t
+
+    replaces = {}
+    i = 0
+    while i < len(self.rng)-1:
+      r0, r1 = self.rng[i], self.rng[i+1]
+      # same axistype and same termination
+      if r0.arg[1] == r1.arg[1] and termination[r0] == termination[r1]:
+        s0, s1 = r0.src[0], r1.src[0]
+        new_range = r0.replace(src=(s0*s1,)).simplify()
+        # this checks the legality of a merge
+        oidx = self.ast.simplify()
+        nidx = graph_rewrite(oidx, _substitute+symbolic+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1}, name=f"check_merge_{i}_{i+1}")
+        # it simplifies
+        if count_divmod(nidx) <= count_divmod(oidx):
+          # it is correct
+          midx = graph_rewrite(nidx, _substitute+symbolic+pm_flatten_range, ctx={new_range:r0*s1+r1}, name=f"correct_merge_{i}_{i+1}")
+          if oidx is midx:
+            termination[new_range] = termination[r0]
+            replaces[r0] = new_range//s1
+            replaces[r1] = new_range%s1
+            self.rng[i] = new_range
+            del self.rng[i+1]
+            continue
+      i += 1
+    self.ast = self.ast.substitute(replaces, name="simplify_merge_adjacent")
+
+  def shift_to(self, axis:int, amount:int, new_type:AxisType, top:bool=False, insert_at:int|None=None):
+    old_sz = self.rng[axis].src[0].arg // amount
+    assert old_sz > 0, f"bad old_sz on {axis} {amount} {self.rng[axis]}"
+
+    maxarg = max([x.arg[0] for x in self.rng])
+    new_rng = UOp.range(dtypes.int, amount, maxarg+1, new_type)
+
+    if old_sz == 1:
+      self.replaces[self.rng[axis]] = new_rng
+      self.rng.insert(insert_at if insert_at is not None else len(self.rng), new_rng)
+      del self.rng[axis]
+    else:
+      replaced_rng = self.rng[axis].replace(src=(UOp.const(dtypes.int, old_sz),))
+      self.replaces[self.rng[axis]] = (new_rng * old_sz + replaced_rng) if top else (replaced_rng * amount + new_rng)
+      self.rng[axis] = replaced_rng
+      self.rng.insert(insert_at if insert_at is not None else len(self.rng), new_rng)
+    return new_rng
+
+  @property
+  def axis_types(self) -> list[AxisType]: return [x.arg[1] for x in self.rng]
+  @property
+  def shape_len(self): return len(self.rng)
+
+  @property
+  def full_shape(self) -> tuple[sint, ...]: return tuple([ssimplify(x.src[0]) for x in self.rng])
+  @property
+  def output_shape(self) -> tuple[sint, ...]: return tuple([ssimplify(x.src[0]) for x in self.ast.src[0].src[2:]])
+
+  def get_optimized_ast(self, name_override:str|None=None) -> UOp:
+    ret = self.ast
+    kernel_name = ret.arg.name if ret.arg is not None and ret.arg.name != "test" else self.name if name_override is None else name_override
+    rarg = KernelInfo(kernel_name, tuple(self.axis_types), self.dont_use_locals, tuple(self.applied_opts))
+    return ret.substitute(self.replaces).replace(arg=rarg)
+
+  # does nothing
+  @axis_types.setter
+  def axis_types(self, value): pass
+
+  def _apply_tc_opt(self, use_tensor_cores:int, axis:int, tc_select:int, opt_level:int) -> bool:
+    reduceop = [x for x in self.ast.toposort() if x.op is Ops.REDUCE][0]
+    if use_tensor_cores and reduceop is not None and reduceop.arg is Ops.ADD:
+      tensor_cores = self.opts.tensor_cores if tc_select == -1 else [self.opts.tensor_cores[tc_select]]
+      for tc in tensor_cores:
+        if tc.dtype_in == dtypes.float and tc.dtype_out == dtypes.float:
+          axes = [1,0]
+
+          # do optimizations and save the ranges
+          ne: list[UOp] = []
+          for opt in tc.opts:
+            ne.append(self.apply_opt(Opt({"u":OptOps.UPCAST, "l":OptOps.LOCAL}[opt[0]], axes[int(opt[1])], 2), append_opt=False))
+          for _, amt in tc.get_reduce_axes():
+            ne.append(self.apply_opt(Opt(OptOps.UNROLL, 0, amt), append_opt=False)) # TODO: this should be the reduce, not 0
+
+          # early realize for TC
+          self.ast = self.ast.substitute(self.replaces)
+          self.replaces = {}
+
+          # fix the srcs
+          reduceop = [x for x in self.ast.toposort() if x.op is Ops.REDUCE][0]
+          tne = [x.replace(tag=1) for x in ne]
+          ret = reduceop.substitute(dict(zip(ne, tne)))
+          srcs = list((ret.src[0] if ret.src[0].op is not Ops.CAST else ret.src[0].src[0]).src)
+          srcs = [x.substitute(dict(zip(tne, [ne[i] for i in p]))) for x,p in zip(srcs, tc.permutes_for_shape_str(tc.base_shape_str()))]
+
+          # get reduce/upcast axes for the tensor cores
+          tc_reduce_axes = self.shape_str_to_axis([f"r{i}" for i in range(len(tc.get_reduce_axes()))])
+          base_upcast_axes = tuple([(s,2) for s in self.shape_str_to_axis(tc.base_upcast_axes())])
+          tc_upcast_axes = tuple([base_upcast_axes[:int(math.log2(tc.elements_per_thread[i]))] for i in range(3)])
+
+          # axes to range number (was done in lowerer)
+          tc_upcast_axes = tuple([tuple([(self.rng[a].arg[0], sz) for a,sz in v]) for v in tc_upcast_axes])
+          tc_reduce_axes = tuple([self.rng[a].arg[0] for a in tc_reduce_axes])
+
+          # construct the op
+          # TODO: remove tc_upcast_axes from the arg
+          wmma_arg = (str(tc), tc.dims, tc.dtype_in, tc.dtype_out, self.opts.device, tc.threads, tc_upcast_axes, tc_reduce_axes)
+          wmma = UOp(Ops.WMMA, dtype=tc.dtype_out.vec(tc.elements_per_thread[2]), src=(
+            UOp(Ops.CONTRACT, dtype=srcs[0].dtype.vec(tc.elements_per_thread[0]), src=(srcs[0],), arg=tc_upcast_axes[0]),
+            UOp(Ops.CONTRACT, dtype=srcs[1].dtype.vec(tc.elements_per_thread[1]), src=(srcs[1],), arg=tc_upcast_axes[1]),
+            UOp.const(tc.dtype_out.vec(tc.elements_per_thread[2]), 0.0)), arg=wmma_arg)
+          tc_uop = UOp(Ops.UNROLL, tc.dtype_out, (wmma,), arg=tc_upcast_axes[2])
+
+          # preserve extra reduces
+          reduce_ranges = [x for x in UOp.sink(*reduceop.src[1:]).toposort() if x.op is Ops.RANGE and x.arg[0] not in tc_reduce_axes]
+          if len(reduce_ranges): tc_uop = UOp(Ops.REDUCE, tc_uop.dtype, (tc_uop,)+tuple(reduce_ranges), Ops.ADD)
+          self.ast = self.ast.substitute({reduceop: tc_uop})
+          return True
+    return False
+
 from dataclasses import replace
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo
 from tinygrad.helpers import colored

tinygrad/renderer/cstyle.py

@@ -316,7 +316,9 @@ class MetalRenderer(CStyleLanguage):
 
   def render_kernel(self, function_name, kernel, bufs, uops, prefix=None):
     prefix = ["#include <metal_stdlib>","using namespace metal;"]
-    for name, _, dtype_in, dtype_out, _, _, _, _ in wmma_args(uops): prefix.append(
+    wargs = wmma_args(uops)
+    if len(wargs) > 0: wargs = wargs[0:1]
+    for name, _, dtype_in, dtype_out, _, _, _, _ in wargs: prefix.append(
   f"""{(dstr_out:=self.render_dtype(dtype_out.vec(2)))} __{name}({(dstr_in:=self.render_dtype(dtype_in.vec(2)))} a, {dstr_in} b, {dstr_out} c){{
   simdgroup_{self.render_dtype(dtype_in)}8x8 mat_a, mat_b; simdgroup_{self.render_dtype(dtype_out)}8x8 mat_c;
   mat_a.thread_elements()[0] = a[0]; mat_b.thread_elements()[0] = b[0]; mat_c.thread_elements()[0] = c[0];

tinygrad/schedule/rangeify.py

@@ -189,6 +189,7 @@ def map_partial_contiguous(ctx:RangeifyContext, x:UOp, idx:UOp):
     ranges.append(ctx.new_range(s) if resolve(s!=1) else UOp.const(dtypes.int, 0))
     new_ranges.append(ranges[-1])
   ret = x.src[0].index(*ranges).bufferize(*[x for x in new_ranges if x.op is not Ops.CONST], arg=x.device)
+  if len(ret.ranges): ret = ret.replace(arg=AddrSpace.LOCAL) # if some ranges are still open, this has to be LOCAL
   return ret.index(*passthrough_idx)
 
 def map_contiguous(ctx:RangeifyContext, x:UOp):
@@ -238,7 +239,10 @@ def index_child(ctx:RangeifyContext, c:UOp, x:UOp, idx:UOp):
   # index based on the shared ranges
   ret = c.index(*out_rngs)
   # if all ranges aren't the same between children, we have to bufferize
-  if len(idx_ranges) > 0: ret = ret.bufferize(*end_ranges, arg=x.device).index(*[idx.src[1+i] for i in idx_ranges])
+  if len(idx_ranges) > 0:
+    ret = ret.bufferize(*end_ranges, arg=x.device)
+    if len(ret.ranges): ret = ret.replace(arg=AddrSpace.LOCAL) # if some ranges are still open, this has to be LOCAL
+    ret = ret.index(*[idx.src[1+i] for i in idx_ranges])
   return ret
 
 def children_gate(ctx:RangeifyContext, idx:UOp, c:UOp):
@@ -334,7 +338,7 @@ def bufferize_to_store(x:UOp, locals_allowed=False):
   shape = tuple([int(r.vmax+1) for r in rngs])
   size = prod(shape)
   assert size > 0, f"no zero sized buffers {shape}"
-  sdtype = x.dtype.ptr(size=size, addrspace=AddrSpace.GLOBAL if not isinstance(x.arg, tuple) else x.arg[0])
+  sdtype = x.dtype.ptr(size=size, addrspace=AddrSpace.GLOBAL if not isinstance(x.arg, AddrSpace) else x.arg)
   if x.src[0].op is Ops.ASSIGN:
     assign_target, assign_src = x.src[0].src
     assert assign_target.op is Ops.INDEX
@@ -344,7 +348,7 @@ def bufferize_to_store(x:UOp, locals_allowed=False):
     buf = UOp.new_buffer(x.arg, size, x.dtype)
   else:
     if not locals_allowed: return None
-    buf = UOp(Ops.DEFINE_LOCAL, sdtype, arg=x.arg[1])
+    buf = UOp(Ops.DEFINE_LOCAL, sdtype, arg=0) #UOp.unique().arg)
   return buf.reshape(shape).index(*rngs, dtype=sdtype).store(x.src[0], *rngs, dtype=sdtype).forced_reshape(shape, dtype=x.dtype)
 
 pm_add_buffers_local = pm_mops+PatternMatcher([

tinygrad/uop/ops.py

@@ -142,6 +142,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     if self.op is Ops.INDEX and self.src[0].op in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG,
                                                    Ops.BUFFER, Ops.BUFFERIZE, Ops.VECTORIZE, Ops.STORE}:
       return None
+    if self.op is Ops.BARRIER: return None
     if self.op in GroupOp.Block: return None
     from tinygrad.shape.shapetracker import ShapeTracker
     # VIEW and MovementOps define a new ShapeTracker from the arg
@@ -206,8 +207,10 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     ret: dict[UOp, None] = {}
     if self.op in range_start.keys():
       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]
+      delete_ranges = self.src[range_start[self.op]:]
+      if len(delete_ranges):
+        for s in UOp.sink(*delete_ranges).ranges:
+          if s in ret: del ret[s]
     else:
       for s in self.src: ret.update(s.ranges)
     return ret