Repository: tetratelabs/wazero
Pull request #2230 (2024-06-05 16:00:04)

backend: removes the map use in ExecutableContext[Instr]
This makes the compilation slightly more faster with less memory:

```
benchstat old.txt new.txt
goos: darwin
goarch: arm64
pkg: github.com/tetratelabs/wazero
               │  old.txt   │             new.txt              │
               │   sec/op   │   sec/op    vs base              │
Compilation-10   2.488 ± 0%   2.438 ± 0%  -2.00% (p=0.001 n=7)

               │   old.txt    │              new.txt               │
               │     B/op     │     B/op      vs base              │
Compilation-10   341.1Mi ± 0%   340.3Mi ± 0%  -0.23% (p=0.001 n=7)

               │   old.txt   │              new.txt              │
               │  allocs/op  │  allocs/op   vs base              │
Compilation-10   607.0k ± 0%   605.5k ± 0%  -0.26% (p=0.001 n=7)
```

#2182 


Patch:
diff --git a/internal/engine/wazevo/backend/executable_context.go b/internal/engine/wazevo/backend/executable_context.go
index 81c6a6b6..8e9571b2 100644
--- a/internal/engine/wazevo/backend/executable_context.go
+++ b/internal/engine/wazevo/backend/executable_context.go
@@ -43,7 +43,7 @@ type ExecutableContextT[Instr any] struct {
 	labelPositionPool wazevoapi.Pool[LabelPosition[Instr]]
 	NextLabel         Label
 	// LabelPositions maps a label to the instructions of the region which the label represents.
-	LabelPositions     map[Label]*LabelPosition[Instr]
+	LabelPositions     []*LabelPosition[Instr]
 	OrderedBlockLabels []*LabelPosition[Instr]
 
 	// PerBlockHead and PerBlockEnd are the head and tail of the instruction list per currently-compiled ssa.BasicBlock.
@@ -67,7 +67,6 @@ func NewExecutableContextT[Instr any](
 		setNext:           setNext,
 		setPrev:           setPrev,
 		labelPositionPool: wazevoapi.NewPool[LabelPosition[Instr]](resetLabelPosition[Instr]),
-		LabelPositions:    make(map[Label]*LabelPosition[Instr]),
 		NextLabel:         LabelInvalid,
 	}
 }
