[wasm] Publish background compiled code in batches

In order to reduce lock contention in the NativeModule, to publish
compiled code in batches.
This is implemented via a new {NativeModule::AddCompiledCode} variant
that takes a {Vector<WasmCompilationResult>}, allocates code space for
all of the results, copies all code over and relocates it, and then
publishes all of it.

R=titzer@chromium.org

Bug: v8:8916
Change-Id: I437bd222dc2471b89b114cdb42049991af36f1f4
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/1532062
Commit-Queue: Clemens Hammacher <clemensh@chromium.org>
Reviewed-by: Ben Titzer <titzer@chromium.org>
Cr-Commit-Position: refs/heads/master@{#60373}

src/wasm/module-compiler.cc

@@ -695,18 +695,46 @@ class BackgroundCompileTask : public CancelableTask {
       }
     }
 
+    std::vector<WasmCompilationResult> results_to_publish;
+
+    auto publish_results =
+        [&results_to_publish](BackgroundCompileScope* compile_scope) {
+          if (results_to_publish.empty()) return;
+          // TODO(clemensh): Refactor {OnFinishedUnit} and remove this.
+          std::vector<ExecutionTier> requested_tiers;
+          requested_tiers.reserve(results_to_publish.size());
+          for (auto& result : results_to_publish) {
+            requested_tiers.push_back(result.requested_tier);
+          }
+
+          std::vector<WasmCode*> generated_code =
+              compile_scope->native_module()->AddCompiledCode(
+                  VectorOf(results_to_publish));
+          results_to_publish.clear();
+
+          // Account for the finished compilation units.
+          // TODO(clemensh): This takes a lock on each invokation. Only do this
+          // once and pass accumulated counts.
+          DCHECK_EQ(generated_code.size(), requested_tiers.size());
+          for (size_t i = 0; i < generated_code.size(); ++i) {
+            compile_scope->compilation_state()->OnFinishedUnit(
+                requested_tiers[i], generated_code[i]);
+          }
+        };
+
     bool compilation_failed = false;
     while (true) {
       // (asynchronous): Execute the compilation.
 
       WasmCompilationResult result = unit->ExecuteCompilation(
           &env.value(), wire_bytes, async_counters_.get(), &detected_features);
+      results_to_publish.emplace_back(std::move(result));
 
       // (synchronized): Publish the compilation result and get the next unit.
       {
         BackgroundCompileScope compile_scope(token_);
         if (compile_scope.cancelled()) return;
-        if (!result.succeeded()) {
+        if (!results_to_publish.back().succeeded()) {
           // Compile error.
           compile_scope.compilation_state()->SetError();
           compile_scope.compilation_state()->OnBackgroundTaskStopped(
@@ -714,14 +742,13 @@ class BackgroundCompileTask : public CancelableTask {
           compilation_failed = true;
           break;
         }
-        WasmCode* code =
-            compile_scope.native_module()->AddCompiledCode(std::move(result));
-        DCHECK_NOT_NULL(code);
+        // Publish TurboFan units immediately to reduce peak memory consumption.
+        if (result.requested_tier == ExecutionTier::kOptimized) {
+          publish_results(&compile_scope);
+        }
 
-        // Successfully finished one unit.
-        compile_scope.compilation_state()->OnFinishedUnit(result.requested_tier,
-                                                          code);
         if (deadline < MonotonicallyIncreasingTimeInMs()) {
+          publish_results(&compile_scope);
           compile_scope.compilation_state()->ReportDetectedFeatures(
               detected_features);
           compile_scope.compilation_state()->RestartBackgroundCompileTask();
@@ -731,6 +758,7 @@ class BackgroundCompileTask : public CancelableTask {
         // Get next unit.
         unit = compile_scope.compilation_state()->GetNextCompilationUnit();
         if (unit == nullptr) {
+          publish_results(&compile_scope);
           compile_scope.compilation_state()->OnBackgroundTaskStopped(
               detected_features);
           return;

src/wasm/wasm-code-manager.cc

@@ -393,8 +393,7 @@ NativeModule::NativeModule(WasmEngine* engine, const WasmFeatures& enabled,
 
   uint32_t num_wasm_functions = module_->num_declared_functions;
   if (num_wasm_functions > 0) {
-    code_table_.reset(new WasmCode*[num_wasm_functions]);
-    memset(code_table_.get(), 0, num_wasm_functions * sizeof(WasmCode*));
+    code_table_.reset(new WasmCode* [num_wasm_functions] {});
 
     jump_table_ = CreateEmptyJumpTable(
         JumpTableAssembler::SizeForNumberOfSlots(num_wasm_functions));
@@ -403,8 +402,7 @@ NativeModule::NativeModule(WasmEngine* engine, const WasmFeatures& enabled,
 
 void NativeModule::ReserveCodeTableForTesting(uint32_t max_functions) {
   DCHECK_LE(num_functions(), max_functions);
-  WasmCode** new_table = new WasmCode*[max_functions];
-  memset(new_table, 0, max_functions * sizeof(*new_table));
+  WasmCode** new_table = new WasmCode* [max_functions] {};
   if (module_->num_declared_functions > 0) {
     memcpy(new_table, code_table_.get(),
            module_->num_declared_functions * sizeof(*new_table));
@@ -534,10 +532,12 @@ WasmCode* NativeModule::AddAndPublishAnonymousCode(Handle<Code> code,
   const size_t code_comments_offset =
       static_cast<size_t>(code->code_comments_offset());
 
+  Vector<uint8_t> code_space = AllocateForCode(instructions.size());
+  memcpy(code_space.begin(), instructions.start(), instructions.size());
   std::unique_ptr<WasmCode> new_code{new WasmCode{
       this,                                     // native_module
       WasmCode::kAnonymousFuncIndex,            // index
-      AllocateForCode(instructions.size()),     // code_space
+      code_space,                               // instructions
       stack_slots,                              // stack_slots
       0,                                        // tagged_parameter_slots
       safepoint_table_offset,                   // safepoint_table_offset
@@ -551,8 +551,6 @@ WasmCode* NativeModule::AddAndPublishAnonymousCode(Handle<Code> code,
       kind,                                     // kind
       WasmCode::kOther}};                       // tier
 
-  memcpy(new_code->instructions().start(), instructions.start(),
-         instructions.size());
   // Apply the relocation delta by iterating over the RelocInfo.
   intptr_t delta = new_code->instruction_start() - code->InstructionStart();
   int mode_mask = RelocInfo::kApplyMask |
@@ -588,6 +586,18 @@ std::unique_ptr<WasmCode> NativeModule::AddCode(
     OwnedVector<trap_handler::ProtectedInstructionData> protected_instructions,
     OwnedVector<const byte> source_position_table, WasmCode::Kind kind,
     WasmCode::Tier tier) {
+  return AddCodeWithCodeSpace(index, desc, stack_slots, tagged_parameter_slots,
+                              std::move(protected_instructions),
+                              std::move(source_position_table), kind, tier,
+                              AllocateForCode(desc.instr_size));
+}
+
+std::unique_ptr<WasmCode> NativeModule::AddCodeWithCodeSpace(
+    uint32_t index, const CodeDesc& desc, uint32_t stack_slots,
+    uint32_t tagged_parameter_slots,
+    OwnedVector<ProtectedInstructionData> protected_instructions,
+    OwnedVector<const byte> source_position_table, WasmCode::Kind kind,
+    WasmCode::Tier tier, Vector<uint8_t> code_space) {
   OwnedVector<byte> reloc_info;
   if (desc.reloc_size > 0) {
     reloc_info = OwnedVector<byte>::New(desc.reloc_size);
@@ -608,15 +618,13 @@ std::unique_ptr<WasmCode> NativeModule::AddCode(
       static_cast<size_t>(desc.code_comments_offset);
   const size_t instr_size = static_cast<size_t>(desc.instr_size);
 
-  std::unique_ptr<WasmCode> code{new WasmCode{
-      this, index, AllocateForCode(desc.instr_size), stack_slots,
-      tagged_parameter_slots, safepoint_table_offset, handler_table_offset,
-      constant_pool_offset, code_comments_offset, instr_size,
-      std::move(protected_instructions), std::move(reloc_info),
-      std::move(source_position_table), kind, tier}};
+  memcpy(code_space.begin(), desc.buffer, static_cast<size_t>(desc.instr_size));
 
-  memcpy(code->instructions().start(), desc.buffer,
-         static_cast<size_t>(desc.instr_size));
+  std::unique_ptr<WasmCode> code{new WasmCode{
+      this, index, code_space, stack_slots, tagged_parameter_slots,
+      safepoint_table_offset, handler_table_offset, constant_pool_offset,
+      code_comments_offset, instr_size, std::move(protected_instructions),
+      std::move(reloc_info), std::move(source_position_table), kind, tier}};
 
   // Apply the relocation delta by iterating over the RelocInfo.
   intptr_t delta = code->instructions().start() - desc.buffer;
@@ -655,6 +663,12 @@ std::unique_ptr<WasmCode> NativeModule::AddCode(
 
 WasmCode* NativeModule::PublishCode(std::unique_ptr<WasmCode> code) {
   base::MutexGuard lock(&allocation_mutex_);
+  return PublishCodeLocked(std::move(code));
+}
+
+WasmCode* NativeModule::PublishCodeLocked(std::unique_ptr<WasmCode> code) {
+  // The caller must hold the {allocation_mutex_}, thus we fail to lock it here.
+  DCHECK(!allocation_mutex_.TryLock());
   // Skip publishing code if there is an active redirection to the interpreter
   // for the given function index, in order to preserve the redirection.
   if (!code->IsAnonymous() && !has_interpreter_redirection(code->index())) {
@@ -691,16 +705,16 @@ WasmCode* NativeModule::AddDeserializedCode(
     OwnedVector<const byte> reloc_info,
     OwnedVector<const byte> source_position_table, WasmCode::Kind kind,
     WasmCode::Tier tier) {
+  Vector<uint8_t> code_space = AllocateForCode(instructions.size());
+  memcpy(code_space.begin(), instructions.start(), instructions.size());
+
   std::unique_ptr<WasmCode> code{new WasmCode{
-      this, index, AllocateForCode(instructions.size()), stack_slots,
-      tagged_parameter_slots, safepoint_table_offset, handler_table_offset,
-      constant_pool_offset, code_comments_offset, unpadded_binary_size,
+      this, index, code_space, stack_slots, tagged_parameter_slots,
+      safepoint_table_offset, handler_table_offset, constant_pool_offset,
+      code_comments_offset, unpadded_binary_size,
       std::move(protected_instructions), std::move(reloc_info),
       std::move(source_position_table), kind, tier}};
 
-  memcpy(code->instructions().start(), instructions.start(),
-         instructions.size());
-
   code->RegisterTrapHandlerData();
 
   // Note: we do not flush the i-cache here, since the code needs to be
@@ -720,10 +734,12 @@ std::vector<WasmCode*> NativeModule::SnapshotCodeTable() const {
 WasmCode* NativeModule::CreateEmptyJumpTable(uint32_t jump_table_size) {
   // Only call this if we really need a jump table.
   DCHECK_LT(0, jump_table_size);
+  Vector<uint8_t> code_space = AllocateForCode(jump_table_size);
+  memset(code_space.begin(), 0, code_space.size());
   std::unique_ptr<WasmCode> code{new WasmCode{
       this,                                     // native_module
       WasmCode::kAnonymousFuncIndex,            // index
-      AllocateForCode(jump_table_size),         // instructions
+      code_space,                               // instructions
       0,                                        // stack_slots
       0,                                        // tagged_parameter_slots
       0,                                        // safepoint_table_offset
@@ -736,7 +752,6 @@ WasmCode* NativeModule::CreateEmptyJumpTable(uint32_t jump_table_size) {
       OwnedVector<const uint8_t>{},             // source_pos
       WasmCode::kJumpTable,                     // kind
       WasmCode::kOther}};                       // tier
-  memset(code->instructions().start(), 0, code->instructions().size());
   return PublishCode(std::move(code));
 }
 
@@ -1205,17 +1220,49 @@ WasmCode::Kind GetCodeKindForExecutionTier(ExecutionTier tier) {
 }  // namespace
 
 WasmCode* NativeModule::AddCompiledCode(WasmCompilationResult result) {
-  DCHECK(result.succeeded());
+  return AddCompiledCode({&result, 1})[0];
+}
 
-  DCHECK_EQ(result.code_desc.buffer, result.instr_buffer.get());
-  std::unique_ptr<WasmCode> code = AddCode(
-      result.func_index, result.code_desc, result.frame_slot_count,
-      result.tagged_parameter_slots, std::move(result.protected_instructions),
-      std::move(result.source_positions),
-      GetCodeKindForExecutionTier(result.result_tier),
-      GetCodeTierForExecutionTier(result.result_tier));
+std::vector<WasmCode*> NativeModule::AddCompiledCode(
+    Vector<WasmCompilationResult> results) {
+  DCHECK(!results.is_empty());
+  // First, allocate code space for all the results.
+  size_t total_code_space = 0;
+  for (auto& result : results) {
+    DCHECK(result.succeeded());
+    total_code_space += RoundUp<kCodeAlignment>(result.code_desc.instr_size);
+  }
+  Vector<byte> code_space = AllocateForCode(total_code_space);
+
+  std::vector<std::unique_ptr<WasmCode>> generated_code;
+  generated_code.reserve(results.size());
+
+  // Now copy the generated code into the code space and relocate it.
+  for (auto& result : results) {
+    DCHECK_EQ(result.code_desc.buffer, result.instr_buffer.get());
+    size_t code_size = RoundUp<kCodeAlignment>(result.code_desc.instr_size);
+    Vector<byte> this_code_space = code_space.SubVector(0, code_size);
+    code_space += code_size;
+    generated_code.emplace_back(AddCodeWithCodeSpace(
+        result.func_index, result.code_desc, result.frame_slot_count,
+        result.tagged_parameter_slots, std::move(result.protected_instructions),
+        std::move(result.source_positions),
+        GetCodeKindForExecutionTier(result.result_tier),
+        GetCodeTierForExecutionTier(result.result_tier), this_code_space));
+  }
+  DCHECK_EQ(0, code_space.size());
+
+  // Under the {allocation_mutex_}, publish the code.
+  std::vector<WasmCode*> returned_code;
+  returned_code.reserve(results.size());
+  {
+    base::MutexGuard lock(&allocation_mutex_);
+    for (auto& result : generated_code) {
+      returned_code.push_back(PublishCodeLocked(std::move(result)));
+    }
+  }
 
-  return PublishCode(std::move(code));
+  return returned_code;
 }
 
 void WasmCodeManager::FreeNativeModule(NativeModule* native_module) {

src/wasm/wasm-code-manager.h

@@ -245,6 +245,8 @@ class V8_EXPORT_PRIVATE NativeModule final {
   // the code table and patching the jump table. It returns a raw pointer to the
   // given {WasmCode} object.
   WasmCode* PublishCode(std::unique_ptr<WasmCode>);
+  // Hold the {allocation_mutex_} when calling {PublishCodeLocked}.
+  WasmCode* PublishCodeLocked(std::unique_ptr<WasmCode>);
 
   WasmCode* AddDeserializedCode(
       uint32_t index, Vector<const byte> instructions, uint32_t stack_slots,
@@ -367,6 +369,7 @@ class V8_EXPORT_PRIVATE NativeModule final {
   void SampleCodeSize(Counters*, CodeSamplingTime) const;
 
   WasmCode* AddCompiledCode(WasmCompilationResult);
+  std::vector<WasmCode*> AddCompiledCode(Vector<WasmCompilationResult>);
 
  private:
   friend class WasmCode;
@@ -380,6 +383,14 @@ class V8_EXPORT_PRIVATE NativeModule final {
                std::shared_ptr<Counters> async_counters,
                std::shared_ptr<NativeModule>* shared_this);
 
+  std::unique_ptr<WasmCode> AddCodeWithCodeSpace(
+      uint32_t index, const CodeDesc& desc, uint32_t stack_slots,
+      uint32_t tagged_parameter_slots,
+      OwnedVector<trap_handler::ProtectedInstructionData>
+          protected_instructions,
+      OwnedVector<const byte> source_position_table, WasmCode::Kind kind,
+      WasmCode::Tier tier, Vector<uint8_t> code_space);
+
   // Add and publish anonymous code.
   WasmCode* AddAndPublishAnonymousCode(Handle<Code>, WasmCode::Kind kind,
                                        const char* name = nullptr);