Files
wazero/internal/engine/wazevo/backend/isa/arm64/machine_regalloc.go
Takeshi Yoneda 5d1d15fb9f wazevo: cleans up debug options (#1867)
Signed-off-by: Takeshi Yoneda <t.y.mathetake@gmail.com>
2023-12-13 07:19:21 -08:00

434 lines
13 KiB
Go

package arm64
// This file implements the interfaces required for register allocations. See regalloc/api.go.
import (
"github.com/tetratelabs/wazero/internal/engine/wazevo/backend/regalloc"
"github.com/tetratelabs/wazero/internal/engine/wazevo/ssa"
)
type (
// regAllocFunctionImpl implements regalloc.Function.
regAllocFunctionImpl struct {
m *machine
// iter is the iterator for reversePostOrderBlocks
iter int
reversePostOrderBlocks []regAllocBlockImpl
// labelToRegAllocBlockIndex maps label to the index of reversePostOrderBlocks.
labelToRegAllocBlockIndex map[label]int
loopNestingForestRoots []ssa.BasicBlock
}
// regAllocBlockImpl implements regalloc.Block.
regAllocBlockImpl struct {
// f is the function this instruction belongs to. Used to reuse the regAllocFunctionImpl.predsSlice slice for Defs() and Uses().
f *regAllocFunctionImpl
sb ssa.BasicBlock
l label
pos *labelPosition
loopNestingForestChildren []ssa.BasicBlock
cur *instruction
id int
cachedLastInstr regalloc.Instr
}
)
func (f *regAllocFunctionImpl) addBlock(sb ssa.BasicBlock, l label, pos *labelPosition) {
i := len(f.reversePostOrderBlocks)
f.reversePostOrderBlocks = append(f.reversePostOrderBlocks, regAllocBlockImpl{
f: f,
sb: sb,
l: l,
pos: pos,
id: int(sb.ID()),
})
f.labelToRegAllocBlockIndex[l] = i
}
func (f *regAllocFunctionImpl) reset() {
f.reversePostOrderBlocks = f.reversePostOrderBlocks[:0]
f.iter = 0
}
var (
_ regalloc.Function = (*regAllocFunctionImpl)(nil)
_ regalloc.Block = (*regAllocBlockImpl)(nil)
)
// PostOrderBlockIteratorBegin implements regalloc.Function PostOrderBlockIteratorBegin.
func (f *regAllocFunctionImpl) PostOrderBlockIteratorBegin() regalloc.Block {
f.iter = len(f.reversePostOrderBlocks) - 1
return f.PostOrderBlockIteratorNext()
}
// PostOrderBlockIteratorNext implements regalloc.Function PostOrderBlockIteratorNext.
func (f *regAllocFunctionImpl) PostOrderBlockIteratorNext() regalloc.Block {
if f.iter < 0 {
return nil
}
b := &f.reversePostOrderBlocks[f.iter]
f.iter--
return b
}
// ReversePostOrderBlockIteratorBegin implements regalloc.Function ReversePostOrderBlockIteratorBegin.
func (f *regAllocFunctionImpl) ReversePostOrderBlockIteratorBegin() regalloc.Block {
f.iter = 0
return f.ReversePostOrderBlockIteratorNext()
}
// ReversePostOrderBlockIteratorNext implements regalloc.Function ReversePostOrderBlockIteratorNext.
func (f *regAllocFunctionImpl) ReversePostOrderBlockIteratorNext() regalloc.Block {
if f.iter >= len(f.reversePostOrderBlocks) {
return nil
}
b := &f.reversePostOrderBlocks[f.iter]
f.iter++
return b
}
// ClobberedRegisters implements regalloc.Function ClobberedRegisters.
func (f *regAllocFunctionImpl) ClobberedRegisters(regs []regalloc.VReg) {
m := f.m
m.clobberedRegs = append(m.clobberedRegs[:0], regs...)
}
// StoreRegisterBefore implements regalloc.Function StoreRegisterBefore.
func (f *regAllocFunctionImpl) StoreRegisterBefore(v regalloc.VReg, instr regalloc.Instr) {
m := f.m
m.insertStoreRegisterAt(v, instr.(*instruction), false)
}
// SwapAtEndOfBlock implements regalloc.Function SwapAtEndOfBlock.
func (f *regAllocFunctionImpl) SwapAtEndOfBlock(x1, x2, tmp regalloc.VReg, block regalloc.Block) {
blk := block.(*regAllocBlockImpl)
cur := blk.LastInstr().(*instruction)
cur = cur.prev
f.m.swap(cur, x1, x2, tmp)
}
func (m *machine) swap(cur *instruction, x1, x2, tmp regalloc.VReg) {
prevNext := cur.next
var mov1, mov2, mov3 *instruction
if x1.RegType() == regalloc.RegTypeInt {
if !tmp.Valid() {
tmp = tmpRegVReg
}
mov1 = m.allocateInstr().asMove64(tmp, x1)
mov2 = m.allocateInstr().asMove64(x1, x2)
mov3 = m.allocateInstr().asMove64(x2, tmp)
cur = linkInstr(cur, mov1)
cur = linkInstr(cur, mov2)
cur = linkInstr(cur, mov3)
linkInstr(cur, prevNext)
} else {
if !tmp.Valid() {
r2 := x2.RealReg()
// Temporarily spill x1 to stack.
cur = m.insertStoreRegisterAt(x1, cur, true).prev
// Then move x2 to x1.
cur = linkInstr(cur, m.allocateInstr().asFpuMov128(x1, x2))
linkInstr(cur, prevNext)
// Then reload the original value on x1 from stack to r2.
m.insertReloadRegisterAt(x1.SetRealReg(r2), cur, true)
} else {
mov1 = m.allocateInstr().asFpuMov128(tmp, x1)
mov2 = m.allocateInstr().asFpuMov128(x1, x2)
mov3 = m.allocateInstr().asFpuMov128(x2, tmp)
cur = linkInstr(cur, mov1)
cur = linkInstr(cur, mov2)
cur = linkInstr(cur, mov3)
linkInstr(cur, prevNext)
}
}
}
// InsertMoveBefore implements regalloc.Function InsertMoveBefore.
func (f *regAllocFunctionImpl) InsertMoveBefore(dst, src regalloc.VReg, instr regalloc.Instr) {
m := f.m
typ := src.RegType()
if typ != dst.RegType() {
panic("BUG: src and dst must have the same type")
}
mov := m.allocateInstr()
if typ == regalloc.RegTypeInt {
mov.asMove64(dst, src)
} else {
mov.asFpuMov128(dst, src)
}
cur := instr.(*instruction).prev
prevNext := cur.next
cur = linkInstr(cur, mov)
linkInstr(cur, prevNext)
}
// StoreRegisterAfter implements regalloc.Function StoreRegisterAfter.
func (f *regAllocFunctionImpl) StoreRegisterAfter(v regalloc.VReg, instr regalloc.Instr) {
m := f.m
m.insertStoreRegisterAt(v, instr.(*instruction), true)
}
// ReloadRegisterBefore implements regalloc.Function ReloadRegisterBefore.
func (f *regAllocFunctionImpl) ReloadRegisterBefore(v regalloc.VReg, instr regalloc.Instr) {
m := f.m
m.insertReloadRegisterAt(v, instr.(*instruction), false)
}
// ReloadRegisterAfter implements regalloc.Function ReloadRegisterAfter.
func (f *regAllocFunctionImpl) ReloadRegisterAfter(v regalloc.VReg, instr regalloc.Instr) {
m := f.m
m.insertReloadRegisterAt(v, instr.(*instruction), true)
}
// Done implements regalloc.Function Done.
func (f *regAllocFunctionImpl) Done() {
m := f.m
// Now that we know the final spill slot size, we must align spillSlotSize to 16 bytes.
m.spillSlotSize = (m.spillSlotSize + 15) &^ 15
}
// ID implements regalloc.Block ID.
func (r *regAllocBlockImpl) ID() int {
return r.id
}
// Preds implements regalloc.Block Preds.
func (r *regAllocBlockImpl) Preds() int {
return r.sb.Preds()
}
// Pred implements regalloc.Block Pred.
func (r *regAllocBlockImpl) Pred(i int) regalloc.Block {
sb := r.sb
pred := sb.Pred(i)
l := r.f.m.ssaBlockIDToLabels[pred.ID()]
index := r.f.labelToRegAllocBlockIndex[l]
return &r.f.reversePostOrderBlocks[index]
}
// Succs implements regalloc.Block Succs.
func (r *regAllocBlockImpl) Succs() int {
return r.sb.Succs()
}
// Succ implements regalloc.Block Succ.
func (r *regAllocBlockImpl) Succ(i int) regalloc.Block {
sb := r.sb
succ := sb.Succ(i)
if succ.ReturnBlock() {
return nil
}
l := r.f.m.ssaBlockIDToLabels[succ.ID()]
index := r.f.labelToRegAllocBlockIndex[l]
return &r.f.reversePostOrderBlocks[index]
}
// LoopHeader implements regalloc.Block LoopHeader.
func (r *regAllocBlockImpl) LoopHeader() bool {
return r.sb.LoopHeader()
}
// LoopNestingForestRoots implements regalloc.Function LoopNestingForestRoots.
func (f *regAllocFunctionImpl) LoopNestingForestRoots() int {
f.loopNestingForestRoots = f.m.compiler.SSABuilder().LoopNestingForestRoots()
return len(f.loopNestingForestRoots)
}
// LoopNestingForestRoot implements regalloc.Function LoopNestingForestRoot.
func (f *regAllocFunctionImpl) LoopNestingForestRoot(i int) regalloc.Block {
blk := f.loopNestingForestRoots[i]
l := f.m.ssaBlockIDToLabels[blk.ID()]
index := f.labelToRegAllocBlockIndex[l]
return &f.reversePostOrderBlocks[index]
}
// LoopNestingForestChildren implements regalloc.Block LoopNestingForestChildren.
func (r *regAllocBlockImpl) LoopNestingForestChildren() int {
r.loopNestingForestChildren = r.sb.LoopNestingForestChildren()
return len(r.loopNestingForestChildren)
}
// LoopNestingForestChild implements regalloc.Block LoopNestingForestChild.
func (r *regAllocBlockImpl) LoopNestingForestChild(i int) regalloc.Block {
blk := r.loopNestingForestChildren[i]
l := r.f.m.ssaBlockIDToLabels[blk.ID()]
index := r.f.labelToRegAllocBlockIndex[l]
return &r.f.reversePostOrderBlocks[index]
}
// InstrIteratorBegin implements regalloc.Block InstrIteratorBegin.
func (r *regAllocBlockImpl) InstrIteratorBegin() regalloc.Instr {
r.cur = r.pos.begin
return r.cur
}
// InstrIteratorNext implements regalloc.Block InstrIteratorNext.
func (r *regAllocBlockImpl) InstrIteratorNext() regalloc.Instr {
for {
if r.cur == r.pos.end {
return nil
}
instr := r.cur.next
r.cur = instr
if instr == nil {
return nil
} else if instr.addedBeforeRegAlloc {
// Only concerned about the instruction added before regalloc.
return instr
}
}
}
// InstrRevIteratorBegin implements regalloc.Block InstrRevIteratorBegin.
func (r *regAllocBlockImpl) InstrRevIteratorBegin() regalloc.Instr {
r.cur = r.pos.end
return r.cur
}
// InstrRevIteratorNext implements regalloc.Block InstrRevIteratorNext.
func (r *regAllocBlockImpl) InstrRevIteratorNext() regalloc.Instr {
for {
if r.cur == r.pos.begin {
return nil
}
instr := r.cur.prev
r.cur = instr
if instr == nil {
return nil
} else if instr.addedBeforeRegAlloc {
// Only concerned about the instruction added before regalloc.
return instr
}
}
}
// BlockParams implements regalloc.Block BlockParams.
func (r *regAllocBlockImpl) BlockParams(regs *[]regalloc.VReg) []regalloc.VReg {
c := r.f.m.compiler
*regs = (*regs)[:0]
for i := 0; i < r.sb.Params(); i++ {
v := c.VRegOf(r.sb.Param(i))
*regs = append(*regs, v)
}
return *regs
}
// Entry implements regalloc.Block Entry.
func (r *regAllocBlockImpl) Entry() bool { return r.sb.EntryBlock() }
// RegisterInfo implements backend.Machine.
func (m *machine) RegisterInfo() *regalloc.RegisterInfo {
return regInfo
}
// Function implements backend.Machine Function.
func (m *machine) Function() regalloc.Function {
m.regAllocStarted = true
return &m.regAllocFn
}
func (m *machine) insertStoreRegisterAt(v regalloc.VReg, instr *instruction, after bool) *instruction {
if !v.IsRealReg() {
panic("BUG: VReg must be backed by real reg to be stored")
}
typ := m.compiler.TypeOf(v)
var prevNext, cur *instruction
if after {
cur, prevNext = instr, instr.next
} else {
cur, prevNext = instr.prev, instr
}
offsetFromSP := m.getVRegSpillSlotOffsetFromSP(v.ID(), typ.Size())
var amode addressMode
cur, amode = m.resolveAddressModeForOffsetAndInsert(cur, offsetFromSP, typ.Bits(), spVReg, true)
store := m.allocateInstr()
store.asStore(operandNR(v), amode, typ.Bits())
cur = linkInstr(cur, store)
return linkInstr(cur, prevNext)
}
func (m *machine) insertReloadRegisterAt(v regalloc.VReg, instr *instruction, after bool) *instruction {
if !v.IsRealReg() {
panic("BUG: VReg must be backed by real reg to be stored")
}
typ := m.compiler.TypeOf(v)
var prevNext, cur *instruction
if after {
cur, prevNext = instr, instr.next
} else {
cur, prevNext = instr.prev, instr
}
offsetFromSP := m.getVRegSpillSlotOffsetFromSP(v.ID(), typ.Size())
var amode addressMode
cur, amode = m.resolveAddressModeForOffsetAndInsert(cur, offsetFromSP, typ.Bits(), spVReg, true)
load := m.allocateInstr()
switch typ {
case ssa.TypeI32, ssa.TypeI64:
load.asULoad(operandNR(v), amode, typ.Bits())
case ssa.TypeF32, ssa.TypeF64:
load.asFpuLoad(operandNR(v), amode, typ.Bits())
case ssa.TypeV128:
load.asFpuLoad(operandNR(v), amode, 128)
default:
panic("TODO")
}
cur = linkInstr(cur, load)
return linkInstr(cur, prevNext)
}
// LowestCommonAncestor implements regalloc.Function LowestCommonAncestor.
func (f *regAllocFunctionImpl) LowestCommonAncestor(blk1, blk2 regalloc.Block) regalloc.Block {
ret := f.m.compiler.SSABuilder().LowestCommonAncestor(blk1.(*regAllocBlockImpl).sb, blk2.(*regAllocBlockImpl).sb)
l := f.m.ssaBlockIDToLabels[ret.ID()]
index := f.labelToRegAllocBlockIndex[l]
return &f.reversePostOrderBlocks[index]
}
func (f *regAllocFunctionImpl) Idom(blk regalloc.Block) regalloc.Block {
builder := f.m.compiler.SSABuilder()
idom := builder.Idom(blk.(*regAllocBlockImpl).sb)
if idom == nil {
panic("BUG: idom must not be nil")
}
l := f.m.ssaBlockIDToLabels[idom.ID()]
index := f.labelToRegAllocBlockIndex[l]
return &f.reversePostOrderBlocks[index]
}
// FirstInstr implements regalloc.Block FirstInstr.
func (r *regAllocBlockImpl) FirstInstr() regalloc.Instr {
return r.pos.begin
}
// LastInstr implements regalloc.Block LastInstr.
func (r *regAllocBlockImpl) LastInstr() regalloc.Instr {
if r.cachedLastInstr == nil {
cur := r.pos.end
for cur.kind == nop0 {
cur = cur.prev
if cur == r.pos.begin {
r.cachedLastInstr = r.pos.end
return r.cachedLastInstr
}
}
switch cur.kind {
case br:
r.cachedLastInstr = cur
default:
r.cachedLastInstr = r.pos.end
}
}
return r.cachedLastInstr
}