wazevo: adds ABIs for amd64 (#1901)

Signed-off-by: Takeshi Yoneda <[email protected]>

internal/engine/wazevo/backend/abi.go

@@ -7,15 +7,24 @@ import (
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/ssa"
 )
 
-// FunctionABI represents an ABI for the specific target combined with the function signature.
-type FunctionABI interface {
-	// CalleeGenFunctionArgsToVRegs generates instructions to move arguments to virtual registers.
-	CalleeGenFunctionArgsToVRegs(regs []ssa.Value)
-	// CalleeGenVRegsToFunctionReturns generates instructions to move virtual registers to a return value locations.
-	CalleeGenVRegsToFunctionReturns(regs []ssa.Value)
+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
+		Initialized bool
+
+		Args, Rets                 []ABIArg
+		ArgStackSize, RetStackSize int64
+
+		ArgRealRegs []regalloc.VReg
+		RetRealRegs []regalloc.VReg
+	}
+
 	// ABIArg represents either argument or return value's location.
 	ABIArg struct {
 		// Index is the index of the argument.
@@ -59,3 +68,86 @@ func (a ABIArgKind) String() string {
 		panic("BUG")
 	}
 }
+
+// Init initializes the abiImpl for the given signature.
+func (a *FunctionABI[M]) Init(sig *ssa.Signature) {
+	argInts, argFloats, resultInts, resultFloats := a.r.ArgsResultsRegs()
+
+	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)
+	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)
+
+	// Gather the real registers usages in arg/return.
+	a.RetRealRegs = a.RetRealRegs[:0]
+	for i := range a.Rets {
+		r := &a.Rets[i]
+		if r.Kind == ABIArgKindReg {
+			a.RetRealRegs = append(a.RetRealRegs, r.Reg)
+		}
+	}
+	a.ArgRealRegs = a.ArgRealRegs[:0]
+	for i := range a.Args {
+		arg := &a.Args[i]
+		if arg.Kind == ABIArgKindReg {
+			reg := arg.Reg
+			a.ArgRealRegs = append(a.ArgRealRegs, reg)
+		}
+	}
+
+	a.Initialized = true
+}
+
+// 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) {
+	il, fl := len(ints), len(floats)
+
+	var stackOffset int64
+	intParamIndex, floatParamIndex := 0, 0
+	for i, typ := range types {
+		arg := &s[i]
+		arg.Index = i
+		arg.Type = typ
+		if typ.IsInt() {
+			if intParamIndex >= il {
+				arg.Kind = ABIArgKindStack
+				const slotSize = 8 // Align 8 bytes.
+				arg.Offset = stackOffset
+				stackOffset += slotSize
+			} else {
+				arg.Kind = ABIArgKindReg
+				arg.Reg = regalloc.FromRealReg(ints[intParamIndex], regalloc.RegTypeInt)
+				intParamIndex++
+			}
+		} else {
+			if floatParamIndex >= fl {
+				arg.Kind = ABIArgKindStack
+				slotSize := int64(8)   // Align at least 8 bytes.
+				if typ.Bits() == 128 { // Vector.
+					slotSize = 16
+				}
+				arg.Offset = stackOffset
+				stackOffset += slotSize
+			} else {
+				arg.Kind = ABIArgKindReg
+				arg.Reg = regalloc.FromRealReg(floats[floatParamIndex], regalloc.RegTypeFloat)
+				floatParamIndex++
+			}
+		}
+	}
+	return stackOffset
+}
+
+func (a *FunctionABI[M]) AlignedArgResultStackSlotSize() int64 {
+	stackSlotSize := a.RetStackSize + a.ArgStackSize
+	// Align stackSlotSize to 16 bytes.
+	stackSlotSize = (stackSlotSize + 15) &^ 15
+	return stackSlotSize
+}

internal/engine/wazevo/backend/compiler_lower.go

@@ -129,12 +129,12 @@ func (c *compiler) lowerFunctionArguments(entry ssa.BasicBlock) {
 			c.tmpVals = append(c.tmpVals, ssa.ValueInvalid)
 		}
 	}
