Repository: tetratelabs/wazero
Pull request #1910 (2024-01-10 21:11:14)

wazevo(amd64): jmp/jmpIf encodings


Patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/cond.go b/internal/engine/wazevo/backend/isa/amd64/cond.go
new file mode 100644
index 00000000..a17b9b9e
--- /dev/null
+++ b/internal/engine/wazevo/backend/isa/amd64/cond.go
@@ -0,0 +1,129 @@
+package amd64
+
+import "github.com/tetratelabs/wazero/internal/engine/wazevo/ssa"
+
+type cond byte
+
+const (
+	// condO represents (overflow) condition.
+	condO cond = iota
+	// condNO represents (no overflow) condition.
+	condNO
+	// condB represents (< unsigned) condition.
+	condB
+	// condNB represents (>= unsigned) condition.
+	condNB
+	// condZ represents (zero) condition.
+	condZ
+	// condNZ represents (not-zero) condition.
+	condNZ
+	// condBE represents (<= unsigned) condition.
+	condBE
+	// condNBE represents (> unsigned) condition.
+	condNBE
+	// condS represents (negative) condition.
+	condS
+	// condNS represents (not-negative) condition.
+	condNS
+	// condP represents (parity) condition.
+	condP
+	// condNP represents (not parity) condition.
+	condNP
+	// condL represents (< signed) condition.
+	condL
+	// condNL represents (>= signed) condition.
+	condNL
+	// condLE represents (<= signed) condition.
+	condLE
+	// condNLE represents (> signed) condition.
+	condNLE
+)
+
+func (c cond) String() string {
+	switch c {
+	case condO:
+		return "o"
+	case condNO:
+		return "no"
+	case condB:
+		return "b"
+	case condNB:
+		return "nb"
+	case condZ:
+		return "z"
+	case condNZ:
+		return "nz"
+	case condBE:
+		return "be"
+	case condNBE:
+		return "nbe"
+	case condS:
+		return "s"
+	case condNS:
+		return "ns"
+	case condL:
+		return "l"
+	case condNL:
+		return "nl"
+	case condLE:
+		return "le"
+	case condNLE:
+		return "nle"
+	case condP:
+		return "p"
+	case condNP:
+		return "np"
+	default:
+		panic("unreachable")
+	}
+}
+
+func condFromSSAIntCmpCond(origin ssa.IntegerCmpCond) cond { // nolint:unused
+	switch origin {
+	case ssa.IntegerCmpCondEqual:
+		return condZ
+	case ssa.IntegerCmpCondNotEqual:
+		return condNZ
+	case ssa.IntegerCmpCondSignedLessThan:
+		return condL
+	case ssa.IntegerCmpCondSignedGreaterThanOrEqual:
+		return condNL
+	case ssa.IntegerCmpCondSignedGreaterThan:
+		return condNLE
+	case ssa.IntegerCmpCondSignedLessThanOrEqual:
+		return condLE
+	case ssa.IntegerCmpCondUnsignedLessThan:
+		return condB
+	case ssa.IntegerCmpCondUnsignedGreaterThanOrEqual:
+		return condNB
+	case ssa.IntegerCmpCondUnsignedGreaterThan:
+		return condNBE
+	case ssa.IntegerCmpCondUnsignedLessThanOrEqual:
+		return condBE
+	default:
+		panic("unreachable")
+	}
+}
+
+func condFromSSAFloatCmpCond(origin ssa.FloatCmpCond) cond { // nolint:unused
+	switch origin {
+	case ssa.FloatCmpCondEqual:
+		return condZ
+	case ssa.FloatCmpCondNotEqual:
+		return condNZ
+	case ssa.FloatCmpCondLessThan:
+		return condB
+	case ssa.FloatCmpCondGreaterThanOrEqual:
+		return condNB
+	case ssa.FloatCmpCondGreaterThan:
+		return condNBE
+	case ssa.FloatCmpCondLessThanOrEqual:
+		return condBE
+	default:
+		panic("unreachable")
+	}
+}
+
+func (c cond) encoding() byte {
+	return byte(c)
+}
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr.go b/internal/engine/wazevo/backend/isa/amd64/instr.go
index 1fd00bcb..fb280562 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr.go
@@ -125,17 +125,13 @@ func (i *instruction) String() string {
 		panic("TODO")
 	case xmmRmRImm:
 		panic("TODO")
-	case jmpKnown:
-		panic("TODO")
+	case jmp:
+		return fmt.Sprintf("jmp %s", i.op1.format(true))
 	case jmpIf:
-		panic("TODO")
-	case jmpCond:
-		panic("TODO")
+		return fmt.Sprintf("j%s %s", cond(i.u1), i.op1.format(true))
 	case jmpTableSeq:
 		panic("TODO")
-	case jmpUnknown:
-		panic("TODO")
-	case trapIf:
+	case exitIf:
 		panic("TODO")
 	case ud2:
 		return "ud2"
@@ -338,34 +334,20 @@ const (
 	// Return.
 	ret
 
-	// Jump to a known target: jmp simm32.
-	jmpKnown
+	// Jump: jmp (reg, mem, imm32 or label)
+	jmp
 
-	// One-way conditional branch: jcond cond target.
-	///
-	// This instruction is useful when we have conditional jumps depending on more than two
-	// conditions, see for instance the lowering of Brz/brnz with Fcmp inputs.
-	///
-	// A note of caution: in contexts where the branch target is another block, this has to be the
-	// same successor as the one specified in the terminator branch of the current block.
-	// Otherwise, this might confuse register allocation by creating new invisible edges.
+	// Jump conditionally: jcond cond label.
 	jmpIf
 
-	// Two-way conditional branch: jcond cond target target.
-	// Emitted as a compound sequence; the MachBuffer will shrink it as appropriate.
-	jmpCond
-
 	// 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.
 	jmpTableSeq
 
-	// Indirect jump: jmpq (reg mem).
-	jmpUnknown
-
-	// Traps if the condition code is set.
-	trapIf
+	// Exits the execution if the condition code is set.
+	exitIf
 
 	// An instruction that will always trigger the illegal instruction exception.
 	ud2
@@ -447,18 +429,14 @@ func (k instructionKind) String() string {
 		return "xmmCmpRmR"
 	case xmmRmRImm:
 		return "xmmRmRImm"
-	case jmpKnown:
-		return "jmpKnown"
 	case jmpIf:
 		return "jmpIf"
-	case jmpCond:
-		return "jmpCond"
+	case jmp:
+		return "jmp"
 	case jmpTableSeq:
 		return "jmpTableSeq"
-	case jmpUnknown:
-		return "jmpUnknown"
-	case trapIf:
-		return "trapIf"
+	case exitIf:
+		return "exitIf"
 	case ud2:
 		return "ud2"
 	default:
@@ -496,6 +474,19 @@ func (a aluRmiROpcode) String() string {
 	}
 }
 
+func (i *instruction) asJmpIf(cond cond, target operand) *instruction {
+	i.kind = jmpIf
+	i.u1 = uint64(cond)
+	i.op1 = target
+	return i
+}
+
+func (i *instruction) asJmp(target operand) *instruction {
+	i.kind = jmp
+	i.op1 = target
+	return i
+}
+
 func (i *instruction) asLEA(a amode, rd regalloc.VReg) *instruction {
 	i.kind = lea
 	i.op1 = newOperandMem(a)
@@ -535,7 +526,7 @@ func (i *instruction) asImm(dst regalloc.VReg, value uint64, _64 bool) *instruct
 }
 
 func (i *instruction) asAluRmiR(op aluRmiROpcode, rm operand, rd regalloc.VReg, _64 bool) *instruction {
-	if rm.kind != operandKindReg && rm.kind != operandKindMem && rm.kind != operandImm32 {
+	if rm.kind != operandKindReg && rm.kind != operandKindMem && rm.kind != operandKindImm32 {
 		panic("BUG")
 	}
 	i.kind = aluRmiR
@@ -597,7 +588,7 @@ func (i *instruction) asPop64(rm regalloc.VReg) *instruction {
 }
 
 func (i *instruction) asPush64(op operand) *instruction {
-	if op.kind != operandKindReg && op.kind != operandKindMem && op.kind != operandImm32 {
+	if op.kind != operandKindReg && op.kind != operandKindMem && op.kind != operandKindImm32 {
 		panic("BUG")
 	}
 	i.kind = push64
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go b/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
index 4838e626..648fe573 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
@@ -71,7 +71,7 @@ func (i *instruction) encode(c backend.Compiler) {
 			case operandKindMem:
 				m := i.op1.amode
 				encodeRegMem(c, legacyPrefixesNone, regMemOpc, regMemOpcNum, dst, m, rex)
-			case operandImm32:
+			case operandKindImm32:
 				imm8 := lower8willSignExtendTo32(op1.imm32)
 				var opc uint32
 				if imm8 {
@@ -85,6 +85,8 @@ func (i *instruction) encode(c backend.Compiler) {
 				} else {
 					c.Emit4Bytes(op1.imm32)
 				}
+			default:
+				panic("BUG: invalid operand kind")
 			}
 		} else {
 			const opcodeNum = 1
@@ -112,7 +114,7 @@ func (i *instruction) encode(c backend.Compiler) {
 			case operandKindMem:
 				m := i.op1.amode
 				encodeRegMem(c, legacyPrefixesNone, opcM, opcodeNum, dst, m, rex)
-			case operandImm32:
+			case operandKindImm32:
 				imm8 := lower8willSignExtendTo32(op1.imm32)
 				var opc uint32
 				if imm8 {
@@ -126,6 +128,8 @@ func (i *instruction) encode(c backend.Compiler) {
 				} else {
 					c.Emit4Bytes(op1.imm32)
 				}
+			default:
+				panic("BUG: invalid operand kind")
 			}
 		}
 
@@ -510,9 +514,11 @@ func (i *instruction) encode(c backend.Compiler) {
 			encodeRegMem(
 				c, legacyPrefixesNone, 0xff, 1, regEnc(6), m, rexInfo(0).clearW(),
 			)
-		case operandImm32:
+		case operandKindImm32:
 			c.EmitByte(0x68)
 			c.Emit4Bytes(op.imm32)
+		default:
+			panic("BUG: invalid operand kind")
 		}
 
 	case pop64:
@@ -521,6 +527,7 @@ func (i *instruction) encode(c backend.Compiler) {
 			c.EmitByte(rexEncodingDefault | 0x1)
 		}
 		c.EmitByte(0x58 | dst.encoding())
+
 	case xmmMovRM:
 		panic("TODO")
 	case xmmLoadConst:
@@ -541,17 +548,56 @@ func (i *instruction) encode(c backend.Compiler) {
 		panic("TODO")
 	case xmmRmRImm:
 		panic("TODO")
-	case jmpKnown:
-		panic("TODO")
+	case jmp:
+		const (
+			regMemOpcode    = 0xff
+			regMemOpcodeNum = 1
+			regMemSubOpcode = 4
+		)
+		op := i.op1
+		switch op.kind {
+		case operandKindLabel:
+			panic("BUG: at this point label should have been resolved to imm32")
+		case operandKindImm32:
+			c.EmitByte(0xe9)
+			c.Emit4Bytes(op.imm32)
+		case operandKindMem:
+			m := op.amode
+			encodeRegMem(c,
+				legacyPrefixesNone,
+				regMemOpcode, regMemOpcodeNum,
+				regMemSubOpcode, m, rexInfo(0).clearW(),
+			)
+		case operandKindReg:
+			r := op.r.RealReg()
+			encodeRegReg(
+				c,
+				legacyPrefixesNone,
+				regMemOpcode, regMemOpcodeNum,
+				regMemSubOpcode,
+				regEncodings[r], rexInfo(0).clearW(),
+			)
+		default:
+			panic("BUG: invalid operand kind")
+		}
+
 	case jmpIf:
-		panic("TODO")
-	case jmpCond:
-		panic("TODO")
+		op := i.op1
+		switch op.kind {
+		case operandKindLabel:
+			panic("BUG: at this point label should have been resolved to imm32")
+		case operandKindImm32:
+			c.EmitByte(0x0f)
+			c.EmitByte(0x80 | cond(i.u1).encoding())
+			c.Emit4Bytes(op.imm32)
+		default:
+			panic("BUG: invalid operand kind")
+		}
+
 	case jmpTableSeq:
 		panic("TODO")
-	case jmpUnknown:
-		panic("TODO")
-	case trapIf:
+
+	case exitIf:
 		panic("TODO")
 
 	case ud2:
diff --git a/internal/engine/wazevo/backend/isa/amd64/operands.go b/internal/engine/wazevo/backend/isa/amd64/operands.go
index af3c383f..6c7a0ef1 100644
--- a/internal/engine/wazevo/backend/isa/amd64/operands.go
+++ b/internal/engine/wazevo/backend/isa/amd64/operands.go
@@ -20,12 +20,15 @@ const (
 	// operandKindReg is an operand which is an integer Register.
 	operandKindReg operandKind = iota + 1
 
-	// operandKindMem is an operand which is either an integer Register or a value in Memory.  This can denote an 8, 16,
+	// operandKindMem is a value in Memory.
 	// 32, 64, or 128 bit value.
 	operandKindMem
 
-	// operandKindRegMemImm is either an integer Register, a value in Memory or an Immediate.
-	operandImm32
+	// operandKindImm32 is a signed-32-bit integer immediate value.
+	operandKindImm32
+
+	// operandKindLabel is a label.
+	operandKindLabel
 )
 
 func (o *operand) format(_64 bool) string {
@@ -34,19 +37,25 @@ func (o *operand) format(_64 bool) string {
 		return formatVRegSized(o.r, _64)
 	case operandKindMem:
 		return o.amode.String()
-	case operandImm32:
+	case operandKindImm32:
 		return fmt.Sprintf("$%d", int32(o.imm32))
+	case operandKindLabel:
+		return backend.Label(o.imm32).String()
 	default:
 		panic("BUG: invalid operand kind")
 	}
 }
 
+func newOperandLabel(label backend.Label) operand { //nolint:unused
+	return operand{kind: operandKindLabel, imm32: uint32(label)}
+}
+
 func newOperandReg(r regalloc.VReg) operand {
 	return operand{kind: operandKindReg, r: r}
 }
 
 func newOperandImm32(imm32 uint32) operand {
-	return operand{kind: operandImm32, imm32: imm32}
+	return operand{kind: operandKindImm32, imm32: imm32}
 }
 
 func newOperandMem(amode amode) operand {
@@ -81,7 +90,7 @@ const (
 	// amodeRipRelative is a memory operand with RIP-relative addressing mode.
 	amodeRipRelative
 
-	// TODO: there are other addressing modes such as the one with base register is absent.
+	// TODO: there are other addressing modes such as the one without base register.
 )
 
 func newAmodeImmReg(imm32 uint32, base regalloc.VReg) amode {


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 7ad63184..f3513e19 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr_encoding_test.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr_encoding_test.go
@@ -1404,6 +1404,101 @@ func TestInstruction_format_encode(t *testing.T) {
 			want:       "ff900000ffff",
 			wantFormat: "callq *-65536(%rax)",
 		},
+		{
+			setup:      func(i *instruction) { i.asJmp(newOperandReg(r12VReg)) },
+			want:       "41ffe4",
+			wantFormat: "jmp %r12",
+		},
+		{
+			setup:      func(i *instruction) { i.asJmp(newOperandMem(newAmodeImmReg(0, raxVReg))) },
+			want:       "ff20",
+			wantFormat: "jmp (%rax)",
+		},
+		{
+			setup:      func(i *instruction) { i.asJmp(newOperandImm32(12345)) },
+			want:       "e939300000",
+			wantFormat: "jmp $12345",
+		},
+		{
+			setup:      func(i *instruction) { i.asJmpIf(condO, newOperandImm32(4)) },
+			want:       "0f8004000000",
+			wantFormat: "jo $4",
+		},
+		{
+			setup:      func(i *instruction) { i.asJmpIf(condNO, newOperandImm32(4)) },
+			want:       "0f8104000000",
+			wantFormat: "jno $4",
+		},
+		{
+			setup:      func(i *instruction) { i.asJmpIf(condB, newOperandImm32(4)) },
+			want:       "0f8204000000",
+			wantFormat: "jb $4",
+		},
+		{
+			setup:      func(i *instruction) { i.asJmpIf(condNB, newOperandImm32(4)) },
+			want:       "0f8304000000",
+			wantFormat: "jnb $4",
+		},
+		{
+			setup:      func(i *instruction) { i.asJmpIf(condZ, newOperandImm32(4)) },
+			want:       "0f8404000000",
+			wantFormat: "jz $4",
+		},
+		{
+			setup:      func(i *instruction) { i.asJmpIf(condNZ, newOperandImm32(4)) },
+			want:       "0f8504000000",
+			wantFormat: "jnz $4",
+		},
+		{
+			setup:      func(i *instruction) { i.asJmpIf(condBE, newOperandImm32(4)) },
+			want:       "0f8604000000",
+			wantFormat: "jbe $4",
+		},
+		{
+			setup:      func(i *instruction) { i.asJmpIf(condNBE, newOperandImm32(4)) },
+			want:       "0f8704000000",
+			wantFormat: "jnbe $4",
+		},
+		{
+			setup:      func(i *instruction) { i.asJmpIf(condS, newOperandImm32(4)) },
+			want:       "0f8804000000",
+			wantFormat: "js $4",
+		},
+		{
+			setup:      func(i *instruction) { i.asJmpIf(condNS, newOperandImm32(4)) },
+			want:       "0f8904000000",
+			wantFormat: "jns $4",
+		},
+		{
+			setup:      func(i *instruction) { i.asJmpIf(condL, newOperandImm32(4)) },
+			want:       "0f8c04000000",
+			wantFormat: "jl $4",
+		},
+		{
+			setup:      func(i *instruction) { i.asJmpIf(condNL, newOperandImm32(4)) },
+			want:       "0f8d04000000",
+			wantFormat: "jnl $4",
+		},
+		{
+			setup:      func(i *instruction) { i.asJmpIf(condLE, newOperandImm32(4)) },
+			want:       "0f8e04000000",
+			wantFormat: "jle $4",
+		},
+		{
+			setup:      func(i *instruction) { i.asJmpIf(condNLE, newOperandImm32(4)) },
+			want:       "0f8f04000000",
+			wantFormat: "jnle $4",
+		},
+		{
+			setup:      func(i *instruction) { i.asJmpIf(condP, newOperandImm32(4)) },
+			want:       "0f8a04000000",
+			wantFormat: "jp $4",
+		},
+		{
+			setup:      func(i *instruction) { i.asJmpIf(condNP, newOperandImm32(4)) },
+			want:       "0f8b04000000",
+			wantFormat: "jnp $4",
+		},
 	} {
 		tc := tc
 		t.Run(tc.wantFormat, func(t *testing.T) {