From 38a9374cb1d7c8b26a592d7f3f79c277a3fcf918 Mon Sep 17 00:00:00 2001
From: Takeshi Yoneda <t.y.mathetake@gmail.com>
Date: Tue, 4 Jun 2024 08:27:03 -0700
Subject: [PATCH] regalloc: simplifies regInUseSet (#2227)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This makes the compilation slightly faster and simplifies the codebase.

```
goos: darwin
goarch: arm64
pkg: github.com/tetratelabs/wazero
               │  old.txt   │             new.txt              │
               │   sec/op   │   sec/op    vs base              │
Compilation-10   2.499 ± 0%   2.478 ± 2%  -0.84% (p=0.026 n=7)

               │   old.txt    │            new.txt            │
               │     B/op     │     B/op      vs base         │
Compilation-10   341.0Mi ± 0%   341.1Mi ± 0%  ~ (p=0.128 n=7)

               │   old.txt   │           new.txt            │
               │  allocs/op  │  allocs/op   vs base         │
Compilation-10   606.9k ± 0%   606.9k ± 0%  ~ (p=0.383 n=7)
```

Signed-off-by: Takeshi Yoneda <t.y.mathetake@gmail.com>
---
 .../wazevo/backend/regalloc/regalloc.go       | 10 ++++-
 .../wazevo/backend/regalloc/regalloc_test.go  |  2 +-
 .../engine/wazevo/backend/regalloc/regset.go  | 44 +++++++------------
 3 files changed, 26 insertions(+), 30 deletions(-)

diff --git a/internal/engine/wazevo/backend/regalloc/regalloc.go b/internal/engine/wazevo/backend/regalloc/regalloc.go
index 6e11c085..731426fc 100644
--- a/internal/engine/wazevo/backend/regalloc/regalloc.go
+++ b/internal/engine/wazevo/backend/regalloc/regalloc.go
@@ -972,6 +972,12 @@ func (a *Allocator) fixMergeState(f Function, blk Block) {
 	bID := blk.ID()
 	blkSt := a.getOrAllocateBlockState(bID)
 	desiredOccupants := &blkSt.startRegs
+	var desiredOccupantsSet RegSet
+	for i, v := range desiredOccupants {
+		if v != VRegInvalid {
+			desiredOccupantsSet = desiredOccupantsSet.add(RealReg(i))
+		}
+	}
 
 	if wazevoapi.RegAllocLoggingEnabled {
 		fmt.Println("fixMergeState", blk.ID(), ":", desiredOccupants.format(a.regInfo))
@@ -993,12 +999,12 @@ func (a *Allocator) fixMergeState(f Function, blk Block) {
 		// Finds the free registers if any.
 		intTmp, floatTmp := VRegInvalid, VRegInvalid
 		if intFree := s.findAllocatable(
-			a.regInfo.AllocatableRegisters[RegTypeInt], desiredOccupants.set,
+			a.regInfo.AllocatableRegisters[RegTypeInt], desiredOccupantsSet,
 		); intFree != RealRegInvalid {
 			intTmp = FromRealReg(intFree, RegTypeInt)
 		}
 		if floatFree := s.findAllocatable(
-			a.regInfo.AllocatableRegisters[RegTypeFloat], desiredOccupants.set,
+			a.regInfo.AllocatableRegisters[RegTypeFloat], desiredOccupantsSet,
 		); floatFree != RealRegInvalid {
 			floatTmp = FromRealReg(floatFree, RegTypeFloat)
 		}
diff --git a/internal/engine/wazevo/backend/regalloc/regalloc_test.go b/internal/engine/wazevo/backend/regalloc/regalloc_test.go
index d0daf2f9..526d2792 100644
--- a/internal/engine/wazevo/backend/regalloc/regalloc_test.go
+++ b/internal/engine/wazevo/backend/regalloc/regalloc_test.go
@@ -366,7 +366,7 @@ func TestAllocator_livenessAnalysis_copy(t *testing.T) {
 
 func Test_findOrSpillAllocatable_prefersSpill(t *testing.T) {
 	t.Run("ok", func(t *testing.T) {
-		s := &state{}
+		s := &state{regsInUse: newRegInUseSet()}
 		s.regsInUse.add(RealReg(1), VReg(2222222))
 		got := s.findOrSpillAllocatable(&Allocator{}, []RealReg{3}, 0, 3)
 		require.Equal(t, RealReg(3), got)
diff --git a/internal/engine/wazevo/backend/regalloc/regset.go b/internal/engine/wazevo/backend/regalloc/regset.go
index e9bf6066..04a8e8f4 100644
--- a/internal/engine/wazevo/backend/regalloc/regset.go
+++ b/internal/engine/wazevo/backend/regalloc/regset.go
@@ -46,23 +46,24 @@ func (rs RegSet) Range(f func(allocatedRealReg RealReg)) {
 	}
 }
 
-type regInUseSet struct {
-	set RegSet
-	vrs [64]VReg
+type regInUseSet [64]VReg
+
+func newRegInUseSet() regInUseSet {
+	var ret regInUseSet
+	ret.reset()
+	return ret
 }
 
 func (rs *regInUseSet) reset() {
-	rs.set = 0
-	for i := range rs.vrs {
-		rs.vrs[i] = VRegInvalid
+	for i := range rs {
+		rs[i] = VRegInvalid
 	}
 }
 
 func (rs *regInUseSet) format(info *RegisterInfo) string { //nolint:unused
 	var ret []string
-	for i := 0; i < 64; i++ {
-		if rs.set&(1<<uint(i)) != 0 {
-			vr := rs.vrs[i]
+	for i, vr := range rs {
+		if vr != VRegInvalid {
 			ret = append(ret, fmt.Sprintf("(%s->v%d)", info.RealRegName(RealReg(i)), vr.ID()))
 		}
 	}
@@ -70,39 +71,28 @@ func (rs *regInUseSet) format(info *RegisterInfo) string { //nolint:unused
 }
 
 func (rs *regInUseSet) has(r RealReg) bool {
-	if r >= 64 {
-		return false
-	}
-	return rs.set&(1<<uint(r)) != 0
+	return r < 64 && rs[r] != VRegInvalid
 }
 
 func (rs *regInUseSet) get(r RealReg) VReg {
-	if r >= 64 {
-		return VRegInvalid
-	}
-	return rs.vrs[r]
+	return rs[r]
 }
 
 func (rs *regInUseSet) remove(r RealReg) {
-	if r >= 64 {
-		return
-	}
-	rs.set &= ^(1 << uint(r))
-	rs.vrs[r] = VRegInvalid
+	rs[r] = VRegInvalid
 }
 
 func (rs *regInUseSet) add(r RealReg, vr VReg) {
 	if r >= 64 {
 		return
 	}
-	rs.set |= 1 << uint(r)
-	rs.vrs[r] = vr
+	rs[r] = vr
 }
 
 func (rs *regInUseSet) range_(f func(allocatedRealReg RealReg, vr VReg)) {
-	for i := 0; i < 64; i++ {
-		if rs.set&(1<<uint(i)) != 0 {
-			f(RealReg(i), rs.vrs[i])
+	for i, vr := range rs {
+		if vr != VRegInvalid {
+			f(RealReg(i), vr)
 		}
 	}
 }