dragonpilot v0.10.3

tinygrad_repo/tinygrad/codegen/__init__.py

@@ -0,0 +1,142 @@
+from typing import cast
+import itertools
+from tinygrad.helpers import DEVECTORIZE, TRANSCENDENTAL, SPEC
+from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat
+from tinygrad.uop.spec import type_verify, program_spec, kernel_spec
+from tinygrad.renderer import Renderer
+from tinygrad.dtype import dtypes, PtrDType
+from tinygrad.helpers import panic
+
+# import all pattern matchers here
+from tinygrad.codegen.gpudims import pm_add_gpudims
+from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing, symbolic, pm_move_where_on_load
+from tinygrad.uop.decompositions import get_late_rewrite_patterns
+from tinygrad.codegen.late.expander import expander, pm_pre_expander, pm_group_for_reduce
+from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
+  ReduceContext, correct_load_store, pm_render, pm_add_loads
+from tinygrad.codegen.opt.postrange import apply_opts
+from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse, pm_split_store
+from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen, pm_mops
+from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize
+
+pm_syntactic_sugar = PatternMatcher([
+  # INDEX on ptr INDEX concats them
+  (UPat(Ops.INDEX, name="i1").f(Ops.INDEX, name="i2", allow_any_len=True),
+   lambda i1,i2: i2.replace(src=i1.src+i2.src[1:]) if isinstance(i1.dtype, PtrDType) and not isinstance(i2.dtype, PtrDType) else None),
+])
+
+def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -> UOp:
+  if ren is None: ren = Renderer()
+
+  if SPEC: type_verify(sink, kernel_spec)
+
+  # preprocess
+  sink = graph_rewrite(sink, pm_mops+pm_syntactic_sugar, name="early movement ops", bottom_up=True)
+
+  # first we optimize
+  if optimize:
+    # collapse loads reduce (indexing by a tensor)
+    sink = graph_rewrite(sink, pm_load_collapse, name="load collapse")
+
+    # split ranges
+    sink = graph_rewrite(sink, pm_split_ranges+pm_flatten_range, ctx={}, name="split ranges")
+
+    # symbolic (NOTE: this is a requirement for pm_simplify_ranges to be correct)
+    sink = graph_rewrite(sink, sym+pm_flatten_range, name="initial symbolic")
+
+    # optimize (schedule) the AST
+    sink = graph_rewrite(sink, pm_simplify_ranges, name="simplify ranges")
+
+    # split store range (only on CPU for now)
+    sink = graph_rewrite(sink, pm_split_store, ctx=ren.device, name="cut store ranges")
+
+    # do postrange optimization, BEAM or hand_coded_optimizations
+    sink = apply_opts(sink, ren)
+
+  # ** expander (expand_rewrite) **
+  sink = graph_rewrite(sink, sym+pm_move_where_on_load, name="postopt symbolic")
+
+  # expand
+  sink = graph_rewrite(sink, sym+pm_pre_expander+pm_group_for_reduce+expander, name="expander")
+
+  # add locals
+  sink = graph_rewrite(sink, pm_add_buffers_local+rangeify_codegen, ctx=itertools.count(0), name="add local buffers")
+
+  # ** devectorizer (full_graph_rewrite) **
+  # remove reduce
+  sink = graph_rewrite(sink, pm_reduce+gep_pushing, ctx=ReduceContext(), name="remove_reduce")
+
+  # add gpu dims (late). this works after devectorize, but it's faster here
+  sink = graph_rewrite(sink, pm_add_gpudims, ctx=ren, name="add gpudims")
+
+  # **** optimizations are done, now we lower to actual code ****
+
+  # add loads
+  sink = graph_rewrite(sink, pm_add_loads, name="** add loads (code)")
+
+  # devectorize (TODO: does this need opts?)
+  if DEVECTORIZE >= 2: pm_devectorize = sym+load_store_folding+load_store_indexing
+  elif DEVECTORIZE: pm_devectorize = sym+devectorize+load_store_folding+correct_load_store+load_store_indexing
+  else: pm_devectorize = sym+load_store_folding+correct_load_store+load_store_indexing
+  sink = graph_rewrite(sink, pm_devectorize, ctx=ren, name="devectorize")
+
+  # lower the index dtype to a concrete int
+  sink = graph_rewrite(sink, pm_lower_index_dtype+load_store_indexing, ctx=ren.device, name="lower all index dtypes")
+  sink = graph_rewrite(sink, symbolic, name="post index symbolic")
+
+  # optional pre matcher
+  if ren.pre_matcher is not None: sink = graph_rewrite(sink, ren.pre_matcher, name="pre_matcher")
+
+  # decompositions
+  supported_ops = tuple(ren.code_for_op.keys())
+  pm_decomp = symbolic_simple+get_late_rewrite_patterns(supported_ops, TRANSCENDENTAL>=2)
+  sink = graph_rewrite(sink, pm_decomp, ctx=ren.device, name="decompositions")
+
+  # final rules for the renderer (without sym)
+  extra_matcher = ren.extra_matcher if ren.extra_matcher is not None else PatternMatcher([])
+  pm_final_rewrite = pm_decomp+pm_render+extra_matcher+pm_split_ends
+  sink = graph_rewrite(sink, pm_final_rewrite, ctx=ren.device, name="final rewrite")
+
+  # this was the linearizer
+  sink = graph_rewrite(sink, pm_add_control_flow, ctx=CFGContext(sink), name="add control flow", bottom_up=True)
+
+  # return the rewritten sink
+  return sink
+
+# inject IF/ENDIF. only needed if device doesn't support gated stores
+pm_linearize_cleanups = PatternMatcher([
+  # if statements are not allowed in the graph
+  (UPat((Ops.IF, Ops.ENDIF)), lambda: panic(RuntimeError("if not allowed in graph"))),
+  # gated INDEX becomes IF-STORE-ENDIF. this is the only use of IF-ENDIF
+  (UPat(Ops.STORE, name="u", src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat(name="gate", dtype=dtypes.bool))).or_casted(), UPat()),
+        allow_any_len=True), lambda u, gate: (u, [mif:=UOp(Ops.IF, src=(gate, u.src[0])), u, UOp(Ops.ENDIF, src=(mif,))]))
+])
+
+# requires lst be toposorted. like graph rewrite, but for lines
+def line_rewrite(lst:list[UOp], pm:PatternMatcher) -> list[UOp]:
+  newlst = []
+  replaced: dict[UOp, UOp] = {}
+  for u in lst:
+    nu = u.replace(src=tuple([replaced[x] for x in u.src]))
+    ret: tuple[UOp, list[UOp]] = cast(tuple[UOp, list[UOp]]|None, pm.rewrite(nu)) or (nu, [nu])
+    replaced[u] = ret[0]
+    newlst.extend(ret[1])
+  return newlst
+
+def full_rewrite(sink:UOp, ren:Renderer|None=None) -> list[UOp]:
+  """
+  Function to transform the Kernel UOp graph into a linearized program.
+
+  Args:
+    sink: The Ops.SINK rooting the Kernel graph.
+    ren: The Renderer (can change how things are processed, fix this).
+
+  Returns:
+    Linear program in UOps.
+  """
+
+  full_sink = full_rewrite_to_sink(sink, ren, optimize=sink.tag is None)
+  assert len(full_sink.ranges) == 0, f"all ranges must end by the sink, {full_sink.ranges}"
+  lst = line_rewrite(linearize(full_sink), pm_linearize_cleanups)
+  if SPEC: type_verify(lst, program_spec)
+  return lst

tinygrad_repo/tinygrad/codegen/gpudims.py

@@ -0,0 +1,108 @@
+import math, functools, operator
+from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType, sint_to_uop
+from tinygrad.helpers import all_int, dedup, get_contraction
+from tinygrad.dtype import dtypes, AddrSpace, Invalid
+from tinygrad.renderer import Renderer
+
+def _group_dims(dims:tuple[sint, ...], max_sizes:tuple[int, ...]):
+  # TODO: symbolic shape
+  if not all_int(dims): return dims
+  while len(dims) > len(max_sizes) or any(d > m for d,m in zip(dims, max_sizes)):
+    for i,m in enumerate(max_sizes):
+      if i < (len(dims)-1) and dims[i] * dims[i+1] <= m:
+        dims = dims[:i] + (dims[i]*dims[i+1],) + dims[i+2:]
+        break
+    else: return None
+  return dims
+
+def _split_dims(dims, max_sizes):
+  if all(d <= m for d,m in zip(dims, max_sizes)): return dims
+  _dims = list(dims) + [1]*(3-len(dims))
+  for i in range(len(_dims)):
+    while _dims[i] > max_sizes[i]:
+      div = next((d for d in range(2, math.ceil(math.sqrt(_dims[i])) + 1) if (_dims[i] % d) == 0), 1)
+      if div == 1: raise RuntimeError(f"cannot limit dim {dims=}, {max_sizes=}")
+      _dims[i], _dims[(i+1)%len(_dims)] = _dims[i]//div, _dims[(i+1)%len(_dims)]*div
+  return tuple(_dims[:2] if _dims[2] == 1 else _dims[0] if _dims[1:3] == [1,1] else _dims)
+
+def get_grouped_dims(prefix, dims:tuple[sint, ...], max_sizes:tuple[int, ...]|None, reverse=False) -> list[UOp]:
+  if reverse: return get_grouped_dims(prefix, dims[::-1], max_sizes)[::-1]
+  if max_sizes is None: limited = dims
+  else:
+    # try to group first: (a, b, c, d) -> (ab, c, d)
+    limited = grouped if (grouped := _group_dims(dims, max_sizes)) else dims
+    # check if grouping failed
+    if len(limited) > len(max_sizes): raise RuntimeError(f"cannot limit dim {dims=}, {max_sizes=}")
+    # try to split up dims: (a,) -> (b, c)
+    if limited == dims: limited = _split_dims(dims, max_sizes)
+  raw_idxs = [UOp(Ops.SPECIAL, dtypes.index, (sint_to_uop(s),), (f"{prefix}{i}")) for i,s in enumerate(limited)]
+  if len(limited) < len(dims):
+    ret = []
+    if (contraction:=get_contraction(dims, limited)) is None: raise RuntimeError(f"get_contraction should not be None {dims=} {limited=}")
+    for idx, contraction_group in zip(raw_idxs, contraction):
+      for c in contraction_group[:-1]:
+        ret.append(idx % dims[c])
+        idx //= dims[c]
+      ret.append(idx)
+    return ret
+  elif (a:=len(limited)) > (b:=len(dims)):
+    if a == 2 and b == 1: return [raw_idxs[0] * limited[1] + raw_idxs[1]]
+    if a == 3 and b == 1: return [(raw_idxs[0] * limited[1] + raw_idxs[1]) * limited[2] + raw_idxs[2]]
+    if a == 3 and b == 2: return [raw_idxs[0] * limited[1] + raw_idxs[1], raw_idxs[2]]
+  elif limited != dims:
+    # Convert to 1D
+    flat = raw_idxs[0]*limited[1]+raw_idxs[1] if len(dims) == 2 else raw_idxs[0]*(limited[1]*limited[2])+raw_idxs[1]*limited[2]+raw_idxs[2]
+    # Get back original indices from 1D
+    return [flat//dims[1], flat%dims[1]] if len(dims) == 2 else [flat//(dims[2]*dims[1]), (flat//dims[2])%dims[1], flat%dims[2]]
+  return raw_idxs
+
+def add_gpudims(ctx:Renderer, s:UOp):
+  if s.arg is None: return None
+  s_topo = list(s.toposort())
+  if any(x.op is Ops.SPECIAL for x in s_topo): return None
+
+  # get ranges
+  all_ranges = {x.arg[0:-1]:x for x in s_topo if x.op is Ops.RANGE}
+
+  # extract global/local dims
+  global_dims = sorted(dedup([x.arg[0:-1] for x in all_ranges.values() if x.arg[-1] in (AxisType.GLOBAL, AxisType.THREAD)]))
+  local_dims = sorted(dedup([x.arg[0:-1] for x in all_ranges.values() if x.arg[-1] in (AxisType.WARP, AxisType.LOCAL, AxisType.GROUP_REDUCE)]))
+  if not global_dims and not local_dims: return None
+
+  # get global and local shape
+  ranges = [all_ranges[r] for r in global_dims+local_dims if r in all_ranges]
+  global_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg[0:-1] in global_dims])
+  local_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg[0:-1] in local_dims])
+
+  # get the idxs
+  ki: KernelInfo = s.arg
+  if ki.dont_use_locals:
+    assert not local_dims, "can't use locals if there's no local dims"
+    idxs = get_grouped_dims("idx", global_shape, ctx.global_max, reverse=True)
+  else:
+    # define indexes for GPU-like execution
+    idxs = get_grouped_dims("gidx", global_shape, ctx.global_max, reverse=True) + get_grouped_dims("lidx", local_shape, ctx.local_max)
+
+  # apply to multiple ranges
+  subs = {}
+  for r in s_topo:
+    # look for local INDEXes that are not used in the GLOBAL store, then add them as an INVALID
+    if r.op is Ops.STORE and r.buf_target().ptrdtype.addrspace == AddrSpace.GLOBAL:
+      idx = r.src[0]
+      missing_locals = [all_ranges[rng] for rng in local_dims if all_ranges[rng] not in idx.ranges]
+      if len(missing_locals):
+        assert len(idx.src) == 2, "index has 2 sources"
+        mask: UOp = functools.reduce(operator.and_, [x.eq(0) for x in missing_locals])
+        subs[idx] = idx.replace(src=(idx.src[0], mask.broadcast(idx.src[1].dtype.count).where(idx.src[1], Invalid)))
+    if r.op is not Ops.RANGE: continue
+    try:
+      ii = (global_dims+local_dims).index(r.arg[0:-1])
+      if r.arg[1] == AxisType.REDUCE: continue
+      subs[r] = idxs[ii]
+    except ValueError: continue
+  return s.substitute(subs)
+
+pm_add_gpudims = PatternMatcher([
+  # add gpudims must be last
+  (UPat(Ops.SINK, name="s"), add_gpudims),
+])

tinygrad_repo/tinygrad/codegen/late/devectorizer.py

