Repository: tetratelabs/wazero
Pull request #1674 (2023-08-30 23:05:05)

wazevo(regalloc): sort liveNodes by begin PC
In preparation for passing "stack-page-guard" spectest,
this refactors a bit of regalloc logic.

#1496 


Patch:
diff --git a/internal/engine/wazevo/backend/regalloc/assign.go b/internal/engine/wazevo/backend/regalloc/assign.go
index f283d8e8..3e8cd5d5 100644
--- a/internal/engine/wazevo/backend/regalloc/assign.go
+++ b/internal/engine/wazevo/backend/regalloc/assign.go
@@ -2,7 +2,6 @@ package regalloc
 
 import (
 	"fmt"
-	"sort"
 
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/wazevoapi"
 )
@@ -13,9 +12,6 @@ func (a *Allocator) assignRegisters(f Function) {
 	for blk := f.ReversePostOrderBlockIteratorBegin(); blk != nil; blk = f.ReversePostOrderBlockIteratorNext() {
 		info := a.blockInfoAt(blk.ID())
 		lns := info.liveNodes
-		sort.SliceStable(lns, func(i, j int) bool {
-			return lns[i].n.v.ID() < lns[j].n.v.ID()
-		})
 		a.assignRegistersPerBlock(f, blk, a.vRegIDToNode, lns)
 	}
 }
