Repository: tetratelabs/wazero
Pull request #1145 (2023-02-21 00:42:52)

wazeroir: removes allocations in wasmOpcodeSignature asymptotically
```
name         old time/op    new time/op    delta
_compile-10     897ms ± 2%     891ms ± 1%  -0.66%  (p=0.000 n=43+49)

name         old alloc/op   new alloc/op   delta
_compile-10     599MB ± 0%     597MB ± 0%  -0.47%  (p=0.000 n=50+50)

name         old allocs/op  new allocs/op  delta
_compile-10     7.53M ± 0%     7.40M ± 0%  -1.70%  (p=0.000 n=50+49)
```


Patch:
diff --git a/internal/wazeroir/compiler.go b/internal/wazeroir/compiler.go
index 96e4fe30..7172bc92 100644
--- a/internal/wazeroir/compiler.go
+++ b/internal/wazeroir/compiler.go
@@ -188,7 +188,8 @@ type compiler struct {
 
 	ensureTermination bool
 	// Pre-allocated bytes.Reader to be used in varous places.
-	br *bytes.Reader
+	br             *bytes.Reader
+	funcTypeToSigs *funcTypeToIRSignatures
 }
 
 //lint:ignore U1000 for debugging only.
@@ -276,10 +277,18 @@ func CompileFunctions(enabledFeatures api.CoreFeatures, callFrameStackSizeInUint
 		tableTypes[i] = tables[i].Type
 	}
 
