Repository: tetratelabs/wazero
Pull request #1904 (2024-01-08 16:08:38)

wazevo: shares more ABI related codes between archs


Patch:
diff --git a/internal/engine/wazevo/backend/abi.go b/internal/engine/wazevo/backend/abi.go
index 67f464a5..bae83848 100644
--- a/internal/engine/wazevo/backend/abi.go
+++ b/internal/engine/wazevo/backend/abi.go
@@ -7,15 +7,9 @@ import (
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/ssa"
 )
 
-type FunctionABIRegInfo interface {
-	// ArgsResultsRegs returns the registers used for passing parameters.
-	ArgsResultsRegs() (argInts, argFloats, resultInt, resultFloats []regalloc.RealReg)
-}
-
 type (
 	// FunctionABI represents the ABI information for a function which corresponds to a ssa.Signature.
-	FunctionABI[R FunctionABIRegInfo] struct {
-		r           R
+	FunctionABI struct {
 		Initialized bool
 
 		Args, Rets                 []ABIArg
@@ -70,19 +64,17 @@ func (a ABIArgKind) String() string {
 }
 
 // Init initializes the abiImpl for the given signature.
-func (a *FunctionABI[M]) Init(sig *ssa.Signature) {
-	argInts, argFloats, resultInts, resultFloats := a.r.ArgsResultsRegs()
-
+func (a *FunctionABI) Init(sig *ssa.Signature, argResultInts, argResultFloats []regalloc.RealReg) {
 	if len(a.Rets) < len(sig.Results) {
 		a.Rets = make([]ABIArg, len(sig.Results))
 	}
 	a.Rets = a.Rets[:len(sig.Results)]
-	a.RetStackSize = a.setABIArgs(a.Rets, sig.Results, argInts, argFloats)
+	a.RetStackSize = a.setABIArgs(a.Rets, sig.Results, argResultInts, argResultFloats)
 	if argsNum := len(sig.Params); len(a.Args) < argsNum {
 		a.Args = make([]ABIArg, argsNum)
 	}
 	a.Args = a.Args[:len(sig.Params)]
-	a.ArgStackSize = a.setABIArgs(a.Args, sig.Params, resultInts, resultFloats)
+	a.ArgStackSize = a.setABIArgs(a.Args, sig.Params, argResultInts, argResultFloats)
 
 	// Gather the real registers usages in arg/return.
 	a.RetRealRegs = a.RetRealRegs[:0]
@@ -106,7 +98,7 @@ func (a *FunctionABI[M]) Init(sig *ssa.Signature) {
 
 // setABIArgs sets the ABI arguments in the given slice. This assumes that len(s) >= len(types)
 // where if len(s) > len(types), the last elements of s is for the multi-return slot.
-func (a *FunctionABI[M]) setABIArgs(s []ABIArg, types []ssa.Type, ints, floats []regalloc.RealReg) (stackSize int64) {
+func (a *FunctionABI) setABIArgs(s []ABIArg, types []ssa.Type, ints, floats []regalloc.RealReg) (stackSize int64) {
 	il, fl := len(ints), len(floats)
 
 	var stackOffset int64
@@ -145,7 +137,7 @@ func (a *FunctionABI[M]) setABIArgs(s []ABIArg, types []ssa.Type, ints, floats [
 	return stackOffset
 }
 
-func (a *FunctionABI[M]) AlignedArgResultStackSlotSize() int64 {
+func (a *FunctionABI) AlignedArgResultStackSlotSize() int64 {
 	stackSlotSize := a.RetStackSize + a.ArgStackSize
 	// Align stackSlotSize to 16 bytes.
 	stackSlotSize = (stackSlotSize + 15) &^ 15
diff --git a/internal/engine/wazevo/backend/compiler.go b/internal/engine/wazevo/backend/compiler.go
index 87cc0c31..c66b5ed3 100644
--- a/internal/engine/wazevo/backend/compiler.go
+++ b/internal/engine/wazevo/backend/compiler.go
@@ -16,9 +16,12 @@ func NewCompiler(ctx context.Context, mach Machine, builder ssa.Builder) Compile
 }
 
 func newCompiler(_ context.Context, mach Machine, builder ssa.Builder) *compiler {
+	argResultInts, argResultFloats := mach.ArgsResultsRegs()
 	c := &compiler{
 		mach: mach, ssaBuilder: builder,
-		nextVRegID: regalloc.VRegIDNonReservedBegin,
+		nextVRegID:      regalloc.VRegIDNonReservedBegin,
+		argResultInts:   argResultInts,
+		argResultFloats: argResultFloats,
 	}
 	mach.SetCompiler(c)
 	return c
@@ -96,6 +99,9 @@ type Compiler interface {
 
 	// Emit4Bytes appends 4 bytes to the buffer. Used during the code emission.
 	Emit4Bytes(b uint32)
+
+	// GetFunctionABI returns the ABI information for the given signature.
+	GetFunctionABI(sig *ssa.Signature) *FunctionABI
 }
 
 // RelocationInfo represents the relocation information for a call instruction.
@@ -135,6 +141,9 @@ type compiler struct {
 	buf             []byte
 	relocations     []RelocationInfo
 	sourceOffsets   []SourceOffsetInfo
+	// abis maps ssa.SignatureID to the ABI implementation.
+	abis                           []FunctionABI
+	argResultInts, argResultFloats []regalloc.RealReg
 }
 
 // SourceOffsetInfo is a data to associate the source offset with the executable offset.
@@ -380,3 +389,17 @@ func (c *compiler) Emit4Bytes(b uint32) {
 func (c *compiler) Buf() []byte {
 	return c.buf
 }
+
+func (c *compiler) GetFunctionABI(sig *ssa.Signature) *FunctionABI {
+	if int(sig.ID) >= len(c.abis) {
+		c.abis = append(c.abis, make([]FunctionABI, int(sig.ID)+1)...)
+	}
+
+	abi := &c.abis[sig.ID]
+	if abi.Initialized {
+		return abi
+	}
+
+	abi.Init(sig, c.argResultInts, c.argResultFloats)
+	return abi
+}
diff --git a/internal/engine/wazevo/backend/compiler_lower.go b/internal/engine/wazevo/backend/compiler_lower.go
index 676c1a72..e59f8675 100644
--- a/internal/engine/wazevo/backend/compiler_lower.go
+++ b/internal/engine/wazevo/backend/compiler_lower.go
@@ -8,7 +8,7 @@ import (
 // Lower implements Compiler.Lower.
 func (c *compiler) Lower() {
 	c.assignVirtualRegisters()
-	c.mach.InitializeABI(c.ssaBuilder.Signature())
+	c.mach.SetCurrentABI(c.GetFunctionABI(c.ssaBuilder.Signature()))
 	c.mach.ExecutableContext().StartLoweringFunction(c.ssaBuilder.BlockIDMax())
 	c.lowerBlocks()
 }
diff --git a/internal/engine/wazevo/backend/isa/amd64/abi.go b/internal/engine/wazevo/backend/isa/amd64/abi.go
index 9771fb4f..19a29824 100644
--- a/internal/engine/wazevo/backend/isa/amd64/abi.go
+++ b/internal/engine/wazevo/backend/isa/amd64/abi.go
@@ -1,7 +1,6 @@
 package amd64
 
 import (
-	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend"
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend/regalloc"
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/ssa"
 )
@@ -47,28 +46,9 @@ var regInfo = &regalloc.RegisterInfo{
 	},
 }
 
-// functionABIRegInfo implements backend.FunctionABIRegInfo.
-type functionABIRegInfo struct{}
-
-// ArgsResultsRegs implements backend.FunctionABIRegInfo.
-func (f functionABIRegInfo) ArgsResultsRegs() (argInts, argFloats, resultInt, resultFloats []regalloc.RealReg) {
-	return intArgResultRegs, floatArgResultRegs, intArgResultRegs, floatArgResultRegs
-}
-
-type abiImpl = backend.FunctionABI[functionABIRegInfo]
-
-func (m *machine) getOrCreateFunctionABI(sig *ssa.Signature) *abiImpl {
-	if int(sig.ID) >= len(m.abis) {
-		m.abis = append(m.abis, make([]abiImpl, int(sig.ID)+1)...)
-	}
-
-	abi := &m.abis[sig.ID]
-	if abi.Initialized {
-		return abi
-	}
-
-	abi.Init(sig)
-	return abi
+// ArgsResultsRegs implements backend.Machine.
+func (m *machine) ArgsResultsRegs() (argResultInts, argResultFloats []regalloc.RealReg) {
+	return intArgResultRegs, floatArgResultRegs
 }
 
 // LowerParams implements backend.Machine.
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine.go b/internal/engine/wazevo/backend/isa/amd64/machine.go
index 26d20878..9815700e 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine.go
@@ -29,11 +29,8 @@ type machine struct {
 	ectx                     *backend.ExecutableContextT[instruction]
 	stackBoundsCheckDisabled bool
 
-	regAlloc regalloc.Allocator
-
-	// abis maps ssa.SignatureID to the ABI implementation.
-	abis       []abiImpl
-	currentABI *abiImpl
+	regAlloc   regalloc.Allocator
+	currentABI *backend.FunctionABI
 }
 
 // Reset implements backend.Machine.
@@ -51,9 +48,9 @@ func (m *machine) DisableStackCheck() { m.stackBoundsCheckDisabled = true }
 // SetCompiler implements backend.Machine.
 func (m *machine) SetCompiler(compiler backend.Compiler) { m.c = compiler }
 
-// InitializeABI implements backend.Machine.
-func (m *machine) InitializeABI(sig *ssa.Signature) {
-	m.currentABI = m.getOrCreateFunctionABI(sig)
+// SetCurrentABI implements backend.Machine.
+func (m *machine) SetCurrentABI(abi *backend.FunctionABI) {
+	m.currentABI = abi
 }
 
 // LowerSingleBranch implements backend.Machine.
diff --git a/internal/engine/wazevo/backend/isa/arm64/abi.go b/internal/engine/wazevo/backend/isa/arm64/abi.go
index edca84d4..de386771 100644
--- a/internal/engine/wazevo/backend/isa/arm64/abi.go
+++ b/internal/engine/wazevo/backend/isa/arm64/abi.go
@@ -56,28 +56,9 @@ var regInfo = &regalloc.RegisterInfo{
 	},
 }
 
-// functionABIRegInfo implements backend.FunctionABIRegInfo.
-type functionABIRegInfo struct{}
-
-// ArgsResultsRegs implements backend.FunctionABIRegInfo.
-func (f functionABIRegInfo) ArgsResultsRegs() (argInts, argFloats, resultInt, resultFloats []regalloc.RealReg) {
-	return intParamResultRegs, floatParamResultRegs, intParamResultRegs, floatParamResultRegs
-}
-
-type functionABI = backend.FunctionABI[functionABIRegInfo]
-
-func (m *machine) getOrCreateFunctionABI(sig *ssa.Signature) *functionABI {
-	if int(sig.ID) >= len(m.abis) {
-		m.abis = append(m.abis, make([]functionABI, int(sig.ID)+1)...)
-	}
-
-	abi := &m.abis[sig.ID]
-	if abi.Initialized {
-		return abi
-	}
-
-	abi.Init(sig)
-	return abi
+// ArgsResultsRegs implements backend.Machine.
+func (m *machine) ArgsResultsRegs() (argResultInts, argResultFloats []regalloc.RealReg) {
+	return intParamResultRegs, floatParamResultRegs
 }
 
 // LowerParams implements backend.FunctionABI.
@@ -198,7 +179,7 @@ func (m *machine) LowerReturns(rets []ssa.Value) {
 
 // callerGenVRegToFunctionArg is the opposite of GenFunctionArgToVReg, which is used to generate the
 // caller side of the function call.
-func (m *machine) callerGenVRegToFunctionArg(a *functionABI, argIndex int, reg regalloc.VReg, def *backend.SSAValueDefinition, slotBegin int64) {
+func (m *machine) callerGenVRegToFunctionArg(a *backend.FunctionABI, argIndex int, reg regalloc.VReg, def *backend.SSAValueDefinition, slotBegin int64) {
 	arg := &a.Args[argIndex]
 	if def != nil && def.IsFromInstr() {
 		// Constant instructions are inlined.
@@ -220,7 +201,7 @@ func (m *machine) callerGenVRegToFunctionArg(a *functionABI, argIndex int, reg r
 	}
 }
 
-func (m *machine) callerGenFunctionReturnVReg(a *functionABI, retIndex int, reg regalloc.VReg, slotBegin int64) {
+func (m *machine) callerGenFunctionReturnVReg(a *backend.FunctionABI, retIndex int, reg regalloc.VReg, slotBegin int64) {
 	r := &a.Rets[retIndex]
 	if r.Kind == backend.ABIArgKindReg {
 		m.InsertMove(reg, r.Reg, r.Type)
@@ -284,7 +265,7 @@ func (m *machine) lowerCall(si *ssa.Instruction) {
 	} else {
 		indirectCalleePtr, sigID, args = si.CallIndirectData()
 	}
-	calleeABI := m.getOrCreateFunctionABI(m.compiler.SSABuilder().ResolveSignature(sigID))
+	calleeABI := m.compiler.GetFunctionABI(m.compiler.SSABuilder().ResolveSignature(sigID))
 
 	stackSlotSize := calleeABI.AlignedArgResultStackSlotSize()
 	if m.maxRequiredStackSizeForCalls < stackSlotSize+16 {
diff --git a/internal/engine/wazevo/backend/isa/arm64/abi_entry_preamble.go b/internal/engine/wazevo/backend/isa/arm64/abi_entry_preamble.go
index 9bdda1a1..bf72976b 100644
--- a/internal/engine/wazevo/backend/isa/arm64/abi_entry_preamble.go
+++ b/internal/engine/wazevo/backend/isa/arm64/abi_entry_preamble.go
@@ -135,8 +135,9 @@ func (m *machine) goEntryPreamblePassResult(cur *instruction, resultSlicePtr reg
 }
 
 func (m *machine) constructEntryPreamble(sig *ssa.Signature) (root *instruction) {
-	abi := functionABI{}
-	abi.Init(sig)
+	abi := backend.FunctionABI{}
+	m.ArgsResultsRegs()
+	abi.Init(sig, intParamResultRegs, floatParamResultRegs)
 
 	root = m.allocateNop()
 
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 198d2f1b..7b10a821 100644
--- a/internal/engine/wazevo/backend/isa/arm64/abi_go_call.go
+++ b/internal/engine/wazevo/backend/isa/arm64/abi_go_call.go
@@ -20,8 +20,8 @@ func (m *machine) CompileGoFunctionTrampoline(exitCode wazevoapi.ExitCode, sig *
 		argBegin++
 	}
 
-	abi := &functionABI{}
-	abi.Init(sig)
+	abi := &backend.FunctionABI{}
+	abi.Init(sig, intParamResultRegs, floatParamResultRegs)
 	m.currentABI = abi
 
 	cur := m.allocateInstr()
diff --git a/internal/engine/wazevo/backend/isa/arm64/instr.go b/internal/engine/wazevo/backend/isa/arm64/instr.go
index bcd1566f..7f9515fa 100644
--- a/internal/engine/wazevo/backend/isa/arm64/instr.go
+++ b/internal/engine/wazevo/backend/isa/arm64/instr.go
@@ -5,6 +5,7 @@ import (
 	"math"
 	"strings"
 
+	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend"
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend/regalloc"
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/ssa"
 )
@@ -26,7 +27,7 @@ type (
 		u1, u2, u3          uint64
 		rd, rm, rn, ra      operand
 		amode               addressMode
-		abi                 *functionABI
+		abi                 *backend.FunctionABI
 		targets             []uint32
 		addedBeforeRegAlloc bool
 	}
@@ -461,13 +462,13 @@ func (i *instruction) AssignUse(index int, reg regalloc.VReg) {
 	}
 }
 
-func (i *instruction) asCall(ref ssa.FuncRef, abi *functionABI) {
+func (i *instruction) asCall(ref ssa.FuncRef, abi *backend.FunctionABI) {
 	i.kind = call
 	i.u1 = uint64(ref)
 	i.abi = abi
 }
 
-func (i *instruction) asCallIndirect(ptr regalloc.VReg, abi *functionABI) {
+func (i *instruction) asCallIndirect(ptr regalloc.VReg, abi *backend.FunctionABI) {
 	i.kind = callInd
 	i.rn = operandNR(ptr)
 	i.abi = abi
@@ -524,7 +525,7 @@ func (i *instruction) nop0Label() label {
 	return label(i.u1)
 }
 
-func (i *instruction) asRet(abi *functionABI) {
+func (i *instruction) asRet(abi *backend.FunctionABI) {
 	i.kind = ret
 	i.abi = abi
 }
diff --git a/internal/engine/wazevo/backend/isa/arm64/machine.go b/internal/engine/wazevo/backend/isa/arm64/machine.go
index bdf88ae2..ad778552 100644
--- a/internal/engine/wazevo/backend/isa/arm64/machine.go
+++ b/internal/engine/wazevo/backend/isa/arm64/machine.go
@@ -16,9 +16,7 @@ type (
 	machine struct {
 		compiler          backend.Compiler
 		executableContext *backend.ExecutableContextT[instruction]
-		currentABI        *functionABI
-		// abis maps ssa.SignatureID to the ABI implementation.
-		abis []functionABI
+		currentABI        *backend.FunctionABI
 
 		regAlloc   regalloc.Allocator
 		regAllocFn *backend.RegAllocFunction[*instruction, *machine]
@@ -149,9 +147,9 @@ func (m *machine) Reset() {
 	m.executableContext.Reset()
 }
 
-// InitializeABI implements backend.Machine InitializeABI.
-func (m *machine) InitializeABI(sig *ssa.Signature) {
-	m.currentABI = m.getOrCreateFunctionABI(sig)
+// SetCurrentABI implements backend.Machine SetCurrentABI.
+func (m *machine) SetCurrentABI(abi *backend.FunctionABI) {
+	m.currentABI = abi
 }
 
 // DisableStackCheck implements backend.Machine DisableStackCheck.
diff --git a/internal/engine/wazevo/backend/machine.go b/internal/engine/wazevo/backend/machine.go
index 34dad417..8b7b0c47 100644
--- a/internal/engine/wazevo/backend/machine.go
+++ b/internal/engine/wazevo/backend/machine.go
@@ -16,8 +16,8 @@ type (
 		// DisableStackCheck disables the stack check for the current compilation for debugging/testing.
 		DisableStackCheck()
 
-		// InitializeABI initializes the FunctionABI for the given signature.
-		InitializeABI(sig *ssa.Signature)
+		// SetCurrentABI initializes the FunctionABI for the given signature.
+		SetCurrentABI(abi *FunctionABI)
 
 		// SetCompiler sets the compilation context used for the lifetime of Machine.
 		// This is only called once per Machine, i.e. before the first compilation.
@@ -88,5 +88,8 @@ type (
 
 		// LowerReturns lowers the given returns.
 		LowerReturns(returns []ssa.Value)
+
+		// ArgsResultsRegs returns the registers used for arguments and return values.
+		ArgsResultsRegs() (argResultInts, argResultFloats []regalloc.RealReg)
 	}
 )


Test patch:
diff --git a/internal/engine/wazevo/backend/compiler_lower_test.go b/internal/engine/wazevo/backend/compiler_lower_test.go
index bfdc8747..851b4409 100644
--- a/internal/engine/wazevo/backend/compiler_lower_test.go
+++ b/internal/engine/wazevo/backend/compiler_lower_test.go
@@ -50,6 +50,10 @@ func TestCompiler_lowerBlockArguments(t *testing.T) {
 							target regalloc.VReg
 						}{instr: instr, target: vr})
 					},
+					argResultInts: []regalloc.RealReg{regalloc.RealReg(0), regalloc.RealReg(1)},
+					argResultFloats: []regalloc.RealReg{
+						regalloc.RealReg(2), regalloc.RealReg(3), regalloc.RealReg(4), regalloc.RealReg(5),
+					},
 				}
 
 				c := newCompiler(context.Background(), m, builder)
diff --git a/internal/engine/wazevo/backend/isa/arm64/abi_go_call_test.go b/internal/engine/wazevo/backend/isa/arm64/abi_go_call_test.go
index d8b5c927..99a81fca 100644
--- a/internal/engine/wazevo/backend/isa/arm64/abi_go_call_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/abi_go_call_test.go
@@ -623,7 +623,7 @@ func Test_goFunctionCallStoreStackResult(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Run(tc.name, func(t *testing.T) {
 				_, _, m := newSetupWithMockContext()
-				m.currentABI = &functionABI{ArgStackSize: 8}
+				m.currentABI = &backend.FunctionABI{ArgStackSize: 8}
 
 				nop := m.allocateNop()
 				m.goFunctionCallStoreStackResult(nop, spVReg, tc.result, tc.resultReg)
diff --git a/internal/engine/wazevo/backend/isa/arm64/abi_test.go b/internal/engine/wazevo/backend/isa/arm64/abi_test.go
index 9d4b8673..f320f175 100644
--- a/internal/engine/wazevo/backend/isa/arm64/abi_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/abi_test.go
@@ -13,12 +13,12 @@ func TestAbiImpl_init(t *testing.T) {
 	for _, tc := range []struct {
 		name string
 		sig  *ssa.Signature
-		exp  functionABI
+		exp  backend.FunctionABI
 	}{
 		{
 			name: "empty sig",
 			sig:  &ssa.Signature{},
-			exp:  functionABI{},
+			exp:  backend.FunctionABI{},
 		},
 		{
 			name: "small sig",
@@ -26,7 +26,7 @@ func TestAbiImpl_init(t *testing.T) {
 				Params:  []ssa.Type{ssa.TypeI32, ssa.TypeF32, ssa.TypeI32},
 				Results: []ssa.Type{ssa.TypeI64, ssa.TypeF64},
 			},
-			exp: functionABI{
+			exp: backend.FunctionABI{
 				Args: []backend.ABIArg{
 					{Index: 0, Kind: backend.ABIArgKindReg, Reg: x0VReg, Type: ssa.TypeI32},
 					{Index: 1, Kind: backend.ABIArgKindReg, Reg: v0VReg, Type: ssa.TypeF32},
@@ -80,7 +80,7 @@ func TestAbiImpl_init(t *testing.T) {
 					ssa.TypeI64, ssa.TypeF64,
 				},
 			},
-			exp: functionABI{
+			exp: backend.FunctionABI{
 				ArgStackSize: 128, RetStackSize: 128,
 				Args: []backend.ABIArg{
 					{Index: 0, Kind: backend.ABIArgKindReg, Reg: x0VReg, Type: ssa.TypeI32},
@@ -178,7 +178,7 @@ func TestAbiImpl_init(t *testing.T) {
 					ssa.TypeI64, ssa.TypeF64,
 				},
 			},
-			exp: functionABI{
+			exp: backend.FunctionABI{
 				Args: []backend.ABIArg{
 					{Index: 0, Kind: backend.ABIArgKindReg, Reg: x0VReg, Type: ssa.TypeI32},
 					{Index: 1, Kind: backend.ABIArgKindReg, Reg: v0VReg, Type: ssa.TypeF32},
@@ -222,8 +222,8 @@ func TestAbiImpl_init(t *testing.T) {
 	} {
 		tc := tc
 		t.Run(tc.name, func(t *testing.T) {
-			abi := functionABI{}
-			abi.Init(tc.sig)
+			abi := backend.FunctionABI{}
+			abi.Init(tc.sig, intParamResultRegs, floatParamResultRegs)
 			require.Equal(t, tc.exp.Args, abi.Args)
 			require.Equal(t, tc.exp.Rets, abi.Rets)
 			require.Equal(t, tc.exp.ArgStackSize, abi.ArgStackSize)
@@ -264,7 +264,8 @@ func TestAbiImpl_callerGenVRegToFunctionArg_constant_inlining(t *testing.T) {
 
 	i64 := builder.AllocateInstruction().AsIconst64(10).Insert(builder)
 	f64 := builder.AllocateInstruction().AsF64const(3.14).Insert(builder)
-	abi := m.getOrCreateFunctionABI(&ssa.Signature{Params: []ssa.Type{ssa.TypeI64, ssa.TypeF64}})
+	abi := &backend.FunctionABI{}
+	abi.Init(&ssa.Signature{Params: []ssa.Type{ssa.TypeI64, ssa.TypeF64}}, intParamResultRegs, floatParamResultRegs)
 	m.callerGenVRegToFunctionArg(abi, 0, regalloc.VReg(100).SetRegType(regalloc.RegTypeInt), &backend.SSAValueDefinition{Instr: i64, RefCount: 1}, 0)
 	m.callerGenVRegToFunctionArg(abi, 1, regalloc.VReg(50).SetRegType(regalloc.RegTypeFloat), &backend.SSAValueDefinition{Instr: f64, RefCount: 1}, 0)
 	require.Equal(t, `movz x100?, #0xa, lsl 0
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 26c6ab1c..13174407 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
@@ -3,6 +3,7 @@ package arm64
 import (
 	"testing"
 
+	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend"
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend/regalloc"
 	"github.com/tetratelabs/wazero/internal/testing/require"
 )
@@ -12,7 +13,7 @@ func TestMachine_SetupPrologue(t *testing.T) {
 		spillSlotSize int64
 		clobberedRegs []regalloc.VReg
 		exp           string
-		abi           functionABI
+		abi           backend.FunctionABI
 	}{
 		{
 			spillSlotSize: 0,
@@ -24,7 +25,7 @@ func TestMachine_SetupPrologue(t *testing.T) {
 		},
 		{
 			spillSlotSize: 0,
-			abi:           functionABI{ArgStackSize: 16, RetStackSize: 16},
+			abi:           backend.FunctionABI{ArgStackSize: 16, RetStackSize: 16},
 			exp: `
 	orr x27, xzr, #0x20
 	sub sp, sp, x27
@@ -74,7 +75,7 @@ func TestMachine_SetupPrologue(t *testing.T) {
 		},
 		{
 			spillSlotSize: 320,
-			abi:           functionABI{ArgStackSize: 320, RetStackSize: 160},
+			abi:           backend.FunctionABI{ArgStackSize: 320, RetStackSize: 160},
 			clobberedRegs: []regalloc.VReg{v18VReg, v19VReg, x18VReg, x25VReg},
 			exp: `
 	orr x27, xzr, #0x1e0
@@ -117,7 +118,7 @@ func TestMachine_SetupPrologue(t *testing.T) {
 func TestMachine_SetupEpilogue(t *testing.T) {
 	for _, tc := range []struct {
 		exp           string
-		abi           functionABI
+		abi           backend.FunctionABI
 		clobberedRegs []regalloc.VReg
 		spillSlotSize int64
 	}{
@@ -148,7 +149,7 @@ func TestMachine_SetupEpilogue(t *testing.T) {
 	add sp, sp, #0x20
 	ret
 `,
-			abi:           functionABI{ArgStackSize: 16, RetStackSize: 16},
+			abi:           backend.FunctionABI{ArgStackSize: 16, RetStackSize: 16},
 			spillSlotSize: 16 * 5,
 			clobberedRegs: nil,
 		},
@@ -201,7 +202,7 @@ func TestMachine_SetupEpilogue(t *testing.T) {
 	ret
 `,
 			spillSlotSize: 16 * 10,
-			abi:           functionABI{ArgStackSize: 16, RetStackSize: 320},
+			abi:           backend.FunctionABI{ArgStackSize: 16, RetStackSize: 320},
 			clobberedRegs: []regalloc.VReg{v18VReg, v27VReg, x18VReg, x25VReg},
 		},
 	} {
diff --git a/internal/engine/wazevo/backend/isa/arm64/machine_test.go b/internal/engine/wazevo/backend/isa/arm64/machine_test.go
index 8d3699bd..a585f1fb 100644
--- a/internal/engine/wazevo/backend/isa/arm64/machine_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/machine_test.go
@@ -3,6 +3,7 @@ package arm64
 import (
 	"testing"
 
+	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend"
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend/regalloc"
 	"github.com/tetratelabs/wazero/internal/testing/require"
 )
@@ -79,7 +80,7 @@ func TestMachine_arg0OffsetFromSP(t *testing.T) {
 func TestMachine_ret0OffsetFromSP(t *testing.T) {
 	m := &machine{
 		clobberedRegs: make([]regalloc.VReg, 10), spillSlotSize: 16 * 8,
-		currentABI: &functionABI{ArgStackSize: 180},
+		currentABI: &backend.FunctionABI{ArgStackSize: 180},
 	}
 	require.Equal(t, int64(16*18)+32+180, m.ret0OffsetFromSP())
 }
diff --git a/internal/engine/wazevo/backend/isa/arm64/util_test.go b/internal/engine/wazevo/backend/isa/arm64/util_test.go
index a033db70..869495f1 100644
--- a/internal/engine/wazevo/backend/isa/arm64/util_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/util_test.go
@@ -61,6 +61,11 @@ type mockCompiler struct {
 	buf         []byte
 }
 
+func (m *mockCompiler) GetFunctionABI(sig *ssa.Signature) *backend.FunctionABI {
+	// TODO implement me
+	panic("implement me")
+}
+
 func (m *mockCompiler) SSABuilder() ssa.Builder { return nil }
 
 func (m *mockCompiler) LoopNestingForestRoots() []ssa.BasicBlock { panic("TODO") }
diff --git a/internal/engine/wazevo/backend/machine_test.go b/internal/engine/wazevo/backend/machine_test.go
index 45ed0f22..e15185d0 100644
--- a/internal/engine/wazevo/backend/machine_test.go
+++ b/internal/engine/wazevo/backend/machine_test.go
@@ -10,19 +10,24 @@ import (
 
 // mockMachine implements Machine for testing.
 type mockMachine struct {
-	startLoweringFunction  func(id ssa.BasicBlockID)
-	startBlock             func(block ssa.BasicBlock)
-	lowerSingleBranch      func(b *ssa.Instruction)
-	lowerConditionalBranch func(b *ssa.Instruction)
-	lowerInstr             func(instruction *ssa.Instruction)
-	endBlock               func()
-	endLoweringFunction    func()
-	reset                  func()
-	insertMove             func(dst, src regalloc.VReg)
-	insertLoadConstant     func(instr *ssa.Instruction, vr regalloc.VReg)
-	format                 func() string
-	linkAdjacentBlocks     func(prev, next ssa.BasicBlock)
-	rinfo                  *regalloc.RegisterInfo
+	argResultInts, argResultFloats []regalloc.RealReg
+	startLoweringFunction          func(id ssa.BasicBlockID)
+	startBlock                     func(block ssa.BasicBlock)
+	lowerSingleBranch              func(b *ssa.Instruction)
+	lowerConditionalBranch         func(b *ssa.Instruction)
+	lowerInstr                     func(instruction *ssa.Instruction)
+	endBlock                       func()
+	endLoweringFunction            func()
+	reset                          func()
+	insertMove                     func(dst, src regalloc.VReg)
+	insertLoadConstant             func(instr *ssa.Instruction, vr regalloc.VReg)
+	format                         func() string
+	linkAdjacentBlocks             func(prev, next ssa.BasicBlock)
+	rinfo                          *regalloc.RegisterInfo
+}
+
+func (m mockMachine) ArgsResultsRegs() (argResultInts, argResultFloats []regalloc.RealReg) {
+	return m.argResultInts, m.argResultFloats
 }
 
 func (m mockMachine) RegAlloc() { panic("implement me") }
@@ -78,8 +83,8 @@ func (m mockMachine) InsertReturn() { panic("TODO") }
 // LinkAdjacentBlocks implements Machine.LinkAdjacentBlocks.
 func (m mockMachine) LinkAdjacentBlocks(prev, next ssa.BasicBlock) { m.linkAdjacentBlocks(prev, next) }
 
-// InitializeABI implements Machine.InitializeABI.
-func (m mockMachine) InitializeABI(*ssa.Signature) {}
+// SetCurrentABI implements Machine.SetCurrentABI.
+func (m mockMachine) SetCurrentABI(*FunctionABI) {}
 
 // SetCompiler implements Machine.SetCompiler.
 func (m mockMachine) SetCompiler(Compiler) {}