Repository: tetratelabs/wazero
Pull request #1969 (2024-01-26 22:21:28)

wazevo(amd64): support for jump table


Patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr.go b/internal/engine/wazevo/backend/isa/amd64/instr.go
index f2edb869..41a019d0 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr.go
@@ -2,6 +2,7 @@ package amd64
 
 import (
 	"fmt"
+	"strings"
 
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend"
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend/regalloc"
@@ -16,6 +17,7 @@ type instruction struct {
 	b1                  bool
 	addedBeforeRegAlloc bool
 	kind                instructionKind
+	targets             []uint32
 }
 
 // Next implements regalloc.Instr.
@@ -197,8 +199,12 @@ func (i *instruction) String() string {
 		return fmt.Sprintf("jmp %s", i.op1.format(true))
 	case jmpIf:
 		return fmt.Sprintf("j%s %s", cond(i.u1), i.op1.format(true))
-	case jmpTableSequence:
-		panic("TODO")
+	case jmpTableIsland:
+		labels := make([]string, len(i.targets))
+		for index, l := range i.targets {
+			labels[index] = backend.Label(l).String()
+		}
+		return fmt.Sprintf("jump_table_island [%s]", strings.Join(labels, ", "))
 	case exitSequence:
 		return fmt.Sprintf("exit_sequence %s", i.op1.format(true))
 	case ud2:
@@ -565,11 +571,8 @@ const (
 	// Jump conditionally: jcond cond label.
 	jmpIf
 
-	// Jump-table sequence, as one compound instruction (see note in lower.rs for rationale).
-	// The generated code sequence is described in the emit's function match arm for this
-	// instruction.
-	// See comment in lowering about the temporaries signedness.
-	jmpTableSequence
+	// jmpTableIsland is to emit the jump table.
+	jmpTableIsland
 
 	// exitSequence exits the execution and go back to the Go world.
 	exitSequence
@@ -663,8 +666,8 @@ func (k instructionKind) String() string {
 		return "jmpIf"
 	case jmp:
 		return "jmp"
-	case jmpTableSequence:
-		return "jmpTableSeq"
+	case jmpTableIsland:
+		return "jmpTableIsland"
 	case exitSequence:
 		return "exit_sequence"
 	case v128ConstIsland:
@@ -713,6 +716,12 @@ func (i *instruction) asJmpIf(cond cond, target operand) *instruction {
 	return i
 }
 
+func (i *instruction) asJmpTableSequence(targets []uint32) *instruction {
+	i.kind = jmpTableIsland
+	i.targets = targets
+	return i
+}
+
 func (i *instruction) asJmp(target operand) *instruction {
 	i.kind = jmp
 	i.op1 = target
@@ -1580,11 +1589,13 @@ var defKinds = [instrMax]defKind{
 	movsxRmR:        defKindOp2,
 	movzxRmR:        defKindOp2,
 	gprToXmm:        defKindOp2,
+	cmove:           defKindNone,
 	call:            defKindCall,
 	callIndirect:    defKindCall,
 	ud2:             defKindNone,
 	jmp:             defKindNone,
 	jmpIf:           defKindNone,
+	jmpTableIsland:  defKindNone,
 	cmpRmiR:         defKindNone,
 	exitSequence:    defKindNone,
 	lea:             defKindOp2,
@@ -1624,6 +1635,7 @@ var useKinds = [instrMax]useKind{
 	movRR:           useKindOp1,
 	movRM:           useKindOp1RegOp2,
 	xmmMovRM:        useKindOp1RegOp2,
+	cmove:           useKindOp1Op2Reg,
 	aluRmiR:         useKindOp1Op2Reg,
 	shiftR:          useKindOp1Op2Reg,
 	imm:             useKindNone,
@@ -1642,6 +1654,7 @@ var useKinds = [instrMax]useKind{
 	exitSequence:    useKindOp1,
 	lea:             useKindOp1,
 	v128ConstIsland: useKindNone,
+	jmpTableIsland:  useKindNone,
 }
 
 func (u useKind) String() string {
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go b/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
index fdcc5c56..86015ea0 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
@@ -952,8 +952,11 @@ func (i *instruction) encode(c backend.Compiler) (needsLabelResolution bool) {
 			panic("BUG: invalid operand kind")
 		}
 
-	case jmpTableSequence:
-		panic("TODO")
+	case jmpTableIsland:
+		needsLabelResolution = true
+		for range i.targets {
+			c.Emit8Bytes(0)
+		}
 
 	case exitSequence:
 		execCtx := i.op1.r
@@ -1006,10 +1009,8 @@ func (i *instruction) encode(c backend.Compiler) (needsLabelResolution bool) {
 
 	case v128ConstIsland:
 		lo, hi := i.u1, i.u2
-		c.Emit4Bytes(uint32(lo))
-		c.Emit4Bytes(uint32(lo >> 32))
-		c.Emit4Bytes(uint32(hi))
-		c.Emit4Bytes(uint32(hi >> 32))
+		c.Emit8Bytes(lo)
+		c.Emit8Bytes(hi)
 
 	default:
 		panic(fmt.Sprintf("TODO: %v", i.kind))
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine.go b/internal/engine/wazevo/backend/isa/amd64/machine.go
index 38099b81..8cfd5c52 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine.go
@@ -53,7 +53,8 @@ type (
 	}
 
 	labelResolutionPend struct {
-		instr *instruction
+		instr       *instruction
+		instrOffset int64
 		// imm32Offset is the offset of the last 4 bytes of the instruction.
 		imm32Offset int64
 	}
@@ -135,7 +136,8 @@ func (m *machine) LowerSingleBranch(b *ssa.Instruction) {
 		}
 		m.insert(jmp)
 	case ssa.OpcodeBrTable:
-		panic("TODO: implement me")
+		index, target := b.BrTableData()
+		m.lowerBrTable(index, target)
 	default:
 		panic("BUG: unexpected branch opcode" + b.Opcode().String())
 	}
@@ -143,6 +145,50 @@ func (m *machine) LowerSingleBranch(b *ssa.Instruction) {
 
 var condBranchMatches = [...]ssa.Opcode{ssa.OpcodeIcmp, ssa.OpcodeFcmp}
 
+func (m *machine) lowerBrTable(index ssa.Value, targets []ssa.BasicBlock) {
+	_v := m.getOperand_Reg(m.c.ValueDefinition(index))
+	v := m.copyToTmp(_v.r)
+
+	// First, we need to do the bounds check.
+	maxIndex := m.c.AllocateVReg(ssa.TypeI32)
+	m.lowerIconst(maxIndex, uint64(len(targets)-1), false)
+	cmp := m.allocateInstr().asCmpRmiR(true, newOperandReg(maxIndex), v, false)
+	m.insert(cmp)
+
+	// Then do the conditional move maxIndex to v if v > maxIndex.
+	cmov := m.allocateInstr().asCmove(condNB, newOperandReg(maxIndex), v, false)
+	m.insert(cmov)
+
+	// Now that v has the correct index. Load the address of the jump table into the addr.
+	addr := m.c.AllocateVReg(ssa.TypeI64)
+	leaJmpTableAddr := m.allocateInstr()
+	m.insert(leaJmpTableAddr)
+
+	// Then add the target's offset into jmpTableAddr.
+	loadTargetOffsetFromJmpTable := m.allocateInstr().asAluRmiR(aluRmiROpcodeAdd,
+		// Shift by 3 because each entry is 8 bytes.
+		newOperandMem(newAmodeRegRegShift(0, addr, v, 3)), addr, true)
+	m.insert(loadTargetOffsetFromJmpTable)
+
+	// Now ready to jump.
+	jmp := m.allocateInstr().asJmp(newOperandReg(addr))
+	m.insert(jmp)
+
+	jmpTableBegin, jmpTableBeginLabel := m.allocateBrTarget()
+	m.insert(jmpTableBegin)
+	leaJmpTableAddr.asLEA(newAmodeRipRelative(jmpTableBeginLabel), addr)
+
+	jmpTable := m.allocateInstr()
+	// TODO: reuse the slice!
+	labels := make([]uint32, len(targets))
+	for j, target := range targets {
+		labels[j] = uint32(m.ectx.GetOrAllocateSSABlockLabel(target))
+	}
+
+	jmpTable.asJmpTableSequence(labels)
+	m.insert(jmpTable)
+}
+
 // LowerConditionalBranch implements backend.Machine.
 func (m *machine) LowerConditionalBranch(b *ssa.Instruction) {
 	exctx := m.ectx
@@ -883,12 +929,13 @@ func (m *machine) Encode(context.Context) {
 			needLabelResolution := cur.encode(m.c)
 			if needLabelResolution {
 				m.labelResolutionPends = append(m.labelResolutionPends,
-					labelResolutionPend{instr: cur, imm32Offset: int64(len(*bufPtr)) - 4},
+					labelResolutionPend{instr: cur, instrOffset: offset, imm32Offset: int64(len(*bufPtr)) - 4},
 				)
 			}
 		}
 	}
 
+	buf := *bufPtr
 	for i := range m.labelResolutionPends {
 		p := &m.labelResolutionPends[i]
 		switch p.instr.kind {
@@ -897,7 +944,15 @@ func (m *machine) Encode(context.Context) {
 			targetOffset := ectx.LabelPositions[target].BinaryOffset
 			imm32Offset := p.imm32Offset
 			jmpOffset := int32(targetOffset - (p.imm32Offset + 4)) // +4 because RIP points to the next instruction.
-			binary.LittleEndian.PutUint32((*bufPtr)[imm32Offset:], uint32(jmpOffset))
+			binary.LittleEndian.PutUint32(buf[imm32Offset:], uint32(jmpOffset))
+		case jmpTableIsland:
+			tableBegin := p.instrOffset
+			// Each entry is the offset from the beginning of the jmpTableIsland instruction in 8 bytes.
+			for i, l := range p.instr.targets {
+				targetOffset := ectx.LabelPositions[backend.Label(l)].BinaryOffset
+				jmpOffset := targetOffset - tableBegin
+				binary.LittleEndian.PutUint64(buf[tableBegin+int64(i)*8:], uint64(jmpOffset))
+			}
 		default:
 			panic("BUG")
 		}


Test patch:
diff --git a/internal/engine/wazevo/e2e_test.go b/internal/engine/wazevo/e2e_test.go
index 5c08b76f..6ee05d9f 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -280,9 +280,8 @@ func TestE2E(t *testing.T) {
 			},
 		},
 		{
-			name:      "br_table",
-			m:         testcases.BrTable.Module,
-			skipAMD64: true,
+			name: "br_table",
+			m:    testcases.BrTable.Module,
 			calls: []callCase{
 				{params: []uint64{0}, expResults: []uint64{11}},
 				{params: []uint64{1}, expResults: []uint64{12}},
@@ -296,9 +295,8 @@ func TestE2E(t *testing.T) {
 			},
 		},
 		{
-			name:      "br_table_with_args",
-			m:         testcases.BrTableWithArg.Module,
-			skipAMD64: true,
+			name: "br_table_with_args",
+			m:    testcases.BrTableWithArg.Module,
 			calls: []callCase{
 				{params: []uint64{0, 100}, expResults: []uint64{11 + 100}},
 				{params: []uint64{1, 100}, expResults: []uint64{12 + 100}},