Repository: tetratelabs/wazero
Pull request #1697 (2023-09-11 07:02:26)

wazevo(arm64): fixes shifted register operands
#1496 


Patch:
diff --git a/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands.go b/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands.go
index 202d4887..2035f419 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands.go
@@ -259,14 +259,14 @@ func (m *machine) getOperand_SR_NR(def *backend.SSAValueDefinition, mode extMode
 	if m.compiler.MatchInstr(def, ssa.OpcodeIshl) {
 		// Check if the shift amount is constant instruction.
 		targetVal, amountVal := def.Instr.Arg2()
+		targetVReg := m.getOperand_NR(m.compiler.ValueDefinition(targetVal), extModeNone).nr()
 		amountDef := m.compiler.ValueDefinition(amountVal)
 		if amountDef.IsFromInstr() && amountDef.Instr.Constant() {
 			// If that is the case, we can use the shifted register operand (SR).
 			c := amountDef.Instr.ConstantVal() & 63 // Clears the unnecessary bits.
-			vreg := m.compiler.VRegOf(targetVal)
 			m.compiler.MarkLowered(def.Instr)
 			m.compiler.MarkLowered(amountDef.Instr)
-			return operandSR(vreg, byte(c), shiftOpLSL)
+			return operandSR(targetVReg, byte(c), shiftOpLSL)
 		}
 	}
 	return m.getOperand_NR(def, mode)


Test patch:
diff --git a/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands_test.go b/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands_test.go
index 4a846f0c..03306820 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands_test.go
@@ -187,7 +187,7 @@ func TestMachine_getOperand_SR_NR(t *testing.T) {
 				ctx.definitions[p2] = &backend.SSAValueDefinition{BlkParamVReg: regalloc.VReg(2), BlockParamValue: p2}
 				ctx.definitions[p3] = &backend.SSAValueDefinition{BlkParamVReg: regalloc.VReg(3), BlockParamValue: p3}
 				ctx.definitions[addResult] = &backend.SSAValueDefinition{Instr: add, N: 0}
-
+				ctx.vRegMap[addResult] = regalloc.VReg(1234) // whatever is fine.
 				ctx.vRegMap[ishl.Return()] = regalloc.VReg(10)
 				def = &backend.SSAValueDefinition{Instr: ishl, N: 0}
 				return def, extModeNone, func(t *testing.T) {}
@@ -199,6 +199,39 @@ func TestMachine_getOperand_SR_NR(t *testing.T) {
 			setup: ishlWithConstAmount,
 			exp:   operandSR(regalloc.VReg(1234), 14, shiftOpLSL),
 		},
+		{
+			name: "ishl with const amount with const shift target",
+			setup: func(
+				ctx *mockCompiler, builder ssa.Builder, m *machine,
+			) (def *backend.SSAValueDefinition, mode extMode, verify func(t *testing.T)) {
+				const nextVReg = 100
+				ctx.vRegCounter = nextVReg - 1
+
+				target := builder.AllocateInstruction().AsIconst64(0xff).Insert(builder)
+				amount := builder.AllocateInstruction().AsIconst32(14).Insert(builder)
+				targetVal, amountVal := target.Return(), amount.Return()
+
+				ishl := builder.AllocateInstruction()
+				ishl.AsIshl(target.Return(), amountVal)
+				builder.InsertInstruction(ishl)
+
+				ctx.definitions[targetVal] = &backend.SSAValueDefinition{Instr: target, N: 0}
+				ctx.definitions[amountVal] = &backend.SSAValueDefinition{Instr: amount, N: 0}
+				ctx.vRegMap[targetVal] = regalloc.VReg(1234)
+				ctx.vRegMap[ishl.Return()] = regalloc.VReg(10)
+				def = &backend.SSAValueDefinition{Instr: ishl, N: 0}
+				mode = extModeNone
+				verify = func(t *testing.T) {
+					_, ok := ctx.lowered[ishl]
+					require.True(t, ok)
+					_, ok = ctx.lowered[amount]
+					require.True(t, ok)
+				}
+				return
+			},
+			exp:          operandSR(regalloc.VReg(100).SetRegType(regalloc.RegTypeInt), 14, shiftOpLSL),
+			instructions: []string{"orr x100?, xzr, #0xff"},
+		},
 		{
 			name: "ishl with const amount but group id is different",
 			setup: func(ctx *mockCompiler, builder ssa.Builder, m *machine) (def *backend.SSAValueDefinition, mode extMode, verify func(t *testing.T)) {