From 64c0a785938ae2689333c77fec2b36cc5ec82148 Mon Sep 17 00:00:00 2001 From: nella Date: Sun, 6 Sep 2026 17:27:26 +0200 Subject: [PATCH] Piece alloc. --- src/backend/codegen/linear_scan_reg_alloc.cpp | 91 +++++ src/backend/codegen/linear_scan_reg_alloc.h | 9 + src/backend/codegen/reg_alloc_base.cpp | 311 +++++++++--------- src/backend/codegen/reg_alloc_base.h | 28 ++ 4 files changed, 282 insertions(+), 157 deletions(-) diff --git a/src/backend/codegen/linear_scan_reg_alloc.cpp b/src/backend/codegen/linear_scan_reg_alloc.cpp index 36193cf..b604c34 100644 --- a/src/backend/codegen/linear_scan_reg_alloc.cpp +++ b/src/backend/codegen/linear_scan_reg_alloc.cpp @@ -329,6 +329,97 @@ namespace rat { } } + void LinearScanRegAllocPass::assignPieces() { + if(!anySpilled()) + return; + List busy(fixedAt.begin(), fixedAt.end()); + for(U32 v = 1; v < fn->nextVReg; ++v) { + const Interval& iv = intervals[v]; + if(!iv.live() || iv.spilled || iv.assigned == kNoReg) + continue; + for(const Seg& sg : iv.segs) + for(I32 p = sg.start; p <= sg.end; ++p) + busy[(U32)p] |= (U64)1 << iv.assigned; + } + // every operand of a spilled vreg, sorted by vreg then program order + List touches; + for(U32 b = 0; b < fn->blocks.size(); ++b) + for(U32 i = 0; i < fn->blocks[b].insts.size(); ++i) { + const MachineInstr& in = fn->blocks[b].insts[i]; + if(in.isCall) + continue; // reads its slot directly + I32 pt = (I32)blkPts[b][i]; + for(const MachineOperand& o : in.uses) + if(o.isVReg() && intervals[o.vreg].spilled) + touches.emplace_back(o.vreg, pt); + for(const MachineOperand& o : in.defs) + if(o.isVReg() && intervals[o.vreg].spilled) + touches.emplace_back(o.vreg, pt); + } + std::sort(touches.begin(), touches.end()); + for(U32 i = 0; i < touches.size();) { + U32 j = i; + while(j < touches.size() && touches[j].first == touches[i].first) + ++j; + cutPieces(&touches[i], j - i, busy); + i = j; + } + } + + void LinearScanRegAllocPass::cutPieces(const Touch* pts, U32 count, List& busy) { + VReg v = pts[0].first; + U32 cls = intervals[v].cls; + U32 i = 0; + while(i < count) { + I32 start = pts[i].second; + I32 end = start; + U32 block = order[(U32)start].block; + U32 j = i + 1; + while(j < count && order[(U32)pts[j].second].block == block) { + auto call = std::upper_bound(callPts.begin(), callPts.end(), end); + if(call != callPts.end() && *call < pts[j].second) + break; + end = pts[j].second; + ++j; + } + i = j; + if(end == start) + continue; + PhysReg reg = pickPieceReg(cls, start, end, busy); + if(reg != kNoReg) + pieces.push_back({v, start, end, reg}); + } + } + + PhysReg LinearScanRegAllocPass::pickPieceReg(U32 cls, I32 start, I32 end, List& busy) { + const RegClass& rc = regClass(cls); + U64 used = 0; + for(I32 p = start; p <= end; ++p) + used |= busy[(U32)p]; + for(PhysReg p : rc.allocatable) { + if((used >> p) & 1) + continue; + if(isCalleeSaved(rc, p) && !usedCallee.count(p)) + continue; + for(I32 q = start; q <= end; ++q) + busy[(U32)q] |= (U64)1 << p; + return p; + } + return kNoReg; + } + + B32 LinearScanRegAllocPass::pieceAfter(const Touch& key, const Piece& pc) { + return key.first < pc.vreg || (key.first == pc.vreg && key.second < pc.start); + } + + RegAllocBase::Piece* LinearScanRegAllocPass::pieceAt(VReg v, I32 pt) { + auto after = std::upper_bound(pieces.begin(), pieces.end(), Touch{v, pt}, pieceAfter); + if(after == pieces.begin()) + return nullptr; + Piece& pc = *(after - 1); + return pc.vreg == v && pt <= pc.end ? &pc : nullptr; + } + VReg LinearScanRegAllocPass::webFind(VReg v) { while(webParent[v] != v) v = webParent[v] = webParent[webParent[v]]; diff --git a/src/backend/codegen/linear_scan_reg_alloc.h b/src/backend/codegen/linear_scan_reg_alloc.h index f5b7459..08a8ddf 100644 --- a/src/backend/codegen/linear_scan_reg_alloc.h +++ b/src/backend/codegen/linear_scan_reg_alloc.h @@ -30,6 +30,8 @@ namespace rat { static constexpr U32 kLoopUseWeight = 3; static constexpr U32 kMaxUseWeight = 100000; + using Touch = Pair; // a spilled vreg and a point reading or writing it + // closed [start, end] struct Seg { I32 start; @@ -71,6 +73,7 @@ namespace rat { } void resetState() override { + pieces.clear(); U32 nv = std::min((U32)intervals.size(), fn->nextVReg); for(U32 v = 0; v < nv; ++v) intervals[v].reset(); @@ -83,6 +86,11 @@ namespace rat { U64 forbidden(const Interval& iv) const; void assignRegs(); void assignSpillSlots(); + void assignPieces() override; + void cutPieces(const Touch* pts, U32 count, List& busy); + static B32 pieceAfter(const Touch& key, const Piece& pc); + PhysReg pickPieceReg(U32 cls, I32 start, I32 end, List& busy); + Piece* pieceAt(VReg v, I32 pt) override; void spillAt(Interval* cur, List& active); void buildSpillWebs(const List& spilled); VReg webFind(VReg v); @@ -94,6 +102,7 @@ namespace rat { } List intervals; + List pieces; // sorted by (vreg, start) List webParent; Map> webMembers; // union-find root -> members List> pinsByPoint; diff --git a/src/backend/codegen/reg_alloc_base.cpp b/src/backend/codegen/reg_alloc_base.cpp index bbbc873..acb60d6 100644 --- a/src/backend/codegen/reg_alloc_base.cpp +++ b/src/backend/codegen/reg_alloc_base.cpp @@ -201,9 +201,149 @@ namespace rat { } } + B32 RegAllocBase::dropsRematDef(const MachineInstr& in) { + if(in.defs.size() != 1 || !in.defs[0].isVReg()) + return false; + VReg v = in.defs[0].vreg; + return rematDef.count(v) && assignmentOf(v).spilled && !slotReadByCall.count(v); + } + + void RegAllocBase::emitReload(List& out, + PhysReg dst, + const MachineOperand& u, + const Assignment& a) { + auto rt = rematDef.find(u.vreg); + if(rt == rematDef.end()) { + out.push_back(hooks->makeReload(dst, a.spillSlot, a.cls, u.width)); + return; + } + MachineInstr m = rt->second; + m.defs[0] = MachineOperand::fixed(dst, u.width); + out.push_back(std::move(m)); + } + + void RegAllocBase::emitStore(List& out, I32 slot, PhysReg src, U32 cls, U32 width) { + out.push_back(hooks->makeSpill(slot, src, cls, width)); + memo = {true, slot, src, cls, width}; + } + + // the register that just stored its slot or target after a reload into it + PhysReg RegAllocBase::sourceOf(List& out, + const MachineOperand& u, + const Assignment& a, + PhysReg target) { + if(memo.on && memo.cls == a.cls && memo.slot == a.spillSlot && memo.width == u.width) + return memo.reg; + emitReload(out, target, u, a); + return target; + } + + B32 RegAllocBase::rewriteCopy(List& out, MachineInstr& in) { + MachineOperand& d = in.defs[0]; + MachineOperand& u = in.uses[0]; + Assignment da = d.isVReg() ? assignmentOf(d.vreg) : Assignment{}; + Assignment ua = u.isVReg() ? assignmentOf(u.vreg) : Assignment{}; + if(!da.spilled && !ua.spilled) + return false; + if(da.spilled && ua.spilled && da.spillSlot == ua.spillSlot) + return true; // slot self-copy + PhysReg dst = kNoReg; // where the value must land + if(d.isPhys()) + dst = d.phys; + else if(!da.spilled) + dst = da.reg; + PhysReg src; + if(ua.spilled) + src = sourceOf(out, u, ua, dst != kNoReg ? dst : scratchAt(ua.cls, 0)); + else + src = u.isPhys() ? u.phys : ua.reg; + if(da.spilled) { + emitStore(out, da.spillSlot, src, da.cls, d.width); + return true; + } + if(src != dst) { + d = MachineOperand::fixed(dst, d.width); + u = MachineOperand::fixed(src, u.width); + out.push_back(in); + } + memo.on = false; + return true; + } + + void RegAllocBase::rewriteInstr(List& out, MachineInstr& in, I32 pt) { + U32 useScratch[kMaxRegClasses] = {0}; + for(MachineOperand& u : in.uses) { + if(!u.isVReg()) + continue; + Assignment a = assignmentOf(u.vreg); + if(!a.spilled) { + u = MachineOperand::fixed(a.reg, u.width); + continue; + } + if(in.isCall) { + u = MachineOperand::frameSlot(a.spillSlot, u.width); + continue; + } + if(Piece* pc = pieceAt(u.vreg, pt)) { + if(!pc->loaded) + emitReload(out, pc->reg, u, a); + pc->loaded = true; + u = MachineOperand::fixed(pc->reg, u.width); + continue; + } + // the previous instruction just stored this slot from a register nothing has + // touched since, reuse + B32 tied = false; + for(const MachineOperand& d : in.defs) + if(d.isVReg() && d.vreg == u.vreg) + tied = true; + if(memo.on && !tied && useScratch[a.cls] == 0 && memo.cls == a.cls && + memo.slot == a.spillSlot && memo.width == u.width) { + u = MachineOperand::fixed(memo.reg, u.width); + ++useScratch[a.cls]; // reserve index 0 in case memo.reg is scratch 0 + continue; + } + PhysReg sc = scratchAt(a.cls, useScratch[a.cls]++); + emitReload(out, sc, u, a); + u = MachineOperand::fixed(sc, u.width); + } + + U32 defScratch[kMaxRegClasses] = {0}; + List stores; + for(MachineOperand& d : in.defs) { + if(!d.isVReg()) + continue; + Assignment a = assignmentOf(d.vreg); + if(!a.spilled) { + d = MachineOperand::fixed(a.reg, d.width); + continue; + } + PhysReg sc; + if(Piece* pc = pieceAt(d.vreg, pt)) { + sc = pc->reg; + pc->loaded = true; + } else { + sc = scratchAt(a.cls, defScratch[a.cls]++); + } + stores.push_back(hooks->makeSpill(a.spillSlot, sc, a.cls, d.width)); + d = MachineOperand::fixed(sc, d.width); + } + + memo.on = false; + out.push_back(in); + for(MachineInstr& s : stores) + out.push_back(std::move(s)); + if(stores.empty() || in.isCall) + return; + // uses[0] = frame slot, uses[1] = source register + const MachineInstr& last = out.back(); + if(last.uses.size() == 2 && last.uses[0].kind == MachineOperand::Kind::FrameSlot && + last.uses[1].kind == MachineOperand::Kind::Phys) + memo = {true, last.uses[0].slot, last.uses[1].phys, last.regClass, last.uses[1].width}; + } + void RegAllocBase::rewrite() { - // spill slots read by a call (stack-passed args): defs must be kept - Set slotReadByCall; + slotReadByCall.clear(); for(U32 b = 0; b < fn->blocks.size(); ++b) for(const MachineInstr& in : fn->blocks[b].insts) if(in.isCall) @@ -213,162 +353,18 @@ namespace rat { for(U32 b = 0; b < fn->blocks.size(); ++b) { List out; - B32 memo = false; - I32 memoSlot = 0; - PhysReg memoReg = kNoReg; - U32 memoCls = 0; - U32 memoWidth = 0; - - auto reloadInto = [&](PhysReg dst, const MachineOperand& u, const Assignment& a) { - if(auto rt = rematDef.find(u.vreg); rt != rematDef.end()) { - MachineInstr m = rt->second; - m.defs[0] = MachineOperand::fixed(dst, u.width); - out.push_back(std::move(m)); - } else { - out.push_back(hooks->makeReload(dst, a.spillSlot, a.cls, u.width)); - } - }; - - for(MachineInstr& in : fn->blocks[b].insts) { - // spilled remat def: drop it, every use re-materializes - if(in.defs.size() == 1 && in.defs[0].isVReg()) { - VReg dv = in.defs[0].vreg; - if(rematDef.find(dv) != rematDef.end() && assignmentOf(dv).spilled && - !slotReadByCall.count(dv)) { - memo = false; - continue; - } - } - - // spill from or reload into it directly instead of bouncing through a scratch register - if(hooks->isCopy && hooks->isCopy(in) && in.defs.size() == 1 && in.uses.size() == 1) { - const MachineOperand& d = in.defs[0]; - const MachineOperand& u = in.uses[0]; - if(d.isVReg() && u.isPhys()) { - Assignment da = assignmentOf(d.vreg); - if(da.spilled) { - out.push_back(hooks->makeSpill(da.spillSlot, u.phys, da.cls, d.width)); - memo = true; - memoSlot = da.spillSlot; - memoReg = u.phys; - memoCls = da.cls; - memoWidth = d.width; - continue; - } - } else if(d.isPhys() && u.isVReg()) { - Assignment ua = assignmentOf(u.vreg); - if(ua.spilled) { - reloadInto(d.phys, u, ua); - memo = false; - continue; - } - } - } - - // spilled-copy peepholes - if(hooks->isCopy && hooks->isCopy(in) && in.defs.size() == 1 && in.uses.size() == 1 && - in.defs[0].isVReg() && in.uses[0].isVReg()) { - Assignment da = assignmentOf(in.defs[0].vreg); - Assignment ua = assignmentOf(in.uses[0].vreg); - if(da.cls == ua.cls && (da.spilled || ua.spilled)) { - if(da.spilled && ua.spilled) { - if(da.spillSlot == ua.spillSlot) - continue; // slot self-copy - PhysReg sc; - if(memo && memoCls == ua.cls && memoSlot == ua.spillSlot && - memoWidth == in.uses[0].width) { - sc = memoReg; // forward the just-stored value - } else { - sc = scratchAt(ua.cls, 0); - reloadInto(sc, in.uses[0], ua); - } - out.push_back(hooks->makeSpill(da.spillSlot, sc, da.cls, in.defs[0].width)); - memo = true; - memoSlot = da.spillSlot; - memoReg = sc; - memoCls = da.cls; - memoWidth = in.defs[0].width; - continue; - } - if(da.spilled) { - // register -> slot - out.push_back(hooks->makeSpill(da.spillSlot, ua.reg, da.cls, in.defs[0].width)); - memo = true; - memoSlot = da.spillSlot; - memoReg = ua.reg; - memoCls = da.cls; - memoWidth = in.defs[0].width; - continue; - } - // slot -> register - reloadInto(da.reg, in.uses[0], ua); - memo = false; - continue; - } - } - - U32 useScratch[kMaxRegClasses] = {0}; - for(MachineOperand& u : in.uses) { - if(!u.isVReg()) - continue; - Assignment a = assignmentOf(u.vreg); - if(a.spilled) { - if(in.isCall) { - u = MachineOperand::frameSlot(a.spillSlot, u.width); - continue; - } - // the previous instruction just stored this slot from a register nothing has - // touched since, reuse - B32 tied = false; - for(const MachineOperand& d : in.defs) - if(d.isVReg() && d.vreg == u.vreg) - tied = true; - if(memo && !tied && useScratch[a.cls] == 0 && memoCls == a.cls && - memoSlot == a.spillSlot && memoWidth == u.width) { - u = MachineOperand::fixed(memoReg, u.width); - ++useScratch[a.cls]; // reserve index 0 in case memoReg is scratch 0 - continue; - } - PhysReg sc = scratchAt(a.cls, useScratch[a.cls]++); - reloadInto(sc, u, a); - u = MachineOperand::fixed(sc, u.width); - } else { - u = MachineOperand::fixed(a.reg, u.width); - } - } - - U32 defScratch[kMaxRegClasses] = {0}; - List spills; - for(MachineOperand& d : in.defs) { - if(!d.isVReg()) - continue; - Assignment a = assignmentOf(d.vreg); - if(a.spilled) { - PhysReg sc = scratchAt(a.cls, defScratch[a.cls]++); - spills.push_back(hooks->makeSpill(a.spillSlot, sc, a.cls, d.width)); - d = MachineOperand::fixed(sc, d.width); - } else { - d = MachineOperand::fixed(a.reg, d.width); - } - } - - memo = false; - - out.push_back(in); - for(MachineInstr& s : spills) - out.push_back(std::move(s)); - if(!spills.empty() && !in.isCall) { - const MachineInstr& last = out.back(); - // makeSpill shape: uses[0] = frame slot, uses[1] = source register - if(last.uses.size() == 2 && last.uses[0].kind == MachineOperand::Kind::FrameSlot && - last.uses[1].kind == MachineOperand::Kind::Phys) { - memo = true; - memoSlot = last.uses[0].slot; - memoReg = last.uses[1].phys; - memoCls = last.regClass; - memoWidth = last.uses[1].width; - } + memo.on = false; + for(U32 i = 0; i < fn->blocks[b].insts.size(); ++i) { + MachineInstr& in = fn->blocks[b].insts[i]; + I32 pt = (I32)blkPts[b][i]; + if(dropsRematDef(in)) { // every use remats instead + memo.on = false; + continue; } + B32 copy = hooks->isCopy && hooks->isCopy(in) && in.defs.size() == 1 && in.uses.size() == 1; + if(copy && rewriteCopy(out, in)) + continue; + rewriteInstr(out, in, pt); } fn->blocks[b].insts = std::move(out); } @@ -438,6 +434,7 @@ namespace rat { } else { solve(); } + assignPieces(); rewrite(); if(usedCalleeSaved) { diff --git a/src/backend/codegen/reg_alloc_base.h b/src/backend/codegen/reg_alloc_base.h index f7a0a51..e6b9ca9 100644 --- a/src/backend/codegen/reg_alloc_base.h +++ b/src/backend/codegen/reg_alloc_base.h @@ -67,6 +67,15 @@ namespace rat { virtual void solve() = 0; // compute assignments (number() already ran) virtual Assignment assignmentOf(VReg v) = 0; // result lookup used by rewrite() virtual B32 anySpilled() const = 0; // did solve() spill anything? + virtual void assignPieces() = 0; + struct Piece { + VReg vreg; + I32 start; + I32 end; + PhysReg reg; + B32 loaded = false; // reloaded once at its first use + }; + virtual Piece* pieceAt(VReg v, I32 pt) = 0; // the piece holding v at pt, if any B32 allocate(MachineFunc& fn, const RegisterInfo& ri, @@ -91,7 +100,24 @@ namespace rat { const RegClass& regClass(U32 cls) const; static B32 isCalleeSaved(const RegClass& rc, PhysReg p); static B32 isAllocatable(const RegClass& rc, PhysReg p); + // a slot just written from a register nothing has touched since + struct Memo { + B32 on = false; + I32 slot = 0; + PhysReg reg = kNoReg; + U32 cls = 0; + U32 width = 0; + }; + void rewrite(); + B32 dropsRematDef(const MachineInstr& in); + void + emitReload(List& out, PhysReg dst, const MachineOperand& u, const Assignment& a); + void emitStore(List& out, I32 slot, PhysReg src, U32 cls, U32 width); + PhysReg + sourceOf(List& out, const MachineOperand& u, const Assignment& a, PhysReg target); + B32 rewriteCopy(List& out, MachineInstr& in); + void rewriteInstr(List& out, MachineInstr& in, I32 pt); PhysReg scratchAt(U32 cls, U32 idx); protected: MachineFunc* fn = nullptr; @@ -107,6 +133,8 @@ namespace rat { Map> copyPinAt; // pin exemptions Map rematDef; // single pure def per remat vreg B32 ok = true; + Memo memo; + Set slotReadByCall; List liveUseScratch; List liveDefScratch; private: