amd64: saves clobbered xmm8-xmm15 beyond runtime.memmove (#2202)

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

internal/engine/wazevo/backend/isa/amd64/instr.go

@@ -298,8 +298,8 @@ func (i *instruction) String() string {
 		}
 		return fmt.Sprintf("lock xadd.%s %s, %s", suffix, i.op1.format(true), i.op2.format(true))
 
-	case keepAlive:
-		return fmt.Sprintf("keepAlive %s", i.op1.format(true))
+	case nopUseReg:
+		return fmt.Sprintf("nop_use_reg %s", i.op1.format(true))
 
 	default:
 		panic(fmt.Sprintf("BUG: %d", int(i.kind)))
@@ -860,8 +860,8 @@ const (
 	// lockxadd is xadd https://www.felixcloutier.com/x86/xadd with a lock prefix.
 	lockxadd
 
-	// keepAlive is a meta instruction that uses one register and does nothing.
-	keepAlive
+	// nopUseReg is a meta instruction that uses one register and does nothing.
+	nopUseReg
 
 	instrMax
 )
@@ -871,8 +871,8 @@ func (i *instruction) asMFence() *instruction {
 	return i
 }
 
-func (i *instruction) asKeepAlive(r regalloc.VReg) *instruction {
-	i.kind = keepAlive
+func (i *instruction) asNopUseReg(r regalloc.VReg) *instruction {
+	i.kind = nopUseReg
 	i.op1 = newOperandReg(r)
 	return i
 }
@@ -2366,7 +2366,7 @@ var defKinds = [instrMax]defKind{
 	lockcmpxchg:            defKindNone,
 	lockxadd:               defKindNone,
 	neg:                    defKindNone,
-	keepAlive:              defKindNone,
+	nopUseReg:              defKindNone,
 }
 
 // String implements fmt.Stringer.
@@ -2449,7 +2449,7 @@ var useKinds = [instrMax]useKind{
 	lockcmpxchg:            useKindRaxOp1RegOp2,
 	lockxadd:               useKindOp1RegOp2,
 	neg:                    useKindOp1,
-	keepAlive:              useKindOp1,
+	nopUseReg:              useKindOp1,
 }
 
 func (u useKind) String() string {

internal/engine/wazevo/backend/isa/amd64/instr_encoding.go

@@ -11,7 +11,7 @@ import (
 
 func (i *instruction) encode(c backend.Compiler) (needsLabelResolution bool) {
 	switch kind := i.kind; kind {
-	case nop0, sourceOffsetInfo, defineUninitializedReg, fcvtToSintSequence, fcvtToUintSequence, keepAlive:
+	case nop0, sourceOffsetInfo, defineUninitializedReg, fcvtToSintSequence, fcvtToUintSequence, nopUseReg:
 	case ret:
 		encodeRet(c)
 	case imm:

internal/engine/wazevo/backend/isa/amd64/machine.go

@@ -1051,7 +1051,7 @@ func (m *machine) lowerAtomicRmw(op ssa.AtomicRmwOp, addr, val ssa.Value, size u
 		}
 
 		// valCopied must be alive at the end of the loop.
-		m.insert(m.allocateInstr().asKeepAlive(valCopied))
+		m.insert(m.allocateInstr().asNopUseReg(valCopied))
 
 		// At this point, accumulator contains the result.
 		m.clearHigherBitsForAtomic(accumulator, size, ret.Type())
@@ -1758,10 +1758,11 @@ func (m *machine) lowerCall(si *ssa.Instruction) {
 	var directCallee ssa.FuncRef
 	var sigID ssa.SignatureID
 	var args []ssa.Value
+	var isMemmove bool
 	if isDirectCall {
 		directCallee, sigID, args = si.CallData()
 	} else {
-		indirectCalleePtr, sigID, args = si.CallIndirectData()
+		indirectCalleePtr, sigID, args, isMemmove = si.CallIndirectData()
 	}
 	calleeABI := m.c.GetFunctionABI(m.c.SSABuilder().ResolveSignature(sigID))
 
@@ -1779,6 +1780,15 @@ func (m *machine) lowerCall(si *ssa.Instruction) {
 		m.callerGenVRegToFunctionArg(calleeABI, i, reg, def, stackSlotSize)
 	}
 
+	if isMemmove {
+		// Go's memmove *might* use all xmm0-xmm15, so we need to release them.
+		// https://github.com/golang/go/blob/49d42128fd8594c172162961ead19ac95e247d24/src/cmd/compile/abi-internal.md#architecture-specifics
+		// https://github.com/golang/go/blob/49d42128fd8594c172162961ead19ac95e247d24/src/runtime/memmove_amd64.s#L271-L286
+		for i := regalloc.RealReg(0); i < 16; i++ {
+			m.insert(m.allocateInstr().asDefineUninitializedReg(regInfo.RealRegToVReg[xmm0+i]))
+		}
+	}
+
 	if isDirectCall {
 		call := m.allocateInstr().asCall(directCallee, calleeABI)
 		m.insert(call)
@@ -1788,6 +1798,12 @@ func (m *machine) lowerCall(si *ssa.Instruction) {
 		m.insert(callInd)
 	}
 
+	if isMemmove {
+		for i := regalloc.RealReg(0); i < 16; i++ {
+			m.insert(m.allocateInstr().asNopUseReg(regInfo.RealRegToVReg[xmm0+i]))
+		}
+	}
+
 	var index int
 	r1, rs := si.Returns()
 	if r1.Valid() {

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

@@ -267,7 +267,7 @@ func (m *machine) lowerCall(si *ssa.Instruction) {
 	if isDirectCall {
 		directCallee, sigID, args = si.CallData()
 	} else {
-		indirectCalleePtr, sigID, args = si.CallIndirectData()
+		indirectCalleePtr, sigID, args, _ /* on arm64, the calling convention is compatible with the Go runtime */ = si.CallIndirectData()
 	}
 	calleeABI := m.compiler.GetFunctionABI(m.compiler.SSABuilder().ResolveSignature(sigID))
 

internal/engine/wazevo/frontend/lower.go

@@ -3743,7 +3743,7 @@ func (c *Compiler) callMemmove(dst, src, size ssa.Value) {
 			wazevoapi.ExecutionContextOffsetMemmoveAddress.U32(),
 			ssa.TypeI64,
 		).Insert(builder).Return()
-	builder.AllocateInstruction().AsCallIndirect(memmovePtr, &c.memmoveSig, args).Insert(builder)
+	builder.AllocateInstruction().AsCallGoRuntimeMemmove(memmovePtr, &c.memmoveSig, args).Insert(builder)
 }
 
 func (c *Compiler) reloadAfterCall() {

internal/engine/wazevo/ssa/instructions.go

@@ -2180,14 +2180,22 @@ func (i *Instruction) AsCallIndirect(funcPtr Value, sig *Signature, args Values)
 	return i
 }
 
+// AsCallGoRuntimeMemmove is the same as AsCallIndirect, but with a special flag set to indicate that it is a call to the Go runtime memmove function.
+func (i *Instruction) AsCallGoRuntimeMemmove(funcPtr Value, sig *Signature, args Values) *Instruction {
+	i.AsCallIndirect(funcPtr, sig, args)
+	i.u2 = 1
+	return i
+}
+
 // CallIndirectData returns the call indirect data for this instruction necessary for backends.
-func (i *Instruction) CallIndirectData() (funcPtr Value, sigID SignatureID, args []Value) {
+func (i *Instruction) CallIndirectData() (funcPtr Value, sigID SignatureID, args []Value, isGoMemmove bool) {
 	if i.opcode != OpcodeCallIndirect {
 		panic("BUG: CallIndirectData only available for OpcodeCallIndirect")
 	}
 	funcPtr = i.v
 	sigID = SignatureID(i.u1)
 	args = i.vs.View()
+	isGoMemmove = i.u2 == 1
 	return
 }
 

internal/integration_test/fuzzcases/fuzzcases_test.go

@@ -1082,8 +1082,5 @@ func Test2201(t *testing.T) {
 	if !platform.CompilerSupported() {
 		return
 	}
-	if runtime.GOARCH == "amd64" {
-		t.Skip("TODO: #2198")
-	}
 	nodiff.RequireNoDiffT(t, getWasmBinary(t, "2201"), false, false)
 }