@@ -0,0 +1,328 @@
+from typing import Any, cast
+import functools, operator, itertools
+from collections import defaultdict
+from dataclasses import dataclass
+from tinygrad.dtype import dtypes, ImageDType, DType, AddrSpace, Invalid, PtrDType
+from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, graph_rewrite, GroupOp, identity_element
+from tinygrad.uop.symbolic import uop_given_valid, parse_valid, sym, symbolic, invalid_gate
+from tinygrad.helpers import getenv, flatten, AMX, prod
+from tinygrad.renderer import Renderer
+
+# ***** image load valid simplification *****
+
+def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
+  idx = uop_given_valid(valid, start_idx)
+  if not isinstance(buf.dtype, ImageDType): return None if idx is start_idx else buf.index(idx.valid(valid), ptr=True)
+
+  # wait for it to be image indexed before running simplification
+  if start_idx.dtype.count != 2: return None
+
+  # can drop valid if idx is out of bound when valid is False
+  drop_stmt = []
+  for stmt in valid.split_uop(Ops.AND):
+    if (res:=parse_valid(stmt)) is None: continue
+    X, is_upper_bound, c = res
+
+    # for X0 + X1 + ... >= 1, check if it's out of bound when Xi = 0 for all i
+    if not is_upper_bound and c == 1 and all(u.op in GroupOp.Irreducible and u.vmin == 0 for u in X.split_uop(Ops.ADD)):
+      testidx = functools.reduce(lambda nowidx,u: nowidx.substitute({u:u.const_like(0)}), X.split_uop(Ops.ADD), idx)
+      testidx = testidx.simplify()
+      if testidx.gep(0).vmax < 0 or testidx.gep(1).vmax < 0:
+        drop_stmt.append(stmt)
+        continue
+
+    # if X <= c, check if it's out of bound when X = c+1
+    # if X >= c, check if it's out of bound when X = c-1
+    test_value = c + 1 if is_upper_bound else c - 1
+    for i,b in zip(idx.src, (buf.dtype.shape[1], buf.dtype.shape[0])):
+      if i.is_increasing():
+        rw = i.substitute({X:X.const_like(test_value)}).simplify()
+        if rw.vmin >= b or rw.vmax < 0:
+          drop_stmt.append(stmt)
+          break
+
+  if not drop_stmt and idx is start_idx: return None
+  new_valid = functools.reduce(operator.and_, ss) if (ss:=[s for s in valid.split_uop(Ops.AND) if s not in drop_stmt]) else None
+  return buf.index(idx.valid(new_valid) if new_valid is not None else idx, ptr=True)
+
+
+load_store_indexing = PatternMatcher([
+  # image load valid idx simplification
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), invalid_gate)), lambda buf,x,i,cond: simplify_valid_load(buf, x, cond)),
+  # simplify away long after index has been lowered
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("x", dtypes.long), UPat.var("c", dtypes.bool))), lambda buf,x,c: simplify_valid_load(buf, x, c)),
+  # drop true gate
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("x"), UPat.const(dtypes.bool, True)),), lambda buf,x: buf.index(x, ptr=True)),
+])
+
+# ***** load/store grouping *****
+
+def expand_index(buf:UOp, vec:UOp):
+  if getenv("UNSAFE_DISABLE_MASK", 0): vec = vec.get_idx()
+  # generate the individual indexes
+  midx = graph_rewrite(UOp.sink(*[buf.index(vec.gep(i), ptr=True) for i in range(vec.dtype.count)]),
+                       symbolic+load_store_indexing, name=f"index_buf_{buf.arg}")
+  # extract all the relevant offsets
+  offsets_rootsrc: defaultdict[Any, dict[int, list[int]]] = defaultdict(dict)
+  for i in range(vec.dtype.count):
+    idx: Any = midx.src[i].src[1].get_idx()
+    if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: root_src, arg = idx.src[0], idx.src[1].arg
+    elif idx.op is Ops.ADD and idx.src[0].op is Ops.CONST: root_src, arg = idx.src[1], idx.src[0].arg
+    elif idx.op is Ops.CONST and idx.arg is Invalid: root_src, arg = "INVALID", 0
+    elif idx.op is Ops.CONST: root_src, arg = "CONST", idx.arg
+    else: root_src, arg = idx, 0
+    root_src = (midx.src[i].src[1].get_valid(), root_src)
+    offsets_rootsrc[root_src].setdefault(arg, []).append(i)
+
+  # then rewrite everything we can into groups
+  ret = []
+  idxs: list[int|None] = [None]*vec.dtype.count
+  global_offset = 0
+  for offsets in offsets_rootsrc.values():
+    grouped_offsets = [[x for _,x in group] for _,group in itertools.groupby(enumerate(sorted(offsets.keys())), lambda x: x[1]-x[0])]
+    for grp in grouped_offsets:
+      # get the index offset for this element. using [0] is okay, because they are the same
+      lidx = midx.src[offsets[grp[0]][0]]
+      if len(grp) > 1: lidx = lidx.cast(buf.ptrdtype.base.vec(len(grp)).ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace))
+      # set the idxs of the output
+      for i,g in enumerate(grp):
+        for oo in offsets[g]: idxs[oo] = global_offset+i
+      # add this lidx to the CAT
+      ret.append(lidx)
+      global_offset += len(grp)
+  assert None not in idxs, f"some idxs are missing {idxs}"
+  # this base thing is for image, we want the CAT to be a normal pointer
+  post_cat = UOp(Ops.PTRCAT, buf.ptrdtype.base.ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace).vec(global_offset), tuple(ret))
+  return post_cat.gep(tuple(cast(list[int], idxs)))
+
+def cat_after_store(cat:UOp, data:UOp, sto:UOp):
+  # TODO: this is written in many places
+  offset = 0
+  ret: list[UOp] = []
+  for s in cat.src:
+    ret.append(s.store(data.gep(tuple(range(offset, offset+s.dtype.count))), *sto.src[2:]))
+    offset += s.dtype.count
+  return UOp.group(*ret)
+
+def gep_on_store(gep:UOp, st:UOp, sto:UOp):
+  # NOTE: we need to invert the gep here, but it may be an expanding gep
+  # fake argsort. TODO: handle duplicates
+  a = {}
+  for i,x in enumerate(gep.arg): a[x] = i
+  new_arg = tuple(x[1] for x in sorted(a.items()))
+  return gep.src[0].store(st.gep(new_arg), *sto.src[2:])
+
+load_store_folding = PatternMatcher([
+  (UPat(Ops.INDEX, src=(UPat(Ops.VECTORIZE, src=UPat(GroupOp.Defines).or_after(name="buf")), UPat.var("vec"))), expand_index),
+  # GEP after LOAD
+  (UPat(Ops.LOAD, src=(UPat(Ops.GEP, name="gep"),), name="ld", allow_any_len=True),
+   lambda gep, ld: ld.replace(dtype=ld.dtype.scalar().vec(gep.dtype.count), src=(gep.src[0],)+ld.src[1:]).gep(gep.arg)),
+  # GEP on data of STORE
+  (UPat(Ops.STORE, src=(UPat(Ops.GEP, name="gep"), UPat.var("st")), allow_any_len=True, name="sto"), gep_on_store),
+  # put PTRCAT after LOAD
+  (UPat(Ops.LOAD, src=(UPat(Ops.PTRCAT, name="cat"),), name="ld", allow_any_len=True),
+   lambda cat,ld: UOp(Ops.CAT, cat.dtype.base.vec(cat.dtype.vcount), tuple(ld.replace(dtype=x.dtype.base, src=(x,)+ld.src[1:]) for x in cat.src))),
+  # put PTRCAT after STORE
+  (UPat(Ops.STORE, src=(UPat(Ops.PTRCAT, name="cat"), UPat(name="data")), allow_any_len=True, name="sto"), cat_after_store),
+])
+
+# *** correct load/store ***
+
+def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):
+  # this splits loads and stores into multiple chunks
+
+  # if there's only one element to load/store, no splitting needed
+  if (sz:=ls.src[0].dtype.count) == 1: return None
+  buf = idx.src[0]
+
+  # determine fold lengths
+  lengths = []
+  must_divide = True
+  if ctx is not None and ctx.device == "DSP":
+    lengths = [128,64,32,16,8,4]
+    must_divide = False
+  elif buf.dtype.base not in (dtypes.float, dtypes.half, *dtypes.fp8s) and not isinstance(buf.dtype, ImageDType):
+    pass
+  elif buf.ptrdtype.addrspace == AddrSpace.REG:
+    pass
+  elif isinstance(buf.dtype, ImageDType):
+    lengths = [4]
+  elif ctx is not None and ctx.supports_float4:
+    # TODO: a better way to get this than ctx
+    lengths = [8,4,2] if buf.dtype.base == dtypes.half and getenv("ALLOW_HALF8") else ([16,8,4,2] if AMX else [4,2])
+  lengths.append(1)  # worst case, it's not folded
+
+  # filter fold lengths that don't divide
+  offset, mask = idx.src[1].get_idx(), idx.src[1].get_valid()
+  if must_divide: lengths = [x for x in lengths if offset.divides(x) is not None]
+
+  # split based on the fold lengths
+  global_offset = 0
+  ret = []
+  while global_offset < sz:
+    # with 1 at the end of the lengths list, this will always hit
+    for fold_length in lengths:
+      if global_offset+fold_length > sz: continue
+      lidx = buf.index((offset + global_offset).valid(mask), ptr=True)
+      if fold_length > 1: lidx = lidx.cast(buf.ptrdtype.base.vec(fold_length).ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace))
+      if ls.op is Ops.STORE: ret.append(ls.replace(src=(lidx,ls.src[1].gep(tuple(range(global_offset, global_offset+fold_length))))+ls.src[2:]))
+      else: ret.append(ls.replace(src=(lidx,)+ls.src[1:], dtype=ls.dtype.scalar().vec(fold_length)))
+      global_offset += fold_length
+      break
+
+  # if it wasn't split, we return None. otherwise we CAT them
+  if len(ret) <= 1: return None
+  return UOp(Ops.CAT, ls.dtype, tuple(ret)) if ls.op is Ops.LOAD else UOp.group(*ret)
+
+def image_fixup(ls:UOp):
+  # normal image load or store, with the CAST from expand_index
+  if ls.src[0].op is Ops.CAST and isinstance(image_dtype:=ls.src[0].src[0].dtype, ImageDType):
+    assert ls.src[0].dtype.count == 4, "image must be casted to 4"
+    idx = ls.src[0].src[0]
+    x, valid = idx.src[1].get_idx(), idx.src[1].get_valid()
+    oidx = UOp(Ops.VECTORIZE, dtypes.index.vec(2), ((x // 4) % image_dtype.shape[1], (x // (4*image_dtype.shape[1]))))
+    idx = idx.replace(src=(idx.src[0], oidx.valid(valid)))
+    return ls.replace(src=(idx,)+ls.src[1:])
+
+  # this is an unprocessed image without a cast, aka unfoldable image load. this doesn't work for stores
+  if isinstance(image_dtype:=ls.src[0].dtype, ImageDType) and ls.src[0].src[1].get_idx().dtype != dtypes.index.vec(2):
+    assert ls.op is Ops.LOAD, "if an image store isn't upcasted to 4, we can't store it"
+    idx = ls.src[0]
+    x, valid = idx.src[1].get_idx(), idx.src[1].get_valid()
+    oidx = UOp(Ops.VECTORIZE, dtypes.index.vec(2), ((x // 4) % image_dtype.shape[1], (x // (4*image_dtype.shape[1]))))
+    idx = idx.replace(src=(idx.src[0], oidx.valid(valid)))
+    vec_load = ls.replace(dtype=ls.dtype.vec(4), src=(idx,)+ls.src[1:])
+    return functools.reduce(lambda ret, i: (x % 4).ne(i).where(ret, vec_load.gep(i)), range(4), ls.const_like(float('nan')))
+
+  return None
+
+correct_load_store = PatternMatcher([
+  # split LOAD/STORE
+  (UPat((Ops.LOAD, Ops.STORE), src=(UPat(Ops.INDEX, name="idx").cast(),), name="ls", allow_any_len=True), split_load_store),
+  # image indexing, including unfoldable images
+  (UPat((Ops.LOAD, Ops.STORE), name="ls"), image_fixup),
+])
+
+# *** uop expander ***
+
+# TODO: there's a lot shared with gep_through_wmma here
+def no_vectorized_wmma(wmma:UOp):
+  out_sz = prod(x[1] for x in wmma.arg[6][-1])
+  if wmma.dtype.count == out_sz: return None
+  tsrcs = []
+  for s,sz in zip(wmma.src, wmma.arg[6]):
+    ssz = prod(x[1] for x in sz)
+    tsrcs.append([s.gep(tuple(range(grp, grp+ssz))) for grp in range(0, s.dtype.count, ssz)])
+  wmmas = [UOp(Ops.WMMA, wmma.dtype.scalar().vec(out_sz), tsrc, wmma.arg) for tsrc in zip(*tsrcs)]
+  wmma_ex = flatten([[e.gep(i) for i in range(out_sz)] for e in wmmas])
+  return UOp(Ops.VECTORIZE, wmma.dtype, tuple(wmma_ex))
+
+def no_vectorized_alu(alu:UOp):
+  if alu.dtype.vcount == 1: return None
+  if alu.op is Ops.WHERE and alu.src[2].arg is Invalid: return None  # image load/store has cond.where(idx.vec(2), Invalid) as the index
+  alus = tuple(UOp(alu.op, alu.dtype.scalar(), tuple(s.gep(i) for s in alu.src), alu.arg) for i in range(alu.dtype.vcount))
+  return UOp(Ops.VECTORIZE, alu.dtype, alus)
+
+def no_vectorized_buf(buf:UOp):
+  return buf.replace(dtype=buf.ptrdtype.base.scalar().ptr(buf.ptrdtype.size*buf.ptrdtype.count, buf.ptrdtype.addrspace)).cast(buf.dtype)
+
+def no_vectorized_index(buf:UOp, cast:UOp, idx:UOp):
+  cnt = cast.dtype.count
+  assert idx.dtype.count == 1, f"idx dtype must be 1 {idx.dtype}"
+  return buf.broadcast(cnt).index(idx.broadcast(cnt)*cnt+UOp.const(dtypes.index.vec(cnt), tuple(range(cnt))), ptr=True)
+
+def no_vectorized_index_broadcast(buf:UOp, cast:UOp, bcast:UOp, idx:UOp):
+  cnt = cast.dtype.count
+  precnt = bcast.dtype.vcount
+  input_gep = bcast.arg if bcast.op is Ops.GEP else ([0]*precnt)
+  gep_arg = tuple(flatten([range(precnt) for _ in range(cnt)]))
+  sum_arg = tuple(flatten([[i+y for y in input_gep] for i in range(cnt)]))
+  return buf.broadcast(cnt*precnt).index(idx.gep(gep_arg)*cnt+UOp.const(dtypes.index.vec(cnt*precnt), sum_arg), ptr=True)
+
+devectorize_buf_and_index = PatternMatcher([
+  (UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG), name="buf"), no_vectorized_buf),
+  (UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG)).or_after(name="buf").cast(name="cast").index(UPat.var("idx")), no_vectorized_index),
+  (UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG)).or_after(name="buf").cast(name="cast").broadcast(name="bcast").index(UPat.var("idx")),
+   no_vectorized_index_broadcast),
+  (UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG)).or_after(name="buf").cast(name="cast").gep(name="bcast").index(UPat.var("idx")),
+   no_vectorized_index_broadcast),
+])
+
+devectorize = PatternMatcher([
+  # CAST after AFTER
+  (UPat(Ops.CAST, name="c").f(Ops.AFTER, allow_any_len=True, name="a"), lambda c,a: c.src[0].after(*a.src[1:]).cast(c.dtype)),
+  # no ALU on vectorized dtypes
+  (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST), name="alu"), no_vectorized_alu),
+  (UPat(Ops.WMMA, name="wmma"), no_vectorized_wmma),
+])+devectorize_buf_and_index
+
+pm_render = PatternMatcher([
+  # for rendering, we use explicit VECTORIZE
+  (UPat(Ops.CONST, name='c'),
+   lambda c: UOp(Ops.VECTORIZE, c.dtype, (UOp.const(c.dtype.scalar(), c.arg),)*c.dtype.vcount) if c.dtype.vcount > 1 else None),
+  (UPat(Ops.VCONST, name='c'), lambda c: UOp(Ops.VECTORIZE, c.dtype, tuple(UOp.const(c.dtype.scalar(), x) for x in c.arg))),
+  (UPat(Ops.GEP, name='gep'), lambda gep: UOp(Ops.VECTORIZE, gep.dtype, tuple(gep.src[0].gep(x) for x in gep.arg)) if len(gep.arg) > 1 else None),
+  (UPat(Ops.GEP, name='gep'), lambda gep: gep.src[0] if gep.src[0].dtype.vcount == 1 and gep.arg == (0,) else None),
+  (UPat(Ops.VECTORIZE, src=(UPat(name='x'),)), lambda x: x),
+  # give any loads that are masked an alt value
+  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat())).or_casted(),), allow_any_len=True, name="x"),
+    lambda x: x.replace(src=(x.src[0], x.const_like(0))+x.src[1:])
+      if len(x.src) == 1 or x.src[1].op in (Ops.CUSTOM, Ops.STORE, Ops.BARRIER) else None),
+  # Where after gated load becomes alt value
+  (UPat.var("c").where(UPat(Ops.LOAD, src=(UPat().index(UPat.var("idx"), UPat.var("c")).or_casted(),), allow_any_len=True, name="l").or_casted(),
+    UPat.var("a")), lambda c,idx,l,a: l.replace(src=(l.src[0], a.cast(l.dtype))+l.src[2:]).cast(a.dtype)),
+  (UPat.var("c").where(UPat.var("a"), UPat(Ops.LOAD, src=(UPat().index(UPat.var("idx"), UPat.var("c").logical_not()).or_casted(),),
+    allow_any_len=True, name="l").or_casted()), lambda c,idx,l,a: l.replace(src=(l.src[0], a.cast(l.dtype))+l.src[2:]).cast(a.dtype)),
+])
+
+# *** Ops.REDUCE -> Ops.DEFINE_ACC ***
+
+@dataclass
+class ReduceContext:
+  acc_num: int = 0
+
+def horizontal_reduce(inp:UOp, out_dtype:DType) -> list[UOp]:
+  # if this has a horizontal reduction component, do that first
+  if inp.dtype != out_dtype:
+    # NOTE: [0 1 2 3 4 5 6 7] -> [0+4, 1+5, 2+6, 3+7]
+    horizontal_amount = inp.dtype.count//out_dtype.count
+    return [inp.gep(tuple(range(i, inp.dtype.count, horizontal_amount))) for i in range(0, horizontal_amount)]
+  return [inp]
+
+def reduce_to_acc(ctx:ReduceContext, red:UOp):
+  inp, reduce_range = red.src[0], red.src[1:]
+  lst = horizontal_reduce(inp, red.dtype)
+  assert all(x.dtype == red.dtype for x in lst), f"horizontal reduction mismatch {lst[0].dtype} != {red.dtype}"
+  # if we have a range
+  if len(reduce_range) != 0:
+    topo = inp.toposort()
+    ended_ranges = flatten([x.ended_ranges for x in topo if x.op is Ops.END])
+    input_ranges = tuple([x for x in topo if x.op is Ops.RANGE and x not in reduce_range and x not in ended_ranges])
+    identity = red.const(red.dtype, identity_element(red.arg, red.dtype.scalar()))
+    acc = UOp(Ops.DEFINE_REG, red.dtype.ptr(size=1, addrspace=AddrSpace.REG), arg=ctx.acc_num)
+    acc_init = acc.after(*input_ranges).index(UOp.const(dtypes.int, 0)).store(identity) if len(input_ranges) else \
+               acc.index(UOp.const(dtypes.int, 0)).store(identity)
+    lst = [acc.after(acc_init, *reduce_range).index(UOp.const(dtypes.int, 0))] + lst  # put acc as the first element
+    ctx.acc_num += 1
+  ret = functools.reduce(lambda x,y: x.alu(red.arg, y), lst)
+  if len(reduce_range) == 0: return ret
+  return acc.after(acc.index(UOp.const(dtypes.int, 0)).store(ret).end(*reduce_range)).index(UOp.const(dtypes.int, 0))
+
+pm_reduce = PatternMatcher([
+  # REDUCE -> DEFINE_ACC+ASSIGN
+  (UPat(Ops.REDUCE, name="red"), reduce_to_acc),
+  # tensor core built in accumulate
+  (UPat(Ops.WMMA, name="wmma") + UPat.var("add"),
+    lambda add, wmma: UOp(wmma.op, wmma.dtype, (wmma.src[0], wmma.src[1], wmma.src[2]+add), wmma.arg)),
+])+sym
+
+# add loads
+
+pm_add_loads = PatternMatcher([
+  # add loads to non ptr index
+  (UPat(Ops.INDEX, name="idx"), lambda idx: None if isinstance(idx.dtype, (PtrDType, ImageDType)) else
+    idx.replace(dtype=idx.src[0].dtype).load(dtype=idx.dtype.base)),
+  # remove loads from stores
+  (UPat(Ops.STORE, src=(UPat(Ops.LOAD),), allow_any_len=True, name="s"), lambda s: s.replace(src=(s.src[0].src[0],)+s.src[1:])),
+])
+

