Repository: tetratelabs/wazero
Pull request #1803 (2023-10-19 16:43:19)

wazevo: reduce memory footprint of regalloc kills
This a small optimization that I noticed while going through the code: the regalloc `kills` are always derived from `lastUses`, which can avoid tracking the program counters in `kills` and turn the collection into a set that is only tracking virtual registers. When we need the `kills` program counters, we can look them back up into `lastUses`.

On the Python compilation, this saves 55MB of a 248MB cumulative allocation:
```
github.com/tetratelabs/wazero/internal/engine/wazevo/backend/regalloc/regalloc.go

  Total:    248.01MB   248.01MB (flat, cum)  1.99%
    249            .          .           			if _, ok := outs[use]; !ok { 
    250     248.01MB   248.01MB           				info.kills[use] = pc 
    251            .          .           			} 
    252            .          .           		}) 
    253            .          .            
    254            .          .           		if wazevoapi.RegAllocLoggingEnabled { 
    255            .          .           			fmt.Printf("\nfinalized info for block[%d]:\n%s\n", blk.ID(), info.Format(a.regInfo)) 
```
```
github.com/tetratelabs/wazero/internal/engine/wazevo/backend/regalloc/regalloc.go

  Total:    -55.36MB   -55.36MB (flat, cum)  0.22%
    249            .          .           			if _, ok := outs[use]; !ok { 
    250     -55.36MB   -55.36MB           				info.kills[use] = struct{}{} 
    251            .          .           			} 
    252            .          .           		}) 
    253            .          .            
    254            .          .           		if wazevoapi.RegAllocLoggingEnabled { 
    255            .          .           			fmt.Printf("\nfinalized info for block[%d]:\n%s\n", blk.ID(), info.Format(a.regInfo)) 
```

As a side note, this also makes the compilation of Python ~1s faster on a M2.

There is likely a lot of room for improvement here by replacing the maps with a bitset, but I wanted to make this optimization first to get a baseline.


Patch:
diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
index cfd11d3a..a57ad662 100644
--- a/.github/CODEOWNERS
+++ b/.github/CODEOWNERS
@@ -1,6 +1,3 @@
 # https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners
 
-/internal/asm/ @mathetake
-/internal/engine/compiler/ @mathetake
-
-* @codefromthecrypt @mathetake
+* @codefromthecrypt @mathetake @evacchi
diff --git a/internal/engine/wazevo/backend/regalloc/regalloc.go b/internal/engine/wazevo/backend/regalloc/regalloc.go
index b336ed08..7642111f 100644
--- a/internal/engine/wazevo/backend/regalloc/regalloc.go
+++ b/internal/engine/wazevo/backend/regalloc/regalloc.go
@@ -79,7 +79,7 @@ type (
 		liveIns  map[VReg]struct{}
 		defs     map[VReg]programCounter
 		lastUses VRegTable
-		kills    map[VReg]programCounter
+		kills    map[VReg]struct{}
 		// Pre-colored real registers can have multiple live ranges in one block.
 		realRegUses [vRegIDReservedForRealNum][]programCounter
 		realRegDefs [vRegIDReservedForRealNum][]programCounter
@@ -247,7 +247,7 @@ func (a *Allocator) livenessAnalysis(f Function) {
 		info.lastUses.Range(func(use VReg, pc programCounter) {
 			// Usage without live-outs is a kill.
 			if _, ok := outs[use]; !ok {
-				info.kills[use] = pc
+				info.kills[use] = struct{}{}
 			}
 		})
 
@@ -348,12 +348,12 @@ func (a *Allocator) buildLiveRangesForNonReals(info *blockInfo) {
 				panic("BUG: v is live-out but also killed")
 			}
 		} else {
-			killPos, ok := kills[v]
+			_, ok := kills[v]
 			if !ok {
 				panic("BUG: v is live-in but not live-out or use")
 			}
 			// v is killed at killPos.
-			begin, end = 0, killPos
+			begin, end = 0, info.lastUses.Lookup(v)
 		}
 		n := a.getOrAllocateNode(v)
 		intervalNode := info.intervalMng.insert(n, begin, end)
@@ -387,13 +387,13 @@ func (a *Allocator) buildLiveRangesForNonReals(info *blockInfo) {
 				panic("BUG: v is killed here but also killed")
 			}
 		} else {
-			killPos, ok := kills[v]
+			_, ok := kills[v]
 			if !ok {
 				// This case the defined value is not used at all.
 				end = defPos
 			} else {
 				// v is killed at pos.
-				end = killPos
+				end = info.lastUses.Lookup(v)
 			}
 		}
 		n := a.getOrAllocateNode(v)
