Reland "[wasm] Support partial serialization of modules"

The original CL was flaky because deserialization did not wait correctly
for the compilation of missing functions to finish. The baseline-finished
event was set even when there were still some functions missing. The
combination of deserialization and lazy compilation was also not handled
correctly.

Original change's description:
> [wasm] Support partial serialization of modules
>
> At the moment a WebAssembly module can be serialized successfully when
> all functions were compiled with TurboFan. However, for some functions
> it may not be necessary to be compiled with TurboFan, e.g. for functions
> where Liftoff code is as good as TurboFan code.
>
> With this CL we allow WebAssembly modules to get serialized even when
> not all functions are compiled with TurboFan. Missing functions are
> marked as missing in the serlialization. Upon deserialization, missing
> functions either get compiled by Liftoff, or initialized with a
> lazy-compilation stub, depending on the V8 configuration.
>
> Bug: v8:11862

Change-Id: I79a9e8e14199cff87fce6ae41a87087e047bbc65
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/3060485
Reviewed-by: Clemens Backes <clemensb@chromium.org>
Commit-Queue: Andreas Haas <ahaas@chromium.org>
Cr-Commit-Position: refs/heads/master@{#76017}

src/wasm/compilation-environment.h

@@ -132,7 +132,8 @@ class V8_EXPORT_PRIVATE CompilationState {
 
   void AddCallback(callback_t);
 
-  void InitializeAfterDeserialization();
+  void InitializeAfterDeserialization(
+      base::Vector<const int> missing_functions);
 
   // Wait until top tier compilation finished, or compilation failed.
   void WaitForTopTierFinished();

src/wasm/module-compiler.cc

@@ -553,7 +553,8 @@ class CompilationStateImpl {
 
   // Initialize the compilation progress after deserialization. This is needed
   // for recompilation (e.g. for tier down) to work later.
-  void InitializeCompilationProgressAfterDeserialization();
+  void InitializeCompilationProgressAfterDeserialization(
+      base::Vector<const int> missing_functions);
 
   // Initializes compilation units based on the information encoded in the
   // {compilation_progress_}.
@@ -660,6 +661,10 @@ class CompilationStateImpl {
   }
 
  private:
+  uint8_t SetupCompilationProgressForFunction(
+      bool lazy_module, const WasmModule* module,
+      const WasmFeatures& enabled_features, int func_index);
+
   // Returns the potentially-updated {function_progress}.
   uint8_t AddCompilationUnitInternal(CompilationUnitBuilder* builder,
                                      int function_index,
@@ -831,8 +836,10 @@ void CompilationState::WaitForTopTierFinished() {
 
 void CompilationState::SetHighPriority() { Impl(this)->SetHighPriority(); }
 
-void CompilationState::InitializeAfterDeserialization() {
-  Impl(this)->InitializeCompilationProgressAfterDeserialization();
+void CompilationState::InitializeAfterDeserialization(
+    base::Vector<const int> missing_functions) {
+  Impl(this)->InitializeCompilationProgressAfterDeserialization(
+      missing_functions);
 }
 
 bool CompilationState::failed() const { return Impl(this)->failed(); }
@@ -2854,6 +2861,38 @@ bool CompilationStateImpl::cancelled() const {
   return compile_cancelled_.load(std::memory_order_relaxed);
 }
 
+uint8_t CompilationStateImpl::SetupCompilationProgressForFunction(
+    bool lazy_module, const WasmModule* module,
+    const WasmFeatures& enabled_features, int func_index) {
+  ExecutionTierPair requested_tiers =
+      GetRequestedExecutionTiers(module, enabled_features, func_index);
+  CompileStrategy strategy =
+      GetCompileStrategy(module, enabled_features, func_index, lazy_module);
+
+  bool required_for_baseline = strategy == CompileStrategy::kEager;
+  bool required_for_top_tier = strategy != CompileStrategy::kLazy;
+  DCHECK_EQ(required_for_top_tier,
+            strategy == CompileStrategy::kEager ||
+                strategy == CompileStrategy::kLazyBaselineEagerTopTier);
+
+  // Count functions to complete baseline and top tier compilation.
+  if (required_for_baseline) outstanding_baseline_units_++;
+  if (required_for_top_tier) outstanding_top_tier_functions_++;
+
+  // Initialize function's compilation progress.
+  ExecutionTier required_baseline_tier = required_for_baseline
+                                             ? requested_tiers.baseline_tier
+                                             : ExecutionTier::kNone;
+  ExecutionTier required_top_tier =
+      required_for_top_tier ? requested_tiers.top_tier : ExecutionTier::kNone;
+  uint8_t function_progress =
+      ReachedTierField::encode(ExecutionTier::kNone) |
+      RequiredBaselineTierField::encode(required_baseline_tier) |
+      RequiredTopTierField::encode(required_top_tier);
+
+  return function_progress;
+}
+
 void CompilationStateImpl::InitializeCompilationProgress(
     bool lazy_module, int num_import_wrappers, int num_export_wrappers) {
   DCHECK(!failed());
@@ -2880,32 +2919,8 @@ void CompilationStateImpl::InitializeCompilationProgress(
       outstanding_top_tier_functions_++;
       continue;
     }
-    ExecutionTierPair requested_tiers =
-        GetRequestedExecutionTiers(module, enabled_features, func_index);
-    CompileStrategy strategy =
-        GetCompileStrategy(module, enabled_features, func_index, lazy_module);
-
-    bool required_for_baseline = strategy == CompileStrategy::kEager;
-    bool required_for_top_tier = strategy != CompileStrategy::kLazy;
-    DCHECK_EQ(required_for_top_tier,
-              strategy == CompileStrategy::kEager ||
-                  strategy == CompileStrategy::kLazyBaselineEagerTopTier);
-
-    // Count functions to complete baseline and top tier compilation.
-    if (required_for_baseline) outstanding_baseline_units_++;
-    if (required_for_top_tier) outstanding_top_tier_functions_++;
-
-    // Initialize function's compilation progress.
-    ExecutionTier required_baseline_tier = required_for_baseline
-                                               ? requested_tiers.baseline_tier
-                                               : ExecutionTier::kNone;
-    ExecutionTier required_top_tier =
-        required_for_top_tier ? requested_tiers.top_tier : ExecutionTier::kNone;
-    uint8_t function_progress = ReachedTierField::encode(ExecutionTier::kNone);
-    function_progress = RequiredBaselineTierField::update(
-        function_progress, required_baseline_tier);
-    function_progress =
-        RequiredTopTierField::update(function_progress, required_top_tier);
+    uint8_t function_progress = SetupCompilationProgressForFunction(
+        lazy_module, module, enabled_features, func_index);
     compilation_progress_.push_back(function_progress);
   }
   DCHECK_IMPLIES(lazy_module, outstanding_baseline_units_ == 0);
@@ -3015,19 +3030,37 @@ void CompilationStateImpl::AddCompilationUnit(CompilationUnitBuilder* builder,
   }
 }
 
-void CompilationStateImpl::InitializeCompilationProgressAfterDeserialization() {
+void CompilationStateImpl::InitializeCompilationProgressAfterDeserialization(
+    base::Vector<const int> missing_functions) {
   auto* module = native_module_->module();
-  base::MutexGuard guard(&callbacks_mutex_);
-  DCHECK(compilation_progress_.empty());
-  constexpr uint8_t kProgressAfterDeserialization =
-      RequiredBaselineTierField::encode(ExecutionTier::kTurbofan) |
-      RequiredTopTierField::encode(ExecutionTier::kTurbofan) |
-      ReachedTierField::encode(ExecutionTier::kTurbofan);
-  finished_events_.Add(CompilationEvent::kFinishedExportWrappers);
-  finished_events_.Add(CompilationEvent::kFinishedBaselineCompilation);
-  finished_events_.Add(CompilationEvent::kFinishedTopTierCompilation);
-  compilation_progress_.assign(module->num_declared_functions,
-                               kProgressAfterDeserialization);
+  auto enabled_features = native_module_->enabled_features();
+  const bool lazy_module = IsLazyModule(module);
+  {
+    base::MutexGuard guard(&callbacks_mutex_);
+    DCHECK(compilation_progress_.empty());
+    constexpr uint8_t kProgressAfterDeserialization =
+        RequiredBaselineTierField::encode(ExecutionTier::kTurbofan) |
+        RequiredTopTierField::encode(ExecutionTier::kTurbofan) |
+        ReachedTierField::encode(ExecutionTier::kTurbofan);
+    finished_events_.Add(CompilationEvent::kFinishedExportWrappers);
+    if (missing_functions.empty() || FLAG_wasm_lazy_compilation) {
+      finished_events_.Add(CompilationEvent::kFinishedBaselineCompilation);
+      finished_events_.Add(CompilationEvent::kFinishedTopTierCompilation);
+    }
+    compilation_progress_.assign(module->num_declared_functions,
+                                 kProgressAfterDeserialization);
+    uint32_t num_imported_functions = module->num_imported_functions;
+    for (auto func_index : missing_functions) {
+      if (FLAG_wasm_lazy_compilation) {
+        native_module_->UseLazyStub(num_imported_functions + func_index);
+      }
+      compilation_progress_[func_index] = SetupCompilationProgressForFunction(
+          lazy_module, module, enabled_features, func_index);
+    }
+  }
+  auto builder = std::make_unique<CompilationUnitBuilder>(native_module_);
+  InitializeCompilationUnits(std::move(builder));
+  WaitForCompilationEvent(CompilationEvent::kFinishedBaselineCompilation);
 }
 
 void CompilationStateImpl::InitializeRecompilation(

src/wasm/wasm-serialization.cc

@@ -304,7 +304,7 @@ NativeModuleSerializer::NativeModuleSerializer(
 size_t NativeModuleSerializer::MeasureCode(const WasmCode* code) const {
   if (code == nullptr) return sizeof(bool);
   DCHECK_EQ(WasmCode::kFunction, code->kind());
-  if (FLAG_wasm_lazy_compilation && code->tier() != ExecutionTier::kTurbofan) {
+  if (code->tier() != ExecutionTier::kTurbofan) {
     return sizeof(bool);
   }
   return kCodeHeaderSize + code->instructions().size() +
@@ -338,11 +338,8 @@ bool NativeModuleSerializer::WriteCode(const WasmCode* code, Writer* writer) {
   // Only serialize TurboFan code, as Liftoff code can contain breakpoints or
   // non-relocatable constants.
   if (code->tier() != ExecutionTier::kTurbofan) {
-    if (FLAG_wasm_lazy_compilation) {
-      writer->Write(false);
-      return true;
-    }
-    return false;
+    writer->Write(false);
+    return true;
   }
   writer->Write(true);
   // Write the size of the entire code section, followed by the code header.
@@ -536,6 +533,10 @@ class V8_EXPORT_PRIVATE NativeModuleDeserializer {
 
   bool Read(Reader* reader);
 
+  base::Vector<const int> missing_functions() {
+    return base::VectorOf(missing_functions_);
+  }
+
  private:
   friend class CopyAndRelocTask;
   friend class PublishTask;
@@ -554,6 +555,7 @@ class V8_EXPORT_PRIVATE NativeModuleDeserializer {
   size_t remaining_code_size_ = 0;
   base::Vector<byte> current_code_space_;
   NativeModule::JumpTablesRef current_jump_tables_;
+  std::vector<int> missing_functions_;
 };
 
 class CopyAndRelocTask : public JobTask {
@@ -687,9 +689,7 @@ DeserializationUnit NativeModuleDeserializer::ReadCode(int fn_index,
                                                        Reader* reader) {
   bool has_code = reader->Read<bool>();
   if (!has_code) {
-    DCHECK(FLAG_wasm_lazy_compilation ||
-           native_module_->enabled_features().has_compilation_hints());
-    native_module_->UseLazyStub(fn_index);
+    missing_functions_.push_back(fn_index);
     return {};
   }
   int constant_pool_offset = reader->Read<int>();
@@ -862,9 +862,14 @@ MaybeHandle<WasmModuleObject> DeserializeNativeModule(
     NativeModuleDeserializer deserializer(shared_native_module.get());
     Reader reader(data + WasmSerializer::kHeaderSize);
     bool error = !deserializer.Read(&reader);
-    shared_native_module->compilation_state()->InitializeAfterDeserialization();
+    if (error) {
+      wasm_engine->UpdateNativeModuleCache(error, &shared_native_module,
+                                           isolate);
+      return {};
+    }
+    shared_native_module->compilation_state()->InitializeAfterDeserialization(
+        deserializer.missing_functions());
     wasm_engine->UpdateNativeModuleCache(error, &shared_native_module, isolate);
-    if (error) return {};
   }
 
   Handle<FixedArray> export_wrappers;

test/mjsunit/wasm/test-partial-serialization.js

@@ -0,0 +1,56 @@
+// Copyright 2021 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.
+
+// Flags: --allow-natives-syntax --liftoff --no-wasm-tier-up --expose-gc
+// Compile functions 0 and 2 with Turbofan, the rest with Liftoff:
+// Flags: --wasm-tier-mask-for-testing=5
+
+d8.file.execute('test/mjsunit/wasm/wasm-module-builder.js');
+
+const num_functions = 5;
+
+function create_builder() {
+  const builder = new WasmModuleBuilder();
+  for (let i = 0; i < num_functions; ++i) {
+    builder.addFunction('f' + i, kSig_i_v)
+        .addBody(wasmI32Const(i))
+        .exportFunc();
+  }
+  return builder;
+}
+
+function check(instance) {
+  for (let i = 0; i < num_functions; ++i) {
+    const expect_liftoff = i != 0 && i != 2;
+    assertEquals(
+        expect_liftoff, %IsLiftoffFunction(instance.exports['f' + i]),
+        'function ' + i);
+  }
+}
+
+const wire_bytes = create_builder().toBuffer();
+
+function testTierTestingFlag() {
+  print(arguments.callee.name);
+  const module = new WebAssembly.Module(wire_bytes);
+  const buff = %SerializeWasmModule(module);
+  const instance = new WebAssembly.Instance(module);
+  check(instance);
+  return buff;
+};
+
+const serialized_module = testTierTestingFlag();
+// Do some GCs to make sure the first module got collected and removed from the
+// module cache.
+gc();
+gc();
+gc();
+
+(function testSerializedModule() {
+  print(arguments.callee.name);
+  const module = %DeserializeWasmModule(serialized_module, wire_bytes);
+
+  const instance = new WebAssembly.Instance(module);
+  check(instance);
+})();

test/mjsunit/wasm/test-serialization-with-lazy-compilation.js

@@ -0,0 +1,43 @@
+// Copyright 2021 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.
+
+// Flags: --allow-natives-syntax --wasm-lazy-compilation --expose-gc
+
+d8.file.execute('test/mjsunit/wasm/wasm-module-builder.js');
+
+const num_functions = 2;
+
+function create_builder() {
+  const builder = new WasmModuleBuilder();
+  for (let i = 0; i < num_functions; ++i) {
+    builder.addFunction('f' + i, kSig_i_v)
+        .addBody(wasmI32Const(i))
+        .exportFunc();
+  }
+  return builder;
+}
+
+const wire_bytes = create_builder().toBuffer();
+
+function serializeModule() {
+  const module = new WebAssembly.Module(wire_bytes);
+  const buff = %SerializeWasmModule(module);
+  return buff;
+};
+
+const serialized_module = serializeModule();
+// Do some GCs to make sure the first module got collected and removed from the
+// module cache.
+gc();
+gc();
+gc();
+
+(function testSerializedModule() {
+  print(arguments.callee.name);
+  const module = %DeserializeWasmModule(serialized_module, wire_bytes);
+
+  const instance = new WebAssembly.Instance(module);
+  assertEquals(0, instance.exports.f0());
+  assertEquals(1, instance.exports.f1());
+})();