Repository: tetratelabs/wazero
Pull request #897 (2022-12-07 04:22:15)

Removes ModuleInstance.TypeIDsIndex to reduce allocations
Ref #889 

This further reduces the allocation during instantiation:

```
$ benchstat old.txt new.txt                                                                                                                    
name                                    old time/op    new time/op    delta
Initialization/interpreter-10             23.7µs ±11%    22.7µs ± 8%    ~     (p=0.421 n=5+5)
Initialization/interpreter-multiple-10    10.8µs ± 7%    11.2µs ± 7%    ~     (p=0.548 n=5+5)
Initialization/compiler-10                33.1µs ± 2%    32.9µs ± 1%    ~     (p=0.905 n=5+4)
Initialization/compiler-multiple-10       15.0µs ± 9%    16.0µs ±16%    ~     (p=0.421 n=5+5)

name                                    old alloc/op   new alloc/op   delta
Initialization/interpreter-10              141kB ± 0%     141kB ± 0%    ~     (p=0.079 n=4+5)
Initialization/interpreter-multiple-10     143kB ± 0%     142kB ± 0%  -0.30%  (p=0.008 n=5+5)
Initialization/compiler-10                 145kB ± 0%     144kB ± 0%  -0.29%  (p=0.008 n=5+5)
Initialization/compiler-multiple-10        145kB ± 0%     145kB ± 0%  -0.29%  (p=0.016 n=5+4)

name                                    old allocs/op  new allocs/op  delta
Initialization/interpreter-10               85.0 ± 0%      83.0 ± 0%  -2.35%  (p=0.008 n=5+5)
Initialization/interpreter-multiple-10      92.0 ± 0%      90.0 ± 0%  -2.17%  (p=0.008 n=5+5)
Initialization/compiler-10                  35.0 ± 0%      33.0 ± 0%  -5.71%  (p=0.008 n=5+5)
Initialization/compiler-multiple-10         41.0 ± 0%      39.0 ± 0%  -4.88%  (p=0.008 n=5+5)


```


Patch:
diff --git a/imports/emscripten/emscripten.go b/imports/emscripten/emscripten.go
index 46980b37..b2eb6d21 100644
--- a/imports/emscripten/emscripten.go
+++ b/imports/emscripten/emscripten.go
@@ -18,7 +18,6 @@ import (
 	"github.com/tetratelabs/wazero"
 	"github.com/tetratelabs/wazero/api"
 	"github.com/tetratelabs/wazero/internal/wasm"
-	"github.com/tetratelabs/wazero/internal/wasmruntime"
 )
 
 // MustInstantiate calls Instantiate or panics on error.