-	c.mach.ABI().CalleeGenFunctionArgsToVRegs(c.tmpVals)
+	c.mach.LowerParams(c.tmpVals)
 	ectx.FlushPendingInstructions()
 }
 
 func (c *compiler) lowerFunctionReturns(returns []ssa.Value) {
-	c.mach.ABI().CalleeGenVRegsToFunctionReturns(returns)
+	c.mach.LowerReturns(returns)
 }
 
 // lowerBlockArguments lowers how to pass arguments to the given successor block.

internal/engine/wazevo/backend/isa/amd64/abi.go

@@ -0,0 +1,84 @@
+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"
+)
+
+// For the details of the ABI, see:
+// https://github.com/golang/go/blob/49d42128fd8594c172162961ead19ac95e247d24/src/cmd/compile/abi-internal.md#amd64-architecture
+
+var (
+	intArgResultRegs   = []regalloc.RealReg{rax, rcx, rbx, rsi, rdi, r8, r9, r10, r11}
+	floatArgResultRegs = []regalloc.RealReg{xmm0, xmm1, xmm2, xmm3, xmm4, xmm5, xmm6, xmm7}
+)
+
+var regInfo = &regalloc.RegisterInfo{
+	AllocatableRegisters: [regalloc.NumRegType][]regalloc.RealReg{
+		regalloc.RegTypeInt: {
+			rax, rcx, rdx, rbx, rsi, rdi, r8, r9, r10, r11, r12, r13, r14, r15,
+		},
+		regalloc.RegTypeFloat: {
+			xmm0, xmm1, xmm2, xmm3, xmm4, xmm5, xmm6, xmm7, xmm8, xmm9, xmm10, xmm11, xmm12, xmm13, xmm14, xmm15,
+		},
+	},
+	CalleeSavedRegisters: regalloc.NewRegSet(
+		rdx, r12, r13, r14, r15,
+		xmm8, xmm9, xmm10, xmm11, xmm12, xmm13, xmm14, xmm15,
+	),
+	CallerSavedRegisters: regalloc.NewRegSet(
+		rax, rcx, rbx, rsi, rdi, r8, r9, r10, r11,
+		xmm0, xmm1, xmm2, xmm3, xmm4, xmm5, xmm6, xmm7,
+	),
+	RealRegToVReg: []regalloc.VReg{
+		rax: raxVReg, rcx: rcxVReg, rdx: rdxVReg, rbx: rbxVReg, rsp: rspVReg, rbp: rbpVReg, rsi: rsiVReg, rdi: rdiVReg,
+		r8: r8VReg, r9: r9VReg, r10: r10VReg, r11: r11VReg, r12: r12VReg, r13: r13VReg, r14: r14VReg, r15: r15VReg,
+		xmm0: xmm0VReg, xmm1: xmm1VReg, xmm2: xmm2VReg, xmm3: xmm3VReg, xmm4: xmm4VReg, xmm5: xmm5VReg, xmm6: xmm6VReg,
+		xmm7: xmm7VReg, xmm8: xmm8VReg, xmm9: xmm9VReg, xmm10: xmm10VReg, xmm11: xmm11VReg, xmm12: xmm12VReg,
+		xmm13: xmm13VReg, xmm14: xmm14VReg, xmm15: xmm15VReg,
+	},
+	RealRegName: func(r regalloc.RealReg) string { return regNames[r] },
+	RealRegType: func(r regalloc.RealReg) regalloc.RegType {
+		if r < xmm0 {
+			return regalloc.RegTypeInt
+		}
+		return regalloc.RegTypeFloat
+	},
+}
+
+// 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
+}
+
+// LowerParams implements backend.Machine.
+func (m *machine) LowerParams(params []ssa.Value) {
+	// TODO implement me
+	panic("implement me")
+}
+
+// LowerReturns implements backend.Machine.
+func (m *machine) LowerReturns(returns []ssa.Value) {
+	// TODO implement me
+	panic("implement me")
+}

internal/engine/wazevo/backend/isa/amd64/instr.go

