[wasm] Compile import wrappers in parallel

R=ahaas@chromium.org
CC=​titzer@chromium.org

Bug: v8:9231
Cq-Include-Trybots: luci.v8.try:v8_linux64_tsan_rel
Change-Id: I3de9c839ad43ab37c69b622ccf221dfc429c2e2d
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/1605732
Commit-Queue: Andreas Haas <ahaas@chromium.org>
Reviewed-by: Clemens Hammacher <clemensh@chromium.org>
Reviewed-by: Andreas Haas <ahaas@chromium.org>
Cr-Commit-Position: refs/heads/master@{#61839}

src/wasm/module-compiler.cc

@@ -1958,7 +1958,6 @@ bool AsyncStreamingProcessor::Deserialize(Vector<const uint8_t> module_bytes,
   return true;
 }
 
-namespace {
 int GetMaxBackgroundTasks() {
   if (NeedsDeterministicCompile()) return 1;
   int num_worker_threads = V8::GetCurrentPlatform()->NumberOfWorkerThreads();
@@ -1966,7 +1965,6 @@ int GetMaxBackgroundTasks() {
       std::min(FLAG_wasm_num_compilation_tasks, num_worker_threads);
   return std::max(1, num_compile_tasks);
 }
-}  // namespace
 
 CompilationStateImpl::CompilationStateImpl(
     const std::shared_ptr<NativeModule>& native_module,
@@ -2273,6 +2271,28 @@ void CompileJsToWasmWrappers(Isolate* isolate, const WasmModule* module,
   }
 }
 
+WasmCode* CompileImportWrapper(
+    WasmEngine* wasm_engine, NativeModule* native_module, Counters* counters,
+    compiler::WasmImportCallKind kind, FunctionSig* sig,
+    WasmImportWrapperCache::ModificationScope* cache_scope) {
+  // Entry should exist, so that we don't insert a new one and invalidate
+  // other threads' iterators/references, but it should not have been compiled
+  // yet.
+  WasmImportWrapperCache::CacheKey key(kind, sig);
+  DCHECK_NULL((*cache_scope)[key]);
+  bool source_positions = native_module->module()->origin == kAsmJsOrigin;
+  // Keep the {WasmCode} alive until we explicitly call {IncRef}.
+  WasmCodeRefScope code_ref_scope;
+  WasmCode* wasm_code = compiler::CompileWasmImportCallWrapper(
+      wasm_engine, native_module, kind, sig, source_positions);
+  (*cache_scope)[key] = wasm_code;
+  wasm_code->IncRef();
+  counters->wasm_generated_code_size()->Increment(
+      wasm_code->instructions().length());
+  counters->wasm_reloc_size()->Increment(wasm_code->reloc_info().length());
+  return wasm_code;
+}
+
 Handle<Script> CreateWasmScript(Isolate* isolate,
                                 const ModuleWireBytes& wire_bytes,
                                 const std::string& source_map_url) {

src/wasm/module-compiler.h

@@ -13,6 +13,7 @@
 #include "src/tasks/cancelable-task.h"
 #include "src/wasm/compilation-environment.h"
 #include "src/wasm/wasm-features.h"
+#include "src/wasm/wasm-import-wrapper-cache.h"
 #include "src/wasm/wasm-module.h"
 
 namespace v8 {
@@ -46,6 +47,15 @@ V8_EXPORT_PRIVATE
 void CompileJsToWasmWrappers(Isolate* isolate, const WasmModule* module,
                              Handle<FixedArray> export_wrappers);
 
+// Compiles the wrapper for this (kind, sig) pair and sets the corresponding
+// cache entry. Assumes the key already exists in the cache but has not been
+// compiled yet.
+V8_EXPORT_PRIVATE
+WasmCode* CompileImportWrapper(
+    WasmEngine* wasm_engine, NativeModule* native_module, Counters* counters,
+    compiler::WasmImportCallKind kind, FunctionSig* sig,
+    WasmImportWrapperCache::ModificationScope* cache_scope);
+
 V8_EXPORT_PRIVATE Handle<Script> CreateWasmScript(
     Isolate* isolate, const ModuleWireBytes& wire_bytes,
     const std::string& source_map_url);
@@ -55,6 +65,8 @@ V8_EXPORT_PRIVATE Handle<Script> CreateWasmScript(
 // also lazy.
 bool CompileLazy(Isolate*, NativeModule*, int func_index);
 
+int GetMaxBackgroundTasks();
+
 // Encapsulates all the state and steps of an asynchronous compilation.
 // An asynchronous compile job consists of a number of tasks that are executed
 // as foreground and background tasks. Any phase that touches the V8 heap or

src/wasm/module-instantiate.cc

@@ -47,6 +47,67 @@ uint32_t EvalUint32InitExpr(Handle<WasmInstanceObject> instance,
       UNREACHABLE();
   }
 }
+
+// Queue of import wrapper keys to compile for an instance.
+class ImportWrapperQueue {
+ public:
+  // Removes an arbitrary cache key from the queue and returns it.
+  // If the queue is empty, returns nullopt.
+  // Thread-safe.
+  base::Optional<WasmImportWrapperCache::CacheKey> pop() {
+    base::Optional<WasmImportWrapperCache::CacheKey> key = base::nullopt;
+    base::LockGuard<base::Mutex> lock(&mutex_);
+    auto it = queue_.begin();
+    if (it != queue_.end()) {
+      key = *it;
+      queue_.erase(it);
+    }
+    return key;
+  }
+
+  // Add the given key to the queue.
+  // Not thread-safe.
+  void insert(const WasmImportWrapperCache::CacheKey& key) {
+    queue_.insert(key);
+  }
+
+ private:
+  base::Mutex mutex_;
+  std::unordered_set<WasmImportWrapperCache::CacheKey,
+                     WasmImportWrapperCache::CacheKeyHash>
+      queue_;
+};
+
+class CompileImportWrapperTask final : public CancelableTask {
+ public:
+  CompileImportWrapperTask(
+      CancelableTaskManager* task_manager, WasmEngine* engine,
+      Counters* counters, NativeModule* native_module,
+      ImportWrapperQueue* queue,
+      WasmImportWrapperCache::ModificationScope* cache_scope)
+      : CancelableTask(task_manager),
+        engine_(engine),
+        counters_(counters),
+        native_module_(native_module),
+        queue_(queue),
+        cache_scope_(cache_scope) {}
+
+  void RunInternal() override {
+    while (base::Optional<WasmImportWrapperCache::CacheKey> key =
+               queue_->pop()) {
+      CompileImportWrapper(engine_, native_module_, counters_, key->first,
+                           key->second, cache_scope_);
+    }
+  }
+
+ private:
+  WasmEngine* const engine_;
+  Counters* const counters_;
+  NativeModule* const native_module_;
+  ImportWrapperQueue* const queue_;
+  WasmImportWrapperCache::ModificationScope* const cache_scope_;
+};
+
 }  // namespace
 
 // A helper class to simplify instantiating a module from a module object.
@@ -169,6 +230,10 @@ class InstanceBuilder {
                                        const WasmGlobal& global,
                                        Handle<WasmGlobalObject> global_object);
 
+  // Compile import wrappers in parallel. The result goes into the native
+  // module's import_wrapper_cache.
+  void CompileImportWrappers(Handle<WasmInstanceObject> instance);
+
   // Process the imports, including functions, tables, globals, and memory, in
   // order, loading them from the {ffi_} object. Returns the number of imported
   // functions.
@@ -819,8 +884,9 @@ bool InstanceBuilder::ProcessImportedFunction(
     default: {
       // The imported function is a callable.
       NativeModule* native_module = instance->module_object().native_module();
-      WasmCode* wasm_code = native_module->import_wrapper_cache()->GetOrCompile(
-          isolate_->wasm_engine(), isolate_->counters(), kind, expected_sig);
+      WasmCode* wasm_code =
+          native_module->import_wrapper_cache()->Get(kind, expected_sig);
+      DCHECK_NOT_NULL(wasm_code);
       ImportedFunctionEntry entry(instance, func_index);
       if (wasm_code->kind() == WasmCode::kWasmToJsWrapper) {
         // Wasm to JS wrappers are treated specially in the import table.
@@ -1129,6 +1195,62 @@ bool InstanceBuilder::ProcessImportedGlobal(Handle<WasmInstanceObject> instance,
   return false;
 }
 
+void InstanceBuilder::CompileImportWrappers(
+    Handle<WasmInstanceObject> instance) {
+  int num_imports = static_cast<int>(module_->import_table.size());
+  NativeModule* native_module = instance->module_object().native_module();
+  WasmImportWrapperCache::ModificationScope cache_scope(
+      native_module->import_wrapper_cache());
+
+  // Compilation is done in two steps:
+  // 1) Insert nullptr entries in the cache for wrappers that need to be
+  // compiled. 2) Compile wrappers in background tasks using the
+  // ImportWrapperQueue. This way the cache won't invalidate other iterators
+  // when inserting a new WasmCode, since the key will already be there.
+  ImportWrapperQueue import_wrapper_queue;
+  for (int index = 0; index < num_imports; ++index) {
+    Handle<Object> value = sanitized_imports_[index].value;
+    if (module_->import_table[index].kind != kExternalFunction ||
+        !value->IsCallable()) {
+      continue;
+    }
+    auto js_receiver = Handle<JSReceiver>::cast(value);
+    uint32_t func_index = module_->import_table[index].index;
+    FunctionSig* sig = module_->functions[func_index].sig;
+    auto kind =
+        compiler::GetWasmImportCallKind(js_receiver, sig, enabled_.bigint);
+    if (kind == compiler::WasmImportCallKind::kWasmToWasm ||
+        kind == compiler::WasmImportCallKind::kLinkError ||
+        kind == compiler::WasmImportCallKind::kWasmToCapi) {
+      continue;
+    }
+    WasmImportWrapperCache::CacheKey key(kind, sig);
+    if (cache_scope[key] != nullptr) {
+      // Cache entry already exists, no need to compile it again.
+      continue;
+    }
+    import_wrapper_queue.insert(key);
+  }
+
+  CancelableTaskManager task_manager;
+  const int max_background_tasks = GetMaxBackgroundTasks();
+  for (int i = 0; i < max_background_tasks; ++i) {
+    auto task = base::make_unique<CompileImportWrapperTask>(
+        &task_manager, isolate_->wasm_engine(), isolate_->counters(),
+        native_module, &import_wrapper_queue, &cache_scope);
+    V8::GetCurrentPlatform()->CallOnWorkerThread(std::move(task));
+  }
+
+  // Also compile in the current thread, in case there are no worker threads.
+  while (base::Optional<WasmImportWrapperCache::CacheKey> key =
+             import_wrapper_queue.pop()) {
+    CompileImportWrapper(isolate_->wasm_engine(), native_module,
+                         isolate_->counters(), key->first, key->second,
+                         &cache_scope);
+  }
+  task_manager.CancelAndWait();
+}
+
 // Process the imports, including functions, tables, globals, and memory, in
 // order, loading them from the {ffi_} object. Returns the number of imported
 // functions.
@@ -1137,6 +1259,8 @@ int InstanceBuilder::ProcessImports(Handle<WasmInstanceObject> instance) {
   int num_imported_tables = 0;
 
   DCHECK_EQ(module_->import_table.size(), sanitized_imports_.size());
+
+  CompileImportWrappers(instance);
   int num_imports = static_cast<int>(module_->import_table.size());
   for (int index = 0; index < num_imports; ++index) {
     const WasmImport& import = module_->import_table[index];

src/wasm/wasm-code-manager.cc

@@ -617,7 +617,7 @@ NativeModule::NativeModule(WasmEngine* engine, const WasmFeatures& enabled,
       enabled_features_(enabled),
       module_(std::move(module)),
       import_wrapper_cache_(std::unique_ptr<WasmImportWrapperCache>(
-          new WasmImportWrapperCache(this))),
+          new WasmImportWrapperCache())),
       engine_(engine),
       use_trap_handler_(trap_handler::IsTrapHandlerEnabled() ? kUseTrapHandler
                                                              : kNoTrapHandler) {

src/wasm/wasm-import-wrapper-cache.cc

@@ -13,32 +13,27 @@ namespace v8 {
 namespace internal {
 namespace wasm {
 
+WasmCode*& WasmImportWrapperCache::ModificationScope::operator[](
+    const CacheKey& key) {
+  return cache_->entry_map_[key];
+}
+
+WasmCode* WasmImportWrapperCache::Get(compiler::WasmImportCallKind kind,
+                                      FunctionSig* sig) const {
+  auto it = entry_map_.find({kind, sig});
+  DCHECK(it != entry_map_.end());
+  return it->second;
+}
+
 WasmImportWrapperCache::~WasmImportWrapperCache() {
   std::vector<WasmCode*> ptrs;
   ptrs.reserve(entry_map_.size());
-  for (auto& e : entry_map_) ptrs.push_back(e.second);
-  WasmCode::DecrementRefCount(VectorOf(ptrs));
-}
-
-WasmCode* WasmImportWrapperCache::GetOrCompile(
-    WasmEngine* wasm_engine, Counters* counters,
-    compiler::WasmImportCallKind kind, FunctionSig* sig) {
-  base::MutexGuard 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.
-    bool source_positions = native_module_->module()->origin == kAsmJsOrigin;
-    // Keep the {WasmCode} alive until we explicitly call {IncRef}.
-    WasmCodeRefScope code_ref_scope;
-    cached = compiler::CompileWasmImportCallWrapper(
-        wasm_engine, native_module_, kind, sig, source_positions);
-    cached->IncRef();
-    counters->wasm_generated_code_size()->Increment(
-        cached->instructions().length());
-    counters->wasm_reloc_size()->Increment(cached->reloc_info().length());
+  for (auto& e : entry_map_) {
+    if (e.second) {
+      ptrs.push_back(e.second);
+    }
   }
-  return cached;
+  WasmCode::DecrementRefCount(VectorOf(ptrs));
 }
 
 }  // namespace wasm

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

@@ -23,23 +23,37 @@ using FunctionSig = Signature<ValueType>;
 // Implements a cache for import wrappers.
 class WasmImportWrapperCache {
  public:
-  ~WasmImportWrapperCache();
+  using CacheKey = std::pair<compiler::WasmImportCallKind, FunctionSig*>;
 
-  V8_EXPORT_PRIVATE WasmCode* GetOrCompile(WasmEngine* wasm_engine,
-                                           Counters* counters,
-                                           compiler::WasmImportCallKind kind,
-                                           FunctionSig* sig);
+  class CacheKeyHash {
+   public:
+    size_t operator()(const CacheKey& key) const {
+      return base::hash_combine(static_cast<uint8_t>(key.first), *key.second);
+    }
+  };
 
- private:
-  friend class NativeModule;
-  using CacheKey = std::pair<uint8_t, FunctionSig>;
+  // Helper class to modify the cache under a lock.
+  class ModificationScope {
+   public:
+    explicit ModificationScope(WasmImportWrapperCache* cache)
+        : cache_(cache), guard_(&cache->mutex_) {}
+
+    V8_EXPORT_PRIVATE WasmCode*& operator[](const CacheKey& key);
+
+   private:
+    WasmImportWrapperCache* const cache_;
+    base::MutexGuard guard_;
+  };
 
-  mutable base::Mutex mutex_;
-  NativeModule* native_module_;
-  std::unordered_map<CacheKey, WasmCode*, base::hash<CacheKey>> entry_map_;
+  // Assumes the key exists in the map.
+  V8_EXPORT_PRIVATE WasmCode* Get(compiler::WasmImportCallKind kind,
+                                  FunctionSig* sig) const;
 
-  explicit WasmImportWrapperCache(NativeModule* native_module)
-      : native_module_(native_module) {}
+  ~WasmImportWrapperCache();
+
+ private:
+  base::Mutex mutex_;
+  std::unordered_map<CacheKey, WasmCode*, CacheKeyHash> entry_map_;
 };
 
 }  // namespace wasm

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

@@ -4,6 +4,7 @@
 
 #include "src/compiler/wasm-compiler.h"
 #include "src/wasm/function-compiler.h"
+#include "src/wasm/module-compiler.h"
 #include "src/wasm/wasm-code-manager.h"
 #include "src/wasm/wasm-engine.h"
 #include "src/wasm/wasm-import-wrapper-cache.h"
@@ -32,17 +33,19 @@ TEST(CacheHit) {
   auto module = NewModule(isolate);
   TestSignatures sigs;
   WasmCodeRefScope wasm_code_ref_scope;
+  WasmImportWrapperCache::ModificationScope cache_scope(
+      module->import_wrapper_cache());
 
   auto kind = compiler::WasmImportCallKind::kJSFunctionArityMatch;
 
-  WasmCode* c1 = module->import_wrapper_cache()->GetOrCompile(
-      isolate->wasm_engine(), isolate->counters(), kind, sigs.i_i());
+  WasmCode* c1 =
+      CompileImportWrapper(isolate->wasm_engine(), module.get(),
+                           isolate->counters(), kind, sigs.i_i(), &cache_scope);
 
   CHECK_NOT_NULL(c1);
   CHECK_EQ(WasmCode::Kind::kWasmToJsWrapper, c1->kind());
 
-  WasmCode* c2 = module->import_wrapper_cache()->GetOrCompile(
-      isolate->wasm_engine(), isolate->counters(), kind, sigs.i_i());
+  WasmCode* c2 = cache_scope[{kind, sigs.i_i()}];
 
   CHECK_NOT_NULL(c2);
   CHECK_EQ(c1, c2);
@@ -53,20 +56,21 @@ TEST(CacheMissSig) {
   auto module = NewModule(isolate);
   TestSignatures sigs;
   WasmCodeRefScope wasm_code_ref_scope;
+  WasmImportWrapperCache::ModificationScope cache_scope(
+      module->import_wrapper_cache());
 
   auto kind = compiler::WasmImportCallKind::kJSFunctionArityMatch;
 
-  WasmCode* c1 = module->import_wrapper_cache()->GetOrCompile(
-      isolate->wasm_engine(), isolate->counters(), kind, sigs.i_i());
+  WasmCode* c1 =
+      CompileImportWrapper(isolate->wasm_engine(), module.get(),
+                           isolate->counters(), kind, sigs.i_i(), &cache_scope);
 
   CHECK_NOT_NULL(c1);
   CHECK_EQ(WasmCode::Kind::kWasmToJsWrapper, c1->kind());
 
-  WasmCode* c2 = module->import_wrapper_cache()->GetOrCompile(
-      isolate->wasm_engine(), isolate->counters(), kind, sigs.i_ii());
+  WasmCode* c2 = cache_scope[{kind, sigs.i_ii()}];
 
-  CHECK_NOT_NULL(c2);
-  CHECK_NE(c1, c2);
+  CHECK_NULL(c2);
 }
 
 TEST(CacheMissKind) {
@@ -74,21 +78,22 @@ TEST(CacheMissKind) {
   auto module = NewModule(isolate);
   TestSignatures sigs;
   WasmCodeRefScope wasm_code_ref_scope;
+  WasmImportWrapperCache::ModificationScope cache_scope(
+      module->import_wrapper_cache());
 
   auto kind1 = compiler::WasmImportCallKind::kJSFunctionArityMatch;
   auto kind2 = compiler::WasmImportCallKind::kJSFunctionArityMismatch;
 
-  WasmCode* c1 = module->import_wrapper_cache()->GetOrCompile(
-      isolate->wasm_engine(), isolate->counters(), kind1, sigs.i_i());
+  WasmCode* c1 = CompileImportWrapper(isolate->wasm_engine(), module.get(),
+                                      isolate->counters(), kind1, sigs.i_i(),
+                                      &cache_scope);
 
   CHECK_NOT_NULL(c1);
   CHECK_EQ(WasmCode::Kind::kWasmToJsWrapper, c1->kind());
 
-  WasmCode* c2 = module->import_wrapper_cache()->GetOrCompile(
-      isolate->wasm_engine(), isolate->counters(), kind2, sigs.i_i());
+  WasmCode* c2 = cache_scope[{kind2, sigs.i_i()}];
 
-  CHECK_NOT_NULL(c2);
-  CHECK_NE(c1, c2);
+  CHECK_NULL(c2);
 }
 
 TEST(CacheHitMissSig) {
@@ -96,29 +101,34 @@ TEST(CacheHitMissSig) {
   auto module = NewModule(isolate);
   TestSignatures sigs;
   WasmCodeRefScope wasm_code_ref_scope;
+  WasmImportWrapperCache::ModificationScope cache_scope(
+      module->import_wrapper_cache());
 
   auto kind = compiler::WasmImportCallKind::kJSFunctionArityMatch;
 
-  WasmCode* c1 = module->import_wrapper_cache()->GetOrCompile(
-      isolate->wasm_engine(), isolate->counters(), kind, sigs.i_i());
+  WasmCode* c1 =
+      CompileImportWrapper(isolate->wasm_engine(), module.get(),
+                           isolate->counters(), kind, sigs.i_i(), &cache_scope);
 
   CHECK_NOT_NULL(c1);
   CHECK_EQ(WasmCode::Kind::kWasmToJsWrapper, c1->kind());
 
-  WasmCode* c2 = module->import_wrapper_cache()->GetOrCompile(
-      isolate->wasm_engine(), isolate->counters(), kind, sigs.i_ii());
+  WasmCode* c2 = cache_scope[{kind, sigs.i_ii()}];
+
+  CHECK_NULL(c2);
+
+  c2 = CompileImportWrapper(isolate->wasm_engine(), module.get(),
+                            isolate->counters(), kind, sigs.i_ii(),
+                            &cache_scope);
 
-  CHECK_NOT_NULL(c2);
   CHECK_NE(c1, c2);
 
-  WasmCode* c3 = module->import_wrapper_cache()->GetOrCompile(
-      isolate->wasm_engine(), isolate->counters(), kind, sigs.i_i());
+  WasmCode* c3 = cache_scope[{kind, sigs.i_i()}];
 
   CHECK_NOT_NULL(c3);
   CHECK_EQ(c1, c3);
 
-  WasmCode* c4 = module->import_wrapper_cache()->GetOrCompile(
-      isolate->wasm_engine(), isolate->counters(), kind, sigs.i_ii());
+  WasmCode* c4 = cache_scope[{kind, sigs.i_ii()}];
 
   CHECK_NOT_NULL(c4);
   CHECK_EQ(c2, c4);

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

@@ -8,6 +8,7 @@
 #include "src/diagnostics/code-tracer.h"
 #include "src/heap/heap-inl.h"
 #include "src/wasm/graph-builder-interface.h"
+#include "src/wasm/module-compiler.h"
 #include "src/wasm/wasm-import-wrapper-cache.h"
 #include "src/wasm/wasm-memory.h"
 #include "src/wasm/wasm-objects-inl.h"
@@ -48,8 +49,15 @@ TestingModuleBuilder::TestingModuleBuilder(
     CodeSpaceMemoryModificationScope modification_scope(isolate_->heap());
     auto kind = compiler::GetWasmImportCallKind(maybe_import->js_function,
                                                 maybe_import->sig, false);
-    auto import_wrapper = native_module_->import_wrapper_cache()->GetOrCompile(
-        isolate_->wasm_engine(), isolate_->counters(), kind, maybe_import->sig);
+    WasmImportWrapperCache::ModificationScope cache_scope(
+        native_module_->import_wrapper_cache());
+    WasmImportWrapperCache::CacheKey key(kind, maybe_import->sig);
+    auto import_wrapper = cache_scope[key];
+    if (import_wrapper == nullptr) {
+      import_wrapper = CompileImportWrapper(
+          isolate_->wasm_engine(), native_module_, isolate_->counters(), kind,
+          maybe_import->sig, &cache_scope);
+    }
 
     ImportedFunctionEntry(instance_object_, maybe_import_index)
         .SetWasmToJs(isolate_, maybe_import->js_function, import_wrapper);