[wasm] Cache import wrappers in NativeModule

Now that import wrappers are no longer specialized to an index, they
can be cached in the native module, keyed by
(WasmImportCallKind, FunctionSig). This saves instantiation time and
also fixes a (slow) memory leak.

R=mstarzinger@chromium.org

Change-Id: I5197bbfae79d6e811a01289b990db445373eea6c
Reviewed-on: https://chromium-review.googlesource.com/c/1270943
Commit-Queue: Ben Titzer <titzer@chromium.org>
Reviewed-by: Michael Starzinger <mstarzinger@chromium.org>
Cr-Commit-Position: refs/heads/master@{#56526}

BUILD.gn

@@ -2589,6 +2589,7 @@ v8_source_set("v8_base") {
     "src/wasm/wasm-feature-flags.h",
     "src/wasm/wasm-features.cc",
     "src/wasm/wasm-features.h",
+    "src/wasm/wasm-import-wrapper-cache-inl.h",
     "src/wasm/wasm-interpreter.cc",
     "src/wasm/wasm-interpreter.h",
     "src/wasm/wasm-js.cc",

src/compiler/wasm-compiler.cc

@@ -4951,9 +4951,10 @@ WasmImportCallKind GetWasmImportCallKind(Handle<JSReceiver> target,
   return WasmImportCallKind::kUseCallBuiltin;
 }
 
-MaybeHandle<Code> CompileWasmImportCallWrapper(
-    Isolate* isolate, WasmImportCallKind kind, wasm::FunctionSig* sig,
-    wasm::ModuleOrigin origin, wasm::UseTrapHandler use_trap_handler) {
+MaybeHandle<Code> CompileWasmImportCallWrapper(Isolate* isolate,
+                                               WasmImportCallKind kind,
+                                               wasm::FunctionSig* sig,
+                                               bool source_positions) {
   DCHECK_NE(WasmImportCallKind::kLinkError, kind);
   DCHECK_NE(WasmImportCallKind::kWasmToWasm, kind);
 
@@ -4975,10 +4976,9 @@ MaybeHandle<Code> CompileWasmImportCallWrapper(
   Node* effect = nullptr;
 
   SourcePositionTable* source_position_table =
-      origin == wasm::kAsmJsOrigin ? new (&zone) SourcePositionTable(&graph)
-                                   : nullptr;
+      source_positions ? new (&zone) SourcePositionTable(&graph) : nullptr;
 
-  wasm::ModuleEnv env(nullptr, use_trap_handler,
+  wasm::ModuleEnv env(nullptr, wasm::kNoTrapHandler,
                       wasm::kRuntimeExceptionSupport);
 
   WasmWrapperGraphBuilder builder(&zone, &env, &jsgraph, sig,

src/compiler/wasm-compiler.h

@@ -71,7 +71,7 @@ class TurbofanWasmCompilationUnit {
 // Calls to WASM imports are handled in several different ways, depending
 // on the type of the target function/callable and whether the signature
 // matches the argument arity.
-enum class WasmImportCallKind {
+enum class WasmImportCallKind : uint8_t {
   kLinkError,                      // static WASM->WASM type error
   kRuntimeTypeError,               // runtime WASM->JS type error
   kWasmToWasm,                     // fast WASM->WASM call
@@ -88,8 +88,7 @@ WasmImportCallKind GetWasmImportCallKind(Handle<JSReceiver> callable,
 // Compiles an import call wrapper, which allows WASM to call imports.
 MaybeHandle<Code> CompileWasmImportCallWrapper(Isolate*, WasmImportCallKind,
                                                wasm::FunctionSig*,
-                                               wasm::ModuleOrigin,
-                                               wasm::UseTrapHandler);
+                                               bool source_positions);
 
 // Creates a code object calling a wasm function with the given signature,
 // callable from JS.

src/wasm/module-compiler.cc

@@ -19,6 +19,7 @@
 #include "src/wasm/streaming-decoder.h"
 #include "src/wasm/wasm-code-manager.h"
 #include "src/wasm/wasm-engine.h"
+#include "src/wasm/wasm-import-wrapper-cache-inl.h"
 #include "src/wasm/wasm-js.h"
 #include "src/wasm/wasm-limits.h"
 #include "src/wasm/wasm-memory.h"
@@ -1516,14 +1517,9 @@ int InstanceBuilder::ProcessImports(Handle<WasmInstanceObject> instance) {
           }
           default: {
             // The imported function is a callable.
-            Handle<Code> wrapper_code = compiler::CompileWasmImportCallWrapper(
-                                            isolate_, kind, expected_sig,
-                                            module_->origin, use_trap_handler())
-                                            .ToHandleChecked();
-            RecordStats(*wrapper_code, isolate_->counters());
-
             WasmCode* wasm_code =
-                native_module->AddImportWrapper(wrapper_code, func_index);
+                native_module->import_wrapper_cache()->GetOrCompile(
+                    isolate_, kind, expected_sig);
             ImportedFunctionEntry entry(instance, func_index);
             entry.SetWasmToJs(isolate_, js_receiver, wasm_code);
             break;

src/wasm/wasm-code-manager.cc

@@ -18,6 +18,7 @@
 #include "src/objects-inl.h"
 #include "src/wasm/function-compiler.h"
 #include "src/wasm/jump-table-assembler.h"
+#include "src/wasm/wasm-import-wrapper-cache-inl.h"
 #include "src/wasm/wasm-module.h"
 #include "src/wasm/wasm-objects-inl.h"
 #include "src/wasm/wasm-objects.h"
@@ -342,6 +343,8 @@ NativeModule::NativeModule(Isolate* isolate, const WasmFeatures& enabled,
     : enabled_features_(enabled),
       module_(std::move(module)),
       compilation_state_(NewCompilationState(isolate, env)),
+      import_wrapper_cache_(std::unique_ptr<WasmImportWrapperCache>(
+          new WasmImportWrapperCache(this))),
       free_code_space_(code_space.region()),
       wasm_code_manager_(code_manager),
       can_request_more_memory_(can_request_more),
@@ -424,15 +427,6 @@ WasmCode* NativeModule::AddOwnedCode(
   return code;
 }
 
-WasmCode* NativeModule::AddImportWrapper(Handle<Code> code, uint32_t index) {
-  // TODO(wasm): Adding instance-specific import wrappers as owned code to
-  // this NativeModule is a memory leak until the whole NativeModule dies.
-  WasmCode* ret = AddAnonymousCode(code, WasmCode::kWasmToJsWrapper);
-  DCHECK_LT(index, module_->num_imported_functions);
-  ret->index_ = index;
-  return ret;
-}
-
 WasmCode* NativeModule::AddInterpreterEntry(Handle<Code> code, uint32_t index) {
   WasmCode* ret = AddAnonymousCode(code, WasmCode::kInterpreterEntry);
   ret->index_ = index;
@@ -463,6 +457,7 @@ void NativeModule::SetLazyBuiltin(Handle<Code> code) {
 }
 
 void NativeModule::SetRuntimeStubs(Isolate* isolate) {
+  HandleScope scope(isolate);
   DCHECK_NULL(runtime_stub_table_[0]);  // Only called once.
 #define COPY_BUILTIN(Name)                                                     \
   runtime_stub_table_[WasmCode::k##Name] =                                     \

src/wasm/wasm-code-manager.h

@@ -31,6 +31,7 @@ namespace wasm {
 class NativeModule;
 class WasmCodeManager;
 class WasmMemoryTracker;
+class WasmImportWrapperCache;
 struct WasmModule;
 
 // Sorted, disjoint and non-overlapping memory regions. A region is of the
@@ -230,14 +231,10 @@ class V8_EXPORT_PRIVATE NativeModule final {
       OwnedVector<const byte> reloc_info,
       OwnedVector<const byte> source_position_table, WasmCode::Tier tier);
 
-  // Add an import wrapper, e.g. for calls to JS functions. This method copies
-  // heap-allocated code because we compile wrappers using a different pipeline.
-  WasmCode* AddImportWrapper(Handle<Code> code, uint32_t index);
-
-  // Add an interpreter entry. For the same reason as AddImportWrapper, we
-  // currently compile these using a different pipeline and we can't get a
-  // CodeDesc here. When adding interpreter wrappers, we do not insert them in
-  // the code_table, however, we let them self-identify as the {index} function.
+  // Add an interpreter entry. We currently compile these using a different
+  // pipeline and we can't get a CodeDesc here. When adding interpreter
+  // wrappers, we do not insert them in the code_table, however, we let them
+  // self-identify as the {index} function.
   WasmCode* AddInterpreterEntry(Handle<Code> code, uint32_t index);
 
   // Adds anonymous code for testing purposes.
@@ -335,6 +332,10 @@ class V8_EXPORT_PRIVATE NativeModule final {
 
   WasmCode* Lookup(Address) const;
 
+  WasmImportWrapperCache* import_wrapper_cache() const {
+    return import_wrapper_cache_.get();
+  }
+
   ~NativeModule();
 
   const WasmFeatures& enabled_features() const { return enabled_features_; }
@@ -343,6 +344,7 @@ class V8_EXPORT_PRIVATE NativeModule final {
   friend class WasmCode;
   friend class WasmCodeManager;
   friend class NativeModuleModificationScope;
+  friend class WasmImportWrapperCache;
 
   NativeModule(Isolate* isolate, const WasmFeatures& enabled_features,
                bool can_request_more, VirtualMemory code_space,
@@ -397,6 +399,9 @@ class V8_EXPORT_PRIVATE NativeModule final {
   // hence needs to be destructed first when this native module dies.
   std::unique_ptr<CompilationState, CompilationStateDeleter> compilation_state_;
 
+  // A cache of the import wrappers, keyed on the kind and signature.
+  std::unique_ptr<WasmImportWrapperCache> import_wrapper_cache_;
+
   // This mutex protects concurrent calls to {AddCode} and friends.
   mutable base::Mutex allocation_mutex_;
 

src/wasm/wasm-import-wrapper-cache-inl.h

+// Copyright 2018 the V8 project authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef V8_WASM_WASM_IMPORT_WRAPPER_CACHE_INL_H_
+#define V8_WASM_WASM_IMPORT_WRAPPER_CACHE_INL_H_
+
+#include "src/compiler/wasm-compiler.h"
+#include "src/counters.h"
+#include "src/handles-inl.h"
+#include "src/objects/code-inl.h"
+#include "src/wasm/value-type.h"
+#include "src/wasm/wasm-code-manager.h"
+
+namespace v8 {
+namespace internal {
+namespace wasm {
+
+using FunctionSig = Signature<ValueType>;
+
+// Implements a cache for import wrappers.
+class WasmImportWrapperCache {
+ public:
+  WasmCode* GetOrCompile(Isolate* isolate, compiler::WasmImportCallKind kind,
+                         FunctionSig* sig) {
+    // TODO(titzer): remove the isolate parameter.
+    base::LockGuard<base::Mutex> lock(&mutex_);
+    CacheKey key(static_cast<uint8_t>(kind), *sig);
+    WasmCode*& cached = entry_map_[key];
+    if (cached == nullptr) {
+      // TODO(wasm): no need to hold the lock while compiling an import wrapper.
+      // TODO(wasm): no need to compile the code onto the heap and copy back.
+      HandleScope scope(isolate);
+      bool source_positions = native_module_->module()->origin == kAsmJsOrigin;
+      Handle<Code> code = compiler::CompileWasmImportCallWrapper(
+                              isolate, kind, sig, source_positions)
+                              .ToHandleChecked();
+      auto counters = isolate->counters();
+      counters->wasm_generated_code_size()->Increment(code->body_size());
+      counters->wasm_reloc_size()->Increment(code->relocation_info()->length());
+      cached =
+          native_module_->AddAnonymousCode(code, WasmCode::kWasmToJsWrapper);
+    }
+    return cached;
+  }
+
+ private:
+  friend class NativeModule;
+  mutable base::Mutex mutex_;
+  NativeModule* native_module_;
+  using CacheKey = std::pair<uint8_t, FunctionSig>;
+  std::unordered_map<CacheKey, WasmCode*, base::hash<CacheKey>> entry_map_;
+
+  explicit WasmImportWrapperCache(NativeModule* native_module)
+      : native_module_(native_module) {}
+};
+
+}  // namespace wasm
+}  // namespace internal
+}  // namespace v8
+
+#endif  // V8_WASM_WASM_IMPORT_WRAPPER_CACHE_INL_H_

test/cctest/BUILD.gn

@@ -255,6 +255,7 @@ v8_source_set("cctest_sources") {
     "wasm/test-streaming-compilation.cc",
     "wasm/test-wasm-breakpoints.cc",
     "wasm/test-wasm-codegen.cc",
+    "wasm/test-wasm-import-wrapper-cache.cc",
     "wasm/test-wasm-interpreter-entry.cc",
     "wasm/test-wasm-serialization.cc",
     "wasm/test-wasm-shared-engine.cc",

test/cctest/wasm/test-wasm-import-wrapper-cache.cc

+// Copyright 2018 the V8 project authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "src/compiler/wasm-compiler.h"
+#include "src/wasm/function-compiler.h"
+#include "src/wasm/wasm-code-manager.h"
+#include "src/wasm/wasm-engine.h"
+#include "src/wasm/wasm-import-wrapper-cache-inl.h"
+#include "src/wasm/wasm-module.h"
+
+#include "test/cctest/cctest.h"
+#include "test/common/wasm/test-signatures.h"
+
+namespace v8 {
+namespace internal {
+namespace wasm {
+namespace test_wasm_import_wrapper_cache {
+
+std::unique_ptr<NativeModule> NewModule(Isolate* isolate) {
+  WasmCodeManager* manager = isolate->wasm_engine()->code_manager();
+  std::shared_ptr<WasmModule> module(new WasmModule);
+  bool can_request_more = false;
+  size_t size = 100;
+  ModuleEnv env(module.get(), UseTrapHandler::kNoTrapHandler,
+                RuntimeExceptionSupport::kNoRuntimeExceptionSupport);
+  auto native_module =
+      manager->NewNativeModule(isolate, kAllWasmFeatures, size,
+                               can_request_more, std::move(module), env);
+  native_module->SetRuntimeStubs(isolate);
+  return native_module;
+}
+
+TEST(CacheHit) {
+  Isolate* isolate = CcTest::InitIsolateOnce();
+  auto module = NewModule(isolate);
+  TestSignatures sigs;
+
+  auto kind = compiler::WasmImportCallKind::kJSFunctionArityMatch;
+
+  WasmCode* c1 =
+      module->import_wrapper_cache()->GetOrCompile(isolate, kind, sigs.i_i());
+
+  CHECK_NOT_NULL(c1);
+  CHECK_EQ(WasmCode::Kind::kWasmToJsWrapper, c1->kind());
+
+  WasmCode* c2 =
+      module->import_wrapper_cache()->GetOrCompile(isolate, kind, sigs.i_i());
+
+  CHECK_NOT_NULL(c2);
+  CHECK_EQ(c1, c2);
+}
+
+TEST(CacheMissSig) {
+  Isolate* isolate = CcTest::InitIsolateOnce();
+  auto module = NewModule(isolate);
+  TestSignatures sigs;
+
+  auto kind = compiler::WasmImportCallKind::kJSFunctionArityMatch;
+
+  WasmCode* c1 =
+      module->import_wrapper_cache()->GetOrCompile(isolate, kind, sigs.i_i());
+
+  CHECK_NOT_NULL(c1);
+  CHECK_EQ(WasmCode::Kind::kWasmToJsWrapper, c1->kind());
+
+  WasmCode* c2 =
+      module->import_wrapper_cache()->GetOrCompile(isolate, kind, sigs.i_ii());
+
+  CHECK_NOT_NULL(c2);
+  CHECK_NE(c1, c2);
+}
+
+TEST(CacheMissKind) {
+  Isolate* isolate = CcTest::InitIsolateOnce();
+  auto module = NewModule(isolate);
+  TestSignatures sigs;
+
+  auto kind1 = compiler::WasmImportCallKind::kJSFunctionArityMatch;
+  auto kind2 = compiler::WasmImportCallKind::kJSFunctionArityMismatch;
+
+  WasmCode* c1 =
+      module->import_wrapper_cache()->GetOrCompile(isolate, kind1, sigs.i_i());
+
+  CHECK_NOT_NULL(c1);
+  CHECK_EQ(WasmCode::Kind::kWasmToJsWrapper, c1->kind());
+
+  WasmCode* c2 =
+      module->import_wrapper_cache()->GetOrCompile(isolate, kind2, sigs.i_i());
+
+  CHECK_NOT_NULL(c2);
+  CHECK_NE(c1, c2);
+}
+
+TEST(CacheHitMissSig) {
+  Isolate* isolate = CcTest::InitIsolateOnce();
+  auto module = NewModule(isolate);
+  TestSignatures sigs;
+
+  auto kind = compiler::WasmImportCallKind::kJSFunctionArityMatch;
+
+  WasmCode* c1 =
+      module->import_wrapper_cache()->GetOrCompile(isolate, kind, sigs.i_i());
+
+  CHECK_NOT_NULL(c1);
+  CHECK_EQ(WasmCode::Kind::kWasmToJsWrapper, c1->kind());
+
+  WasmCode* c2 =
+      module->import_wrapper_cache()->GetOrCompile(isolate, kind, sigs.i_ii());
+
+  CHECK_NOT_NULL(c2);
+  CHECK_NE(c1, c2);
+
+  WasmCode* c3 =
+      module->import_wrapper_cache()->GetOrCompile(isolate, kind, sigs.i_i());
+
+  CHECK_NOT_NULL(c3);
+  CHECK_EQ(c1, c3);
+
+  WasmCode* c4 =
+      module->import_wrapper_cache()->GetOrCompile(isolate, kind, sigs.i_ii());
+
+  CHECK_NOT_NULL(c4);
+  CHECK_EQ(c2, c4);
+}
+
+}  // namespace test_wasm_import_wrapper_cache
+}  // namespace wasm
+}  // namespace internal
+}  // namespace v8

test/cctest/wasm/wasm-run-utils.cc

@@ -6,6 +6,7 @@
 
 #include "src/assembler-inl.h"
 #include "src/code-tracer.h"
+#include "src/wasm/wasm-import-wrapper-cache-inl.h"
 #include "src/wasm/wasm-memory.h"
 #include "src/wasm/wasm-objects-inl.h"
 
@@ -43,12 +44,8 @@ TestingModuleBuilder::TestingModuleBuilder(
     CodeSpaceMemoryModificationScope modification_scope(isolate_->heap());
     auto kind = compiler::GetWasmImportCallKind(maybe_import->js_function,
                                                 maybe_import->sig);
-    MaybeHandle<Code> code = compiler::CompileWasmImportCallWrapper(
-        isolate_, kind, maybe_import->sig, test_module_->origin,
-        trap_handler::IsTrapHandlerEnabled() ? kUseTrapHandler
-                                             : kNoTrapHandler);
-    auto import_wrapper = native_module_->AddImportWrapper(
-        code.ToHandleChecked(), maybe_import_index);
+    auto import_wrapper = native_module_->import_wrapper_cache()->GetOrCompile(
+        isolate_, kind, maybe_import->sig);
 
     ImportedFunctionEntry(instance_object_, maybe_import_index)
         .SetWasmToJs(isolate_, maybe_import->js_function, import_wrapper);