@@ -34,11 +30,26 @@ func (a *Allocator) assignRegistersPerBlock(f Function, blk Block, vRegIDToNode
 }
 
 func (a *Allocator) assignRegistersPerInstr(f Function, pc programCounter, instr Instr, vRegIDToNode []*node, liveNodes []liveNodeInBlock) {
+	if wazevoapi.RegAllocValidationEnabled {
+		// Check if the liveNodes are sorted by the start program counter.
+		for i := 1; i < len(liveNodes); i++ {
+			n, m := liveNodes[i-1], liveNodes[i]
+			if n.n.ranges[n.rangeIndex].begin > m.n.ranges[m.rangeIndex].begin {
+				panic(fmt.Sprintf("BUG: liveNodes are not sorted by the start program counter: %d > %d",
+					n.n.ranges[n.rangeIndex].begin, m.n.ranges[m.rangeIndex].begin,
+				))
+			}
+		}
+	}
+
 	if direct := instr.IsCall(); direct || instr.IsIndirectCall() {
 		// Only take care of non-real VRegs (e.g. VReg.IsRealReg() == false) since
 		// the real VRegs are already placed in the right registers at this point.
-		actives := a.activeNonRealVRegsAt(pc+pcUseOffset, liveNodes)
-		for _, active := range actives {
+		a.nodes2 = a.collectActiveNonRealVRegsAt(a.nodes2,
+			// To find the all the live registers "after" call, we need to add pcDefOffset for search.
+			pc+pcDefOffset,
+			liveNodes)
+		for _, active := range a.nodes2 {
 			if r := active.r; a.regInfo.isCallerSaved(r) {
 				v := active.v.SetRealReg(r)
 				f.StoreRegisterBefore(v, instr)
@@ -104,20 +115,24 @@ func (a *Allocator) assignRegistersPerInstr(f Function, pc programCounter, instr
 	}
 }
 
-// activeRegistersAt returns the set of active registers at the given program counter.
-// This excludes the VRegs backed by a real register since this is used to list the registers
+// collectActiveNonRealVRegsAt collects the set of active registers at the given program counter into `nodes` slice by appending
+// the found registers from its beginning. This excludes the VRegs backed by a real register since this is used to list the registers
 // alive but not used by a call instruction.
-func (a *Allocator) activeNonRealVRegsAt(pc programCounter, liveNodes []liveNodeInBlock) []*node {
-	a.nodes1 = a.nodes1[:0]
+func (a *Allocator) collectActiveNonRealVRegsAt(nodes []*node, pc programCounter, liveNodes []liveNodeInBlock) []*node {
+	nodes = nodes[:0]
 	for _, live := range liveNodes {
 		n := live.n
 		if n.spill() || n.v.IsRealReg() {
 			continue
 		}
 		r := &n.ranges[live.rangeIndex]
-		if r.contains(pc) {
-			a.nodes1 = append(a.nodes1, n)
+		if r.begin > pc {
+			// liveNodes are sorted by the start program counter, so we can break here.
+			break
+		}
+		if pc <= r.end { // pc is in the range.
+			nodes = append(nodes, n)
 		}
 	}
-	return a.nodes1
+	return nodes
 }
diff --git a/internal/engine/wazevo/backend/regalloc/regalloc.go b/internal/engine/wazevo/backend/regalloc/regalloc.go
index 11c3d9a7..782e8ea2 100644
--- a/internal/engine/wazevo/backend/regalloc/regalloc.go
+++ b/internal/engine/wazevo/backend/regalloc/regalloc.go
@@ -289,6 +289,15 @@ func (a *Allocator) buildLiveRanges(f Function) {
 		info := a.blockInfoAt(blkID)
 		a.buildLiveRangesForNonReals(blkID, info)
 		a.buildLiveRangesForReals(blkID, info)
+		// Sort the live range for a fast lookup to find live registers at a given program counter.
+		sort.Slice(info.liveNodes, func(i, j int) bool {
+			inode, jnode := &info.liveNodes[i], &info.liveNodes[j]
+			irange, jrange := inode.n.ranges[inode.rangeIndex], jnode.n.ranges[jnode.rangeIndex]
+			if irange.begin == jrange.begin {
+				return irange.end < jrange.end
+			}
+			return irange.begin < jrange.begin
+		})
 	}
 }
 
@@ -624,11 +633,6 @@ func (l *liveRange) intersects(other *liveRange) bool {
 	return other.begin <= l.end && l.begin <= other.end
 }
 
-// contains returns true if the live range contains the given program counter.
-func (l *liveRange) contains(pc programCounter) bool {
-	return l.begin <= pc && pc < l.end
-}
-
 func (r *RegisterInfo) isCalleeSaved(reg RealReg) bool {
 	_, ok := r.CalleeSavedRegisters[reg]
 	return ok
diff --git a/internal/engine/wazevo/wazevoapi/debug_consts.go b/internal/engine/wazevo/wazevoapi/debug_consts.go
index 12c6e14c..ce357025 100644
--- a/internal/engine/wazevo/wazevoapi/debug_consts.go
+++ b/internal/engine/wazevo/wazevoapi/debug_consts.go
@@ -57,13 +57,13 @@ const (
 )
 
 type (
-	verifierContext struct {
+	verifierState struct {
 		randomizedIndexes []int
 		r                 *rand.Rand
 		values            map[string]string
 	}
-	existingValuesContextKey struct{}
-	currentFunctionNameKey   struct{}
+	verifierStateContextKey struct{}
+	currentFunctionNameKey  struct{}
 )
 
 // NewDeterministicCompilationVerifierContext creates a new context with the deterministic compilation verifier used per wasm.Module.
@@ -73,7 +73,7 @@ func NewDeterministicCompilationVerifierContext(ctx context.Context, localFuncti
 		randomizedIndexes[i] = i
 	}
 	r := rand.New(rand.NewSource(time.Now().UnixNano()))
-	return context.WithValue(ctx, existingValuesContextKey{}, &verifierContext{
+	return context.WithValue(ctx, verifierStateContextKey{}, &verifierState{
 		r: r, randomizedIndexes: randomizedIndexes, values: map[string]string{},
 	})
 }
@@ -81,7 +81,7 @@ func NewDeterministicCompilationVerifierContext(ctx context.Context, localFuncti
 // DeterministicCompilationVerifierRandomizeIndexes randomizes the indexes for the deterministic compilation verifier.
 // To get the randomized index, use DeterministicCompilationVerifierGetRandomizedLocalFunctionIndex.
 func DeterministicCompilationVerifierRandomizeIndexes(ctx context.Context) {
-	verifierCtx := ctx.Value(existingValuesContextKey{}).(*verifierContext)
+	verifierCtx := ctx.Value(verifierStateContextKey{}).(*verifierState)
 	r := verifierCtx.r
 	r.Shuffle(len(verifierCtx.randomizedIndexes), func(i, j int) {
 		verifierCtx.randomizedIndexes[i], verifierCtx.randomizedIndexes[j] = verifierCtx.randomizedIndexes[j], verifierCtx.randomizedIndexes[i]
@@ -91,7 +91,7 @@ func DeterministicCompilationVerifierRandomizeIndexes(ctx context.Context) {
 // DeterministicCompilationVerifierGetRandomizedLocalFunctionIndex returns the randomized index for the given `index`
 // which is assigned by DeterministicCompilationVerifierRandomizeIndexes.
 func DeterministicCompilationVerifierGetRandomizedLocalFunctionIndex(ctx context.Context, index int) int {
-	verifierCtx := ctx.Value(existingValuesContextKey{}).(*verifierContext)
+	verifierCtx := ctx.Value(verifierStateContextKey{}).(*verifierState)
 	ret := verifierCtx.randomizedIndexes[index]
 	return ret
 }
@@ -108,7 +108,7 @@ func DeterministicCompilationVerifierSetCurrentFunctionName(ctx context.Context,
 func VerifyOrSetDeterministicCompilationContextValue(ctx context.Context, scope string, newValue string) {
 	fn := ctx.Value(currentFunctionNameKey{}).(string)
 	key := fn + ": " + scope
-	verifierCtx := ctx.Value(existingValuesContextKey{}).(*verifierContext)
+	verifierCtx := ctx.Value(verifierStateContextKey{}).(*verifierState)
 	oldValue, ok := verifierCtx.values[key]
 	if !ok {
 		verifierCtx.values[key] = newValue
@@ -118,7 +118,7 @@ func VerifyOrSetDeterministicCompilationContextValue(ctx context.Context, scope
 		fmt.Printf(
 			`BUG: Deterministic compilation failed for function%s at scope="%s".
 
-This is mostly due to (but might not be not limited to):
+This is mostly due to (but might not be limited to):
 	* Resetting ssa.Builder, backend.Compiler or frontend.Compiler, etc doens't work as expected, and the compilation has been affected by the previous iterations.
 	* Using a map with non-deterministic iteration order.
 
@@ -128,8 +128,7 @@ This is mostly due to (but might not be not limited to):
 ---------- [new] ----------
 %s
 `,
-			fn,
-			scope, oldValue, newValue,
+			fn, scope, oldValue, newValue,
 		)
 		os.Exit(1)
 	}


Test patch:
diff --git a/internal/engine/wazevo/backend/backend_test.go b/internal/engine/wazevo/backend/backend_test.go
index 24ed8a09..67c90b65 100644
--- a/internal/engine/wazevo/backend/backend_test.go
+++ b/internal/engine/wazevo/backend/backend_test.go
@@ -2054,13 +2054,13 @@ L4 (SSA Block: blk3):
 	str x8, [x9, #0x8]
 	mov x0, x9
 	mov x1, x8
-	str x9, [sp]
-	str x8, [sp, #0x8]
-	str w11, [sp, #0x10]
+	str w11, [sp]
+	str x9, [sp, #0x4]
+	str x8, [sp, #0xc]
 	bl f0
-	ldr w11, [sp, #0x10]
-	ldr x8, [sp, #0x8]
-	ldr x9, [sp]
+	ldr x8, [sp, #0xc]
+	ldr x9, [sp, #0x4]
+	ldr w11, [sp]
 	mov x10, x0
 	sub w2, w11, #0x2
 	str x8, [x9, #0x8]
diff --git a/internal/engine/wazevo/backend/regalloc/assign_test.go b/internal/engine/wazevo/backend/regalloc/assign_test.go
index e4d7f6a3..494fabe0 100644
--- a/internal/engine/wazevo/backend/regalloc/assign_test.go
+++ b/internal/engine/wazevo/backend/regalloc/assign_test.go
@@ -70,8 +70,8 @@ func TestAllocator_activeNonRealVRegsAt(t *testing.T) {
 			name: "one live",
 			pc:   10,
 			lives: []liveNodeInBlock{
-				{n: &node{r: 1, v: 0xa, ranges: []liveRange{{begin: 100, end: 2000}}}},
 				{n: &node{r: 2, v: 0xf, ranges: []liveRange{{begin: 5, end: 20}}}},
+				{n: &node{r: 1, v: 0xa, ranges: []liveRange{{begin: 100, end: 2000}}}},
 			},
 			want: []VReg{0xf},
 		},
@@ -99,7 +99,7 @@ func TestAllocator_activeNonRealVRegsAt(t *testing.T) {
 		tc := tc
 		t.Run(tc.name, func(t *testing.T) {
 			a := NewAllocator(&RegisterInfo{})
-			ans := a.activeNonRealVRegsAt(tc.pc, tc.lives)
+			ans := a.collectActiveNonRealVRegsAt(nil, tc.pc, tc.lives)
 
 			actual := make([]VReg, len(ans))
 			for i, n := range ans {