[wasm] Do not store the NativeModule in compilation units

Instead, pass it as a parameter to the compilation.
This makes compilation units slimmer with the end goal of them being
just the function index and execution tier.
It also makes ownership handling of the NativeModule easier.

R=titzer@chromium.org

Bug: v8:8343, v8:7921
Change-Id: I0522c894569c71d8b7245f5ed5612ab2a249e1ad
Reviewed-on: https://chromium-review.googlesource.com/c/1406668Reviewed-by: Ben Titzer <titzer@chromium.org>
Commit-Queue: Clemens Hammacher <clemensh@chromium.org>
Cr-Commit-Position: refs/heads/master@{#58776}

src/compiler/wasm-compiler.cc

@@ -5717,9 +5717,10 @@ TurbofanWasmCompilationUnit::TurbofanWasmCompilationUnit(
 TurbofanWasmCompilationUnit::~TurbofanWasmCompilationUnit() = default;
 
 bool TurbofanWasmCompilationUnit::BuildGraphForWasmFunction(
-    wasm::CompilationEnv* env, const wasm::FunctionBody& func_body,
-    wasm::WasmFeatures* detected, double* decode_ms, MachineGraph* mcgraph,
-    NodeOriginTable* node_origins, SourcePositionTable* source_positions) {
+    wasm::CompilationEnv* env, wasm::NativeModule* native_module,
+    const wasm::FunctionBody& func_body, wasm::WasmFeatures* detected,
+    double* decode_ms, MachineGraph* mcgraph, NodeOriginTable* node_origins,
+    SourcePositionTable* source_positions) {
   base::ElapsedTimer decode_timer;
   if (FLAG_trace_wasm_decode_time) {
     decode_timer.Start();
@@ -5729,15 +5730,14 @@ bool TurbofanWasmCompilationUnit::BuildGraphForWasmFunction(
   WasmGraphBuilder builder(env, mcgraph->zone(), mcgraph, func_body.sig,
                            source_positions);
   wasm::VoidResult graph_construction_result = wasm::BuildTFGraph(
-      wasm_unit_->wasm_engine_->allocator(),
-      wasm_unit_->native_module_->enabled_features(), env->module, &builder,
-      detected, func_body, node_origins);
+      wasm_unit_->wasm_engine_->allocator(), native_module->enabled_features(),
+      env->module, &builder, detected, func_body, node_origins);
   if (graph_construction_result.failed()) {
     if (FLAG_trace_wasm_compiler) {
       StdoutStream{} << "Compilation failed: "
                      << graph_construction_result.error_msg() << std::endl;
     }
-    wasm_unit_->native_module()->compilation_state()->SetError(
+    native_module->compilation_state()->SetError(
         wasm_unit_->func_index_, std::move(graph_construction_result));
     return false;
   }
@@ -5778,8 +5778,9 @@ Vector<const char> GetDebugName(Zone* zone, int index) {
 }  // namespace
 
 void TurbofanWasmCompilationUnit::ExecuteCompilation(
-    wasm::CompilationEnv* env, const wasm::FunctionBody& func_body,
-    Counters* counters, wasm::WasmFeatures* detected) {
+    wasm::CompilationEnv* env, wasm::NativeModule* native_module,
+    const wasm::FunctionBody& func_body, Counters* counters,
+    wasm::WasmFeatures* detected) {
   TRACE_EVENT0(TRACE_DISABLED_BY_DEFAULT("v8.wasm"),
                "ExecuteTurbofanCompilation");
   double decode_ms = 0;
@@ -5810,8 +5811,9 @@ void TurbofanWasmCompilationUnit::ExecuteCompilation(
                                       : nullptr;
   SourcePositionTable* source_positions =
       new (mcgraph->zone()) SourcePositionTable(mcgraph->graph());
-  if (!BuildGraphForWasmFunction(env, func_body, detected, &decode_ms, mcgraph,
-                                 node_origins, source_positions)) {
+  if (!BuildGraphForWasmFunction(env, native_module, func_body, detected,
+                                 &decode_ms, mcgraph, node_origins,
+                                 source_positions)) {
     // Compilation failed.
     return;
   }
@@ -5834,7 +5836,7 @@ void TurbofanWasmCompilationUnit::ExecuteCompilation(
 
   if (wasm::WasmCode* wasm_code = Pipeline::GenerateCodeForWasmFunction(
           &info, wasm_unit_->wasm_engine_, mcgraph, call_descriptor,
-          source_positions, node_origins, func_body, wasm_unit_->native_module_,
+          source_positions, node_origins, func_body, native_module,
           wasm_unit_->func_index_)) {
     wasm_unit_->SetResult(wasm_code, counters);
   }

src/compiler/wasm-compiler.h

@@ -51,14 +51,16 @@ class TurbofanWasmCompilationUnit {
   ~TurbofanWasmCompilationUnit();
 
   bool BuildGraphForWasmFunction(wasm::CompilationEnv* env,
+                                 wasm::NativeModule* native_module,
                                  const wasm::FunctionBody& func_body,
                                  wasm::WasmFeatures* detected,
                                  double* decode_ms, MachineGraph* mcgraph,
                                  NodeOriginTable* node_origins,
                                  SourcePositionTable* source_positions);
 
-  void ExecuteCompilation(wasm::CompilationEnv*, const wasm::FunctionBody&,
-                          Counters*, wasm::WasmFeatures* detected);
+  void ExecuteCompilation(wasm::CompilationEnv*, wasm::NativeModule*,
+                          const wasm::FunctionBody&, Counters*,
+                          wasm::WasmFeatures* detected);
 
  private:
   wasm::WasmCompilationUnit* const wasm_unit_;

src/wasm/baseline/liftoff-compiler.cc

@@ -1947,6 +1947,7 @@ class LiftoffCompiler {
 }  // namespace
 
 bool LiftoffCompilationUnit::ExecuteCompilation(CompilationEnv* env,
+                                                NativeModule* native_module,
                                                 const FunctionBody& func_body,
                                                 Counters* counters,
                                                 WasmFeatures* detected) {
@@ -1963,8 +1964,8 @@ bool LiftoffCompilationUnit::ExecuteCompilation(CompilationEnv* env,
   base::Optional<TimedHistogramScope> liftoff_compile_time_scope(
       base::in_place, counters->liftoff_compile_time());
   WasmFullDecoder<Decoder::kValidate, LiftoffCompiler> decoder(
-      &zone, module, wasm_unit_->native_module_->enabled_features(), detected,
-      func_body, call_descriptor, env, &zone);
+      &zone, module, native_module->enabled_features(), detected, func_body,
+      call_descriptor, env, &zone);
   decoder.Decode();
   liftoff_compile_time_scope.reset();
   LiftoffCompiler* compiler = &decoder.interface();
@@ -1993,7 +1994,7 @@ bool LiftoffCompilationUnit::ExecuteCompilation(CompilationEnv* env,
   uint32_t frame_slot_count = compiler->GetTotalFrameSlotCount();
   int safepoint_table_offset = compiler->GetSafepointTableOffset();
 
-  WasmCode* code = wasm_unit_->native_module_->AddCode(
+  WasmCode* code = native_module->AddCode(
       wasm_unit_->func_index_, desc, frame_slot_count, safepoint_table_offset,
       0, std::move(protected_instructions), std::move(source_positions),
       WasmCode::kFunction, WasmCode::kLiftoff);

src/wasm/baseline/liftoff-compiler.h

@@ -16,6 +16,7 @@ namespace wasm {
 
 struct CompilationEnv;
 struct FunctionBody;
+class NativeModule;
 class WasmCompilationUnit;
 struct WasmFeatures;
 
@@ -24,8 +25,8 @@ class LiftoffCompilationUnit final {
   explicit LiftoffCompilationUnit(WasmCompilationUnit* wasm_unit)
       : wasm_unit_(wasm_unit) {}
 
-  bool ExecuteCompilation(CompilationEnv*, const FunctionBody&, Counters*,
-                          WasmFeatures* detected);
+  bool ExecuteCompilation(CompilationEnv*, NativeModule*, const FunctionBody&,
+                          Counters*, WasmFeatures* detected);
 
  private:
   WasmCompilationUnit* const wasm_unit_;

src/wasm/function-compiler.cc

@@ -31,24 +31,16 @@ const char* GetExecutionTierAsString(ExecutionTier tier) {
 }  // namespace
 
 // static
-ExecutionTier WasmCompilationUnit::GetDefaultExecutionTier() {
-  return FLAG_liftoff ? ExecutionTier::kBaseline : ExecutionTier::kOptimized;
+ExecutionTier WasmCompilationUnit::GetDefaultExecutionTier(
+    const WasmModule* module) {
+  return FLAG_liftoff && module->origin == kWasmOrigin
+             ? ExecutionTier::kBaseline
+             : ExecutionTier::kOptimized;
 }
 
-WasmCompilationUnit::WasmCompilationUnit(WasmEngine* wasm_engine,
-                                         NativeModule* native_module, int index,
+WasmCompilationUnit::WasmCompilationUnit(WasmEngine* wasm_engine, int index,
                                          ExecutionTier tier)
-    : wasm_engine_(wasm_engine),
-      func_index_(index),
-      native_module_(native_module),
-      tier_(tier) {
-  const WasmModule* module = native_module->module();
-  DCHECK_GE(index, module->num_imported_functions);
-  DCHECK_LT(index, module->functions.size());
-  // Always disable Liftoff for asm.js, for two reasons:
-  //    1) asm-specific opcodes are not implemented, and
-  //    2) tier-up does not work with lazy compilation.
-  if (module->origin == kAsmJsOrigin) tier = ExecutionTier::kOptimized;
+    : wasm_engine_(wasm_engine), func_index_(index), tier_(tier) {
   if (V8_UNLIKELY(FLAG_wasm_tier_mask_for_testing) && index < 32 &&
       (FLAG_wasm_tier_mask_for_testing & (1 << index))) {
     tier = ExecutionTier::kOptimized;
@@ -61,20 +53,18 @@ WasmCompilationUnit::WasmCompilationUnit(WasmEngine* wasm_engine,
 WasmCompilationUnit::~WasmCompilationUnit() = default;
 
 void WasmCompilationUnit::ExecuteCompilation(
-    CompilationEnv* env, std::shared_ptr<WireBytesStorage> wire_bytes_storage,
-    Counters* counters, WasmFeatures* detected) {
-  const WasmModule* module = native_module_->module();
-  DCHECK_EQ(module, env->module);
-
+    CompilationEnv* env, NativeModule* native_module,
+    std::shared_ptr<WireBytesStorage> wire_bytes_storage, Counters* counters,
+    WasmFeatures* detected) {
   auto* func = &env->module->functions[func_index_];
   Vector<const uint8_t> code = wire_bytes_storage->GetCode(func->code);
   wasm::FunctionBody func_body{func->sig, func->code.offset(), code.start(),
                                code.end()};
 
-  auto size_histogram =
-      SELECT_WASM_COUNTER(counters, module->origin, wasm, function_size_bytes);
+  auto size_histogram = SELECT_WASM_COUNTER(counters, env->module->origin, wasm,
+                                            function_size_bytes);
   size_histogram->AddSample(static_cast<int>(func_body.end - func_body.start));
-  auto timed_histogram = SELECT_WASM_COUNTER(counters, module->origin,
+  auto timed_histogram = SELECT_WASM_COUNTER(counters, env->module->origin,
                                              wasm_compile, function_time);
   TimedHistogramScope wasm_compile_function_time_scope(timed_histogram);
 
@@ -85,8 +75,8 @@ void WasmCompilationUnit::ExecuteCompilation(
 
   switch (tier_) {
     case ExecutionTier::kBaseline:
-      if (liftoff_unit_->ExecuteCompilation(env, func_body, counters,
-                                            detected)) {
+      if (liftoff_unit_->ExecuteCompilation(env, native_module, func_body,
+                                            counters, detected)) {
         break;
       }
       // Otherwise, fall back to turbofan.
@@ -95,7 +85,8 @@ void WasmCompilationUnit::ExecuteCompilation(
       // function to avoid compiling it twice with TurboFan.
       V8_FALLTHROUGH;
     case ExecutionTier::kOptimized:
-      turbofan_unit_->ExecuteCompilation(env, func_body, counters, detected);
+      turbofan_unit_->ExecuteCompilation(env, native_module, func_body,
+                                         counters, detected);
       break;
     case ExecutionTier::kInterpreter:
       UNREACHABLE();  // TODO(titzer): compile interpreter entry stub.
@@ -135,18 +126,18 @@ void WasmCompilationUnit::CompileWasmFunction(Isolate* isolate,
                              wire_bytes.start() + function->code.offset(),
                              wire_bytes.start() + function->code.end_offset()};
 
-  WasmCompilationUnit unit(isolate->wasm_engine(), native_module,
-                           function->func_index, tier);
+  WasmCompilationUnit unit(isolate->wasm_engine(), function->func_index, tier);
   CompilationEnv env = native_module->CreateCompilationEnv();
   unit.ExecuteCompilation(
-      &env, native_module->compilation_state()->GetWireBytesStorage(),
+      &env, native_module,
+      native_module->compilation_state()->GetWireBytesStorage(),
       isolate->counters(), detected);
 }
 
 void WasmCompilationUnit::SetResult(WasmCode* code, Counters* counters) {
   DCHECK_NULL(result_);
   result_ = code;
-  native_module()->PublishCode(code);
+  code->native_module()->PublishCode(code);
 
   counters->wasm_generated_code_size()->Increment(
       static_cast<int>(code->instructions().size()));

src/wasm/function-compiler.h

@@ -30,29 +30,27 @@ struct WasmFunction;
 
 class WasmCompilationUnit final {
  public:
-  static ExecutionTier GetDefaultExecutionTier();
+  static ExecutionTier GetDefaultExecutionTier(const WasmModule*);
 
   // If constructing from a background thread, pass in a Counters*, and ensure
   // that the Counters live at least as long as this compilation unit (which
   // typically means to hold a std::shared_ptr<Counters>).
   // If used exclusively from a foreground thread, Isolate::counters() may be
   // used by callers to pass Counters.
-  WasmCompilationUnit(WasmEngine*, NativeModule*, int index,
-                      ExecutionTier = GetDefaultExecutionTier());
+  WasmCompilationUnit(WasmEngine*, int index, ExecutionTier);
 
   ~WasmCompilationUnit();
 
-  void ExecuteCompilation(CompilationEnv*, std::shared_ptr<WireBytesStorage>,
-                          Counters*, WasmFeatures* detected);
+  void ExecuteCompilation(CompilationEnv*, NativeModule*,
+                          std::shared_ptr<WireBytesStorage>, Counters*,
+                          WasmFeatures* detected);
 
-  NativeModule* native_module() const { return native_module_; }
   ExecutionTier tier() const { return tier_; }
   WasmCode* result() const { return result_; }
 
-  static void CompileWasmFunction(Isolate* isolate, NativeModule* native_module,
-                                  WasmFeatures* detected,
-                                  const WasmFunction* function,
-                                  ExecutionTier = GetDefaultExecutionTier());
+  static void CompileWasmFunction(Isolate*, NativeModule*,
+                                  WasmFeatures* detected, const WasmFunction*,
+                                  ExecutionTier);
 
  private:
   friend class LiftoffCompilationUnit;
@@ -60,7 +58,6 @@ class WasmCompilationUnit final {
 
   WasmEngine* const wasm_engine_;
   const int func_index_;
-  NativeModule* const native_module_;
   ExecutionTier tier_;
   WasmCode* result_ = nullptr;
 

src/wasm/module-compiler.cc

@@ -550,10 +550,13 @@ WasmCode* LazyCompileFunction(Isolate* isolate, NativeModule* native_module,
                          module_start + func->code.offset(),
                          module_start + func->code.end_offset()};
 
-  WasmCompilationUnit unit(isolate->wasm_engine(), native_module, func_index);
+  WasmCompilationUnit unit(
+      isolate->wasm_engine(), func_index,
+      WasmCompilationUnit::GetDefaultExecutionTier(native_module->module()));
   CompilationEnv env = native_module->CreateCompilationEnv();
   unit.ExecuteCompilation(
-      &env, native_module->compilation_state()->GetWireBytesStorage(),
+      &env, native_module,
+      native_module->compilation_state()->GetWireBytesStorage(),
       isolate->counters(),
       Impl(native_module->compilation_state())->detected_features());
 
@@ -607,7 +610,10 @@ class CompilationUnitBuilder {
  public:
   explicit CompilationUnitBuilder(NativeModule* native_module,
                                   WasmEngine* wasm_engine)
-      : native_module_(native_module), wasm_engine_(wasm_engine) {}
+      : native_module_(native_module),
+        wasm_engine_(wasm_engine),
+        default_tier_(WasmCompilationUnit::GetDefaultExecutionTier(
+            native_module->module())) {}
 
   void AddUnit(uint32_t func_index) {
     switch (compilation_state()->compile_mode()) {
@@ -618,8 +624,7 @@ class CompilationUnitBuilder {
             CreateUnit(func_index, ExecutionTier::kBaseline));
         return;
       case CompileMode::kRegular:
-        baseline_units_.emplace_back(CreateUnit(
-            func_index, WasmCompilationUnit::GetDefaultExecutionTier()));
+        baseline_units_.emplace_back(CreateUnit(func_index, default_tier_));
         return;
     }
     UNREACHABLE();
@@ -640,8 +645,8 @@ class CompilationUnitBuilder {
  private:
   std::unique_ptr<WasmCompilationUnit> CreateUnit(uint32_t func_index,
                                                   ExecutionTier tier) {
-    return base::make_unique<WasmCompilationUnit>(wasm_engine_, native_module_,
-                                                  func_index, tier);
+    return base::make_unique<WasmCompilationUnit>(wasm_engine_, func_index,
+                                                  tier);
   }
 
   CompilationStateImpl* compilation_state() const {
@@ -650,6 +655,7 @@ class CompilationUnitBuilder {
 
   NativeModule* const native_module_;
   WasmEngine* const wasm_engine_;
+  const ExecutionTier default_tier_;
   std::vector<std::unique_ptr<WasmCompilationUnit>> baseline_units_;
   std::vector<std::unique_ptr<WasmCompilationUnit>> tiering_units_;
 };
@@ -685,6 +691,7 @@ double MonotonicallyIncreasingTimeInMs() {
 // within the result_mutex_ lock when no finishing task is running, i.e. when
 // the finisher_is_running_ flag is not set.
 bool FetchAndExecuteCompilationUnit(CompilationEnv* env,
+                                    NativeModule* native_module,
                                     CompilationStateImpl* compilation_state,
                                     WasmFeatures* detected,
                                     Counters* counters) {
@@ -697,7 +704,8 @@ bool FetchAndExecuteCompilationUnit(CompilationEnv* env,
   // Get the tier before starting compilation, as compilation can switch tiers
   // if baseline bails out.
   ExecutionTier tier = unit->tier();
-  unit->ExecuteCompilation(env, compilation_state->GetSharedWireBytesStorage(),
+  unit->ExecuteCompilation(env, native_module,
+                           compilation_state->GetSharedWireBytesStorage(),
                            counters, detected);
   compilation_state->OnFinishedUnit(tier, unit->result());
 
@@ -770,7 +778,7 @@ void CompileInParallel(Isolate* isolate, NativeModule* native_module) {
   //    a time and execute the parallel phase of the compilation unit.
   WasmFeatures detected_features;
   CompilationEnv env = native_module->CreateCompilationEnv();
-  while (FetchAndExecuteCompilationUnit(&env, compilation_state,
+  while (FetchAndExecuteCompilationUnit(&env, native_module, compilation_state,
                                         &detected_features,
                                         isolate->counters()) &&
          !compilation_state->baseline_compilation_finished()) {
@@ -812,12 +820,14 @@ void CompileSequentially(Isolate* isolate, NativeModule* native_module,
   const WasmModule* module = native_module->module();
   WasmFeatures detected = kNoWasmFeatures;
   auto* comp_state = Impl(native_module->compilation_state());
+  ExecutionTier tier =
+      WasmCompilationUnit::GetDefaultExecutionTier(native_module->module());
   for (const WasmFunction& func : module->functions) {
     if (func.imported) continue;  // Imports are compiled at instantiation time.
 
     // Compile the function.
     WasmCompilationUnit::CompileWasmFunction(isolate, native_module, &detected,
-                                             &func);
+                                             &func, tier);
     if (comp_state->failed()) {
       thrower->CompileFailed(comp_state->GetCompileError());
       break;
@@ -975,8 +985,9 @@ class BackgroundCompileTask : public CancelableTask {
     auto* compilation_state = Impl(native_module_->compilation_state());
     WasmFeatures detected_features = kNoWasmFeatures;
     while (!compilation_state->failed()) {
-      if (!FetchAndExecuteCompilationUnit(&env, compilation_state,
-                                          &detected_features, counters_)) {
+      if (!FetchAndExecuteCompilationUnit(&env, native_module_,
+                                          compilation_state, &detected_features,
+                                          counters_)) {
         break;
       }
     }

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

@@ -419,11 +419,12 @@ void WasmFunctionCompiler::Build(const byte* start, const byte* end) {
                          func_wire_bytes.start(), func_wire_bytes.end()};
   NativeModule* native_module =
       builder_->instance_object()->module_object()->native_module();
-  WasmCompilationUnit unit(isolate()->wasm_engine(), native_module,
-                           function_->func_index, tier);
+  WasmCompilationUnit unit(isolate()->wasm_engine(), function_->func_index,
+                           tier);
   WasmFeatures unused_detected_features;
   unit.ExecuteCompilation(
-      &env, native_module->compilation_state()->GetWireBytesStorage(),
+      &env, native_module,
+      native_module->compilation_state()->GetWireBytesStorage(),
       isolate()->counters(), &unused_detected_features);
   WasmCode* result = unit.result();
   DCHECK_NOT_NULL(result);