tinygrad 0.10.2__py3-none-any.whl → 0.11.0__py3-none-any.whl

tinygrad/codegen/lowerer.py

@@ -1,161 +1,114 @@
 # the job of the lowerer is to do indexing
-import functools, itertools, operator, math
-from dataclasses import dataclass
+import functools, operator
 from typing import cast
-from tinygrad.dtype import dtypes, PtrDType
-from tinygrad.ops import KernelInfo, UOp, Ops, graph_rewrite, PatternMatcher, UPat, sint, identity_element, sint_to_uop
-from tinygrad.renderer import Renderer
-from tinygrad.helpers import all_int, prod, partition, flatten, unwrap
-from tinygrad.codegen.expander import expand_rewrite
-
-# returns the axes to create new_shape if new_shape can be created by combining axis from old_shape
-def get_contraction(old_shape:tuple[sint, ...], new_shape:tuple[sint, ...]) -> list[list[int]]|None:
-  acc_old, acc_new = list(itertools.accumulate(old_shape, operator.mul)), list(itertools.accumulate(new_shape, operator.mul))
-  try: split = [acc_old.index(acc)+1 if acc != 1 else 0 for acc in acc_new]
-  except ValueError: return None
-  return [list(range(st,ed)) for st,ed in zip([0]+split[:-1], split[:-1]+[len(old_shape)])]
+from dataclasses import dataclass
+from tinygrad.dtype import dtypes, AddrSpace, PtrDType
+from tinygrad.uop.ops import KernelInfo, UOp, Ops, PatternMatcher, UPat, sint_to_uop, AxisType, graph_rewrite
+from tinygrad.helpers import prod, partition, flatten
 
 # ***** indexing *****
-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: dims = dims[::-1]
-  # try to group first: (a, b, c, d) -> (ab, c, d)
-  limited = (grouped if (grouped := _group_dims(dims, max_sizes)) else dims) if max_sizes is not None else dims
-  # check if grouping failed
-  if max_sizes is not None and 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) if max_sizes is not None else dims
-  ret = raw_idxs = [UOp(Ops.SPECIAL, dtypes.int, (), (f"{prefix}{i}", s)) for i,s in enumerate(limited)]
-  if len(limited) < len(dims):
-    ret = []
-    if (contraction:=get_contraction(dims, limited)) is None: raise AssertionError(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)
-  elif len(limited) > len(dims):
-    a, b = len(limited), len(dims)
-    if a == 2 and b == 1: ret = [raw_idxs[0] * limited[1] + raw_idxs[1]]
-    if a == 3 and b == 1: ret = [raw_idxs[0] * (limited[1] * limited[2]) + raw_idxs[1] * limited[2] + raw_idxs[2]]
-    if a == 3 and b == 2: ret = [raw_idxs[0] * limited[1] + raw_idxs[1], raw_idxs[2]]
-  return ret[::-1] if reverse else ret
 
 @dataclass
 class IndexContext:
+  axis_types: tuple[AxisType, ...]
   idxs: list[UOp]
-  ridxs: list[UOp]
-  acc_num: int = 0
-
-def get_index(ast:UOp, opts:Renderer) -> IndexContext:
-  ki = ast.arg if isinstance(ast.arg, KernelInfo) else KernelInfo()
-  # NOTE: assumes the shape is <global dims> <local dims> <group_for_reduces> <reduces> <upcasts/unrolls>
-  full_shape = ast.full_shape
-  first_upcasted = len(full_shape)-ki.upcasted
-  # if there's no reduce, this is first_upcasted. assumes reduces are at the end
-  first_reduce = min([first_upcasted]+flatten(x.axis_arg for x in ast.toposort if x.op is Ops.REDUCE_AXIS))
-  local_loads = [x for x in ast.toposort if x.op is Ops.LOAD and x.src[0].op is Ops.DEFINE_LOCAL]
-  # NOTE: sum up the reduced axes looking across all local loads, yields the number of grouped reduces
-  group_for_reduces = sum([any(l.st_arg.shape[i]!=ast.src[0].st_arg.shape[i] for l in local_loads) for i in range(first_reduce,first_upcasted)])
-  global_dims = first_reduce-ki.local_dims
-
-  if opts.has_local:
-    if ki.dont_use_locals:
-      assert ki.local_dims == 0, "can't use locals if there's no local dims"
-      idxs = get_grouped_dims("idx", full_shape[:global_dims], opts.global_max, reverse=True)
+  start: int = 0
+
+def shape_to_idx(s, axis_types, start=0):
+  # indexes
+  idxs = []
+  for i, (s, at) in enumerate(zip(s, axis_types)):
+    if at in (AxisType.UPCAST, AxisType.UNROLL):
+      assert isinstance(s, int), "needs to be int to upcast/unroll"
+      idxs.append(UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(s), tuple(range(s))),), ((i,s),), tag=1))
     else:
