Repository: tetratelabs/wazero
Pull request #1344 (2023-04-06 07:34:04)

compiler: avoid alloc with stack pointer ceil and reuse bytes reader
```
goos: darwin
goarch: arm64
pkg: github.com/tetratelabs/wazero/internal/integration_test/bench
                                    │   old.txt   │             new2.txt              │
                                    │   sec/op    │   sec/op     vs base              │
Compilation/with_extern_cache-10      90.60µ ± 3%   89.70µ ± 2%  -0.99% (p=0.007 n=7)
Compilation/without_extern_cache-10   1.423m ± 1%   1.415m ± 1%  -0.59% (p=0.001 n=7)
Compilation/interpreter-10            408.3µ ± 0%   408.0µ ± 0%       ~ (p=0.710 n=7)
geomean                               374.8µ        372.7µ       -0.55%

                                    │   old.txt    │              new2.txt              │
                                    │     B/op     │     B/op      vs base              │
Compilation/with_extern_cache-10      38.25Ki ± 0%   38.24Ki ± 0%  -0.03% (p=0.015 n=7)
Compilation/without_extern_cache-10   769.0Ki ± 0%   766.1Ki ± 0%  -0.37% (p=0.001 n=7)
Compilation/interpreter-10            595.9Ki ± 0%   595.9Ki ± 0%       ~ (p=0.219 n=7)
geomean                               259.8Ki        259.4Ki       -0.14%

                                    │   old.txt   │              new2.txt               │
                                    │  allocs/op  │  allocs/op   vs base                │
Compilation/with_extern_cache-10       515.0 ± 0%    515.0 ± 0%       ~ (p=1.000 n=7) ¹
Compilation/without_extern_cache-10   2.042k ± 0%   1.961k ± 0%  -3.97% (p=0.001 n=7)
Compilation/interpreter-10            1.147k ± 0%   1.147k ± 0%       ~ (p=1.000 n=7) ¹
geomean                               1.064k        1.050k       -1.34%

```


Patch:
diff --git a/internal/engine/compiler/impl_amd64.go b/internal/engine/compiler/impl_amd64.go
index 0a903dd1..77c36b2f 100644
--- a/internal/engine/compiler/impl_amd64.go
+++ b/internal/engine/compiler/impl_amd64.go
@@ -92,10 +92,11 @@ type amd64Compiler struct {
 	labels [wazeroir.LabelKindNum][]amd64LabelInfo
 	// stackPointerCeil is the greatest stack pointer value (from runtimeValueLocationStack) seen during compilation.
 	stackPointerCeil uint64
-	// onStackPointerCeilDeterminedCallBack hold a callback which are called when the max stack pointer is determined BEFORE generating native code.
-	onStackPointerCeilDeterminedCallBack func(stackPointerCeil uint64)
-	withListener                         bool
-	typ                                  *wasm.FunctionType
+	// assignStackPointerCeilNeeded holds an asm.Node whose AssignDestinationConstant must be called with the determined stack pointer ceiling.
+	assignStackPointerCeilNeeded asm.Node
+	withListener                 bool
+	typ                          *wasm.FunctionType
+	br                           *bytes.Reader
 }
 
 func newAmd64Compiler() compiler {
@@ -103,6 +104,7 @@ func newAmd64Compiler() compiler {
 		assembler:     amd64.NewAssembler(),
 		locationStack: newRuntimeValueLocationStack(),
 		cpuFeatures:   platform.CpuFeatures,
+		br:            bytes.NewReader(nil),
 	}
 	return c
 }
@@ -123,6 +125,7 @@ func (c *amd64Compiler) Init(typ *wasm.FunctionType, ir *wazeroir.CompilationRes
 		withListener:  withListener,
 		labels:        c.labels,
 		typ:           typ,
+		br:            c.br,
 	}
 }
 
@@ -248,17 +251,15 @@ func (c *amd64Compiler) compile() (code []byte, stackPointerCeil uint64, err err
 
 	// Now that the max stack pointer is determined, we are invoking the callback.
 	// Note this MUST be called before Assemble() below.
-	if c.onStackPointerCeilDeterminedCallBack != nil {
-		c.onStackPointerCeilDeterminedCallBack(stackPointerCeil)
-		c.onStackPointerCeilDeterminedCallBack = nil
-	}
+	c.assignStackPointerCeil(stackPointerCeil)
 
 	code, err = c.assembler.Assemble()
 	if err != nil {
 		return
 	}
 
-	code, err = platform.MmapCodeSegment(bytes.NewReader(code), len(code))
+	c.br.Reset(code)
+	code, err = platform.MmapCodeSegment(c.br, len(code))
 	return
 }
 
@@ -268,6 +269,13 @@ func (c *amd64Compiler) compileUnreachable() error {
 	return nil
 }
 
