IQ.Pilot Prebuilt Release @ 27f668a

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),
+])