-      # define indexes for GPU-like execution
-      idxs = get_grouped_dims("gidx", full_shape[:global_dims], opts.global_max, reverse=True) + \
-             get_grouped_dims("lidx", full_shape[global_dims:first_reduce+group_for_reduces], opts.local_max)
-  else:
-    # all loops are RANGES
-    idxs = [UOp(Ops.RANGE, dtypes.int, (sint_to_uop(0), sint_to_uop(g)), i) for i,g in enumerate(full_shape[:first_reduce])]
+      # all others are RANGES
+      idxs.append(UOp(Ops.RANGE, dtypes.int, (sint_to_uop(s),), start+i))
+  return idxs
 
-  # reduce loops
-  idxs += [UOp(Ops.RANGE, dtypes.int, (sint_to_uop(0), sint_to_uop(g)), i)
-    for i,g in enumerate(full_shape[first_reduce+group_for_reduces:first_upcasted], start=first_reduce+group_for_reduces)]
+def get_index(ast:UOp) -> IndexContext:
+  axis_types = ast.arg.axis_types if isinstance(ast.arg, KernelInfo) else ()
+  if len(ast.full_shape) != len(axis_types): axis_types = (AxisType.LOOP,)*len(ast.full_shape)
+  return IndexContext(axis_types, [], 0)
 
-  # upcast loops
-  for i,g in enumerate(full_shape[first_upcasted:], start=first_upcasted):
-    assert isinstance(g, int), "needs to be int to upcast/unroll"
-    idxs.append(UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(g), tuple(range(g))),), ((i,g),)))
+# ***** lowering (given index) *****
 
-  # late indexes (group for reduce)
-  ridxs = idxs[:]
-  for a in range(first_reduce, first_reduce+group_for_reduces):
-    ridxs[a] = UOp(Ops.RANGE, dtypes.int, (sint_to_uop(0), sint_to_uop(full_shape[a])), 1000+a)
+def subblock(ctx: IndexContext, full_new_idx: list[UOp], src: UOp):
+  lc = IndexContext(ctx.axis_types, full_new_idx, ctx.start+1000)
+  ctx.start = lc.start
+  return graph_rewrite(src, pm_lowerer, lc, name="subblock", bottom_up=True)
 
-  return IndexContext(idxs, ridxs)
+def lower_reduce_axis(ctx: IndexContext, x: UOp):
+  new_idxs = shape_to_idx(x.src[0].shape, ctx.axis_types, ctx.start)
+  full_new_idx = list(ctx.idxs)
+  for a in x.axis_arg: full_new_idx[a] = new_idxs[a]
 
-# ***** lowering (given index) *****
+  ret = subblock(ctx, full_new_idx, x.src[0])
 
-def lower_reduce_axis(ctx: IndexContext, x: UOp):
   # NOTE: always using ridxs is fine here
-  reduce_range, reduce_expand = partition([ctx.ridxs[i] for i in x.axis_arg], lambda y: y.op is Ops.RANGE)
+  reduce_range, reduce_expand = partition([full_new_idx[i] for i in x.axis_arg], lambda y: y.op is Ops.RANGE)
   assert all(x.op is Ops.UNROLL for x in reduce_expand), f"not all UNROLLS in {reduce_expand} for {x.axis_arg}"
