Repository: tetratelabs/wazero
Pull request #2092 (2024-02-26 01:09:05)

wazevo(backend): block arg as one instruction before regalloc
on arm64, constant block args were loaded via multiple instructions,
which makes it harder for me to optimize the regalloc's bottleneck.
Therefore, this introduces load_const_block_arg pseudo instruction,
which will be lowered at post-regalloc into multiple instructions when needed.


Patch:
diff --git a/internal/engine/wazevo/backend/compiler_lower.go b/internal/engine/wazevo/backend/compiler_lower.go
index 165786a9..80e65668 100644
--- a/internal/engine/wazevo/backend/compiler_lower.go
+++ b/internal/engine/wazevo/backend/compiler_lower.go
@@ -221,6 +221,6 @@ func (c *compiler) lowerBlockArguments(args []ssa.Value, succ ssa.BasicBlock) {
 	// Finally, move the constants.
 	for _, edge := range c.constEdges {
 		cInst, dst := edge.cInst, edge.dst
-		c.mach.InsertLoadConstant(cInst, dst)
+		c.mach.InsertLoadConstantBlockArg(cInst, dst)
 	}
 }
diff --git a/internal/engine/wazevo/backend/isa/amd64/abi.go b/internal/engine/wazevo/backend/isa/amd64/abi.go
index 6bd3dc01..130f8c62 100644
--- a/internal/engine/wazevo/backend/isa/amd64/abi.go
+++ b/internal/engine/wazevo/backend/isa/amd64/abi.go
@@ -136,7 +136,7 @@ func (m *machine) LowerReturn(ret ssa.Value, r *backend.ABIArg) {
 	if def := m.c.ValueDefinition(ret); def.IsFromInstr() {
 		// Constant instructions are inlined.
 		if inst := def.Instr; inst.Constant() {
-			m.InsertLoadConstant(inst, reg)
+			m.insertLoadConstant(inst, reg)
 		}
 	}
 	if r.Kind == backend.ABIArgKindReg {
diff --git a/internal/engine/wazevo/backend/isa/amd64/lower_constant.go b/internal/engine/wazevo/backend/isa/amd64/lower_constant.go
index 01e0d1db..55d05ef6 100644
--- a/internal/engine/wazevo/backend/isa/amd64/lower_constant.go
+++ b/internal/engine/wazevo/backend/isa/amd64/lower_constant.go
@@ -11,12 +11,16 @@ func (m *machine) lowerConstant(instr *ssa.Instruction) (vr regalloc.VReg) {
 	valType := val.Type()
 
 	vr = m.c.AllocateVReg(valType)
-	m.InsertLoadConstant(instr, vr)
+	m.insertLoadConstant(instr, vr)
 	return
 }
 
-// InsertLoadConstant implements backend.Machine.
-func (m *machine) InsertLoadConstant(instr *ssa.Instruction, vr regalloc.VReg) {
+// InsertLoadConstantBlockArg implements backend.Machine.
+func (m *machine) InsertLoadConstantBlockArg(instr *ssa.Instruction, vr regalloc.VReg) {
+	m.insertLoadConstant(instr, vr)
+}
+
+func (m *machine) insertLoadConstant(instr *ssa.Instruction, vr regalloc.VReg) {
 	val := instr.Return()
 	valType := val.Type()
 	v := instr.ConstantVal()
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine.go b/internal/engine/wazevo/backend/isa/amd64/machine.go
index c1df70c5..09ccbe00 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine.go
@@ -1642,7 +1642,7 @@ func (m *machine) callerGenVRegToFunctionArg(a *backend.FunctionABI, argIndex in
 	if def != nil && def.IsFromInstr() {
 		// Constant instructions are inlined.
 		if inst := def.Instr; inst.Constant() {
-			m.InsertLoadConstant(inst, reg)
+			m.insertLoadConstant(inst, reg)
 		}
 	}
 	if arg.Kind == backend.ABIArgKindReg {
diff --git a/internal/engine/wazevo/backend/isa/arm64/abi.go b/internal/engine/wazevo/backend/isa/arm64/abi.go
index f7aeb573..ce868eec 100644
--- a/internal/engine/wazevo/backend/isa/arm64/abi.go
+++ b/internal/engine/wazevo/backend/isa/arm64/abi.go
@@ -130,7 +130,10 @@ func (m *machine) LowerReturns(rets []ssa.Value) {
 		if def := m.compiler.ValueDefinition(ret); def.IsFromInstr() {
 			// Constant instructions are inlined.
 			if inst := def.Instr; inst.Constant() {
-				m.InsertLoadConstant(inst, reg)
+				val := inst.Return()
+				valType := val.Type()
+				v := inst.ConstantVal()
+				m.insertLoadConstant(v, valType, reg)
 			}
 		}
 		if r.Kind == backend.ABIArgKindReg {
@@ -182,7 +185,10 @@ func (m *machine) callerGenVRegToFunctionArg(a *backend.FunctionABI, argIndex in
 	if def != nil && def.IsFromInstr() {
 		// Constant instructions are inlined.
 		if inst := def.Instr; inst.Constant() {
-			m.InsertLoadConstant(inst, reg)
+			val := inst.Return()
+			valType := val.Type()
+			v := inst.ConstantVal()
+			m.insertLoadConstant(v, valType, reg)
 		}
 	}
 	if arg.Kind == backend.ABIArgKindReg {
diff --git a/internal/engine/wazevo/backend/isa/arm64/instr.go b/internal/engine/wazevo/backend/isa/arm64/instr.go
index 10d28223..8aabc599 100644
--- a/internal/engine/wazevo/backend/isa/arm64/instr.go
+++ b/internal/engine/wazevo/backend/isa/arm64/instr.go
@@ -164,6 +164,7 @@ var defKinds = [numInstructionKinds]defKind{
 	atomicLoad:           defKindRD,
 	atomicStore:          defKindNone,
 	dmb:                  defKindNone,
+	loadConstBlockArg:    defKindRD,
 }
 
 // Defs returns the list of regalloc.VReg that are defined by the instruction.
@@ -300,6 +301,7 @@ var useKinds = [numInstructionKinds]useKind{
 	atomicCas:            useKindRDRewrite,
 	atomicLoad:           useKindRN,
 	atomicStore:          useKindRNRM,
+	loadConstBlockArg:    useKindNone,
 	dmb:                  useKindNone,
 }
 
@@ -1516,6 +1518,8 @@ func (i *instruction) String() (str string) {
 		str = fmt.Sprintf("source_offset_info %d", ssa.SourceOffset(i.u1))
 	case vecLoad1R:
 		str = fmt.Sprintf("ld1r {%s}, [%s]", formatVRegVec(i.rd.nr(), vecArrangement(i.u1), vecIndexNone), formatVRegSized(i.rn.nr(), 64))
+	case loadConstBlockArg:
+		str = fmt.Sprintf("load_const_block_arg %s, %#x", formatVRegSized(i.rd.nr(), 64), i.u1)
 	default:
 		panic(i.kind)
 	}
@@ -1736,6 +1740,8 @@ const (
 	dmb
 	// UDF is the undefined instruction. For debugging only.
 	udf
+	// loadConstBlockArg represents a load of a constant block argument.
+	loadConstBlockArg
 
 	// emitSourceOffsetInfo is a dummy instruction to emit source offset info.
 	// The existence of this instruction does not affect the execution.
@@ -1745,6 +1751,18 @@ const (
 	numInstructionKinds
 )
 
+func (i *instruction) asLoadConstBlockArg(v uint64, typ ssa.Type, dst regalloc.VReg) *instruction {
+	i.kind = loadConstBlockArg
+	i.u1 = v
+	i.u2 = uint64(typ)
+	i.rd = operandNR(dst)
+	return i
+}
+
+func (i *instruction) loadConstBlockArgData() (v uint64, typ ssa.Type, dst regalloc.VReg) {
+	return i.u1, ssa.Type(i.u2), i.rd.nr()
+}
+
 func (i *instruction) asEmitSourceOffsetInfo(l ssa.SourceOffset) *instruction {
 	i.kind = emitSourceOffsetInfo
 	i.u1 = uint64(l)
@@ -2365,7 +2383,7 @@ func (i *instruction) size() int64 {
 	switch i.kind {
 	case exitSequence:
 		return exitSequenceSize // 5 instructions as in encodeExitSequence.
-	case nop0:
+	case nop0, loadConstBlockArg:
 		return 0
 	case emitSourceOffsetInfo:
 		return 0
diff --git a/internal/engine/wazevo/backend/isa/arm64/instr_encoding.go b/internal/engine/wazevo/backend/isa/arm64/instr_encoding.go
index 471a8399..513ad6fb 100644
--- a/internal/engine/wazevo/backend/isa/arm64/instr_encoding.go
+++ b/internal/engine/wazevo/backend/isa/arm64/instr_encoding.go
@@ -24,7 +24,7 @@ func (m *machine) encode(root *instruction) {
 func (i *instruction) encode(m *machine) {
 	c := m.compiler
 	switch kind := i.kind; kind {
-	case nop0, emitSourceOffsetInfo:
+	case nop0, emitSourceOffsetInfo, loadConstBlockArg:
 	case exitSequence:
 		encodeExitSequence(c, i.rn.reg())
 	case ret:
diff --git a/internal/engine/wazevo/backend/isa/arm64/lower_constant.go b/internal/engine/wazevo/backend/isa/arm64/lower_constant.go
index 5516c0b3..698b382d 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_constant.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_constant.go
@@ -11,16 +11,27 @@ func (m *machine) lowerConstant(instr *ssa.Instruction) (vr regalloc.VReg) {
 	valType := val.Type()
 
 	vr = m.compiler.AllocateVReg(valType)
-	m.InsertLoadConstant(instr, vr)
+	v := instr.ConstantVal()
+	m.insertLoadConstant(v, valType, vr)
 	return
 }
 
-// InsertLoadConstant implements backend.Machine.
-func (m *machine) InsertLoadConstant(instr *ssa.Instruction, vr regalloc.VReg) {
+// InsertLoadConstantBlockArg implements backend.Machine.
+func (m *machine) InsertLoadConstantBlockArg(instr *ssa.Instruction, vr regalloc.VReg) {
 	val := instr.Return()
 	valType := val.Type()
 	v := instr.ConstantVal()
+	load := m.allocateInstr()
+	load.asLoadConstBlockArg(v, valType, vr)
+	m.insert(load)
+}
+
+func (m *machine) lowerLoadConstantBlockArgAfterRegAlloc(i *instruction) {
+	v, typ, dst := i.loadConstBlockArgData()
+	m.insertLoadConstant(v, typ, dst)
+}
 
+func (m *machine) insertLoadConstant(v uint64, valType ssa.Type, vr regalloc.VReg) {
 	if valType.Bits() < 64 { // Clear the redundant bits just in case it's unexpectedly sign-extended, etc.
 		v = v & ((1 << valType.Bits()) - 1)
 	}
diff --git a/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue.go b/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue.go
index df1ce7d6..466fac46 100644
--- a/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue.go
+++ b/internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue.go
@@ -198,18 +198,28 @@ func (m *machine) createFrameSizeSlot(cur *instruction, s int64) *instruction {
 func (m *machine) postRegAlloc() {
 	ectx := m.executableContext
 	for cur := ectx.RootInstr; cur != nil; cur = cur.next {
-		if cur.kind == ret {
+		switch cur.kind {
+		case ret:
 			m.setupEpilogueAfter(cur.prev)
-			continue
-		}
-
-		// Removes the redundant copy instruction.
-		if cur.IsCopy() && cur.rn.realReg() == cur.rd.realReg() {
-			prev, next := cur.prev, cur.next
-			// Remove the copy instruction.
-			prev.next = next
-			if next != nil {
-				next.prev = prev
+		case loadConstBlockArg:
+			lc := cur
+			next := lc.next
+			m.executableContext.PendingInstructions = m.executableContext.PendingInstructions[:0]
+			m.lowerLoadConstantBlockArgAfterRegAlloc(lc)
+			for _, instr := range m.executableContext.PendingInstructions {
+				cur = linkInstr(cur, instr)
+			}
+			linkInstr(cur, next)
+			m.executableContext.PendingInstructions = m.executableContext.PendingInstructions[:0]
+		default:
+			// Removes the redundant copy instruction.
+			if cur.IsCopy() && cur.rn.realReg() == cur.rd.realReg() {
+				prev, next := cur.prev, cur.next
+				// Remove the copy instruction.
+				prev.next = next
+				if next != nil {
+					next.prev = prev
+				}
 			}
 		}
 	}
diff --git a/internal/engine/wazevo/backend/machine.go b/internal/engine/wazevo/backend/machine.go
index c746cf1c..f56e3480 100644
--- a/internal/engine/wazevo/backend/machine.go
+++ b/internal/engine/wazevo/backend/machine.go
@@ -45,8 +45,8 @@ type (
 		// InsertReturn inserts the return instruction to return from the current function.
 		InsertReturn()
 
-		// InsertLoadConstant inserts the instruction(s) to load the constant value into the given regalloc.VReg.
-		InsertLoadConstant(instr *ssa.Instruction, vr regalloc.VReg)
+		// InsertLoadConstantBlockArg inserts the instruction(s) to load the constant value into the given regalloc.VReg.
+		InsertLoadConstantBlockArg(instr *ssa.Instruction, vr regalloc.VReg)
 
 		// Format returns the string representation of the currently compiled machine code.
 		// This is only for testing purpose.


Test patch:
diff --git a/internal/engine/wazevo/backend/backend_test.go b/internal/engine/wazevo/backend/backend_test.go
index 0cf8384f..895e8471 100644
--- a/internal/engine/wazevo/backend/backend_test.go
+++ b/internal/engine/wazevo/backend/backend_test.go
@@ -603,7 +603,7 @@ L4 (SSA Block: blk5):
 	ret
 L3 (SSA Block: blk4):
 L5 (SSA Block: blk3):
-	orr w131?, wzr, #0x1
+	load_const_block_arg x131?, 0x1
 	b L2
 `,
 			afterFinalizeARM64: `
@@ -619,6 +619,7 @@ L4 (SSA Block: blk5):
 	ret
 L3 (SSA Block: blk4):
 L5 (SSA Block: blk3):
+	load_const_block_arg x8, 0x1
 	orr w8, wzr, #0x1
 	mov x2, x8
 	b #-0x1c (L2)
diff --git a/internal/engine/wazevo/backend/machine_test.go b/internal/engine/wazevo/backend/machine_test.go
index d9b2b093..b2d8efe6 100644
--- a/internal/engine/wazevo/backend/machine_test.go
+++ b/internal/engine/wazevo/backend/machine_test.go
@@ -23,7 +23,6 @@ type mockMachine struct {
 	insertLoadConstant             func(instr *ssa.Instruction, vr regalloc.VReg)
 	format                         func() string
 	linkAdjacentBlocks             func(prev, next ssa.BasicBlock)
-	rinfo                          *regalloc.RegisterInfo
 }
 
 func (m mockMachine) ArgsResultsRegs() (argResultInts, argResultFloats []regalloc.RealReg) {
@@ -60,17 +59,6 @@ func (m mockMachine) ResolveRelocations(map[ssa.FuncRef]int, []byte, []Relocatio
 // PostRegAlloc implements Machine.SetupPrologue.
 func (m mockMachine) PostRegAlloc() {}
 
-// Function implements Machine.Function.
-func (m mockMachine) Function() (f regalloc.Function) { return }
-
-// RegisterInfo implements Machine.RegisterInfo.
-func (m mockMachine) RegisterInfo() *regalloc.RegisterInfo {
-	if m.rinfo != nil {
-		return m.rinfo
-	}
-	return &regalloc.RegisterInfo{}
-}
-
 // InsertReturn implements Machine.InsertReturn.
 func (m mockMachine) InsertReturn() { panic("TODO") }
 
@@ -131,8 +119,8 @@ func (m mockMachine) InsertMove(dst, src regalloc.VReg, typ ssa.Type) {
 	m.insertMove(dst, src)
 }
 
-// InsertLoadConstant implements Machine.InsertLoadConstant.
-func (m mockMachine) InsertLoadConstant(instr *ssa.Instruction, vr regalloc.VReg) {
+// InsertLoadConstantBlockArg implements Machine.InsertLoadConstantBlockArg.
+func (m mockMachine) InsertLoadConstantBlockArg(instr *ssa.Instruction, vr regalloc.VReg) {
 	m.insertLoadConstant(instr, vr)
 }