Repository: tetratelabs/wazero
Pull request #726 (2022-07-28 07:51:27)

amd64: fixes i64x2.shr_s on CX register allocation
Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>


Patch:
diff --git a/internal/engine/compiler/arch.go b/internal/engine/compiler/arch.go
index ce6c69f4..b6182514 100644
--- a/internal/engine/compiler/arch.go
+++ b/internal/engine/compiler/arch.go
@@ -1,5 +1,7 @@
 package compiler
 
+import "github.com/tetratelabs/wazero/internal/asm"
+
 // newArchContext returns a new archContext which is architecture-specific type to be embedded in callEngine.
 // This must be initialized in init() function in architecture-specific arch_*.go file which is guarded by build tag.
 var newArchContext func() archContext
@@ -10,3 +12,6 @@ var newArchContext func() archContext
 //
 // Note: this is implemented in per-arch Go assembler file. For example, arch_amd64.s implements this for amd64.
 func nativecall(codeSegment, ce uintptr, moduleInstanceAddress uintptr)
+
+// registerNameFn is used for debugging purpose to have register symbols in the string of runtimeValueLocation.
+var registerNameFn func(register asm.Register) string
diff --git a/internal/engine/compiler/arch_amd64.go b/internal/engine/compiler/arch_amd64.go
index 40080b30..c33f80c6 100644
--- a/internal/engine/compiler/arch_amd64.go
+++ b/internal/engine/compiler/arch_amd64.go
@@ -1,12 +1,14 @@
 package compiler
 
 import (
+	"github.com/tetratelabs/wazero/internal/asm/amd64"
 	"github.com/tetratelabs/wazero/internal/wazeroir"
 )
 
 // init initializes variables for the amd64 architecture
 func init() {
 	newArchContext = newArchContextImpl
+	registerNameFn = amd64.RegisterName
 	unreservedGeneralPurposeRegisters = amd64UnreservedGeneralPurposeRegisters
 	unreservedVectorRegisters = amd64UnreservedVectorRegisters
 }
diff --git a/internal/engine/compiler/arch_arm64.go b/internal/engine/compiler/arch_arm64.go
index 6aa771ca..d975d9bd 100644
--- a/internal/engine/compiler/arch_arm64.go
+++ b/internal/engine/compiler/arch_arm64.go
@@ -3,12 +3,14 @@ package compiler
 import (
 	"math"
 
+	"github.com/tetratelabs/wazero/internal/asm/arm64"
 	"github.com/tetratelabs/wazero/internal/wazeroir"
 )
 
 // init initializes variables for the arm64 architecture
 func init() {
 	newArchContext = newArchContextImpl
+	registerNameFn = arm64.RegisterName
 	unreservedGeneralPurposeRegisters = arm64UnreservedGeneralPurposeRegisters
 	unreservedVectorRegisters = arm64UnreservedVectorRegisters
 }
diff --git a/internal/engine/compiler/compiler_value_location.go b/internal/engine/compiler/compiler_value_location.go
index bcf16e74..41a6a692 100644
--- a/internal/engine/compiler/compiler_value_location.go
+++ b/internal/engine/compiler/compiler_value_location.go
@@ -104,13 +104,17 @@ func (v *runtimeValueLocation) String() string {
 	} else if v.onConditionalRegister() {
 		location = fmt.Sprintf("conditional(%d)", v.conditionalRegister)
 	} else if v.onRegister() {
-		location = fmt.Sprintf("register(%d)", v.register)
+		location = fmt.Sprintf("register(%s)", registerNameFn(v.register))
 	}
 	return fmt.Sprintf("{type=%s,location=%s}", v.valueType, location)
 }
 
 func newRuntimeValueLocationStack() *runtimeValueLocationStack {
-	return &runtimeValueLocationStack{usedRegisters: map[asm.Register]struct{}{}}
+	return &runtimeValueLocationStack{
+		usedRegisters:                     map[asm.Register]struct{}{},
+		unreservedVectorRegisters:         unreservedVectorRegisters,
+		unreservedGeneralPurposeRegisters: unreservedGeneralPurposeRegisters,
+	}
 }
 
 // runtimeValueLocationStack represents the wazeroir virtual stack