-  alu_op: Ops = x.arg[0]
-  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))
-    ret = functools.reduce(lambda x,y: x.alu(alu_op, y), [ret.gep(i) for i in range(ret.dtype.count)])
-  if not len(reduce_range): return ret
-  # create ACC and assign
-  acc = UOp(Ops.DEFINE_ACC, x.dtype, (x.const_like(identity_element(alu_op, x.dtype.scalar())),) + tuple(reduce_range), (ctx.acc_num,))
-  ctx.acc_num += 1
-  return acc.assign(acc.alu(alu_op, ret))
-
-def lower_load_store(ctx: IndexContext, x: UOp):
-  idx, valid = x.st_arg.to_indexed_uops(ctx.ridxs if x.op is Ops.LOAD and x.src[0].op is Ops.DEFINE_LOCAL else ctx.idxs)
-  buf = x.src[0]
-  if x.op is Ops.LOAD:
-    barrier = (UOp(Ops.BARRIER, dtypes.void, (x.src[2],)),) if x.src[0].op is Ops.DEFINE_LOCAL else ()
-    return UOp(Ops.LOAD, x.dtype, (buf.index(idx, valid),) + barrier)
-  # NOTE: only store the local reduceop in the threads that are actually doing the reduce
-  if cast(PtrDType, x.src[0].dtype).local and x.src[2].op is Ops.ASSIGN:
-    reduce_input = x.src[2].src[1].src[1] if x.src[2].src[1].src[1] is not x.src[2].src[0] else x.src[2].src[1].src[0]
-    store_back = reduce_input.op is Ops.LOAD and cast(PtrDType, reduce_input.src[0].dtype).local
-  else: store_back = False
-  # NOTE: If we're storing the reduced value back into each thread, need to zero-out the reduced axes
-  if store_back: idx, _ = x.st_arg.to_indexed_uops([u.const_like(0) if u in x.src[2].src else u for u in ctx.idxs])
-  if (not cast(PtrDType, x.src[0].dtype).local) or store_back:
-    for oidx, ridx in zip(ctx.idxs, ctx.ridxs):
-      if oidx is not ridx: valid = valid * oidx.eq(0)
-  return UOp(Ops.STORE, dtypes.void, (buf.index(idx, valid), x.src[2]))
-
-def lower_const(x:UOp):
-  assert all(v.mask is None for v in unwrap(x.st).views), f"VIEW in CONST/DEFINE_VAR source must be unmasked, got {x.st}"
-  return x.replace(src=())
+    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 UOp(Ops.REDUCE, x.dtype, (ret,)+tuple(reduce_range), x.arg[0])
+
+def lower_store(ctx: IndexContext, x: UOp, buf: UOp):
+  # TODO: reenable after REDUCE_AXIS is fixed
+  #assert x.src[1].shape == x.src[0].shape, f"shape mismatch on store {x.src[1].shape} != {x.src[0].shape}"
+
+  new_idxs = shape_to_idx(x.src[0].shape, ctx.axis_types, ctx.start)
+  idx, valid = x.st_arg.to_indexed_uops(new_idxs)
+  used_idxs = [x for x in UOp.sink(idx, valid).toposort() if x in new_idxs]
+  real_new_idxs = []
+  for i in range(len(x.src[0].shape)):
+    if new_idxs[i] in used_idxs or len(ctx.idxs) <= i: real_new_idxs.append(new_idxs[i])
+    else: real_new_idxs.append(ctx.idxs[i])
+
+  stored = subblock(ctx, real_new_idxs, x.src[1])
+  used_ranges = [x for x in used_idxs if x.op is Ops.RANGE]
+  ret = buf.index(idx, valid).store(stored, *used_ranges)
+
+  # insert BARRIER if we are ending a LOCAL, IF if we are ending a GROUP_REDUCE
+  if cast(PtrDType, buf.dtype).addrspace == AddrSpace.LOCAL and \
+      any(ctx.axis_types[x.arg%1000] in {AxisType.GROUP_REDUCE, AxisType.LOCAL} for x in used_ranges):
+    ret = ret.barrier()
+    range_gates = [x.eq(0) for x in used_ranges if ctx.axis_types[x.arg%1000] == AxisType.GROUP_REDUCE]
+    if len(range_gates): ret = UOp(Ops.IF, src=(functools.reduce(operator.and_, range_gates), ret))
+  return ret
+
+def fixup_wmma(ctx:IndexContext, x:UOp):
+  if x.tag is not None: return None
+  new_idxs = shape_to_idx(x.src[0].shape, ctx.axis_types, ctx.start)
+  full_new_idx = list(ctx.idxs)
+  for a in x.arg[-1]: full_new_idx[a] = new_idxs[a]
+
+  srcs = subblock(ctx, full_new_idx, UOp.sink(*x.src)).src
+
+  # NOTE: this assumes these are expanded. which now shouldn't change anything
+  new_x_arg_m2 = tuple([tuple([(full_new_idx[a].arg[0][0], sz) for a,sz in v]) for v in x.arg[-2]])
+  new_x_arg_m1 = tuple([full_new_idx[a].arg[0][0] for a in x.arg[-1]])
+  return x.replace(src=srcs, arg=x.arg[:-2]+(new_x_arg_m2, new_x_arg_m1), tag=1)
 
 pm_lowerer = PatternMatcher([
+  # TODO: remove these hacks
+  # hack for old style CONST(VIEW) (now it's just VIEW(CONST))
+  (UPat((Ops.DEFINE_VAR, Ops.CONST), src=(UPat(Ops.VIEW, name="v"),), name="c"), lambda c,v: c.replace(src=()).view(v.arg)),
+  # hack for old style VALID (now it's just VIEW(CONST))
+  (UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar("c"), UPat(Ops.CONST, arg=0)), lambda c,v: c.replace(src=()).view(v.arg)),
+
+  # consts and loads
+  (UPat(Ops.VIEW, src=(UPat((Ops.CONST, Ops.DEFINE_VAR), name="c"),), name="view"),
+   lambda ctx,view,c: c if all(x.mask is None for x in view.arg.views) else view.arg.to_indexed_uops(ctx.idxs)[1].where(c, c.const_like(0))),
+  (UPat(Ops.LOAD, src=(UPat.var("buf").view(),), allow_any_len=True, name="x"),
+   lambda ctx,buf,x: UOp(Ops.LOAD, x.dtype, (buf.index(*x.st_arg.to_indexed_uops(ctx.idxs)),)+x.src[1:])),
+
+  # reduce/view_const
   (UPat(Ops.REDUCE_AXIS, name="x"), lower_reduce_axis),
-  (UPat((Ops.CONST, Ops.DEFINE_VAR), src=(UPat(Ops.VIEW),), name="x"), lower_const),
-  (UPat(Ops.VALID, src=(UPat(Ops.VIEW),), name="x"), lambda ctx,x: x.st_arg.to_indexed_uops(ctx.idxs)[1]),
-  # rewrite LOAD/STORE VIEW to LOAD/STORE with indexed
-  (UPat((Ops.LOAD, Ops.STORE), src=(UPat(), UPat(Ops.VIEW)), allow_any_len=True, name="x"), lower_load_store),
-  (UPat(Ops.INDEX, src=(UPat.var("b"), UPat.var("idx"), UPat.const(dtypes.bool, True))), lambda b, idx: b.index(idx)),
-])
+  (UPat(Ops.STORE, src=(UPat.var("buf").view(),), allow_any_len=True, name="x"), lower_store),
+  (UPat(Ops.WMMA, name="x"), fixup_wmma),
 