+// assignStackPointerCeil implements compilerImpl.assignStackPointerCeil for the amd64 architecture.
+func (c *amd64Compiler) assignStackPointerCeil(ceil uint64) {
+	if c.assignStackPointerCeilNeeded != nil {
+		c.assignStackPointerCeilNeeded.AssignDestinationConstant(int64(ceil) << 3)
+	}
+}
+
 // compileSet implements compiler.compileSet for the amd64 architecture.
 func (c *amd64Compiler) compileSet(o *wazeroir.UnionOperation) error {
 	depth := int(o.U1)
@@ -4949,9 +4957,7 @@ func (c *amd64Compiler) compileMaybeGrowStack() error {
 
 	// If stack base pointer + max stack pointer > stackLen, we need to grow the stack.
 	cmpWithStackPointerCeil := c.assembler.CompileRegisterToConst(amd64.CMPQ, tmpRegister, 0)
-	c.onStackPointerCeilDeterminedCallBack = func(stackPointerCeil uint64) {
-		cmpWithStackPointerCeil.AssignDestinationConstant(int64(stackPointerCeil) << 3)
-	}
+	c.assignStackPointerCeilNeeded = cmpWithStackPointerCeil
 
 	// Jump if we have no need to grow.
 	jmpIfNoNeedToGrowStack := c.assembler.CompileJump(amd64.JCC)
diff --git a/internal/engine/compiler/impl_arm64.go b/internal/engine/compiler/impl_arm64.go
index 745d6d76..24c2d34a 100644
--- a/internal/engine/compiler/impl_arm64.go
+++ b/internal/engine/compiler/impl_arm64.go
@@ -26,16 +26,18 @@ type arm64Compiler struct {
 	labels [wazeroir.LabelKindNum][]arm64LabelInfo
 	// stackPointerCeil is the greatest stack pointer value (from runtimeValueLocationStack) seen during compilation.
 	stackPointerCeil uint64
-	// onStackPointerCeilDeterminedCallBack hold a callback which are called when the ceil of stack pointer is determined before generating native code.
-	onStackPointerCeilDeterminedCallBack func(stackPointerCeil uint64)
-	withListener                         bool
-	typ                                  *wasm.FunctionType
+	// assignStackPointerCeilNeeded holds an asm.Node whose AssignDestinationConstant must be called with the determined stack pointer ceiling.
+	assignStackPointerCeilNeeded asm.Node
+	withListener                 bool
+	typ                          *wasm.FunctionType
+	br                           *bytes.Reader
 }
 
 func newArm64Compiler() compiler {
 	return &arm64Compiler{
 		assembler:     arm64.NewAssembler(arm64ReservedRegisterForTemporary),
 		locationStack: newRuntimeValueLocationStack(),
+		br:            bytes.NewReader(nil),
 	}
 }
 
@@ -51,6 +53,7 @@ func (c *arm64Compiler) Init(typ *wasm.FunctionType, ir *wazeroir.CompilationRes
 		assembler: assembler, locationStack: locationStack,
 		ir: ir, withListener: withListener, labels: c.labels,
 		typ: typ,
+		br:  c.br,
 	}
 }
 
@@ -121,9 +124,7 @@ func (c *arm64Compiler) compile() (code []byte, stackPointerCeil uint64, err err
 
 	// Now that the ceil of stack pointer is determined, we are invoking the callback.
 	// Note: this must be called before Assemble() below.
-	if c.onStackPointerCeilDeterminedCallBack != nil {
-		c.onStackPointerCeilDeterminedCallBack(stackPointerCeil)
-	}
+	c.assignStackPointerCeil(stackPointerCeil)
 
 	var original []byte
 	original, err = c.assembler.Assemble()
@@ -131,7 +132,8 @@ func (c *arm64Compiler) compile() (code []byte, stackPointerCeil uint64, err err
 		return
 	}
 
-	code, err = platform.MmapCodeSegment(bytes.NewReader(original), len(original))
+	c.br.Reset(original)
+	code, err = platform.MmapCodeSegment(c.br, len(original))
 	return
 }
 
@@ -143,6 +145,13 @@ type arm64LabelInfo struct {
 	initialStack runtimeValueLocationStack
 }
 
+// assignStackPointerCeil implements compilerImpl.assignStackPointerCeil for the arm64 architecture.
+func (c *arm64Compiler) assignStackPointerCeil(ceil uint64) {
+	if c.assignStackPointerCeilNeeded != nil {
+		c.assignStackPointerCeilNeeded.AssignSourceConstant(int64(ceil) << 3)
+	}
+}
+
 func (c *arm64Compiler) label(label wazeroir.Label) *arm64LabelInfo {
 	kind := label.Kind()
 	frames := c.labels[kind]
@@ -268,9 +277,7 @@ func (c *arm64Compiler) compileMaybeGrowStack() error {
 	)
 	// At this point of compilation, we don't know the value of stack point ceil,
 	// so we lazily resolve the value later.
-	c.onStackPointerCeilDeterminedCallBack = func(stackPointerCeil uint64) {
-		loadStackPointerCeil.AssignSourceConstant(int64(stackPointerCeil) << 3)
-	}
+	c.assignStackPointerCeilNeeded = loadStackPointerCeil
 
 	// Compare tmpX (len(ce.stack) - ce.stackBasePointer) and tmpY (ce.stackPointerCeil)
 	c.assembler.CompileTwoRegistersToNone(arm64.CMP, tmpX, tmpY)


Test patch:
diff --git a/internal/engine/compiler/compiler_initialization_test.go b/internal/engine/compiler/compiler_initialization_test.go
index 690f8765..c2c153f7 100644
--- a/internal/engine/compiler/compiler_initialization_test.go
+++ b/internal/engine/compiler/compiler_initialization_test.go
@@ -184,11 +184,9 @@ func TestCompiler_compileMaybeGrowStack(t *testing.T) {
 				err := compiler.compilePreamble()
 				require.NoError(t, err)
 
-				require.NotNil(t, compiler.getOnStackPointerCeilDeterminedCallBack())
-
 				stackLen := uint64(len(env.stack()))
 				stackBasePointer := stackLen - baseOffset // Ceil <= stackLen - stackBasePointer = no need to grow!
-				compiler.getOnStackPointerCeilDeterminedCallBack()(stackPointerCeil)
+				compiler.assignStackPointerCeil(stackPointerCeil)
 				env.setStackBasePointer(stackBasePointer)
 
 				compiler.compileExitFromNativeCode(nativeCallStatusCodeReturned)
diff --git a/internal/engine/compiler/compiler_test.go b/internal/engine/compiler/compiler_test.go
index ce4dcd9b..2cc096e9 100644
--- a/internal/engine/compiler/compiler_test.go
+++ b/internal/engine/compiler/compiler_test.go
@@ -237,14 +237,14 @@ func (j *compilerEnv) requireNewCompiler(t *testing.T, functionType *wasm.Functi
 	return ret
 }
 
-// CompilerImpl is the interface used for architecture-independent unit tests in this pkg.
+// compilerImpl is the interface used for architecture-independent unit tests in this pkg.
 // This is currently implemented by amd64 and arm64.
 type compilerImpl interface {
 	compiler
 	compileExitFromNativeCode(nativeCallStatusCode)
 	compileMaybeGrowStack() error
 	compileReturnFunction() error
-	getOnStackPointerCeilDeterminedCallBack() func(uint64)
+	assignStackPointerCeil(uint64)
 	setStackPointerCeil(uint64)
 	compileReleaseRegisterToStack(loc *runtimeValueLocation)
 	setRuntimeValueLocationStack(runtimeValueLocationStack)
diff --git a/internal/engine/compiler/impl_amd64_test.go b/internal/engine/compiler/impl_amd64_test.go
index 4acbf3ee..26a3b52e 100644
--- a/internal/engine/compiler/impl_amd64_test.go
+++ b/internal/engine/compiler/impl_amd64_test.go
@@ -577,11 +577,6 @@ func collectRegistersFromRuntimeValues(locs []*runtimeValueLocation) []asm.Regis
 	return out
 }
 
-// compile implements compilerImpl.getOnStackPointerCeilDeterminedCallBack for the amd64 architecture.
-func (c *amd64Compiler) getOnStackPointerCeilDeterminedCallBack() func(uint64) {
-	return c.onStackPointerCeilDeterminedCallBack
-}
-
 // compile implements compilerImpl.setStackPointerCeil for the amd64 architecture.
 func (c *amd64Compiler) setStackPointerCeil(v uint64) {
 	c.stackPointerCeil = v
diff --git a/internal/engine/compiler/impl_arm64_test.go b/internal/engine/compiler/impl_arm64_test.go
index cab0d2b8..f385fbd6 100644
--- a/internal/engine/compiler/impl_arm64_test.go
+++ b/internal/engine/compiler/impl_arm64_test.go
@@ -100,11 +100,6 @@ func TestArm64Compiler_readInstructionAddress(t *testing.T) {
 	require.Equal(t, nativeCallStatusCodeReturned, env.compilerStatus())
 }
 
-// compile implements compilerImpl.getOnStackPointerCeilDeterminedCallBack for the amd64 architecture.
-func (c *arm64Compiler) getOnStackPointerCeilDeterminedCallBack() func(uint64) {
-	return c.onStackPointerCeilDeterminedCallBack
-}
-
 // compile implements compilerImpl.setStackPointerCeil for the amd64 architecture.
 func (c *arm64Compiler) setStackPointerCeil(v uint64) {
 	c.stackPointerCeil = v