tinygrad_repo/tinygrad/codegen/late/expander.py

@@ -0,0 +1,156 @@
+# this converts a lowerer program into a vectorized program
+import functools, itertools
+from tinygrad.dtype import dtypes, PtrDType, AddrSpace
+from tinygrad.helpers import AMX, dedup, flatten, all_same, prod, partition
+from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp, AxisType, range_start
+from tinygrad.schedule.rangeify import BufferizeOpts
+
+def _expand_arg_to_idx(args:tuple[tuple[int, int], ...], rpk:dict[int, int]) -> int:
+  idx, mul = 0, 1
+  for axis,m in args[::-1]:
+    idx += rpk[axis] * mul
+    mul *= m
+  return idx
+
+def _choices_from_args(args:tuple[tuple[int, int], ...]) -> list[dict[int, int]]:
+  return [dict(x) for x in itertools.product(*[zip(itertools.repeat(axis), range(m)) for axis,m in args])]
+
+@functools.cache
+def _swizzle_args(cargs:tuple[tuple[int, int], ...], eargs:tuple[tuple[int, int], ...], exclude_args:tuple[int, ...]) -> list[int]:
+  return [_expand_arg_to_idx(eargs, {**rpk, **{x:0 for x in exclude_args}} if exclude_args else rpk) for rpk in _choices_from_args(cargs)]
+
+def do_expand(root:UOp):
+  expands = [x for x in root.src if x.op is Ops.UNROLL]
+  if len(expands) == 0: return None
+  # NOTE: we 0 out the reduce axis for WMMA. in theory they should all be the same, but is this always correct?
+  exclude_args = tuple(dedup(root.arg[-1] + tuple(y[0] for y in flatten(root.arg[-2])))) if root.op is Ops.WMMA else ()
+  if all_same(expands_args:=[x.arg for x in expands]) and len(exclude_args) == 0:
+    # if there's only one expand arg, it's okay to use it (optimization)
+    expand_args = expands[0].arg
+  else:
+    # otherwise, we sort them and GEP
+    expand_args = tuple(x for x in sorted(dedup(flatten(expands_args))) if x[0] not in exclude_args)
+  expand_sz = prod([x[1] for x in expand_args])
+  new_srcs = []
+  for i,src in enumerate(root.src):
+    if src.op is Ops.UNROLL:
+      if expand_args == src.arg:
+        # just remove the expand
+        new_srcs.append(src.src[0])
+      else:
+        lst = _swizzle_args(expand_args, src.arg, exclude_args)
+        # if the base dtype is > 1, put those at the end
+        if src.dtype.count > 1: lst = flatten([[i*src.dtype.count+j for j in range(src.dtype.count)] for i in lst])
+        new_srcs.append(src.src[0].gep(tuple(lst)))
+    else:
+      # non-UNROLL input
+      if root.op in range_start and i >= range_start[root.op]:
+        # for any range args of STORE/REDUCE, pass them through
+        new_srcs.append(src)
+      elif root.op is Ops.INDEX and i >= 1 and not isinstance(root.dtype, PtrDType):
+        new_srcs.append(src)
+      elif src.dtype.count > 1:
+        # put any input dtype > 1 grouped together
+        new_srcs.append(UOp(Ops.CAT, src.dtype.scalar().vec(expand_sz*src.dtype.count), (src,)*expand_sz))
+      else:
+        # repeat the arg
+        new_srcs.append(src.broadcast(expand_sz))
+
+  new_arg = root.arg
+  if root.op is Ops.GEP:
+    assert root.dtype.count == 1
+    # is this right?
+    new_arg = tuple(range(root.arg[0], new_srcs[0].dtype.count, new_srcs[0].dtype.count // expand_sz))
+  nsrc = UOp(root.op, root.dtype.scalar().vec(root.dtype.count*expand_sz), tuple(new_srcs), new_arg)
+  return UOp(Ops.UNROLL, root.dtype, (nsrc,), expand_args)
+
+def do_contract(con:UOp):
+  ex = con.src[0]
+  # CONTRACT without UNROLL repeats the element VECTORIZED
+  if ex.op is not Ops.UNROLL: return UOp(Ops.VECTORIZE, con.dtype, con.src*con.dtype.count)
+  # CONTRACT may remove several axes from UNROLL
+  assert con.dtype == dtypes.void or con.dtype.count == prod([x[1] for x in con.arg]), "dtype is wrong"
+  idxs = []
+  for rpk in _choices_from_args(new_ex_args:=tuple(x for x in ex.arg if x not in con.arg)):
+    idxs += [_expand_arg_to_idx(ex.arg, {**rpk, **lrpk}) for lrpk in _choices_from_args(con.arg)]
+  return UOp(Ops.UNROLL, con.dtype, (ex.src[0].gep(tuple(idxs)),), new_ex_args)
+
+def end_unrolls(u:UOp):
+  unrolls, src = partition(u.src[1:], lambda x: x.op is Ops.UNROLL)
+  if not len(unrolls): return None
+  ret = UOp(Ops.CONTRACT, dtypes.void, (u.src[0],), sum([x.arg for x in unrolls], start=()))
+  return u.replace(src=(ret,)+tuple(src))
+
+expander = PatternMatcher([
+  # push broadcast through AFTER
+  (UPat.var("x").broadcast(name="b").after(name="a", allow_any_len=True), lambda x,b,a: x.after(*a.src[1:]).broadcast(len(b.src))),
+  (UPat.var("x").broadcast(name="b").end(name="a", allow_any_len=True), lambda x,b,a: x.end(*a.src[1:]).broadcast(len(b.src))),
+  # END on UNROLL ends the UNROLL
+  (UPat(Ops.END, name="u"), end_unrolls),
+  # BUFFERIZE puts UNROLLs for ranges as contract
+  (UPat(Ops.BUFFERIZE, src=(UPat(Ops.UNROLL), UPat(Ops.UNROLL)), name="x"),
+    lambda x: x.replace(src=tuple(UOp(Ops.CONTRACT, dtype=s.dtype.vec(x.src[1].src[0].dtype.count), src=(s,), arg=x.src[1].arg) for s in x.src))),
+  # double expand
+  (UPat(Ops.UNROLL, name="outer", src=(UPat(Ops.UNROLL, name="inner"),)),
+   lambda outer, inner: UOp(Ops.UNROLL, outer.dtype, (inner.src[0],), inner.arg+outer.arg)),
+  # do expansion
+  (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST, Ops.GEP, Ops.WMMA, Ops.LOAD, Ops.STORE, Ops.INDEX, Ops.BUFFERIZE,
+         Ops.VECTORIZE, Ops.REDUCE, Ops.END, Ops.AFTER), name="root", custom_early_reject=set([Ops.UNROLL])), do_expand),
+  (UPat(Ops.CONTRACT, name="con"), do_contract),
+  # BARRIERs aren't actually expanded
+  (UPat(Ops.BARRIER, src=(UPat(Ops.UNROLL, name="ex"),)),
+   lambda ex: UOp(Ops.UNROLL, src=(UOp(Ops.BARRIER, src=ex.src),)*len(ex.src), arg=ex.arg)),
+  # empty UNROLL is NOOP
+  (UPat(Ops.UNROLL, src=(UPat.var('x'),), arg=()), lambda x: x),
+  # UNROLL GEP (needed for WMMA, generalize this) -> vectorized ALU
+  (UPat(Ops.UNROLL, name="ex", src=tuple(UPat.var('x').gep(i)+UPat.var('y').gep(i) for i in range(256 if AMX else 8))),
+    lambda ex,x,y: UOp(Ops.UNROLL, ex.dtype, tuple((x+y).gep(i) for i in range(256 if AMX else 8)), ex.arg)),
+])
+
+# ****
+
+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=BufferizeOpts(reduce_gfr[0].arg[0], AddrSpace.LOCAL)).index(*upstream_locals, *reduce_loop)
+
+  # do the final reduce (if/barrier are added in gpudims step)
+  return buf.reduce(*reduce_loop, arg=x.arg)
+
+pm_pre_expander = PatternMatcher([
+  # rewrite UPCAST/UNROLL range to something to be expanded
+  (UPat(Ops.RANGE, name="r"),
+   lambda r: UOp(Ops.UNROLL, r.dtype, (UOp.const(r.dtype.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),
+])
+
+pm_group_for_reduce = PatternMatcher([
+  # fix group for reduce
+  (UPat(Ops.REDUCE, name="x"), fix_group_for_reduce),
+])

tinygrad_repo/tinygrad/codegen/late/linearizer.py

@@ -0,0 +1,101 @@
+import heapq
+from typing import Any
+from collections import defaultdict
+from tinygrad.uop.ops import PatternMatcher, UOp, Ops, UPat, multirange_str
+from tinygrad.helpers import prod, getenv, TUPLE_ORDER
+
+def linearize(sink:UOp) -> list[UOp]:
+  # this is a toposort with priority
+  lst = list(sink.toposort())
+  consumers: defaultdict[UOp, list[UOp]] = defaultdict(list)
+  in_degree:dict[UOp, int] = {}
+  out_degree:dict[UOp, int] = {}
+  priorities:dict[UOp, tuple[int, int, Any]] = {}
+
+  # get consumers and assign priorities
+  # NOTE: this requires the lst be locally toposorted
+  for u in reversed(lst):
+    for s in u.src: consumers[s].append(u)
+    in_degree[u] = len(u.src)
+    out_degree[u] = len(consumers[u])
+
+    # we place UOps with higher run_counts later
+    run_count = prod([int(r.vmax)+1 for r in u.ranges])
+
+    # simple priority override. this is all bottom up now, smaller numbers will be closer to the top
+    extra = None
+    match u.op:
+      # the order and placement of these defines is important
+      case Ops.DEFINE_GLOBAL: priority, extra = -20, u.arg
+      case Ops.DEFINE_VAR: priority, extra = -19, u.arg
+      case Ops.DEFINE_LOCAL: priority = -18
+      case Ops.DEFINE_REG: priority = -17
+      case Ops.CONST: priority = -10  # early consts
+      case Ops.LOAD: priority = -1    # place loads early
+      case Ops.STORE: priority = 1    # place stores late
+      case Ops.RANGE: priority = 5    # placing RANGE is good
+      case Ops.END: priority = -5     # placing END is bad
+      case _: priority = 0            # everything else has priority 0
+    priorities[u] = (run_count, priority, extra)
+
+  # number the uops in "ideal" order
+  nkey = {u:i for i,u in enumerate(sorted(lst, key=lambda x: priorities[x]+(x.tuplize if TUPLE_ORDER else ())))}
+
+  # then force them to be toposorted in as close to the ideal order as possible
+  heap = [(-nkey[sink], sink)]
+  newlst = []
+  while heap:
+    newlst.append(u:=heapq.heappop(heap)[1])
+    for v in u.src:
+      out_degree[v] -= 1
+      if out_degree[v] == 0: heapq.heappush(heap, (-nkey[v],v))
+  newlst = newlst[::-1]
+
+  if getenv("DEBUG_LINEARIZE"):
+    for i,u in enumerate(newlst):
+      print(f"{i:4d} {str(u.op):20s} {multirange_str(u.ranges, color=True, pad=10)} {priorities[u]}")
+  return newlst
+
+class CFGContext:
+  def __init__(self, sink:UOp):
+    # there are 3 relationships between ranges:
+    # nested, meaning endrange y is a dependency of endrange x and range x is a dependency of endrange y
+    # dependent, meaning endrange y is a dependency of endrange x and range x is not a dependency of endrange y
+    # independent, endrange y is not a dependency of endrange x
+    # everything is nested inside the sink
+    deps: dict[UOp, dict[UOp, None]] = {}
+    nesting: dict[UOp, UOp] = {}
+    for u in sink.toposort():
+      # get the deps from the src
+      deps[u] = {}
+      for s in u.src: deps[u] |= deps[s]
+
+      if u.op in (Ops.END, Ops.SINK):
+        nesting |= {x:u for x in deps[u] if x.op is Ops.END and (u.op is Ops.SINK or u.src[1] in deps[x]) and x not in nesting}
+      if u.op in (Ops.RANGE, Ops.END): deps[u][u] = None
+
+    self.edges: dict[UOp, UOp] = {}
+    siblings: dict[UOp, list[UOp]] = {}
+    for k,vv in nesting.items(): siblings.setdefault(vv, []).append(k)
+    for k,v in siblings.items():
+      # ranges that have dependencies on other siblings need to be scheduled after them
+      order = sorted(v, key=lambda x: len([u for u in v if u in deps[x]]))
+      zipped = zip(order, order[1:]) if k.op is Ops.SINK else zip([k.src[1]] + order, order)
+      for x,y in zipped:
+        # TODO: this can happen! it causes infinite loop in shufflenet
+        assert y.src[1] not in x.backward_slice_with_self
+        self.edges[y.src[1]] = x
+
+pm_add_control_flow = PatternMatcher([
+  (UPat(Ops.RANGE, name="x"), lambda ctx,x: x.replace(src=x.src+(y,)) if (y:=ctx.edges.get(x)) is not None else None),
+])
+
+def do_split_ends(e:UOp):
+  ret = e.src[0]
+  for r in sorted(UOp.sink(*e.src[1:]).ranges, key=lambda x: x.arg, reverse=True): ret = ret.end(r)
+  return ret
+
+pm_split_ends = PatternMatcher([
+  # split the ends
+  (UPat(Ops.END, name="e"), do_split_ends),
+])

tinygrad_repo/tinygrad/codegen/opt/__init__.py

@@ -0,0 +1,20 @@
+# opt opinionatedly transforms an ast into an optimized ast using either heuristics or beam search
+from __future__ import annotations
+from enum import Enum, auto
+from dataclasses import dataclass
+
+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
+  def __lt__(self, x:OptOps): return self.value < x.value
+
+@dataclass(frozen=True, order=True)
+class Opt:
+  op: OptOps
+  axis: int|None = None
+  arg: int|tuple|None = None
+  def __repr__(self): return f"Opt(op={self.op}, axis={self.axis}, arg={self.arg})"
+
+class KernelOptError(Exception): pass
+def check(cond:bool, msg:str=""):
+  if not cond: raise KernelOptError(msg)

tinygrad_repo/tinygrad/codegen/opt/heuristic.py

@@ -0,0 +1,190 @@
+import itertools
+from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
+from tinygrad.helpers import getenv, DEBUG, prod, NOLOCALS, TC_OPT, TC_SELECT, USE_TC, AMX
+from tinygrad.dtype import ImageDType
+from tinygrad.uop.ops import Ops, resolve, AxisType
+from tinygrad.codegen.opt.postrange import Scheduler
+
+def hand_coded_optimizations(k:Scheduler) -> Scheduler:
+  # first try the tensor cores
+  """ Attempts to apply a tensor core optimization to the kernel. If one exists and applies properly, return true, otherwise return false.
+  Tensor cores are optimized instructions that matrix multiply-accumulate across a wave of threads: D(M, N) = A(M, K) * B(K, N) + C(M, N).
+
+  Keyword arguments:
+  use_tensor_cores -- controls how tensor cores are applied (default 1)
+    0: will disable any tensor core matching
+    1: enable tensor cores
+    2: apply tensor core shape but don't use UOp.WMMA
+  extra_opts -- additional Opt's to apply after the tensor core instead of the hand-coded additional Opt's (default None)
+  tc_select -- specifies which tensor core(s) to use for optimization (default -1)
+    -1: iterates through all available tensor cores in order and uses the first one that matches the requirements (dims and dtypes)
+    [0-N]: uses only the n'th tensor core available; useful for search
+  tc_opt -- controls which kinds of kernels may be eligible for tensor cores application (default 2 during BEAM, 0 otherwise)
+    0: applies to only kernels with a single reduce axis and direct Ops.LOAD into Ops.MUL
+    1: allows kernels with multiple reduce axes and also multiplication of Ops.CAST'd buffers
+    2: allows kernels with M, N, K axes that are not multiples of the tensor core dimensions by applying padding those axes as needed
+  """
+  # NOTE: unless TC_OPT is > 0, we only trigger tensor cores if there's only one reduce axis
+  if USE_TC > 0 and (len(k.axes_of(AxisType.GROUP_REDUCE, AxisType.REDUCE)) == 1 or (TC_OPT.value >= 1)):
+    good_tc_opt = False
+    tk = k.copy()
+    try: # check TC first and apply hand-coded opts if successful
+      rngs = tk.apply_opt(Opt(OptOps.TC, 0, (TC_SELECT.value, TC_OPT.value, USE_TC.value)))
+      good_tc_opt = True
+    except KernelOptError:
+      pass
+    # skip hand-coded TC opts if AMX, upcasting will make kernel slower
+    if good_tc_opt and not AMX:
+      if rngs is not None:
+        for tc_dim in [1,0]: # attempt to upcast M and N
+          szs = [sz for sz in [5,4,3,2] if rngs[tc_dim].src[0].divides(sz) is not None]
+          if szs:
+            # set it to the replaced range
+            rngs[tc_dim] = tk.apply_opt(Opt(OptOps.UPCAST, tk.rngs.index(rngs[tc_dim]), szs[0]))[0]
+        if (szs := [sz for sz in [4,2] if rngs[0].src[0].divides(sz) is not None]): # attempt to local N
+          tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(rngs[0]), szs[0]))
+      return tk
+
+  # make a copy so it does not mutate the input
+  k = k.copy()
+
+  # upcast float4 images, this must be early so we don't accidentally add locals before the upcast
+  for buf_index,buf in enumerate(k.bufs):
+    if isinstance(buf.src[0].dtype, ImageDType):
+      # part of is_expanded
+      unit_stride_axes_mul_4 = [k.rngs.index(c) for c in k.bufs[buf_index].src[1].get_idx().split_uop(Ops.ADD) if
+        c.op is Ops.RANGE and (c.vmax+1)%4 == 0]
+      if len(unit_stride_axes_mul_4):
+        if (axis:=unit_stride_axes_mul_4[0]) in k.upcastable_dims:
+          k.apply_opt(Opt(OptOps.UPCAST, axis, 4))
+        elif axis in k.unrollable_dims:
+          k.apply_opt(Opt(OptOps.UNROLL, k.unrollable_dims.index(axis), 4))
+
+  # should use matvec - TODO: adjust/tune based on the wide vs tall/large vs small mat
+  MV_BLOCKSIZE, MV_THREADS_PER_ROW, MV_ROWS_PER_THREAD = getenv("MV_BLOCKSIZE", 4), getenv("MV_THREADS_PER_ROW", 8), getenv("MV_ROWS_PER_THREAD", 4)
+  if k.ren.has_local and getenv("MV",1) != 0 and (MV_BLOCKSIZE > 1 or MV_THREADS_PER_ROW > 1 or MV_ROWS_PER_THREAD > 1) and  \
+    k.reduceop is not None and k.reduceop.arg[0] is Ops.ADD and len(k.full_shape) >= 2 and k.ren.has_shared and \
+    (mulop:=k.reduceop.src[0]).op is Ops.MUL and mulop.src[0].op is Ops.INDEX and mulop.src[1].op is Ops.INDEX:
+    idx0, idx1 = mulop.src[0].src[1].get_idx(), mulop.src[1].src[1].get_idx()
+    if k.ranges_of(AxisType.REDUCE):
+      first_reduce_rng = k.ranges_of(AxisType.REDUCE)[0]
+      if any(u is first_reduce_rng for u in idx0.split_uop(Ops.ADD)) and all(r in idx1.ranges for r in idx0.ranges):
+        for global_idx in k.axes_of(AxisType.GLOBAL):
+          if first_reduce_rng.src[0].divides(MV_THREADS_PER_ROW) is not None and k.full_shape[global_idx]%(MV_BLOCKSIZE*MV_ROWS_PER_THREAD) == 0:
+            if DEBUG >= 3:
+              print(f"MATVEC: {k.full_shape=} {first_reduce_rng.render()} {MV_BLOCKSIZE=} {MV_THREADS_PER_ROW=} {MV_ROWS_PER_THREAD=}")
+            try:
+              if MV_THREADS_PER_ROW > 1: k.apply_opt(Opt(OptOps.GROUP, 0, MV_THREADS_PER_ROW))
+            except KernelOptError: pass
+            if MV_BLOCKSIZE > 1: k.apply_opt(Opt(OptOps.LOCAL, global_idx, MV_BLOCKSIZE))
+            if MV_ROWS_PER_THREAD > 1: k.apply_opt(Opt(OptOps.UPCAST, global_idx, MV_ROWS_PER_THREAD))
+            return k
+
+  # are we grouping? (requires local shape support)
+  if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= (240 if NOLOCALS else 2048), False):
+    for axis, sz in itertools.product((0, 1, 2), (16,)):
+      try:
+        k.apply_opt(Opt(OptOps.GROUPTOP, axis, sz))
+        break
+      except KernelOptError: pass
+
+  # no more opt if we are grouping
+  if k.group_for_reduces: return k
+
+  # **** below this line need to be optional and benchmarked ****
+
+  # if there are small dims with lots of valid masks, upcast them (they might be from Tensor.stack)
+  to_upcast: list[int] = []
+  # upcast leading axes first (hack-ish for winograd; we actually want to upcast masked axes with low stride first)
+  for axis in k.upcastable_dims:
+    # for Schedule, we check if the range is used in INDEX gates or WHERE gates
+    is_masked = any(any(o is k.rngs[axis] for o in u.src[0].backward_slice) for u in k.ast.backward_slice if u.op is Ops.WHERE)
+    if k.full_shape[axis] <= 7 and is_masked and prod(k.full_shape[j] for j in to_upcast) * k.full_shape[axis] <= 7 * 7:
+      if DEBUG >= 4: print(f"upcasting masked axis : {axis}")
+      to_upcast.append(axis)
+  for axis in to_upcast[::-1]: k.apply_opt(Opt(OptOps.UPCAST, axis, 0))
+
+  # potentially do more upcasts of non reduce axes based on a heuristic
+  is_dsp = k.ren is not None and k.ren.device == "DSP"
+  upcasted_axis: set[int] = set()
+  while resolve(prod(k.output_shape[i] for i in k.upcastable_dims) >= 1024) and (k.upcast_size() < 32):
+    xb_choices = []
+    # consider all upcastable axes with 3 or 4 upcast (128 on the DSP)
+    for axis, upcast_amount in itertools.product(k.upcastable_dims, ([128] if not len(upcasted_axis) else []) if is_dsp else [3,4]):
+      # if we haven't upcasted it, it mods, and buffer has stride 0 on axis while having no stride 0 in the upcasted axis already
+      if axis in upcasted_axis or k.full_shape[axis]%upcast_amount != 0: continue
+      rng = k.rngs[axis]
+      if any(rng not in b.src[1].get_idx().backward_slice and all(r2 in b.src[1].get_idx().backward_slice
+          for r2 in k.ranges_of(AxisType.UPCAST, AxisType.UNROLL)) for b in k.bufs):
+        num_strides, sum_strides = 0, 0
+        for b in k.bufs:
+          idx = b.src[1].get_idx()
+          if rng in idx.backward_slice: num_strides += 1
+          for c in idx.split_uop(Ops.ADD):
+            if c is rng: sum_strides += 1
+            if c.op is Ops.MUL and c.src[0] is rng and c.src[1].op is Ops.CONST: sum_strides += c.src[1].arg
+            if c.op is Ops.MUL and c.src[1] is rng and c.src[0].op is Ops.CONST: sum_strides += c.src[0].arg
+        xb_choices.append((num_strides, sum_strides, axis, upcast_amount))
+    if xb_choices:
+      xb_choices = sorted(xb_choices)
+      if DEBUG >= 4: print(f"more upcast axis : {xb_choices}")
+      k.apply_opt(Opt(OptOps.UPCAST, xb_choices[0][2], xb_choices[0][3]))
+      upcasted_axis.add(xb_choices[0][2])
+    else: break
+
+  # if last reduce dim is small(ish), loop unroll the reduce
+  # NOTE: this can fail on multireduce with mismatching dimensions, this is okay
+  try:
+    if k.unrollable_dims and (k.upcast_size() <= 4 or not k.axes_of(AxisType.UNROLL)) and (k.upcast_size() < 64):
+      if (s:=k.full_shape[k.unrollable_dims[-1]]) <= 32:
+        k.apply_opt(Opt(OptOps.UNROLL, len(k.unrollable_dims)-1, 0))
+        # if it's small, upcast a second reduce dimension too
+        if k.unrollable_dims and s <= 3 and k.full_shape[k.unrollable_dims[-1]] <= 3:
+          k.apply_opt(Opt(OptOps.UNROLL, len(k.unrollable_dims)-1, 0))
+      else:
+        for splits in [4]:
+          if k.full_shape[axis:=k.unrollable_dims[-1]]%splits == 0:
+            k.apply_opt(Opt(OptOps.UNROLL, len(k.unrollable_dims)-1, splits))
+            break
+  except KernelOptError: pass
+
+  # if nothing at all is upcasted and it's easy to, do an upcast
+  for splits in [4]:
+    if not k.upcasted and k.upcastable_dims and k.full_shape[k.upcastable_dims[-1]] % splits == 0:
+      k.apply_opt(Opt(OptOps.UPCAST, k.upcastable_dims[-1], splits))
+
+  # **** local groups ****
+
+  if k.ren.has_local:
+    if NOLOCALS:
+      k.apply_opt(Opt(OptOps.NOLOCALS))
+    else:
+      # prioritize making expand axes local
+      local_axis_ranking = [(any(k.rngs[axis] not in b.src[1].get_idx().backward_slice for b in k.bufs), axis) \
+                              for axis in k.axes_of(AxisType.GLOBAL, AxisType.LOOP) if k.rngs[axis].src[0].op is Ops.CONST]
+      to_local: list[tuple[int, int]] = []
+      for _, axis in sorted(local_axis_ranking, key=lambda x: (-x[0], -x[1])):
+        local_size = prod(sz for _, sz in to_local)
+        local_sz: int|None = next((x for x in ([32] * (axis == 0) + [16,8,4,3,2]) if k.full_shape[axis] % x == 0 and local_size * x <= 128), None)
+        if local_sz is not None: to_local.append((axis, local_sz))
+      deleted_shape = 0
+      for axis, local_sz in sorted(to_local[:3]):
+        axis = axis - deleted_shape
+        will_delete_shape = local_sz == k.full_shape[axis]
+        k.apply_opt(Opt(OptOps.LOCAL, axis, local_sz))
+        if will_delete_shape: deleted_shape += 1
+
+  # **** threading ****
+
+  if k.ren.has_threads and k.ren.global_max is not None:
+    for threads in [32,16,12,8,6,5,4,3,2]:
+      # Skip if too many threads. Heuristic: use about 128K ops per thread
+      if threads > k.ren.global_max[0] or resolve(prod(k.full_shape) // (128 << 10) < threads): continue
+      for axis in k.axes_of(AxisType.LOOP):
+        if k.full_shape[axis] % threads == 0:
+          try: k.apply_opt(Opt(OptOps.THREAD, axis, threads))
+          except KernelOptError: pass
+          break
+      if k.applied_opts and k.applied_opts[-1].op is OptOps.THREAD: break
+
+  return k

tinygrad_repo/tinygrad/codegen/opt/postrange.py

@@ -0,0 +1,350 @@
+from __future__ import annotations
+import math, itertools
+from collections import defaultdict
+from typing import cast, Final
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, GroupOp, axis_letters, axis_colors
+from tinygrad.device import Buffer
+from tinygrad.dtype import dtypes, ImageDType
+from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod, merge_dicts, get_single_element, flatten
+from tinygrad.codegen.opt import Opt, OptOps, KernelOptError, check
+from tinygrad.codegen.simplify import pm_flatten_range
+from tinygrad.renderer import Renderer
+
+remove_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
+
+# NOTE: LOCAL and GROUP_REDUCE have the same priority. the order here matters
+axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3,
+               AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
+
+class Scheduler:
+  def __init__(self, ast:UOp, ren:Renderer):
+    self.ast, self.ren = ast, ren
+    self.dont_use_locals = self.ast.arg.dont_use_locals if self.ast.arg is not None else False
+    self.applied_opts = list(self.ast.arg.applied_opts) if self.ast.arg is not None else []
+
+  @property
+  def rngs(self):
+    # always in order by axistype
+    return sorted([u for u in self.ast.backward_slice if u.op is Ops.RANGE and u.vmax > 0], key=lambda x: (axis_to_pos[x.arg[-1]],) + x.arg[0:-1])
+  @property
+  def shape_len(self): return len(self.rngs)
+  @property
+  def full_shape(self): return [ssimplify(x.src[0]) for x in self.rngs]
+  @property
+  def axis_types(self): return [x.arg[-1] for x in self.rngs]
+  @property
+  def maxarg(self): return max([x.arg[0] for x in self.rngs], default=0)
+
+  # strings like ['g0', 'g1', 'l0', 'l1', 'l2', 'l3', 'l4', 'l5', 'R0', 'r0', 'r1', 'r2', 'u0', 'u1', 'u2']
+  def shape_str(self) -> list[str]:
+    ret: list[str] = []
+    cnt: dict[AxisType, int] = {}
+    for x in self.axis_types:
+      cnt[x] = (cnt[x] + 1) if x in cnt else 0
+      ret.append(f"{axis_letters[x]}{cnt[x]}")
+    return ret
+  def shape_str_to_axis(self, nms:list[str]) -> tuple[int, ...]: return tuple([self.shape_str().index(x) for x in nms])
+
+  def copy(self):
+    ret = Scheduler(self.ast, self.ren)
+    ret.dont_use_locals = self.dont_use_locals
+    ret.applied_opts = self.applied_opts[:]
+    return ret
+
+  kernel_cnt: Final[defaultdict[str, int]] = defaultdict(int)
+  def get_optimized_ast(self, name_override:str|None=None):
+    if name_override is not None: name = name_override
+    else:
+      kernel_type = "r" if self.reduceop is not None else "E"
+      name = kernel_type + colored('_', 'BLACK').join(['']+[colored(x.src[0].render(), color) for x,color in zip(self.rngs, self.colors())])
+      Scheduler.kernel_cnt[(function_name := to_function_name(name))] += 1
+      num = f"n{Scheduler.kernel_cnt[function_name]-1}" if Scheduler.kernel_cnt[function_name] > 1 else ""
+      name += colored(num, 'BLACK')
+    self.ast = graph_rewrite(self.ast, pm_flatten_range, name="flatten range")
+    return self.ast.replace(arg=KernelInfo(name=name, applied_opts=tuple(self.applied_opts), dont_use_locals=self.dont_use_locals), tag=1)
+
+  def _output_rngs(self) -> list[UOp]:
+    return flatten([[r for r in UOp.sink(*s.src[1:]).ranges if r.arg[-1] != AxisType.REDUCE] for s in self.ast.src if s.op is Ops.END])
+  def _globalizable_rngs(self) -> list[UOp]:
+    ret = [r for r in self._output_rngs() if r.arg[-1] == AxisType.LOOP]
+    # exclude any output ranges from global that don't appear in all BUFFERIZE
+    for x in self.ast.toposort():
+      if x.op is Ops.BUFFERIZE:
+        ret = [r for r in ret if r in x.ranges]
+    return ret
+
+  def convert_loop_to_global(self):
+    if not self.ren.has_local: return None
+
+    globalizible_rngs = self._globalizable_rngs()
+    rng = [x.replace(arg=x.arg[0:-1]+(AxisType.GLOBAL,)) if x in globalizible_rngs else x for x in self.rngs]
+
+    self.ast = self.ast.substitute(dict(zip(self.rngs, rng)))
+
+  def colors(self) -> list[str]:
+    output_rngs = self._output_rngs()
+    globalizible_rngs = self._globalizable_rngs()
+    ret = []
+    for x,r in zip(self.axis_types, self.rngs):
+      if self.dont_use_locals and x == AxisType.GLOBAL: ret.append("BLUE")
+      elif r not in output_rngs and x == AxisType.LOOP: ret.append("BLACK")
+      elif r not in globalizible_rngs and x == AxisType.LOOP: ret.append("white")
+      else: ret.append(axis_colors[x])
+    return ret
+  def colored_shape(self) -> str: return ' '.join([colored(f'{x.src[0].render():>4s}', color) for x,color in zip(self.rngs, self.colors())])
+
+  def shift_to(self, rng:UOp, amount:int, new_type:AxisType, top:bool=False, input_new_rng=None):
+    if (old_sz:=rng.src[0].divides(amount)) is None:
+      raise KernelOptError(f"{amount} can't divide {rng.src[0]} in {self.colored_shape()}")
+    new_rng = UOp.range(amount, self.maxarg+1, new_type) if input_new_rng is None else input_new_rng
+    replaced_rng = rng.replace(src=(UOp.const(dtypes.int, old_sz),))
+    sub_axis = (new_rng * old_sz + replaced_rng) if top else (replaced_rng * amount + new_rng)
+    self.ast = self.ast.substitute({rng:sub_axis}, name=f"shift {rng.arg[:-1]} {amount} {str(new_type).split('.')[1].lower()}")
+    return replaced_rng, new_rng
+
+  def ranges_of(self, *axis_type:AxisType) -> list[UOp]: return [r for r in self.rngs if r.arg[-1] in axis_type]
+  def axes_of(self, *axis_type:AxisType) -> list[int]: return [i for i,t in enumerate(self.axis_types) if t in axis_type]
+
+  def upcast_size(self) -> int: return prod(self.full_shape[a] for a in self.axes_of(AxisType.UPCAST, AxisType.UNROLL))
+
+  # copied from kernel.py
+  @property
+  def upcastable_dims(self) -> list[int]: return [i for i in self.axes_of(AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP) \
+                                                  if isinstance(s:=self.full_shape[i], int) and s > 1]
+  @property
+  def unrollable_dims(self) -> list[int]: return [i for i in self.axes_of(AxisType.GROUP_REDUCE, AxisType.REDUCE) \
+                                                  if isinstance(s:=self.full_shape[i], int) and s > 1]
+
+  def real_axis(self, op:OptOps, axis:int|None):
+    try:
+      if axis is None or op is OptOps.TC: return -1
+      if op is OptOps.UNROLL: return self.unrollable_dims[axis]
+      if op in {OptOps.GROUP, OptOps.GROUPTOP}: return self.axes_of(AxisType.REDUCE)[axis]
+      check(axis < self.shape_len, f"invalid axis on {axis=} {op=} {self.shape_len=}")
+      return axis
+    except IndexError as e: raise KernelOptError from e
+
+  def apply_opt(self, opt:Opt, append_opt:bool=True):
+    if opt.op is OptOps.NOLOCALS:
+      check(all(x not in {AxisType.WARP, AxisType.LOCAL, AxisType.GROUP_REDUCE} for x in self.axis_types), "no locals can't have locals")
+      if append_opt: self.applied_opts.append(opt)
+      self.dont_use_locals = True
+      return
+
+    if opt.op in {OptOps.LOCAL, OptOps.GROUP, OptOps.GROUPTOP}:
+      check(self.ren.has_local, "locals needed for opt")
+
+    rng = self.rngs[real_axis] if (real_axis:=self.real_axis(opt.op, opt.axis)) >= 0 else UOp(Ops.NOOP)
+
+    opt_to_at = {
+      OptOps.LOCAL: AxisType.LOCAL, OptOps.UPCAST: AxisType.UPCAST,
+      OptOps.UNROLL: AxisType.UNROLL, OptOps.GROUP: AxisType.GROUP_REDUCE,
+      OptOps.GROUPTOP: AxisType.GROUP_REDUCE, OptOps.THREAD: AxisType.THREAD}
+
+    ret = None
+    if opt.op in opt_to_at:
+      amt:int = int(rng.vmax+1) if opt.arg == 0 else cast(int, opt.arg)
+
+      # copied from kernel.py. prevents METAL compiler hangs
+      if self.reduceop is not None and (opt.op in {OptOps.GROUP, OptOps.GROUPTOP} or \
+                                        (self.group_for_reduces and opt.op not in {OptOps.NOLOCALS, OptOps.PADTO})):
+        upcast_local_sz = prod([self.full_shape[a] for a in self.axes_of(AxisType.UPCAST, AxisType.WARP, AxisType.LOCAL, AxisType.GROUP_REDUCE)])
+        smem_sz = amt*upcast_local_sz*self.reduceop.dtype.itemsize
+        check(smem_sz <= self.ren.shared_max, f"exceeds maximum shared memory size: needs {smem_sz}, max {self.ren.shared_max}")
+      if self.reduceop is not None and (opt.op in {OptOps.GROUP, OptOps.GROUPTOP}):
+        # We currently dont support a group within another rudece, TODO: fix if-contexts
+        reduce = [u for u in self.ast.backward_slice if u.op is Ops.REDUCE and rng in merge_dicts([r.ranges for r in u.src[1:]])][0]
+        check(not any(u.arg[-1] in (AxisType.REDUCE, AxisType.UNROLL, AxisType.GROUP_REDUCE) for u in reduce.ranges),
+          "cannot have a GROUP_REDUCE inside another reduce")
+
+      if opt.op is OptOps.UNROLL:
+        check(amt <= 32, "don't unroll more than 32")
+        check(rng.arg[-1] in {AxisType.GROUP_REDUCE, AxisType.REDUCE}, "unroll is for GROUP_REDUCE/REDUCE")
+      if opt.op is OptOps.UPCAST:
+        check((self.ren is not None and self.ren.device == "DSP") or amt <= 16, "don't upcast more than 16")
+        check(rng.arg[-1] in {AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP}, f"upcast is for GLOBAL/LOCAL/LOOP, not {rng.arg[-1]}")
+      if opt.op is OptOps.LOCAL:
+        check(not self.dont_use_locals, "can't use locals")
+        check(rng.arg[-1] in {AxisType.GLOBAL, AxisType.LOOP}, "local is for globals")
+      if opt.op is OptOps.THREAD:
+        check(self.ren is not None and self.ren.has_threads, "target does not support threads")
+        check(self.ren is not None and self.ren.global_max is not None and amt <= self.ren.global_max[0], "too many threads")
+        check(all(x is not AxisType.THREAD for x in self.axis_types), "already threaded")
+        check(rng in self._globalizable_rngs(), "can't apply range to this dim")
+      if opt.op in {OptOps.GROUP, OptOps.GROUPTOP}:
+        check(all(x.op is not OptOps.TC for x in self.applied_opts), "no grouping with tensor cores")  # TODO: why is this wrong?
+        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.TC:
+      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(-1 <= (tc_select:=cast(tuple, opt.arg)[0]) < len(self.ren.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)
+      except ValueError as e: raise KernelOptError(str(e))
+      check(ret is not None, "no tensor core available")
+    elif opt.op is OptOps.PADTO:
+      check(rng.src[0].op is Ops.CONST, "only pad const axes")
+      check(rng.arg[-1] not in {AxisType.UPCAST, AxisType.UNROLL}, "cannot pad upcasted") # TODO: why is this wrong?
+      check(rng.arg[-1] is not AxisType.THREAD, "cannot pad thread")
+      # ok to pad SUM if all parent ALU ops have f(0) = 0
+      if (r:=self.reduceop) is not None and rng.arg[-1] in (AxisType.GROUP_REDUCE, AxisType.REDUCE):
+        check(r.arg[0] is Ops.ADD and not r.op_in_backward_slice_with_self(*GroupOp.UnsafePad), f"cannot pad {r}")
+      new_sz = round_up(int(rng.vmax+1), cast(int, opt.arg))
+      check(rng.vmax+1 > new_sz//4, "pad adds more than quadruple the work")
+      replaced_rng = UOp.range(new_sz, *rng.arg)
+      replaces = {rng:replaced_rng}
+      valid = replaced_rng < rng.vmax+1
+      for b in self.bufs:
+        if rng in (i:=b.src[1].get_idx()).backward_slice_with_self:
+          replaces[b] = b.replace(src=(b.src[0],(valid&b.src[1].get_valid()).where(i, UOp.invalid())))
+      self.ast = self.ast.substitute(replaces, f"padto {rng.arg[:-1]} {opt.arg}")
+    elif opt.op is OptOps.SWAP:
+      try:
+        altrng:UOp = self.rngs[opt.arg]
+      except IndexError:
+        raise KernelOptError
+      check(rng.arg[-1] == AxisType.GLOBAL and altrng.arg[-1] == AxisType.GLOBAL, "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)},
+                                      name=f"swap {rng.arg[:-1]} {altrng.arg[:-1]}")
+      self.ast = graph_rewrite(self.ast, remove_tags, name="swap remove tags")
+    else:
+      raise KernelOptError(f"unsupported opt {opt.op}")
+
+    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]:
+    if not (reduceops := self.reduceops): raise KernelOptError("no reduce ops for TensorCore")
+    reduceop = reduceops[0]
+    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
+      in0, in1 = mul.src
+      try:
+        tensor_cores = self.ren.tensor_cores if tc_select == -1 else [self.ren.tensor_cores[tc_select]]
+      except IndexError:
+        raise KernelOptError(f"invalid tensor core choice {tc_select}")
+      for tc in tensor_cores:
+        if tc.dtype_in == in0.dtype.scalar() and tc.dtype_in == in1.dtype.scalar() and tc.dtype_out == reduceop.dtype.scalar():
+          # 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])
+          red_ranges = sorted(reduceop.src[1:], key=lambda x: -x.arg[0])
+          if DEBUG >= 3:
+            print(f"TC({axis}): {[(x.arg[0],x.vmax+1) for x in in0_ranges]}",
+                              f"{[(x.arg[0],x.vmax+1) for x in in1_ranges]} {[(x.arg[0],x.vmax+1) for x in red_ranges]}")
+          if not len(in0_ranges) or not len(in1_ranges) or not len(red_ranges): continue
+
+          # pick ranges
+          # NOTE: why are in1 and in0 switched?
+          axis_choices = list(itertools.product(in1_ranges, in0_ranges, red_ranges))
+          if not (axis < len(axis_choices)): continue
+          axes = list(axis_choices[axis])
+
+          # tag the reduceop
+          self.ast = self.ast.substitute({reduceop: reduceop.replace(tag="TC")})
+
+          # do optimizations and save the ranges
+          try:
+            for i,a in enumerate(axes):
+              idx = self.rngs.index(a)
+              if (a.vmax+1) % tc.dims[i] != 0:
+                if opt_level < 2: raise KernelOptError("tc padding requires opt_level >= 2")
+                # apply_opt should return the updated range?
+                self.apply_opt(Opt(OptOps.PADTO, idx, tc.dims[i]), append_opt=False) # PADTO might fail
+                axes[i] = self.rngs[idx]
+          except KernelOptError: continue
+
+          # we create the warp as a whole thing, in case some of these ranges are moved/removed later
+          warp = UOp.range(tc.threads, -1, AxisType.WARP)
+          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)
+              warp //= 2
+            elif opt[0] == "u":
+              axes[int(opt[1])], new_range = self.shift_to(axes[int(opt[1])], 2, AxisType.UPCAST)
+            else: raise RuntimeError(f"unsupported opt {opt[0]} in tensor cores")
+            ne.append(new_range)
+
+          for _, amt in tc.get_reduce_axes():
+            axes[2], new_range = self.shift_to(axes[2], amt, AxisType.UNROLL)
+            ne.append(new_range)
+
+          if use_tensor_cores != 2:
+            # fix the srcs
+            reduceop = get_single_element([x for x in self.ast.toposort() if x.op is Ops.REDUCE and x.tag == "TC"])
+            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 argsort(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.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])
+
+            # construct the op
+            # TODO: remove tc_upcast_axes from the arg
+            # do the reduce_axes always disappear? i think they don't
+            # they need to be moved into the WMMA srcs
+            wmma_arg = (str(tc), tc.dims, tc.dtype_in, tc.dtype_out, self.ren.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], tag=1),
+              UOp(Ops.CONTRACT, dtype=srcs[1].dtype.vec(tc.elements_per_thread[1]), src=(srcs[1],), arg=tc_upcast_axes[1], tag=1),
+              UOp.const(tc.dtype_out.vec(tc.elements_per_thread[2]), 0.0)), arg=wmma_arg, tag=1)
+            tc_uop = UOp(Ops.UNROLL, tc.dtype_out, (wmma,), arg=tc_upcast_axes[2], tag=1)
+
+            # 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 axes
+    return None
+
+  # helpers for hand_coded_optimizations
+  @property
+  def reduceops(self) -> list[UOp]: return [x for x in self.ast.backward_slice if x.op is Ops.REDUCE]
+  @property
+  def reduceop(self) -> UOp|None:
+    if not (red := self.reduceops): return None
+    return UOp(Ops.REDUCE_AXIS, red[0].dtype, red[0].src, (red[0].arg, ()))
+  @property
+  def bufs(self) -> list[UOp]: return [x for x in self.ast.toposort() if x.op is Ops.INDEX][::-1]
+  @property
+  def output_shape(self):
+    return [s if at not in {AxisType.REDUCE, AxisType.UNROLL, AxisType.GROUP_REDUCE} else 1 for s,at in zip(self.full_shape, self.axis_types)]
+  @property
+  def upcasted(self) -> int: return len(self.axes_of(AxisType.UPCAST, AxisType.UNROLL))
+  @property
+  def group_for_reduces(self) -> int: return len(self.axes_of(AxisType.GROUP_REDUCE))
+
+def bufs_from_ast(ast:UOp, dname:str) -> list[Buffer]:
+  glbls = sorted([x for x in ast.backward_slice if x.op is Ops.DEFINE_GLOBAL], key=lambda x: x.arg)
+  return [Buffer(dname, x.ptrdtype.size, x.dtype.base if not isinstance(x.dtype, ImageDType) else x.dtype) for x in glbls]
+
+def apply_opts(ast:UOp, ren:Renderer) -> UOp:
+  if ast.tag is not None: return ast
+  k = Scheduler(ast, ren)
+  k.convert_loop_to_global()
+  if ast.arg is not None and ast.arg.opts_to_apply is not None:
+    for opt in ast.arg.opts_to_apply: k.apply_opt(opt)
+  elif BEAM >= 1:
+    from tinygrad.codegen.opt.search import beam_search
+    rawbufs = bufs_from_ast(ast, ren.device)
+    k = beam_search(k, rawbufs, BEAM.value, bool(getenv("BEAM_ESTIMATE", 1)))
+  elif not NOOPT and (ast.arg is None or ast.arg.applied_opts == ()):
+    from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
+    # NOTE: hand_coded_optimizations doesn't support multiblock opts yet
+    if not any(u.op is Ops.BUFFERIZE for u in ast.backward_slice):
+      k = hand_coded_optimizations(k)
+  return k.get_optimized_ast(name_override=ast.arg.name if ast.arg is not None and ast.arg.name != "test" else None)