-def rewrite_shapetracker_with_index(ast:UOp, opts:Renderer) -> UOp:
-  sink = graph_rewrite(ast, pm_lowerer, ctx=get_index(ast, opts))
-  # expand_rewrite turns this into a vectorized program
-  return expand_rewrite(sink)
+  # axis fixups for WMMA
+  (UPat((Ops.CONTRACT, Ops.UNROLL), name="x"),
+   lambda ctx,x: x.replace(tag=1, arg=tuple([(ctx.idxs[a].arg[0][0], sz) for a,sz in x.arg])) if x.tag is None else None),
+])

tinygrad/codegen/opt/__init__.py

@@ -0,0 +1,38 @@
+# opt opinionatedly transforms an ast into an optimized ast using either heuristics or beam search
+
+from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
+from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops
+from tinygrad.helpers import NOOPT, BEAM, USE_TC, getenv
+from tinygrad.renderer import Renderer
+from tinygrad.uop.spec import type_verify
+
+def get_optimized_ast(ast:UOp, renderer:Renderer) -> UOp:
+  """
+  Optimize an AST based on heuristics or BEAM search.
+
+  Args:
+    ast: The Ops.SINK rooted AST
+    renderer: The renderer used to generate the code
+
+  Returns:
+    The Ops.SINK rooted AST transformed to apply the opts and with a KernelInfo in the arg.
+  """
+
+  k = Kernel(ast, opts=renderer)
+  if ast.arg is not None and ast.arg.opts_to_apply is not None: k.apply_opts(ast.arg.opts_to_apply)
+  elif not NOOPT:
+    if not k.apply_tensor_cores(USE_TC.value): k.apply_opts(hand_coded_optimizations(k))
+    if BEAM >= 1:
+      from tinygrad.codegen.opt.search import beam_search, bufs_from_lin
+      kb = Kernel(ast, opts=renderer)
+      rawbufs = bufs_from_lin(kb, allocate=False)
+      k = beam_search(kb, rawbufs, BEAM.value, bool(getenv("BEAM_ESTIMATE", 1)))
+  ret = k.get_optimized_ast()
+  if __debug__: type_verify(list(ret.toposort()))
+  return ret
+
+pm_optimize = PatternMatcher([
+  (UPat(Ops.SINK, name="ast"), lambda ctx,ast:
+   get_optimized_ast(ast, ctx) if (ast.arg is None or ast.arg.opts_to_apply is not None) and ast.src[0].st is not None else None),
+])

