Repository: tetratelabs/wazero
Pull request #1950 (2024-01-24 17:01:08)

wazevo(amd64): lowering for more arithmetic operations
- wazevo(amd64): lowering for more arithmetic operations
- wazevo(amd64): lowering for more arithmetic operations (shifts)
- wazevo(amd64): lowering for more arithmetic operations (rot)


evacchi: rebased to resolve conflict, added small change to e2e test to exercise alternative lowering branch.

Patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr.go b/internal/engine/wazevo/backend/isa/amd64/instr.go
index 807b1dc0..0138c224 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr.go
@@ -998,11 +998,11 @@ func (u unaryRmROpcode) String() string {
 type shiftROp byte
 
 const (
-	shiftROpRotateLeft           = 0
-	shiftROpRotateRight          = 1
-	shiftROpShiftLeft            = 4
-	shiftROpShiftRightLogical    = 5
-	shiftROpShiftRightArithmetic = 7
+	shiftROpRotateLeft           shiftROp = 0
+	shiftROpRotateRight          shiftROp = 1
+	shiftROpShiftLeft            shiftROp = 4
+	shiftROpShiftRightLogical    shiftROp = 5
+	shiftROpShiftRightArithmetic shiftROp = 7
 )
 
 func (s shiftROp) String() string {
@@ -1535,6 +1535,7 @@ var defKinds = [instrMax]defKind{
 	movRM:        defKindNone,
 	xmmMovRM:     defKindNone,
 	aluRmiR:      defKindNone,
+	shiftR:       defKindNone,
 	imm:          defKindOp2,
 	xmmUnaryRmR:  defKindOp2,
 	mov64MR:      defKindOp2,
@@ -1582,6 +1583,7 @@ var useKinds = [instrMax]useKind{
 	movRM:        useKindOp1RegOp2,
 	xmmMovRM:     useKindOp1RegOp2,
 	aluRmiR:      useKindOp1Op2Reg,
+	shiftR:       useKindOp1Op2Reg,
 	imm:          useKindNone,
 	xmmUnaryRmR:  useKindOp1,
 	mov64MR:      useKindOp1,
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine.go b/internal/engine/wazevo/backend/isa/amd64/machine.go
index c6fbb65b..b758ae7d 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine.go
@@ -206,6 +206,26 @@ func (m *machine) LowerInstr(instr *ssa.Instruction) {
 		m.lowerAluRmiROp(instr, aluRmiROpcodeAdd)
 	case ssa.OpcodeIsub:
 		m.lowerAluRmiROp(instr, aluRmiROpcodeSub)
+	case ssa.OpcodeImul:
+		m.lowerAluRmiROp(instr, aluRmiROpcodeMul)
+	case ssa.OpcodeBand:
+		m.lowerAluRmiROp(instr, aluRmiROpcodeAnd)
+	case ssa.OpcodeBor:
+		m.lowerAluRmiROp(instr, aluRmiROpcodeOr)
+	case ssa.OpcodeBxor:
+		m.lowerAluRmiROp(instr, aluRmiROpcodeXor)
+	case ssa.OpcodeIshl:
+		m.lowerShiftR(instr, shiftROpShiftLeft)
+	case ssa.OpcodeSshr:
+		m.lowerShiftR(instr, shiftROpShiftRightArithmetic)
+	case ssa.OpcodeRotl:
+		m.lowerShiftR(instr, shiftROpRotateLeft)
+	case ssa.OpcodeRotr:
+		m.lowerShiftR(instr, shiftROpRotateRight)
+	case ssa.OpcodeUshr:
+		m.lowerShiftR(instr, shiftROpShiftRightLogical)
+	case ssa.OpcodeUndefined:
+		m.insert(m.allocateInstr().asUD2())
 	case ssa.OpcodeExitWithCode:
 		execCtx, code := instr.ExitWithCodeData()
 		m.lowerExitWithCode(m.c.VRegOf(execCtx), code)
@@ -277,7 +297,7 @@ func (m *machine) lowerAluRmiROp(si *ssa.Instruction, op aluRmiROpcode) {
 	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,
+	// rn 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)
@@ -293,6 +313,51 @@ func (m *machine) lowerAluRmiROp(si *ssa.Instruction, op aluRmiROpcode) {
 	m.insert(mov2)
 }
 
+func (m *machine) lowerShiftR(si *ssa.Instruction, op shiftROp) {
+	x, amt := si.Arg2()
+	if !x.Type().IsInt() {
+		panic("BUG?")
+	}
+	_64 := x.Type().Bits() == 64
+
+	xDef, amtDef := m.c.ValueDefinition(x), m.c.ValueDefinition(amt)
+
+	opAmt := m.getOperand_Imm32_Reg(amtDef)
+	rx := m.getOperand_Reg(xDef)
+	rd := m.c.VRegOf(si.Return())
+	tmpDst := m.c.AllocateVReg(si.Return().Type())
+
+	// rx 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(rx.r, tmpDst, _64)
+	m.insert(mov)
+
+	if opAmt.r != regalloc.VRegInvalid {
+		// If opAmt is a register we must copy its value to rcx,
+		// because shiftR encoding mandates that the shift amount is in rcx.
+		mov := m.allocateInstr()
+		mov.asMovRR(opAmt.r, rcxVReg, _64)
+		m.insert(mov)
+
+		alu := m.allocateInstr()
+		alu.asShiftR(op, newOperandReg(rcxVReg), tmpDst, _64)
+		m.insert(alu)
+
+	} else {
+		println("::::::SONO QUI")
+
+		alu := m.allocateInstr()
+		alu.asShiftR(op, opAmt, tmpDst, _64)
+		m.insert(alu)
+	}
+
+	// tmp now contains the result, we copy it to the dest register.
+	mov2 := m.allocateInstr()
+	mov2.asMovRR(tmpDst, rd, _64)
+	m.insert(mov2)
+}
+
 func (m *machine) lowerStore(si *ssa.Instruction) {
 	value, ptr, offset, storeSizeInBits := si.StoreData()
 	rm := m.c.VRegOf(value)


Test patch:
diff --git a/internal/engine/wazevo/e2e_test.go b/internal/engine/wazevo/e2e_test.go
index 965d6a68..b1b39ff5 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -102,6 +102,27 @@ func TestE2E(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "arithm_return", m: testcases.ArithmReturn.Module,
+			calls: []callCase{
+				{
+					params: []uint64{
+						21, 10, 0xf0,
+						21, 10, 0xf0,
+					},
+					expResults: []uint64{
+						21 * 10, 21 & 10, 21 | 10, 21 ^ 10,
+						21 << 10, 21 >> 10, 0xf0 >> 10,
+						88080384, 20971520,
+
+						21 * 10, 21 & 10, 21 | 10, 21 ^ 10,
+						21 << 10, 21 >> 10, 0xf0 >> 10,
+						378302368699121664, 20971520,
+					},
+				},
+			},
+		},
+
 		{
 			name: "many_params_many_results",
 			m:    testcases.ManyParamsManyResults.Module,
diff --git a/internal/engine/wazevo/testcases/testcases.go b/internal/engine/wazevo/testcases/testcases.go
index fcc09f0d..7e54e753 100644
--- a/internal/engine/wazevo/testcases/testcases.go
+++ b/internal/engine/wazevo/testcases/testcases.go
@@ -59,7 +59,7 @@ var (
 			wasm.OpcodeI64Const, 6,
 			wasm.OpcodeI64Sub,
 
-			// Large i64 constants are load from register on arm64, amd64.
+			// Large i64 constants are loaded from register on arm64, amd64.
 			wasm.OpcodeI64Const, 3,
 			wasm.OpcodeI64Const, 0xff, 0xff, 0xff, 0xff, 0xff, 0,
 			wasm.OpcodeI64Add,
@@ -80,6 +80,89 @@ var (
 			wasm.OpcodeEnd,
 		}, nil),
 	}
+	ArithmReturn = TestCase{
+		Name: "add_sub_params_return",
+		Module: SingleFunctionModule(
+			wasm.FunctionType{
+				Params: []wasm.ValueType{i32, i32, i32, i64, i64, i64},
+				Results: []wasm.ValueType{
+					i32, i32, i32, i32,
+					i32, i32, i32,
+					i32, i32,
+
+					i64, i64, i64, i64,
+					i64, i64, i64,
+					i64, i64,
+				},
+			},
+			[]byte{
+				wasm.OpcodeLocalGet, 0,
+				wasm.OpcodeLocalGet, 1,
+				wasm.OpcodeI32Mul,
+				wasm.OpcodeLocalGet, 0,
+				wasm.OpcodeLocalGet, 1,
+				wasm.OpcodeI32And,
+				wasm.OpcodeLocalGet, 0,
+				wasm.OpcodeLocalGet, 1,
+				wasm.OpcodeI32Or,
+				wasm.OpcodeLocalGet, 0,
+				wasm.OpcodeLocalGet, 1,
+				wasm.OpcodeI32Xor,
+
+				wasm.OpcodeLocalGet, 0,
+				wasm.OpcodeLocalGet, 1,
+				wasm.OpcodeI32Shl,
+				wasm.OpcodeLocalGet, 0,
+				wasm.OpcodeLocalGet, 2,
+				wasm.OpcodeI32ShrS,
+				wasm.OpcodeLocalGet, 0,
+				wasm.OpcodeLocalGet, 1,
+				wasm.OpcodeI32ShrU,
+
+				wasm.OpcodeLocalGet, 0,
+				wasm.OpcodeLocalGet, 1,
+				wasm.OpcodeI32Rotr,
+				// Exercise the case where the amount is an immediate.
+				wasm.OpcodeI32Const, 10,
+				wasm.OpcodeLocalGet, 0,
+				wasm.OpcodeI32Rotl,
+
+				// i64
+
+				wasm.OpcodeLocalGet, 3,
+				wasm.OpcodeLocalGet, 4,
+				wasm.OpcodeI64Mul,
+				wasm.OpcodeLocalGet, 3,
+				wasm.OpcodeLocalGet, 4,
+				wasm.OpcodeI64And,
+				wasm.OpcodeLocalGet, 3,
+				wasm.OpcodeLocalGet, 4,
+				wasm.OpcodeI64Or,
+				wasm.OpcodeLocalGet, 3,
+				wasm.OpcodeLocalGet, 4,
+				wasm.OpcodeI64Xor,
+
+				wasm.OpcodeLocalGet, 3,
+				wasm.OpcodeLocalGet, 4,
+				wasm.OpcodeI64Shl,
+				wasm.OpcodeLocalGet, 5,
+				wasm.OpcodeLocalGet, 4,
+				wasm.OpcodeI64ShrS,
+				wasm.OpcodeLocalGet, 3,
+				wasm.OpcodeLocalGet, 4,
+				wasm.OpcodeI64ShrU,
+
+				wasm.OpcodeLocalGet, 3,
+				wasm.OpcodeLocalGet, 4,
+				wasm.OpcodeI64Rotr,
+				// Exercise the case where the amount is an immediate.
+				wasm.OpcodeI64Const, 10,
+				wasm.OpcodeLocalGet, 3,
+				wasm.OpcodeI64Rotl,
+
+				wasm.OpcodeEnd,
+			}, nil),
+	}
 	Locals       = TestCase{Name: "locals", Module: SingleFunctionModule(vv, []byte{wasm.OpcodeEnd}, []wasm.ValueType{i32, i64, f32, f64})}
 	LocalsParams = TestCase{
 		Name: "locals_params",