@@ -97,11 +96,7 @@ func (e *ExecutableContextT[Instr]) StartBlock(blk ssa.BasicBlock) {
 	end := e.allocateNop0()
 	e.PerBlockHead, e.PerBlockEnd = end, end
 
-	labelPos, ok := e.LabelPositions[l]
-	if !ok {
-		labelPos = e.AllocateLabelPosition(l)
-		e.LabelPositions[l] = labelPos
-	}
+	labelPos := e.GetOrAllocateLabelPosition(l)
 	e.OrderedBlockLabels = append(e.OrderedBlockLabels, labelPos)
 	labelPos.Begin, labelPos.End = end, end
 	labelPos.SB = blk
@@ -146,8 +141,8 @@ func (e *ExecutableContextT[T]) FlushPendingInstructions() {
 func (e *ExecutableContextT[T]) Reset() {
 	e.labelPositionPool.Reset()
 	e.InstructionPool.Reset()
-	for l := Label(0); l <= e.NextLabel; l++ {
-		delete(e.LabelPositions, l)
+	for i := range e.LabelPositions {
+		e.LabelPositions[i] = nil
 	}
 	e.PendingInstructions = e.PendingInstructions[:0]
 	e.OrderedBlockLabels = e.OrderedBlockLabels[:0]
@@ -163,10 +158,17 @@ func (e *ExecutableContextT[T]) AllocateLabel() Label {
 	return e.NextLabel
 }
 
-func (e *ExecutableContextT[T]) AllocateLabelPosition(la Label) *LabelPosition[T] {
-	l := e.labelPositionPool.Allocate()
-	l.L = la
-	return l
+func (e *ExecutableContextT[T]) GetOrAllocateLabelPosition(l Label) *LabelPosition[T] {
+	if len(e.LabelPositions) <= int(l) {
+		e.LabelPositions = append(e.LabelPositions, make([]*LabelPosition[T], int(l)+1-len(e.LabelPositions))...)
+	}
+	ret := e.LabelPositions[l]
+	if ret == nil {
+		ret = e.labelPositionPool.Allocate()
+		ret.L = l
+		e.LabelPositions[l] = ret
+	}
+	return ret
 }
 
 func (e *ExecutableContextT[T]) GetOrAllocateSSABlockLabel(blk ssa.BasicBlock) Label {
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine.go b/internal/engine/wazevo/backend/isa/amd64/machine.go
index 310ad220..61ae6f40 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine.go
+++ b/internal/engine/wazevo/backend/isa/amd64/machine.go
@@ -1906,8 +1906,10 @@ func (m *machine) InsertMove(dst, src regalloc.VReg, typ ssa.Type) {
 func (m *machine) Format() string {
 	ectx := m.ectx
 	begins := map[*instruction]backend.Label{}
-	for l, pos := range ectx.LabelPositions {
-		begins[pos.Begin] = l
+	for _, pos := range ectx.LabelPositions {
+		if pos != nil {
+			begins[pos.Begin] = pos.L
+		}
 	}
 
 	irBlocks := map[backend.Label]ssa.BasicBlockID{}
@@ -1950,7 +1952,10 @@ func (m *machine) encodeWithoutSSA(root *instruction) {
 		offset := int64(len(*bufPtr))
 		if cur.kind == nop0 {
 			l := cur.nop0Label()
-			if pos, ok := ectx.LabelPositions[l]; ok {
+			if int(l) >= len(ectx.LabelPositions) {
+				continue
+			}
+			if pos := ectx.LabelPositions[l]; pos != nil {
 				pos.BinaryOffset = offset
 			}
 		}
@@ -2005,7 +2010,7 @@ func (m *machine) Encode(ctx context.Context) (err error) {
 			switch cur.kind {
 			case nop0:
 				l := cur.nop0Label()
-				if pos, ok := ectx.LabelPositions[l]; ok {
+				if pos := ectx.LabelPositions[l]; pos != nil {
 					pos.BinaryOffset = offset
 				}
 			case sourceOffsetInfo:
@@ -2165,8 +2170,7 @@ func (m *machine) allocateBrTarget() (nop *instruction, l backend.Label) { //nol
 func (m *machine) allocateLabel() *labelPosition {
 	ectx := m.ectx
 	l := ectx.AllocateLabel()
-	pos := ectx.AllocateLabelPosition(l)
-	ectx.LabelPositions[l] = pos
+	pos := ectx.GetOrAllocateLabelPosition(l)
 	return pos
 }
 
diff --git a/internal/engine/wazevo/backend/isa/arm64/machine.go b/internal/engine/wazevo/backend/isa/arm64/machine.go
index 311b34f4..5f584f92 100644
--- a/internal/engine/wazevo/backend/isa/arm64/machine.go
+++ b/internal/engine/wazevo/backend/isa/arm64/machine.go
@@ -187,9 +187,8 @@ func (m *machine) allocateBrTarget() (nop *instruction, l label) {
 	l = ectx.AllocateLabel()
 	nop = m.allocateInstr()
 	nop.asNop0WithLabel(l)
-	pos := ectx.AllocateLabelPosition(l)
+	pos := ectx.GetOrAllocateLabelPosition(l)
 	pos.Begin, pos.End = nop, nop
-	ectx.LabelPositions[l] = pos
 	return
 }
 
@@ -285,7 +284,7 @@ func (m *machine) resolveRelativeAddresses(ctx context.Context) {
 				switch cur.kind {
 				case nop0:
 					l := cur.nop0Label()
-					if pos, ok := ectx.LabelPositions[l]; ok {
+					if pos := ectx.LabelPositions[l]; pos != nil {
 						pos.BinaryOffset = offset + size
 					}
 				case condBr:
@@ -432,8 +431,10 @@ func (m *machine) insertConditionalJumpTrampoline(cbr *instruction, currentBlk *
 func (m *machine) Format() string {
 	ectx := m.executableContext
 	begins := map[*instruction]label{}
-	for l, pos := range ectx.LabelPositions {
-		begins[pos.Begin] = l
+	for _, pos := range ectx.LabelPositions {
+		if pos != nil {
+			begins[pos.Begin] = pos.L
+		}
 	}
 
 	irBlocks := map[label]ssa.BasicBlockID{}


Test patch:
diff --git a/internal/engine/wazevo/backend/isa/arm64/machine_test.go b/internal/engine/wazevo/backend/isa/arm64/machine_test.go
index 882b7384..943acff8 100644
--- a/internal/engine/wazevo/backend/isa/arm64/machine_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/machine_test.go
@@ -181,10 +181,10 @@ L200:
 
 			ectx := m.executableContext
 
-			originLabelPos := ectx.AllocateLabelPosition(originLabel)
+			originLabelPos := ectx.GetOrAllocateLabelPosition(originLabel)
 			originLabelPos.Begin = cbr
 			originLabelPos.End = linkInstr(cbr, end)
-			originNextLabelPos := ectx.AllocateLabelPosition(originLabelNext)
+			originNextLabelPos := ectx.GetOrAllocateLabelPosition(originLabelNext)
 			originNextLabelPos.Begin = originalEndNext
 			linkInstr(originLabelPos.End, originalEndNext)