abi.go

  1package backend
  2
  3import (
  4	"fmt"
  5
  6	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend/regalloc"
  7	"github.com/tetratelabs/wazero/internal/engine/wazevo/ssa"
  8)
  9
 10type (
 11	// FunctionABI represents the ABI information for a function which corresponds to a ssa.Signature.
 12	FunctionABI struct {
 13		Initialized bool
 14
 15		Args, Rets                 []ABIArg
 16		ArgStackSize, RetStackSize int64
 17
 18		ArgIntRealRegs   byte
 19		ArgFloatRealRegs byte
 20		RetIntRealRegs   byte
 21		RetFloatRealRegs byte
 22	}
 23
 24	// ABIArg represents either argument or return value's location.
 25	ABIArg struct {
 26		// Index is the index of the argument.
 27		Index int
 28		// Kind is the kind of the argument.
 29		Kind ABIArgKind
 30		// Reg is valid if Kind == ABIArgKindReg.
 31		// This VReg must be based on RealReg.
 32		Reg regalloc.VReg
 33		// Offset is valid if Kind == ABIArgKindStack.
 34		// This is the offset from the beginning of either arg or ret stack slot.
 35		Offset int64
 36		// Type is the type of the argument.
 37		Type ssa.Type
 38	}
 39
 40	// ABIArgKind is the kind of ABI argument.
 41	ABIArgKind byte
 42)
 43
 44const (
 45	// ABIArgKindReg represents an argument passed in a register.
 46	ABIArgKindReg = iota
 47	// ABIArgKindStack represents an argument passed in the stack.
 48	ABIArgKindStack
 49)
 50
 51// String implements fmt.Stringer.
 52func (a *ABIArg) String() string {
 53	return fmt.Sprintf("args[%d]: %s", a.Index, a.Kind)
 54}
 55
 56// String implements fmt.Stringer.
 57func (a ABIArgKind) String() string {
 58	switch a {
 59	case ABIArgKindReg:
 60		return "reg"
 61	case ABIArgKindStack:
 62		return "stack"
 63	default:
 64		panic("BUG")
 65	}
 66}
 67
 68// Init initializes the abiImpl for the given signature.
 69func (a *FunctionABI) Init(sig *ssa.Signature, argResultInts, argResultFloats []regalloc.RealReg) {
 70	if len(a.Rets) < len(sig.Results) {
 71		a.Rets = make([]ABIArg, len(sig.Results))
 72	}
 73	a.Rets = a.Rets[:len(sig.Results)]
 74	a.RetStackSize = a.setABIArgs(a.Rets, sig.Results, argResultInts, argResultFloats)
 75	if argsNum := len(sig.Params); len(a.Args) < argsNum {
 76		a.Args = make([]ABIArg, argsNum)
 77	}
 78	a.Args = a.Args[:len(sig.Params)]
 79	a.ArgStackSize = a.setABIArgs(a.Args, sig.Params, argResultInts, argResultFloats)
 80
 81	// Gather the real registers usages in arg/return.
 82	a.ArgIntRealRegs, a.ArgFloatRealRegs = 0, 0
 83	a.RetIntRealRegs, a.RetFloatRealRegs = 0, 0
 84	for i := range a.Rets {
 85		r := &a.Rets[i]
 86		if r.Kind == ABIArgKindReg {
 87			if r.Type.IsInt() {
 88				a.RetIntRealRegs++
 89			} else {
 90				a.RetFloatRealRegs++
 91			}
 92		}
 93	}
 94	for i := range a.Args {
 95		arg := &a.Args[i]
 96		if arg.Kind == ABIArgKindReg {
 97			if arg.Type.IsInt() {
 98				a.ArgIntRealRegs++
 99			} else {
100				a.ArgFloatRealRegs++
101			}
102		}
103	}
104
105	a.Initialized = true
106}
107
108// setABIArgs sets the ABI arguments in the given slice. This assumes that len(s) >= len(types)
109// where if len(s) > len(types), the last elements of s is for the multi-return slot.
110func (a *FunctionABI) setABIArgs(s []ABIArg, types []ssa.Type, ints, floats []regalloc.RealReg) (stackSize int64) {
111	il, fl := len(ints), len(floats)
112
113	var stackOffset int64
114	intParamIndex, floatParamIndex := 0, 0
115	for i, typ := range types {
116		arg := &s[i]
117		arg.Index = i
118		arg.Type = typ
119		if typ.IsInt() {
120			if intParamIndex >= il {
121				arg.Kind = ABIArgKindStack
122				const slotSize = 8 // Align 8 bytes.
123				arg.Offset = stackOffset
124				stackOffset += slotSize
125			} else {
126				arg.Kind = ABIArgKindReg
127				arg.Reg = regalloc.FromRealReg(ints[intParamIndex], regalloc.RegTypeInt)
128				intParamIndex++
129			}
130		} else {
131			if floatParamIndex >= fl {
132				arg.Kind = ABIArgKindStack
133				slotSize := int64(8)   // Align at least 8 bytes.
134				if typ.Bits() == 128 { // Vector.
135					slotSize = 16
136				}
137				arg.Offset = stackOffset
138				stackOffset += slotSize
139			} else {
140				arg.Kind = ABIArgKindReg
141				arg.Reg = regalloc.FromRealReg(floats[floatParamIndex], regalloc.RegTypeFloat)
142				floatParamIndex++
143			}
144		}
145	}
146	return stackOffset
147}
148
149func (a *FunctionABI) AlignedArgResultStackSlotSize() uint32 {
150	stackSlotSize := a.RetStackSize + a.ArgStackSize
151	// Align stackSlotSize to 16 bytes.
152	stackSlotSize = (stackSlotSize + 15) &^ 15
153	// Check overflow 32-bit.
154	if stackSlotSize > 0xFFFFFFFF {
155		panic("ABI stack slot size overflow")
156	}
157	return uint32(stackSlotSize)
158}
159
160func (a *FunctionABI) ABIInfoAsUint64() uint64 {
161	return uint64(a.ArgIntRealRegs)<<56 |
162		uint64(a.ArgFloatRealRegs)<<48 |
163		uint64(a.RetIntRealRegs)<<40 |
164		uint64(a.RetFloatRealRegs)<<32 |
165		uint64(a.AlignedArgResultStackSlotSize())
166}
167
168func ABIInfoFromUint64(info uint64) (argIntRealRegs, argFloatRealRegs, retIntRealRegs, retFloatRealRegs byte, stackSlotSize uint32) {
169	return byte(info >> 56), byte(info >> 48), byte(info >> 40), byte(info >> 32), uint32(info)
170}