Repository: tetratelabs/wazero
Pull request #2020 (2024-02-06 20:06:56)

wazevo(amd64): efficient amode struct 


Patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/abi_go_call.go b/internal/engine/wazevo/backend/isa/amd64/abi_go_call.go
index b1fb32dd..b8786658 100644
--- a/internal/engine/wazevo/backend/isa/amd64/abi_go_call.go
+++ b/internal/engine/wazevo/backend/isa/amd64/abi_go_call.go
@@ -331,7 +331,7 @@ func (m *machine) storeReturnAddressAndExit(cur *instruction, execCtx regalloc.V
 
 	nop, l := m.allocateBrTarget()
 	cur = linkInstr(cur, nop)
-	readRip.asLEA(m.newAmodeRipRelative(l), ripReg)
+	readRip.asLEA(newOperandLabel(l), ripReg)
 	return cur
 }
 
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr.go b/internal/engine/wazevo/backend/isa/amd64/instr.go
index a86af5d1..a10dd81c 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr.go
@@ -992,9 +992,9 @@ func (i *instruction) jmpLabel() backend.Label {
 	}
 }
 
-func (i *instruction) asLEA(a *amode, rd regalloc.VReg) *instruction {
+func (i *instruction) asLEA(target operand, rd regalloc.VReg) *instruction {
 	i.kind = lea
-	i.op1 = newOperandMem(a)
+	i.op1 = target
 	i.op2 = newOperandReg(rd)
 	return i
 }
diff --git a/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go b/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
index d37ee4e9..b3a24243 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr_encoding.go
@@ -681,9 +681,26 @@ func (i *instruction) encode(c backend.Compiler) (needsLabelResolution bool) {
 
 	case lea:
 		needsLabelResolution = true
-		a := i.op1.addressMode()
 		dst := regEncodings[i.op2.reg().RealReg()]
-		encodeRegMem(c, legacyPrefixesNone, 0x8d, 1, dst, a, rexInfo(0).setW())
+		rex := rexInfo(0).setW()
+		const opcode, opcodeNum = 0x8d, 1
+		switch i.op1.kind {
+		case operandKindMem:
+			a := i.op1.addressMode()
+			encodeRegMem(c, legacyPrefixesNone, opcode, opcodeNum, dst, a, rex)
+		case operandKindLabel:
+			rex.encode(c, regRexBit(byte(dst)), 0)
+			c.EmitByte(byte((opcode) & 0xff))
+
+			// Indicate "LEAQ [RIP + 32bit displacement].
+			// https://wiki.osdev.org/X86-64_Instruction_Encoding#32.2F64-bit_addressing
+			c.EmitByte(encodeModRM(0b00, dst.encoding(), 0b101))
+
+			// This will be resolved later, so we just emit a placeholder (0xffffffff for testing).
+			c.Emit4Bytes(0xffffffff)
+		default:
+			panic("BUG: invalid operand kind")
+		}
 
 	case movRM:
 		m := i.op2.addressMode()
@@ -1087,7 +1104,7 @@ func (i *instruction) encode(c backend.Compiler) (needsLabelResolution bool) {
 
 		// Restore the RBP, RSP, and return to the Go code:
 		*allocatedAmode = amode{
-			kind: amodeImmReg, base: execCtx,
+			kindWithShift: uint32(amodeImmReg), base: execCtx,
 			imm32: wazevoapi.ExecutionContextOffsetOriginalFramePointer.U32(),
 		}
 		encodeLoad64(c, allocatedAmode, rbp)
@@ -1229,7 +1246,7 @@ func encodeEncMem(
 		useSBI = 4 // the encoding of rsp or r12 register.
 	)
 
-	switch m.kind {
+	switch m.kind() {
 	case amodeImmReg, amodeImmRBP:
 		base := m.base.RealReg()
 		baseEnc := regEncodings[base]
@@ -1287,30 +1304,17 @@ func encodeEncMem(
 		immZero, baseRbp, baseR13 := m.imm32 == 0, base == rbp, base == r13
 		if immZero && !baseRbp && !baseR13 { // rbp or r13 can't be used as base for without displacement encoding. (curious why? because it's interpreted as RIP relative addressing).
 			c.EmitByte(encodeModRM(modNoDisplacement, regEncoding(r), useSBI))
-			c.EmitByte(encodeSIB(m.shift, indexEnc.encoding(), baseEnc.encoding()))
+			c.EmitByte(encodeSIB(m.shift(), indexEnc.encoding(), baseEnc.encoding()))
 		} else if lower8willSignExtendTo32(m.imm32) {
 			c.EmitByte(encodeModRM(modShortDisplacement, regEncoding(r), useSBI))
-			c.EmitByte(encodeSIB(m.shift, indexEnc.encoding(), baseEnc.encoding()))
+			c.EmitByte(encodeSIB(m.shift(), indexEnc.encoding(), baseEnc.encoding()))
 			c.EmitByte(byte(m.imm32))
 		} else {
 			c.EmitByte(encodeModRM(modLongDisplacement, regEncoding(r), useSBI))
-			c.EmitByte(encodeSIB(m.shift, indexEnc.encoding(), baseEnc.encoding()))
+			c.EmitByte(encodeSIB(m.shift(), indexEnc.encoding(), baseEnc.encoding()))
 			c.Emit4Bytes(m.imm32)
 		}
 
-	case amodeRipRelative:
-		rex.encode(c, regRexBit(r), 0)
-		for opcodeNum > 0 {
-			opcodeNum--
-			c.EmitByte(byte((opcodes >> (opcodeNum << 3)) & 0xff))
-		}
-
-		// Indicate "LEAQ [RIP + 32bit displacement].
-		// https://wiki.osdev.org/X86-64_Instruction_Encoding#32.2F64-bit_addressing
-		c.EmitByte(encodeModRM(0b00, regEncoding(r), 0b101))
-
-		// This will be resolved later, so we just emit a placeholder. (m.imm32 for testing).
-		c.Emit4Bytes(m.imm32)
 	default:
 		panic("BUG: invalid addressing mode")
 	}
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine.go b/internal/engine/wazevo/backend/isa/amd64/machine.go
index ca9b38dc..91aedc9e 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine.go
@@ -27,7 +27,7 @@ func NewBackend() backend.Machine {
 		cpuFeatures: platform.CpuFeatures,
 		regAlloc:    regalloc.NewAllocator(regInfo),
 		spillSlots:  map[regalloc.VRegID]int64{},
-		amodePool:   wazevoapi.NewPool[amode](resetAmode),
+		amodePool:   wazevoapi.NewPool[amode](nil),
 	}
 }
 
@@ -196,7 +196,7 @@ func (m *machine) lowerBrTable(index ssa.Value, targets []ssa.BasicBlock) {
 
 	jmpTableBegin, jmpTableBeginLabel := m.allocateBrTarget()
 	m.insert(jmpTableBegin)
-	leaJmpTableAddr.asLEA(m.newAmodeRipRelative(jmpTableBeginLabel), addr)
+	leaJmpTableAddr.asLEA(newOperandLabel(jmpTableBeginLabel), addr)
 
 	jmpTable := m.allocateInstr()
 	targetSliceIndex := m.addJmpTableTarget(targets)
@@ -648,7 +648,7 @@ func (m *machine) lowerVconst(res ssa.Value, lo, hi uint64) {
 	m.insert(constIsland)
 	m.insert(afterLoadNop)
 
-	lea.asLEA(m.newAmodeRipRelative(constLabel), islandAddr)
+	lea.asLEA(newOperandLabel(constLabel), islandAddr)
 	jmp.asJmp(newOperandLabel(afterLoadLabel))
 }
 
@@ -891,7 +891,7 @@ func (m *machine) lowerExitWithCode(execCtx regalloc.VReg, code wazevoapi.ExitCo
 
 	// Insert the label for the return address.
 	nop, l := m.allocateBrTarget()
-	readRip.asLEA(m.newAmodeRipRelative(l), ripReg)
+	readRip.asLEA(newOperandLabel(l), ripReg)
 	m.insert(nop)
 	return l
 }
diff --git a/internal/engine/wazevo/backend/isa/amd64/operands.go b/internal/engine/wazevo/backend/isa/amd64/operands.go
index 2efc9860..e6a5a1f8 100644
--- a/internal/engine/wazevo/backend/isa/amd64/operands.go
+++ b/internal/engine/wazevo/backend/isa/amd64/operands.go
@@ -89,13 +89,9 @@ func (o *operand) label() backend.Label {
 	switch o.kind {
 	case operandKindLabel:
 		return backend.Label(o.data)
-	case operandKindMem:
-		mem := o.addressMode()
-		if mem.kind == amodeRipRelative {
-			return mem.label
-		}
+	default:
+		panic("BUG: invalid operand kind")
 	}
-	panic("BUG: invalid operand kind")
 }
 
 func newOperandLabel(label backend.Label) operand {
@@ -116,22 +112,14 @@ func newOperandMem(amode *amode) operand {
 
 // amode is a memory operand (addressing mode).
 type amode struct {
-	kind  amodeKind
-	imm32 uint32
-	base  regalloc.VReg
+	kindWithShift uint32
+	imm32         uint32
+	base          regalloc.VReg
 
 	// For amodeRegRegShift:
 	index regalloc.VReg
-	shift byte // 0, 1, 2, 3
-
-	// For amodeRipRelative.
-	// If kind == amodeRipRelative, and label is invalid,
-	// then imm32 should represent the resolved address.
-	label backend.Label
 }
 
-func resetAmode(am *amode) {}
-
 type amodeKind byte
 
 const (
@@ -146,31 +134,36 @@ const (
 	// amodeRegRegShift calculates sign-extend-32-to-64(Immediate) + base + (Register2 << Shift)
 	amodeRegRegShift
 
-	// amodeRipRelative is a memory operand with RIP-relative addressing mode.
-	amodeRipRelative
-
 	// TODO: there are other addressing modes such as the one without base register.
 )
 
+func (a *amode) kind() amodeKind {
+	return amodeKind(a.kindWithShift & 0xff)
+}
+
+func (a *amode) shift() byte {
+	return byte(a.kindWithShift >> 8)
+}
+
 func (a *amode) uses(rs *[]regalloc.VReg) {
-	switch a.kind {
+	switch a.kind() {
 	case amodeImmReg:
 		*rs = append(*rs, a.base)
 	case amodeRegRegShift:
 		*rs = append(*rs, a.base, a.index)
-	case amodeRipRelative, amodeImmRBP:
+	case amodeImmRBP:
 	default:
 		panic("BUG: invalid amode kind")
 	}
 }
 
 func (a *amode) nregs() int {
-	switch a.kind {
+	switch a.kind() {
 	case amodeImmReg:
 		return 1
 	case amodeRegRegShift:
 		return 2
-	case amodeRipRelative, amodeImmRBP:
+	case amodeImmRBP:
 		return 0
 	default:
 		panic("BUG: invalid amode kind")
@@ -178,7 +171,7 @@ func (a *amode) nregs() int {
 }
 
 func (a *amode) assignUses(i int, reg regalloc.VReg) {
-	switch a.kind {
+	switch a.kind() {
 	case amodeImmReg:
 		if i == 0 {
 			a.base = reg
@@ -193,7 +186,6 @@ func (a *amode) assignUses(i int, reg regalloc.VReg) {
 		} else {
 			panic("BUG: invalid amode assignment")
 		}
-	case amodeRipRelative:
 	default:
 		panic("BUG: invalid amode assignment")
 	}
@@ -201,13 +193,13 @@ func (a *amode) assignUses(i int, reg regalloc.VReg) {
 
 func (m *machine) newAmodeImmReg(imm32 uint32, base regalloc.VReg) *amode {
 	ret := m.amodePool.Allocate()
-	*ret = amode{kind: amodeImmReg, imm32: imm32, base: base}
+	*ret = amode{kindWithShift: uint32(amodeImmReg), imm32: imm32, base: base}
 	return ret
 }
 
 func (m *machine) newAmodeImmRBPReg(imm32 uint32) *amode {
 	ret := m.amodePool.Allocate()
-	*ret = amode{kind: amodeImmRBP, imm32: imm32, base: rbpVReg}
+	*ret = amode{kindWithShift: uint32(amodeImmRBP), imm32: imm32, base: rbpVReg}
 	return ret
 }
 
@@ -216,47 +208,28 @@ func (m *machine) newAmodeRegRegShift(imm32 uint32, base, index regalloc.VReg, s
 		panic(fmt.Sprintf("BUG: invalid shift (must be 3>=): %d", shift))
 	}
 	ret := m.amodePool.Allocate()
-	*ret = amode{kind: amodeRegRegShift, imm32: imm32, base: base, index: index, shift: shift}
-	return ret
-}
-
-func (a *amode) resolveRipRelative(imm32 uint32) {
-	if a.kind != amodeRipRelative {
-		panic("BUG: invalid amode kind")
-	}
-	a.imm32 = imm32
-	a.label = backend.LabelInvalid
-}
-
-func (m *machine) newAmodeRipRelative(label backend.Label) *amode {
-	ret := m.amodePool.Allocate()
-	*ret = amode{kind: amodeRipRelative, label: label}
+	*ret = amode{kindWithShift: uint32(amodeRegRegShift) | uint32(shift)<<8, imm32: imm32, base: base, index: index}
 	return ret
 }
 
 // String implements fmt.Stringer.
 func (a *amode) String() string {
-	switch a.kind {
+	switch a.kind() {
 	case amodeImmReg, amodeImmRBP:
 		if a.imm32 == 0 {
 			return fmt.Sprintf("(%s)", formatVRegSized(a.base, true))
 		}
 		return fmt.Sprintf("%d(%s)", int32(a.imm32), formatVRegSized(a.base, true))
 	case amodeRegRegShift:
+		shift := 1 << a.shift()
 		if a.imm32 == 0 {
 			return fmt.Sprintf(
 				"(%s,%s,%d)",
-				formatVRegSized(a.base, true), formatVRegSized(a.index, true), 1<<a.shift)
+				formatVRegSized(a.base, true), formatVRegSized(a.index, true), shift)
 		}
 		return fmt.Sprintf(
 			"%d(%s,%s,%d)",
-			int32(a.imm32), formatVRegSized(a.base, true), formatVRegSized(a.index, true), 1<<a.shift)
-	case amodeRipRelative:
-		if a.label != backend.LabelInvalid {
-			return fmt.Sprintf("%s(%%rip)", a.label)
-		} else {
-			return fmt.Sprintf("%d(%%rip)", int32(a.imm32))
-		}
+			int32(a.imm32), formatVRegSized(a.base, true), formatVRegSized(a.index, true), shift)
 	default:
 		panic("BUG: invalid amode kind")
 	}
diff --git a/internal/engine/wazevo/wazevoapi/pool.go b/internal/engine/wazevo/wazevoapi/pool.go
index 1f50c6a9..7adcd79c 100644
--- a/internal/engine/wazevo/wazevoapi/pool.go
+++ b/internal/engine/wazevo/wazevoapi/pool.go
@@ -39,7 +39,9 @@ func (p *Pool[T]) Allocate() *T {
 		p.index = 0
 	}
 	ret := &p.pages[len(p.pages)-1][p.index]
-	p.resetFn(ret)
+	if p.resetFn != nil {
+		p.resetFn(ret)
+	}
 	p.index++
 	p.allocated++
 	return ret


Test patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/abi_go_call_test.go b/internal/engine/wazevo/backend/isa/amd64/abi_go_call_test.go
index 45da2427..67db16b2 100644
--- a/internal/engine/wazevo/backend/isa/amd64/abi_go_call_test.go
+++ b/internal/engine/wazevo/backend/isa/amd64/abi_go_call_test.go
@@ -72,7 +72,7 @@ L2:
 	mov.l %r12, (%rax)
 	mov.q %rsp, 56(%rax)
 	mov.q %rbp, 1152(%rax)
-	lea L3(%rip), %r12
+	lea L3, %r12
 	mov.q %r12, 48(%rax)
 	exit_sequence %rax
 L3:
@@ -147,7 +147,7 @@ L2:
 	mov.l %r12, (%rax)
 	mov.q %rsp, 56(%rax)
 	mov.q %rbp, 1152(%rax)
-	lea L3(%rip), %r12
+	lea L3, %r12
 	mov.q %r12, 48(%rax)
 	exit_sequence %rax
 L3:
@@ -213,7 +213,7 @@ L2:
 	mov.l %r12, (%rax)
 	mov.q %rsp, 56(%rax)
 	mov.q %rbp, 1152(%rax)
-	lea L3(%rip), %r12
+	lea L3, %r12
 	mov.q %r12, 48(%rax)
 	exit_sequence %rax
 L3:
@@ -309,7 +309,7 @@ L2:
 	mov.l %r12, (%rax)
 	mov.q %rsp, 56(%rax)
 	mov.q %rbp, 1152(%rax)
-	lea L3(%rip), %r12
+	lea L3, %r12
 	mov.q %r12, 48(%rax)
 	exit_sequence %rax
 L3:
@@ -435,7 +435,7 @@ func TestMachine_CompileStackGrowCallSequence(t *testing.T) {
 	mov.l %r12, (%rax)
 	mov.q %rsp, 56(%rax)
 	mov.q %rbp, 1152(%rax)
-	lea L1(%rip), %r12
+	lea L1, %r12
 	mov.q %r12, 48(%rax)
 	exit_sequence %rax
 L1:
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 a34b37fd..7140b82c 100644
--- a/internal/engine/wazevo/backend/isa/amd64/instr_encoding_test.go
+++ b/internal/engine/wazevo/backend/isa/amd64/instr_encoding_test.go
@@ -5,6 +5,7 @@ import (
 	"runtime"
 	"testing"
 
+	"github.com/tetratelabs/wazero/internal/engine/wazevo/backend"
 	"github.com/tetratelabs/wazero/internal/testing/require"
 )
 
@@ -13,7 +14,6 @@ func TestInstruction_format_encode(t *testing.T) {
 	defer func() { runtime.KeepAlive(m) }()
 	newAmodeImmReg := m.newAmodeImmReg
 	newAmodeRegRegShift := m.newAmodeRegRegShift
-	newAmodeRipRelative := m.newAmodeRipRelative
 	allocateExitSeq := m.allocateExitSeq
 
 	for _, tc := range []struct {
@@ -1842,37 +1842,28 @@ func TestInstruction_format_encode(t *testing.T) {
 			wantFormat: "xor %r11d, %r15d",
 		},
 		{
-			setup:      func(i *instruction) { i.asLEA(newAmodeImmReg(0, rdiVReg), rdxVReg) },
+			setup:      func(i *instruction) { i.asLEA(newOperandMem(newAmodeImmReg(0, rdiVReg)), rdxVReg) },
 			want:       "488d17",
 			wantFormat: "lea (%rdi), %rdx",
 		},
 		{
-			setup:      func(i *instruction) { i.asLEA(newAmodeImmReg(0xffff, rdiVReg), rdxVReg) },
+			setup:      func(i *instruction) { i.asLEA(newOperandMem(newAmodeImmReg(0xffff, rdiVReg)), rdxVReg) },
 			want:       "488d97ffff0000",
 			wantFormat: "lea 65535(%rdi), %rdx",
 		},
-		{
-			setup:      func(i *instruction) { i.asLEA(newAmodeRegRegShift(0xffff, rspVReg, r13VReg, 3), rdxVReg) },
-			want:       "4a8d94ecffff0000",
-			wantFormat: "lea 65535(%rsp,%r13,8), %rdx",
-		},
 		{
 			setup: func(i *instruction) {
-				a := newAmodeRipRelative(1)
-				a.resolveRipRelative(1234)
-				i.asLEA(a, r11VReg)
+				i.asLEA(newOperandMem(newAmodeRegRegShift(0xffff, rspVReg, r13VReg, 3)), rdxVReg)
 			},
-			want:       "4c8d1dd2040000",
-			wantFormat: "lea 1234(%rip), %r11",
+			want:       "4a8d94ecffff0000",
+			wantFormat: "lea 65535(%rsp,%r13,8), %rdx",
 		},
 		{
 			setup: func(i *instruction) {
-				a := newAmodeRipRelative(1)
-				a.resolveRipRelative(8)
-				i.asLEA(a, r11VReg)
+				i.asLEA(newOperandLabel(backend.Label(1234)), r11VReg)
 			},
-			want:       "4c8d1d08000000",
-			wantFormat: "lea 8(%rip), %r11",
+			want:       "4c8d1dffffffff",
+			wantFormat: "lea L1234, %r11",
 		},
 		{
 			setup:      func(i *instruction) { i.kind = ud2 },
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine_test.go b/internal/engine/wazevo/backend/isa/amd64/machine_test.go
index 5d03eb50..588cc8cb 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine_test.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine_test.go
@@ -288,7 +288,7 @@ func TestMachine_lowerExitWithCode(t *testing.T) {
 	mov.q %rbp, 1152(%r15)
 	movl $6, %ebp
 	mov.l %rbp, (%r15)
-	lea L1(%rip), %rbp
+	lea L1, %rbp
 	mov.q %rbp, 48(%r15)
 	exit_sequence %r15
 L1: