Repository: tetratelabs/wazero
Pull request #1942 (2024-01-23 18:50:38)

wazevo(amd64): support for conditional jmp


Patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/cond.go b/internal/engine/wazevo/backend/isa/amd64/cond.go
index a17b9b9e..c422e72a 100644
--- a/internal/engine/wazevo/backend/isa/amd64/cond.go
+++ b/internal/engine/wazevo/backend/isa/amd64/cond.go
@@ -78,7 +78,7 @@ func (c cond) String() string {
 	}
 }
 
-func condFromSSAIntCmpCond(origin ssa.IntegerCmpCond) cond { // nolint:unused
+func condFromSSAIntCmpCond(origin ssa.IntegerCmpCond) cond {
 	switch origin {
 	case ssa.IntegerCmpCondEqual:
 		return condZ
@@ -127,3 +127,42 @@ func condFromSSAFloatCmpCond(origin ssa.FloatCmpCond) cond { // nolint:unused
 func (c cond) encoding() byte {
 	return byte(c)
 }
+
+func (c cond) invert() cond {
+	switch c {
+	case condO:
+		return condNO
+	case condNO:
+		return condO
+	case condB:
+		return condNB
+	case condNB:
+		return condB
+	case condZ:
+		return condNZ
+	case condNZ:
+		return condZ
+	case condBE:
+		return condNBE
+	case condNBE:
+		return condBE
+	case condS:
+		return condNS
+	case condNS:
+		return condS
+	case condP:
+		return condNP
+	case condNP:
+		return condP
+	case condL:
+		return condNL
+	case condNL:
+		return condL
+	case condLE:
+		return condNLE
+	case condNLE:
+		return condLE
+	default:
+		panic("unreachable")
+	}
+}
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr.go b/internal/engine/wazevo/backend/isa/amd64/instr.go
index 6c95b8a9..a20c90a4 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr.go
@@ -688,6 +688,13 @@ func (i *instruction) asJmp(target operand) *instruction {
 	return i
 }
 
+func (i *instruction) jmpLabel() backend.Label {
+	if i.kind != jmp && i.kind != jmpIf {
+		panic("BUG")
+	}
+	return i.op1.label()
+}
+
 func (i *instruction) asLEA(a amode, rd regalloc.VReg) *instruction {
 	i.kind = lea
 	i.op1 = newOperandMem(a)
@@ -886,14 +893,14 @@ func (i *instruction) asXmmRmiReg(op sseOpcode, rm operand, rd regalloc.VReg) *i
 	return i
 }
 
-func (i *instruction) asCmpRmiR(cmpOrTest bool, rm operand, rd regalloc.VReg, _64 bool) *instruction {
+func (i *instruction) asCmpRmiR(cmp bool, rm operand, rd regalloc.VReg, _64 bool) *instruction {
 	if rm.kind != operandKindReg && rm.kind != operandKindImm32 && rm.kind != operandKindMem {
 		panic("BUG")
 	}
 	i.kind = cmpRmiR
 	i.op1 = rm
 	i.op2 = newOperandReg(rd)
-	if cmpOrTest {
+	if cmp {
 		i.u1 = 1
 	}
 	i.b1 = _64
@@ -1527,6 +1534,9 @@ var defKinds = [instrMax]defKind{
 	gprToXmm:    defKindOp2,
 	call:        defKindCall,
 	ud2:         defKindNone,
+	jmp:         defKindNone,
+	jmpIf:       defKindNone,
+	cmpRmiR:     defKindNone,
 }
 
 // String implements fmt.Stringer.
@@ -1569,6 +1579,9 @@ var useKinds = [instrMax]useKind{
 	gprToXmm:    useKindOp1,
 	call:        useKindCall,
 	ud2:         useKindNone,
+	jmpIf:       useKindOp1,
+	jmp:         useKindOp1,
+	cmpRmiR:     useKindOp1Op2Reg,
 }
 
 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 36456fdd..c38bb683 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
@@ -907,7 +907,8 @@ func (i *instruction) encode(c backend.Compiler) (needsLabelResolution bool) {
 		op := i.op1
 		switch op.kind {
 		case operandKindLabel:
-			panic("BUG: at this point label should have been resolved to imm32")
+			needsLabelResolution = true
+			fallthrough
 		case operandKindImm32:
 			c.EmitByte(0xe9)
 			c.Emit4Bytes(op.imm32)
@@ -936,7 +937,7 @@ func (i *instruction) encode(c backend.Compiler) (needsLabelResolution bool) {
 		switch op.kind {
 		case operandKindLabel:
 			needsLabelResolution = true
-			panic("BUG: at this point label should have been resolved to imm32")
+			fallthrough
 		case operandKindImm32:
 			c.EmitByte(0x0f)
 			c.EmitByte(0x80 | cond(i.u1).encoding())
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine.go b/internal/engine/wazevo/backend/isa/amd64/machine.go
index 8059d9e9..831c72b8 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine.go
@@ -2,6 +2,7 @@ package amd64
 
 import (
 	"context"
+	"encoding/binary"
 	"fmt"
 	"strings"
 
@@ -134,10 +135,57 @@ func (m *machine) LowerSingleBranch(b *ssa.Instruction) {
 	}
 }
 
+var condBranchMatches = [...]ssa.Opcode{ssa.OpcodeIcmp, ssa.OpcodeFcmp}
+
 // LowerConditionalBranch implements backend.Machine.
 func (m *machine) LowerConditionalBranch(b *ssa.Instruction) {
-	// TODO implement me
-	panic("implement me")
+	exctx := m.ectx
+	cval, args, targetBlk := b.BranchData()
+	if len(args) > 0 {
+		panic(fmt.Sprintf(
+			"conditional branch shouldn't have args; likely a bug in critical edge splitting: from %s to %s",
+			exctx.CurrentSSABlk,
+			targetBlk,
+		))
+	}
+
+	target := exctx.GetOrAllocateSSABlockLabel(targetBlk)
+	cvalDef := m.c.ValueDefinition(cval)
+
+	switch m.c.MatchInstrOneOf(cvalDef, condBranchMatches[:]) {
+	case ssa.OpcodeIcmp:
+		cvalInstr := cvalDef.Instr
+		x, y, c := cvalInstr.IcmpData()
+
+		cc := condFromSSAIntCmpCond(c)
+		if b.Opcode() == ssa.OpcodeBrz {
+			cc = cc.invert()
+		}
+
+		// First, perform the comparison and set the flag.
+		xd, yd := m.c.ValueDefinition(x), m.c.ValueDefinition(y)
+		m.lowerIcmpToFlag(xd, yd, x.Type() == ssa.TypeI64)
+
+		// Then perform the conditional branch.
+		m.insert(m.allocateInstr().asJmpIf(cc, newOperandLabel(target)))
+		cvalDef.Instr.MarkLowered()
+	case ssa.OpcodeFcmp:
+		panic("TODO")
+	default:
+		v := m.getOperand_Reg(cvalDef)
+
+		var cc cond
+		if b.Opcode() == ssa.OpcodeBrz {
+			cc = condZ
+		} else {
+			cc = condNZ
+		}
+
+		// Perform test %v, %v to set the flag.
+		cmp := m.allocateInstr().asCmpRmiR(false, v, v.r, false)
+		m.insert(cmp)
+		m.insert(m.allocateInstr().asJmpIf(cc, newOperandLabel(target)))
+	}
 }
 
 // LowerInstr implements backend.Machine.
@@ -356,7 +404,7 @@ func (m *machine) Encode(context.Context) {
 			needLabelResolution := cur.encode(m.c)
 			if needLabelResolution {
 				m.labelResolutionPends = append(m.labelResolutionPends,
-					labelResolutionPend{instr: cur, offset: int64(offset)},
+					labelResolutionPend{instr: cur, offset: offset},
 				)
 			}
 		}
@@ -364,11 +412,15 @@ func (m *machine) Encode(context.Context) {
 
 	for i := range m.labelResolutionPends {
 		p := &m.labelResolutionPends[i]
+		target := p.instr.jmpLabel()
+		targetOffset := ectx.LabelPositions[target].BinaryOffset
 		switch p.instr.kind {
 		case jmp:
 			panic("TODO")
 		case jmpIf:
-			panic("TODO")
+			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.
+			binary.LittleEndian.PutUint32((*bufPtr)[imm32Offset:], uint32(jmpOffset))
 		default:
 			panic("BUG")
 		}
@@ -390,6 +442,13 @@ func (m *machine) ResolveRelocations(refToBinaryOffset map[ssa.FuncRef]int, bina
 	}
 }
 
+func (m *machine) lowerIcmpToFlag(xd, yd *backend.SSAValueDefinition, _64 bool) {
+	x := m.getOperand_Reg(xd)
+	y := m.getOperand_Mem_Imm32_Reg(yd)
+	cmp := m.allocateInstr().asCmpRmiR(true, y, x.r, _64)
+	m.insert(cmp)
+}
+
 // allocateInstr allocates an instruction.
 func (m *machine) allocateInstr() *instruction {
 	instr := m.ectx.InstructionPool.Allocate()
diff --git a/internal/engine/wazevo/backend/isa/amd64/operands.go b/internal/engine/wazevo/backend/isa/amd64/operands.go
index 20cbee47..0f5c7d27 100644
--- a/internal/engine/wazevo/backend/isa/amd64/operands.go
+++ b/internal/engine/wazevo/backend/isa/amd64/operands.go
@@ -69,6 +69,13 @@ func newOperandLabel(label backend.Label) operand { //nolint:unused
 	return operand{kind: operandKindLabel, imm32: uint32(label)}
 }
 
+func (o *operand) label() backend.Label {
+	if o.kind != operandKindLabel {
+		panic("BUG: invalid operand kind")
+	}
+	return backend.Label(o.imm32)
+}
+
 func newOperandReg(r regalloc.VReg) operand {
 	return operand{kind: operandKindReg, r: r}
 }


Test patch:
diff --git a/internal/engine/wazevo/e2e_test.go b/internal/engine/wazevo/e2e_test.go
index 5404bfaa..baf57875 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -79,7 +79,6 @@ func TestE2E(t *testing.T) {
 		},
 		{
 			name: "swap", m: testcases.SwapParamAndReturn.Module,
-			skipAMD64: true,
 			calls: []callCase{
 				{params: []uint64{math.MaxUint32, math.MaxInt32}, expResults: []uint64{math.MaxInt32, math.MaxUint32}},
 			},
@@ -126,7 +125,6 @@ func TestE2E(t *testing.T) {
 		},
 		{
 			name: "fibonacci_recursive", m: testcases.FibonacciRecursive.Module,
-			skipAMD64: true,
 			calls: []callCase{
 				{params: []uint64{0}, expResults: []uint64{0}},
 				{params: []uint64{1}, expResults: []uint64{1}},