Repository: tetratelabs/wazero
Pull request #1676 (2023-08-31 05:50:15)

wazevo(arm64): fixes spill/reload logic
#1496 


Patch:
diff --git a/internal/engine/wazevo/backend/isa/arm64/machine_regalloc.go b/internal/engine/wazevo/backend/isa/arm64/machine_regalloc.go
index 5819de40..3940ada0 100644
--- a/internal/engine/wazevo/backend/isa/arm64/machine_regalloc.go
+++ b/internal/engine/wazevo/backend/isa/arm64/machine_regalloc.go
@@ -105,68 +105,27 @@ func (f *regAllocFunctionImpl) ClobberedRegisters(regs []regalloc.VReg) {
 }
 
 // StoreRegisterBefore implements regalloc.Function StoreRegisterBefore.
-func (f *regAllocFunctionImpl) StoreRegisterBefore(v regalloc.VReg, _instr regalloc.Instr) {
-	if !v.IsRealReg() {
-		panic("BUG: VReg must be backed by real reg to be stored")
-	}
-
+func (f *regAllocFunctionImpl) StoreRegisterBefore(v regalloc.VReg, instr regalloc.Instr) {
 	m := f.m
-	typ := m.compiler.TypeOf(v)
-
-	offsetFromSP := m.getVRegSpillSlotOffset(v.ID(), typ.Size()) + m.clobberedRegSlotSize()
-	admode := m.resolveAddressModeForOffset(offsetFromSP, typ.Bits(), spVReg)
-	store := m.allocateInstrAfterLowering()
-	store.asStore(operandNR(v), admode, typ.Bits())
-
-	instr := _instr.(*regAllocInstrImpl).i
-	prev := instr.prev
-
-	prev.next = store
-	store.prev = prev
+	m.insertStoreRegisterAt(v, instr.(*regAllocInstrImpl).i, false)
+}
 
-	store.next = instr
-	instr.prev = store
+// StoreRegisterAfter implements regalloc.Function StoreRegisterAfter.
+func (f *regAllocFunctionImpl) StoreRegisterAfter(v regalloc.VReg, instr regalloc.Instr) {
+	m := f.m
+	m.insertStoreRegisterAt(v, instr.(*regAllocInstrImpl).i, true)
 }
 
 // ReloadRegisterBefore implements regalloc.Function ReloadRegisterBefore.
 func (f *regAllocFunctionImpl) ReloadRegisterBefore(v regalloc.VReg, instr regalloc.Instr) {
-	panic("TODO")
+	m := f.m
+	m.reloadRegister(v, instr.(*regAllocInstrImpl).i, false)
 }
 
 // ReloadRegisterAfter implements regalloc.Function ReloadRegisterAfter.
-func (f *regAllocFunctionImpl) ReloadRegisterAfter(v regalloc.VReg, _instr regalloc.Instr) {
-	if !v.IsRealReg() {
-		panic("BUG: VReg must be backed by real reg to be stored")
-	}
-
+func (f *regAllocFunctionImpl) ReloadRegisterAfter(v regalloc.VReg, instr regalloc.Instr) {
 	m := f.m
-	typ := m.compiler.TypeOf(v)
-
-	offsetFromSP := m.getVRegSpillSlotOffset(v.ID(), typ.Size()) + m.clobberedRegSlotSize()
-	admode := m.resolveAddressModeForOffset(offsetFromSP, typ.Bits(), spVReg)
-	load := m.allocateInstrAfterLowering()
-	switch typ {
-	case ssa.TypeI32, ssa.TypeI64:
-		load.asULoad(operandNR(v), admode, typ.Bits())
-	case ssa.TypeF32, ssa.TypeF64:
-		load.asFpuLoad(operandNR(v), admode, typ.Bits())
-	default:
-		panic("TODO")
-	}
-
-	instr := _instr.(*regAllocInstrImpl).i
-	next := instr.next
-
-	instr.next = load
-	load.prev = instr
-
-	load.next = next
-	next.prev = load
-}
-
-// StoreRegisterAfter implements regalloc.Function StoreRegisterAfter.
-func (f *regAllocFunctionImpl) StoreRegisterAfter(v regalloc.VReg, _instr regalloc.Instr) {
-	panic("TODO")
+	m.reloadRegister(v, instr.(*regAllocInstrImpl).i, true)
 }
 
 // Done implements regalloc.Function Done.
@@ -284,3 +243,80 @@ func (r *regAllocInstrImpl) AssignUses(vs []regalloc.VReg) {
 func (r *regAllocInstrImpl) AssignDef(v regalloc.VReg) {
 	r.i.assignDef(v)
 }
+
+func (m *machine) insertStoreRegisterAt(v regalloc.VReg, instr *instruction, after bool) {
+	if !v.IsRealReg() {
+		panic("BUG: VReg must be backed by real reg to be stored")
+	}
+
+	typ := m.compiler.TypeOf(v)
+
+	offsetFromSP := m.getVRegSpillSlotOffset(v.ID(), typ.Size()) + m.clobberedRegSlotSize()
+	m.pendingInstructions = m.pendingInstructions[:0]
+	admode := m.resolveAddressModeForOffset(offsetFromSP, typ.Bits(), spVReg)
+	store := m.allocateInstrAfterLowering()
+	store.asStore(operandNR(v), admode, typ.Bits())
+
+	var prevNext, cur *instruction
+	if after {
+		cur, prevNext = instr, instr.next
+	} else {
+		cur, prevNext = instr.prev, instr
+	}
+
+	// If the offset is large, we might end up with having multiple instructions inserted in resolveAddressModeForOffset.
+	for _, instr := range m.pendingInstructions {
+		instr.addedAfterLowering = true
+		cur.next = instr
+		instr.prev = cur
+		cur = instr
+	}
+
+	cur.next = store
+	store.prev = cur
+
+	store.next = prevNext
+	prevNext.prev = store
+}
+
+func (m *machine) reloadRegister(v regalloc.VReg, instr *instruction, after bool) {
+	if !v.IsRealReg() {
+		panic("BUG: VReg must be backed by real reg to be stored")
+	}
+
+	typ := m.compiler.TypeOf(v)
+
+	offsetFromSP := m.getVRegSpillSlotOffset(v.ID(), typ.Size()) + m.clobberedRegSlotSize()
+	m.pendingInstructions = m.pendingInstructions[:0]
+	admode := m.resolveAddressModeForOffset(offsetFromSP, typ.Bits(), spVReg)
+	load := m.allocateInstrAfterLowering()
+	switch typ {
+	case ssa.TypeI32, ssa.TypeI64:
+		load.asULoad(operandNR(v), admode, typ.Bits())
+	case ssa.TypeF32, ssa.TypeF64:
+		load.asFpuLoad(operandNR(v), admode, typ.Bits())
+	default:
+		panic("TODO")
+	}
+
+	var prevNext, cur *instruction
+	if after {
+		cur, prevNext = instr, instr.next
+	} else {
+		cur, prevNext = instr.prev, instr
+	}
+
+	// If the offset is large, we might end up with having multiple instructions inserted in resolveAddressModeForOffset.
+	for _, instr := range m.pendingInstructions {
+		instr.addedAfterLowering = true
+		cur.next = instr
+		instr.prev = cur
+		cur = instr
+	}
+
+	cur.next = load
+	load.prev = cur
+
+	load.next = prevNext
+	prevNext.prev = load
+}


Test patch:
diff --git a/internal/engine/wazevo/backend/isa/arm64/machine_regalloc_test.go b/internal/engine/wazevo/backend/isa/arm64/machine_regalloc_test.go
index 2676fdb7..2dcde7b5 100644
--- a/internal/engine/wazevo/backend/isa/arm64/machine_regalloc_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/machine_regalloc_test.go
@@ -122,6 +122,56 @@ func TestRegAllocFunctionImpl_StoreRegisterBefore(t *testing.T) {
 `, m.Format())
 }
 
+func TestMachine_StoreRegisterAfter(t *testing.T) {
+	ctx, _, m := newSetupWithMockContext()
+	m.spillSlotSize = 0xffff
+
+	ctx.typeOf = map[regalloc.VReg]ssa.Type{x1VReg: ssa.TypeI64, v1VReg: ssa.TypeF64}
+	i1, i2 := m.allocateInstr().asUDF(), m.allocateInstr().asUDF()
+	i1.next = i2
+	i2.prev = i1
+
+	f := &regAllocFunctionImpl{m: m}
+	f.StoreRegisterAfter(x1VReg, &regAllocInstrImpl{i: i1})
+	f.StoreRegisterAfter(v1VReg, &regAllocInstrImpl{i: i1})
+
+	m.rootInstr = i1
+	require.Equal(t, `
+	udf
+	movz x27, #0x7, lsl 0
+	movk x27, #0x1, lsl 16
+	str d1, [sp, x27]
+	movz x27, #0xffff, lsl 0
+	str x1, [sp, x27]
+	udf
+`, m.Format())
+}
+
+func TestMachine_ReloadRegisterBefore(t *testing.T) {
+	ctx, _, m := newSetupWithMockContext()
+	m.spillSlotSize = 0xffff
+
+	ctx.typeOf = map[regalloc.VReg]ssa.Type{x1VReg: ssa.TypeI64, v1VReg: ssa.TypeF64}
+	i1, i2 := m.allocateInstr().asUDF(), m.allocateInstr().asUDF()
+	i1.next = i2
+	i2.prev = i1
+
+	f := &regAllocFunctionImpl{m: m}
+	f.ReloadRegisterBefore(x1VReg, &regAllocInstrImpl{i: i2})
+	f.ReloadRegisterBefore(v1VReg, &regAllocInstrImpl{i: i2})
+
+	m.rootInstr = i1
+	require.Equal(t, `
+	udf
+	movz x27, #0xffff, lsl 0
+	ldr x1, [sp, x27]
+	movz x27, #0x7, lsl 0
+	movk x27, #0x1, lsl 16
+	ldr d1, [sp, x27]
+	udf
+`, m.Format())
+}
+
 func TestRegAllocFunctionImpl_ClobberedRegisters(t *testing.T) {
 	_, _, m := newSetupWithMockContext()
 	f := &regAllocFunctionImpl{m: m}