Repository: tetratelabs/wazero
Pull request #898 (2022-12-07 05:54:50)

compiler: make moduleContext.functions non-pointer slice
Ref #889 

This further reduces the allocations during instantiation

```
$ benchstat old_re2.txt new_re2.txt
name                                    old time/op    new time/op    delta
Initialization/interpreter-10              136µs ± 2%     137µs ± 2%    ~     (p=0.222 n=5+5)
Initialization/interpreter-multiple-10    62.5µs ±17%    67.4µs ±13%    ~     (p=0.421 n=5+5)
Initialization/compiler-10                 167µs ± 5%     156µs ± 4%  -6.66%  (p=0.008 n=5+5)
Initialization/compiler-multiple-10       44.1µs ± 8%    43.0µs ± 9%    ~     (p=0.548 n=5+5)

name                                    old alloc/op   new alloc/op   delta
Initialization/interpreter-10              550kB ± 0%     550kB ± 0%    ~     (p=0.881 n=5+5)
Initialization/interpreter-multiple-10     551kB ± 0%     551kB ± 0%    ~     (p=0.548 n=5+5)
Initialization/compiler-10                 506kB ± 0%     498kB ± 0%  -1.62%  (p=0.008 n=5+5)
Initialization/compiler-multiple-10        507kB ± 0%     498kB ± 0%  -1.60%  (p=0.008 n=5+5)

name                                    old allocs/op  new allocs/op  delta
Initialization/interpreter-10              1.88k ± 0%     1.88k ± 0%    ~     (all equal)
Initialization/interpreter-multiple-10     1.88k ± 0%     1.88k ± 0%    ~     (all equal)
Initialization/compiler-10                  37.0 ± 0%      36.0 ± 0%  -2.70%  (p=0.008 n=5+5)
Initialization/compiler-multiple-10         44.0 ± 0%      43.0 ± 0%  -2.27%  (p=0.008 n=5+5)
```


Patch:
diff --git a/internal/engine/compiler/engine.go b/internal/engine/compiler/engine.go
index e22fcac4..a2f24614 100644
--- a/internal/engine/compiler/engine.go
+++ b/internal/engine/compiler/engine.go
@@ -48,7 +48,7 @@ type (
 		// The index is module instance-scoped. We intentionally avoid using map
 		// as the underlying memory region is accessed by assembly directly by using
 		// codesElement0Address.
-		functions []*function
+		functions []function
 	}
 
 	// callEngine holds context per moduleEngine.Call, and shared across all the
@@ -302,21 +302,6 @@ type (
 	}
 )
 