@@ -556,7 +556,7 @@ func (a *Allocator) initBlockInfo(i *blockInfo) {
 		resetMap(a, i.defs)
 	}
 	if i.kills == nil {
-		i.kills = make(map[VReg]programCounter)
+		i.kills = make(map[VReg]struct{})
 	} else {
 		resetMap(a, i.kills)
 	}


Test patch:
diff --git a/internal/engine/wazevo/backend/regalloc/regalloc_test.go b/internal/engine/wazevo/backend/regalloc/regalloc_test.go
index f3ab10fa..a0205af2 100644
--- a/internal/engine/wazevo/backend/regalloc/regalloc_test.go
+++ b/internal/engine/wazevo/backend/regalloc/regalloc_test.go
@@ -41,7 +41,7 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 				0: {
 					defs:     map[VReg]programCounter{2: pcDefOffset + pcStride, 1: pcDefOffset},
 					lastUses: makeVRegTable(map[VReg]programCounter{1: pcStride + pcUseOffset}),
-					kills:    map[VReg]programCounter{1: pcStride + pcUseOffset},
+					kills:    map[VReg]struct{}{1: {}},
 				},
 			},
 		},
@@ -81,10 +81,10 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 						2:    pcStride*2 + pcUseOffset,
 					}),
 					liveOuts: map[VReg]struct{}{3: {}},
-					kills: map[VReg]programCounter{
-						1000: pcStride + pcUseOffset,
-						1:    pcStride*2 + pcUseOffset,
-						2:    pcStride*2 + pcUseOffset,
+					kills: map[VReg]struct{}{
+						1000: {},
+						1:    {},
+						2:    {},
 					},
 				},
 				1: {
@@ -107,7 +107,7 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 				2: {
 					liveIns:  map[VReg]struct{}{3: {}, 4: {}, 5: {}},
 					lastUses: makeVRegTable(map[VReg]programCounter{3: pcUseOffset, 4: pcUseOffset, 5: pcUseOffset}),
-					kills:    map[VReg]programCounter{3: pcUseOffset, 4: pcUseOffset, 5: pcUseOffset},
+					kills:    map[VReg]struct{}{3: {}, 4: {}, 5: {}},
 				},
 			},
 		},
@@ -158,7 +158,7 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 					liveIns:  map[VReg]struct{}{1000: {}, 1: {}},
 					liveOuts: map[VReg]struct{}{1000: {}},
 					lastUses: makeVRegTable(map[VReg]programCounter{1: pcUseOffset}),
-					kills:    map[VReg]programCounter{1: pcUseOffset},
+					kills:    map[VReg]struct{}{1: {}},
 					realRegDefs: [vRegIDReservedForRealNum][]programCounter{
 						realRegID:  {pcDefOffset, pcStride*4 + pcDefOffset},
 						realRegID2: {pcStride*2 + pcDefOffset},
@@ -172,14 +172,14 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 					liveIns:     map[VReg]struct{}{1000: {}, 2: {}},
 					liveOuts:    map[VReg]struct{}{1000: {}},
 					lastUses:    makeVRegTable(map[VReg]programCounter{2: pcUseOffset}),
-					kills:       map[VReg]programCounter{2: pcUseOffset},
+					kills:       map[VReg]struct{}{2: {}},
 					realRegUses: [vRegIDReservedForRealNum][]programCounter{realRegID2: {pcUseOffset}},
 					realRegDefs: [vRegIDReservedForRealNum][]programCounter{realRegID2: {0}},
 				},
 				3: {
 					liveIns:  map[VReg]struct{}{1000: {}},
 					lastUses: makeVRegTable(map[VReg]programCounter{1000: pcUseOffset}),
-					kills:    map[VReg]programCounter{1000: pcUseOffset},
+					kills:    map[VReg]struct{}{1000: {}},
 				},
 			},
 		},