tinygrad_repo/tinygrad/codegen/opt/search.py

@@ -0,0 +1,191 @@
+import functools, math, time, multiprocessing, traceback, signal, atexit
+from dataclasses import replace
+from tinygrad.uop.ops import sym_infer, AxisType, pyrender
+from tinygrad.device import Device, Buffer, Compiler
+from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, diskcache_put, getenv, Context, colored, time_to_str, unwrap
+from tinygrad.helpers import IGNORE_BEAM_CACHE
+from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
+from tinygrad.tensor import Tensor
+from tinygrad.engine.realize import CompiledRunner, get_program
+from tinygrad.renderer import ProgramSpec
+from tinygrad.codegen.opt.postrange import Scheduler
+
+actions = [Opt(op=OptOps.UPCAST, axis=axis, arg=amt) for amt in [0,2,3,4,5,7] for axis in range(8)]
+actions += [Opt(op=OptOps.UNROLL, axis=axis, arg=amt) for amt in [0,4,7] for axis in range(5)]
+actions += [Opt(op=OptOps.LOCAL, axis=axis, arg=amt) for amt in [2,3,4,8,13,16,29] for axis in range(6)]
+actions += [Opt(op=OptOps.GROUPTOP, axis=axis, arg=amt) for amt in [13,16,28,29,32,49,64,256] for axis in range(3)]
+actions += [Opt(op=OptOps.GROUP, axis=axis, arg=amt) for amt in [0,4,8,16] for axis in range(3)]
+if getenv("BEAM_PADTO", 0): actions += [Opt(op=OptOps.PADTO, axis=axis, arg=amt) for amt in [32] for axis in range(7)]
+actions += [Opt(op=OptOps.LOCAL, axis=0, arg=32), Opt(op=OptOps.LOCAL, axis=6, arg=2)]
+actions += [Opt(op=OptOps.TC, axis=0, arg=(-1, 0, getenv("TC", 1)))]
+# covers resnet kernels (3 global * 3 reduce)
+actions += [Opt(op=OptOps.TC, axis=axis, arg=(-1, getenv("TC_OPT", 2), getenv("TC", 1))) for axis in range(9)]
+actions += [Opt(op=OptOps.SWAP, axis=axis_0, arg=axis_1) for axis_0 in range(5) for axis_1 in range(axis_0+1, 5)]
+actions += [Opt(op=OptOps.THREAD, axis=axis, arg=amt) for amt in [2,3,4,5,8,12,16,24,32,64] for axis in range(3)]
+if getenv("NOLOCALS"): actions += [Opt(op=OptOps.NOLOCALS)]
+
+def get_test_global_size(global_size, max_global_size, var_vals):
+  test_global_size = [sym_infer(sz, var_vals) for sz in global_size]
+  input_size = prod(test_global_size)
+  while prod(test_global_size) > max_global_size:
+    for j in range(len(global_size)-1,-1,-1):
+      if test_global_size[j] > 16:
+        test_global_size[j] //= 2
+        break
+  return test_global_size, input_size / prod(test_global_size)
+
+def _time_program(p:ProgramSpec, lib:bytes, var_vals:dict[str, int], rawbufs:list[Buffer], early_stop:float|None=None,
+                  allow_test_size:int=True, max_global_size:int|None=65536, clear_l2=False, cnt=3, name="test") -> list[float]:
+  factor = 1
+  if allow_test_size and p.global_size is not None and max_global_size is not None:
+    global_size, factor = get_test_global_size(p.global_size, max_global_size, var_vals)
+    p = replace(p, global_size=global_size)
+  try: car = CompiledRunner(p, precompiled=lib)
+  except AssertionError: return [math.inf] * cnt
+  tms = []
+  input_bufs = [rawbufs[i] for i in car.p.globals]
+  for _ in range(cnt):
+    if clear_l2:
+      if hasattr(dev:=Device[p.device], 'invalidate_caches'): dev.invalidate_caches()
+      else:
+        with Context(DEBUG=0, BEAM=0, CAPTURING=0, TRACK_MATCH_STATS=0): Tensor.ones(1024,1024).contiguous().realize(do_update_stats=False)
+    tms.append(unwrap(car(input_bufs, var_vals, wait=True))*factor)
+    if early_stop is not None and early_stop < min(tms): break
+  return tms
+
+class TimeoutException(Exception): pass
+def timeout_handler(signum, frame):
+  if DEBUG >= 2: print("*** BEAM COMPILE TIMEOUT")
+  raise TimeoutException()
+
+def _try_compile(x:tuple[int,Scheduler], compiler:Compiler) -> tuple[int, tuple[ProgramSpec, bytes, float]|None]:
+  if hasattr(signal, "alarm"):
+    signal.signal(getattr(signal, 'SIGALRM'), timeout_handler)
+    # set timeout
+    signal.alarm(getenv("BEAM_TIMEOUT_SEC", 10))
+  ret = None
+  try:
+    p = get_program(x[1].copy().get_optimized_ast(name_override="test"), x[1].ren)
+    assert p.uops is not None, "uop list wasn't generated?"
+    if len(p.uops) >= (uops_max:=getenv("BEAM_UOPS_MAX", 3000)) > 0:
+      if getenv("BEAM_LOG_SURPASS_MAX"): print(f"too many uops. {len(p.uops)=}, {uops_max=}")
+      raise RuntimeError("too many uops")
+    st = time.perf_counter()
+    prog = compiler.compile(p.src)
+    et = time.perf_counter() - st
+    ret = (p, prog, et)
+  except RuntimeError:
+    if DEBUG >= 4: traceback.print_exc()
+  except Exception as e:
+    if getenv("BEAM_STRICT_MODE"): raise e
+  finally:
+    if hasattr(signal, "alarm"): signal.alarm(0)
+  return x[0], ret
+
+# workers should not open devices and should ignore ctrl c and should not launch VIZ
+def _init_worker():
+  Context(ALLOW_DEVICE_USAGE=0, VIZ=0, TRACK_MATCH_STATS=0).__enter__()
+  signal.signal(signal.SIGINT, signal.SIG_IGN)
+
+def _ensure_buffer_alloc(bufs:list[Buffer]) -> list[Buffer]: return [buf.ensure_allocated() if buf is not None else buf for buf in bufs]
+
+# *** external API ***
+
+# get dictionary of all possible actions
+def get_kernel_actions(s:Scheduler, include_0=True, candidates:list[Opt]|None=None) -> dict[int, Scheduler]:
+  acted, max_up, max_lcl = {0:s} if include_0 else {}, getenv("BEAM_UPCAST_MAX", 256), getenv("BEAM_LOCAL_MAX", 1024)
+  kernel_actions = (actions if candidates is None else candidates).copy()
+
+  for i,a in enumerate(kernel_actions):
+    if a.axis is not None and a.op is not OptOps.TC:
+      try: ax = s.real_axis(a.op, a.axis)
+      except KernelOptError: continue
+      if (ax >= s.shape_len) or (s.full_shape[ax] == a.arg and Opt(a.op, a.axis, 0) in kernel_actions): continue
+    s2 = s.copy()
+    try:
+      s2.apply_opt(a)
+      up, lcl, tc_up = 1, 1, prod(tc.dims)//tc.threads if hasattr(s2, 'tensor_core') and (tc:=s2.tensor_core) else 1
+      for x,t in zip(s2.full_shape, s2.axis_types):
+        if t in (AxisType.UPCAST, AxisType.UNROLL): up *= x
+        elif t in (AxisType.WARP, AxisType.LOCAL, AxisType.GROUP_REDUCE): lcl *= x
+      if up//tc_up > max_up or lcl > max_lcl:
+        if getenv("BEAM_LOG_SURPASS_MAX"): print(f"too many upcast/local. {up//tc_up=}, {max_up=}, {lcl=}, {max_lcl=}")
+        continue
+      acted[i+1] = s2
+    except KernelOptError: pass
+  return acted
+
+beam_pool, BEAM_DEBUG = None, getenv("BEAM_DEBUG")
+def beam_search(s:Scheduler, rawbufs:list[Buffer], amt:int, allow_test_size=True, disable_cache=IGNORE_BEAM_CACHE.value):
+  global beam_pool
+  key = {"ast": s.ast.key, "amt": amt, "allow_test_size": allow_test_size, "device": s.ren.device, "suffix": s.ren.suffix}
+  if not disable_cache and CACHELEVEL >= 1 and (val:=diskcache_get("beam_search", key)) is not None:
+    ret = s.copy()
+    for o in val[len(s.applied_opts):]: ret.apply_opt(o)
+    return ret
+
+  beam: list[tuple[Scheduler, float]] = [(s, float("inf"))]
+  seen_libs = set()
+
+  default_parallel = multiprocessing.cpu_count() if s.ren.device in {"CUDA", "AMD", "NV", "METAL", "HIP"} else 0
+  if beam_pool is None and (workers := getenv("PARALLEL", default_parallel)):
+    beam_pool = multiprocessing.get_context("spawn").Pool(workers, _init_worker, (), getenv("BEAM_MAX_TASKS_PER_CHILD", 16))
+    @atexit.register
+    def close_pool(): beam_pool.close()
+
+  min_progress = getenv("BEAM_MIN_PROGRESS", 0.01)/1e6
+  if BEAM_DEBUG:
+    print("BEAM_SEARCH:")
+    print(pyrender(s.ast.replace(arg=None)))
+  if DEBUG >= 2: print(f"   0.00s:                from   1 ->   1 actions {s.colored_shape()}")
+
+  try:
+    rawbufs = _ensure_buffer_alloc(rawbufs)
+    var_vals: dict[str, int] = {k.expr:int(k.vmax+k.vmin)//2 for k in s.ast.variables()}
+    exiting, st = False, time.perf_counter()
+    dev = Device[s.ren.device]
+    while not exiting:
+      candidates: list[Scheduler] = flatten([get_kernel_actions(si, include_0=False).values() for si,_ in beam])
+      timed: list[tuple[Scheduler, float]] = []
+      _compile_fn = functools.partial(_try_compile, compiler=dev.compiler)
+      least_compute_ops = math.inf
+      for i,proc in (map(_compile_fn, enumerate(candidates)) if beam_pool is None else beam_pool.imap_unordered(_compile_fn, enumerate(candidates))):
+        if proc is None: continue
+        p, lib, compile_et = proc
+        if lib in seen_libs: continue
+        # filter out kernels that use 1000x more compute than the smallest
+        least_compute_ops = min(this_compute_ops:=sym_infer(p.estimates.ops, var_vals), least_compute_ops)
+        if least_compute_ops*1000 < this_compute_ops:
+          if getenv("BEAM_LOG_SURPASS_MAX"): print(f"too much compute. {this_compute_ops} when least is {least_compute_ops}")
+          continue
+        seen_libs.add(lib)
+        try: tms = _time_program(p, lib, var_vals, rawbufs, early_stop=beam[0][1]*3 if len(beam) else 1.0,
+                                 allow_test_size=allow_test_size, clear_l2=hasattr(dev, 'invalidate_caches'))
+        except Exception as e:
+          if BEAM_DEBUG: print(f"BEAM failed for opts: {candidates[i].applied_opts}\n{e}")
+          if isinstance(e, RuntimeError): continue
+          raise
+        timed.append((candidates[i], min(tms)))
+        if BEAM_DEBUG > 1:
+          print(f"{time.perf_counter() - st:7.2f}s: {i:5d} {len(unwrap(p.uops)):5d} uops",
+                f"{time_to_str(compile_et, w=12)} compile/{time_to_str(timed[-1][1], w=12)} run",
+                f"      {len(timed):4d}/{len(candidates):4d}         {timed[-1][0].colored_shape()}")
+        elif DEBUG >= 2:
+          print(f"\r{time.perf_counter() - st:7.2f}s: {time_to_str(timed[-1][1], w=12)}",
+                f"      {len(timed):4d}/{len(candidates):4d}         {timed[-1][0].colored_shape()}\033[K", end="")
+
+      # done
+      opts = sorted(timed, key=lambda x: x[1])
+      exiting = len(opts) == 0 or (opts[0][1] < min_progress) or (len(beam) > 0 and ((beam[0][1]-opts[0][1]) < min_progress))
+      if not exiting: beam = opts[:amt]
+      elif len(opts) > 0 and opts[0][1] < beam[0][1]: beam = opts[:1]
+      if DEBUG >= 2:
+        print(f"\r{time.perf_counter() - st:7.2f}s:", colored(time_to_str(beam[0][1], w=12), "green" if exiting else None),
+              f"from {len(candidates):3d} -> {len(opts):3d} actions\033[K", beam[0][0].colored_shape())
+  except KeyboardInterrupt as e:
+    if beam_pool is not None: beam_pool.terminate()
+    raise e
+
+  if CACHELEVEL >= 1: diskcache_put("beam_search", key, beam[0][0].applied_opts)
+  if BEAM_DEBUG: print(f"BEAM_SEARCH: final tm={time_to_str(beam[0][1], w=0)}, applied_opts={beam[0][0].applied_opts}")
+  return beam[0][0]

tinygrad_repo/tinygrad/codegen/opt/tc.py

@@ -0,0 +1,151 @@
+import math, functools
+from dataclasses import dataclass
+from tinygrad.dtype import DType, dtypes
+from tinygrad.helpers import getenv
+
+@dataclass(frozen=True)
+class TensorCore: # D = A * B + C, A is (M x K), B is (K x N), C and D are (M x N)
+  dims: tuple[int,int,int] # N, M, K
+  threads: int # number of threads that construct the warp
+  elements_per_thread: tuple[int, int, int] # elements per-thread to load/store from A/B/C
+  dtype_in: DType # dtype for A and B
+  dtype_out: DType # dtype for C and D
+  opts: tuple[str, ...] # ordered tuple of "ux" or "lx" specifying kernel opts to perform. "ux" upcasts dim x and "lx" localizes dim x
+  # (local_swizzle, upcast_swizzle, reduce_swizzle)
+  # l<num> is the num axis of the locals, similar for u<num> and upcasts, r<num> and reduces
+  swizzle: tuple[tuple[tuple[str, ...], tuple[str, ...], tuple[str, ...]], tuple[tuple[str, ...], tuple[str, ...], tuple[str, ...]]]
+  @functools.cache  # pylint: disable=method-cache-max-size-none
+  def _remaps(self) -> list[dict[str, str]]:
+    local_axes, upcast_axes, reduce_axes = len(self.get_local_axes()), len(self.get_upcast_axes()), len(self.get_reduce_axes())
+    fwd_st = [f"l{i}" for i in range(local_axes)] + [f"u{i}" for i in range(upcast_axes)] + [f"r{i}" for i in range(reduce_axes)]
+    return [dict(zip(fwd_st, sum(s, ()))) for s in self.swizzle]
+  def permutes_for_shape_str(self, shape_str:list[str]) -> tuple[tuple[int, ...], tuple[int, ...]]:
+    ret = [[shape_str.index(remap[ss]) if ss in remap else i for i,ss in enumerate(shape_str)] for remap in self._remaps()]
+    return tuple(ret[0]), tuple(ret[1])
+  @functools.cache  # pylint: disable=method-cache-max-size-none
+  def base_shape_str(self) -> list[str]:
+    ret = []
+    cnt = {'u': 0, 'l': 0}
+    for opt in self.opts:
+      ret.append(f"{opt[0]}{cnt[opt[0]]}")
+      cnt[opt[0]] += 1
+    # assumes you do the UNROLL after the opts
+    return ret + [f"r{i}" for i in range(len(self.get_reduce_axes()))]
+  def get_reduce_axes(self): return [(i, 2) for i in range(int(math.log2(self.dims[2])))]
+  def get_upcast_axes(self): return [opt for opt in self.opts if opt[0] == "u"]
+  def get_local_axes(self): return [opt for opt in self.opts if opt[0] == "l"]
+  def base_upcast_axes(self):
+    # this is defined in the swizzle. first we use the upcast axes, then the reduce
+    return ([f"r{i}" for i in range(len(self.get_reduce_axes()))] + [f"u{i}" for i in range(len(self.get_upcast_axes()))])[::-1]
+  def __str__(self): return "_".join(["WMMA"] + list(map(str, self.dims)) + [self.dtype_in.name, self.dtype_out.name])
+  def __post_init__(self):
+    # all axes have size 2, <local> <reduce> <upcast> is the order
+    local_axes, upcast_axes, reduce_axes = len(self.get_local_axes()), len(self.get_upcast_axes()), len(self.get_reduce_axes())
+    assert self.dims[0] * self.dims[1] == 2**(local_axes + upcast_axes), \
+      f"N({self.dims[0]}) x M({self.dims[1]}) != local({2**local_axes}) x upcast({2**upcast_axes}) with opts({self.opts})"
+    assert 2**local_axes == self.threads, f"{self.threads} threads construct the warp but found {2**local_axes} in {self.opts}"
+    assert 2**upcast_axes == self.elements_per_thread[2], \
+      f"{self.elements_per_thread[2]} elements from C are processed per thread but found {2**upcast_axes} in {self.opts}"
+    # check dims match opts
+    assert self.dims[0] == 2**len(gd:=[x for x in self.opts if x[1] == '0']), f"opts wrong on dims[0], {self.dims[0]} vs {gd}"
+    assert self.dims[1] == 2**len(gd:=[x for x in self.opts if x[1] == '1']), f"opts wrong on dims[1], {self.dims[1]} vs {gd}"
+    # NOTE: the K opts is implictly set by the dim
+    # check swizzle
+    assert len(self.swizzle[0]) == 3 and len(self.swizzle[1]) == 3, "swizzle has wrong part count"
+    assert len(self.swizzle[0][0]) == len(self.swizzle[1][0]) == local_axes, "local swizzle size is wrong"
+    assert len(self.swizzle[0][1]) == len(self.swizzle[1][1]) == upcast_axes, "upcast swizzle size is wrong"
+    assert len(self.swizzle[0][2]) == len(self.swizzle[1][2]) == reduce_axes, "reduce swizzle size is wrong"
+    assert all(len(s) == local_axes+upcast_axes+reduce_axes for s in self._remaps()), "remaps are the wrong size"
+    # check elements_per_thread
+    un, ln = 0, 0
+    zero_stride_0 = []
+    zero_stride_1 = []
+    for o in self.opts:
+      if o[1] == '0': zero_stride_0.append(o[0] + str(un if o[0] == 'u' else ln))
+      if o[1] == '1': zero_stride_1.append(o[0] + str(un if o[0] == 'u' else ln))
+      if o[0] == 'u': un += 1
+      if o[0] == 'l': ln += 1
+    # NOTE: all the zero_stride dims can be placed in any order in the swizzle
+    upcasted_0 = [x for x in (self.swizzle[0][1] + self.swizzle[0][2]) if x not in zero_stride_0 and x[0] != 'l']
+    upcasted_1 = [x for x in (self.swizzle[1][1] + self.swizzle[1][2]) if x not in zero_stride_1 and x[0] != 'l']
+    assert 2**len(upcasted_0) == self.elements_per_thread[0], f"mismatch in elements_per_thread[0], {upcasted_0} vs {self.elements_per_thread[0]}"
+    assert 2**len(upcasted_1) == self.elements_per_thread[1], f"mismatch in elements_per_thread[1], {upcasted_1} vs {self.elements_per_thread[1]}"
+
+# ***** NVIDIA *****
+
+cuda_tc_opts = ("u0","l0","l0","l1","l1","l1","u1")  # shared by all shapes with M=16 N=8
+
+# https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-multiply-accumulate-instructions
+cuda_81616 = [TensorCore(dims=(8,16,16), threads=32, elements_per_thread=(8,4,4), dtype_in=di, dtype_out=do, opts=cuda_tc_opts,
+  swizzle=((('r1', 'r2', 'l2', 'l3', 'l4'), ('u1', 'r3'), ('l0', 'l1', 'u0', 'r0')),
+           (('r1', 'r2', 'u0', 'l0', 'l1'), ('r0', 'r3'), ('l2', 'l3', 'l4', 'u1'))))
+  for di,do in [(dtypes.half,dtypes.float), (dtypes.bfloat16,dtypes.float), (dtypes.half,dtypes.half)]]
+cuda_81632_f8 = [TensorCore(dims=(8,16,32), threads=32, elements_per_thread=(16,8,4), dtype_in=di, dtype_out=do, opts=cuda_tc_opts,
+  swizzle=((('r2', 'r3', 'l2', 'l3', 'l4'), ('u1', 'r4'), ('l0', 'l1', 'u0', 'r0', 'r1')),
+           (('r2', 'r3', 'u0', 'l0', 'l1'), ('r1', 'r4'), ('l2', 'l3', 'l4', 'u1', 'r0'))))
+  for di,do in [(dtypes.fp8e4m3,dtypes.float),(dtypes.fp8e5m2,dtypes.float)]]
+cuda_8168_f16 = [TensorCore(dims=(8,16,8), threads=32, elements_per_thread=(4,2,4), dtype_in=di, dtype_out=do, opts=cuda_tc_opts,
+  swizzle=((('r1', 'r2', 'l2', 'l3', 'l4'), ('r0', 'u1'), ('l0', 'l1', 'u0')),
+           (('r1', 'r2', 'u0', 'l0', 'l1'), ('u1', 'r0'), ('l2', 'l3', 'l4'))))
+  for di,do in [(dtypes.half,dtypes.float), (dtypes.half,dtypes.half)]]
+cuda_8168_tf32 = [TensorCore(dims=(8,16,8), threads=32, elements_per_thread=(4,2,4), dtype_in=dtypes.float, dtype_out=dtypes.float, opts=cuda_tc_opts,
+  swizzle=((('r0', 'r1', 'l2', 'l3', 'l4'), ('u1', 'r2'), ('l0', 'l1', 'u0')),
+           (('r0', 'r1', 'u0', 'l0', 'l1'), ('u1', 'r2'), ('l2', 'l3', 'l4'))))]
+cuda_sm75: list[TensorCore] = cuda_8168_f16
+cuda_sm80: list[TensorCore] = cuda_81616 + cuda_8168_f16
+if getenv("ALLOW_TF32", 0): cuda_sm80 += cuda_8168_tf32
+cuda_sm89: list[TensorCore] = cuda_sm80 + cuda_81632_f8
+
+# ***** AMD *****
+
+# https://gpuopen.com/learn/wmma_on_rdna3/
+amd_rdna3 = [TensorCore(dims=(16,16,16), threads=32, elements_per_thread=(16,16,8), dtype_in=di, dtype_out=do,
+  opts=("l0","l0","l0","l0","l1","u1","u1","u1"),
+  swizzle=((('l4', 'u0', 'u1', 'u2', 'l0'), ('r1', 'r2', 'r3'), ('l1', 'l2', 'l3', 'r0')),
+           (('l0', 'l1', 'l2', 'l3', 'l4'), ('r1', 'r2', 'r3'), ('u0', 'u1', 'u2', 'r0'))))
+  for di,do in [(dtypes.half,dtypes.float),(dtypes.half,dtypes.half),(dtypes.bfloat16,dtypes.float)]]
+amd_rdna4 = [TensorCore(dims=(16,16,16), threads=32, elements_per_thread=(8,8,8), dtype_in=di, dtype_out=do,
+  opts=("l0","l0","l0","l0","u1","u1","u1","l1"),
+  swizzle=((('u0', 'u1', 'u2', 'l4', 'r2'), ('r0', 'r1', 'r3'), ('l0', 'l1', 'l2', 'l3')),
+           (('l0', 'l1', 'l2', 'l3', 'r2'), ('r0', 'r1', 'r3'), ('l4', 'u0', 'u1', 'u2'))))
+  for di,do in [(dtypes.half,dtypes.float),(dtypes.half,dtypes.half),(dtypes.bfloat16,dtypes.float),(dtypes.bfloat16,dtypes.bfloat16)]]
+
+# https://gpuopen.com/learn/amd-lab-notes/amd-lab-notes-matrix-cores-readme
+amd_cdna_161616 = [TensorCore(dims=(16,16,16), threads=64, elements_per_thread=(4,4,4), dtype_in=di, dtype_out=do,
+  opts=("l0","l0","l0","l0","u1","u1","l1","l1"),
+  swizzle=((('u0', 'u1', 'l4', 'l5', 'r2', 'r3'), ('r0', 'r1'), ('l0', 'l1', 'l2', 'l3')),
+           (('l0', 'l1', 'l2', 'l3', 'r2', 'r3'), ('r0', 'r1'), ('l4', 'l5', 'u0', 'u1'))))
+  for di,do in [(dtypes.half,dtypes.float),(dtypes.bfloat16,dtypes.float)]]
+
+amd_cdna_161632 = [TensorCore(dims=(16,16,32), threads=64, elements_per_thread=(8,8,4), dtype_in=di, dtype_out=do,
+  opts=("l0","l0","l0","l0","u1","u1","l1","l1"),
+  swizzle=((('u0', 'u1', 'l4', 'l5', 'r3', 'r4'), ('r0', 'r1'), ('l0', 'l1', 'l2', 'l3', 'r2')),
+           (('l0', 'l1', 'l2', 'l3', 'r3', 'r4'), ('r0', 'r1'), ('l4', 'l5', 'u0', 'u1', 'r2'))))
+  for di,do in [(dtypes.fp8e5m2,dtypes.float),(dtypes.fp8e4m3,dtypes.float),(dtypes.half,dtypes.float),(dtypes.bfloat16,dtypes.float)]]
+
+amd_cdna3 = amd_cdna_161632[:2] + amd_cdna_161616
+
+amd_cdna4 = amd_cdna_161632 + amd_cdna_161616
+
+# ***** Apple Metal *****
+
+metal = [TensorCore(dims=(8,8,8), threads=32, elements_per_thread=(2,2,2), dtype_in=di, dtype_out=do,
+  opts=("u0","l0","l1","l1","l0","l1"),
+  swizzle=((('r1', 'l1', 'l2', 'r2', 'l4'), ('r0',), ('u0', 'l0', 'l3')),
+           (('l0', 'r0', 'r1', 'l3', 'r2'), ('u0',), ('l1', 'l2', 'l4'))))
+  for di,do in [(dtypes.float,dtypes.float),(dtypes.half,dtypes.float),
+                (dtypes.half,dtypes.half),(dtypes.bfloat16,dtypes.float),(dtypes.bfloat16,dtypes.bfloat16)]]
+
+# ***** Apple AMX *****
+
+amx = [TensorCore(dims=(sz,sz,1), threads=1, elements_per_thread=(sz,sz,sz*sz), dtype_in=dt, dtype_out=dt,
+                  swizzle=(((), ('u0', 'u1', 'u2', 'u3', 'u4', 'u5', 'u6', 'u7'), ()),
+                           ((), ('u4', 'u5', 'u6', 'u7', 'u0', 'u1', 'u2', 'u3'), ())),
+                  opts=("u0","u0","u0","u0","u1","u1","u1","u1")) for dt,sz in [(dt, 64 // dt.itemsize) for dt in [dtypes.float]]]
+
+# ***** Intel ****
+
+intel = [TensorCore(dims=(8,8,16), threads=8, elements_per_thread=(16,16,8), dtype_in=dtypes.half, dtype_out=dtypes.float,
+                    opts=("l0","l0","l0","u1","u1","u1"),
+                    swizzle=((('r1', 'r2', 'r3'), ('u0', 'u1', 'u2'), ('l0', 'l1', 'l2', 'r0')),
+                             (('l0', 'l1', 'l2'), ('r1', 'r2', 'r3'), ('u0', 'u1', 'u2', 'r0'))))]

tinygrad_repo/tinygrad/codegen/simplify.py

@@ -0,0 +1,161 @@
+import itertools
+from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute, range_start, ImageDType
+from tinygrad.uop.symbolic import symbolic
+from tinygrad.helpers import partition, dedup
+from tinygrad.dtype import dtypes
+
+def flatten_range(r:UOp):
+  off = range_start[r.op]
+  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, Ops.END), 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}])
+def simplify_merge_adjacent(u:UOp) -> UOp|None:
+  reduce_ranges = [x.ranges for x in u.backward_slice_with_self if x.op is Ops.REDUCE]
+  # on END we only want to merge adjacent ranges, on REDUCE we want to try all combinations
+  for r0, r1 in (zip(u.ended_ranges, u.ended_ranges[1:]) if u.op is Ops.END else itertools.permutations(u.ended_ranges, 2)):
+    # check same type
+    if r0.arg[-1] == r1.arg[-1]:
+      # check if the ranges to merge are in the same reduces
+      if all((r0 in rngs) == (r1 in rngs) for rngs in reduce_ranges):
+        s0, s1 = r0.src[0], r1.src[0]
+        # do the merge
+        new_range = r0.replace(src=(s0*s1,))
+        nidx = graph_rewrite(u, _substitute+symbolic+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1},
+                             name=f"check_merge_{r0.arg[0]}_{r1.arg[0]}")
+
+        # check if it simplifies
+        if count_divmod(nidx) <= count_divmod(u):
+          u = nidx
+  return u
+
+pm_simplify_ranges = PatternMatcher([
+  (UPat((Ops.END, Ops.REDUCE), name="u"), simplify_merge_adjacent),
+])
+
+def mark_range_mod(ctx, r:UOp, c:UOp):
+  if r not in ctx and r.src[0].op is Ops.CONST and r.src[0].divides(c.arg) is not None: ctx[r] = c
+
+def do_substitute(ctx, x: UOp):
+  subs = {}
+  for k,v in ctx.items():
+    if v is not None:
+      subs[k] = k.replace(src=(k.src[0]//v,), arg=k.arg[0:-1]+(0,k.arg[-1]))*v + k.replace(src=(v,), arg=k.arg[0:-1]+(1,k.arg[-1]))
+  if not len(subs): return None
+  ret = x.substitute(subs).simplify()
+  ctx.clear()
+  return ret
+
+def dont_sub_ranges_for_image(ctx, x:UOp):
+  if isinstance(x.src[0].src[0].dtype, ImageDType):
+    for s in x.src[0].ranges: ctx[s] = None
+
+pm_split_ranges = PatternMatcher([
+  (UPat(Ops.RANGE, name="r")%UPat.cvar("c"), mark_range_mod),
+  (UPat(Ops.STORE, name="x"), dont_sub_ranges_for_image),
+  (UPat(Ops.SINK, name="x"), do_substitute),
+])
+
+# **** reduce simplification ****
+
+def no_range(u:UOp) -> bool: return not any(x.op is Ops.RANGE for x in u.backward_slice_with_self)
+
+def reduce_unparented(red:UOp):
+  if red.arg not in {Ops.ADD, Ops.MAX, Ops.MUL}: return None
+  assert all(x.op is Ops.RANGE for x in red.src[1:]), "some reduce srcs aren't ranges"
+  reduce_parented, reduce_unparented = partition(red.src[1:], lambda x: x in red.src[0].ranges)
+  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[0].cast(ret.dtype.scalar()).broadcast(ret.dtype.count)
+  if red.arg is Ops.MUL:
+    for r in reduce_unparented: ret = ret ** r.src[0].cast(ret.dtype.scalar()).broadcast(ret.dtype.count)
+  return ret
+
+pm_reduce_unparented = PatternMatcher([
+  # remove any ranges from a REDUCE that aren't referenced in the reduce source
+  (UPat(Ops.REDUCE, name="red"), reduce_unparented),
+])
+
+pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([
+  # lift x+y out of reduce on lt
+  ((UPat.var("x")+UPat.var("y")).or_casted() < UPat.var("c"), lambda x,y,c: (x < (c.cast(y.dtype)-y)) if no_range(y) and no_range(c) else None),
+  # lift x*y out of reduce
+  ((UPat.var("x")*UPat.var("y")) < UPat.var("c"), lambda x,y,c: (x < ((c+y-1) // y)) if no_range(y) and no_range(c) and y.vmin > 0 else None),
+  # fold the range
+  # bound from below
+  ((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(0, UPat.var("val")).reduce(UPat.var("r"), arg=Ops.ADD),
+   lambda r,cut,val: (r.src[0]-cut).maximum(0).minimum(r.src[0]).cast(val.dtype) * val if no_range(val) else None),
+  # bound from two sides
+  (((UPat.var("r")<UPat.var("lower")).logical_not()&(UPat(Ops.RANGE, name="r")<UPat.var("upper"))).where(UPat.var("val"), 0).reduce(UPat.var("r"),
+    arg=Ops.ADD), lambda r,lower,upper,val:
+      (upper.minimum(r.src[0])-lower.maximum(0)).maximum(0).minimum(r.src[0]).cast(val.dtype) * val if no_range(val) else None),
+  # bound from above
+  ((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(UPat.var("val"), 0).reduce(UPat.var("r"), arg=Ops.ADD),
+   lambda r,cut,val: cut.maximum(0).minimum(r.src[0]).cast(val.dtype) * val if no_range(val) else None),
+  # REDUCE on 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)),
+  # AND on WHERE
+  ((UPat(Ops.DEFINE_VAR, name="x") & UPat.var("y")).where(UPat.var("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)),
+  # MUL casted bool
+  ((UPat.var("x") * UPat.var("gate", dtype=dtypes.bool).cast()), lambda x,gate: gate.where(x, 0)),
+])+symbolic
+
+pm_reduce_load_collapse = pm_reduce_collapse + PatternMatcher([
+  # lift x+y out of reduce on ne
+  ((UPat.var("x")+UPat.var("y")).or_casted() != UPat.var("c"), lambda x,y,c: (x != (c.cast(y.dtype)-y)) if no_range(y) and no_range(c) else None),
+  # reduce on gated load becomes can substitute the range and remove the reduce
+  ((UPat.var("idx")!=(UPat(Ops.RANGE, name="r").or_casted())).where(0, UPat.var("expr")).reduce(UPat.var("r"), arg=Ops.ADD),
+   lambda r,idx,expr: (v:=(idx.cast(r.dtype) >= 0) & (idx.cast(r.dtype) < r.src[0])).where(expr.substitute({r:idx.cast(r.dtype).valid(v)}),0)),
+])
+
+def reduce_collapse(red:UOp, u:UOp, pm=pm_reduce_collapse):
+  for r in red.src[1:]:
+    included = u.toposort(gate=lambda x: r in x.ranges)
+    if any(x.op in {Ops.STORE, Ops.REDUCE} for x in included): return None
+    replaces: dict[UOp, UOp] = {}
+    for u in included:
+      for s in u.src:
+        if s in included or s in replaces or s.op in {Ops.CONST, Ops.VCONST, Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR}: continue
+        replaces[s] = UOp(Ops.DEFINE_VAR, dtype=s.dtype, arg=(f'in{len(replaces)}', s.vmin, s.vmax))
+    collapse_fxn = u.substitute(replaces).reduce(r, arg=Ops.ADD)
+    sink = graph_rewrite(collapse_fxn, pm, name="reduce_collapse")
+    if not no_range(sink): return None
+    u = sink.substitute({v:k for k,v in replaces.items()})
+  return u
+
+def reduce_load_collapse(red:UOp, u:UOp): return reduce_collapse(red, u, pm=pm_reduce_load_collapse)
+
+# remove REDUCE without loads (generic arange opt / indexing).
+pm_reduce_simplify = pm_reduce_unparented + PatternMatcher([
+  (UPat(Ops.REDUCE, src=(UPat.var("u"),), allow_any_len=True, arg=Ops.ADD, name="red"), reduce_collapse),
+])
+# remove REDUCE on load, comes from indexing a tensor with another tensor
+def no_load(u:UOp) -> bool: return not any(x.op is Ops.INDEX for x in u.backward_slice_with_self)
+pm_load_collapse = PatternMatcher([
+  (UPat(Ops.REDUCE, src=(UPat.var("u"), UPat()), name="red"), reduce_load_collapse),
+  # we want to make sure we dont do math on a loaded index since that can cause overflow, this undoes the rule in pm_reduce_load_collapse
+  ((UPat.var("x", dtypes.index)+UPat.var("y"))<UPat.var("c"), lambda x,y,c: x < c-y if no_load(y) and no_load(c) and not no_load(x) else None),
+])
+
+def cut_store_range(ctx, store:UOp, r:UOp):
+  # only cut ranges on CPU for now
+  if r.src[0].op is not Ops.CONST or ctx!="CPU": return None
+  if not (cuts:=[c.src[1].arg for c in store.get_consumer_map()[r] if c.op is Ops.CMPLT and r is c.src[0] and c.src[1].op is Ops.CONST]): return None
+  cuts = sorted(dedup([0] + cuts + [r.src[0].arg]))
+  ranges = [UOp.range((end-start), *(r.arg[0:-1]+(i,r.arg[-1]))) for i,(start,end) in enumerate(zip(cuts[:-1], cuts[1:]))]
+
+  return UOp.group(*[store.substitute({r: new_r+start}).end(new_r) for new_r, start in zip(ranges, cuts[:-1])])
+
+pm_split_store = pm_flatten_range+PatternMatcher([
+  (UPat(Ops.END, src=(UPat(Ops.STORE, name="store"), UPat.var("r"))), cut_store_range),
+])