Repository: tetratelabs/wazero
Pull request #1943 (2024-01-23 19:36:12)

wazevo(amd64): support for unconditional jmp


Patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine.go b/internal/engine/wazevo/backend/isa/amd64/machine.go
index 831c72b8..eaa7aa9f 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine.go
@@ -416,7 +416,9 @@ func (m *machine) Encode(context.Context) {
 		targetOffset := ectx.LabelPositions[target].BinaryOffset
 		switch p.instr.kind {
 		case jmp:
-			panic("TODO")
+			imm32Offset := p.offset + 1                       // +1 because p.offset is the beginning of the instruction.
+			jmpOffset := int32(targetOffset - (p.offset + 5)) // +5 because jmp is 5 bytes long.
+			binary.LittleEndian.PutUint32((*bufPtr)[imm32Offset:], uint32(jmpOffset))
 		case jmpIf:
 			imm32Offset := p.offset + 2                       // +2 because p.offset is the beginning of the instruction.
 			jmpOffset := int32(targetOffset - (p.offset + 6)) // +6 because jmpIf is 6 bytes long.
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine_regalloc.go b/internal/engine/wazevo/backend/isa/amd64/machine_regalloc.go
index 8f961301..01759d08 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine_regalloc.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine_regalloc.go
@@ -8,8 +8,22 @@ import (
 
 // InsertMoveBefore implements backend.RegAllocFunctionMachine.
 func (m *machine) InsertMoveBefore(dst, src regalloc.VReg, instr *instruction) {
-	// TODO implement me
-	panic("implement me")
+	typ := src.RegType()
+	if typ != dst.RegType() {
+		panic("BUG: src and dst must have the same type")
+	}
+
+	mov := m.allocateInstr()
+	if typ == regalloc.RegTypeInt {
+		mov.asMovRR(src, dst, true)
+	} else {
+		mov.asXmmUnaryRmR(sseOpcodeMovdqu, newOperandReg(src), dst)
+	}
+
+	cur := instr.prev
+	prevNext := cur.next
+	cur = linkInstr(cur, mov)
+	linkInstr(cur, prevNext)
 }
 
 // InsertStoreRegisterAt implements backend.RegAllocFunctionMachine.
@@ -97,8 +111,19 @@ func (m *machine) Swap(cur *instruction, x1, x2, tmp regalloc.VReg) {
 
 // LastInstrForInsertion implements backend.RegAllocFunctionMachine.
 func (m *machine) LastInstrForInsertion(begin, end *instruction) *instruction {
-	// TODO implement me
-	panic("implement me")
+	cur := end
+	for cur.kind == nop0 {
+		cur = cur.prev
+		if cur == begin {
+			return end
+		}
+	}
+	switch cur.kind {
+	case jmp:
+		return cur
+	default:
+		return end
+	}
 }
 
 // SSABlockLabel implements backend.RegAllocFunctionMachine.
diff --git a/internal/engine/wazevo/backend/regalloc/api.go b/internal/engine/wazevo/backend/regalloc/api.go
index 941808ad..cb3ad7b7 100644
--- a/internal/engine/wazevo/backend/regalloc/api.go
+++ b/internal/engine/wazevo/backend/regalloc/api.go
@@ -72,9 +72,10 @@ type (
 		EndInstr() Instr
 		// LastInstrForInsertion returns the last instruction in this block where instructions will be inserted before it.
 		// Such insertions only happen when we need to insert spill/reload instructions to adjust the merge edges.
-		// If the very last instruction is the unconditional branching, then the returned instruction is the one before it.
-		// Note that at the time of register allocation, all the critical edges are already split, so there is no need
+		// At the time of register allocation, all the critical edges are already split, so there is no need
 		// to worry about the case where branching instruction has multiple successors.
+		// Therefore, usually, it is the nop instruction, but if the block ends with an unconditional branching, then it returns
+		// the unconditional branch, not the nop. In other words it is either nop or unconditional branch.
 		LastInstrForInsertion() Instr
 		// Preds returns the number of predecessors of this block in the CFG.
 		Preds() int


Test patch:
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 802d0363..5d01c018 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine_regalloc_test.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine_regalloc_test.go
@@ -17,8 +17,8 @@ func TestMachine_InsertStoreRegisterAt(t *testing.T) {
 			spillSlotSize: 16,
 			expected: `
 	ud2
-	movsd 16(%rsp), %xmm1
-	movq 24(%rsp), %rax
+	movsd %xmm1, 16(%rsp)
+	mov.q %rax, 24(%rsp)
 	ret
 `,
 		},
@@ -26,8 +26,8 @@ func TestMachine_InsertStoreRegisterAt(t *testing.T) {
 			spillSlotSize: 160,
 			expected: `
 	ud2
-	movsd 160(%rsp), %xmm1
-	movq 168(%rsp), %rax
+	movsd %xmm1, 160(%rsp)
+	mov.q %rax, 168(%rsp)
 	ret
 `,
 		},
@@ -52,11 +52,11 @@ func TestMachine_InsertStoreRegisterAt(t *testing.T) {
 					i2.prev = i1
 
 					if after {
-						m.InsertReloadRegisterAt(raxVReg, i1, after)
-						m.InsertReloadRegisterAt(xmm1VReg, i1, after)
+						m.InsertStoreRegisterAt(raxVReg, i1, after)
+						m.InsertStoreRegisterAt(xmm1VReg, i1, after)
 					} else {
-						m.InsertReloadRegisterAt(xmm1VReg, i2, after)
-						m.InsertReloadRegisterAt(raxVReg, i2, after)
+						m.InsertStoreRegisterAt(xmm1VReg, i2, after)
+						m.InsertStoreRegisterAt(raxVReg, i2, after)
 					}
 					m.ectx.RootInstr = i1
 					require.Equal(t, tc.expected, m.Format())
@@ -123,3 +123,40 @@ func TestMachine_InsertReloadRegisterAt(t *testing.T) {
 		})
 	}
 }
+
+func TestMachine_InsertMoveBefore(t *testing.T) {
+	for _, tc := range []struct {
+		src, dst regalloc.VReg
+		expected string
+	}{
+		{
+			src: raxVReg,
+			dst: rdxVReg,
+			expected: `
+	ud2
+	movq %rax, %rdx
+	ret
+`,
+		},
+		{
+			src: xmm1VReg,
+			dst: xmm2VReg,
+			expected: `
+	ud2
+	movdqu %xmm1, %xmm2
+	ret
+`,
+		},
+	} {
+		t.Run(tc.expected, func(t *testing.T) {
+			_, _, m := newSetupWithMockContext()
+			i1, i2 := m.allocateInstr().asUD2(), m.allocateInstr().asRet(nil)
+			i1.next = i2
+			i2.prev = i1
+
+			m.InsertMoveBefore(tc.dst, tc.src, i2)
+			m.ectx.RootInstr = i1
+			require.Equal(t, tc.expected, m.Format())
+		})
+	}
+}
diff --git a/internal/engine/wazevo/e2e_test.go b/internal/engine/wazevo/e2e_test.go
index baf57875..1fc83580 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -295,9 +295,8 @@ func TestE2E(t *testing.T) {
 			},
 		},
 		{
-			name:      "multi_predecessor_local_ref",
-			m:         testcases.MultiPredecessorLocalRef.Module,
-			skipAMD64: true,
+			name: "multi_predecessor_local_ref",
+			m:    testcases.MultiPredecessorLocalRef.Module,
 			calls: []callCase{
 				{params: []uint64{0, 100}, expResults: []uint64{100}},
 				{params: []uint64{1, 100}, expResults: []uint64{1}},