Repository: tetratelabs/wazero
Pull request #1725 (2023-09-19 22:40:00)

wazevo(arm64): remove allocInstrAfterRegalloc API for simplicity
#1496 


Patch:
diff --git a/internal/engine/wazevo/backend/isa/arm64/abi.go b/internal/engine/wazevo/backend/isa/arm64/abi.go
index aa777dcc..80a96a3d 100644
--- a/internal/engine/wazevo/backend/isa/arm64/abi.go
+++ b/internal/engine/wazevo/backend/isa/arm64/abi.go
@@ -305,9 +305,6 @@ func (m *machine) resolveAddressModeForOffsetAndInsert(cur *instruction, offset
 	m.pendingInstructions = m.pendingInstructions[:0]
 	mode := m.resolveAddressModeForOffset(offset, dstBits, rn)
 	for _, instr := range m.pendingInstructions {
-		// This is called after/during alloc, so set the flag so that
-		// this insertion won't interfere with the register allocation.
-		instr.addedAfterLowering = true
 		cur = linkInstr(cur, instr)
 	}
 	return cur, mode
@@ -351,7 +348,7 @@ func (m *machine) lowerCall(si *ssa.Instruction) {
 
 	stackSlotSize := calleeABI.alignedArgResultStackSlotSize()
 	if stackSlotSize > 0 {
-		m.insertAddOrSubStackPointer(spVReg, stackSlotSize, false /* == sub */, false)
+		m.insertAddOrSubStackPointer(spVReg, stackSlotSize, false /* == sub */)
 	}
 
 	if m.maxRequiredStackSizeForCalls < stackSlotSize {
@@ -388,18 +385,13 @@ func (m *machine) lowerCall(si *ssa.Instruction) {
 	}
 
 	if stackSlotSize > 0 {
-		m.insertAddOrSubStackPointer(spVReg, stackSlotSize, true /* add */, false)
+		m.insertAddOrSubStackPointer(spVReg, stackSlotSize, true /* add */)
 	}
 }
 
-func (m *machine) insertAddOrSubStackPointer(rd regalloc.VReg, diff int64, add bool, afterLowering bool) {
+func (m *machine) insertAddOrSubStackPointer(rd regalloc.VReg, diff int64, add bool) {
 	if imm12Operand, ok := asImm12Operand(uint64(diff)); ok {
-		var alu *instruction
-		if afterLowering {
-			alu = m.allocateInstrAfterLowering()
-		} else {
-			alu = m.allocateInstr()
-		}
+		alu := m.allocateInstr()
 		var ao aluOp
 		if add {
 			ao = aluOpAdd
@@ -410,12 +402,7 @@ func (m *machine) insertAddOrSubStackPointer(rd regalloc.VReg, diff int64, add b
 		m.insert(alu)
 	} else {
 		m.lowerConstantI64(tmpRegVReg, diff)
-		var alu *instruction
-		if afterLowering {
-			alu = m.allocateInstrAfterLowering()
-		} else {
-			alu = m.allocateInstr()
-		}
+		alu := m.allocateInstr()
 		var ao aluOp
 		if add {
 			ao = aluOpAdd
diff --git a/internal/engine/wazevo/backend/isa/arm64/abi_go_call.go b/internal/engine/wazevo/backend/isa/arm64/abi_go_call.go
index dfc402d9..6d300d50 100644
--- a/internal/engine/wazevo/backend/isa/arm64/abi_go_call.go
+++ b/internal/engine/wazevo/backend/isa/arm64/abi_go_call.go
@@ -94,7 +94,7 @@ func (m *machine) CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *
 	}
 
 	// Advances the stack pointer.
-	cur = m.addsAddOrSubStackPointer(cur, spVReg, goCallStackSize, false, true)
+	cur = m.addsAddOrSubStackPointer(cur, spVReg, goCallStackSize, false)
 
 	// Copy the pointer to x15VReg.
 	arg0ret0AddrReg := x15VReg // Caller save, so we can use it for whatever we want.
@@ -158,11 +158,11 @@ func (m *machine) CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *
 
 	// Get the pointer to the arg[0]/ret[0]: We need to skip `frame_size + sliceSize`.
 	if len(abi.rets) > 0 {
-		cur = m.addsAddOrSubStackPointer(cur, arg0ret0AddrReg, 16, true, true)
+		cur = m.addsAddOrSubStackPointer(cur, arg0ret0AddrReg, 16, true)
 	}
 
 	// Removes the stack frame! --> Stack pointer now points to the original arg 0 slot.
-	cur = m.addsAddOrSubStackPointer(cur, spVReg, goCallStackSize+callFrameSize, true, true)
+	cur = m.addsAddOrSubStackPointer(cur, spVReg, goCallStackSize+callFrameSize, true)
 
 	for i := range abi.rets {
 		r := &abi.rets[i]
@@ -224,7 +224,7 @@ func (m *machine) CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *
 		}
 	}
 
-	ret := m.allocateInstrAfterLowering()
+	ret := m.allocateInstr()
 	ret.asRet(nil)
 	linkInstr(cur, ret)
 
@@ -235,7 +235,7 @@ func (m *machine) CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *
 func (m *machine) saveRegistersInExecutionContext(cur *instruction, regs []regalloc.VReg) *instruction {
 	offset := wazevoapi.ExecutionContextOffsetSavedRegistersBegin.I64()
 	for _, v := range regs {
-		store := m.allocateInstrAfterLowering()
+		store := m.allocateInstr()
 		var sizeInBits byte
 		switch v.RegType() {
 		case regalloc.RegTypeInt:
@@ -260,7 +260,7 @@ func (m *machine) saveRegistersInExecutionContext(cur *instruction, regs []regal
 func (m *machine) restoreRegistersInExecutionContext(cur *instruction, regs []regalloc.VReg) *instruction {
 	offset := wazevoapi.ExecutionContextOffsetSavedRegistersBegin.I64()
 	for _, v := range regs {
-		load := m.allocateInstrAfterLowering()
+		load := m.allocateInstr()
 		var as func(dst operand, amode addressMode, sizeInBits byte)
 		var sizeInBits byte
 		switch v.RegType() {
@@ -306,7 +306,7 @@ func (m *machine) setExitCode(cur *instruction, execCtr regalloc.VReg, exitCode
 	cur = m.lowerConstantI32AndInsert(cur, constReg, int32(exitCode))
 
 	// Set the exit status on the execution context.
-	setExistStatus := m.allocateInstrAfterLowering()
+	setExistStatus := m.allocateInstr()
 	setExistStatus.asStore(operandNR(constReg),
 		addressMode{
 			kind: addressModeKindRegUnsignedImm12,
@@ -318,11 +318,11 @@ func (m *machine) setExitCode(cur *instruction, execCtr regalloc.VReg, exitCode
 
 func (m *machine) storeReturnAddressAndExit(cur *instruction) *instruction {
 	// Read the return address into tmp, and store it in the execution context.
-	adr := m.allocateInstrAfterLowering()
+	adr := m.allocateInstr()
 	adr.asAdr(tmpRegVReg, exitSequenceSize+8)
 	cur = linkInstr(cur, adr)
 
-	storeReturnAddr := m.allocateInstrAfterLowering()
+	storeReturnAddr := m.allocateInstr()
 	storeReturnAddr.asStore(operandNR(tmpRegVReg),
 		addressMode{
 			kind: addressModeKindRegUnsignedImm12,
@@ -332,7 +332,7 @@ func (m *machine) storeReturnAddressAndExit(cur *instruction) *instruction {
 	cur = linkInstr(cur, storeReturnAddr)
 
 	// Exit the execution.
-	trapSeq := m.allocateInstrAfterLowering()
+	trapSeq := m.allocateInstr()
 	trapSeq.asExitSequence(x0VReg)
 	cur = linkInstr(cur, trapSeq)
 	return cur
@@ -342,11 +342,11 @@ func (m *machine) saveCurrentStackPointer(cur *instruction, execCtr regalloc.VRe
 	// Save the current stack pointer:
 	// 	mov tmp, sp,
 	// 	str tmp, [exec_ctx, #stackPointerBeforeGoCall]
-	movSp := m.allocateInstrAfterLowering()
+	movSp := m.allocateInstr()
 	movSp.asMove64(tmpRegVReg, spVReg)
 	cur = linkInstr(cur, movSp)
 
-	strSp := m.allocateInstrAfterLowering()
+	strSp := m.allocateInstr()
 	strSp.asStore(operandNR(tmpRegVReg),
 		addressMode{
 			kind: addressModeKindRegUnsignedImm12,
diff --git a/internal/engine/wazevo/backend/isa/arm64/instr.go b/internal/engine/wazevo/backend/isa/arm64/instr.go
index 7961a121..fa953850 100644
--- a/internal/engine/wazevo/backend/isa/arm64/instr.go
+++ b/internal/engine/wazevo/backend/isa/arm64/instr.go
@@ -21,14 +21,14 @@ type (
 	//
 	// TODO: optimize the layout later once the impl settles.
 	instruction struct {
-		kind               instructionKind
-		prev, next         *instruction
-		u1, u2, u3         uint64
-		rd, rm, rn, ra     operand
-		amode              addressMode
-		abi                *abiImpl
-		targets            []uint32
-		addedAfterLowering bool
+		kind                instructionKind
+		prev, next          *instruction
+		u1, u2, u3          uint64
+		rd, rm, rn, ra      operand
+		amode               addressMode
+		abi                 *abiImpl
+		targets             []uint32
+		addedBeforeRegAlloc bool
 	}
 
 	// instructionKind represents the kind of instruction.
diff --git a/internal/engine/wazevo/backend/isa/arm64/machine.go b/internal/engine/wazevo/backend/isa/arm64/machine.go
index 5c21c59b..86c2e394 100644
--- a/internal/engine/wazevo/backend/isa/arm64/machine.go
+++ b/internal/engine/wazevo/backend/isa/arm64/machine.go
@@ -81,6 +81,8 @@ type (
 
 		maxRequiredStackSizeForCalls int64
 		stackBoundsCheckDisabled     bool
+
+		regAllocStarted bool
 	}
 
 	addend32 struct {
@@ -123,6 +125,7 @@ func NewBackend() backend.Machine {
 
 // Reset implements backend.Machine.
 func (m *machine) Reset() {
+	m.regAllocStarted = false
 	m.instrPool.Reset()
 	m.labelPositionPool.Reset()
 	m.currentSSABlk = nil
@@ -275,13 +278,9 @@ func (l label) String() string {
 func (m *machine) allocateInstr() *instruction {
 	instr := m.instrPool.Allocate()
 	*instr = instruction{}
-	return instr
-}
-
-// allocateInstrAfterLowering allocates an instruction that is added after lowering.
-func (m *machine) allocateInstrAfterLowering() *instruction {
-	instr := m.instrPool.Allocate()
-	instr.addedAfterLowering = true
+	if !m.regAllocStarted {
+		instr.addedBeforeRegAlloc = true
+	}
 	return instr
 }
 
diff --git a/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue.go b/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue.go
index 64410aa1..b6c02a35 100644
--- a/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue.go
+++ b/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue.go
@@ -47,7 +47,7 @@ func (m *machine) SetupPrologue() {
 			panic(fmt.Errorf("BUG: spill slot size %d is not 16-byte aligned", size))
 		}
 
-		cur = m.addsAddOrSubStackPointer(cur, spVReg, size, false, true)
+		cur = m.addsAddOrSubStackPointer(cur, spVReg, size, false)
 
 		// At this point, the stack looks like:
 		//
@@ -105,7 +105,7 @@ func (m *machine) SetupPrologue() {
 		)
 		for _, vr := range regs {
 			// TODO: pair stores to reduce the number of instructions.
-			store := m.allocateInstrAfterLowering()
+			store := m.allocateInstr()
 			store.asStore(operandNR(vr), _amode, regTypeToRegisterSizeInBits(vr.RegType()))
 			cur = linkInstr(cur, store)
 		}
@@ -156,7 +156,7 @@ func (m *machine) createReturnAddrAndSizeOfArgRetSlot(cur *instruction) *instruc
 	} else {
 		sizeOfArgRetReg = xzrVReg
 	}
-	pstr := m.allocateInstrAfterLowering()
+	pstr := m.allocateInstr()
 	amode := addressModePreOrPostIndex(spVReg, -16, true /* decrement before store */)
 	pstr.asStorePair64(lrVReg, sizeOfArgRetReg, amode)
 	cur = linkInstr(cur, pstr)
@@ -175,7 +175,7 @@ func (m *machine) createFrameSizeSlot(cur *instruction, s int64) *instruction {
 		-16,  // stack pointer must be 16-byte aligned.
 		true, // Decrement before store.
 	)
-	store := m.allocateInstrAfterLowering()
+	store := m.allocateInstr()
 	store.asStore(operandNR(frameSizeReg), _amode, 64)
 	cur = linkInstr(cur, store)
 	return cur
@@ -206,7 +206,7 @@ func (m *machine) setupEpilogueAfter(cur *instruction) {
 	prevNext := cur.next
 
 	// We've stored the frame size in the prologue, and now that we are about to return from this function, we won't need it anymore.
-	cur = m.addsAddOrSubStackPointer(cur, spVReg, 16, true, true)
+	cur = m.addsAddOrSubStackPointer(cur, spVReg, 16, true)
 
 	// First we need to restore the clobbered registers.
 	if len(m.clobberedRegs) > 0 {
@@ -237,7 +237,7 @@ func (m *machine) setupEpilogueAfter(cur *instruction) {
 		l := len(m.clobberedRegs) - 1
 		for i := range m.clobberedRegs {
 			vr := m.clobberedRegs[l-i] // reverse order to restore.
-			load := m.allocateInstrAfterLowering()
+			load := m.allocateInstr()
 			amode := addressModePreOrPostIndex(spVReg,
 				16,    // stack pointer must be 16-byte aligned.
 				false, // Increment after store.
@@ -276,7 +276,7 @@ func (m *machine) setupEpilogueAfter(cur *instruction) {
 		//   SP---> +-----------------+
 		//             (low address)
 		//
-		cur = m.addsAddOrSubStackPointer(cur, spVReg, s, true, true)
+		cur = m.addsAddOrSubStackPointer(cur, spVReg, s, true)
 	}
 
 	// Reload the return address (lr).
@@ -294,7 +294,7 @@ func (m *machine) setupEpilogueAfter(cur *instruction) {
 	//            |  ReturnAddress  |
 	//    SP----> +-----------------+
 
-	ldr := m.allocateInstrAfterLowering()
+	ldr := m.allocateInstr()
 	amode := addressModePreOrPostIndex(spVReg, 16 /* stack pointer must be 16-byte aligned. */, false /* increment after loads */)
 	ldr.asULoad(operandNR(lrVReg), amode, 64)
 	cur = linkInstr(cur, ldr)
@@ -324,14 +324,14 @@ func (m *machine) insertStackBoundsCheck(requiredStackSize int64, cur *instructi
 
 	if immm12op, ok := asImm12Operand(uint64(requiredStackSize)); ok {
 		// sub tmp, sp, #requiredStackSize
-		sub := m.allocateInstrAfterLowering()
+		sub := m.allocateInstr()
 		sub.asALU(aluOpSub, operandNR(tmpRegVReg), operandNR(spVReg), immm12op, true)
 		cur = linkInstr(cur, sub)
 	} else {
 		// This case, we first load the requiredStackSize into the temporary register,
 		cur = m.lowerConstantI64AndInsert(cur, tmpRegVReg, requiredStackSize)
 		// Then subtract it.
-		sub := m.allocateInstrAfterLowering()
+		sub := m.allocateInstr()
 		sub.asALU(aluOpSub, operandNR(tmpRegVReg), operandNR(spVReg), operandNR(tmpRegVReg), true)
 		cur = linkInstr(cur, sub)
 	}
@@ -339,7 +339,7 @@ func (m *machine) insertStackBoundsCheck(requiredStackSize int64, cur *instructi
 	tmp2 := x11VReg // Callee save, so it is safe to use it here in the prologue.
 
 	// ldr tmp2, [executionContext #StackBottomPtr]
-	ldr := m.allocateInstrAfterLowering()
+	ldr := m.allocateInstr()
 	ldr.asULoad(operandNR(tmp2), addressMode{
 		kind: addressModeKindRegUnsignedImm12,
 		rn:   x0VReg, // execution context is always the first argument.
@@ -348,7 +348,7 @@ func (m *machine) insertStackBoundsCheck(requiredStackSize int64, cur *instructi
 	cur = linkInstr(cur, ldr)
 
 	// subs xzr, tmp, tmp2
-	subs := m.allocateInstrAfterLowering()
+	subs := m.allocateInstr()
 	subs.asALU(aluOpSubS, operandNR(xzrVReg), operandNR(tmpRegVReg), operandNR(tmp2), true)
 	cur = linkInstr(cur, subs)
 
@@ -361,7 +361,7 @@ func (m *machine) insertStackBoundsCheck(requiredStackSize int64, cur *instructi
 	{
 		// First load the requiredStackSize into the temporary register,
 		cur = m.lowerConstantI64AndInsert(cur, tmpRegVReg, requiredStackSize)
-		setRequiredStackSize := m.allocateInstrAfterLowering()
+		setRequiredStackSize := m.allocateInstr()
 		setRequiredStackSize.asStore(operandNR(tmpRegVReg),
 			addressMode{
 				kind: addressModeKindRegUnsignedImm12,
@@ -429,9 +429,9 @@ func (m *machine) CompileStackGrowCallSequence() []byte {
 	return m.compiler.Buf()
 }
 
-func (m *machine) addsAddOrSubStackPointer(cur *instruction, rd regalloc.VReg, diff int64, add bool, afterLowering bool) *instruction {
+func (m *machine) addsAddOrSubStackPointer(cur *instruction, rd regalloc.VReg, diff int64, add bool) *instruction {
 	m.pendingInstructions = m.pendingInstructions[:0]
-	m.insertAddOrSubStackPointer(rd, diff, add, afterLowering)
+	m.insertAddOrSubStackPointer(rd, diff, add)
 	for _, inserted := range m.pendingInstructions {
 		cur = linkInstr(cur, inserted)
 	}
diff --git a/internal/engine/wazevo/backend/isa/arm64/machine_regalloc.go b/internal/engine/wazevo/backend/isa/arm64/machine_regalloc.go
index 51047034..76349491 100644
--- a/internal/engine/wazevo/backend/isa/arm64/machine_regalloc.go
+++ b/internal/engine/wazevo/backend/isa/arm64/machine_regalloc.go
@@ -163,8 +163,8 @@ func (r *regAllocBlockImpl) InstrIteratorNext() regalloc.Instr {
 		instr := r.instrIteratorNext()
 		if instr == nil {
 			return nil
-		} else if !instr.i.addedAfterLowering {
-			// Skips the instruction added after lowering.
+		} else if instr.i.addedBeforeRegAlloc {
+			// Only concerned about the instruction added before regalloc.
 			return instr
 		}
 	}
@@ -243,6 +243,7 @@ func (m *machine) RegisterInfo(debug bool) *regalloc.RegisterInfo {
 
 // Function implements backend.Machine Function.
 func (m *machine) Function() regalloc.Function {
+	m.regAllocStarted = true
 	return &m.regAllocFn
 }
 
@@ -288,7 +289,7 @@ func (m *machine) insertStoreRegisterAt(v regalloc.VReg, instr *instruction, aft
 	offsetFromSP := m.getVRegSpillSlotOffsetFromSP(v.ID(), typ.Size())
 	var amode addressMode
 	cur, amode = m.resolveAddressModeForOffsetAndInsert(cur, offsetFromSP, typ.Bits(), spVReg)
-	store := m.allocateInstrAfterLowering()
+	store := m.allocateInstr()
 	store.asStore(operandNR(v), amode, typ.Bits())
 
 	cur = linkInstr(cur, store)
@@ -312,7 +313,7 @@ func (m *machine) insertReloadRegisterAt(v regalloc.VReg, instr *instruction, af
 	offsetFromSP := m.getVRegSpillSlotOffsetFromSP(v.ID(), typ.Size())
 	var amode addressMode
 	cur, amode = m.resolveAddressModeForOffsetAndInsert(cur, offsetFromSP, typ.Bits(), spVReg)
-	load := m.allocateInstrAfterLowering()
+	load := m.allocateInstr()
 	switch typ {
 	case ssa.TypeI32, ssa.TypeI64:
 		load.asULoad(operandNR(v), amode, typ.Bits())


Test patch:
diff --git a/internal/engine/wazevo/backend/isa/arm64/abi_test.go b/internal/engine/wazevo/backend/isa/arm64/abi_test.go
index f100ea77..c17daabe 100644
--- a/internal/engine/wazevo/backend/isa/arm64/abi_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/abi_test.go
@@ -238,23 +238,23 @@ func TestAbiImpl_init(t *testing.T) {
 func TestMachine_insertAddOrSubStackPointer(t *testing.T) {
 	t.Run("add/small", func(t *testing.T) {
 		_, _, m := newSetupWithMockContext()
-		m.insertAddOrSubStackPointer(spVReg, 0x10, true, false)
+		m.insertAddOrSubStackPointer(spVReg, 0x10, true)
 		require.Equal(t, `add sp, sp, #0x10`, formatEmittedInstructionsInCurrentBlock(m))
 	})
 	t.Run("add/large stack", func(t *testing.T) {
 		_, _, m := newSetupWithMockContext()
-		m.insertAddOrSubStackPointer(spVReg, 0xffffffff_8, true, false)
+		m.insertAddOrSubStackPointer(spVReg, 0xffffffff_8, true)
 		require.Equal(t, `orr x27, xzr, #0xffffffff8
 add sp, sp, x27`, formatEmittedInstructionsInCurrentBlock(m))
 	})
 	t.Run("sub/small", func(t *testing.T) {
 		_, _, m := newSetupWithMockContext()
-		m.insertAddOrSubStackPointer(spVReg, 0x10, false, false)
+		m.insertAddOrSubStackPointer(spVReg, 0x10, false)
 		require.Equal(t, `sub sp, sp, #0x10`, formatEmittedInstructionsInCurrentBlock(m))
 	})
 	t.Run("sub/large stack", func(t *testing.T) {
 		_, _, m := newSetupWithMockContext()
-		m.insertAddOrSubStackPointer(spVReg, 0xffffffff_8, false, false)
+		m.insertAddOrSubStackPointer(spVReg, 0xffffffff_8, false)
 		require.Equal(t, `orr x27, xzr, #0xffffffff8
 sub sp, sp, x27`, formatEmittedInstructionsInCurrentBlock(m))
 	})
diff --git a/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue_test.go b/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue_test.go
index 75ddbf24..f2899bac 100644
--- a/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue_test.go
@@ -210,7 +210,7 @@ func TestMachine_insertStackBoundsCheck(t *testing.T) {
 		tc := tc
 		t.Run(tc.exp, func(t *testing.T) {
 			ctx, _, m := newSetupWithMockContext()
-			m.rootInstr = m.allocateInstrAfterLowering()
+			m.rootInstr = m.allocateInstr()
 			m.rootInstr.asNop0()
 			m.insertStackBoundsCheck(tc.requiredStackSize, m.rootInstr)
 			fmt.Println(m.Format())