wazevo(arm64): fixes tmp register types (#1831)

Signed-off-by: Takeshi Yoneda <t.y.mathetake@gmail.com>

internal/engine/wazevo/backend/backend_test.go

@@ -1095,14 +1095,14 @@ L1 (SSA Block: blk0):
 	stp x30, xzr, [sp, #-0x10]!
 	str xzr, [sp, #-0x10]!
 	mov x8, x0
-	mov x10.8b, v0.8b
 	msr fpsr, xzr
-	fcvtzs x0, d10
+	fcvtzs x0, d0
 	mrs x9 fpsr
 	subs xzr, x9, #0x1
 	mov x9, x8
+	mov d8, d0
 	b.ne #0x70, (L17)
-	fcmp x10, x10
+	fcmp d8, d8
 	mov x10, x9
 	b.vc #0x34, (L16)
 	movz x11, #0xc, lsl 0
@@ -1126,9 +1126,9 @@ L17:
 	mrs x9 fpsr
 	subs xzr, x9, #0x1
 	mov x9, x8
-	mov x10, d1
+	mov d8, d1
 	b.ne #0x70, (L15)
-	fcmp w10, w10
+	fcmp s8, s8
 	mov x10, x9
 	b.vc #0x34, (L14)
 	movz x11, #0xc, lsl 0
@@ -1148,12 +1148,13 @@ L14:
 	exit_sequence x9
 L15:
 	msr fpsr, xzr
-	fcvtzs w2, d10
+	fcvtzs w2, d0
 	mrs x9 fpsr
 	subs xzr, x9, #0x1
 	mov x9, x8
+	mov d8, d0
 	b.ne #0x70, (L13)
-	fcmp x10, x10
+	fcmp d8, d8
 	mov x10, x9
 	b.vc #0x34, (L12)
 	movz x11, #0xc, lsl 0
@@ -1177,9 +1178,9 @@ L13:
 	mrs x9 fpsr
 	subs xzr, x9, #0x1
 	mov x9, x8
-	mov x10, d1
+	mov d8, d1
 	b.ne #0x70, (L11)
-	fcmp w10, w10
+	fcmp s8, s8
 	mov x10, x9
 	b.vc #0x34, (L10)
 	movz x11, #0xc, lsl 0
@@ -1199,12 +1200,13 @@ L10:
 	exit_sequence x9
 L11:
 	msr fpsr, xzr
-	fcvtzu x4, d10
+	fcvtzu x4, d0
 	mrs x9 fpsr
 	subs xzr, x9, #0x1
 	mov x9, x8
+	mov d8, d0
 	b.ne #0x70, (L9)
-	fcmp x10, x10
+	fcmp d8, d8
 	mov x10, x9
 	b.vc #0x34, (L8)
 	movz x11, #0xc, lsl 0
@@ -1228,9 +1230,9 @@ L9:
 	mrs x9 fpsr
 	subs xzr, x9, #0x1
 	mov x9, x8
-	mov x10, d1
+	mov d8, d1
 	b.ne #0x70, (L7)
-	fcmp w10, w10
+	fcmp s8, s8
 	mov x10, x9
 	b.vc #0x34, (L6)
 	movz x11, #0xc, lsl 0
@@ -1250,12 +1252,13 @@ L6:
 	exit_sequence x9
 L7:
 	msr fpsr, xzr
-	fcvtzu w6, d10
+	fcvtzu w6, d0
 	mrs x9 fpsr
 	subs xzr, x9, #0x1
 	mov x9, x8
+	mov d8, d0
 	b.ne #0x70, (L5)
-	fcmp x10, x10
+	fcmp d8, d8
 	mov x10, x9
 	b.vc #0x34, (L4)
 	movz x11, #0xc, lsl 0
@@ -1278,9 +1281,9 @@ L5:
 	fcvtzu w7, s1
 	mrs x9 fpsr
 	subs xzr, x9, #0x1
-	mov x9, d1
+	mov d8, d1
 	b.ne #0x70, (L3)
-	fcmp w9, w9
+	fcmp s8, s8
 	mov x9, x8
 	b.vc #0x34, (L2)
 	movz x10, #0xc, lsl 0
@@ -1299,7 +1302,7 @@ L2:
 	str x9, [x8, #0x30]
 	exit_sequence x8
 L3:
-	fcvt s0, x10
+	fcvt s0, d0
 	fcvt d1, s1
 	add sp, sp, #0x10
 	ldr x30, [sp], #0x10

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

@@ -1189,7 +1189,7 @@ func (m *machine) lowerIDiv(execCtxVReg regalloc.VReg, rd, rn, rm operand, _64bi
 // If `c` (cond type) is a register, `cond64bit` must be chosen to indicate whether the register is 32-bit or 64-bit.
 // Otherwise, `cond64bit` is ignored.
 func (m *machine) exitIfNot(execCtxVReg regalloc.VReg, c cond, cond64bit bool, code wazevoapi.ExitCode) {
-	execCtxTmp := m.copyToTmp(execCtxVReg, ssa.TypeI64)
+	execCtxTmp := m.copyToTmp(execCtxVReg)
 
 	cbr := m.allocateInstr()
 	m.insert(cbr)
@@ -1317,8 +1317,8 @@ func (m *machine) lowerFpuToInt(rd, rn operand, ctx regalloc.VReg, signed, src64
 		alu.asALU(aluOpSubS, operandNR(xzrVReg), operandNR(tmpReg), operandImm12(1, 0), true)
 		m.insert(alu)
 
-		execCtx := m.copyToTmp(ctx, ssa.TypeI64)
+		execCtx := m.copyToTmp(ctx)
-		_rn := operandNR(m.copyToTmp(rn.nr(), ssa.TypeI64))
+		_rn := operandNR(m.copyToTmp(rn.nr()))
 
 		// If it is not undefined, we can return the result.
 		ok := m.allocateInstr()
@@ -1843,7 +1843,7 @@ func (m *machine) lowerExitIfTrueWithCode(execCtxVReg regalloc.VReg, cond ssa.Va
 	signed := c.Signed()
 	m.lowerIcmpToFlag(x, y, signed)
 
-	execCtxTmp := m.copyToTmp(execCtxVReg, ssa.TypeI64)
+	execCtxTmp := m.copyToTmp(execCtxVReg)
 
 	// We have to skip the entire exit sequence if the condition is false.
 	cbr := m.allocateInstr()
@@ -1939,7 +1939,8 @@ func (m *machine) lowerSelectVec(rc, rn, rm, rd operand) {
 
 // copyToTmp copies the given regalloc.VReg to a temporary register. This is called before cbr to avoid the regalloc issue
 // e.g. reload happening in the middle of the exit sequence which is not the path the normal path executes
-func (m *machine) copyToTmp(v regalloc.VReg, typ ssa.Type) regalloc.VReg {
+func (m *machine) copyToTmp(v regalloc.VReg) regalloc.VReg {
+	typ := m.compiler.TypeOf(v)
 	mov := m.allocateInstr()
 	tmp := m.compiler.AllocateVReg(typ)
 	mov.asMove64(tmp, v)

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

@@ -378,7 +378,8 @@ L2:
 			rd, rn, rm := regalloc.VReg(1).SetRegType(regalloc.RegTypeInt),
 				regalloc.VReg(2).SetRegType(regalloc.RegTypeInt),
 				regalloc.VReg(3).SetRegType(regalloc.RegTypeInt)
-			_, _, m := newSetupWithMockContext()
+			mc, _, m := newSetupWithMockContext()
+			mc.typeOf = map[regalloc.VReg]ssa.Type{execCtx: ssa.TypeI64, 2: ssa.TypeI64, 3: ssa.TypeI64}
 			m.lowerIDiv(execCtx, operandNR(rd), operandNR(rn), operandNR(rm), tc._64bit, tc.signed)
 			require.Equal(t, tc.exp, "\n"+formatEmittedInstructionsInCurrentBlock(m)+"\n")
 		})
@@ -416,7 +417,7 @@ fcvtzu w1, s2
 mrs x1? fpsr
 subs xzr, x1?, #0x1
 mov x2?, x15
-mov x3?, x2
+mov x3?, d2
 b.ne L2
 fcmp w3?, w3?
 mov x4?, x2?
@@ -448,8 +449,9 @@ fcvtzu w1, s2
 		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
-			_, _, m := newSetupWithMockContext()
+			mc, _, m := newSetupWithMockContext()
-			m.lowerFpuToInt(operandNR(x1VReg), operandNR(x2VReg), x15VReg, false, false, false, tc.nontrapping)
+			mc.typeOf = map[regalloc.VReg]ssa.Type{v2VReg: ssa.TypeI64, x15VReg: ssa.TypeI64}
+			m.lowerFpuToInt(operandNR(x1VReg), operandNR(v2VReg), x15VReg, false, false, false, tc.nontrapping)
 			require.Equal(t, tc.expectedAsm, "\n"+formatEmittedInstructionsInCurrentBlock(m)+"\n")
 
 			m.FlushPendingInstructions()

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

@@ -90,6 +90,7 @@ func newMockCompilationContext() *mockCompiler {
 	return &mockCompiler{
 		vRegMap:     make(map[ssa.Value]regalloc.VReg),
 		definitions: make(map[ssa.Value]*backend.SSAValueDefinition),
+		typeOf:      map[regalloc.VReg]ssa.Type{},
 	}
 }
 
@@ -102,7 +103,9 @@ func (m *mockCompiler) ResolveSignature(id ssa.SignatureID) *ssa.Signature {
 func (m *mockCompiler) AllocateVReg(typ ssa.Type) regalloc.VReg {
 	m.vRegCounter++
 	regType := regalloc.RegTypeOf(typ)
-	return regalloc.VReg(m.vRegCounter).SetRegType(regType)
+	ret := regalloc.VReg(m.vRegCounter).SetRegType(regType)
+	m.typeOf[ret] = typ
+	return ret
 }
 
 // ValueDefinition implements backend.Compiler.