@@ -132,6 +136,9 @@ type runtimeValueLocationStack struct {
 	usedRegisters map[asm.Register]struct{}
 	// stackPointerCeil tracks max(.sp) across the lifespan of this struct.
 	stackPointerCeil uint64
+	// unreservedGeneralPurposeRegisters and unreservedVectorRegisters hold
+	// architecture dependent unreserved register list.
+	unreservedGeneralPurposeRegisters, unreservedVectorRegisters []asm.Register
 }
 
 func (v *runtimeValueLocationStack) String() string {
@@ -141,7 +148,7 @@ func (v *runtimeValueLocationStack) String() string {
 	}
 	var usedRegisters []string
 	for reg := range v.usedRegisters {
-		usedRegisters = append(usedRegisters, fmt.Sprintf("%d", reg))
+		usedRegisters = append(usedRegisters, registerNameFn(reg))
 	}
 	return fmt.Sprintf("sp=%d, stack=[%s], used_registers=[%s]", v.sp, strings.Join(stackStr, ","), strings.Join(usedRegisters, ","))
 }
@@ -163,6 +170,8 @@ func (v *runtimeValueLocationStack) clone() *runtimeValueLocationStack {
 		}
 	}
 	ret.stackPointerCeil = v.stackPointerCeil
+	ret.unreservedGeneralPurposeRegisters = v.unreservedGeneralPurposeRegisters
+	ret.unreservedVectorRegisters = v.unreservedVectorRegisters
 	return ret
 }
 
