[wasm] Improve code size estimate for streaming

In streaming compilation, we were computing a way too low code size
estimate, since all {WasmFunction::code} fields were still zero when we
were calling {EstimateNativeModuleCodeSize}. This lead to many separate
code spaces being created during compilation, creating significant
performance and memory overhead.

This CL fixes this by passing the code section length when creating the
{NativeModule}. From this, we can compute the code size estimate just as
before.

Drive-by: Rename "functions_count" to "num_functions" in
{ProcessCodeSectionHeader} to be consistent with the declaration.

R=ahaas@chromium.org

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

src/wasm/module-compiler.cc

@@ -1426,8 +1426,9 @@ class AsyncStreamingProcessor final : public StreamingProcessor {
   bool ProcessSection(SectionCode section_code, Vector<const uint8_t> bytes,
                       uint32_t offset) override;
 
-  bool ProcessCodeSectionHeader(int functions_count, uint32_t offset,
-                                std::shared_ptr<WireBytesStorage>) override;
+  bool ProcessCodeSectionHeader(int num_functions, uint32_t offset,
+                                std::shared_ptr<WireBytesStorage>,
+                                int code_section_length) override;
 
   bool ProcessFunctionBody(Vector<const uint8_t> bytes,
                            uint32_t offset) override;
@@ -1484,7 +1485,7 @@ AsyncCompileJob::~AsyncCompileJob() {
 }
 
 void AsyncCompileJob::CreateNativeModule(
-    std::shared_ptr<const WasmModule> module) {
+    std::shared_ptr<const WasmModule> module, size_t code_size_estimate) {
   // Embedder usage count for declared shared memories.
   if (module->has_shared_memory) {
     isolate_->CountUsage(v8::Isolate::UseCounterFeature::kWasmSharedMemory);
@@ -1500,8 +1501,6 @@ void AsyncCompileJob::CreateNativeModule(
   // breakpoints on a (potentially empty) subset of the instances.
   // Create the module object.
 
-  size_t code_size_estimate =
-      wasm::WasmCodeManager::EstimateNativeModuleCodeSize(module.get());
   native_module_ = isolate_->wasm_engine()->NewNativeModule(
       isolate_, enabled_features_, std::move(module), code_size_estimate);
   native_module_->SetWireBytes({std::move(bytes_copy_), wire_bytes_.length()});
@@ -1823,7 +1822,11 @@ class AsyncCompileJob::DecodeModule : public AsyncCompileJob::CompileStep {
       job->DoSync<DecodeFail>(std::move(result).error());
     } else {
       // Decode passed.
-      job->DoSync<PrepareAndStartCompile>(std::move(result).value(), true);
+      std::shared_ptr<WasmModule> module = std::move(result).value();
+      size_t code_size_estimate =
+          wasm::WasmCodeManager::EstimateNativeModuleCodeSize(module.get());
+      job->DoSync<PrepareAndStartCompile>(std::move(module), true,
+                                          code_size_estimate);
     }
   }
 
@@ -1854,12 +1857,15 @@ class AsyncCompileJob::DecodeFail : public CompileStep {
 class AsyncCompileJob::PrepareAndStartCompile : public CompileStep {
  public:
   PrepareAndStartCompile(std::shared_ptr<const WasmModule> module,
-                         bool start_compilation)
-      : module_(std::move(module)), start_compilation_(start_compilation) {}
+                         bool start_compilation, size_t code_size_estimate)
+      : module_(std::move(module)),
+        start_compilation_(start_compilation),
+        code_size_estimate_(code_size_estimate) {}
 
  private:
-  std::shared_ptr<const WasmModule> module_;
-  bool start_compilation_;
+  const std::shared_ptr<const WasmModule> module_;
+  const bool start_compilation_;
+  const size_t code_size_estimate_;
 
   void RunInForeground(AsyncCompileJob* job) override {
     TRACE_COMPILE("(2) Prepare and start compile...\n");
@@ -1868,7 +1874,7 @@ class AsyncCompileJob::PrepareAndStartCompile : public CompileStep {
     // is done.
     job->background_task_manager_.CancelAndWait();
 
-    job->CreateNativeModule(module_);
+    job->CreateNativeModule(module_, code_size_estimate_);
 
     CompilationStateImpl* compilation_state =
         Impl(job->native_module_->compilation_state());
@@ -2033,19 +2039,25 @@ bool AsyncStreamingProcessor::ProcessSection(SectionCode section_code,
 
 // Start the code section.
 bool AsyncStreamingProcessor::ProcessCodeSectionHeader(
-    int functions_count, uint32_t offset,
-    std::shared_ptr<WireBytesStorage> wire_bytes_storage) {
+    int num_functions, uint32_t offset,
+    std::shared_ptr<WireBytesStorage> wire_bytes_storage,
+    int code_section_length) {
   TRACE_STREAMING("Start the code section with %d functions...\n",
-                  functions_count);
-  if (!decoder_.CheckFunctionsCount(static_cast<uint32_t>(functions_count),
+                  num_functions);
+  if (!decoder_.CheckFunctionsCount(static_cast<uint32_t>(num_functions),
                                     offset)) {
     FinishAsyncCompileJobWithError(decoder_.FinishDecoding(false).error());
     return false;
   }
   // Execute the PrepareAndStartCompile step immediately and not in a separate
   // task.
+  int num_imported_functions =
+      static_cast<int>(decoder_.module()->num_imported_functions);
+  size_t code_size_estimate =
+      wasm::WasmCodeManager::EstimateNativeModuleCodeSize(
+          num_functions, num_imported_functions, code_section_length);
   job_->DoImmediately<AsyncCompileJob::PrepareAndStartCompile>(
-      decoder_.shared_module(), false);
+      decoder_.shared_module(), false, code_size_estimate);
   auto* compilation_state = Impl(job_->native_module_->compilation_state());
   compilation_state->SetWireBytesStorage(std::move(wire_bytes_storage));
   DCHECK_EQ(job_->native_module_->module()->origin, kWasmOrigin);
@@ -2154,7 +2166,8 @@ void AsyncStreamingProcessor::OnFinishedStream(OwnedVector<uint8_t> bytes) {
   if (job_->native_module_ == nullptr) {
     // We are processing a WebAssembly module without code section. Create the
     // runtime objects now (would otherwise happen in {PrepareAndStartCompile}).
-    job_->CreateNativeModule(std::move(result).value());
+    constexpr size_t kCodeSizeEstimate = 0;
+    job_->CreateNativeModule(std::move(result).value(), kCodeSizeEstimate);
     DCHECK(needs_finish);
   }
   job_->wire_bytes_ = ModuleWireBytes(bytes.as_vector());

src/wasm/module-compiler.h

@@ -143,7 +143,8 @@ class AsyncCompileJob {
     return outstanding_finishers_.fetch_sub(1) == 1;
   }
 
-  void CreateNativeModule(std::shared_ptr<const WasmModule> module);
+  void CreateNativeModule(std::shared_ptr<const WasmModule> module,
+                          size_t code_size_estimate);
   void PrepareRuntimeObjects();
 
   void FinishCompile();

src/wasm/streaming-decoder.cc

@@ -437,9 +437,12 @@ StreamingDecoder::DecodeNumberOfFunctions::NextWithValue(
     return std::make_unique<DecodeSectionID>(streaming->module_offset());
   }
 
+  DCHECK_GE(kMaxInt, payload_buf.length());
+  int code_section_len = static_cast<int>(payload_buf.length());
   DCHECK_GE(kMaxInt, value_);
   streaming->StartCodeSection(static_cast<int>(value_),
-                              streaming->section_buffers_.back());
+                              streaming->section_buffers_.back(),
+                              code_section_len);
   if (!streaming->ok()) return nullptr;
   return std::make_unique<DecodeFunctionLength>(
       section_buffer_, section_buffer_->payload_offset() + bytes_consumed_,

src/wasm/streaming-decoder.h

@@ -37,7 +37,8 @@ class V8_EXPORT_PRIVATE StreamingProcessor {
   // Process the start of the code section. Returns true if the processing
   // finished successfully and the decoding should continue.
   virtual bool ProcessCodeSectionHeader(int num_functions, uint32_t offset,
-                                        std::shared_ptr<WireBytesStorage>) = 0;
+                                        std::shared_ptr<WireBytesStorage>,
+                                        int code_section_length) = 0;
 
   // Process a function body. Returns true if the processing finished
   // successfully and the decoding should continue.
@@ -228,13 +229,14 @@ class V8_EXPORT_PRIVATE StreamingDecoder {
   }
 
   void StartCodeSection(int num_functions,
-                        std::shared_ptr<WireBytesStorage> wire_bytes_storage) {
+                        std::shared_ptr<WireBytesStorage> wire_bytes_storage,
+                        int code_section_length) {
     if (!ok()) return;
     // The offset passed to {ProcessCodeSectionHeader} is an error offset and
     // not the start offset of a buffer. Therefore we need the -1 here.
-    if (!processor_->ProcessCodeSectionHeader(num_functions,
-                                              module_offset() - 1,
-                                              std::move(wire_bytes_storage))) {
+    if (!processor_->ProcessCodeSectionHeader(
+            num_functions, module_offset() - 1, std::move(wire_bytes_storage),
+            code_section_length)) {
       Fail();
     }
   }

src/wasm/wasm-code-manager.cc

@@ -1540,18 +1540,33 @@ VirtualMemory WasmCodeManager::TryAllocate(size_t size, void* hint) {
 
 // static
 size_t WasmCodeManager::EstimateNativeModuleCodeSize(const WasmModule* module) {
+  int num_functions = static_cast<int>(module->num_declared_functions);
+  int num_imported_functions = static_cast<int>(module->num_imported_functions);
+  int code_section_length = 0;
+  if (num_functions > 0) {
+    DCHECK_EQ(module->functions.size(), num_imported_functions + num_functions);
+    auto* first_fn = &module->functions[module->num_imported_functions];
+    auto* last_fn = &module->functions.back();
+    code_section_length =
+        static_cast<int>(last_fn->code.end_offset() - first_fn->code.offset());
+  }
+  return EstimateNativeModuleCodeSize(num_functions, num_imported_functions,
+                                      code_section_length);
+}
+
+// static
+size_t WasmCodeManager::EstimateNativeModuleCodeSize(int num_functions,
+                                                     int num_imported_functions,
+                                                     int code_section_length) {
   constexpr size_t kCodeSizeMultiplier = 4;
   constexpr size_t kCodeOverhead = 32;     // for prologue, stack check, ...
   constexpr size_t kStaticCodeSize = 512;  // runtime stubs, ...
   constexpr size_t kImportSize = 64 * kSystemPointerSize;
 
-  size_t estimate = kStaticCodeSize;
-  for (auto& function : module->functions) {
-    estimate += kCodeOverhead + kCodeSizeMultiplier * function.code.length();
-  }
-  estimate += kImportSize * module->num_imported_functions;
-
-  return estimate;
+  return kStaticCodeSize                              // static
+         + kCodeOverhead * num_functions              // per function
+         + kCodeSizeMultiplier * code_section_length  // per opcode (~ byte)
+         + kImportSize * num_imported_functions;      // per import
 }
 
 // static

src/wasm/wasm-code-manager.h

@@ -694,6 +694,11 @@ class V8_EXPORT_PRIVATE WasmCodeManager final {
 
   // Estimate the needed code space from a completely decoded module.
   static size_t EstimateNativeModuleCodeSize(const WasmModule* module);
+  // Estimate the needed code space from the number of functions and total code
+  // section length.
+  static size_t EstimateNativeModuleCodeSize(int num_functions,
+                                             int num_imported_functions,
+                                             int code_section_length);
   // Estimate the size of meta data needed for the NativeModule, excluding
   // generated code. This data still be stored on the C++ heap.
   static size_t EstimateNativeModuleMetaDataSize(const WasmModule* module);

test/unittests/wasm/streaming-decoder-unittest.cc

@@ -57,7 +57,8 @@ class MockStreamingProcessor : public StreamingProcessor {
   }
 
   bool ProcessCodeSectionHeader(int num_functions, uint32_t offset,
-                                std::shared_ptr<WireBytesStorage>) override {
+                                std::shared_ptr<WireBytesStorage>,
+                                int code_section_length) override {
     return true;
   }