[wasm] Introduce AssemblerBufferCache

This introduces an AssemblerBufferCache class which will (later) cache
the backing store of AssemblerBuffers. This is needed for PKU-protected
assembler buffers, which are expensive to allocate and deallocate.

For now, the AssemblerBufferCache does not do any caching, this will be
added in a follow-up CL.

R=thibaudm@chromium.org

Bug: v8:12809
Change-Id: I4a7ccff49c9930584a9fcda8899cfe38cfc61419
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/3593136Reviewed-by: Thibaud Michaud <thibaudm@chromium.org>
Commit-Queue: Clemens Backes <clemensb@chromium.org>
Cr-Commit-Position: refs/heads/main@{#80145}

BUILD.bazel

@@ -2435,6 +2435,8 @@ filegroup(
             "src/debug/debug-wasm-objects-inl.h",
             "src/runtime/runtime-test-wasm.cc",
             "src/runtime/runtime-wasm.cc",
+            "src/wasm/assembler-buffer-cache.cc",
+            "src/wasm/assembler-buffer-cache.h",
             "src/wasm/baseline/liftoff-assembler.cc",
             "src/wasm/baseline/liftoff-assembler-defs.h",
             "src/wasm/baseline/liftoff-assembler.h",

BUILD.gn

@@ -3521,6 +3521,7 @@ v8_header_set("v8_internal_headers") {
       "src/debug/debug-wasm-objects.h",
       "src/trap-handler/trap-handler-internal.h",
       "src/trap-handler/trap-handler.h",
+      "src/wasm/assembler-buffer-cache.h",
       "src/wasm/baseline/liftoff-assembler-defs.h",
       "src/wasm/baseline/liftoff-assembler.h",
       "src/wasm/baseline/liftoff-compiler.h",
@@ -4583,6 +4584,7 @@ v8_source_set("v8_base_without_compiler") {
       "src/trap-handler/handler-inside.cc",
       "src/trap-handler/handler-outside.cc",
       "src/trap-handler/handler-shared.cc",
+      "src/wasm/assembler-buffer-cache.cc",
       "src/wasm/baseline/liftoff-assembler.cc",
       "src/wasm/baseline/liftoff-compiler.cc",
       "src/wasm/canonical-types.cc",

src/wasm/assembler-buffer-cache.cc

+// Copyright 2022 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.
+
+#include "src/wasm/assembler-buffer-cache.h"
+
+#include "src/codegen/assembler.h"
+
+namespace v8::internal::wasm {
+
+std::unique_ptr<AssemblerBuffer> AssemblerBufferCache::GetAssemblerBuffer(
+    int size) {
+  // TODO(12809): Return PKU-protected buffers, and cache them.
+  return NewAssemblerBuffer(size);
+}
+
+}  // namespace v8::internal::wasm

src/wasm/assembler-buffer-cache.h

+// Copyright 2022 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.
+
+#if !V8_ENABLE_WEBASSEMBLY
+#error This header should only be included if WebAssembly is enabled.
+#endif  // !V8_ENABLE_WEBASSEMBLY
+
+#ifndef V8_WASM_ASSEMBLER_BUFFER_CACHE_H_
+#define V8_WASM_ASSEMBLER_BUFFER_CACHE_H_
+
+#include <memory>
+
+namespace v8::internal {
+class AssemblerBuffer;
+}
+
+namespace v8::internal::wasm {
+
+// Creating assembler buffers can be expensive, in particular if PKU is used,
+// which requires an {mmap} and {pkey_protect} system call for each new buffer.
+// Hence pool-allocate a larger memory region and reuse it if assembler buffers
+// are freed.
+// For now, this class only implements the interface without actually caching
+// anything.
+// TODO(12809): Actually cache the assembler buffers.
+class AssemblerBufferCache final {
+ public:
+  std::unique_ptr<AssemblerBuffer> GetAssemblerBuffer(int size);
+};
+
+}  // namespace v8::internal::wasm
+
+#endif  // V8_WASM_ASSEMBLER_BUFFER_CACHE_H_

src/wasm/baseline/liftoff-compiler.cc

@@ -21,6 +21,7 @@
 #include "src/tracing/trace-event.h"
 #include "src/utils/ostreams.h"
 #include "src/utils/utils.h"
+#include "src/wasm/assembler-buffer-cache.h"
 #include "src/wasm/baseline/liftoff-assembler.h"
 #include "src/wasm/baseline/liftoff-register.h"
 #include "src/wasm/function-body-decoder-impl.h"
@@ -6585,6 +6586,21 @@ constexpr WasmOpcode LiftoffCompiler::kNoOutstandingOp;
 // static
 constexpr base::EnumSet<ValueKind> LiftoffCompiler::kUnconditionallySupported;
 
+std::unique_ptr<AssemblerBuffer> NewLiftoffAssemblerBuffer(
+    AssemblerBufferCache* assembler_buffer_cache, int func_body_size) {
+  size_t code_size_estimate =
+      WasmCodeManager::EstimateLiftoffCodeSize(func_body_size);
+  // Allocate the initial buffer a bit bigger to avoid reallocation during code
+  // generation. Overflows when casting to int are fine, as we will allocate at
+  // least {AssemblerBase::kMinimalBufferSize} anyway, so in the worst case we
+  // have to grow more often.
+  int initial_buffer_size = static_cast<int>(128 + code_size_estimate * 4 / 3);
+
+  return assembler_buffer_cache
+             ? assembler_buffer_cache->GetAssemblerBuffer(initial_buffer_size)
+             : NewAssemblerBuffer(initial_buffer_size);
+}
+
 }  // namespace
 
 WasmCompilationResult ExecuteLiftoffCompilation(
@@ -6602,13 +6618,7 @@ WasmCompilationResult ExecuteLiftoffCompilation(
 
   Zone zone(GetWasmEngine()->allocator(), "LiftoffCompilationZone");
   auto call_descriptor = compiler::GetWasmCallDescriptor(&zone, func_body.sig);
-  size_t code_size_estimate =
-      WasmCodeManager::EstimateLiftoffCodeSize(func_body_size);
-  // Allocate the initial buffer a bit bigger to avoid reallocation during code
-  // generation. Overflows when casting to int are fine, as we will allocate at
-  // least {AssemblerBase::kMinimalBufferSize} anyway, so in the worst case we
-  // have to grow more often.
-  int initial_buffer_size = static_cast<int>(128 + code_size_estimate * 4 / 3);
+
   std::unique_ptr<DebugSideTableBuilder> debug_sidetable_builder;
   if (compiler_options.debug_sidetable) {
     debug_sidetable_builder = std::make_unique<DebugSideTableBuilder>();
@@ -6616,13 +6626,15 @@ WasmCompilationResult ExecuteLiftoffCompilation(
   DCHECK_IMPLIES(compiler_options.max_steps,
                  compiler_options.for_debugging == kForDebugging);
   WasmFeatures unused_detected_features;
+
   WasmFullDecoder<Decoder::kBooleanValidation, LiftoffCompiler> decoder(
       &zone, env->module, env->enabled_features,
       compiler_options.detected_features ? compiler_options.detected_features
                                          : &unused_detected_features,
       func_body, call_descriptor, env, &zone,
-      NewAssemblerBuffer(initial_buffer_size), debug_sidetable_builder.get(),
-      compiler_options);
+      NewLiftoffAssemblerBuffer(compiler_options.assembler_buffer_cache,
+                                func_body_size),
+      debug_sidetable_builder.get(), compiler_options);
   decoder.Decode();
   LiftoffCompiler* compiler = &decoder.interface();
   if (decoder.failed()) compiler->OnFirstError(&decoder);

src/wasm/baseline/liftoff-compiler.h

@@ -57,6 +57,7 @@ struct LiftoffOptions {
   int func_index = -1;
   ForDebugging for_debugging = kNoDebugging;
   Counters* counters = nullptr;
+  AssemblerBufferCache* assembler_buffer_cache = nullptr;
   WasmFeatures* detected_features = nullptr;
   base::Vector<const int> breakpoints = {};
   std::unique_ptr<DebugSideTable>* debug_sidetable = nullptr;
@@ -77,6 +78,7 @@ struct LiftoffOptions {
   SETTER(func_index)
   SETTER(for_debugging)
   SETTER(counters)
+  SETTER(assembler_buffer_cache)
   SETTER(detected_features)
   SETTER(breakpoints)
   SETTER(debug_sidetable)

src/wasm/function-compiler.cc

@@ -34,13 +34,14 @@ ExecutionTier WasmCompilationUnit::GetBaselineExecutionTier(
 
 WasmCompilationResult WasmCompilationUnit::ExecuteCompilation(
     CompilationEnv* env, const WireBytesStorage* wire_bytes_storage,
-    Counters* counters, WasmFeatures* detected) {
+    Counters* counters, AssemblerBufferCache* buffer_cache,
+    WasmFeatures* detected) {
   WasmCompilationResult result;
   if (func_index_ < static_cast<int>(env->module->num_imported_functions)) {
     result = ExecuteImportWrapperCompilation(env);
   } else {
-    result =
-        ExecuteFunctionCompilation(env, wire_bytes_storage, counters, detected);
+    result = ExecuteFunctionCompilation(env, wire_bytes_storage, counters,
+                                        buffer_cache, detected);
   }
 
   if (result.succeeded() && counters) {
@@ -70,7 +71,8 @@ WasmCompilationResult WasmCompilationUnit::ExecuteImportWrapperCompilation(
 
 WasmCompilationResult WasmCompilationUnit::ExecuteFunctionCompilation(
     CompilationEnv* env, const WireBytesStorage* wire_bytes_storage,
-    Counters* counters, WasmFeatures* detected) {
+    Counters* counters, AssemblerBufferCache* buffer_cache,
+    WasmFeatures* detected) {
   auto* func = &env->module->functions[func_index_];
   base::Vector<const uint8_t> code = wire_bytes_storage->GetCode(func->code);
   wasm::FunctionBody func_body{func->sig, func->code.offset(), code.begin(),
@@ -126,6 +128,7 @@ WasmCompilationResult WasmCompilationUnit::ExecuteFunctionCompilation(
                 .set_for_debugging(for_debugging_)
                 .set_counters(counters)
                 .set_detected_features(detected)
+                .set_assembler_buffer_cache(buffer_cache)
                 .set_debug_sidetable(debug_sidetable_ptr));
         if (result.succeeded()) break;
       }
@@ -134,7 +137,7 @@ WasmCompilationResult WasmCompilationUnit::ExecuteFunctionCompilation(
       // failed.
       if (FLAG_liftoff_only) break;
 
-      // If Liftoff failed, fall back to turbofan.
+      // If Liftoff failed, fall back to TurboFan.
       // TODO(wasm): We could actually stop or remove the tiering unit for this
       // function to avoid compiling it twice with TurboFan.
       V8_FALLTHROUGH;
@@ -166,7 +169,7 @@ void WasmCompilationUnit::CompileWasmFunction(Isolate* isolate,
   CompilationEnv env = native_module->CreateCompilationEnv();
   WasmCompilationResult result = unit.ExecuteCompilation(
       &env, native_module->compilation_state()->GetWireBytesStorage().get(),
-      isolate->counters(), detected);
+      isolate->counters(), nullptr, detected);
   if (result.succeeded()) {
     WasmCodeRefScope code_ref_scope;
     native_module->PublishCode(

src/wasm/function-compiler.h

@@ -28,6 +28,7 @@ class TurbofanCompilationJob;
 
 namespace wasm {
 
+class AssemblerBufferCache;
 class NativeModule;
 class WasmCode;
 class WasmEngine;
@@ -69,6 +70,7 @@ class V8_EXPORT_PRIVATE WasmCompilationUnit final {
 
   WasmCompilationResult ExecuteCompilation(CompilationEnv*,
                                            const WireBytesStorage*, Counters*,
+                                           AssemblerBufferCache*,
                                            WasmFeatures* detected);
 
   ExecutionTier tier() const { return tier_; }
@@ -83,6 +85,7 @@ class V8_EXPORT_PRIVATE WasmCompilationUnit final {
   WasmCompilationResult ExecuteFunctionCompilation(CompilationEnv*,
                                                    const WireBytesStorage*,
                                                    Counters*,
+                                                   AssemblerBufferCache*,
                                                    WasmFeatures* detected);
 
   WasmCompilationResult ExecuteImportWrapperCompilation(CompilationEnv*);

src/wasm/module-compiler.cc

@@ -25,6 +25,7 @@
 #include "src/tracing/trace-event.h"
 #include "src/trap-handler/trap-handler.h"
 #include "src/utils/identity-map.h"
+#include "src/wasm/assembler-buffer-cache.h"
 #include "src/wasm/code-space-access.h"
 #include "src/wasm/module-decoder.h"
 #include "src/wasm/streaming-decoder.h"
@@ -1174,10 +1175,12 @@ bool CompileLazy(Isolate* isolate, Handle<WasmInstanceObject> instance,
                                     kNoDebugging};
   CompilationEnv env = native_module->CreateCompilationEnv();
   WasmEngine* engine = GetWasmEngine();
+  // TODO(wasm): Use an assembler buffer cache for lazy compilation.
+  AssemblerBufferCache* assembler_buffer_cache = nullptr;
   WasmFeatures detected_features;
   WasmCompilationResult result = baseline_unit.ExecuteCompilation(
       &env, compilation_state->GetWireBytesStorage().get(), counters,
-      &detected_features);
+      assembler_buffer_cache, &detected_features);
   compilation_state->OnCompilationStopped(detected_features);
   if (!thread_ticks.IsNull()) {
     native_module->UpdateCPUDuration(
@@ -1579,6 +1582,16 @@ CompilationExecutionResult ExecuteCompilationUnits(
   }
   TRACE_COMPILE("ExecuteCompilationUnits (task id %d)\n", task_id);
 
+  // If PKU is enabled, use an assembler buffer cache to avoid many expensive
+  // buffer allocations. Otherwise, malloc/free is efficient enough to prefer
+  // that bit of overhead over the memory consumption increase by the cache.
+  base::Optional<AssemblerBufferCache> optional_assembler_buffer_cache;
+  AssemblerBufferCache* assembler_buffer_cache = nullptr;
+  if (GetWasmCodeManager()->MemoryProtectionKeysEnabled()) {
+    optional_assembler_buffer_cache.emplace();
+    assembler_buffer_cache = &*optional_assembler_buffer_cache;
+  }
+
   std::vector<WasmCompilationResult> results_to_publish;
   while (true) {
     ExecutionTier current_tier = unit->tier();
@@ -1586,8 +1599,9 @@ CompilationExecutionResult ExecuteCompilationUnits(
     TRACE_EVENT0("v8.wasm", event_name);
     while (unit->tier() == current_tier) {
       // (asynchronous): Execute the compilation.
-      WasmCompilationResult result = unit->ExecuteCompilation(
-          &env.value(), wire_bytes.get(), counters, &detected_features);
+      WasmCompilationResult result =
+          unit->ExecuteCompilation(&env.value(), wire_bytes.get(), counters,
+                                   assembler_buffer_cache, &detected_features);
       results_to_publish.emplace_back(std::move(result));
 
       bool yield = delegate && delegate->ShouldYield();

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

@@ -593,7 +593,7 @@ void WasmFunctionCompiler::Build(const byte* start, const byte* end) {
                              for_debugging);
     result.emplace(unit.ExecuteCompilation(
         &env, native_module->compilation_state()->GetWireBytesStorage().get(),
-        nullptr, nullptr));
+        nullptr, nullptr, nullptr));
   }
   WasmCode* code = native_module->PublishCode(
       native_module->AddCompiledCode(std::move(*result)));