wazevo(arm64): simplifies lowerLoadSplat

Signed-off-by: Takeshi Yoneda <[email protected]>

internal/engine/wazevo/backend/isa/arm64/lower_mem.go

@@ -233,55 +233,16 @@ func (m *machine) lowerLoad(ptr ssa.Value, offset uint32, typ ssa.Type, ret ssa.
 }
 
 func (m *machine) lowerLoadSplat(ptr ssa.Value, offset uint32, lane ssa.VecLane, ret ssa.Value) {
-	var opSize byte
-	switch lane {
-	case ssa.VecLaneI8x16:
-		opSize = 8
-	case ssa.VecLaneI16x8:
-		opSize = 16
-	case ssa.VecLaneI32x4:
-		opSize = 32
-	case ssa.VecLaneI64x2:
-		opSize = 64
-	}
-	amode := m.lowerToAddressMode(ptr, offset, opSize)
-	rd := operandNR(m.compiler.VRegOf(ret))
-	m.lowerLoadSplatFromAddressMode(rd, amode, lane)
-}
-
-// lowerLoadSplatFromAddressMode is extracted from lowerLoadSplat for testing.
-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 we simply add the offset to the base.
+	base := m.getOperand_NR(m.compiler.ValueDefinition(ptr), extModeNone).nr()
+	offsetReg := m.compiler.AllocateVReg(ssa.TypeI64)
+	m.lowerConstantI64(offsetReg, int64(offset))
+	addedBase := m.addReg64ToReg64(base, offsetReg)
 
-	// vecLoad1R has offset address mode (base+imm) only for post index, so the only addressing mode
-	// we can use here is "no-offset" register addressing mode, i.e. `addressModeKindRegReg`.
-	switch amode.kind {
-	case addressModeKindRegReg:
-		add := m.allocateInstr()
-		add.asALU(aluOpAdd, tmpReg, operandNR(amode.rn), operandNR(amode.rm), true)
-		m.insert(add)
-	case addressModeKindRegSignedImm9:
-		add := m.allocateInstr()
-		add.asALU(aluOpAdd, tmpReg, operandNR(amode.rn), operandImm12(uint16(amode.imm), 0), true)
-		m.insert(add)
-	case addressModeKindRegUnsignedImm12:
-		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")
-	}
-
-	arr := ssaLaneToArrangement(lane)
+	rd := operandNR(m.compiler.VRegOf(ret))
 
 	ld1r := m.allocateInstr()
-	ld1r.asVecLoad1R(rd, tmpReg, arr)
+	ld1r.asVecLoad1R(rd, operandNR(addedBase), ssaLaneToArrangement(lane))
 	m.insert(ld1r)
 }
 

internal/engine/wazevo/backend/isa/arm64/lower_mem_test.go

@@ -3,7 +3,6 @@ package arm64
 import (
 	"fmt"
 	"math"
-	"strconv"
 	"strings"
 	"testing"
 
@@ -831,72 +830,3 @@ func Test_extLoadSizeSign(t *testing.T) {
 		require.Equal(t, tc.signed, signed)
 	}
 }
-
-func Test_lowerLoadSplatFromAddressMode(t *testing.T) {
-	positiveTests := make(map[addressModeKind]bool)
-	nextVReg := regalloc.VReg(100).SetRegType(regalloc.RegTypeInt)
-
-	for _, tc := range []struct {
-		amode       addressMode
-		expected    string
-		expectPanic bool
-	}{
-		{
-			amode: addressMode{kind: addressModeKindRegReg, rn: x0VReg, rm: x1VReg},
-			expected: `
-add x100?, x0, x1
-ld1r {x10.4s}, [x100?]
-`,
-		},
-		{
-			amode: addressMode{kind: addressModeKindRegUnsignedImm12, rn: x0VReg, imm: 15616},
-			expected: `
-movz x101?, #0x3d00, lsl 0
-add x100?, x0, x101?
-ld1r {x10.4s}, [x100?]
-`,
-		},
-		{
-			amode: addressMode{kind: addressModeKindRegUnsignedImm12, rn: x15VReg, imm: 0},
-			expected: `
-ld1r {x10.4s}, [x15]
-`,
-		},
-		{
-			amode: addressMode{kind: addressModeKindRegSignedImm9, rn: x0VReg, imm: 42},
-			expected: `
-add x100?, x0, #0x2a
-ld1r {x10.4s}, [x100?]
-`,
-		},
-	} {
-		tc := tc
-		t.Run("address mode "+strconv.Itoa(int(tc.amode.kind)), func(t *testing.T) {
-			ctx, _, m := newSetupWithMockContext()
-			ctx.vRegCounter = int(nextVReg.ID()) - 1
-			positiveTests[tc.amode.kind] = true
-
-			m.lowerLoadSplatFromAddressMode(operandNR(x10VReg), tc.amode, ssa.VecLaneI32x4)
-			require.Equal(t, tc.expected, "\n"+formatEmittedInstructionsInCurrentBlock(m)+"\n")
-		})
-	}
-
-	// Must panic for all other addressModeKinds.
-	for k := 0; k <= int(addressModeKindResultStackSpace); k++ {
-		amk := addressModeKind(k)
-		if positiveTests[amk] {
-			continue
-		}
-
-		ctx, _, m := newSetupWithMockContext()
-		ctx.vRegCounter = int(nextVReg.ID()) - 1
-
-		t.Run("address mode "+strconv.Itoa(k), func(t *testing.T) {
-			err := require.CapturePanic(func() {
-				m.lowerLoadSplatFromAddressMode(operandNR(x10VReg), addressMode{kind: amk}, ssa.VecLaneI32x4)
-			})
-			require.Contains(t, err.Error(), "unsupported address mode for LoadSplat")
-		})
-
-	}
-}