Repository: tetratelabs/wazero
Pull request #1877 (2023-12-16 03:04:34)

wazevo(arm64): fixes lowerLoadSplatFromAddressMode for zero offset


Patch:
diff --git a/internal/engine/wazevo/backend/isa/arm64/lower_mem.go b/internal/engine/wazevo/backend/isa/arm64/lower_mem.go
index a2484929..3179abda 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_mem.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_mem.go
@@ -245,11 +245,11 @@ func (m *machine) lowerLoadSplat(ptr ssa.Value, offset uint32, lane ssa.VecLane,
 	}
 	amode := m.lowerToAddressMode(ptr, offset, opSize)
 	rd := operandNR(m.compiler.VRegOf(ret))
-	m.lowerLoadSplatFromAddressMode(rd, amode, opSize, lane)
+	m.lowerLoadSplatFromAddressMode(rd, amode, lane)
 }
 
 // lowerLoadSplatFromAddressMode is extracted from lowerLoadSplat for testing.
-func (m *machine) lowerLoadSplatFromAddressMode(rd operand, amode addressMode, opSize byte, lane ssa.VecLane) {
+func (m *machine) lowerLoadSplatFromAddressMode(rd operand, amode addressMode, lane ssa.VecLane) {
 	tmpReg := operandNR(m.compiler.AllocateVReg(ssa.TypeI64))
 
 	// vecLoad1R has offset address mode (base+imm) only for post index, so the only addressing mode
@@ -264,12 +264,15 @@ func (m *machine) lowerLoadSplatFromAddressMode(rd operand, amode addressMode, o
 		add.asALU(aluOpAdd, tmpReg, operandNR(amode.rn), operandImm12(uint16(amode.imm), 0), true)
 		m.insert(add)
 	case addressModeKindRegUnsignedImm12:
-		offsetReg := operandNR(m.compiler.AllocateVReg(ssa.TypeI64))
-		m.load64bitConst(amode.imm, offsetReg.nr())
-
-		add := m.allocateInstr()
-		m.insert(add)
-		add.asALU(aluOpAdd, tmpReg, operandNR(amode.rn), offsetReg, true)
+		if amode.imm != 0 {
+			offsetReg := m.compiler.AllocateVReg(ssa.TypeI64)
+			m.load64bitConst(amode.imm, offsetReg)
+			add := m.allocateInstr()
+			m.insert(add)
+			add.asALU(aluOpAdd, tmpReg, operandNR(amode.rn), operandNR(offsetReg), true)
+		} else {
+			tmpReg = operandNR(amode.rn)
+		}
 	default:
 		panic("unsupported address mode for LoadSplat")
 	}


Test patch:
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 14a2b55f..34b44462 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_mem_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_mem_test.go
@@ -856,24 +856,30 @@ func Test_lowerLoadSplatFromAddressMode(t *testing.T) {
 		expectPanic bool
 	}{
 		{
-			amode: addressMode{kind: addressModeKindRegReg, rn: v0VReg, rm: v1VReg},
+			amode: addressMode{kind: addressModeKindRegReg, rn: x0VReg, rm: x1VReg},
 			expected: `
-add x100?, d0, d1
+add x100?, x0, x1
 ld1r {x10.4s}, [x100?]
 `,
 		},
 		{
-			amode: addressMode{kind: addressModeKindRegUnsignedImm12, rn: v0VReg, imm: 15616},
+			amode: addressMode{kind: addressModeKindRegUnsignedImm12, rn: x0VReg, imm: 15616},
 			expected: `
 movz x101?, #0x3d00, lsl 0
-add x100?, d0, x101?
+add x100?, x0, x101?
 ld1r {x10.4s}, [x100?]
 `,
 		},
 		{
-			amode: addressMode{kind: addressModeKindRegSignedImm9, rn: v0VReg, imm: 42},
+			amode: addressMode{kind: addressModeKindRegUnsignedImm12, rn: x15VReg, imm: 0},
 			expected: `
-add x100?, d0, #0x2a
+ld1r {x10.4s}, [x15]
+`,
+		},
+		{
+			amode: addressMode{kind: addressModeKindRegSignedImm9, rn: x0VReg, imm: 42},
+			expected: `
+add x100?, x0, #0x2a
 ld1r {x10.4s}, [x100?]
 `,
 		},
@@ -884,7 +890,7 @@ ld1r {x10.4s}, [x100?]
 			ctx.vRegCounter = int(nextVReg.ID()) - 1
 			positiveTests[tc.amode.kind] = true
 
-			m.lowerLoadSplatFromAddressMode(operandNR(x10VReg), tc.amode, 32, ssa.VecLaneI32x4)
+			m.lowerLoadSplatFromAddressMode(operandNR(x10VReg), tc.amode, ssa.VecLaneI32x4)
 			require.Equal(t, tc.expected, "\n"+formatEmittedInstructionsInCurrentBlock(m)+"\n")
 		})
 	}
@@ -901,7 +907,7 @@ ld1r {x10.4s}, [x100?]
 
 		t.Run("address mode "+strconv.Itoa(k), func(t *testing.T) {
 			err := require.CapturePanic(func() {
-				m.lowerLoadSplatFromAddressMode(operandNR(x10VReg), addressMode{kind: amk}, 32, ssa.VecLaneI32x4)
+				m.lowerLoadSplatFromAddressMode(operandNR(x10VReg), addressMode{kind: amk}, ssa.VecLaneI32x4)
 			})
 			require.Contains(t, err.Error(), "unsupported address mode for LoadSplat")
 		})