@@ -225,7 +225,7 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 					liveOuts: map[VReg]struct{}{phiVReg: {}, 3000: {}},
 					defs:     map[VReg]programCounter{phiVReg: pcDefOffset},
 					lastUses: makeVRegTable(map[VReg]programCounter{2000: pcUseOffset}),
-					kills:    map[VReg]programCounter{2000: pcUseOffset},
+					kills:    map[VReg]struct{}{2000: {}},
 				},
 				2: {
 					liveIns:  map[VReg]struct{}{phiVReg: {}, 3000: {}},
@@ -237,12 +237,12 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 					liveOuts: map[VReg]struct{}{phiVReg: {}, 3000: {}},
 					defs:     map[VReg]programCounter{phiVReg: pcDefOffset},
 					lastUses: makeVRegTable(map[VReg]programCounter{1000: pcUseOffset}),
-					kills:    map[VReg]programCounter{1000: pcUseOffset},
+					kills:    map[VReg]struct{}{1000: {}},
 				},
 				4: {
 					liveIns:  map[VReg]struct{}{phiVReg: {}, 3000: {}},
 					lastUses: makeVRegTable(map[VReg]programCounter{phiVReg: pcUseOffset, 3000: pcUseOffset}),
-					kills:    map[VReg]programCounter{phiVReg: pcUseOffset, 3000: pcUseOffset},
+					kills:    map[VReg]struct{}{phiVReg: {}, 3000: {}},
 				},
 			},
 		},
@@ -295,8 +295,8 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 					lastUses: makeVRegTable(map[VReg]programCounter{
 						1: pcStride + pcUseOffset,
 					}),
-					kills: map[VReg]programCounter{
-						1: pcStride + pcUseOffset,
+					kills: map[VReg]struct{}{
+						1: {},
 					},
 				},
 				1: {
@@ -304,28 +304,28 @@ func TestAllocator_livenessAnalysis(t *testing.T) {
 					liveOuts: map[VReg]struct{}{phiVReg: {}, 9999: {}},
 					defs:     map[VReg]programCounter{9999: pcDefOffset},
 					lastUses: makeVRegTable(map[VReg]programCounter{}),
-					kills:    map[VReg]programCounter{},
+					kills:    map[VReg]struct{}{},
 				},
 				2: {
 					liveIns:  map[VReg]struct{}{phiVReg: {}, 9999: {}},
 					liveOuts: map[VReg]struct{}{100: {}},
 					defs:     map[VReg]programCounter{100: pcDefOffset},
 					lastUses: makeVRegTable(map[VReg]programCounter{phiVReg: pcUseOffset, 9999: pcUseOffset}),
-					kills:    map[VReg]programCounter{phiVReg: pcUseOffset, 9999: pcUseOffset},
+					kills:    map[VReg]struct{}{phiVReg: {}, 9999: {}},
 				},
 				3: {
 					liveIns:  map[VReg]struct{}{100: {}},
 					liveOuts: map[VReg]struct{}{54321: {}},
 					defs:     map[VReg]programCounter{54321: pcDefOffset},
 					lastUses: makeVRegTable(map[VReg]programCounter{100: pcStride + pcUseOffset}),
-					kills:    map[VReg]programCounter{100: pcStride + pcUseOffset},
+					kills:    map[VReg]struct{}{100: {}},
 				},
 				4: {
 					liveIns:  map[VReg]struct{}{54321: {}},
 					liveOuts: map[VReg]struct{}{phiVReg: {}},
 					defs:     map[VReg]programCounter{phiVReg: pcDefOffset},
 					lastUses: makeVRegTable(map[VReg]programCounter{54321: pcUseOffset}),
-					kills:    map[VReg]programCounter{54321: pcUseOffset},
+					kills:    map[VReg]struct{}{54321: {}},
 				},
 			},
 		},
@@ -569,7 +569,7 @@ func initMapInInfo(info *blockInfo) {
 		info.defs = make(map[VReg]programCounter)
 	}
 	if info.kills == nil {
-		info.kills = make(map[VReg]programCounter)
+		info.kills = make(map[VReg]struct{})
 	}
 }