Repository: tetratelabs/wazero
Pull request #1639 (2023-08-18 23:47:55)

regalloc: enforces unique definition of virtual regs
FYI @evacchi 


Patch:
diff --git a/internal/engine/wazevo/backend/compiler.go b/internal/engine/wazevo/backend/compiler.go
index 0e752a81..19493457 100644
--- a/internal/engine/wazevo/backend/compiler.go
+++ b/internal/engine/wazevo/backend/compiler.go
@@ -259,9 +259,9 @@ func (c *compiler) AllocateVRegWithSSAType(regType regalloc.RegType, typ ssa.Typ
 
 // Init implements Compiler.Init.
 func (c *compiler) Init(needGoEntryPreamble bool) {
-	for i := regalloc.VRegID(0); i < c.nextVRegID; i++ {
+	for i, v := range c.ssaValueToVRegs {
 		c.ssaValueToVRegs[i] = regalloc.VRegInvalid
-		delete(c.ssaTypeOfVRegID, i)
+		delete(c.ssaTypeOfVRegID, v.ID())
 	}
 	c.currentGID = 0
 	c.nextVRegID = 0
diff --git a/internal/engine/wazevo/backend/isa/arm64/lower_instr.go b/internal/engine/wazevo/backend/isa/arm64/lower_instr.go
index 789b37bb..55d61c75 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_instr.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_instr.go
@@ -315,11 +315,11 @@ func (m *machine) lowerCtz(x, result ssa.Value) {
 	rd := m.compiler.VRegOf(result)
 	rn := m.getOperand_NR(m.compiler.ValueDefinition(x), extModeNone)
 	rbit := m.allocateInstr()
-	rbit.asBitRR(bitOpRbit, rd, rn.nr(), x.Type().Bits() == 64)
+	rbit.asBitRR(bitOpRbit, tmpRegVReg, rn.nr(), x.Type().Bits() == 64)
 	m.insert(rbit)
 
 	clz := m.allocateInstr()
-	clz.asBitRR(bitOpClz, rd, rd, x.Type().Bits() == 64)
+	clz.asBitRR(bitOpClz, rd, tmpRegVReg, x.Type().Bits() == 64)
 	m.insert(clz)
 }
 
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 cbd82fb0..d16e6324 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_instr_operands.go
@@ -301,20 +301,25 @@ func (m *machine) getOperand_NR(def *backend.SSAValueDefinition, mode extMode) (
 		}
 	}
 
+	r := v
 	switch inBits := def.SSAValue().Type().Bits(); {
 	case mode == extModeNone:
 	case inBits == 32 && (mode == extModeZeroExtend32 || mode == extModeSignExtend32):
 	case inBits == 32 && mode == extModeZeroExtend64:
+		extended := m.compiler.AllocateVReg(regalloc.RegTypeInt)
 		ext := m.allocateInstr()
-		ext.asExtend(v, v, 32, 64, false)
+		ext.asExtend(extended, v, 32, 64, false)
 		m.insert(ext)
+		r = extended
 	case inBits == 32 && mode == extModeSignExtend64:
+		extended := m.compiler.AllocateVReg(regalloc.RegTypeInt)
 		ext := m.allocateInstr()
-		ext.asExtend(v, v, 32, 64, true)
+		ext.asExtend(extended, v, 32, 64, true)
 		m.insert(ext)
+		r = extended
 	case inBits == 64 && (mode == extModeZeroExtend64 || mode == extModeSignExtend64):
 	}
-	return operandNR(v)
+	return operandNR(r)
 }
 
 func asImm12Operand(val uint64) (op operand, ok bool) {
diff --git a/internal/engine/wazevo/backend/regalloc/reg.go b/internal/engine/wazevo/backend/regalloc/reg.go
index ab7fa404..31e39732 100644
--- a/internal/engine/wazevo/backend/regalloc/reg.go
+++ b/internal/engine/wazevo/backend/regalloc/reg.go
@@ -87,7 +87,7 @@ func (v VReg) String() string {
 	if v.IsRealReg() {
 		return fmt.Sprintf("r%d", v.ID()-vRegIDReservedForRealBegin)
 	}
-	return fmt.Sprintf("r%d?", v.ID())
+	return fmt.Sprintf("v%d?", v.ID())
 }
 
 // RegType represents the type of a register.
diff --git a/internal/engine/wazevo/backend/regalloc/regalloc.go b/internal/engine/wazevo/backend/regalloc/regalloc.go
index f00961ee..c0a24448 100644
--- a/internal/engine/wazevo/backend/regalloc/regalloc.go
+++ b/internal/engine/wazevo/backend/regalloc/regalloc.go
@@ -188,6 +188,19 @@ func (a *Allocator) livenessAnalysis(f Function) {
 				if def.IsRealReg() {
 					info.realRegDefs[def] = append(info.realRegDefs[def], pos)
 				} else {
+					if _, ok := info.defs[def]; ok {
+						// This must be a bug in lowering logic in the ISA-specific code.
+						// In short, a virtual register must be defined only once because that's the invariant of
+						// liveness analysis result used in the current implementation. If you reach here, you can either
+						// use "temporary" physical register (like x27 in AArch64), or simply allocate a new virtual
+						// register for the additional definition. In any ways, such a new virtual register doesn't cost
+						// at all in the final code because they are instantly killed and won't interfere with other
+						// virtual registers later on.
+						//
+						// TODO: this should be enforced globally, not just per block. But that can be costly to check.
+						panic(fmt.Sprintf("BUG: multiple definitions found for a virtual register %s", def.String()))
+					}
+
 					info.defs[def] = pos
 					a.vs = append(a.vs, def)
 				}


Test patch:
diff --git a/internal/engine/wazevo/backend/backend_test.go b/internal/engine/wazevo/backend/backend_test.go
index 2df19b39..409fff70 100644
--- a/internal/engine/wazevo/backend/backend_test.go
+++ b/internal/engine/wazevo/backend/backend_test.go
@@ -1424,11 +1424,11 @@ L1 (SSA Block: blk0):
 	mov x2?, x2
 	mov x3?, x3
 	clz w4?, w2?
-	rbit w5?, w2?
-	clz w5?, w5?
+	rbit w27, w2?
+	clz w5?, w27
 	clz x6?, x3?
-	rbit x7?, x3?
-	clz x7?, x7?
+	rbit x27, x3?
+	clz x7?, x27
 	mov x3, x7?
 	mov x2, x6?
 	mov x1, x5?
@@ -1439,11 +1439,11 @@ L1 (SSA Block: blk0):
 L1 (SSA Block: blk0):
 	str x30, [sp, #-0x10]!
 	clz w0, w2
-	rbit w1, w2
-	clz w1, w1
+	rbit w27, w2
+	clz w1, w27
 	clz x2, x3
-	rbit x3, x3
-	clz x3, x3
+	rbit x27, x3
+	clz x3, x27
 	ldr x30, [sp], #0x10
 	ret
 `,
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 b6d0f2dc..ce7664c5 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
@@ -53,8 +53,8 @@ func TestMachine_getOperand_NR(t *testing.T) {
 				def = &backend.SSAValueDefinition{BlkParamVReg: regToVReg(x4), BlockParamValue: v}
 				return def, extModeZeroExtend64
 			},
-			exp:          operandNR(regToVReg(x4)),
-			instructions: []string{"uxtw x4, w4"},
+			exp:          operandNR(regalloc.VReg(1).SetRegType(regalloc.RegTypeInt)),
+			instructions: []string{"uxtw x1?, w4"},
 		},
 		{
 			name: "block param - sign extend",
@@ -64,8 +64,8 @@ func TestMachine_getOperand_NR(t *testing.T) {
 				def = &backend.SSAValueDefinition{BlkParamVReg: regToVReg(x4), BlockParamValue: v}
 				return def, extModeSignExtend64
 			},
-			exp:          operandNR(regToVReg(x4)),
-			instructions: []string{"sxtw x4, w4"},
+			exp:          operandNR(regalloc.VReg(1).SetRegType(regalloc.RegTypeInt)),
+			instructions: []string{"sxtw x1?, w4"},
 		},
 		{
 			name: "const instr",
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 0c3b6aa1..31ca7353 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_mem_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_mem_test.go
@@ -394,10 +394,10 @@ func TestMachine_collectAddends(t *testing.T) {
 						require.True(t, lowerd)
 					}
 					// Param must be zero-extended.
-					require.Equal(t, "uxtw x1000?, w1000?", formatEmittedInstructionsInCurrentBlock(m))
+					require.Equal(t, "uxtw x1?, w1000?", formatEmittedInstructionsInCurrentBlock(m))
 				}
 			},
-			exp64s: []regalloc.VReg{v1000},
+			exp64s: []regalloc.VReg{regalloc.VReg(1).SetRegType(regalloc.RegTypeInt)},
 			offset: 1 + 2 + 3 - 1,
 		},
 		{
diff --git a/internal/engine/wazevo/backend/regalloc/coloring_test.go b/internal/engine/wazevo/backend/regalloc/coloring_test.go
index 965cb0a4..88b02af8 100644
--- a/internal/engine/wazevo/backend/regalloc/coloring_test.go
+++ b/internal/engine/wazevo/backend/regalloc/coloring_test.go
@@ -1,7 +1,6 @@
 package regalloc
 
 import (
-	"fmt"
 	"testing"
 
 	"github.com/tetratelabs/wazero/internal/testing/require"
@@ -219,7 +218,6 @@ func TestAllocator_coloringFor(t *testing.T) {
 			var actual []string
 			for _, n := range testNodes {
 				actual = append(actual, n.r.String())
-				fmt.Println(n)
 			}
 			var exp []string
 			for _, r := range tc.expRegs {
diff --git a/internal/engine/wazevo/backend/regalloc/reg_test.go b/internal/engine/wazevo/backend/regalloc/reg_test.go
index 16c67fe0..ba311e61 100644
--- a/internal/engine/wazevo/backend/regalloc/reg_test.go
+++ b/internal/engine/wazevo/backend/regalloc/reg_test.go
@@ -15,8 +15,8 @@ func TestRegTypeOf(t *testing.T) {
 }
 
 func TestVReg_String(t *testing.T) {
-	require.Equal(t, "r0?", VReg(0).String())
-	require.Equal(t, "r100?", VReg(100).String())
+	require.Equal(t, "v0?", VReg(0).String())
+	require.Equal(t, "v100?", VReg(100).String())
 	require.Equal(t, "r5", FromRealReg(5, RegTypeInt).String())
 }