mirror of
https://jihulab.com/mr-one/openpilot.git
synced 2026-10-28 05:23:46 +08:00
1okko
129 lines
6.7 KiB
Python
129 lines
6.7 KiB
Python
import itertools
|
|
from tinygrad.helpers import dedup
|
|
from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat
|
|
from tinygrad.renderer.isa import ISARenderer, Register, rdef, LinearContext
|
|
from typing import Any
|
|
from dataclasses import replace
|
|
|
|
PSEUDO_OPS = {Ops.CONST, Ops.CAST, Ops.BITCAST, Ops.NOOP, Ops.AFTER, Ops.BARRIER, 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, ctx:LinearContext, uops:list[UOp], ren:ISARenderer):
|
|
self.uops = uops
|
|
self.ren = ren
|
|
self.idx = itertools.count()
|
|
|
|
# compute live ranges
|
|
self.live_range: dict[Register, list[int]] = {}
|
|
lr = self.live_range
|
|
loops: dict[int, int] = {} # the interval of each loop, from its RANGE to the last uop that reads that RANGE
|
|
for idx,u in reversed(list(enumerate(uops))):
|
|
if u.op in PSEUDO_OPS: continue
|
|
defs = u.tag if isinstance(u.tag, tuple) else ()
|
|
for v in defs + tuple(rdef(s) for s in dedup(u.src)):
|
|
if isinstance(v, Register): lr.setdefault(v, []).insert(0, idx)
|
|
for v in defs:
|
|
if v in lr and (n:=max((e for s,e in loops.items() if s <= lr[v][-1] < e), default=None)): lr[v].append(n)
|
|
if u.op is Ops.RANGE: loops[idx] = max(j for j,x in enumerate(uops) if u in x.src)
|
|
|
|
# allocate registers
|
|
self.spills: dict[Register, Any] = {} # mapping from virtual to arbitrary spill 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:
|
|
self.spills[v] = ctx.assign_spill_slot(v, self.vdef(v))
|
|
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 in (Ops.END, Ops.BACKEDGE): continue
|
|
if not isinstance(v:=rdef(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):
|
|
if v not in live:
|
|
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(rdef(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]
|
|
|
|
# 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 < loops[i] 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 in (Ops.END, Ops.BACKEDGE):
|
|
# 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
|
|
live_in = live_ins.pop()
|
|
for v,r in live_in.items():
|
|
if v not in live or live[v] != r: live[v] = fill(v, i, (r,))
|
|
# a spilled var that only got a register inside the loop must be forced to reload from spill slot
|
|
for v in list(live.keys()):
|
|
if v in self.spills and v not in live_in: live.pop(v)
|
|
|
|
def regalloc_rewrite(ctx:LinearScanRegallocContext, x:UOp):
|
|
i = next(ctx.idx)
|
|
if x.op in PSEUDO_OPS: return None
|
|
nsrc = []
|
|
def retag(x:UOp, r:Register) -> UOp:
|
|
return x.replace(src=(retag(x.src[0],r),*x.src[1:])) if x.op in {Ops.AFTER, Ops.BITCAST} else x.replace(tag=(r,))
|
|
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:=rdef(ctx.uops[i].src[j])) in ctx.reals.get(i, {}):
|
|
reg = replace(ctx.reals[i][v], size=v.size)
|
|
if v in ctx.spills:
|
|
fill = ctx.ren.fill(ctx.spills[v], ctx.vdef(v), reg)
|
|
# an AFTER-wrapped value (e.g. a RANGE bound) keeps its ordering deps around the fill
|
|
nsrc.append(s.replace(src=(fill,)+s.src[1:]) if s.op is Ops.AFTER else fill)
|
|
else: nsrc.append(retag(s, reg))
|
|
else: nsrc.append(s)
|
|
ndefs = tuple(replace(ctx.reals[i][v], size=v.size) for v in x.tag) if isinstance(x.tag, tuple) else x.tag
|
|
nx = x.replace(src=tuple(nsrc), tag=ndefs)
|
|
|
|
before = [ctx.ren.fill(ctx.spills[v], ctx.vdef(v), replace(r, size=v.size)) 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 []
|
|
|
|
return nx, before + [nx] + after
|
|
|
|
pm_regalloc_rewrite = PatternMatcher([
|
|
(UPat({Ops.INS, Ops.RANGE, Ops.END, Ops.BACKEDGE, Ops.BUFFER, Ops.PARAM, Ops.SPECIAL} | PSEUDO_OPS, name="x"), regalloc_rewrite),
|
|
])
|