tinygrad/codegen/opt/heuristic.py

@@ -0,0 +1,125 @@
+import itertools
+from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps, KernelOptError, AxisType
+from tinygrad.helpers import getenv, DEBUG, prod, NOLOCALS
+from tinygrad.dtype import ImageDType
+from tinygrad.uop.ops import Ops, resolve
+
+def hand_coded_optimizations(k:Kernel) -> list[Opt]:
+  # make a copy so it does not mutate the input
+  k = k.copy()
+
+  # 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.opts.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.opts.has_shared and \
+    (mulop:=k.reduceop.src[0]).op is Ops.MUL and mulop.src[0].op is Ops.LOAD and mulop.src[1].op is Ops.LOAD:
+    st0, st1 = k.sts[k.bufs.index(mulop.src[0])], k.sts[k.bufs.index(mulop.src[1])]
+    strides0, strides1 = st0.real_strides(), st1.real_strides()
+    def has_expanded_axis(shape, strides): return any(resolve(s > 1) and not resolve(st != 0) for s,st in zip(shape,strides))
+    if strides0[first_reduce:=(k.axes_of(AxisType.REDUCE)[0])] == 1 and \
+      not (has_expanded_axis(st0.shape, strides0) and has_expanded_axis(st1.shape, strides1)):
+      for global_idx in k.axes_of(AxisType.GLOBAL):
+        if k.full_shape[first_reduce]%MV_THREADS_PER_ROW == 0 and k.full_shape[global_idx]%(MV_BLOCKSIZE*MV_ROWS_PER_THREAD) == 0:
+          if DEBUG >= 3:
+            print(f"MATVEC: {k.full_shape=} {first_reduce=} {strides0=} {MV_BLOCKSIZE=} {MV_THREADS_PER_ROW=} {MV_ROWS_PER_THREAD=}")
+          if MV_THREADS_PER_ROW > 1: k.apply_opt(Opt(OptOps.GROUP, 0, MV_THREADS_PER_ROW))
+          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.applied_opts
+
+  # are we grouping? (requires local shape support)
+  if resolve(prod(k.sts[0].shape[i] for i in k.upcastable_dims) <= 2048, False):
+    for sz in [16]:
+      try:
+        k.apply_opt(Opt(OptOps.GROUPTOP, 0, sz))
+        break
+      except KernelOptError: pass
+
+  # upcast float4 images
+  for buf_index,buf in enumerate(k.bufs):
+    if isinstance(buf.src[0].dtype, ImageDType):
+      if (unit_stride_axes_mul_4 := [i for i in k.sts[buf_index].unit_stride_axes(ignore_valid=True) if k.sts[buf_index].shape[i]%4 == 0]):
+        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))
+
+  # no more opt if we are grouping
+  if k.group_for_reduces: return k.applied_opts
+
+  # **** 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:
+    if k.full_shape[axis] <= 7 and any(st.axis_is_masked(axis) for st in k.sts) 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.opts is not None and k.opts.device == "DSP"
+  upcasted_axis: set[int] = set()
+  while resolve(prod(k.sts[0].shape[i] for i in k.upcastable_dims) >= 1024):
+    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
+      if any(st.views[-1].strides[axis] == 0 and \
+             all(x != 0 for t,x in zip(k.axis_types, st.real_strides()) if t in (AxisType.UPCAST, AxisType.UNROLL)) for st in k.sts):
+        xb_choices.append((sum(st.views[-1].strides[axis]>0 for st in k.sts),
+                           sum(st.views[-1].strides[axis] for st in k.sts), 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:
+    upcast_size = prod(k.full_shape[a] for a in k.axes_of(AxisType.UPCAST, AxisType.UNROLL))
+    if k.unrollable_dims and (upcast_size <= 4 or not k.axes_of(AxisType.UNROLL)) and (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]:
+    # TODO: somehow this never hits a reduce
+    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.opts.has_local:
+    if NOLOCALS:
+      k.apply_opt(Opt(OptOps.NOLOCALS))
+    else:
+      # prioritize making expand axes local
+      local_axis_ranking = [(any(st.views[-1].strides[axis] == 0 for st in k.sts), axis) for axis in k.axes_of(AxisType.GLOBAL, AxisType.LOOP)]
+      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
+
+  return k.applied_opts