+	types := module.TypeSection
+
+	funcTypeToSigs := &funcTypeToIRSignatures{
+		indirectCalls: make([]*signature, len(types)),
+		directCalls:   make([]*signature, len(types)),
+		wasmTypes:     types,
+	}
+
 	var ret []*CompilationResult
 	for funcIndex := range module.FunctionSection {
 		typeID := module.FunctionSection[funcIndex]
-		sig := module.TypeSection[typeID]
+		sig := types[typeID]
 		code := module.CodeSection[funcIndex]
 		if code.GoFunc != nil {
 			// Assume the function might use memory if it has a parameter for the api.Module
@@ -300,8 +309,8 @@ func CompileFunctions(enabledFeatures api.CoreFeatures, callFrameStackSizeInUint
 			continue
 		}
 		r, err := compile(enabledFeatures, callFrameStackSizeInUint64, sig, code.Body,
-			code.LocalTypes, module.TypeSection, functions, globals, code.BodyOffsetInCodeSection,
-			module.DWARFLines != nil, ensureTermination)
+			code.LocalTypes, types, functions, globals, code.BodyOffsetInCodeSection,
+			module.DWARFLines != nil, ensureTermination, funcTypeToSigs)
 		if err != nil {
 			def := module.FunctionDefinitionSection[uint32(funcIndex)+module.ImportFuncCount()]
 			return nil, fmt.Errorf("failed to lower func[%s] to wazeroir: %w", def.DebugName(), err)
@@ -309,7 +318,7 @@ func CompileFunctions(enabledFeatures api.CoreFeatures, callFrameStackSizeInUint
 		r.IsHostFunction = code.IsHostFunction
 		r.Globals = globals
 		r.Functions = functions
-		r.Types = module.TypeSection
+		r.Types = types
 		r.HasMemory = hasMemory
 		r.HasTable = hasTable
 		r.HasDataInstances = hasDataInstances
@@ -335,6 +344,7 @@ func compile(enabledFeatures api.CoreFeatures,
 	bodyOffsetInCodeSection uint64,
 	needSourceOffset bool,
 	ensureTermination bool,
+	funcTypeToSigs *funcTypeToIRSignatures,
 ) (*CompilationResult, error) {
 	c := compiler{
 		enabledFeatures:            enabledFeatures,
@@ -351,6 +361,7 @@ func compile(enabledFeatures api.CoreFeatures,
 		bodyOffsetInCodeSection:    bodyOffsetInCodeSection,
 		ensureTermination:          ensureTermination,
 		br:                         bytes.NewReader(nil),
+		funcTypeToSigs:             funcTypeToSigs,
 	}
 
 	c.initializeStack()
diff --git a/internal/wazeroir/signature.go b/internal/wazeroir/signature.go
index 0b80ffa6..0437dbb6 100644
--- a/internal/wazeroir/signature.go
+++ b/internal/wazeroir/signature.go
@@ -253,11 +253,9 @@ func (c *compiler) wasmOpcodeSignature(op wasm.Opcode, index uint32) (*signature
 	case wasm.OpcodeReturn:
 		return signature_None_None, nil
 	case wasm.OpcodeCall:
-		return funcTypeToSignature(c.types[c.funcs[index]]), nil
+		return c.funcTypeToSigs.get(c.funcs[index], false /* direct */), nil
 	case wasm.OpcodeCallIndirect:
-		ret := funcTypeToSignature(c.types[index])
-		ret.in = append(ret.in, UnsignedTypeI32)
-		return ret, nil
+		return c.funcTypeToSigs.get(index, true /* call_indirect */), nil
 	case wasm.OpcodeDrop:
 		return signature_Unknown_None, nil
 	case wasm.OpcodeSelect, wasm.OpcodeTypedSelect:
@@ -593,15 +591,53 @@ func (c *compiler) wasmOpcodeSignature(op wasm.Opcode, index uint32) (*signature
 	}
 }
 
-func funcTypeToSignature(tps *wasm.FunctionType) *signature {
-	ret := &signature{}
-	for _, vt := range tps.Params {
-		ret.in = append(ret.in, wasmValueTypeToUnsignedType(vt))
+// funcTypeToIRSignatures is the central cache for a module to get the *signature
+// for function calls.
+type funcTypeToIRSignatures struct {
+	directCalls   []*signature
+	indirectCalls []*signature
+	wasmTypes     []*wasm.FunctionType
+}
+
+// get returns the *signature for the direct or indirect function call against functions whose type is at `typeIndex`.
+func (f *funcTypeToIRSignatures) get(typeIndex wasm.Index, indirect bool) *signature {
+	var sig *signature
+	if indirect {
+		sig = f.indirectCalls[typeIndex]
+	} else {
+		sig = f.directCalls[typeIndex]
+	}
+	if sig != nil {
+		return sig
 	}
-	for _, vt := range tps.Results {
-		ret.out = append(ret.out, wasmValueTypeToUnsignedType(vt))
+
+	tp := f.wasmTypes[typeIndex]
+	if indirect {
+		sig = &signature{
+			in:  make([]UnsignedType, 0, len(tp.Params)+1), // +1 to reserve space for call indirect index.
+			out: make([]UnsignedType, 0, len(tp.Results)),
+		}
+	} else {
+		sig = &signature{
+			in:  make([]UnsignedType, 0, len(tp.Params)),
+			out: make([]UnsignedType, 0, len(tp.Results)),
+		}
+	}
+
+	for _, vt := range tp.Params {
+		sig.in = append(sig.in, wasmValueTypeToUnsignedType(vt))
+	}
+	for _, vt := range tp.Results {
+		sig.out = append(sig.out, wasmValueTypeToUnsignedType(vt))
+	}
+
+	if indirect {
+		sig.in = append(sig.in, UnsignedTypeI32)
+		f.indirectCalls[typeIndex] = sig
+	} else {
+		f.directCalls[typeIndex] = sig
 	}
-	return ret
+	return sig
 }
 
 func wasmValueTypeToUnsignedType(vt wasm.ValueType) UnsignedType {


Test patch:
diff --git a/internal/wazeroir/signature_test.go b/internal/wazeroir/signature_test.go
index 819658ab..dbf4c9a3 100644
--- a/internal/wazeroir/signature_test.go
+++ b/internal/wazeroir/signature_test.go
@@ -100,3 +100,24 @@ func TestCompiler_wasmOpcodeSignature(t *testing.T) {
 		})
 	}
 }
+
+func Test_funcTypeToIRSignatures(t *testing.T) {
+	f := &funcTypeToIRSignatures{
+		wasmTypes:     []*wasm.FunctionType{v_v, i32_i32, v_f64f64},
+		directCalls:   make([]*signature, 3),
+		indirectCalls: make([]*signature, 3),
+	}
+
+	require.Equal(t, &signature{in: make([]UnsignedType, 0), out: make([]UnsignedType, 0)}, f.get(0, false))
+	require.Equal(t, &signature{in: []UnsignedType{UnsignedTypeI32}, out: make([]UnsignedType, 0)}, f.get(0, true))
+	require.NotNil(t, f.directCalls[0])
+	require.NotNil(t, f.indirectCalls[0])
+	require.Equal(t, &signature{in: []UnsignedType{UnsignedTypeI32}, out: []UnsignedType{UnsignedTypeI32}}, f.get(1, false))
+	require.Equal(t, &signature{in: []UnsignedType{UnsignedTypeI32, UnsignedTypeI32}, out: []UnsignedType{UnsignedTypeI32}}, f.get(1, true))
+	require.NotNil(t, f.directCalls[1])
+	require.NotNil(t, f.indirectCalls[1])
+	require.Equal(t, &signature{in: make([]UnsignedType, 0), out: []UnsignedType{UnsignedTypeF64, UnsignedTypeF64}}, f.get(2, false))
+	require.Equal(t, &signature{in: []UnsignedType{UnsignedTypeI32}, out: []UnsignedType{UnsignedTypeF64, UnsignedTypeF64}}, f.get(2, true))
+	require.NotNil(t, f.directCalls[2])
+	require.NotNil(t, f.indirectCalls[2])
+}