@@ -0,0 +1,21 @@
+package amd64
+
+type instruction struct {
+	prev, next *instruction
+}
+
+func resetInstruction(i *instruction) {
+	i.prev = nil
+	i.next = nil
+}
+
+func setNext(i *instruction, next *instruction) {
+	i.next = next
+}
+
+func setPrev(i *instruction, prev *instruction) {
+	i.prev = prev
+}
+
+func asNop(*instruction) {
+}

internal/engine/wazevo/backend/isa/amd64/machine.go

@@ -11,54 +11,49 @@ import (
 
 // NewBackend returns a new backend for arm64.
 func NewBackend() backend.Machine {
-	m := &machine{}
-	return m
+	ectx := backend.NewExecutableContextT[instruction](
+		resetInstruction,
+		setNext,
+		setPrev,
+		asNop,
+	)
+	return &machine{
+		ectx: ectx,
+	}
 }
 
 // machine implements backend.Machine for amd64.
 type machine struct {
+	c                        backend.Compiler
+	ectx                     *backend.ExecutableContextT[instruction]
 	stackBoundsCheckDisabled bool
+
+	// abis maps ssa.SignatureID to the ABI implementation.
+	abis       []abiImpl
+	currentABI *abiImpl
 }
 
-// ExecutableContext implements backend.Machine.
-func (m *machine) ExecutableContext() backend.ExecutableContext {
-	// TODO implement me
-	panic("implement me")
+// Reset implements backend.Machine.
+func (m *machine) Reset() {
+	m.stackBoundsCheckDisabled = false
+	m.ectx.Reset()
 }
 
+// ExecutableContext implements backend.Machine.
+func (m *machine) ExecutableContext() backend.ExecutableContext { return m.ectx }
+
 // DisableStackCheck implements backend.Machine.
-func (m *machine) DisableStackCheck() {
-	m.stackBoundsCheckDisabled = true
-}
+func (m *machine) DisableStackCheck() { m.stackBoundsCheckDisabled = true }
+
+// SetCompiler implements backend.Machine.
+func (m *machine) SetCompiler(compiler backend.Compiler) { m.c = compiler }
 
 // RegisterInfo implements backend.Machine.
-func (m *machine) RegisterInfo() *regalloc.RegisterInfo {
-	// TODO implement me
-	panic("implement me")
-}
+func (m *machine) RegisterInfo() *regalloc.RegisterInfo { return regInfo }
 
 // InitializeABI implements backend.Machine.
 func (m *machine) InitializeABI(sig *ssa.Signature) {
-	// TODO implement me
-	panic("implement me")
-}
-
-// ABI implements backend.Machine.
-func (m *machine) ABI() backend.FunctionABI {
-	// TODO implement me
-	panic("implement me")
-}
-
-// SetCompiler implements backend.Machine.
-func (m *machine) SetCompiler(compiler backend.Compiler) {
-	// TODO implement me
-	panic("implement me")
-}
-
-// StartLoweringFunction implements backend.Machine.
-func (m *machine) StartLoweringFunction(maximumBlockID ssa.BasicBlockID) {
-	// TODO implement me
-	panic("implement me")
+	m.currentABI = m.getOrCreateFunctionABI(sig)
 }
 
 // StartBlock implements backend.Machine.
@@ -85,24 +80,6 @@ func (m *machine) LowerInstr(instruction *ssa.Instruction) {
 	panic("implement me")
 }
 
-// EndBlock implements backend.Machine.
-func (m *machine) EndBlock() {
-	// TODO implement me
-	panic("implement me")
-}
-
-// LinkAdjacentBlocks implements backend.Machine.
-func (m *machine) LinkAdjacentBlocks(prev, next ssa.BasicBlock) {
-	// TODO implement me
-	panic("implement me")
-}
-
-// Reset implements backend.Machine.
-func (m *machine) Reset() {
-	// TODO implement me
-	panic("implement me")
-}
-
 // FlushPendingInstructions implements backend.Machine.
 func (m *machine) FlushPendingInstructions() {
 	// TODO implement me