From 2ca59ecee89b8b0cd2424f4cd85f4c8c46908f6f Mon Sep 17 00:00:00 2001 From: Takeshi Yoneda Date: Wed, 18 Oct 2023 16:09:54 +0900 Subject: [PATCH] wazevo: pass reset func to NewPool (#1796) Signed-off-by: Takeshi Yoneda --- internal/engine/wazevo/backend/compiler.go | 1 + internal/engine/wazevo/backend/compiler_lower.go | 5 +++++ .../engine/wazevo/backend/isa/arm64/machine.go | 16 +++++++++++----- .../wazevo/backend/isa/arm64/machine_test.go | 3 +-- .../wazevo/backend/regalloc/interval_tree.go | 16 ++++++++++++---- .../engine/wazevo/backend/regalloc/regalloc.go | 16 +++++++++++----- .../wazevo/backend/regalloc/regalloc_test.go | 6 +++--- internal/engine/wazevo/ssa/basic_block.go | 2 +- internal/engine/wazevo/ssa/builder.go | 8 +++----- internal/engine/wazevo/ssa/instructions.go | 4 ++-- internal/engine/wazevo/wazevoapi/pool.go | 13 +++++-------- 11 files changed, 55 insertions(+), 35 deletions(-) diff --git a/internal/engine/wazevo/backend/compiler.go b/internal/engine/wazevo/backend/compiler.go index 2b1ed5b1..04bb9683 100644 --- a/internal/engine/wazevo/backend/compiler.go +++ b/internal/engine/wazevo/backend/compiler.go @@ -135,6 +135,7 @@ type compiler struct { dst regalloc.VReg } vRegSet []bool + vRegIDs []regalloc.VRegID tempRegs []regalloc.VReg tmpVals []ssa.Value ssaTypeOfVRegID [] /* VRegID to */ ssa.Type diff --git a/internal/engine/wazevo/backend/compiler_lower.go b/internal/engine/wazevo/backend/compiler_lower.go index 0b3d03a7..8f513423 100644 --- a/internal/engine/wazevo/backend/compiler_lower.go +++ b/internal/engine/wazevo/backend/compiler_lower.go @@ -160,12 +160,14 @@ func (c *compiler) lowerBlockArguments(args []ssa.Value, succ ssa.BasicBlock) { } // Check if there's an overlap among the dsts and srcs in varEdges. + c.vRegIDs = c.vRegIDs[:0] for _, edge := range c.varEdges { src := edge[0].ID() if int(src) >= len(c.vRegSet) { c.vRegSet = append(c.vRegSet, make([]bool, src+1)...) } c.vRegSet[src] = true + c.vRegIDs = append(c.vRegIDs, src) } separated := true for _, edge := range c.varEdges { @@ -179,6 +181,9 @@ func (c *compiler) lowerBlockArguments(args []ssa.Value, succ ssa.BasicBlock) { } } } + for _, id := range c.vRegIDs { + c.vRegSet[id] = false // reset for the next use. + } if separated { // If there's no overlap, we can simply move the source to destination. diff --git a/internal/engine/wazevo/backend/isa/arm64/machine.go b/internal/engine/wazevo/backend/isa/arm64/machine.go index 65bae30b..046bf6be 100644 --- a/internal/engine/wazevo/backend/isa/arm64/machine.go +++ b/internal/engine/wazevo/backend/isa/arm64/machine.go @@ -112,8 +112,8 @@ const ( // NewBackend returns a new backend for arm64. func NewBackend() backend.Machine { m := &machine{ - instrPool: wazevoapi.NewPool[instruction](), - labelPositionPool: wazevoapi.NewPool[labelPosition](), + instrPool: wazevoapi.NewPool[instruction](resetInstruction), + labelPositionPool: wazevoapi.NewPool[labelPosition](resetLabelPosition), labelPositions: make(map[label]*labelPosition), spillSlots: make(map[regalloc.VRegID]int64), nextLabel: invalidLabel, @@ -243,10 +243,13 @@ func (m *machine) insertBrTargetLabel() label { func (m *machine) allocateLabelPosition() *labelPosition { l := m.labelPositionPool.Allocate() - *l = labelPosition{} return l } +func resetLabelPosition(l *labelPosition) { + *l = labelPosition{} +} + func (m *machine) FlushPendingInstructions() { l := len(m.pendingInstructions) if l == 0 { @@ -277,15 +280,18 @@ func (l label) String() string { // allocateInstr allocates an instruction. func (m *machine) allocateInstr() *instruction { instr := m.instrPool.Allocate() - *instr = instruction{} if !m.regAllocStarted { instr.addedBeforeRegAlloc = true } return instr } +func resetInstruction(i *instruction) { + *i = instruction{} +} + func (m *machine) allocateNop() *instruction { - instr := m.instrPool.Allocate() + instr := m.allocateInstr() instr.asNop0() return instr } diff --git a/internal/engine/wazevo/backend/isa/arm64/machine_test.go b/internal/engine/wazevo/backend/isa/arm64/machine_test.go index 7c54a436..9edd9331 100644 --- a/internal/engine/wazevo/backend/isa/arm64/machine_test.go +++ b/internal/engine/wazevo/backend/isa/arm64/machine_test.go @@ -57,8 +57,7 @@ func TestMachine_resolveAddressingMode(t *testing.T) { }) t.Run("tmp reg", func(t *testing.T) { - // 0x89705f4136b4a598 - m := &machine{instrPool: wazevoapi.NewPool[instruction]()} + m := &machine{instrPool: wazevoapi.NewPool[instruction](resetInstruction)} root := &instruction{kind: udf} i := &instruction{prev: root} i.asULoad(operandNR(x17VReg), addressMode{ diff --git a/internal/engine/wazevo/backend/regalloc/interval_tree.go b/internal/engine/wazevo/backend/regalloc/interval_tree.go index 7d74378c..9dd9475f 100644 --- a/internal/engine/wazevo/backend/regalloc/interval_tree.go +++ b/internal/engine/wazevo/backend/regalloc/interval_tree.go @@ -32,7 +32,7 @@ func (t *intervalTree) reset() { func newIntervalTree() *intervalTree { return &intervalTree{ - allocator: wazevoapi.NewPool[intervalTreeNode](), + allocator: wazevoapi.NewPool[intervalTreeNode](resetIntervalTreeNode), intervals: make(map[uint64]*intervalTreeNode), } } @@ -48,16 +48,24 @@ type intervalTreeNode struct { // TODO: color for red-black balancing. } +func resetIntervalTreeNode(i *intervalTreeNode) { + i.begin = 0 + i.end = 0 + i.nodes = i.nodes[:0] + i.maxEnd = 0 + i.neighbors = i.neighbors[:0] + i.left = nil + i.right = nil +} + func (i *intervalTreeNode) insert(t *intervalTree, n *node, begin, end programCounter) *intervalTreeNode { if i == nil { intervalNode := t.allocator.Allocate() - intervalNode.right = nil - intervalNode.left = nil intervalNode.nodes = append(intervalNode.nodes, n) intervalNode.maxEnd = end intervalNode.begin = begin intervalNode.end = end - key := uint64(begin) | uint64(end)<<32 + key := intervalTreeNodeKey(begin, end) t.intervals[key] = intervalNode return intervalNode } diff --git a/internal/engine/wazevo/backend/regalloc/regalloc.go b/internal/engine/wazevo/backend/regalloc/regalloc.go index 9c9efbbf..0dc91216 100644 --- a/internal/engine/wazevo/backend/regalloc/regalloc.go +++ b/internal/engine/wazevo/backend/regalloc/regalloc.go @@ -21,7 +21,7 @@ import ( func NewAllocator(allocatableRegs *RegisterInfo) Allocator { a := Allocator{ regInfo: allocatableRegs, - nodePool: wazevoapi.NewPool[node](), + nodePool: wazevoapi.NewPool[node](resetNode), phis: make(map[VReg]Block), } for _, regs := range allocatableRegs.AllocatableRegisters { @@ -487,16 +487,22 @@ func (a *Allocator) getOrAllocateNode(v VReg) (n *node) { return } -func (a *Allocator) allocateNode() (n *node) { - n = a.nodePool.Allocate() - n.id = a.nodePool.Allocated() - 1 +func resetNode(n *node) { + n.r = RealRegInvalid + n.v = VRegInvalid n.ranges = n.ranges[:0] + n.neighbors = n.neighbors[:0] n.copyFromVReg = nil n.copyToVReg = nil n.copyFromReal = RealRegInvalid n.copyToReal = RealRegInvalid - n.neighbors = n.neighbors[:0] + n.degree = 0 n.visited = false +} + +func (a *Allocator) allocateNode() (n *node) { + n = a.nodePool.Allocate() + n.id = a.nodePool.Allocated() - 1 return } diff --git a/internal/engine/wazevo/backend/regalloc/regalloc_test.go b/internal/engine/wazevo/backend/regalloc/regalloc_test.go index af9736c0..2e51478a 100644 --- a/internal/engine/wazevo/backend/regalloc/regalloc_test.go +++ b/internal/engine/wazevo/backend/regalloc/regalloc_test.go @@ -445,7 +445,10 @@ func TestAllocator_livenessAnalysis(t *testing.T) { actual := &a.blockInfos[blockID] exp := tc.exp[blockID] initMapInInfo(exp) + saved := actual.intervalTree + actual.intervalTree = nil // Don't compare intervalTree. require.Equal(t, exp, actual, "\n[exp for block[%d]]\n%s\n[actual for block[%d]]\n%s", blockID, exp, blockID, actual) + actual.intervalTree = saved } // Sanity check: buildLiveRanges should not panic. @@ -517,9 +520,6 @@ func TestAllocator_recordCopyRelation(t *testing.T) { } func initMapInInfo(info *blockInfo) { - if info.intervalTree == nil { - info.intervalTree = newIntervalTree() - } if info.liveIns == nil { info.liveIns = make(map[VReg]struct{}) } diff --git a/internal/engine/wazevo/ssa/basic_block.go b/internal/engine/wazevo/ssa/basic_block.go index 13ba0963..7472cbc2 100644 --- a/internal/engine/wazevo/ssa/basic_block.go +++ b/internal/engine/wazevo/ssa/basic_block.go @@ -245,7 +245,7 @@ func (bb *basicBlock) Tail() *Instruction { } // reset resets the basicBlock to its initial state so that it can be reused for another function. -func (bb *basicBlock) reset() { +func resetBasicBlock(bb *basicBlock) { bb.params = bb.params[:0] bb.rootInstr, bb.currentInstr = nil, nil bb.preds = bb.preds[:0] diff --git a/internal/engine/wazevo/ssa/builder.go b/internal/engine/wazevo/ssa/builder.go index 8abe0a0b..e7a02b47 100644 --- a/internal/engine/wazevo/ssa/builder.go +++ b/internal/engine/wazevo/ssa/builder.go @@ -126,8 +126,8 @@ type Builder interface { // NewBuilder returns a new Builder implementation. func NewBuilder() Builder { return &builder{ - instructionsPool: wazevoapi.NewPool[Instruction](), - basicBlocksPool: wazevoapi.NewPool[basicBlock](), + instructionsPool: wazevoapi.NewPool[Instruction](resetInstruction), + basicBlocksPool: wazevoapi.NewPool[basicBlock](resetBasicBlock), valueAnnotations: make(map[ValueID]string), signatures: make(map[SignatureID]*Signature), blkVisited: make(map[*basicBlock]int), @@ -196,7 +196,7 @@ func (b *builder) ReturnBlock() BasicBlock { // Init implements Builder.Reset. func (b *builder) Init(s *Signature) { b.currentSignature = s - b.returnBlk.reset() + resetBasicBlock(b.returnBlk) b.instructionsPool.Reset() b.basicBlocksPool.Reset() b.donePasses = false @@ -211,7 +211,6 @@ func (b *builder) Init(s *Signature) { for i := 0; i < b.basicBlocksPool.Allocated(); i++ { blk := b.basicBlocksPool.View(i) - blk.reset() delete(b.blkVisited, blk) } b.basicBlocksPool.Reset() @@ -250,7 +249,6 @@ func (b *builder) AnnotateValue(value Value, a string) { // AllocateInstruction implements Builder.AllocateInstruction. func (b *builder) AllocateInstruction() *Instruction { instr := b.instructionsPool.Allocate() - instr.reset() return instr } diff --git a/internal/engine/wazevo/ssa/instructions.go b/internal/engine/wazevo/ssa/instructions.go index 694cb2da..f04be1b2 100644 --- a/internal/engine/wazevo/ssa/instructions.go +++ b/internal/engine/wazevo/ssa/instructions.go @@ -74,8 +74,8 @@ func (i *Instruction) Lowered() bool { return i.alreadyLowered } -// reset resets this instruction to the initial state. -func (i *Instruction) reset() { +// resetInstruction resets this instruction to the initial state. +func resetInstruction(i *Instruction) { *i = Instruction{} i.v = ValueInvalid i.v2 = ValueInvalid diff --git a/internal/engine/wazevo/wazevoapi/pool.go b/internal/engine/wazevo/wazevoapi/pool.go index bcbadb9a..500b3287 100644 --- a/internal/engine/wazevo/wazevoapi/pool.go +++ b/internal/engine/wazevo/wazevoapi/pool.go @@ -6,12 +6,15 @@ const poolPageSize = 128 // This is useful to avoid unnecessary allocations. type Pool[T any] struct { pages []*[poolPageSize]T + resetFn func(*T) allocated, index int } // NewPool returns a new Pool. -func NewPool[T any]() Pool[T] { +// resetFn is called when a new T is allocated in Pool.Allocate. +func NewPool[T any](resetFn func(*T)) Pool[T] { var ret Pool[T] + ret.resetFn = resetFn ret.Reset() return ret } @@ -36,6 +39,7 @@ func (p *Pool[T]) Allocate() *T { p.index = 0 } ret := &p.pages[len(p.pages)-1][p.index] + p.resetFn(ret) p.index++ p.allocated++ return ret @@ -49,13 +53,6 @@ func (p *Pool[T]) View(i int) *T { // Reset resets the pool. func (p *Pool[T]) Reset() { - for _, ns := range p.pages { - pages := ns[:] - for i := range pages { - var v T - pages[i] = v - } - } p.pages = p.pages[:0] p.index = poolPageSize p.allocated = 0