-// createFunction creates a new function which uses the native code compiled.
-func (c *code) createFunction(f *wasm.FunctionInstance) function {
-	// This object is only referenced from a slice. Instead of creating a heap object
-	// here and returning a pointer, we return the struct directly, so it can be
-	// stored directly in the slice. This reduces the number of heap objects which
-	// improves GC performance.
-	return function{
-		codeInitialAddress:    uintptr(unsafe.Pointer(&c.codeSegment[0])),
-		stackPointerCeil:      c.stackPointerCeil,
-		moduleInstanceAddress: uintptr(unsafe.Pointer(f.Module)),
-		source:                f,
-		parent:                c,
-	}
-}
-
 // Native code reads/writes Go's structs with the following constants.
 // See TestVerifyOffsetValue for how to derive these values.
 const (
@@ -351,6 +336,7 @@ const (
 	functionCodeInitialAddressOffset    = 0
 	functionSourceOffset                = 16
 	functionModuleInstanceAddressOffset = 24
+	functionSize                        = 40
 
 	// Offsets for wasm.ModuleInstance.
 	moduleInstanceGlobalsOffset          = 48
@@ -546,7 +532,7 @@ func (e *engine) CompileModule(ctx context.Context, module *wasm.Module, listene
 func (e *engine) NewModuleEngine(name string, module *wasm.Module, functions []wasm.FunctionInstance) (wasm.ModuleEngine, error) {
 	me := &moduleEngine{
 		name:      name,
-		functions: make([]*function, len(functions)),
+		functions: make([]function, len(functions)),
 	}
 
 	imported := int(module.ImportFuncCount())
@@ -562,14 +548,16 @@ func (e *engine) NewModuleEngine(name string, module *wasm.Module, functions []w
 		return nil, err
 	}
 
-	// TODO: Ideally we only have one []function slice, instead of a memory store and a slice of pointers.
-	// This will require refactoring in assembly, so for now we simply have two slices.
-	functionsMem := make([]function, len(functions[imported:]))
 	for i, c := range codes {
 		offset := imported + i
 		f := &functions[offset]
-		functionsMem[i] = c.createFunction(f)
-		me.functions[offset] = &functionsMem[i]
+		me.functions[offset] = function{
+			codeInitialAddress:    uintptr(unsafe.Pointer(&c.codeSegment[0])),
+			stackPointerCeil:      c.stackPointerCeil,
+			moduleInstanceAddress: uintptr(unsafe.Pointer(f.Module)),
+			source:                f,
+			parent:                c,
+		}
 	}
 	return me, nil
 }
@@ -581,7 +569,7 @@ func (e *moduleEngine) Name() string {
 
 // FunctionInstanceReference implements the same method as documented on wasm.ModuleEngine.
 func (e *moduleEngine) FunctionInstanceReference(funcIndex wasm.Index) wasm.Reference {
-	return uintptr(unsafe.Pointer(e.functions[funcIndex]))
+	return uintptr(unsafe.Pointer(&e.functions[funcIndex]))
 }
 
 // CreateFuncElementInstance implements the same method as documented on wasm.ModuleEngine.
@@ -589,7 +577,7 @@ func (e *moduleEngine) CreateFuncElementInstance(indexes []*wasm.Index) *wasm.El
 	refs := make([]wasm.Reference, len(indexes))
 	for i, index := range indexes {
 		if index != nil {
-			refs[i] = uintptr(unsafe.Pointer(e.functions[*index]))
+			refs[i] = uintptr(unsafe.Pointer(&e.functions[*index]))
 		}
 	}
 	return &wasm.ElementInstance{
@@ -598,16 +586,10 @@ func (e *moduleEngine) CreateFuncElementInstance(indexes []*wasm.Index) *wasm.El
 	}
 }
 
-func (e *moduleEngine) NewCallEngine(callCtx *wasm.CallContext, f *wasm.FunctionInstance) (ce wasm.CallEngine, err error) {
+func (e *moduleEngine) NewCallEngine(_ *wasm.CallContext, f *wasm.FunctionInstance) (ce wasm.CallEngine, err error) {
 	// Note: The input parameters are pre-validated, so a compiled function is only absent on close. Updates to
 	// code on close aren't locked, neither is this read.
-	compiled := e.functions[f.Idx]
-	if compiled == nil { // Lazy check the cause as it could be because the module was already closed.
-		if err = callCtx.FailIfClosed(); err == nil {
-			panic(fmt.Errorf("BUG: %s.func[%d] was nil before close", e.name, f.Idx))
-		}
-		return
-	}
+	compiled := &e.functions[f.Idx]
 
 	initStackSize := initialStackSize
 	if initialStackSize < compiled.stackPointerCeil {
diff --git a/internal/engine/compiler/impl_amd64.go b/internal/engine/compiler/impl_amd64.go
index b45b0f38..8d799306 100644
--- a/internal/engine/compiler/impl_amd64.go
+++ b/internal/engine/compiler/impl_amd64.go
@@ -696,19 +696,13 @@ func (c *amd64Compiler) compileCall(o *wazeroir.OperationCall) error {
 		return err
 	}
 
-	// First, we read the address of the first item of callEngine.functions slice (= &callEngine.functions[0])
-	// into tmpRegister.
-	c.assembler.CompileMemoryToRegister(amd64.MOVQ, amd64ReservedRegisterForCallEngine,
-		callEngineModuleContextFunctionsElement0AddressOffset, targetAddressRegister)
+	// First, push the index to the callEngine.functionsElement0Address into the target register.
+	c.assembler.CompileConstToRegister(amd64.MOVQ, int64(o.FunctionIndex)*functionSize, targetAddressRegister)
 
-	// next, read the address of the target function (= &callEngine.codes[offset])
-	// into targetAddressRegister.
-	c.assembler.CompileMemoryToRegister(amd64.MOVQ,
-		// Note: FunctionIndex is limited up to 2^27 so this offset never exceeds 32-bit integer.
-		// *8 because the size of *code equals 8 bytes.
-		targetAddressRegister, int64(o.FunctionIndex)*8,
-		targetAddressRegister,
-	)
+	// Next, we add the address of the first item of callEngine.functions slice (= &callEngine.functions[0])
+	// to the target register.
+	c.assembler.CompileMemoryToRegister(amd64.ADDQ, amd64ReservedRegisterForCallEngine,
+		callEngineModuleContextFunctionsElement0AddressOffset, targetAddressRegister)
 
 	if err := c.compileCallFunctionImpl(targetAddressRegister, targetType); err != nil {
 		return err
@@ -4231,20 +4225,15 @@ func (c *amd64Compiler) compileRefFunc(o *wazeroir.OperationRefFunc) error {
 		return err
 	}
 
-	// ref = [amd64ReservedRegisterForCallEngine + callEngineModuleContextFunctionsElement0AddressOffset]
-	//     = &moduleEngine.functions[0]
+	c.assembler.CompileConstToRegister(amd64.MOVQ, int64(o.FunctionIndex)*functionSize, ref)
+
+	// ref = [amd64ReservedRegisterForCallEngine + callEngineModuleContextFunctionsElement0AddressOffset + int64(o.FunctionIndex)*functionSize]
+	//     = &moduleEngine.functions[index]
 	c.assembler.CompileMemoryToRegister(
-		amd64.MOVQ, amd64ReservedRegisterForCallEngine, callEngineModuleContextFunctionsElement0AddressOffset,
+		amd64.ADDQ, amd64ReservedRegisterForCallEngine, callEngineModuleContextFunctionsElement0AddressOffset,
 		ref,
 	)
 
-	// ref = [ref +  int64(o.FunctionIndex)*8]
-	//     = [&moduleEngine.functions[0] + sizeOf(*function) * index]
-	//     = moduleEngine.functions[index]
-	c.assembler.CompileMemoryToRegister(amd64.MOVQ,
-		ref, int64(o.FunctionIndex)*8, // * 8 because the size of *code equals 8 bytes.
-		ref,
-	)
 	c.pushRuntimeValueLocationOnRegister(ref, runtimeValueTypeI64)
 	return nil
 }
diff --git a/internal/engine/compiler/impl_arm64.go b/internal/engine/compiler/impl_arm64.go
index ec9035f0..3da4238b 100644
--- a/internal/engine/compiler/impl_arm64.go
+++ b/internal/engine/compiler/impl_arm64.go
@@ -898,9 +898,9 @@ func (c *arm64Compiler) compileCall(o *wazeroir.OperationCall) error {
 		arm64ReservedRegisterForCallEngine, callEngineModuleContextFunctionsElement0AddressOffset,
 		targetFunctionAddressReg)
 
-	c.assembler.CompileMemoryToRegister(
-		arm64.LDRD,
-		targetFunctionAddressReg, int64(o.FunctionIndex)*8, // * 8 because the size of *function equals 8 bytes.
+	c.assembler.CompileConstToRegister(
+		arm64.ADD,
+		int64(o.FunctionIndex)*functionSize, // * 8 because the size of *function equals 8 bytes.
 		targetFunctionAddressReg)
 
 	return c.compileCallImpl(targetFunctionAddressReg, tp)
@@ -3577,13 +3577,12 @@ func (c *arm64Compiler) compileRefFunc(o *wazeroir.OperationRefFunc) error {
 	//                                   = &moduleEngine.functions[0]
 	c.assembler.CompileMemoryToRegister(arm64.LDRD,
 		arm64ReservedRegisterForCallEngine, callEngineModuleContextFunctionsElement0AddressOffset,
-		arm64ReservedRegisterForTemporary)
+		ref)
 
-	// ref = [arm64ReservedRegisterForTemporary +  int64(o.FunctionIndex)*8]
-	//     = [&moduleEngine.functions[0] + sizeOf(*function) * index]
-	//     = moduleEngine.functions[index]
-	c.assembler.CompileMemoryToRegister(arm64.LDRD,
-		arm64ReservedRegisterForTemporary, int64(o.FunctionIndex)*8, // * 8 because the size of *code equals 8 bytes.
+	// ref = ref + int64(o.FunctionIndex)*sizeOf(function)
+	//     = &moduleEngine.functions[index]
+	c.assembler.CompileConstToRegister(arm64.ADD,
+		int64(o.FunctionIndex)*functionSize,
 		ref,
 	)
 


Test patch:
diff --git a/internal/engine/compiler/compiler_controlflow_test.go b/internal/engine/compiler/compiler_controlflow_test.go
index 5a6b3e69..2a07d97b 100644
--- a/internal/engine/compiler/compiler_controlflow_test.go
+++ b/internal/engine/compiler/compiler_controlflow_test.go
@@ -651,7 +651,7 @@ func TestCompiler_compileCallIndirect(t *testing.T) {
 		// and the typeID matches the table[targetOffset]'s type ID.
 		env.module().TypeIDs = make([]wasm.FunctionTypeID, 100)
 		env.module().TypeIDs[operation.TypeIndex] = targetTypeID
-		env.module().Engine = &moduleEngine{functions: []*function{}}
+		env.module().Engine = &moduleEngine{functions: []function{}}
 
 		me := env.moduleEngine()
 		for i := 0; i < len(table); i++ {
@@ -679,16 +679,15 @@ func TestCompiler_compileCallIndirect(t *testing.T) {
 
 			// Now that we've generated the code for this function,
 			// add it to the module engine and assign its pointer to the table index.
-			f := &function{
+			me.functions = append(me.functions, function{
 				parent:                &code{codeSegment: c},
 				codeInitialAddress:    uintptr(unsafe.Pointer(&c[0])),
 				moduleInstanceAddress: uintptr(unsafe.Pointer(env.moduleInstance)),
 				source: &wasm.FunctionInstance{
 					TypeID: targetTypeID,
 				},
-			}
-			me.functions = append(me.functions, f)
-			table[i] = uintptr(unsafe.Pointer(f))
+			})
+			table[i] = uintptr(unsafe.Pointer(&me.functions[len(me.functions)-1]))
 		}
 
 		// Test to ensure that we can call all the functions stored in the table.
@@ -746,7 +745,7 @@ func TestCompiler_callIndirect_largeTypeIndex(t *testing.T) {
 	operation := &wazeroir.OperationCallIndirect{TypeIndex: typeIndex}
 	env.module().TypeIDs = make([]wasm.FunctionTypeID, typeIndex+1)
 	env.module().TypeIDs[typeIndex] = typeID
-	env.module().Engine = &moduleEngine{functions: []*function{}}
+	env.module().Engine = &moduleEngine{functions: []function{}}
 
 	types := make([]*wasm.FunctionType, typeIndex+1)
 	types[typeIndex] = &wasm.FunctionType{}
@@ -762,14 +761,14 @@ func TestCompiler_callIndirect_largeTypeIndex(t *testing.T) {
 		c, _, err := compiler.compile()
 		require.NoError(t, err)
 
-		f := &function{
+		f := function{
 			parent:                &code{codeSegment: c},
 			codeInitialAddress:    uintptr(unsafe.Pointer(&c[0])),
 			moduleInstanceAddress: uintptr(unsafe.Pointer(env.moduleInstance)),
 			source:                &wasm.FunctionInstance{TypeID: 0},
 		}
 		me.functions = append(me.functions, f)
-		table[0] = uintptr(unsafe.Pointer(f))
+		table[0] = uintptr(unsafe.Pointer(&f))
 	}
 
 	compiler := env.requireNewCompiler(t, newCompiler, &wazeroir.CompilationResult{
@@ -835,7 +834,7 @@ func TestCompiler_compileCall(t *testing.T) {
 		c, _, err := compiler.compile()
 		require.NoError(t, err)
 		index := wasm.Index(i)
-		me.functions = append(me.functions, &function{
+		me.functions = append(me.functions, function{
 			parent:                &code{codeSegment: c},
 			codeInitialAddress:    uintptr(unsafe.Pointer(&c[0])),
 			moduleInstanceAddress: uintptr(unsafe.Pointer(env.moduleInstance)),
diff --git a/internal/engine/compiler/compiler_initialization_test.go b/internal/engine/compiler/compiler_initialization_test.go
index fd396d5f..da580f1f 100644
--- a/internal/engine/compiler/compiler_initialization_test.go
+++ b/internal/engine/compiler/compiler_initialization_test.go
@@ -120,7 +120,7 @@ func TestCompiler_compileModuleContextInitialization(t *testing.T) {
 				ir.Globals = append(ir.Globals, g.Type)
 			}
 			compiler := env.requireNewCompiler(t, newCompiler, ir)
-			me := &moduleEngine{functions: make([]*function, 10)}
+			me := &moduleEngine{functions: make([]function, 10)}
 			tc.moduleInstance.Engine = me
 
 			err := compiler.compileModuleContextInitialization()
diff --git a/internal/engine/compiler/compiler_post1_0_test.go b/internal/engine/compiler/compiler_post1_0_test.go
index 788a52de..fbf87017 100644
--- a/internal/engine/compiler/compiler_post1_0_test.go
+++ b/internal/engine/compiler/compiler_post1_0_test.go
@@ -949,7 +949,7 @@ func TestCompiler_compileRefFunc(t *testing.T) {
 	me := env.moduleEngine()
 	const numFuncs = 20
 	for i := 0; i < numFuncs; i++ {
-		me.functions = append(me.functions, &function{source: &wasm.FunctionInstance{}})
+		me.functions = append(me.functions, function{source: &wasm.FunctionInstance{}})
 	}
 
 	for i := 0; i < numFuncs; i++ {
@@ -974,7 +974,7 @@ func TestCompiler_compileRefFunc(t *testing.T) {
 
 			require.Equal(t, nativeCallStatusCodeReturned, env.compilerStatus())
 			require.Equal(t, uint64(1), env.stackPointer())
-			require.Equal(t, uintptr(unsafe.Pointer(me.functions[i])), uintptr(env.stackTopAsUint64()))
+			require.Equal(t, uintptr(unsafe.Pointer(&me.functions[i])), uintptr(env.stackTopAsUint64()))
 		})
 	}
 }
diff --git a/internal/engine/compiler/compiler_test.go b/internal/engine/compiler/compiler_test.go
index c9f906b5..7c74190a 100644
--- a/internal/engine/compiler/compiler_test.go
+++ b/internal/engine/compiler/compiler_test.go
@@ -68,6 +68,7 @@ func init() {
 	requireEqual(int(unsafe.Offsetof(compiledFunc.codeInitialAddress)), functionCodeInitialAddressOffset, "functionCodeInitialAddressOffset")
 	requireEqual(int(unsafe.Offsetof(compiledFunc.source)), functionSourceOffset, "functionSourceOffset")
 	requireEqual(int(unsafe.Offsetof(compiledFunc.moduleInstanceAddress)), functionModuleInstanceAddressOffset, "functionModuleInstanceAddressOffset")
+	requireEqual(int(unsafe.Sizeof(compiledFunc)), functionSize, "functionModuleInstanceAddressOffset")
 
 	// Offsets for wasm.ModuleInstance.
 	var moduleInstance wasm.ModuleInstance
diff --git a/internal/engine/compiler/engine_test.go b/internal/engine/compiler/engine_test.go
index f448010c..1fa91c83 100644
--- a/internal/engine/compiler/engine_test.go
+++ b/internal/engine/compiler/engine_test.go
@@ -49,7 +49,7 @@ func (e *engineTester) NewEngine(enabledFeatures api.CoreFeatures) wasm.Engine {
 // CompiledFunctionPointerValue implements the same method as documented on enginetest.EngineTester.
 func (e engineTester) CompiledFunctionPointerValue(me wasm.ModuleEngine, funcIndex wasm.Index) uint64 {
 	internal := me.(*moduleEngine)
-	return uint64(uintptr(unsafe.Pointer(internal.functions[funcIndex])))
+	return uint64(uintptr(unsafe.Pointer(&internal.functions[funcIndex])))
 }
 
 func TestCompiler_Engine_NewModuleEngine(t *testing.T) {
diff --git a/internal/engine/compiler/impl_amd64_test.go b/internal/engine/compiler/impl_amd64_test.go
index 71b1e58f..cbf20ecd 100644
--- a/internal/engine/compiler/impl_amd64_test.go
+++ b/internal/engine/compiler/impl_amd64_test.go
@@ -22,7 +22,7 @@ func TestAmd64Compiler_indirectCallWithTargetOnCallingConvReg(t *testing.T) {
 	// and the typeID  matches the table[targetOffset]'s type ID.
 	operation := &wazeroir.OperationCallIndirect{TypeIndex: 0}
 	env.module().TypeIDs = []wasm.FunctionTypeID{0}
-	env.module().Engine = &moduleEngine{functions: []*function{}}
+	env.module().Engine = &moduleEngine{functions: []function{}}
 
 	me := env.moduleEngine()
 	{ // Compiling call target.
@@ -35,14 +35,14 @@ func TestAmd64Compiler_indirectCallWithTargetOnCallingConvReg(t *testing.T) {
 		c, _, err := compiler.compile()
 		require.NoError(t, err)
 
-		f := &function{
+		f := function{
 			parent:                &code{codeSegment: c},
 			codeInitialAddress:    uintptr(unsafe.Pointer(&c[0])),
 			moduleInstanceAddress: uintptr(unsafe.Pointer(env.moduleInstance)),
 			source:                &wasm.FunctionInstance{TypeID: 0},
 		}
 		me.functions = append(me.functions, f)
-		table[0] = uintptr(unsafe.Pointer(f))
+		table[0] = uintptr(unsafe.Pointer(&f))
 	}
 
 	compiler := env.requireNewCompiler(t, newCompiler, &wazeroir.CompilationResult{
diff --git a/internal/engine/compiler/impl_arm64_test.go b/internal/engine/compiler/impl_arm64_test.go
index 130f7c3e..b3181b96 100644
--- a/internal/engine/compiler/impl_arm64_test.go
+++ b/internal/engine/compiler/impl_arm64_test.go
@@ -21,7 +21,7 @@ func TestArm64Compiler_indirectCallWithTargetOnCallingConvReg(t *testing.T) {
 	// and the typeID  matches the table[targetOffset]'s type ID.
 	operation := &wazeroir.OperationCallIndirect{TypeIndex: 0}
 	env.module().TypeIDs = []wasm.FunctionTypeID{0}
-	env.module().Engine = &moduleEngine{functions: []*function{}}
+	env.module().Engine = &moduleEngine{functions: []function{}}
 
 	me := env.moduleEngine()
 	{ // Compiling call target.
@@ -34,14 +34,14 @@ func TestArm64Compiler_indirectCallWithTargetOnCallingConvReg(t *testing.T) {
 		c, _, err := compiler.compile()
 		require.NoError(t, err)
 
-		f := &function{
+		f := function{
 			parent:                &code{codeSegment: c},
 			codeInitialAddress:    uintptr(unsafe.Pointer(&c[0])),
 			moduleInstanceAddress: uintptr(unsafe.Pointer(env.moduleInstance)),
 			source:                &wasm.FunctionInstance{TypeID: 0},
 		}
 		me.functions = append(me.functions, f)
-		table[0] = uintptr(unsafe.Pointer(f))
+		table[0] = uintptr(unsafe.Pointer(&f))
 	}
 
 	compiler := env.requireNewCompiler(t, newCompiler, &wazeroir.CompilationResult{