wazevo: add VRegTable data structure (#1802)

Signed-off-by: Achille Roussel <achille.roussel@gmail.com>
This commit is contained in:
Achille
2023-10-19 09:36:56 -07:00
committed by GitHub
parent 2084866060
commit 4b6f8f7ccf
6 changed files with 250 additions and 49 deletions

View File

@@ -13,7 +13,7 @@ import (
const xArgRetRegMax, vArgRetRegMax = x7, v7 // x0-x7 & v0-v7.
var regInfo = &regalloc.RegisterInfo{
AllocatableRegisters: [regalloc.RegTypeNum][]regalloc.RealReg{
AllocatableRegisters: [regalloc.NumRegType][]regalloc.RealReg{
// We don't allocate:
// - x18: Reserved by the macOS: https://developer.apple.com/documentation/xcode/writing-arm64-code-for-apple-platforms#Respect-the-purpose-of-specific-CPU-registers
// - x28: Reserved by Go runtime.

View File

@@ -2,6 +2,7 @@ package regalloc
import (
"fmt"
"math/bits"
"github.com/tetratelabs/wazero/internal/engine/wazevo/ssa"
)
@@ -13,6 +14,10 @@ type VReg uint64
// VRegID is the lower 32bit of VReg, which is the pure identifier of VReg without RealReg info.
type VRegID uint32
const (
MaxVRegID = ^VRegID(0)
)
// RealReg returns the RealReg of this VReg.
func (v VReg) RealReg() RealReg {
return RealReg(v >> 32)
@@ -58,6 +63,130 @@ func (v VReg) Valid() bool {
return v.ID() != vRegIDInvalid && v.RegType() != RegTypeInvalid
}
// VRegIDMinSet is used to collect the minimum ID by type in a collection of
// virtual registers.
//
// We store the min values + 1 so the zero-value of the VRegIDMinSet is valid.
type VRegIDMinSet [NumRegType]VRegID
func (mins *VRegIDMinSet) Min(t RegType) VRegID {
return mins[t] - 1
}
func (mins *VRegIDMinSet) Observe(v VReg) {
if rt, id := v.RegType(), v.ID(); id < (mins[rt] - 1) {
mins[rt] = id + 1
}
}
// VRegTable is a data structure designed for fast association of program
// counters to virtual registers.
type VRegTable [NumRegType]VRegTypeTable
func (t *VRegTable) Contains(v VReg) bool {
return t[v.RegType()].Contains(v.ID())
}
func (t *VRegTable) Lookup(v VReg) programCounter {
return t[v.RegType()].Lookup(v.ID())
}
func (t *VRegTable) Insert(v VReg, p programCounter) {
if v.IsRealReg() {
panic("BUG: cannot insert real registers in virtual register table")
}
t[v.RegType()].Insert(v.ID(), p)
}
func (t *VRegTable) Range(f func(VReg, programCounter)) {
for i := range t {
t[i].Range(func(id VRegID, p programCounter) {
f(VReg(id).SetRegType(RegType(i)), p)
})
}
}
func (t *VRegTable) Reset(minVRegIDs VRegIDMinSet) {
for i := range t {
t[i].Reset(minVRegIDs.Min(RegType(i)))
}
}
// VRegTypeTable implements a table for virtual registers of a specific type.
//
// The parent VRegTable uses 4 instances of this type to maintain a table for
// each virtual register type.
//
// The virtual register type table uses a bitset to accelerate checking whether
// a virtual register exists in the table. The bitset also helps identify which
// slots in the program counter array are used, since the table is sparse and
// some locations may not be occupied and will have the value zero even though
// zero is a valid program counter value.
type VRegTypeTable struct {
min VRegID
set bitset
pcs []programCounter
}
func (t *VRegTypeTable) Contains(id VRegID) bool {
return t.set.has(uint(id - t.min))
}
func (t *VRegTypeTable) Lookup(id VRegID) programCounter {
if id := int(id - t.min); t.set.has(uint(id)) {
return t.pcs[id]
}
return -1
}
func (t *VRegTypeTable) Insert(id VRegID, p programCounter) {
if p < 0 {
panic("BUG: cannot insert negative program counter in virtual register table")
}
i := int(id - t.min)
if len(t.pcs) <= i {
t.pcs = append(t.pcs, make([]programCounter, (i+1)-len(t.pcs))...)
}
t.set.set(uint(i))
t.pcs[i] = p
}
func (t *VRegTypeTable) Range(f func(VRegID, programCounter)) {
t.set.scan(func(i uint) { f(VRegID(i)+t.min, t.pcs[i]) })
}
func (t *VRegTypeTable) Reset(minVRegID VRegID) {
t.min = minVRegID
t.set = nil
t.pcs = nil
}
type bitset []uint64
func (b bitset) scan(f func(uint)) {
for i, v := range b {
for j := uint(i * 64); v != 0; j++ {
n := uint(bits.TrailingZeros64(v))
j += n
v >>= (n + 1)
f(j)
}
}
}
func (b bitset) has(i uint) bool {
index, shift := i/64, i%64
return index < uint(len(b)) && ((b[index] & (1 << shift)) != 0)
}
func (b *bitset) set(i uint) {
index, shift := i/64, i%64
if index >= uint(len(*b)) {
*b = append(*b, make([]uint64, (index+1)-uint(len(*b)))...)
}
(*b)[index] |= 1 << shift
}
// RealReg represents a physical register.
type RealReg byte
@@ -95,7 +224,7 @@ const (
RegTypeInvalid RegType = iota
RegTypeInt
RegTypeFloat
RegTypeNum
NumRegType
)
// String implements fmt.Stringer.

View File

@@ -25,3 +25,33 @@ func Test_FromRealReg(t *testing.T) {
require.Equal(t, RealReg(5), r.RealReg())
require.Equal(t, VRegID(5), r.ID())
}
func TestVRegTable(t *testing.T) {
min := VRegIDMinSet{}
min.Observe(VReg(vRegIDReservedForRealNum + 2))
min.Observe(VReg(vRegIDReservedForRealNum + 1))
min.Observe(VReg(vRegIDReservedForRealNum + 0))
table := VRegTable{}
table.Reset(min)
table.Insert(VReg(vRegIDReservedForRealNum+0), 1)
table.Insert(VReg(vRegIDReservedForRealNum+1), 10)
table.Insert(VReg(vRegIDReservedForRealNum+2), 100)
vregs := map[VReg]programCounter{}
table.Range(func(v VReg, p programCounter) {
vregs[v] = p
})
require.Equal(t, 3, len(vregs))
for v, p := range vregs {
require.True(t, table.Contains(v))
require.Equal(t, p, table.Lookup(v))
}
table.Range(func(v VReg, p programCounter) {
require.Equal(t, vregs[v], p)
delete(vregs, v)
})
require.Equal(t, 0, len(vregs))
}

View File

@@ -36,7 +36,7 @@ type (
RegisterInfo struct {
// AllocatableRegisters is a 2D array of allocatable RealReg, indexed by regTypeNum and regNum.
// The order matters: the first element is the most preferred one when allocating.
AllocatableRegisters [RegTypeNum][]RealReg
AllocatableRegisters [NumRegType][]RealReg
CalleeSavedRegisters [RealRegsNumMax]bool
CallerSavedRegisters [RealRegsNumMax]bool
RealRegToVReg []VReg
@@ -78,7 +78,7 @@ type (
liveOuts map[VReg]struct{}
liveIns map[VReg]struct{}
defs map[VReg]programCounter
lastUses map[VReg]programCounter
lastUses VRegTable
kills map[VReg]programCounter
// Pre-colored real registers can have multiple live ranges in one block.
realRegUses [vRegIDReservedForRealNum][]programCounter
@@ -172,6 +172,19 @@ func (a *Allocator) livenessAnalysis(f Function) {
for blk := f.PostOrderBlockIteratorBegin(); blk != nil; blk = f.PostOrderBlockIteratorNext() {
info := a.blockInfoAt(blk.ID())
// We have to do a first pass to find the lowest VRegID in the block;
// this is used to reduce memory utilization in the VRegTable, which
// can avoid allocating memory for registers zero to minVRegID-1.
minVRegID := VRegIDMinSet{}
for instr := blk.InstrIteratorBegin(); instr != nil; instr = blk.InstrIteratorNext() {
for _, use := range instr.Uses() {
if !use.IsRealReg() {
minVRegID.Observe(use)
}
}
}
info.lastUses.Reset(minVRegID)
var pc programCounter
for instr := blk.InstrIteratorBegin(); instr != nil; instr = blk.InstrIteratorNext() {
var srcVR, dstVR VReg
@@ -181,7 +194,7 @@ func (a *Allocator) livenessAnalysis(f Function) {
if use.IsRealReg() {
info.addRealRegUsage(use, pos)
} else {
info.lastUses[use] = pos
info.lastUses.Insert(use, pos)
}
}
for _, def := range instr.Defs() {
@@ -209,7 +222,6 @@ func (a *Allocator) livenessAnalysis(f Function) {
}
pc += pcStride
}
if wazevoapi.RegAllocLoggingEnabled {
fmt.Printf("prepared block info for block[%d]:\n%s\n\n", blk.ID(), info.Format(a.regInfo))
}
@@ -231,13 +243,13 @@ func (a *Allocator) livenessAnalysis(f Function) {
// Now that we finished gathering liveIns, liveOuts, defs, and lastUses, the only thing left is to construct kills.
for blk := f.PostOrderBlockIteratorBegin(); blk != nil; blk = f.PostOrderBlockIteratorNext() { // Order doesn't matter.
info := a.blockInfoAt(blk.ID())
lastUses, outs := info.lastUses, info.liveOuts
for use, pc := range lastUses {
outs := info.liveOuts
info.lastUses.Range(func(use VReg, pc programCounter) {
// Usage without live-outs is a kill.
if _, ok := outs[use]; !ok {
info.kills[use] = pc
}
}
})
if wazevoapi.RegAllocLoggingEnabled {
fmt.Printf("\nfinalized info for block[%d]:\n%s\n", blk.ID(), info.Format(a.regInfo))
@@ -251,7 +263,7 @@ func (a *Allocator) beginUpAndMarkStack(f Function, v VReg, isPhi bool, phiDefin
panic(fmt.Sprintf("block without predecessor must be optimized out by the compiler: %d", blk.ID()))
}
info := a.blockInfoAt(blk.ID())
if _, ok := info.lastUses[v]; !ok {
if !info.lastUses.Contains(v) {
continue
}
// TODO: we might want to avoid recursion here.
@@ -463,7 +475,7 @@ func (a *Allocator) Reset() {
func (a *Allocator) allocateBlockInfo(blockID int) *blockInfo {
if blockID >= len(a.blockInfos) {
a.blockInfos = append(a.blockInfos, make([]blockInfo, blockID+1)...)
a.blockInfos = append(a.blockInfos, make([]blockInfo, (blockID+1)-len(a.blockInfos))...)
}
info := &a.blockInfos[blockID]
a.initBlockInfo(info)
@@ -543,11 +555,6 @@ func (a *Allocator) initBlockInfo(i *blockInfo) {
} else {
resetMap(a, i.defs)
}
if i.lastUses == nil {
i.lastUses = make(map[VReg]programCounter)
} else {
resetMap(a, i.lastUses)
}
if i.kills == nil {
i.kills = make(map[VReg]programCounter)
} else {
@@ -587,9 +594,9 @@ func (i *blockInfo) Format(ri *RegisterInfo) string {
buf.WriteString(fmt.Sprintf("%v@%v ", v, pos))
}
buf.WriteString("\n\tlastUses: ")
for v, pos := range i.lastUses {
i.lastUses.Range(func(v VReg, pos programCounter) {
buf.WriteString(fmt.Sprintf("%v@%v ", v, pos))
}
})
buf.WriteString("\n\tkills: ")
for v, pos := range i.kills {
buf.WriteString(fmt.Sprintf("%v@%v ", v, pos))

View File

@@ -6,6 +6,18 @@ import (
"github.com/tetratelabs/wazero/internal/testing/require"
)
func makeVRegTable(vregs map[VReg]programCounter) (table VRegTable) {
min := VRegIDMinSet{}
for v := range vregs {
min.Observe(v)
}
table.Reset(min)
for v, p := range vregs {
table.Insert(v, p)
}
return table
}
func TestAllocator_livenessAnalysis(t *testing.T) {
const realRegID, realRegID2 = 50, 100
realReg, realReg2 := FromRealReg(realRegID, RegTypeInt), FromRealReg(realRegID2, RegTypeInt)
@@ -28,11 +40,12 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
exp: map[int]*blockInfo{
0: {
defs: map[VReg]programCounter{2: pcDefOffset + pcStride, 1: pcDefOffset},
lastUses: map[VReg]programCounter{1: pcStride + pcUseOffset},
lastUses: makeVRegTable(map[VReg]programCounter{1: pcStride + pcUseOffset}),
kills: map[VReg]programCounter{1: pcStride + pcUseOffset},
},
},
},
{
name: "straight",
// b0 -> b1 -> b2
@@ -62,11 +75,11 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
2: pcDefOffset,
3: pcStride*2 + pcDefOffset,
},
lastUses: map[VReg]programCounter{
lastUses: makeVRegTable(map[VReg]programCounter{
1000: pcStride + pcUseOffset,
1: pcStride*2 + pcUseOffset,
2: pcStride*2 + pcUseOffset,
},
}),
liveOuts: map[VReg]struct{}{3: {}},
kills: map[VReg]programCounter{
1000: pcStride + pcUseOffset,
@@ -77,9 +90,9 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
1: {
liveIns: map[VReg]struct{}{3: {}},
liveOuts: map[VReg]struct{}{3: {}, 4: {}, 5: {}},
lastUses: map[VReg]programCounter{
lastUses: makeVRegTable(map[VReg]programCounter{
3: pcStride + pcUseOffset,
},
}),
defs: map[VReg]programCounter{
4: pcStride + pcDefOffset,
5: pcStride + pcDefOffset,
@@ -93,7 +106,7 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
},
2: {
liveIns: map[VReg]struct{}{3: {}, 4: {}, 5: {}},
lastUses: map[VReg]programCounter{3: pcUseOffset, 4: pcUseOffset, 5: pcUseOffset},
lastUses: makeVRegTable(map[VReg]programCounter{3: pcUseOffset, 4: pcUseOffset, 5: pcUseOffset}),
kills: map[VReg]programCounter{3: pcUseOffset, 4: pcUseOffset, 5: pcUseOffset},
},
},
@@ -139,11 +152,12 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
2: pcStride*2 + pcDefOffset,
},
liveOuts: map[VReg]struct{}{1000: {}, 1: {}, 2: {}},
lastUses: makeVRegTable(nil),
},
1: {
liveIns: map[VReg]struct{}{1000: {}, 1: {}},
liveOuts: map[VReg]struct{}{1000: {}},
lastUses: map[VReg]programCounter{1: pcUseOffset},
lastUses: makeVRegTable(map[VReg]programCounter{1: pcUseOffset}),
kills: map[VReg]programCounter{1: pcUseOffset},
realRegDefs: [vRegIDReservedForRealNum][]programCounter{
realRegID: {pcDefOffset, pcStride*4 + pcDefOffset},
@@ -157,18 +171,19 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
2: {
liveIns: map[VReg]struct{}{1000: {}, 2: {}},
liveOuts: map[VReg]struct{}{1000: {}},
lastUses: map[VReg]programCounter{2: pcUseOffset},
lastUses: makeVRegTable(map[VReg]programCounter{2: pcUseOffset}),
kills: map[VReg]programCounter{2: pcUseOffset},
realRegUses: [vRegIDReservedForRealNum][]programCounter{realRegID2: {pcUseOffset}},
realRegDefs: [vRegIDReservedForRealNum][]programCounter{realRegID2: {0}},
},
3: {
liveIns: map[VReg]struct{}{1000: {}},
lastUses: map[VReg]programCounter{1000: pcUseOffset},
lastUses: makeVRegTable(map[VReg]programCounter{1000: pcUseOffset}),
kills: map[VReg]programCounter{1000: pcUseOffset},
},
},
},
{
name: "phis",
// 0
@@ -203,32 +218,35 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
0: {
defs: map[VReg]programCounter{1000: pcDefOffset, 2000: pcDefOffset, 3000: pcDefOffset},
liveOuts: map[VReg]struct{}{1000: {}, 2000: {}, 3000: {}},
lastUses: makeVRegTable(nil),
},
1: {
liveIns: map[VReg]struct{}{2000: {}, 3000: {}},
liveOuts: map[VReg]struct{}{phiVReg: {}, 3000: {}},
defs: map[VReg]programCounter{phiVReg: pcDefOffset},
lastUses: map[VReg]programCounter{2000: pcUseOffset},
lastUses: makeVRegTable(map[VReg]programCounter{2000: pcUseOffset}),
kills: map[VReg]programCounter{2000: pcUseOffset},
},
2: {
liveIns: map[VReg]struct{}{phiVReg: {}, 3000: {}},
liveOuts: map[VReg]struct{}{phiVReg: {}, 3000: {}},
lastUses: makeVRegTable(nil),
},
3: {
liveIns: map[VReg]struct{}{1000: {}, 3000: {}},
liveOuts: map[VReg]struct{}{phiVReg: {}, 3000: {}},
defs: map[VReg]programCounter{phiVReg: pcDefOffset},
lastUses: map[VReg]programCounter{1000: pcUseOffset},
lastUses: makeVRegTable(map[VReg]programCounter{1000: pcUseOffset}),
kills: map[VReg]programCounter{1000: pcUseOffset},
},
4: {
liveIns: map[VReg]struct{}{phiVReg: {}, 3000: {}},
lastUses: map[VReg]programCounter{phiVReg: pcUseOffset, 3000: pcUseOffset},
lastUses: makeVRegTable(map[VReg]programCounter{phiVReg: pcUseOffset, 3000: pcUseOffset}),
kills: map[VReg]programCounter{phiVReg: pcUseOffset, 3000: pcUseOffset},
},
},
},
{
name: "loop",
// 0 -> 1 -> 2
@@ -274,9 +292,9 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
1: pcDefOffset,
phiVReg: pcStride + pcDefOffset,
},
lastUses: map[VReg]programCounter{
lastUses: makeVRegTable(map[VReg]programCounter{
1: pcStride + pcUseOffset,
},
}),
kills: map[VReg]programCounter{
1: pcStride + pcUseOffset,
},
@@ -285,32 +303,33 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
liveIns: map[VReg]struct{}{phiVReg: {}},
liveOuts: map[VReg]struct{}{phiVReg: {}, 9999: {}},
defs: map[VReg]programCounter{9999: pcDefOffset},
lastUses: map[VReg]programCounter{},
lastUses: makeVRegTable(map[VReg]programCounter{}),
kills: map[VReg]programCounter{},
},
2: {
liveIns: map[VReg]struct{}{phiVReg: {}, 9999: {}},
liveOuts: map[VReg]struct{}{100: {}},
defs: map[VReg]programCounter{100: pcDefOffset},
lastUses: map[VReg]programCounter{phiVReg: pcUseOffset, 9999: pcUseOffset},
lastUses: makeVRegTable(map[VReg]programCounter{phiVReg: pcUseOffset, 9999: pcUseOffset}),
kills: map[VReg]programCounter{phiVReg: pcUseOffset, 9999: pcUseOffset},
},
3: {
liveIns: map[VReg]struct{}{100: {}},
liveOuts: map[VReg]struct{}{54321: {}},
defs: map[VReg]programCounter{54321: pcDefOffset},
lastUses: map[VReg]programCounter{100: pcStride + pcUseOffset},
lastUses: makeVRegTable(map[VReg]programCounter{100: pcStride + pcUseOffset}),
kills: map[VReg]programCounter{100: pcStride + pcUseOffset},
},
4: {
liveIns: map[VReg]struct{}{54321: {}},
liveOuts: map[VReg]struct{}{phiVReg: {}},
defs: map[VReg]programCounter{phiVReg: pcDefOffset},
lastUses: map[VReg]programCounter{54321: pcUseOffset},
lastUses: makeVRegTable(map[VReg]programCounter{54321: pcUseOffset}),
kills: map[VReg]programCounter{54321: pcUseOffset},
},
},
},
{
// -----+
// v |
@@ -336,23 +355,29 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
0: {
defs: map[VReg]programCounter{99999: pcDefOffset},
liveOuts: map[VReg]struct{}{99999: {}},
lastUses: makeVRegTable(nil),
},
1: {
liveIns: map[VReg]struct{}{99999: {}},
liveOuts: map[VReg]struct{}{99999: {}},
lastUses: map[VReg]programCounter{99999: pcUseOffset},
lastUses: makeVRegTable(map[VReg]programCounter{99999: pcUseOffset}),
},
2: {
liveIns: map[VReg]struct{}{99999: {}},
liveOuts: map[VReg]struct{}{99999: {}},
lastUses: makeVRegTable(nil),
},
3: {
liveIns: map[VReg]struct{}{99999: {}},
liveOuts: map[VReg]struct{}{99999: {}},
lastUses: makeVRegTable(nil),
},
4: {
lastUses: makeVRegTable(nil),
},
4: {},
},
},
// 2
// ^ +----+
// | v |
@@ -405,35 +430,48 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
return newMockFunction(b0, b1, b2, b3, b4, b7, b8, b5, b6, b9)
},
exp: map[int]*blockInfo{
0: {},
1: {},
2: {},
0: {
lastUses: makeVRegTable(nil),
},
1: {
lastUses: makeVRegTable(nil),
},
2: {
lastUses: makeVRegTable(nil),
},
3: {
liveOuts: map[VReg]struct{}{100: {}},
defs: map[VReg]programCounter{100: pcDefOffset},
liveOuts: map[VReg]struct{}{100: {}},
lastUses: makeVRegTable(nil),
},
4: {
liveIns: map[VReg]struct{}{100: {}},
liveOuts: map[VReg]struct{}{100: {}},
lastUses: makeVRegTable(nil),
},
5: {
liveIns: map[VReg]struct{}{100: {}},
liveOuts: map[VReg]struct{}{100: {}},
lastUses: map[VReg]programCounter{100: pcUseOffset},
lastUses: makeVRegTable(map[VReg]programCounter{100: pcUseOffset}),
},
6: {
liveIns: map[VReg]struct{}{100: {}},
liveOuts: map[VReg]struct{}{100: {}},
lastUses: makeVRegTable(nil),
},
7: {
liveIns: map[VReg]struct{}{100: {}},
liveOuts: map[VReg]struct{}{100: {}},
lastUses: makeVRegTable(nil),
},
8: {
liveIns: map[VReg]struct{}{100: {}},
liveOuts: map[VReg]struct{}{100: {}},
lastUses: makeVRegTable(nil),
},
9: {
lastUses: makeVRegTable(nil),
},
9: {},
},
},
} {
@@ -448,7 +486,7 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
initMapInInfo(exp)
saved := actual.intervalMng
actual.intervalMng = nil // Don't compare intervalManager.
require.Equal(t, exp, actual, "\n[exp for block[%d]]\n%s\n[actual for block[%d]]\n%s", blockID, exp, blockID, actual)
require.Equal(t, exp, actual, "\n[exp for block[%d]]\n%v\n[actual for block[%d]]\n%v", blockID, exp, blockID, actual)
actual.intervalMng = saved
}
@@ -533,9 +571,6 @@ func initMapInInfo(info *blockInfo) {
if info.kills == nil {
info.kills = make(map[VReg]programCounter)
}
if info.lastUses == nil {
info.lastUses = make(map[VReg]programCounter)
}
}
func TestNode_assignedRealReg(t *testing.T) {

View File

@@ -29,7 +29,7 @@ func TestSpillHandler_getUnusedOrEvictReg(t *testing.T) {
require.Equal(t, 4, len(s.activeRegs))
regInfo := RegisterInfo{
AllocatableRegisters: [RegTypeNum][]RealReg{
AllocatableRegisters: [NumRegType][]RealReg{
RegTypeInt: {
RealReg(0xff), // unused.
RealReg(0), RealReg(1), RealReg(2),