Repository: tetratelabs/wazero
Pull request #1970 (2024-01-29 17:43:35)

wazevo(amd64): implement swap in regalloc


Patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr.go b/internal/engine/wazevo/backend/isa/amd64/instr.go
index b34e8ffb..7136689f 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr.go
@@ -235,6 +235,8 @@ func (i *instruction) String() string {
 		return fmt.Sprintf("callq *%s", i.op1.format(true))
 	case v128ConstIsland:
 		return fmt.Sprintf("v128ConstIsland (%#x, %#x)", i.u1, i.u2)
+	case xchg:
+		return fmt.Sprintf("xchg %s, %s", i.op1.format(true), i.op2.format(true))
 	default:
 		panic(fmt.Sprintf("BUG: %d", int(i.kind)))
 	}
@@ -642,6 +644,11 @@ const (
 	// An instruction that will always trigger the illegal instruction exception.
 	ud2
 
+	// xchg swaps the contents of two gp registers.
+	// The instruction doesn't make sense before register allocation, so it doensn't
+	// have useKinds and defKinds to avoid being used by the register allocator.
+	xchg
+
 	// v128ConstIsland is 16 bytes (128-bit) constant that will be loaded into an XMM.
 	v128ConstIsland
 
@@ -736,6 +743,8 @@ func (k instructionKind) String() string {
 		return "v128ConstIsland"
 	case ud2:
 		return "ud2"
+	case xchg:
+		return "xchg"
 	default:
 		panic("BUG")
 	}
@@ -1092,6 +1101,13 @@ func (i *instruction) asPush64(op operand) *instruction {
 	return i
 }
 
+func (i *instruction) asXCHG(rm, rd regalloc.VReg) *instruction {
+	i.kind = xchg
+	i.op1 = newOperandReg(rm)
+	i.op2 = newOperandReg(rd)
+	return i
+}
+
 type unaryRmROpcode byte
 
 const (
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go b/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
index 54c5f9f6..1bc03ab5 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
@@ -1040,6 +1040,12 @@ func (i *instruction) encode(c backend.Compiler) (needsLabelResolution bool) {
 		c.Emit8Bytes(lo)
 		c.Emit8Bytes(hi)
 
+	case xchg:
+		r1, r2 := i.op1.r.RealReg(), i.op2.r.RealReg()
+
+		enc1, enc2 := regEncodings[r1], regEncodings[r2]
+		encodeRegReg(c, legacyPrefixesNone, 0x87, 1, enc2, enc1, rexInfo(0).setW())
+
 	default:
 		panic(fmt.Sprintf("TODO: %v", i.kind))
 	}
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine_regalloc.go b/internal/engine/wazevo/backend/isa/amd64/machine_regalloc.go
index 01759d08..aa7f87eb 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine_regalloc.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine_regalloc.go
@@ -105,8 +105,29 @@ func (m *machine) ClobberedRegisters(regs []regalloc.VReg) {
 
 // Swap implements backend.RegAllocFunctionMachine.
 func (m *machine) Swap(cur *instruction, x1, x2, tmp regalloc.VReg) {
-	// TODO implement me
-	panic("implement me")
+	if x1.RegType() == regalloc.RegTypeInt {
+		prevNext := cur.next
+		xc := m.allocateInstr().asXCHG(x1, x2)
+		cur = linkInstr(cur, xc)
+		linkInstr(cur, prevNext)
+	} else {
+		if tmp.Valid() {
+			prevNext := cur.next
+			m.InsertMoveBefore(tmp, x1, prevNext)
+			m.InsertMoveBefore(x1, x2, prevNext)
+			m.InsertMoveBefore(x2, tmp, prevNext)
+		} else {
+			prevNext := cur.next
+			r2 := x2.RealReg()
+			// Temporarily spill x1 to stack.
+			cur = m.InsertStoreRegisterAt(x1, cur, true).prev
+			// Then move x2 to x1.
+			cur = linkInstr(cur, m.allocateInstr().asXmmUnaryRmR(sseOpcodeMovdqa, newOperandReg(x2), x1))
+			linkInstr(cur, prevNext)
+			// Then reload the original value on x1 from stack to r2.
+			m.InsertReloadRegisterAt(x1.SetRealReg(r2), cur, true)
+		}
+	}
 }
 
 // LastInstrForInsertion implements backend.RegAllocFunctionMachine.


Test patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr_encoding_test.go b/internal/engine/wazevo/backend/isa/amd64/instr_encoding_test.go
index c7aefd4c..526133ef 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr_encoding_test.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr_encoding_test.go
@@ -2876,6 +2876,21 @@ func TestInstruction_format_encode(t *testing.T) {
 			want:       "eeeeeeeeffffffffbbbbbbbbaaaaaaaa",
 			wantFormat: "v128ConstIsland (0xffffffffeeeeeeee, 0xaaaaaaaabbbbbbbb)",
 		},
+		{
+			setup:      func(i *instruction) { i.asXCHG(r11VReg, r14VReg) },
+			want:       "4d87f3",
+			wantFormat: "xchg %r11, %r14",
+		},
+		{
+			setup:      func(i *instruction) { i.asXCHG(r15VReg, raxVReg) },
+			want:       "4987c7",
+			wantFormat: "xchg %r15, %rax",
+		},
+		{
+			setup:      func(i *instruction) { i.asXCHG(rbxVReg, rsiVReg) },
+			want:       "4887f3",
+			wantFormat: "xchg %rbx, %rsi",
+		},
 	} {
 		tc := tc
 		t.Run(tc.wantFormat, func(t *testing.T) {
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine_regalloc_test.go b/internal/engine/wazevo/backend/isa/amd64/machine_regalloc_test.go
index 5d01c018..f86d17b8 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine_regalloc_test.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine_regalloc_test.go
@@ -160,3 +160,73 @@ func TestMachine_InsertMoveBefore(t *testing.T) {
 		})
 	}
 }
+
+func TestMachineSwap(t *testing.T) {
+	for _, tc := range []struct {
+		x1, x2, tmp regalloc.VReg
+		expected    string
+	}{
+		{
+			x1:  r15VReg,
+			x2:  raxVReg,
+			tmp: rdiVReg,
+			expected: `
+	ud2
+	xchg %r15, %rax
+	exit_sequence %rsi
+`,
+		},
+		{
+			x1: r15VReg,
+			x2: raxVReg,
+			expected: `
+	ud2
+	xchg %r15, %rax
+	exit_sequence %rsi
+`,
+		},
+		{
+			x1:  xmm1VReg,
+			x2:  xmm12VReg,
+			tmp: xmm11VReg,
+			expected: `
+	ud2
+	movdqu %xmm1, %xmm11
+	movdqu %xmm12, %xmm1
+	movdqu %xmm11, %xmm12
+	exit_sequence %rsi
+`,
+		},
+		{
+			x1: xmm1VReg,
+			x2: xmm12VReg,
+			// Tmp not given.
+			expected: `
+	ud2
+	movsd %xmm1, (%rsp)
+	movdqa %xmm12, %xmm1
+	movsd (%rsp), %xmm12
+	exit_sequence %rsi
+`,
+		},
+	} {
+		t.Run(tc.expected, func(t *testing.T) {
+			ctx, _, m := newSetupWithMockContext()
+
+			ctx.typeOf = map[regalloc.VRegID]ssa.Type{
+				r15VReg.ID(): ssa.TypeI64, raxVReg.ID(): ssa.TypeI64,
+				xmm1VReg.ID(): ssa.TypeF64, xmm12VReg.ID(): ssa.TypeF64,
+			}
+			cur, i2 := m.allocateInstr().asUD2(), m.allocateInstr().asExitSeq(rsiVReg)
+			cur.next = i2
+			i2.prev = cur
+
+			m.Swap(cur, tc.x1, tc.x2, tc.tmp)
+			m.ectx.RootInstr = cur
+
+			require.Equal(t, tc.expected, m.Format())
+
+			m.encodeWithoutSSA(cur)
+		})
+	}
+}