Add prefetching in arena blocks.

PiperOrigin-RevId: 604441990

src/google/protobuf/BUILD.bazel

@@ -319,7 +319,7 @@ cc_library(
     ],
     deps = [
         "@com_google_absl//absl/base:core_headers",
-        "@com_google_absl//absl/log:absl_log",
+        "@com_google_absl//absl/base:prefetch",
     ],
 )
 
@@ -388,6 +388,7 @@ cc_library(
         "@com_google_absl//absl/log:absl_log",
         "@com_google_absl//absl/numeric:bits",
         "@com_google_absl//absl/synchronization",
+        "@com_google_absl//absl/types:span",
         "@com_google_absl//absl/utility:if_constexpr",
     ],
 )

src/google/protobuf/arena.cc

@@ -13,17 +13,21 @@
 #include <cstdint>
 #include <limits>
 #include <string>
-#include <typeinfo>
 #include <vector>
 
 #include "absl/base/attributes.h"
+#include "absl/base/prefetch.h"
 #include "absl/container/internal/layout.h"
+#include "absl/log/absl_check.h"
+#include "absl/log/absl_log.h"
 #include "absl/synchronization/mutex.h"
+#include "absl/types/span.h"
 #include "google/protobuf/arena_allocation_policy.h"
 #include "google/protobuf/arena_cleanup.h"
 #include "google/protobuf/arenaz_sampler.h"
 #include "google/protobuf/port.h"
 #include "google/protobuf/serial_arena.h"
+#include "google/protobuf/string_block.h"
 #include "google/protobuf/thread_safe_arena.h"
 
 
@@ -299,6 +303,7 @@ size_t SerialArena::FreeStringBlocks(StringBlock* string_block,
                                      size_t unused_bytes) {
   ABSL_DCHECK(string_block != nullptr);
   StringBlock* next = string_block->next();
+  absl::PrefetchToLocalCacheNta(next);
   std::string* end = string_block->end();
   for (std::string* s = string_block->AtOffset(unused_bytes); s != end; ++s) {
     s->~basic_string();
@@ -307,6 +312,7 @@ size_t SerialArena::FreeStringBlocks(StringBlock* string_block,
 
   while ((string_block = next) != nullptr) {
     next = string_block->next();
+    absl::PrefetchToLocalCacheNta(next);
     for (std::string& s : *string_block) {
       s.~basic_string();
     }
@@ -324,6 +330,18 @@ void SerialArena::CleanupList() {
     char* limit = b->Limit();
     char* it = reinterpret_cast<char*>(b->cleanup_nodes);
     ABSL_DCHECK(!b->IsSentry() || it == limit);
+    // A prefetch distance of 8 here was chosen arbitrarily.
+    char* prefetch = it;
+    int prefetch_dist = 8;
+    for (; prefetch < limit && --prefetch_dist; prefetch += cleanup::Size()) {
+      cleanup::PrefetchNode(prefetch);
+    }
+    for (; prefetch < limit;
+         it += cleanup::Size(), prefetch += cleanup::Size()) {
+      cleanup::DestroyNode(it);
+      cleanup::PrefetchNode(prefetch);
+    }
+    absl::PrefetchToLocalCacheNta(b->next);
     for (; it < limit; it += cleanup::Size()) {
       cleanup::DestroyNode(it);
     }
@@ -780,6 +798,8 @@ void ThreadSafeArena::WalkConstSerialArenaChunk(Functor fn) const {
   const SerialArenaChunk* chunk = head_.load(std::memory_order_acquire);
 
   for (; !chunk->IsSentry(); chunk = chunk->next_chunk()) {
+    // Prefetch the next chunk.
+    absl::PrefetchToLocalCache(chunk->next_chunk());
     fn(chunk);
   }
 }
@@ -794,6 +814,8 @@ void ThreadSafeArena::WalkSerialArenaChunk(Functor fn) {
   while (!chunk->IsSentry()) {
     // Cache next chunk in case this chunk is destroyed.
     SerialArenaChunk* next_chunk = chunk->next_chunk();
+    // Prefetch the next chunk.
+    absl::PrefetchToLocalCache(next_chunk);
     fn(chunk);
     chunk = next_chunk;
   }

src/google/protobuf/arena_cleanup.h

@@ -9,15 +9,11 @@
 #define GOOGLE_PROTOBUF_ARENA_CLEANUP_H__
 
 #include <cstddef>
-#include <cstdint>
-#include <string>
+#include <cstring>
 #include <vector>
 
 #include "absl/base/attributes.h"
-#include "absl/log/absl_check.h"
-#include "absl/log/absl_log.h"
-#include "absl/strings/cord.h"
-
+#include "absl/base/prefetch.h"
 
 // Must be included last.
 #include "google/protobuf/port_def.inc"
@@ -54,6 +50,10 @@ inline ABSL_ATTRIBUTE_ALWAYS_INLINE void CreateNode(void* pos, void* elem,
 
 // Optimization: performs a prefetch on the elem for the cleanup node at `pos`.
 inline ABSL_ATTRIBUTE_ALWAYS_INLINE void PrefetchNode(void* pos) {
+  // We explicitly use NTA prefetch here to avoid polluting remote caches: we
+  // are destroying these instances, there is no purpose for these cache lines
+  // to linger around in remote caches.
+  absl::PrefetchToLocalCacheNta(ToCleanup(pos)->elem);
 }
 
 // Destroys the object referenced by the cleanup node.

src/google/protobuf/serial_arena.h

@@ -15,12 +15,8 @@
 #include <cstddef>
 #include <cstdint>
 #include <string>
-#include <type_traits>
-#include <typeinfo>
-#include <utility>
 #include <vector>
 
-#include "google/protobuf/stubs/common.h"
 #include "absl/base/attributes.h"
 #include "absl/base/optimization.h"
 #include "absl/base/prefetch.h"
@@ -28,7 +24,6 @@
 #include "absl/numeric/bits.h"
 #include "google/protobuf/arena_align.h"
 #include "google/protobuf/arena_cleanup.h"
-#include "google/protobuf/arenaz_sampler.h"
 #include "google/protobuf/port.h"
 #include "google/protobuf/string_block.h"