@@ -275,9 +284,9 @@ func (v *runtimeValueLocationStack) takeFreeRegister(tp registerType) (reg asm.R
 	var targetRegs []asm.Register
 	switch tp {
 	case registerTypeVector:
-		targetRegs = unreservedVectorRegisters
+		targetRegs = v.unreservedVectorRegisters
 	case registerTypeGeneralPurpose:
-		targetRegs = unreservedGeneralPurposeRegisters
+		targetRegs = v.unreservedGeneralPurposeRegisters
 	}
 	for _, candidate := range targetRegs {
 		if _, ok := v.usedRegisters[candidate]; ok {
@@ -292,9 +301,9 @@ func (v *runtimeValueLocationStack) takeFreeRegisters(tp registerType, num int)
 	var targetRegs []asm.Register
 	switch tp {
 	case registerTypeVector:
-		targetRegs = unreservedVectorRegisters
+		targetRegs = v.unreservedVectorRegisters
 	case registerTypeGeneralPurpose:
-		targetRegs = unreservedGeneralPurposeRegisters
+		targetRegs = v.unreservedGeneralPurposeRegisters
 	}
 
 	regs = make([]asm.Register, 0, num)
diff --git a/internal/engine/compiler/impl_vec_amd64.go b/internal/engine/compiler/impl_vec_amd64.go
index f030316c..0aedbee0 100644
--- a/internal/engine/compiler/impl_vec_amd64.go
+++ b/internal/engine/compiler/impl_vec_amd64.go
@@ -835,22 +835,25 @@ func (c *amd64Compiler) compileV128ShrImpl(o *wazeroir.OperationV128Shr) error {
 // PSRAQ instruction requires AVX, so we emulate it without AVX instructions. https://www.felixcloutier.com/x86/psraw:psrad:psraq
 func (c *amd64Compiler) compileV128ShrI64x2SignedImpl() error {
 	const shiftCountRegister = amd64.RegCX
-	// If another value lives on the CX register, we release it to the stack.
-	c.onValueReleaseRegisterToStack(shiftCountRegister)
 
 	s := c.locationStack.pop()
-	if s.onStack() {
-		s.setRegister(shiftCountRegister)
-		c.compileLoadValueOnStackToRegister(s)
-	} else if s.onConditionalRegister() {
-		c.compileMoveConditionalToGeneralPurposeRegister(s, shiftCountRegister)
-	} else { // already on register.
-		old := s.register
-		c.assembler.CompileRegisterToRegister(amd64.MOVL, old, shiftCountRegister)
-		s.setRegister(shiftCountRegister)
-		c.locationStack.markRegisterUnused(old)
-	}
-	c.locationStack.markRegisterUnused(shiftCountRegister)
+	if s.register != shiftCountRegister {
+		// If another value lives on the CX register, we release it to the stack.
+		c.onValueReleaseRegisterToStack(shiftCountRegister)
+		if s.onStack() {
+			s.setRegister(shiftCountRegister)
+			c.compileLoadValueOnStackToRegister(s)
+		} else if s.onConditionalRegister() {
+			c.compileMoveConditionalToGeneralPurposeRegister(s, shiftCountRegister)
+		} else { // already on register.
+			old := s.register
+			c.assembler.CompileRegisterToRegister(amd64.MOVL, old, shiftCountRegister)
+			s.setRegister(shiftCountRegister)
+			c.locationStack.markRegisterUnused(old)
+		}
+	}
+
+	c.locationStack.markRegisterUsed(shiftCountRegister)
 	tmp, err := c.allocateRegister(registerTypeGeneralPurpose)
 	if err != nil {
 		return err


Test patch:
diff --git a/internal/engine/compiler/compiler_stack_test.go b/internal/engine/compiler/compiler_stack_test.go
index 261e1cb8..7e3f6d3f 100644
--- a/internal/engine/compiler/compiler_stack_test.go
+++ b/internal/engine/compiler/compiler_stack_test.go
@@ -33,11 +33,13 @@ func TestCompiler_releaseRegisterToStack(t *testing.T) {
 			err := compiler.compilePreamble()
 			require.NoError(t, err)
 
-			// Setup the location stack so that we push the const on the specified height.
+			// Set up the location stack so that we push the const on the specified height.
 			s := &runtimeValueLocationStack{
-				sp:            tc.stackPointer,
-				stack:         make([]*runtimeValueLocation, tc.stackPointer),
-				usedRegisters: map[asm.Register]struct{}{},
+				sp:                                tc.stackPointer,
+				stack:                             make([]*runtimeValueLocation, tc.stackPointer),
+				usedRegisters:                     map[asm.Register]struct{}{},
+				unreservedVectorRegisters:         unreservedVectorRegisters,
+				unreservedGeneralPurposeRegisters: unreservedGeneralPurposeRegisters,
 			}
 			// Peek must be non-nil. Otherwise, compileConst* would fail.
 			s.stack[s.sp-1] = &runtimeValueLocation{}
diff --git a/internal/engine/compiler/compiler_value_location_test.go b/internal/engine/compiler/compiler_value_location_test.go
index 25158754..f86a57d9 100644
--- a/internal/engine/compiler/compiler_value_location_test.go
+++ b/internal/engine/compiler/compiler_value_location_test.go
@@ -44,6 +44,8 @@ func TestRuntimeValueLocationStack_basic(t *testing.T) {
 	// Clone.
 	cloned := s.clone()
 	require.Equal(t, s.usedRegisters, cloned.usedRegisters)
+	require.Equal(t, s.unreservedGeneralPurposeRegisters, cloned.unreservedGeneralPurposeRegisters)
+	require.Equal(t, s.unreservedVectorRegisters, cloned.unreservedVectorRegisters)
 	require.Equal(t, len(s.stack), len(cloned.stack))
 	require.Equal(t, s.sp, cloned.sp)
 	for i := 0; i < int(s.sp); i++ {
diff --git a/internal/engine/compiler/impl_vec_amd64_test.go b/internal/engine/compiler/impl_vec_amd64_test.go
index f1671ad3..f842e283 100644
--- a/internal/engine/compiler/impl_vec_amd64_test.go
+++ b/internal/engine/compiler/impl_vec_amd64_test.go
@@ -5,6 +5,7 @@ import (
 	"math"
 	"testing"
 
+	"github.com/tetratelabs/wazero/internal/asm"
 	"github.com/tetratelabs/wazero/internal/asm/amd64"
 	"github.com/tetratelabs/wazero/internal/testing/require"
 	"github.com/tetratelabs/wazero/internal/wasm"
@@ -93,6 +94,31 @@ func TestAmd64Compiler_compileV128ShrI64x2SignedImpl(t *testing.T) {
 			},
 			verifyFn: func(t *testing.T, env *compilerEnv) {},
 		},
+		{
+			name: "RegR10/CX not in use and CX is next free register",
+			shiftAmountSetupFn: func(t *testing.T, c *amd64Compiler) {
+				// Move the shift amount to R10.
+				loc := c.locationStack.peek()
+				oldReg, newReg := loc.register, amd64.RegR10
+				c.assembler.CompileRegisterToRegister(amd64.MOVQ, oldReg, newReg)
+				loc.setRegister(newReg)
+				c.locationStack.markRegisterUnused(oldReg)
+				c.locationStack.markRegisterUsed(newReg)
+
+				// Ensures that the next free becomes CX.
+				newUnreservedRegs := make([]asm.Register, len(c.locationStack.unreservedVectorRegisters))
+				copy(newUnreservedRegs, c.locationStack.unreservedGeneralPurposeRegisters)
+				for i, r := range newUnreservedRegs {
+					// If CX register is found, we swap it with the first register in the list.
+					// This forces runtimeLocationStack to take CX as a first free register.
+					if r == amd64.RegCX {
+						newUnreservedRegs[0], newUnreservedRegs[i] = newUnreservedRegs[i], newUnreservedRegs[0]
+					}
+				}
+				c.locationStack.unreservedGeneralPurposeRegisters = newUnreservedRegs
+			},
+			verifyFn: func(t *testing.T, env *compilerEnv) {},
+		},
 		{
 			name: "RegR10/CX in use",
 			shiftAmountSetupFn: func(t *testing.T, c *amd64Compiler) {