Repository: tetratelabs/wazero
Pull request #1926 (2024-01-17 15:00:23)

wazevo(amd64): Iadd, Isub lowering
add admits various operand combinations, but `asAluRmiR()` accepts only `rm operand, rd reg` (at least for now).

This PR also sets up clobber/declobber in the prologue/epilogue as it is made necessary in the e2e test case. 

mathetake: will take it from here

Patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr.go b/internal/engine/wazevo/backend/isa/amd64/instr.go
index bbe7506b..284b8f34 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr.go
@@ -239,6 +239,24 @@ func (i *instruction) Uses(regs *[]regalloc.VReg) []regalloc.VReg {
 	*regs = (*regs)[:0]
 	switch uk := useKinds[i.kind]; uk {
 	case useKindNone:
+	case useKindOp1Op2:
+		op1, op2 := &i.op1, &i.op2
+		// The destination operand (op2) can be only reg,
+		// the source operand (op1) can be imm32, reg or mem.
+		switch op1.kind {
+		case operandKindReg:
+			*regs = append(*regs, op1.r)
+		case operandKindMem:
+			op1.amode.uses(regs)
+		case operandKindImm32:
+			// Valid, no reg.
+		default:
+			panic(fmt.Sprintf("BUG: invalid operand pair: %s -- %s, %s", i, op1.kind, op2.kind))
+		}
+		if op2.kind != operandKindReg {
+			panic(fmt.Sprintf("BUG: unsupported operand op1 in %s -- %s", i, op1.kind))
+		}
+		*regs = append(*regs, op2.r)
 	case useKindOp1:
 		op := i.op1
 		switch op.kind {
@@ -262,6 +280,48 @@ func (i *instruction) Uses(regs *[]regalloc.VReg) []regalloc.VReg {
 func (i *instruction) AssignUse(index int, v regalloc.VReg) {
 	switch uk := useKinds[i.kind]; uk {
 	case useKindNone:
+	case useKindOp1Op2:
+		op1, op2 := &i.op1, &i.op2
+		switch op1.kind {
+		case operandKindReg:
+			if index == 0 {
+				if op1.r.IsRealReg() {
+					panic("BUG already assigned: " + i.String())
+				}
+				op1.r = v
+			} else if index == 1 {
+				if op2.r.IsRealReg() {
+					panic("BUG already assigned: " + i.String())
+				}
+				op2.r = v
+			} else {
+				panic("BUG")
+			}
+		case operandKindMem:
+			nregs := op1.amode.nregs()
+			if index < nregs {
+				op1.amode.assignUses(index, v)
+			} else if index == nregs {
+				if op2.r.IsRealReg() {
+					panic("BUG already assigned: " + i.String())
+				}
+				op2.r = v
+			} else {
+				panic("BUG")
+			}
+		case operandKindImm32:
+			if index == 0 {
+				if op2.r.IsRealReg() {
+					panic("BUG already assigned: " + i.String())
+				}
+				op2.r = v
+			} else {
+				panic("BUG")
+			}
+		default:
+			panic(fmt.Sprintf("BUG: invalid operand pair: %s -- %s, %s", i, op1.kind, op2.kind))
+		}
+
 	case useKindOp1:
 		op := &i.op1
 		switch op.kind {
@@ -1454,6 +1514,7 @@ var defKinds = [instrMax]defKind{
 	ret:         defKindNone,
 	movRR:       defKindOp2,
 	movRM:       defKindNone,
+	aluRmiR:     defKindNone,
 	imm:         defKindOp2,
 	xmmUnaryRmR: defKindOp2,
 	gprToXmm:    defKindOp2,
@@ -1479,6 +1540,7 @@ type useKind byte
 const (
 	useKindNone useKind = iota + 1
 	useKindOp1
+	useKindOp1Op2
 	useKindCall
 )
 
@@ -1487,6 +1549,7 @@ var useKinds = [instrMax]useKind{
 	ret:         useKindNone,
 	movRR:       useKindOp1,
 	movRM:       useKindOp1,
+	aluRmiR:     useKindOp1Op2,
 	imm:         useKindNone,
 	xmmUnaryRmR: useKindOp1,
 	gprToXmm:    useKindOp1,
@@ -1499,6 +1562,8 @@ func (u useKind) String() string {
 		return "none"
 	case useKindOp1:
 		return "op1"
+	case useKindOp1Op2:
+		return "op1op2"
 	case useKindCall:
 		return "call"
 	default:
diff --git a/internal/engine/wazevo/backend/isa/amd64/lower_constant.go b/internal/engine/wazevo/backend/isa/amd64/lower_constant.go
index b2305da8..3f2e60af 100644
--- a/internal/engine/wazevo/backend/isa/amd64/lower_constant.go
+++ b/internal/engine/wazevo/backend/isa/amd64/lower_constant.go
@@ -5,6 +5,16 @@ import (
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/ssa"
 )
 
+// lowerConstant allocates a new VReg and inserts the instruction to load the constant value.
+func (m *machine) lowerConstant(instr *ssa.Instruction) (vr regalloc.VReg) {
+	val := instr.Return()
+	valType := val.Type()
+
+	vr = m.c.AllocateVReg(valType)
+	m.InsertLoadConstant(instr, vr)
+	return
+}
+
 // InsertLoadConstant implements backend.Machine.
 func (m *machine) InsertLoadConstant(instr *ssa.Instruction, vr regalloc.VReg) {
 	val := instr.Return()
diff --git a/internal/engine/wazevo/backend/isa/amd64/lower_mem.go b/internal/engine/wazevo/backend/isa/amd64/lower_mem.go
new file mode 100644
index 00000000..5a58ca1d
--- /dev/null
+++ b/internal/engine/wazevo/backend/isa/amd64/lower_mem.go
@@ -0,0 +1,10 @@
+package amd64
+
+import (
+	"github.com/tetratelabs/wazero/internal/engine/wazevo/ssa"
+)
+
+// lowerToAddressMode converts a pointer to an addressMode that can be used as an operand for load/store instructions.
+func (m *machine) lowerToAddressMode(ptr ssa.Value, offsetBase uint32) (am amode) {
+	panic("TODO")
+}
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine.go b/internal/engine/wazevo/backend/isa/amd64/machine.go
index b954240f..c8df65e9 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine.go
@@ -138,11 +138,47 @@ func (m *machine) LowerInstr(instr *ssa.Instruction) {
 		m.lowerCall(instr)
 	case ssa.OpcodeStore:
 		m.lowerStore(instr)
+	case ssa.OpcodeIadd:
+		m.lowerAluRmiROp(instr, aluRmiROpcodeAdd)
+	case ssa.OpcodeIsub:
+		m.lowerAluRmiROp(instr, aluRmiROpcodeSub)
 	default:
 		panic("TODO: lowering " + op.String())
 	}
 }
 
+func (m *machine) lowerAluRmiROp(si *ssa.Instruction, op aluRmiROpcode) {
+	x, y := si.Arg2()
+	if !x.Type().IsInt() {
+		panic("BUG?")
+	}
+
+	_64 := x.Type().Bits() == 64
+
+	xDef, yDef := m.c.ValueDefinition(x), m.c.ValueDefinition(y)
+
+	// TODO: commutative args can be swapped if one of them is an immediate.
+	rn := m.getOperand_Reg(xDef)
+	rm := m.getOperand_Mem_Imm32_Reg(yDef)
+	rd := m.c.VRegOf(si.Return())
+	tmp := m.c.AllocateVReg(si.Return().Type())
+
+	// rm is being overwritten, so we first copy its value to a temp register,
+	// in case it is referenced again later.
+	mov := m.allocateInstr()
+	mov.asMovRR(rn.r, tmp, _64)
+	m.insert(mov)
+
+	alu := m.allocateInstr()
+	alu.asAluRmiR(op, rm, tmp, _64)
+	m.insert(alu)
+
+	// tmp now contains the result, we copy it to the dest register.
+	mov2 := m.allocateInstr()
+	mov2.asMovRR(tmp, rd, _64)
+	m.insert(mov2)
+}
+
 func (m *machine) lowerStore(si *ssa.Instruction) {
 	value, ptr, offset, storeSizeInBits := si.StoreData()
 	rm := m.c.VRegOf(value)
diff --git a/internal/engine/wazevo/backend/isa/amd64/operands.go b/internal/engine/wazevo/backend/isa/amd64/operands.go
index 59c7b824..abac7abd 100644
--- a/internal/engine/wazevo/backend/isa/amd64/operands.go
+++ b/internal/engine/wazevo/backend/isa/amd64/operands.go
@@ -5,6 +5,7 @@ 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"
 )
 
 type operand struct {
@@ -99,7 +100,7 @@ type amode struct {
 type amodeKind byte
 
 const (
-	// amodeRegRegShift calcualtes sign-extend-32-to-64(Immediate) + base
+	// amodeRegRegShift calculates sign-extend-32-to-64(Immediate) + base
 	amodeImmReg amodeKind = iota + 1
 
 	// amodeRegRegShift calculates sign-extend-32-to-64(Immediate) + base + (Register2 << Shift)
@@ -124,6 +125,19 @@ func (a *amode) uses(rs *[]regalloc.VReg) {
 	}
 }
 
+func (a *amode) nregs() int {
+	switch a.kind {
+	case amodeImmReg:
+		return 1
+	case amodeRegRegShift:
+		return 2
+	case amodeRipRelative:
+		return 0
+	default:
+		panic("BUG: invalid amode kind")
+	}
+}
+
 func (a *amode) assignUses(i int, reg regalloc.VReg) {
 	switch a.kind {
 	case amodeImmReg:
@@ -198,3 +212,79 @@ func (a *amode) String() string {
 	}
 	panic("BUG: invalid amode kind")
 }
+
+func (m *machine) getOperand_Mem_Imm32_Reg(def *backend.SSAValueDefinition) (op operand) {
+	if def.IsFromBlockParam() {
+		return newOperandReg(def.BlkParamVReg)
+	}
+
+	if opcode := m.c.MatchInstrOneOf(def, []ssa.Opcode{
+		ssa.OpcodeLoad,
+		ssa.OpcodeUload8, ssa.OpcodeUload16, ssa.OpcodeUload32,
+		ssa.OpcodeSload8, ssa.OpcodeSload16, ssa.OpcodeSload32,
+	}); opcode != ssa.OpcodeInvalid {
+		switch opcode {
+		case ssa.OpcodeLoad, ssa.OpcodeUload32, ssa.OpcodeSload32:
+			instr := def.Instr
+			ptr, offset, _ := instr.LoadData()
+			op = newOperandMem(m.lowerToAddressMode(ptr, offset))
+			instr.MarkLowered()
+			return op
+		default:
+			panic("TODO: " + opcode.String())
+		}
+	}
+	return m.getOperand_Imm32_Reg(def)
+}
+
+func (m *machine) getOperand_Imm32_Reg(def *backend.SSAValueDefinition) (op operand) {
+	if def.IsFromBlockParam() {
+		return newOperandReg(def.BlkParamVReg)
+	}
+
+	instr := def.Instr
+	if instr.Constant() {
+		if op, ok := asImm32Operand(instr.ConstantVal()); ok {
+			instr.MarkLowered()
+			return op
+		}
+	}
+	return m.getOperand_Reg(def)
+}
+
+func asImm32(val uint64) (uint32, bool) {
+	u32val := uint32(val)
+	if uint64(u32val) != val {
+		return 0, false
+	}
+	return u32val, true
+}
+
+func asImm32Operand(val uint64) (operand, bool) {
+	if u32val, ok := asImm32(val); ok {
+		return newOperandImm32(u32val), true
+	}
+	return operand{}, false
+}
+
+func (m *machine) getOperand_Reg(def *backend.SSAValueDefinition) (op operand) {
+	var v regalloc.VReg
+	if def.IsFromBlockParam() {
+		v = def.BlkParamVReg
+	} else {
+		instr := def.Instr
+		if instr.Constant() {
+			// We inline all the constant instructions so that we could reduce the register usage.
+			v = m.lowerConstant(instr)
+			instr.MarkLowered()
+		} else {
+			if n := def.N; n == 0 {
+				v = m.c.VRegOf(instr.Return())
+			} else {
+				_, rs := instr.Returns()
+				v = m.c.VRegOf(rs[n-1])
+			}
+		}
+	}
+	return newOperandReg(v)
+}
diff --git a/internal/engine/wazevo/backend/isa/arm64/lower_mem.go b/internal/engine/wazevo/backend/isa/arm64/lower_mem.go
index 3179abda..e64b8824 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_mem.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_mem.go
@@ -5,6 +5,7 @@ import (
 
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend/regalloc"
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/ssa"
+	"github.com/tetratelabs/wazero/internal/engine/wazevo/wazevoapi"
 )
 
 type (
@@ -310,35 +311,35 @@ func (m *machine) lowerToAddressMode(ptr ssa.Value, offsetBase uint32, size byte
 // During the construction, this might emit additional instructions.
 //
 // Extracted as a separate function for easy testing.
-func (m *machine) lowerToAddressModeFromAddends(a32s *queue[addend32], a64s *queue[regalloc.VReg], size byte, offset int64) (amode addressMode) {
-	switch a64sExist, a32sExist := !a64s.empty(), !a32s.empty(); {
+func (m *machine) lowerToAddressModeFromAddends(a32s *wazevoapi.Queue[addend32], a64s *wazevoapi.Queue[regalloc.VReg], size byte, offset int64) (amode addressMode) {
+	switch a64sExist, a32sExist := !a64s.Empty(), !a32s.Empty(); {
 	case a64sExist && a32sExist:
 		var base regalloc.VReg
-		base = a64s.dequeue()
+		base = a64s.Dequeue()
 		var a32 addend32
-		a32 = a32s.dequeue()
+		a32 = a32s.Dequeue()
 		amode = addressMode{kind: addressModeKindRegExtended, rn: base, rm: a32.r, extOp: a32.ext}
 	case a64sExist && offsetFitsInAddressModeKindRegUnsignedImm12(size, offset):
 		var base regalloc.VReg
-		base = a64s.dequeue()
+		base = a64s.Dequeue()
 		amode = addressMode{kind: addressModeKindRegUnsignedImm12, rn: base, imm: offset}
 		offset = 0
 	case a64sExist && offsetFitsInAddressModeKindRegSignedImm9(offset):
 		var base regalloc.VReg
-		base = a64s.dequeue()
+		base = a64s.Dequeue()
 		amode = addressMode{kind: addressModeKindRegSignedImm9, rn: base, imm: offset}
 		offset = 0
 	case a64sExist:
 		var base regalloc.VReg
-		base = a64s.dequeue()
-		if !a64s.empty() {
-			index := a64s.dequeue()
+		base = a64s.Dequeue()
+		if !a64s.Empty() {
+			index := a64s.Dequeue()
 			amode = addressMode{kind: addressModeKindRegReg, rn: base, rm: index, extOp: extendOpUXTX /* indicates index reg is 64-bit */}
 		} else {
 			amode = addressMode{kind: addressModeKindRegUnsignedImm12, rn: base, imm: 0}
 		}
 	case a32sExist:
-		base32 := a32s.dequeue()
+		base32 := a32s.Dequeue()
 
 		// First we need 64-bit base.
 		base := m.compiler.AllocateVReg(ssa.TypeI64)
@@ -350,8 +351,8 @@ func (m *machine) lowerToAddressModeFromAddends(a32s *queue[addend32], a64s *que
 		baseExt.asExtend(base, base32.r, 32, 64, signed)
 		m.insert(baseExt)
 
-		if !a32s.empty() {
-			index := a32s.dequeue()
+		if !a32s.Empty() {
+			index := a32s.Dequeue()
 			amode = addressMode{kind: addressModeKindRegExtended, rn: base, rm: index.r, extOp: index.ext}
 		} else {
 			amode = addressMode{kind: addressModeKindRegUnsignedImm12, rn: base, imm: 0}
@@ -368,13 +369,13 @@ func (m *machine) lowerToAddressModeFromAddends(a32s *queue[addend32], a64s *que
 		baseReg = m.addConstToReg64(baseReg, offset) // baseReg += offset
 	}
 
-	for !a64s.empty() {
-		a64 := a64s.dequeue()
+	for !a64s.Empty() {
+		a64 := a64s.Dequeue()
 		baseReg = m.addReg64ToReg64(baseReg, a64) // baseReg += a64
 	}
 
-	for !a32s.empty() {
-		a32 := a32s.dequeue()
+	for !a32s.Empty() {
+		a32 := a32s.Dequeue()
 		baseReg = m.addRegToReg64Ext(baseReg, a32.r, a32.ext) // baseReg += (a32 extended to 64-bit)
 	}
 	amode.rn = baseReg
@@ -383,22 +384,22 @@ func (m *machine) lowerToAddressModeFromAddends(a32s *queue[addend32], a64s *que
 
 var addendsMatchOpcodes = [4]ssa.Opcode{ssa.OpcodeUExtend, ssa.OpcodeSExtend, ssa.OpcodeIadd, ssa.OpcodeIconst}
 
-func (m *machine) collectAddends(ptr ssa.Value) (addends32 *queue[addend32], addends64 *queue[regalloc.VReg], offset int64) {
-	m.addendsWorkQueue.reset()
-	m.addends32.reset()
-	m.addends64.reset()
-	m.addendsWorkQueue.enqueue(ptr)
+func (m *machine) collectAddends(ptr ssa.Value) (addends32 *wazevoapi.Queue[addend32], addends64 *wazevoapi.Queue[regalloc.VReg], offset int64) {
+	m.addendsWorkQueue.Reset()
+	m.addends32.Reset()
+	m.addends64.Reset()
+	m.addendsWorkQueue.Enqueue(ptr)
 
-	for !m.addendsWorkQueue.empty() {
-		v := m.addendsWorkQueue.dequeue()
+	for !m.addendsWorkQueue.Empty() {
+		v := m.addendsWorkQueue.Dequeue()
 
 		def := m.compiler.ValueDefinition(v)
 		switch op := m.compiler.MatchInstrOneOf(def, addendsMatchOpcodes[:]); op {
 		case ssa.OpcodeIadd:
 			// If the addend is an add, we recursively collect its operands.
 			x, y := def.Instr.Arg2()
-			m.addendsWorkQueue.enqueue(x)
-			m.addendsWorkQueue.enqueue(y)
+			m.addendsWorkQueue.Enqueue(x)
+			m.addendsWorkQueue.Enqueue(y)
 			def.Instr.MarkLowered()
 		case ssa.OpcodeIconst:
 			// If the addend is constant, we just statically merge it into the offset.
@@ -414,7 +415,7 @@ func (m *machine) collectAddends(ptr ssa.Value) (addends32 *queue[addend32], add
 			switch input := def.Instr.Arg(); input.Type().Bits() {
 			case 64:
 				// If the input is already 64-bit, this extend is a no-op. TODO: shouldn't this be optimized out at much earlier stage? no?
-				m.addends64.enqueue(m.getOperand_NR(m.compiler.ValueDefinition(input), extModeNone).nr())
+				m.addends64.Enqueue(m.getOperand_NR(m.compiler.ValueDefinition(input), extModeNone).nr())
 				def.Instr.MarkLowered()
 				continue
 			case 32:
@@ -435,7 +436,7 @@ func (m *machine) collectAddends(ptr ssa.Value) (addends32 *queue[addend32], add
 					// Sign-extension of a 32-bit constant can be merged into the offset.
 					offset += int64(int32(inputDef.Instr.ConstantVal())) // sign-extend!
 				default:
-					m.addends32.enqueue(addend32{r: m.getOperand_NR(inputDef, extModeNone).nr(), ext: ext})
+					m.addends32.Enqueue(addend32{r: m.getOperand_NR(inputDef, extModeNone).nr(), ext: ext})
 				}
 				def.Instr.MarkLowered()
 				continue
@@ -446,7 +447,7 @@ func (m *machine) collectAddends(ptr ssa.Value) (addends32 *queue[addend32], add
 			panic("TODO: add tests")
 		default:
 			// If the addend is not one of them, we simply use it as-is (without merging!), optionally zero-extending it.
-			m.addends64.enqueue(m.getOperand_NR(def, extModeZeroExtend64 /* optional zero ext */).nr())
+			m.addends64.Enqueue(m.getOperand_NR(def, extModeZeroExtend64 /* optional zero ext */).nr())
 		}
 	}
 	return &m.addends32, &m.addends64, offset
@@ -483,28 +484,3 @@ func (m *machine) addRegToReg64Ext(rn, rm regalloc.VReg, ext extendOp) (rd regal
 	m.insert(alu)
 	return
 }
-
-// queue is the resettable queue where the underlying slice is reused.
-type queue[T any] struct {
-	index int
-	data  []T
-}
-
-func (q *queue[T]) enqueue(v T) {
-	q.data = append(q.data, v)
-}
-
-func (q *queue[T]) dequeue() (ret T) {
-	ret = q.data[q.index]
-	q.index++
-	return
-}
-
-func (q *queue[T]) empty() bool {
-	return q.index >= len(q.data)
-}
-
-func (q *queue[T]) reset() {
-	q.index = 0
-	q.data = q.data[:0]
-}
diff --git a/internal/engine/wazevo/backend/isa/arm64/machine.go b/internal/engine/wazevo/backend/isa/arm64/machine.go
index 94668446..971df393 100644
--- a/internal/engine/wazevo/backend/isa/arm64/machine.go
+++ b/internal/engine/wazevo/backend/isa/arm64/machine.go
@@ -22,10 +22,10 @@ type (
 		regAllocFn *backend.RegAllocFunction[*instruction, *machine]
 
 		// addendsWorkQueue is used during address lowering, defined here for reuse.
-		addendsWorkQueue queue[ssa.Value]
-		addends32        queue[addend32]
+		addendsWorkQueue wazevoapi.Queue[ssa.Value]
+		addends32        wazevoapi.Queue[addend32]
 		// addends64 is used during address lowering, defined here for reuse.
-		addends64              queue[regalloc.VReg]
+		addends64              wazevoapi.Queue[regalloc.VReg]
 		unresolvedAddressModes []*instruction
 
 		// condBrRelocs holds the conditional branches which need offset relocation.
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 36436dac..3ec69038 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
@@ -259,8 +259,8 @@ func (m *machine) setupEpilogueAfter(cur *instruction) {
 		//          |   ReturnAddress |                      |   ReturnAddress |
 		//          +-----------------+      ========>       +-----------------+ <---- SP
 		//          |   clobbered M   |
-		//          |   clobbered 1   |
 		//          |   ...........   |
+		//          |   clobbered 1   |
 		//          |   clobbered 0   |
 		//   SP---> +-----------------+
 		//             (low address)
diff --git a/internal/engine/wazevo/wazevoapi/queue.go b/internal/engine/wazevo/wazevoapi/queue.go
new file mode 100644
index 00000000..e3118fa6
--- /dev/null
+++ b/internal/engine/wazevo/wazevoapi/queue.go
@@ -0,0 +1,26 @@
+package wazevoapi
+
+// Queue is the resettable queue where the underlying slice is reused.
+type Queue[T any] struct {
+	index int
+	Data  []T
+}
+
+func (q *Queue[T]) Enqueue(v T) {
+	q.Data = append(q.Data, v)
+}
+
+func (q *Queue[T]) Dequeue() (ret T) {
+	ret = q.Data[q.index]
+	q.index++
+	return
+}
+
+func (q *Queue[T]) Empty() bool {
+	return q.index >= len(q.Data)
+}
+
+func (q *Queue[T]) Reset() {
+	q.index = 0
+	q.Data = q.Data[:0]
+}


Test patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr_test.go b/internal/engine/wazevo/backend/isa/amd64/instr_test.go
new file mode 100644
index 00000000..633e6fb5
--- /dev/null
+++ b/internal/engine/wazevo/backend/isa/amd64/instr_test.go
@@ -0,0 +1,49 @@
+package amd64
+
+import (
+	"testing"
+
+	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend/regalloc"
+	"github.com/tetratelabs/wazero/internal/testing/require"
+)
+
+func TestMachine_lowerAluRmiROp_Uses_AssignUse(t *testing.T) {
+	vr0 := regalloc.VReg(0).SetRegType(regalloc.RegTypeInt)
+	vr1 := regalloc.VReg(1).SetRegType(regalloc.RegTypeInt)
+	tests := []struct {
+		name     string
+		instr    func(*instruction)
+		expected string
+	}{
+		{
+			name:     "reg_reg",
+			instr:    func(i *instruction) { i.asAluRmiR(aluRmiROpcodeAdd, newOperandReg(vr0), vr1, false) },
+			expected: "add %eax, %ecx",
+		},
+		{
+			name: "mem_reg",
+			instr: func(i *instruction) {
+				i.asAluRmiR(aluRmiROpcodeAdd, newOperandMem(newAmodeImmReg(123, vr0)), vr1, false)
+			},
+			expected: "add 123(%rax), %ecx",
+		},
+		{
+			name:     "imm_reg",
+			instr:    func(i *instruction) { i.asAluRmiR(aluRmiROpcodeAdd, newOperandImm32(123), vr1, false) },
+			expected: "add $123, %eax",
+		},
+	}
+	rs := []regalloc.RealReg{rax, rcx}
+	for _, tt := range tests {
+		tc := tt
+		t.Run(tt.name, func(t *testing.T) {
+			regs := &[]regalloc.VReg{}
+			instr := &instruction{}
+			tc.instr(instr)
+			for i, use := range instr.Uses(regs) {
+				instr.AssignUse(i, use.SetRealReg(rs[i]))
+			}
+			require.Equal(t, tc.expected, instr.String())
+		})
+	}
+}
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine_test.go b/internal/engine/wazevo/backend/isa/amd64/machine_test.go
new file mode 100644
index 00000000..2efbefdc
--- /dev/null
+++ b/internal/engine/wazevo/backend/isa/amd64/machine_test.go
@@ -0,0 +1,148 @@
+package amd64
+
+import (
+	"strings"
+	"testing"
+
+	"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"
+	"github.com/tetratelabs/wazero/internal/testing/require"
+)
+
+func Test_asImm32(t *testing.T) {
+	v, ok := asImm32(0xffffffff)
+	require.True(t, ok)
+	require.Equal(t, uint32(0xffffffff), v)
+
+	_, ok = asImm32(0xffffffff << 1)
+	require.False(t, ok)
+}
+
+func TestMachine_getOperand_Reg(t *testing.T) {
+	for _, tc := range []struct {
+		name         string
+		setup        func(*mockCompiler, ssa.Builder, *machine) *backend.SSAValueDefinition
+		exp          operand
+		instructions []string
+	}{
+		{
+			name: "block param",
+			setup: func(ctx *mockCompiler, builder ssa.Builder, m *machine) *backend.SSAValueDefinition {
+				return &backend.SSAValueDefinition{BlkParamVReg: raxVReg, Instr: nil, N: 0}
+			},
+			exp: newOperandReg(raxVReg),
+		},
+
+		{
+			name: "const instr",
+			setup: func(ctx *mockCompiler, builder ssa.Builder, m *machine) *backend.SSAValueDefinition {
+				instr := builder.AllocateInstruction()
+				instr.AsIconst32(0xf00000f)
+				builder.InsertInstruction(instr)
+				ctx.vRegCounter = 99
+				return &backend.SSAValueDefinition{Instr: instr, N: 0}
+			},
+			exp:          newOperandReg(regalloc.VReg(100).SetRegType(regalloc.RegTypeInt)),
+			instructions: []string{"movl $251658255, %r100d?"},
+		},
+		{
+			name: "non const instr (single-return)",
+			setup: func(ctx *mockCompiler, builder ssa.Builder, m *machine) *backend.SSAValueDefinition {
+				c := builder.AllocateInstruction()
+				sig := &ssa.Signature{Results: []ssa.Type{ssa.TypeI64}}
+				builder.DeclareSignature(sig)
+				c.AsCall(ssa.FuncRef(0), sig, nil)
+				builder.InsertInstruction(c)
+				r := c.Return()
+				ctx.vRegMap[r] = regalloc.VReg(50)
+				return &backend.SSAValueDefinition{Instr: c, N: 0}
+			},
+			exp: newOperandReg(regalloc.VReg(50)),
+		},
+		{
+			name: "non const instr (multi-return)",
+			setup: func(ctx *mockCompiler, builder ssa.Builder, m *machine) *backend.SSAValueDefinition {
+				c := builder.AllocateInstruction()
+				sig := &ssa.Signature{Results: []ssa.Type{ssa.TypeI64, ssa.TypeF64, ssa.TypeF64}}
+				builder.DeclareSignature(sig)
+				c.AsCall(ssa.FuncRef(0), sig, nil)
+				builder.InsertInstruction(c)
+				_, rs := c.Returns()
+				ctx.vRegMap[rs[1]] = regalloc.VReg(50)
+				return &backend.SSAValueDefinition{Instr: c, N: 2}
+			},
+			exp: newOperandReg(regalloc.VReg(50)),
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx, b, m := newSetupWithMockContext()
+			def := tc.setup(ctx, b, m)
+			actual := m.getOperand_Reg(def)
+			require.Equal(t, tc.exp, actual)
+			require.Equal(t, strings.Join(tc.instructions, "\n"), formatEmittedInstructionsInCurrentBlock(m))
+		})
+	}
+}
+
+func TestMachine_getOperand_Imm32_Reg(t *testing.T) {
+	for _, tc := range []struct {
+		name         string
+		setup        func(*mockCompiler, ssa.Builder, *machine) *backend.SSAValueDefinition
+		exp          operand
+		instructions []string
+	}{
+		{
+			name: "block param falls back to getOperand_Reg",
+			setup: func(ctx *mockCompiler, builder ssa.Builder, m *machine) *backend.SSAValueDefinition {
+				return &backend.SSAValueDefinition{BlkParamVReg: raxVReg, Instr: nil, N: 0}
+			},
+			exp: newOperandReg(raxVReg),
+		},
+		{
+			name: "const imm 32",
+			setup: func(ctx *mockCompiler, builder ssa.Builder, m *machine) *backend.SSAValueDefinition {
+				instr := builder.AllocateInstruction()
+				instr.AsIconst32(0xf00000f)
+				builder.InsertInstruction(instr)
+				ctx.vRegCounter = 99
+				ctx.currentGID = 0xff // const can be merged anytime, regardless of the group id.
+				return &backend.SSAValueDefinition{Instr: instr, N: 0}
+			},
+			exp: newOperandImm32(0xf00000f),
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx, b, m := newSetupWithMockContext()
+			def := tc.setup(ctx, b, m)
+			actual := m.getOperand_Imm32_Reg(def)
+			require.Equal(t, tc.exp, actual)
+			require.Equal(t, strings.Join(tc.instructions, "\n"), formatEmittedInstructionsInCurrentBlock(m))
+		})
+	}
+}
+
+func Test_machine_getOperand_Mem_Imm32_Reg(t *testing.T) {
+	for _, tc := range []struct {
+		name         string
+		setup        func(*mockCompiler, ssa.Builder, *machine) *backend.SSAValueDefinition
+		exp          operand
+		instructions []string
+	}{
+		{
+			name: "block param falls back to getOperand_Imm32_Reg",
+			setup: func(ctx *mockCompiler, builder ssa.Builder, m *machine) *backend.SSAValueDefinition {
+				return &backend.SSAValueDefinition{BlkParamVReg: raxVReg, Instr: nil, N: 0}
+			},
+			exp: newOperandReg(raxVReg),
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx, b, m := newSetupWithMockContext()
+			def := tc.setup(ctx, b, m)
+			actual := m.getOperand_Mem_Imm32_Reg(def)
+			require.Equal(t, tc.exp, actual)
+			require.Equal(t, strings.Join(tc.instructions, "\n"), formatEmittedInstructionsInCurrentBlock(m))
+		})
+	}
+}
diff --git a/internal/engine/wazevo/backend/isa/amd64/util_test.go b/internal/engine/wazevo/backend/isa/amd64/util_test.go
index 98b939c4..e7a91dc5 100644
--- a/internal/engine/wazevo/backend/isa/amd64/util_test.go
+++ b/internal/engine/wazevo/backend/isa/amd64/util_test.go
@@ -2,6 +2,7 @@ package amd64
 
 import (
 	"context"
+	"strings"
 
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend"
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend/regalloc"
@@ -135,3 +136,12 @@ func (m *mockCompiler) MatchInstrOneOf(def *backend.SSAValueDefinition, opcodes
 func (m *mockCompiler) Compile(context.Context) (_ []byte, _ []backend.RelocationInfo, _ error) {
 	return
 }
+
+func formatEmittedInstructionsInCurrentBlock(m *machine) string {
+	m.ectx.FlushPendingInstructions()
+	var strs []string
+	for cur := m.ectx.PerBlockHead; cur != nil; cur = cur.next {
+		strs = append(strs, cur.String())
+	}
+	return strings.Join(strs, "\n")
+}
diff --git a/internal/engine/wazevo/backend/isa/arm64/lower_mem_test.go b/internal/engine/wazevo/backend/isa/arm64/lower_mem_test.go
index 34b44462..e50f2d5e 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_mem_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_mem_test.go
@@ -10,6 +10,7 @@ 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"
+	"github.com/tetratelabs/wazero/internal/engine/wazevo/wazevoapi"
 	"github.com/tetratelabs/wazero/internal/testing/require"
 )
 
@@ -473,8 +474,8 @@ func TestMachine_collectAddends(t *testing.T) {
 			ctx, b, m := newSetupWithMockContext()
 			ptr, verify := tc.setup(ctx, b, m)
 			actual32sQ, actual64sQ, actualOffset := m.collectAddends(ptr)
-			require.Equal(t, tc.exp32s, actual32sQ.data)
-			require.Equal(t, tc.exp64s, actual64sQ.data)
+			require.Equal(t, tc.exp32s, actual32sQ.Data)
+			require.Equal(t, tc.exp64s, actual64sQ.Data)
 			require.Equal(t, tc.offset, actualOffset)
 			verify(t)
 		})
@@ -490,7 +491,7 @@ func TestMachine_addConstToReg64(t *testing.T) {
 		m.addConstToReg64(regalloc.FromRealReg(x15, regalloc.RegTypeInt), c)
 		require.Equal(t, `add x100?, x15, #0xaaa`, formatEmittedInstructionsInCurrentBlock(m))
 	})
-	t.Run("positive imm12", func(t *testing.T) {
+	t.Run("negative imm12", func(t *testing.T) {
 		c := int64(-0xaaa)
 		ctx, _, m := newSetupWithMockContext()
 		ctx.vRegCounter = nextVRegID - 1
@@ -812,13 +813,13 @@ func TestMachine_lowerToAddressModeFromAddends(t *testing.T) {
 			ctx, _, m := newSetupWithMockContext()
 			ctx.vRegCounter = int(nextVReg.ID()) - 1
 
-			var a32s queue[addend32]
-			var a64s queue[regalloc.VReg]
+			var a32s wazevoapi.Queue[addend32]
+			var a64s wazevoapi.Queue[regalloc.VReg]
 			for _, a32 := range tc.a32s {
-				a32s.enqueue(a32)
+				a32s.Enqueue(a32)
 			}
 			for _, a64 := range tc.a64s {
-				a64s.enqueue(a64)
+				a64s.Enqueue(a64)
 			}
 			actual := m.lowerToAddressModeFromAddends(&a32s, &a64s, tc.dstSizeInBits, tc.offset)
 			require.Equal(t, strings.Join(tc.insts, "\n"), formatEmittedInstructionsInCurrentBlock(m))
diff --git a/internal/engine/wazevo/e2e_test.go b/internal/engine/wazevo/e2e_test.go
index ad50e630..a8605555 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -95,6 +95,16 @@ func TestE2E(t *testing.T) {
 			skipAMD64: true,
 			calls:     []callCase{{expErr: "unreachable"}},
 		},
+		{
+			name: "add_sub_return", m: testcases.AddSubReturn.Module,
+			skipAMD64: false,
+			calls: []callCase{
+				{
+					params:     []uint64{},
+					expResults: []uint64{3, 3, 3, 3},
+				},
+			},
+		},
 		{
 			name: "fibonacci_recursive", m: testcases.FibonacciRecursive.Module,
 			skipAMD64: true,
diff --git a/internal/engine/wazevo/testcases/testcases.go b/internal/engine/wazevo/testcases/testcases.go
index 7b44687c..fcc09f0d 100644
--- a/internal/engine/wazevo/testcases/testcases.go
+++ b/internal/engine/wazevo/testcases/testcases.go
@@ -32,9 +32,43 @@ var (
 		byte(math.Float64bits(64.0) >> 56),
 		wasm.OpcodeEnd,
 	}, nil)}
-	Unreachable        = TestCase{Name: "unreachable", Module: SingleFunctionModule(vv, []byte{wasm.OpcodeUnreachable, wasm.OpcodeEnd}, nil)}
-	OnlyReturn         = TestCase{Name: "only_return", Module: SingleFunctionModule(vv, []byte{wasm.OpcodeReturn, wasm.OpcodeEnd}, nil)}
-	Params             = TestCase{Name: "params", Module: SingleFunctionModule(i32f32f64_v, []byte{wasm.OpcodeReturn, wasm.OpcodeEnd}, nil)}
+	Unreachable  = TestCase{Name: "unreachable", Module: SingleFunctionModule(vv, []byte{wasm.OpcodeUnreachable, wasm.OpcodeEnd}, nil)}
+	OnlyReturn   = TestCase{Name: "only_return", Module: SingleFunctionModule(vv, []byte{wasm.OpcodeReturn, wasm.OpcodeEnd}, nil)}
+	Params       = TestCase{Name: "params", Module: SingleFunctionModule(i32f32f64_v, []byte{wasm.OpcodeReturn, wasm.OpcodeEnd}, nil)}
+	AddSubReturn = TestCase{
+		Name: "add_sub_params_return_const",
+		Module: SingleFunctionModule(wasm.FunctionType{Results: []wasm.ValueType{i32, i32, i64, i64}}, []byte{
+			// Small i32 constants should be inlined on arm64, amd64.
+			wasm.OpcodeI32Const, 4,
+			wasm.OpcodeI32Const, 5,
+			wasm.OpcodeI32Add,
+			wasm.OpcodeI32Const, 6,
+			wasm.OpcodeI32Sub,
+
+			// Large i32 constants should be inlined on amd4, load from register on arm64.
+			wasm.OpcodeI32Const, 3,
+			wasm.OpcodeI32Const, 0xff, 0xff, 0xff, 0xff, 0,
+			wasm.OpcodeI32Add,
+			wasm.OpcodeI32Const, 0xff, 0xff, 0xff, 0xff, 0,
+			wasm.OpcodeI32Sub,
+
+			// Small i64 constants should be inlined on arm64, amd64.
+			wasm.OpcodeI64Const, 4,
+			wasm.OpcodeI64Const, 5,
+			wasm.OpcodeI64Add,
+			wasm.OpcodeI64Const, 6,
+			wasm.OpcodeI64Sub,
+
+			// Large i64 constants are load from register on arm64, amd64.
+			wasm.OpcodeI64Const, 3,
+			wasm.OpcodeI64Const, 0xff, 0xff, 0xff, 0xff, 0xff, 0,
+			wasm.OpcodeI64Add,
+			wasm.OpcodeI64Const, 0xff, 0xff, 0xff, 0xff, 0xff, 0,
+			wasm.OpcodeI64Sub,
+
+			wasm.OpcodeEnd,
+		}, nil),
+	}
 	AddSubParamsReturn = TestCase{
 		Name: "add_sub_params_return",
 		Module: SingleFunctionModule(i32i32_i32, []byte{