Repository: tetratelabs/wazero
Pull request #1795 (2023-10-18 05:43:00)

wazevo: reduces map usages
#1496 


Patch:
diff --git a/internal/engine/wazevo/backend/compiler.go b/internal/engine/wazevo/backend/compiler.go
index d6f321f2..2b1ed5b1 100644
--- a/internal/engine/wazevo/backend/compiler.go
+++ b/internal/engine/wazevo/backend/compiler.go
@@ -23,11 +23,8 @@ func newCompiler(ctx context.Context, mach Machine, builder ssa.Builder) *compil
 
 	c := &compiler{
 		mach: mach, ssaBuilder: builder,
-		alreadyLowered:  make(map[*ssa.Instruction]bool),
-		vRegSet:         make(map[regalloc.VReg]bool),
-		ssaTypeOfVRegID: make(map[regalloc.VRegID]ssa.Type),
-		nextVRegID:      0,
-		regAlloc:        regalloc.NewAllocator(mach.RegisterInfo(registerSetDebug)),
+		nextVRegID: 0,
+		regAlloc:   regalloc.NewAllocator(mach.RegisterInfo(registerSetDebug)),
 	}
 	mach.SetCompiler(c)
 	return c
@@ -75,10 +72,6 @@ type Compiler interface {
 	// AllocateVReg allocates a new virtual register of the given type.
 	AllocateVReg(typ ssa.Type) regalloc.VReg
 
-	// MarkLowered is used to mark the given instruction as already lowered
-	// which tells the compiler to skip it when traversing.
-	MarkLowered(inst *ssa.Instruction)
-
 	// ValueDefinition returns the definition of the given value.
 	ValueDefinition(ssa.Value) *SSAValueDefinition
 
@@ -127,25 +120,24 @@ type compiler struct {
 	// nextVRegID is the next virtual register ID to be allocated.
 	nextVRegID regalloc.VRegID
 	// ssaValueToVRegs maps ssa.ValueID to regalloc.VReg.
-	ssaValueToVRegs []regalloc.VReg
+	ssaValueToVRegs [] /* VRegID to */ regalloc.VReg
 	// ssaValueDefinitions maps ssa.ValueID to its definition.
 	ssaValueDefinitions []SSAValueDefinition
 	// ssaValueRefCounts is a cached list obtained by ssa.Builder.ValueRefCounts().
 	ssaValueRefCounts []int
 	// returnVRegs is the list of virtual registers that store the return values.
-	returnVRegs    []regalloc.VReg
-	alreadyLowered map[*ssa.Instruction]bool
-	regAlloc       regalloc.Allocator
-	varEdges       [][2]regalloc.VReg
-	varEdgeTypes   []ssa.Type
-	constEdges     []struct {
+	returnVRegs  []regalloc.VReg
+	regAlloc     regalloc.Allocator
+	varEdges     [][2]regalloc.VReg
+	varEdgeTypes []ssa.Type
+	constEdges   []struct {
 		cInst *ssa.Instruction
 		dst   regalloc.VReg
 	}
-	vRegSet         map[regalloc.VReg]bool
+	vRegSet         []bool
 	tempRegs        []regalloc.VReg
 	tmpVals         []ssa.Value
-	ssaTypeOfVRegID map[regalloc.VRegID]ssa.Type
+	ssaTypeOfVRegID [] /* VRegID to */ ssa.Type
 	buf             []byte
 	relocations     []RelocationInfo
 	sourceOffsets   []SourceOffsetInfo
@@ -284,45 +276,30 @@ func (c *compiler) assignVirtualRegisters() {
 func (c *compiler) AllocateVReg(typ ssa.Type) regalloc.VReg {
 	regType := regalloc.RegTypeOf(typ)
 	r := regalloc.VReg(c.nextVRegID).SetRegType(regType)
-	c.ssaTypeOfVRegID[r.ID()] = typ
+
+	id := r.ID()
+	if int(id) >= len(c.ssaTypeOfVRegID) {
+		c.ssaTypeOfVRegID = append(c.ssaTypeOfVRegID, make([]ssa.Type, id+1)...)
+	}
+	c.ssaTypeOfVRegID[id] = typ
 	c.nextVRegID++
 	return r
 }
 
 // Init implements Compiler.Init.
 func (c *compiler) Init() {
-	for i, v := range c.ssaValueToVRegs {
-		c.ssaValueToVRegs[i] = regalloc.VRegInvalid
-		delete(c.ssaTypeOfVRegID, v.ID())
-	}
 	c.currentGID = 0
 	c.nextVRegID = 0
 	c.returnVRegs = c.returnVRegs[:0]
 	c.mach.Reset()
 	c.varEdges = c.varEdges[:0]
 	c.constEdges = c.constEdges[:0]
-
-	for key := range c.alreadyLowered {
-		c.alreadyLowered[key] = false
-	}
-	c.resetVRegSet()
 	c.regAlloc.Reset()
 	c.buf = c.buf[:0]
 	c.sourceOffsets = c.sourceOffsets[:0]
 	c.relocations = c.relocations[:0]
 }
 
-func (c *compiler) resetVRegSet() {
-	for key := range c.vRegSet {
-		c.vRegSet[key] = false
-	}
-}
-
-// MarkLowered implements Compiler.MarkLowered.
-func (c *compiler) MarkLowered(inst *ssa.Instruction) {
-	c.alreadyLowered[inst] = true
-}
-
 // ValueDefinition implements Compiler.ValueDefinition.
 func (c *compiler) ValueDefinition(value ssa.Value) *SSAValueDefinition {
 	return &c.ssaValueDefinitions[value.ID()]
@@ -340,11 +317,7 @@ func (c *compiler) Format() string {
 
 // TypeOf implements Compiler.Format.
 func (c *compiler) TypeOf(v regalloc.VReg) ssa.Type {
-	typ, ok := c.ssaTypeOfVRegID[v.ID()]
-	if !ok {
-		panic(fmt.Sprintf("BUG: v%d is not a valid vreg", v.ID()))
-	}
-	return typ
+	return c.ssaTypeOfVRegID[v.ID()]
 }
 
 // MatchInstr implements Compiler.MatchInstr.
diff --git a/internal/engine/wazevo/backend/compiler_lower.go b/internal/engine/wazevo/backend/compiler_lower.go
index 3740ddd6..0b3d03a7 100644
--- a/internal/engine/wazevo/backend/compiler_lower.go
+++ b/internal/engine/wazevo/backend/compiler_lower.go
@@ -60,7 +60,7 @@ func (c *compiler) lowerBlock(blk ssa.BasicBlock) {
 	// Now start lowering the non-branching instructions.
 	for ; cur != nil; cur = cur.Prev() {
 		c.setCurrentGroupID(cur.GroupID())
-		if c.alreadyLowered[cur] {
+		if cur.Lowered() {
 			continue
 		}
 
@@ -160,17 +160,23 @@ func (c *compiler) lowerBlockArguments(args []ssa.Value, succ ssa.BasicBlock) {
 	}
 
 	// Check if there's an overlap among the dsts and srcs in varEdges.
-	c.resetVRegSet()
 	for _, edge := range c.varEdges {
-		src := edge[0]
+		src := edge[0].ID()
+		if int(src) >= len(c.vRegSet) {
+			c.vRegSet = append(c.vRegSet, make([]bool, src+1)...)
+		}
 		c.vRegSet[src] = true
 	}
 	separated := true
 	for _, edge := range c.varEdges {
-		dst := edge[1]
-		if c.vRegSet[dst] {
-			separated = false
-			break
+		dst := edge[1].ID()
+		if int(dst) >= len(c.vRegSet) {
+			c.vRegSet = append(c.vRegSet, make([]bool, dst+1)...)
+		} else {
+			if c.vRegSet[dst] {
+				separated = false
+				break
+			}
 		}
 	}
 
diff --git a/internal/engine/wazevo/backend/isa/arm64/lower_instr.go b/internal/engine/wazevo/backend/isa/arm64/lower_instr.go
index f3490b6a..319df83d 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_instr.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_instr.go
@@ -97,7 +97,7 @@ func (m *machine) LowerConditionalBranch(b *ssa.Instruction) {
 		cbr := m.allocateInstr()
 		cbr.asCondBr(cc.asCond(), target, false /* ignored */)
 		m.insert(cbr)
-		m.compiler.MarkLowered(cvalDef.Instr)
+		cvalDef.Instr.MarkLowered()
 	case m.compiler.MatchInstr(cvalDef, ssa.OpcodeFcmp): // This case we can use the Fpu flag directly.
 		cvalInstr := cvalDef.Instr
 		x, y, c := cvalInstr.FcmpData()
@@ -109,7 +109,7 @@ func (m *machine) LowerConditionalBranch(b *ssa.Instruction) {
 		cbr := m.allocateInstr()
 		cbr.asCondBr(cc.asCond(), target, false /* ignored */)
 		m.insert(cbr)
-		m.compiler.MarkLowered(cvalDef.Instr)
+		cvalDef.Instr.MarkLowered()
 	default:
 		rn := m.getOperand_NR(cvalDef, extModeNone)
 		var c cond
@@ -1807,7 +1807,7 @@ func (m *machine) lowerExitIfTrueWithCode(execCtxVReg regalloc.VReg, cond ssa.Va
 	if !m.compiler.MatchInstr(condDef, ssa.OpcodeIcmp) {
 		panic("TODO: OpcodeExitIfTrueWithCode must come after Icmp at the moment: " + condDef.Instr.Opcode().String())
 	}
-	m.compiler.MarkLowered(condDef.Instr)
+	condDef.Instr.MarkLowered()
 
 	cvalInstr := condDef.Instr
 	x, y, c := cvalInstr.IcmpData()
@@ -1833,13 +1833,13 @@ func (m *machine) lowerSelect(c, x, y, result ssa.Value) {
 		x, y, c := cvalInstr.IcmpData()
 		cc = condFlagFromSSAIntegerCmpCond(c)
 		m.lowerIcmpToFlag(x, y, c.Signed())
-		m.compiler.MarkLowered(cvalDef.Instr)
+		cvalDef.Instr.MarkLowered()
 	case m.compiler.MatchInstr(cvalDef, ssa.OpcodeFcmp): // This case we can use the Fpu flag directly.
 		cvalInstr := cvalDef.Instr
 		x, y, c := cvalInstr.FcmpData()
 		cc = condFlagFromSSAFloatCmpCond(c)
 		m.lowerFcmpToFlag(x, y)
-		m.compiler.MarkLowered(cvalDef.Instr)
+		cvalDef.Instr.MarkLowered()
 	default:
 		rn := m.getOperand_NR(cvalDef, extModeNone)
 		if c.Type() != ssa.TypeI32 && c.Type() != ssa.TypeI64 {
diff --git a/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands.go b/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands.go
index 2e1fa4a0..d9fbf178 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands.go
@@ -170,7 +170,7 @@ func (m *machine) getOperand_Imm12_ER_SR_NR(def *backend.SSAValueDefinition, mod
 	instr := def.Instr
 	if instr.Opcode() == ssa.OpcodeIconst {
 		if imm12Op, ok := asImm12Operand(instr.ConstantVal()); ok {
-			m.compiler.MarkLowered(instr)
+			instr.MarkLowered()
 			return imm12Op
 		}
 	}
@@ -188,7 +188,7 @@ func (m *machine) getOperand_MaybeNegatedImm12_ER_SR_NR(def *backend.SSAValueDef
 	if instr.Opcode() == ssa.OpcodeIconst {
 		c := instr.ConstantVal()
 		if imm12Op, ok := asImm12Operand(c); ok {
-			m.compiler.MarkLowered(instr)
+			instr.MarkLowered()
 			return imm12Op, false
 		}
 
@@ -198,7 +198,7 @@ func (m *machine) getOperand_MaybeNegatedImm12_ER_SR_NR(def *backend.SSAValueDef
 		}
 		negatedWithoutSign := -signExtended
 		if imm12Op, ok := asImm12Operand(uint64(negatedWithoutSign)); ok {
-			m.compiler.MarkLowered(instr)
+			instr.MarkLowered()
 			return imm12Op, true
 		}
 	}
@@ -223,7 +223,7 @@ func (m *machine) getOperand_ER_SR_NR(def *backend.SSAValueDefinition, mode extM
 			eop := extendOpFrom(signed, innerExtFromBits)
 			extArg := m.getOperand_NR(m.compiler.ValueDefinition(extInstr.Arg()), extModeNone)
 			op = operandER(extArg.nr(), eop, innerExtToBits)
-			m.compiler.MarkLowered(extInstr) // We merged the instruction in the operand.
+			extInstr.MarkLowered()
 			return
 		}
 
@@ -236,7 +236,7 @@ func (m *machine) getOperand_ER_SR_NR(def *backend.SSAValueDefinition, mode extM
 			// Two sign/zero extensions are equivalent to one sign/zero extension for the larger size.
 			eop := extendOpFrom(modeSigned, innerExtFromBits)
 			op = operandER(m.compiler.VRegOf(extInstr.Arg()), eop, modeBits)
-			m.compiler.MarkLowered(extInstr) // We merged the instruction in the operand.
+			extInstr.MarkLowered()
 		case (signed && !modeSigned) || (!signed && modeSigned):
 			// We need to {sign, zero}-extend the result of the {zero,sign} extension.
 			eop := extendOpFrom(modeSigned, innerExtToBits)
@@ -264,8 +264,8 @@ func (m *machine) getOperand_SR_NR(def *backend.SSAValueDefinition, mode extMode
 		if amountDef.IsFromInstr() && amountDef.Instr.Constant() {
 			// If that is the case, we can use the shifted register operand (SR).
 			c := byte(amountDef.Instr.ConstantVal()) & (targetVal.Type().Bits() - 1) // Clears the unnecessary bits.
-			m.compiler.MarkLowered(def.Instr)
-			m.compiler.MarkLowered(amountDef.Instr)
+			def.Instr.MarkLowered()
+			amountDef.Instr.MarkLowered()
 			return operandSR(targetVReg, c, shiftOpLSL)
 		}
 	}
@@ -298,7 +298,7 @@ func (m *machine) getOperand_NR(def *backend.SSAValueDefinition, mode extMode) (
 		if instr.Constant() {
 			// We inline all the constant instructions so that we could reduce the register usage.
 			v = m.lowerConstant(instr)
-			m.compiler.MarkLowered(instr)
+			instr.MarkLowered()
 		} else {
 			if n := def.N; n == 0 {
 				v = m.compiler.VRegOf(instr.Return())
diff --git a/internal/engine/wazevo/backend/isa/arm64/lower_mem.go b/internal/engine/wazevo/backend/isa/arm64/lower_mem.go
index c104a88b..a38bef5f 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_mem.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_mem.go
@@ -346,7 +346,7 @@ func (m *machine) collectAddends(ptr ssa.Value) (addends32 []addend32, addends64
 			// If the addend is an add, we recursively collect its operands.
 			x, y := def.Instr.Arg2()
 			m.addendsWorkQueue = append(m.addendsWorkQueue, x, y)
-			m.compiler.MarkLowered(def.Instr)
+			def.Instr.MarkLowered()
 		case ssa.OpcodeIconst:
 			// If the addend is constant, we just statically merge it into the offset.
 			ic := def.Instr
@@ -356,13 +356,13 @@ func (m *machine) collectAddends(ptr ssa.Value) (addends32 []addend32, addends64
 			} else {
 				offset += int64(u64)
 			}
-			m.compiler.MarkLowered(def.Instr)
+			def.Instr.MarkLowered()
 		case ssa.OpcodeUExtend, ssa.OpcodeSExtend:
 			switch input := def.Instr.Arg(); input.Type().Bits() {
 			case 64:
 				// If the input is already 64-bit, this extend is a no-op. TODO: shouldn't this be optimized out at much earlier stage? no?
 				m.addends64 = append(m.addends64, m.getOperand_NR(m.compiler.ValueDefinition(input), extModeNone).nr())
-				m.compiler.MarkLowered(def.Instr)
+				def.Instr.MarkLowered()
 				continue
 			case 32:
 				var ext extendOp
@@ -384,7 +384,7 @@ func (m *machine) collectAddends(ptr ssa.Value) (addends32 []addend32, addends64
 				default:
 					m.addends32 = append(m.addends32, addend32{r: m.getOperand_NR(inputDef, extModeNone).nr(), ext: ext})
 				}
-				m.compiler.MarkLowered(def.Instr)
+				def.Instr.MarkLowered()
 				continue
 			}
 			// If this is the extension smaller than 32 bits, this cannot be merged into addressing mode since
diff --git a/internal/engine/wazevo/ssa/instructions.go b/internal/engine/wazevo/ssa/instructions.go
index 8324f68e..694cb2da 100644
--- a/internal/engine/wazevo/ssa/instructions.go
+++ b/internal/engine/wazevo/ssa/instructions.go
@@ -27,11 +27,12 @@ type Instruction struct {
 	targets    []BasicBlock
 	prev, next *Instruction
 
-	rValue       Value
-	rValues      []Value
-	gid          InstructionGroupID
-	sourceOffset SourceOffset
-	live         bool
+	rValue         Value
+	rValues        []Value
+	gid            InstructionGroupID
+	sourceOffset   SourceOffset
+	live           bool
+	alreadyLowered bool
 }
 
 // SourceOffset represents the offset of the source of an instruction.
@@ -63,6 +64,16 @@ func (i *Instruction) GroupID() InstructionGroupID {
 	return i.gid
 }
 
+// MarkLowered marks this instruction as already lowered.
+func (i *Instruction) MarkLowered() {
+	i.alreadyLowered = true
+}
+
+// Lowered returns true if this instruction is already lowered.
+func (i *Instruction) Lowered() bool {
+	return i.alreadyLowered
+}
+
 // reset resets this instruction to the initial state.
 func (i *Instruction) reset() {
 	*i = Instruction{}


Test patch:
diff --git a/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands_test.go b/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands_test.go
index 4d25ca5a..914d256b 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands_test.go
@@ -142,10 +142,8 @@ func TestMachine_getOperand_SR_NR(t *testing.T) {
 		def = &backend.SSAValueDefinition{Instr: ishl, N: 0}
 		mode = extModeNone
 		verify = func(t *testing.T) {
-			_, ok := ctx.lowered[ishl]
-			require.True(t, ok)
-			_, ok = ctx.lowered[amount]
-			require.True(t, ok)
+			require.True(t, ishl.Lowered())
+			require.True(t, amount.Lowered())
 		}
 		return
 	}
@@ -233,10 +231,8 @@ func TestMachine_getOperand_SR_NR(t *testing.T) {
 				def = &backend.SSAValueDefinition{Instr: ishl, N: 0}
 				mode = extModeNone
 				verify = func(t *testing.T) {
-					_, ok := ctx.lowered[ishl]
-					require.True(t, ok)
-					_, ok = ctx.lowered[amount]
-					require.True(t, ok)
+					require.True(t, ishl.Lowered())
+					require.True(t, amount.Lowered())
 				}
 				return
 			},
@@ -265,10 +261,8 @@ func TestMachine_getOperand_SR_NR(t *testing.T) {
 				def = &backend.SSAValueDefinition{Instr: ishl, N: 0}
 				mode = extModeNone
 				verify = func(t *testing.T) {
-					_, ok := ctx.lowered[ishl]
-					require.True(t, ok)
-					_, ok = ctx.lowered[amount]
-					require.True(t, ok)
+					require.True(t, ishl.Lowered())
+					require.True(t, amount.Lowered())
 				}
 				return
 			},
@@ -280,7 +274,6 @@ func TestMachine_getOperand_SR_NR(t *testing.T) {
 			setup: func(ctx *mockCompiler, builder ssa.Builder, m *machine) (def *backend.SSAValueDefinition, mode extMode, verify func(t *testing.T)) {
 				def, mode, _ = ishlWithConstAmount(ctx, builder, m)
 				ctx.currentGID = 1230
-				verify = func(t *testing.T) { require.Equal(t, 0, len(ctx.lowered)) }
 				return
 			},
 			exp: operandNR(regalloc.VReg(10)),
@@ -290,7 +283,6 @@ func TestMachine_getOperand_SR_NR(t *testing.T) {
 			setup: func(ctx *mockCompiler, builder ssa.Builder, m *machine) (def *backend.SSAValueDefinition, mode extMode, verify func(t *testing.T)) {
 				def, mode, _ = ishlWithConstAmount(ctx, builder, m)
 				def.RefCount = 10
-				verify = func(t *testing.T) { require.Equal(t, 0, len(ctx.lowered)) }
 				return
 			},
 			exp: operandNR(regalloc.VReg(10)),
@@ -301,7 +293,9 @@ func TestMachine_getOperand_SR_NR(t *testing.T) {
 			def, mode, verify := tc.setup(ctx, b, m)
 			actual := m.getOperand_SR_NR(def, mode)
 			require.Equal(t, tc.exp, actual)
-			verify(t)
+			if verify != nil {
+				verify(t)
+			}
 			require.Equal(t, strings.Join(tc.instructions, "\n"), formatEmittedInstructionsInCurrentBlock(m))
 		})
 	}
@@ -369,8 +363,7 @@ func TestMachine_getOperand_ER_SR_NR(t *testing.T) {
 					ctx.definitions[v] = &backend.SSAValueDefinition{BlkParamVReg: regalloc.VReg(10), BlockParamValue: extArg}
 					def = &backend.SSAValueDefinition{Instr: ext, N: 0}
 					return def, extModeNone, func(t *testing.T) {
-						_, ok := ctx.lowered[ext]
-						require.True(t, ok)
+						require.True(t, ext.Lowered())
 					}
 				},
 				exp: c.exp,
@@ -589,8 +582,7 @@ func TestMachine_getOperand_ER_SR_NR(t *testing.T) {
 						}
 						def = &backend.SSAValueDefinition{Instr: ext, N: 0}
 						return def, c.mode, func(t *testing.T) {
-							_, ok := ctx.lowered[ext]
-							require.Equal(t, c.lowered, ok)
+							require.Equal(t, c.lowered, ext.Lowered())
 						}
 					},
 					exp:          c.exp,
diff --git a/internal/engine/wazevo/backend/isa/arm64/lower_instr_test.go b/internal/engine/wazevo/backend/isa/arm64/lower_instr_test.go
index 87da9174..278bc429 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_instr_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_instr_test.go
@@ -58,8 +58,7 @@ func TestMachine_LowerConditionalBranch(t *testing.T) {
 		}
 		builder.InsertInstruction(b)
 		return b, func(t *testing.T) {
-			_, ok := ctx.lowered[cmpInstr]
-			require.True(t, ok)
+			require.True(t, cmpInstr.Lowered())
 		}
 	}
 
@@ -92,8 +91,7 @@ func TestMachine_LowerConditionalBranch(t *testing.T) {
 		}
 		builder.InsertInstruction(b)
 		return b, func(t *testing.T) {
-			_, ok := ctx.lowered[icmp]
-			require.True(t, ok)
+			require.True(t, icmp.Lowered())
 		}
 	}
 
@@ -123,8 +121,7 @@ func TestMachine_LowerConditionalBranch(t *testing.T) {
 				// Indicate that currently compiling in the different group.
 				ctx.currentGID = 1000
 				return brz, func(t *testing.T) {
-					_, ok := ctx.lowered[icmp]
-					require.False(t, ok)
+					require.False(t, icmp.Lowered())
 				}
 			},
 			instructions: []string{"cbz w3?, (L1)"},
diff --git a/internal/engine/wazevo/backend/isa/arm64/lower_mem_test.go b/internal/engine/wazevo/backend/isa/arm64/lower_mem_test.go
index d5d1887d..6d61ed4a 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_mem_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_mem_test.go
@@ -355,8 +355,7 @@ func TestMachine_collectAddends(t *testing.T) {
 				iadd3 := insertIadd(ctx, b, iadd1.Return(), iadd2.Return())
 				return iadd3.Return(), func(t *testing.T) {
 					for _, instr := range []*ssa.Instruction{iadd1, iadd2, iadd3} {
-						_, lowerd := ctx.lowered[instr]
-						require.True(t, lowerd)
+						require.True(t, instr.Lowered())
 					}
 				}
 			},
@@ -371,8 +370,7 @@ func TestMachine_collectAddends(t *testing.T) {
 				iadd3 := insertIadd(ctx, b, iadd1.Return(), iadd2.Return())
 				return iadd3.Return(), func(t *testing.T) {
 					for _, instr := range []*ssa.Instruction{iadd1, iadd2, iadd3} {
-						_, lowerd := ctx.lowered[instr]
-						require.True(t, lowerd)
+						require.True(t, instr.Lowered())
 					}
 				}
 			},
@@ -390,8 +388,7 @@ func TestMachine_collectAddends(t *testing.T) {
 
 				return iadd4.Return(), func(t *testing.T) {
 					for _, instr := range []*ssa.Instruction{iadd1, iadd2, iadd3, iadd4} {
-						_, lowerd := ctx.lowered[instr]
-						require.True(t, lowerd)
+						require.True(t, instr.Lowered())
 					}
 					// Param must be zero-extended.
 					require.Equal(t, "uxtw x1?, w1000?", formatEmittedInstructionsInCurrentBlock(m))
@@ -411,8 +408,7 @@ func TestMachine_collectAddends(t *testing.T) {
 				iadd4 := insertIadd(ctx, b, param, ext.Return())
 				return iadd4.Return(), func(t *testing.T) {
 					for _, instr := range []*ssa.Instruction{ext, iadd4} {
-						_, lowerd := ctx.lowered[instr]
-						require.True(t, lowerd)
+						require.True(t, instr.Lowered())
 					}
 				}
 			},
@@ -430,8 +426,7 @@ func TestMachine_collectAddends(t *testing.T) {
 				iadd4 := insertIadd(ctx, b, param, ext.Return())
 				return iadd4.Return(), func(t *testing.T) {
 					for _, instr := range []*ssa.Instruction{ext, iadd4} {
-						_, lowerd := ctx.lowered[instr]
-						require.True(t, lowerd)
+						require.True(t, instr.Lowered())
 					}
 				}
 			},
@@ -449,8 +444,7 @@ func TestMachine_collectAddends(t *testing.T) {
 				iadd4 := insertIadd(ctx, b, param, ext.Return())
 				return iadd4.Return(), func(t *testing.T) {
 					for _, instr := range []*ssa.Instruction{ext, iadd4} {
-						_, lowerd := ctx.lowered[instr]
-						require.True(t, lowerd)
+						require.True(t, instr.Lowered())
 					}
 				}
 			},
@@ -466,8 +460,7 @@ func TestMachine_collectAddends(t *testing.T) {
 				iadd4 := insertIadd(ctx, b, param, ext.Return())
 				return iadd4.Return(), func(t *testing.T) {
 					for _, instr := range []*ssa.Instruction{ext, iadd4} {
-						_, lowerd := ctx.lowered[instr]
-						require.True(t, lowerd)
+						require.True(t, instr.Lowered())
 					}
 				}
 			},
diff --git a/internal/engine/wazevo/backend/isa/arm64/util_test.go b/internal/engine/wazevo/backend/isa/arm64/util_test.go
index 1ab1ae76..4a1e7200 100644
--- a/internal/engine/wazevo/backend/isa/arm64/util_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/util_test.go
@@ -55,7 +55,6 @@ type mockCompiler struct {
 	vRegCounter int
 	vRegMap     map[ssa.Value]regalloc.VReg
 	definitions map[ssa.Value]*backend.SSAValueDefinition
-	lowered     map[*ssa.Instruction]bool
 	sigs        map[ssa.SignatureID]*ssa.Signature
 	typeOf      map[regalloc.VReg]ssa.Type
 	relocs      []backend.RelocationInfo
@@ -89,7 +88,6 @@ func newMockCompilationContext() *mockCompiler {
 	return &mockCompiler{
 		vRegMap:     make(map[ssa.Value]regalloc.VReg),
 		definitions: make(map[ssa.Value]*backend.SSAValueDefinition),
-		lowered:     make(map[*ssa.Instruction]bool),
 	}
 }
 
@@ -105,11 +103,6 @@ func (m *mockCompiler) AllocateVReg(typ ssa.Type) regalloc.VReg {
 	return regalloc.VReg(m.vRegCounter).SetRegType(regType)
 }
 
-// MarkLowered implements backend.Compiler.
-func (m *mockCompiler) MarkLowered(inst *ssa.Instruction) {
-	m.lowered[inst] = true
-}
-
 // ValueDefinition implements backend.Compiler.
 func (m *mockCompiler) ValueDefinition(value ssa.Value) *backend.SSAValueDefinition {
 	definition, exists := m.definitions[value]