IQ.Pilot Prebuilt Release @ 27f668a

artifacts/package_runtime/tinygrad/codegen/late/coalesce.py

@@ -0,0 +1,165 @@
+import itertools, functools
+from collections import defaultdict
+from tinygrad.dtype import dtypes, AddrSpace, Invalid, DType
+from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp, shape_to_shape_arg
+from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate
+from tinygrad.helpers import getenv, IMAGE, OSX, ceildiv, is_image_shape
+from tinygrad.renderer import Renderer
+
+# ***** image load valid simplification *****
+
+@functools.cache
+def _drop_valid_stmts(valid:UOp, idx:UOp, height:int, width:int) -> list[UOp]:
+  # can drop valid if idx is out of bound when valid is False
+  drop_stmt = []
+  for i,stmt in enumerate(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)
+      if testidx.index(0).vmax < 0 or testidx.index(1).vmax < 0:
+        drop_stmt.append(stmt)
+        continue
+
+    # check if idx is out of bound when X is on the wrong side of the bound: X in [c+1, vmax] or [vmin, c-1]
+    lo, hi = (c + 1, X.vmax) if is_upper_bound else (X.vmin, c - 1)
+    if lo <= hi:
+      fake = UOp.variable(f"fake{i}", lo, hi, X.dtype)
+      for coord,b in zip(idx.src, (width, height)):
+        rw = coord.substitute({X:fake}).simplify()
+        if rw.vmin >= b or rw.vmax < 0:
+          drop_stmt.append(stmt)
+          break
+  return drop_stmt
+
+def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
+  idx = uop_given_valid(valid, start_idx)
+  return None if idx is start_idx or idx is start_idx.simplify() else buf.index(idx.valid(valid))
+
+def simplify_valid_image_load(buf:UOp, idx_y:UOp, idx_x:UOp, valid:UOp) -> UOp|None:
+  if not is_image_shape(buf._shape): return None
+  if idx_x.dtype != idx_y.dtype: idx_x, idx_y = idx_x.cast(dtypes.int), idx_y.cast(dtypes.int)
+  start_idx = idx_x.stack(idx_y)
+  idx = uop_given_valid(valid, start_idx)
+  drop_stmt = _drop_valid_stmts(valid, idx, buf._shape[0], buf._shape[1])
+
+  if not drop_stmt and idx is start_idx: return None
+  new_valid = UOp.uprod(*ss) if (ss:=[s for s in valid.split_uop(Ops.AND) if s not in drop_stmt]) else None
+  idx_y, idx_x = idx.index(1), idx.index(0)
+  if new_valid is not None: return buf.index(idx_y.valid(new_valid), idx_x.valid(new_valid), dtype=dtypes.float)
+  return buf.index(idx_y, idx_x, dtype=dtypes.float)
+
+indexing_simplify = 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)),
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("valid").where(UPat.var("idx_y"), UPat(arg=Invalid)),
+                                         UPat.var("valid").where(UPat.var("idx_x"), UPat(arg=Invalid)))), simplify_valid_image_load),
+])
+
+# get list of (height, width) that do not require pitch padding
+def image_valid_dims(base:DType, size:int, arch:str) -> list[tuple[int,int]]:
+  if (ALIGN:=next((int(p.split('=')[1]) for p in arch.split(',') if p.startswith("IMAGE_PITCH_ALIGNMENT=")), 0)) == 0: return []
+  MAXW, pxls = 16384, size // 4
+  if base not in (dtypes.half, dtypes.float) or size > 4*MAXW*MAXW: return []
+  # height=1 images just need to abide by alignment requirements in bytes, not pixels!
+  if size % (ALIGN * 4) != 0: return [] if (base.itemsize * size) % (64 if OSX else ALIGN) != 0 or pxls > MAXW else [(1, pxls)]
+  return [(pxls//ALIGN//k, ALIGN*k) for k in range(ceildiv(pxls//ALIGN, MAXW), min(pxls//ALIGN, MAXW//ALIGN)+1) if (pxls//ALIGN)%k == 0]
+
+def transform_to_image(ctx, buf:UOp, x:UOp) -> UOp|None:
+  shapes, ren = ctx
+  if not IMAGE or ren.target.device not in {"QCOM", "CL", "PYTHON", "NULL"}: return None
+  valid, x = x.get_valid(), x.get_idx()
+  # search for dims that drop the most valid statements
+  best_drop, cands = -1, []
+  for ch, cw in [shapes[buf.arg.slot]] if buf.arg.slot in shapes else image_valid_dims(buf.dtype, buf.max_numel(), ren.target.arch):
+    cidx = uop_given_valid(valid, ((x//4)%cw).stack(x//(4*cw)))
+    dropped = len(_drop_valid_stmts(valid, cidx, ch, cw))
+    if dropped > best_drop: best_drop, cands = dropped, [(ch, cw, cidx)]
+    elif dropped == best_drop: cands.append((ch, cw, cidx))
+  # if no candidates, we don't rewrite
+  if len(cands) == 0: return None
+  # and tiebreak with indexing complexity (ie. number of nodes)
+  h, w, cidx = cands[0] if len(cands) == 1 else min(cands, key=lambda cand: len(cand[2].index(1).simplify().backward_slice))
+  buf = buf.replace(src=(shape_to_shape_arg((h, w, 4)),))
+  shapes[buf.arg.slot] = (h, w)
+  if valid.op is not Ops.CONST or valid.val is not True:
+    return buf.index(cidx.src[1].valid(valid), cidx.src[0].valid(valid), dtype=dtypes.float)
+  else:
+    return buf.index(cidx.src[1], cidx.src[0], dtype=dtypes.float)
+
+pm_simplify_add_image = PatternMatcher([
+  (UPat(Ops.SHRINK, src=(UPat(Ops.PARAM, name="buf"), UPat(name="x"), UPat(arg=4))), transform_to_image),
+  # image load/store is always float
+  (UPat(Ops.INDEX, dtype=dtypes.float, name="x").load(dtype=dtypes.half), lambda x: x.load().cast(dtypes.half)),
+  (UPat(Ops.INDEX, dtype=dtypes.float, name="x").store(UPat(name="d", dtype=dtypes.half)), lambda x,d: x.store(d.cast(dtypes.float))),
+  (UPat.var("x", dtype=dtypes.float).cast(dtypes.half).cast(dtypes.float), lambda x: x),
+])
+
+def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
+  if getenv("DMC"): return sink
+
+  # collect
+  memory: defaultdict[tuple[Ops, UOp, UOp|str, UOp], dict[int, list[UOp]]] = defaultdict(dict)
+  for u in sink.toposort():
+    # TODO: this should handle images too, it's just memory coalescing
+    if u.op in {Ops.LOAD, Ops.STORE}:
+      assert len(u.src) == (2 if u.op is Ops.STORE else 1), "memory coalescing does not support gated loads/stores"
+      assert u.src[0].op is Ops.INDEX, f"memory coalescing should be on INDEX, not {u.src[0].op}"
+      buf, idx_u = u.src[0].src
+      if buf.addrspace == AddrSpace.REG: continue
+      idx, valid = idx_u.get_idx(), idx_u.get_valid()
+      root_src: UOp|str
+      if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: root_src, arg = idx.src[0], idx.src[1].val
+      elif idx.op is Ops.ADD and idx.src[0].op is Ops.CONST: root_src, arg = idx.src[1], idx.src[0].val
+      elif idx.op is Ops.CONST and idx.val is Invalid: root_src, arg = "INVALID", 0
+      elif idx.op is Ops.CONST: root_src, arg = "CONST", idx.val
+      else: root_src, arg = idx, 0
+      memory[(u.op, buf, root_src, valid)].setdefault(arg, []).append(u)
+
+  # build replacements
+  replacements = {}
+  for (op,buf,base,valid),offsets in memory.items():
+    # allowed lengths (copied in)
+    lengths = []
+    must_divide = True
+    if ctx is not None and ctx.target.device == "DSP":
+      lengths = [128,64,32,16,8,4]
+      must_divide = False
+    elif buf.dtype not in (dtypes.float, dtypes.half, dtypes.int, dtypes.uint, *dtypes.fp8s) and not is_image_shape(buf._shape):
+      pass
+    elif buf.addrspace == AddrSpace.REG:
+      pass
+    elif is_image_shape(buf._shape):
+      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 == dtypes.half and getenv("ALLOW_HALF8") else [4,2]
+    lengths.append(1)  # worst case, it's not folded
+    # do the grouping
+    grouped_offsets = [[x for _,x in group] for _,group in itertools.groupby(enumerate(sorted(offsets.keys())), lambda x: x[1]-x[0])]
+    for full_grp in grouped_offsets:
+      while len(full_grp):
+        offset = (base+full_grp[0]) if isinstance(base, UOp) else UOp.const(full_grp[0])
+        length = [l for l in lengths if l <= len(full_grp) and (not must_divide or offset.divides(l) is not None)][0]
+        grp = full_grp[:length]
+        # NOTE: we apply the valid again after we determine the length
+        offset = offset.valid(valid) if valid is not None else offset
+        idx = UOp(Ops.SHRINK, src=(buf, offset, UOp.const(len(grp)))) if len(grp) > 1 else buf.index(offset)
+        if op == Ops.STORE:
+          datas = []
+          for i,g in enumerate(grp):
+            assert len(offsets[g]) == 1, f"attempting multiple stores: {len(offsets[g])}"
+            datas.append(offsets[g][0].src[1])
+          store = idx.store(UOp.stack(*datas) if len(datas) > 1 else datas[0])
+          for i,g in enumerate(grp): replacements[offsets[g][0]] = store
+        else:
+          ld = idx.load()
+          for i,g in enumerate(grp):
+            for oo in offsets[g]:
+              replacements[oo] = ld.index(i) if len(grp) > 1 else ld
+        full_grp = full_grp[length:]
+
+  # apply
+  return sink.substitute(replacements, name="memory coalescing")

artifacts/package_runtime/tinygrad/codegen/late/gater.py

@@ -0,0 +1,29 @@
+# this is a temporary intermediate step while we remove this index style
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops
+from tinygrad.dtype import Invalid, dtypes
+
+def move_where_load(gate, l, a, w):
+  return l.replace(src=(l.src[0], l.vconst_like(0) if a.is_invalid else
+    a.src[0] if a.op is Ops.CAST and a.src[0].dtype == l.dtype else a.cast(l.dtype), l.src[2])).cast(w.dtype)
+
+pm_move_gates_from_index = PatternMatcher([
+  # for image idx (must be first)
+  (UPat.var("buf").index(UPat.var("gate").where(UPat.var("idx_y"), UPat(arg=Invalid)),
+                         UPat.var("gate").where(UPat.var("idx_x"), UPat(arg=Invalid))).load(name="l"),
+   lambda buf,gate,idx_y,idx_x,l: buf.index(idx_y, idx_x, dtype=dtypes.float).load(l.vconst_like(0), gate)),
+  (UPat.var("buf").index(UPat.var("gate").where(UPat.var("idx_y"), UPat(arg=Invalid)),
+                         UPat.var("gate").where(UPat.var("idx_x"), UPat(arg=Invalid))).store(UPat.var("data")),
+   lambda buf,gate,idx_y,idx_x,data: buf.index(idx_y, idx_x, dtype=dtypes.float).store(data, gate)),
+
+  # here we create the alt value for load to be 0s and remove the where Invalid
+  (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat(), UPat.var("gate").where(UPat.var("idx"), UPat(arg=Invalid)),), name="mop", allow_any_len=True) \
+   .load(name="l"), lambda mop,gate,idx,l: mop.replace(src=(mop.src[0],idx)+mop.src[2:]).load(l.vconst_like(0), gate)),
+  (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat(), UPat.var("gate").where(UPat.var("idx"), UPat(arg=Invalid)),), name="mop", allow_any_len=True) \
+   .store(UPat.var("data")), lambda mop,gate,idx,data: mop.replace(src=(mop.src[0],idx)+mop.src[2:]).store(data, gate)),
+
+  # Where after gated load becomes alt value
+  (UPat.var("gate").where(UPat().load(UPat(), UPat.var("gate", dtype=dtypes.bool), name="l").or_casted(), UPat.var("a")).named("w"),
+   move_where_load),
+  (UPat.var("gate").where(UPat.var("a"), UPat().load(UPat(), ~UPat.var("gate", dtype=dtypes.bool), name="l").or_casted()).named("w"),
+   move_where_load),
+])

artifacts/package_runtime/tinygrad/codegen/late/linearizer.py

@@ -0,0 +1,95 @@
+import heapq
+from typing import Any
+from collections import defaultdict
+from tinygrad.uop.ops import PatternMatcher, UOp, Ops, UPat, multirange_str
+from tinygrad.dtype import AddrSpace, dtypes
+from tinygrad.helpers import prod, getenv, TUPLE_ORDER
+
+def linearize(sink:UOp) -> list[UOp]:
+  # this is a toposort with priority
+  lst = list(sink.toposort())
+  out_degree:defaultdict[UOp, int] = defaultdict(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: out_degree[s] += 1
+
+    # 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.PARAM: priority, extra = -20, u.arg.slot
+      case Ops.BUFFER: priority = -17 if u.addrspace == AddrSpace.LOCAL else -18
+      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, backedge = e.src[0], tuple(x for x in e.src[1:] if x.dtype in (dtypes.void, dtypes.bool))
+  for r in sorted(UOp.sink(*[x for x in e.src[1:] if x not in backedge]).ranges, key=lambda x: x.arg, reverse=True): ret = ret.end(r)
+  return ret.end(*backedge) if len(backedge) else ret
+
+pm_split_ends = PatternMatcher([
+  # split the ends
+  (UPat(Ops.END, name="e"), do_split_ends),
+])

artifacts/package_runtime/tinygrad/codegen/late/regalloc.py

@@ -0,0 +1,136 @@
+import itertools
+from tinygrad.helpers import dedup
+from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat
+from tinygrad.renderer.isa import ISARenderer, Register, greg
+from tinygrad.dtype import dtypes
+
+PSEUDO_OPS = {Ops.CONST, Ops.NOOP, Ops.AFTER, Ops.BARRIER, Ops.GROUP, Ops.STACK}
+
+class LinearScanRegallocContext:
+  # returns the uop that defines the virtual register
+  def vdef(self, v:Register) -> UOp: return self.uops[self.live_range[v][0]]
+  def __init__(self, uops:list[UOp], ren:ISARenderer):
+    self.uops = uops
+    self.ren = ren
+    self.idx = itertools.count()
+    # the label associated with each loop NOTE: this is only used post regalloc and should be removed
+    self.loop_label: dict[UOp, str] = {}
+
+    # compute live ranges
+    self.live_range: dict[Register, list[int]] = {}
+    lr = self.live_range
+    ranges: list[Register] = []
+    for i,u in enumerate(reversed(uops)):
+      if u.op in PSEUDO_OPS: continue
+      defs = u.tag if isinstance(u.tag, tuple) else ()
+      for v in defs + tuple(greg(s) for s in dedup(u.src)):
+        if isinstance(v, Register): lr.setdefault(v, []).insert(0, len(uops) - 1 - i)
+      for v in defs:
+        if v in lr and (n:=max((lr[rng][-1] for rng in ranges if lr[rng][0] <= lr[v][-1] < lr[rng][-1]), default=None)): lr[v].append(n)
+      if u.op is Ops.RANGE: ranges.append(greg(u))
+
+    # allocate registers
+    self.stack_size: int = 0
+    self.locals: dict[UOp, UOp] = {}
+    self.spills: dict[Register, UOp] = {} # mapping from virtual to stack slot
+    self.reals: dict[int, dict[Register, Register]] = {} # mapping from virtual to real at each program point
+    self.insert_before: dict[int, list[tuple[Register, Register]]] = {} # fills to be inserted at each program point
+    live: dict[Register, Register] = {} # mapping from virtual to real that's currently assigned to it
+    live_ins: list[dict[Register, Register]] = [] # mapping from virtual to real at loop entry
+
+    def alloc(cons:tuple[Register, ...], i:int) -> Register:
+      live_inv = {v:k for k,v in live.items()}
+      # allocate the best register. Registers not in live or not used again are free and have priority,
+      # otherwise pick the one with the furthest next use. Regs that appear first in cons have priority in case of a tie
+      reg,vreg = max(((r,live_inv.get(r)) for r in cons),
+                    key=lambda rv: next((j-i for j in ([] if rv[1] is None else lr[rv[1]]) if j >= i), len(uops)))
+      return live.pop(vreg) if vreg is not None else reg
+
+    # assign register to spilled virtual and record load to be emitted before current uop, also assign it a stack slot
+    def fill(v:Register, i:int, cons:tuple[Register, ...]|None=None) -> Register:
+      if v not in self.spills:
+        # the value of a BUFFER is its 64bit address, XMM registers need 16 bytes
+        sz = 16 if v.cons[0].size == 16 else (8 if self.vdef(v).op is Ops.BUFFER else self.vdef(v).dtype.itemsize)
+        offset = self.stack_size + (sz - self.stack_size % sz) % sz
+        self.spills[v] = UOp.const(offset, dtypes.int32)
+        self.stack_size = offset + sz
+      r = alloc(cons if cons is not None else v.cons, i)
+      self.insert_before.setdefault(i, []).append((v, r))
+      return r
+
+    for i,u in enumerate(uops):
+      if u.op in PSEUDO_OPS: continue
+      # allocate uses
+      for s in u.src:
+        # HACK: cause of later hacks to lower range
+        if u.op is Ops.END: continue
+        if not isinstance(v:=greg(s), Register): continue
+        if v not in live: live[v] = fill(v, i)
+        self.reals.setdefault(i, {})[v] = live[v]
+
+      # allocate defs
+      if isinstance(u.tag, tuple):
+        for j,v in enumerate(u.tag):
+          # register should only be defined once
+          assert isinstance(v, Register) and lr[v][0] == i
+          cons = v.cons
+          # two address instructions (src is reused by def) can only coalesce reused src. reused src goes first to get priority in case of a tiebreak
+          if ren.is_two_address(u) and j == 0:
+            uses = tuple(live.get(greg(s)) for s in u.src)
+            cons = ((uses[0],) if uses[0] in cons else ()) + tuple(r for r in cons if r not in uses)
+          # HACK: cause the range is missing the comparison
+          live[v] = alloc(cons, i+1 if u.op is not Ops.RANGE else i)
+          self.reals.setdefault(i, {})[v] = live[v]
+
+      # allocate stack array
+      if u.op is Ops.BUFFER:
+        self.locals[u] = UOp.const(self.stack_size, dtypes.int32)
+        self.stack_size += u.max_numel() * u.dtype.itemsize
+
+      # loop prologue, avoid loading inside the loop
+      if u.op is Ops.RANGE:
+        # we move to registers vars used in the loop sorted by next use, vars not used in the loop will not be reloaded in the epilogue
+        used_in_loop = [v for v in live.keys() | self.spills.keys() if any(i <= l < lr[greg(u)][-1] for l in lr[v])]
+        sorted_uses = sorted(used_in_loop, key=lambda k: (next(l-i for l in lr[k] if l >= i), lr[k][0], k.name, k.index))
+        live_in: dict[Register, Register] = {}
+        for v in sorted_uses:
+          # if all the possible registers are already in live_in there's no space for this var
+          if set(v.cons).issubset(live_in.values()): continue
+          if v not in live: live[v] = fill(v, i)
+          live_in[v] = live[v]
+        live_ins.append(live_in)
+
+      # loop epilogue, reload registers that were live at loop entry
+      if u.op is Ops.END:
+        # TODO: if a uop is in a different reg in live out vs live in move between registers instead of loading
+        # TODO: don't reload if first use in loop is a load
+        for v,r in live_ins.pop().items():
+          if v not in live or live[v] != r: live[v] = fill(v, i, (r,))
+
+def regalloc_rewrite(ctx:LinearScanRegallocContext, x:UOp):
+  i = next(ctx.idx)
+  if x.op in PSEUDO_OPS: return None
+  nsrc = []
+  for j,s in enumerate(x.src):
+    # v here is the virtual defined by the original s as s is the rewritten version
+    if i in ctx.reals and (v:=greg(ctx.uops[i].src[j])) in ctx.spills: nsrc.append(ctx.ren.fill(ctx.spills[v], ctx.vdef(v), ctx.reals[i][v]))
+    else: nsrc.append(s)
+  ndefs = tuple(ctx.reals[i][v] for v in x.tag) if isinstance(x.tag, tuple) else x.tag
+  if x.op is Ops.BUFFER: nx = ctx.ren.isel_matcher.rewrite(ctx.ren.stack_pointer().index(ctx.locals[x], tag=ndefs))
+  else: nx = x.replace(src=tuple(nsrc), tag=ndefs)
+
+  before = [ctx.ren.fill(ctx.spills[v], ctx.vdef(v), r) for v,r in ctx.insert_before.get(i, [])]
+  after = [ctx.ren.spill(ctx.spills[v], nx) for v in x.tag if v in ctx.spills] if isinstance(x.tag, tuple) else []
+
+  # alloc/dealloc stack
+  if ctx.stack_size > 0:
+    sp = ctx.ren.stack_pointer()
+    offset = UOp.const(ctx.stack_size, sp.dtype)
+    if i == 0: before = [ctx.ren.isel_matcher.rewrite(UOp(Ops.SUB, src=(sp, offset), tag=sp.tag))] + before
+    elif i == len(ctx.uops) - 2: before += [ctx.ren.isel_matcher.rewrite(UOp(Ops.ADD, src=(sp, offset), tag=sp.tag))]
+
+  return nx, before + [nx] + after
+
+pm_regalloc_rewrite = PatternMatcher([
+  (UPat({Ops.INS, Ops.RANGE, Ops.END, Ops.BUFFER, Ops.PARAM, Ops.SPECIAL} | PSEUDO_OPS, name="x"), regalloc_rewrite),
+])