@@ -148,7 +147,7 @@ var invokeI = &wasm.HostFunc{
 }
 
 func invokeIFn(ctx context.Context, mod api.Module, stack []uint64) {
-	ret, err := callDynamic(ctx, mod.(*wasm.CallContext), "v_i32", wasm.Index(stack[0]), nil)
+	ret, err := callDynamic(ctx, mod.(*wasm.CallContext), wasm.PreAllocatedTypeID_v_i32, wasm.Index(stack[0]), nil)
 	if err != nil {
 		panic(err)
 	}
@@ -168,7 +167,7 @@ var invokeIi = &wasm.HostFunc{
 }
 
 func invokeIiFn(ctx context.Context, mod api.Module, stack []uint64) {
-	ret, err := callDynamic(ctx, mod.(*wasm.CallContext), "i32_i32", wasm.Index(stack[0]), stack[1:])
+	ret, err := callDynamic(ctx, mod.(*wasm.CallContext), wasm.PreAllocatedTypeID_i32_i32, wasm.Index(stack[0]), stack[1:])
 	if err != nil {
 		panic(err)
 	}
@@ -188,7 +187,7 @@ var invokeIii = &wasm.HostFunc{
 }
 
 func invokeIiiFn(ctx context.Context, mod api.Module, stack []uint64) {
-	ret, err := callDynamic(ctx, mod.(*wasm.CallContext), "i32i32_i32", wasm.Index(stack[0]), stack[1:])
+	ret, err := callDynamic(ctx, mod.(*wasm.CallContext), wasm.PreAllocatedTypeID_i32i32_i32, wasm.Index(stack[0]), stack[1:])
 	if err != nil {
 		panic(err)
 	}
@@ -208,7 +207,7 @@ var invokeIiii = &wasm.HostFunc{
 }
 
 func invokeIiiiFn(ctx context.Context, mod api.Module, stack []uint64) {
-	ret, err := callDynamic(ctx, mod.(*wasm.CallContext), "i32i32i32_i32", wasm.Index(stack[0]), stack[1:])
+	ret, err := callDynamic(ctx, mod.(*wasm.CallContext), wasm.PreAllocatedTypeID_i32i32i32_i32, wasm.Index(stack[0]), stack[1:])
 	if err != nil {
 		panic(err)
 	}
@@ -228,7 +227,7 @@ var invokeIiiii = &wasm.HostFunc{
 }
 
 func invokeIiiiiFn(ctx context.Context, mod api.Module, stack []uint64) {
-	ret, err := callDynamic(ctx, mod.(*wasm.CallContext), "i32i32i32i32_i32", wasm.Index(stack[0]), stack[1:])
+	ret, err := callDynamic(ctx, mod.(*wasm.CallContext), wasm.PreAllocatedTypeID_i32i32i32i32_i32, wasm.Index(stack[0]), stack[1:])
 	if err != nil {
 		panic(err)
 	}
@@ -248,7 +247,7 @@ var invokeV = &wasm.HostFunc{
 }
 
 func invokeVFn(ctx context.Context, mod api.Module, stack []uint64) {
-	_, err := callDynamic(ctx, mod.(*wasm.CallContext), "v_v", wasm.Index(stack[0]), nil)
+	_, err := callDynamic(ctx, mod.(*wasm.CallContext), wasm.PreAllocatedTypeID_v_v, wasm.Index(stack[0]), nil)
 	if err != nil {
 		panic(err)
 	}
@@ -267,7 +266,7 @@ var invokeVi = &wasm.HostFunc{
 }
 
 func invokeViFn(ctx context.Context, mod api.Module, stack []uint64) {
-	_, err := callDynamic(ctx, mod.(*wasm.CallContext), "i32_v", wasm.Index(stack[0]), stack[1:])
+	_, err := callDynamic(ctx, mod.(*wasm.CallContext), wasm.PreAllocatedTypeID_i32_v, wasm.Index(stack[0]), stack[1:])
 	if err != nil {
 		panic(err)
 	}
@@ -286,7 +285,7 @@ var invokeVii = &wasm.HostFunc{
 }
 
 func invokeViiFn(ctx context.Context, mod api.Module, stack []uint64) {
-	_, err := callDynamic(ctx, mod.(*wasm.CallContext), "i32i32_v", wasm.Index(stack[0]), stack[1:])
+	_, err := callDynamic(ctx, mod.(*wasm.CallContext), wasm.PreAllocatedTypeID_i32i32_v, wasm.Index(stack[0]), stack[1:])
 	if err != nil {
 		panic(err)
 	}
@@ -305,7 +304,7 @@ var invokeViii = &wasm.HostFunc{
 }
 
 func invokeViiiFn(ctx context.Context, mod api.Module, stack []uint64) {
-	_, err := callDynamic(ctx, mod.(*wasm.CallContext), "i32i32i32_v", wasm.Index(stack[0]), stack[1:])
+	_, err := callDynamic(ctx, mod.(*wasm.CallContext), wasm.PreAllocatedTypeID_i32i32i32_v, wasm.Index(stack[0]), stack[1:])
 	if err != nil {
 		panic(err)
 	}
@@ -324,7 +323,7 @@ var invokeViiii = &wasm.HostFunc{
 }
 
 func invokeViiiiFn(ctx context.Context, mod api.Module, stack []uint64) {
-	_, err := callDynamic(ctx, mod.(*wasm.CallContext), "i32i32i32i32_v", wasm.Index(stack[0]), stack[1:])
+	_, err := callDynamic(ctx, mod.(*wasm.CallContext), wasm.PreAllocatedTypeID_i32i32i32i32_v, wasm.Index(stack[0]), stack[1:])
 	if err != nil {
 		panic(err)
 	}
@@ -337,18 +336,14 @@ func invokeViiiiFn(ctx context.Context, mod api.Module, stack []uint64) {
 //
 //   - ctx: the propagated go context.
 //   - callCtx: the incoming context of the `invoke_` function.
-//   - typeName: used to look up the function type. ex "i32i32_i32" or "v_i32"
+//   - typeID: used to type check on indirect calls.
 //   - tableOffset: position in the module's only table
 //   - params: parameters to the funcref
-func callDynamic(ctx context.Context, callCtx *wasm.CallContext, typeName string, tableOffset wasm.Index, params []uint64) (results []uint64, err error) {
+func callDynamic(ctx context.Context, callCtx *wasm.CallContext, typeID wasm.FunctionTypeID, tableOffset wasm.Index, params []uint64) (results []uint64, err error) {
 	m := callCtx.Module()
-	typeId, ok := m.TypeIDIndex[typeName]
-	if !ok {
-		return nil, wasmruntime.ErrRuntimeIndirectCallTypeMismatch
-	}
 
 	t := m.Tables[0] // Emscripten doesn't use multiple tables
-	idx, err := m.Engine.LookupFunction(t, typeId, tableOffset)
+	idx, err := m.Engine.LookupFunction(t, typeID, tableOffset)
 	if err != nil {
 		return nil, err
 	}
diff --git a/internal/engine/compiler/engine.go b/internal/engine/compiler/engine.go
index f17ec58b..e22fcac4 100644
--- a/internal/engine/compiler/engine.go
+++ b/internal/engine/compiler/engine.go
@@ -358,8 +358,8 @@ const (
 	moduleInstanceTablesOffset           = 80
 	moduleInstanceEngineOffset           = 112
 	moduleInstanceTypeIDsOffset          = 128
-	moduleInstanceDataInstancesOffset    = 160
-	moduleInstanceElementInstancesOffset = 184
+	moduleInstanceDataInstancesOffset    = 152
+	moduleInstanceElementInstancesOffset = 176
 
 	// Offsets for wasm.TableInstance.
 	tableInstanceTableOffset    = 0
diff --git a/internal/wasm/store.go b/internal/wasm/store.go
index 60b91d90..c779ec20 100644
--- a/internal/wasm/store.go
+++ b/internal/wasm/store.go
@@ -69,8 +69,7 @@ type (
 
 		// TypeIDs is index-correlated with types and holds typeIDs which is uniquely assigned to a type by store.
 		// This is necessary to achieve fast runtime type checking for indirect function calls at runtime.
-		TypeIDs     []FunctionTypeID
-		TypeIDIndex map[string]FunctionTypeID
+		TypeIDs []FunctionTypeID
 
 		// DataInstances holds data segments bytes of the module.
 		// This is only used by bulk memory operations.
@@ -157,13 +156,7 @@ type (
 const maximumFunctionTypes = 1 << 27
 
 // addSections adds section elements to the ModuleInstance
-func (m *ModuleInstance) addSections(module *Module, importedGlobals, globals []*GlobalInstance, tables []*TableInstance, memory, importedMemory *MemoryInstance,
-	types []*FunctionType,
-) {
-	m.TypeIDIndex = make(map[string]FunctionTypeID, len(types))
-	for i, t := range types {
-		m.TypeIDIndex[t.string] = m.TypeIDs[i]
-	}
+func (m *ModuleInstance) addSections(module *Module, importedGlobals, globals []*GlobalInstance, tables []*TableInstance, memory, importedMemory *MemoryInstance) {
 	m.Globals = append(importedGlobals, globals...)
 	m.Tables = tables
 
@@ -272,11 +265,16 @@ func (m *ModuleInstance) getExport(name string, et ExternType) (ExportInstance,
 
 func NewStore(enabledFeatures api.CoreFeatures, engine Engine) (*Store, *Namespace) {
 	ns := newNamespace()
+
+	typeIDs := make(map[string]FunctionTypeID, len(preAllocatedTypeIDs))
+	for k, v := range preAllocatedTypeIDs {
+		typeIDs[k] = v
+	}
 	return &Store{
 		EnabledFeatures:  enabledFeatures,
 		Engine:           engine,
 		namespaces:       []*Namespace{ns},
-		typeIDs:          map[string]FunctionTypeID{},
+		typeIDs:          typeIDs,
 		functionMaxTypes: maximumFunctionTypes,
 	}, ns
 }
@@ -371,7 +369,7 @@ func (s *Store) instantiate(
 	globals, memory := module.buildGlobals(importedGlobals, m.Engine.FunctionInstanceReference), module.buildMemory()
 
 	// Now we have all instances from imports and local ones, so ready to create a new ModuleInstance.
-	m.addSections(module, importedGlobals, globals, tables, importedMemory, memory, module.TypeSection)
+	m.addSections(module, importedGlobals, globals, tables, importedMemory, memory)
 
 	// As of reference types proposal, data segment validation must happen after instantiation,
 	// and the side effect must persist even if there's out of bounds error after instantiation.
@@ -590,6 +588,45 @@ func (s *Store) getFunctionTypeIDs(ts []*FunctionType) ([]FunctionTypeID, error)
 	return ret, nil
 }
 
+// preAllocatedTypeIDs maps several "well-known" FunctionType strings to the pre allocated FunctionID.
+// This is used by emscripten integration, but it is harmless to have this all the time as it's only
+// used during Store creation.
+var preAllocatedTypeIDs = map[string]FunctionTypeID{
+	"i32i32i32i32_v":   PreAllocatedTypeID_i32i32i32i32_v,
+	"i32i32i32_v":      PreAllocatedTypeID_i32i32i32_v,
+	"i32i32_v":         PreAllocatedTypeID_i32i32_v,
+	"i32_v":            PreAllocatedTypeID_i32_v,
+	"v_v":              PreAllocatedTypeID_v_v,
+	"i32i32i32i32_i32": PreAllocatedTypeID_i32i32i32i32_i32,
+	"i32i32i32_i32":    PreAllocatedTypeID_i32i32i32_i32,
+	"i32i32_i32":       PreAllocatedTypeID_i32i32_i32,
+	"i32_i32":          PreAllocatedTypeID_i32_i32,
+	"v_i32":            PreAllocatedTypeID_v_i32,
+}
+
+const (
+	// PreAllocatedTypeID_i32i32i32i32_v is FunctionTypeID for i32i32i32i32_v.
+	PreAllocatedTypeID_i32i32i32i32_v FunctionTypeID = iota
+	// PreAllocatedTypeID_i32i32i32_v is FunctionTypeID for i32i32i32_v
+	PreAllocatedTypeID_i32i32i32_v
+	// PreAllocatedTypeID_i32i32_v is FunctionTypeID for i32i32_v
+	PreAllocatedTypeID_i32i32_v
+	// PreAllocatedTypeID_i32_v is FunctionTypeID for i32_v
+	PreAllocatedTypeID_i32_v
+	// PreAllocatedTypeID_v_v is FunctionTypeID for v_v
+	PreAllocatedTypeID_v_v
+	// PreAllocatedTypeID_i32i32i32i32_i32 is FunctionTypeID for i32i32i32i32_i32
+	PreAllocatedTypeID_i32i32i32i32_i32
+	// PreAllocatedTypeID_i32i32i32_i32 is FunctionTypeID for i32i32i32_i32
+	PreAllocatedTypeID_i32i32i32_i32
+	// PreAllocatedTypeID_i32i32_i32 is FunctionTypeID for i32i32_i32
+	PreAllocatedTypeID_i32i32_i32
+	// PreAllocatedTypeID_i32_i32 is FunctionTypeID for i32_i32
+	PreAllocatedTypeID_i32_i32
+	// PreAllocatedTypeID_v_i32 is FunctionTypeID for v_i32
+	PreAllocatedTypeID_v_i32
+)
+
 func (s *Store) getFunctionTypeID(t *FunctionType) (FunctionTypeID, error) {
 	key := t.key()
 	s.mux.RLock()


Test patch:
diff --git a/internal/wasm/store_test.go b/internal/wasm/store_test.go
index dfa4ef97..c6849da9 100644
--- a/internal/wasm/store_test.go
+++ b/internal/wasm/store_test.go
@@ -90,6 +90,16 @@ func TestModuleInstance_Memory(t *testing.T) {
 	}
 }
 
+func TestNewStore(t *testing.T) {
+	s, _ := NewStore(api.CoreFeaturesV1, &mockEngine{shouldCompileFail: false, callFailIndex: -1})
+	// Ensures that a newly created store has the pre allocated type IDs.
+	for k, v := range preAllocatedTypeIDs {
+		actual, ok := s.typeIDs[k]
+		require.True(t, ok)
+		require.Equal(t, v, actual)
+	}
+}
+
 func TestStore_Instantiate(t *testing.T) {
 	s, ns := newStore()
 	m, err := NewHostModule("", map[string]interface{}{"fn": func() {}}, map[string]*HostFuncNames{"fn": {}}, api.CoreFeaturesV1)
@@ -910,3 +920,13 @@ func TestModuleInstance_applyTableInits(t *testing.T) {
 			tables[0].References)
 	})
 }
+
+// TestPreAllocatedTypeIDs ensures that PreAllocatedTypeIDs has no duplication on the values (FunctionTypeID).
+func TestPreAllocatedTypeIDs(t *testing.T) {
+	exists := make(map[FunctionTypeID]struct{}, len(preAllocatedTypeIDs))
+	for _, v := range preAllocatedTypeIDs {
+		_, ok := exists[v]
+		require.False(t, ok)
+		exists[v] = struct{}{}
+	}
+}