Add a single-tag implementation for CheckedPtr perf testing

We suspect that CheckedPtr2 and MTECheckedPtr suffer from cache misses.
This solution will help us verify this hypothesis by satisfying all
requests from a single tag, which is very likely in a hot cacheline.

Bug: 1073933,1092288
Change-Id: I5ca13baed9e183d509d7546dffc5fa921cac252d
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2306131
Commit-Queue: Bartek Nowierski <bartekn@chromium.org>
Auto-Submit: Bartek Nowierski <bartekn@chromium.org>
Reviewed-by: Takashi Sakamoto <tasak@google.com>
Reviewed-by: Kentaro Hara <haraken@chromium.org>
Cr-Commit-Position: refs/heads/master@{#789866}

base/allocator/partition_allocator/partition_alloc.cc

@@ -73,17 +73,25 @@ static_assert(kGenericMaxDirectMapped <=
                   (1UL << 31) + kPageAllocationGranularity,
               "maximum direct mapped allocation");
 // Check that some of our zanier calculations worked out as expected.
-#if !ENABLE_TAG_FOR_MTE_CHECKED_PTR
-static_assert(kGenericSmallestBucket == alignof(std::max_align_t),
+#if ENABLE_TAG_FOR_MTE_CHECKED_PTR
+static_assert(kGenericSmallestBucket >= alignof(std::max_align_t),
               "generic smallest bucket");
 #else
-static_assert(kGenericSmallestBucket >= alignof(std::max_align_t),
+static_assert(kGenericSmallestBucket == alignof(std::max_align_t),
               "generic smallest bucket");
 #endif
 static_assert(kGenericMaxBucketed == 983040, "generic max bucketed");
 static_assert(kMaxSystemPagesPerSlotSpan < (1 << 8),
               "System pages per slot span must be less than 128.");
 
+#if ENABLE_TAG_FOR_SINGLE_TAG_CHECKED_PTR
+namespace internal {
+BASE_EXPORT PartitionTagWrapper g_checked_ptr_single_tag = {{},
+                                                            kFixedTagValue,
+                                                            {}};
+}
+#endif
+
 Lock& GetHooksLock() {
   static NoDestructor<Lock> lock;
   return *lock;

base/allocator/partition_allocator/partition_alloc.h

@@ -823,9 +823,7 @@ ALWAYS_INLINE void* PartitionRoot<thread_safe>::AllocFlags(
   }
   size_t requested_size = size;
   size = internal::PartitionSizeAdjustAdd(allow_extras, size);
-#if ENABLE_TAG_FOR_CHECKED_PTR2
-  PA_CHECK(size >= requested_size);
-#endif
+  PA_CHECK(size >= requested_size);  // check for overflows
   auto* bucket = SizeToBucket(size);
   PA_DCHECK(bucket);
   {

base/allocator/partition_allocator/partition_alloc_unittest.cc

@@ -2490,7 +2490,8 @@ TEST_F(PartitionAllocTest, AlignedAllocations) {
   }
 }
 
-#if ENABLE_TAG_FOR_CHECKED_PTR2 || ENABLE_TAG_FOR_MTE_CHECKED_PTR
+#if ENABLE_TAG_FOR_CHECKED_PTR2 || ENABLE_TAG_FOR_MTE_CHECKED_PTR || \
+    ENABLE_TAG_FOR_SINGLE_TAG_CHECKED_PTR
 
 TEST_F(PartitionAllocTest, TagBasic) {
   size_t alloc_size = 64 - kExtraAllocSize;
@@ -2513,9 +2514,16 @@ TEST_F(PartitionAllocTest, TagBasic) {
   EXPECT_EQ(char_ptr1 + page->bucket->slot_size, char_ptr2);
   EXPECT_EQ(char_ptr2 + page->bucket->slot_size, char_ptr3);
 
+#if !ENABLE_TAG_FOR_SINGLE_TAG_CHECKED_PTR
   constexpr PartitionTag kTag1 = static_cast<PartitionTag>(0xBADA);
   constexpr PartitionTag kTag2 = static_cast<PartitionTag>(0xDB8A);
   constexpr PartitionTag kTag3 = static_cast<PartitionTag>(0xA3C4);
+#else
+  // The in-memory tag will always be kFixedTagValue no matter what we set.
+  constexpr PartitionTag kTag1 = static_cast<PartitionTag>(kFixedTagValue);
+  constexpr PartitionTag kTag2 = static_cast<PartitionTag>(kFixedTagValue);
+  constexpr PartitionTag kTag3 = static_cast<PartitionTag>(kFixedTagValue);
+#endif
   PartitionTagSetValue(ptr1, page->bucket->slot_size, kTag1);
   PartitionTagSetValue(ptr2, page->bucket->slot_size, kTag2);
   PartitionTagSetValue(ptr3, page->bucket->slot_size, kTag3);
@@ -2545,11 +2553,8 @@ TEST_F(PartitionAllocTest, TagBasic) {
   EXPECT_EQ(ptr2, new_ptr2);
   EXPECT_EQ(kTag3, PartitionTagGetValue(ptr3));
 
-  request_size =
-      page->bucket->slot_size - kExtraAllocSize + kInSlotTagBufferSize;
-#if ENABLE_TAG_FOR_MTE_CHECKED_PTR
-  request_size += kGenericSmallestBucket;
-#endif
+  // Add 1B to ensure the object is rellocated to a larger slot.
+  request_size = page->bucket->slot_size - kExtraAllocSize + 1;
   new_ptr2 = allocator.root()->Realloc(ptr2, request_size, type_name);
   EXPECT_TRUE(new_ptr2);
   EXPECT_NE(ptr2, new_ptr2);

base/allocator/partition_allocator/partition_tag.h

@@ -11,6 +11,7 @@
 #include "base/allocator/partition_allocator/partition_alloc_constants.h"
 #include "base/allocator/partition_allocator/partition_cookie.h"
 #include "base/allocator/partition_allocator/partition_tag_bitmap.h"
+#include "base/base_export.h"
 #include "base/notreached.h"
 #include "build/build_config.h"
 
@@ -150,7 +151,53 @@ ALWAYS_INLINE void PartitionTagClearValue(void* ptr, size_t size) {
   memset(PartitionTagPointer(ptr), 0, tag_region_size);
 }
 
-#else  // !ENABLE_TAG_FOR_CHECKED_PTR2 && !ENABLE_TAG_FOR_MTE_CHECKED_PTR
+#elif ENABLE_TAG_FOR_SINGLE_TAG_CHECKED_PTR
+
+using PartitionTag = uint8_t;
+
+static constexpr PartitionTag kFixedTagValue = 0xAD;
+
+struct PartitionTagWrapper {
+  // Add padding before and after the tag, to avoid cacheline false sharing.
+  // Assume cacheline is 64B.
+  uint8_t unused1[64];
+  PartitionTag partition_tag;
+  uint8_t unused2[64];
+};
+extern BASE_EXPORT PartitionTagWrapper g_checked_ptr_single_tag;
+
+static constexpr size_t kInSlotTagBufferSize = 0;
+
+ALWAYS_INLINE size_t PartitionTagSizeAdjustAdd(size_t size) {
+  return size;
+}
+
+ALWAYS_INLINE size_t PartitionTagSizeAdjustSubtract(size_t size) {
+  return size;
+}
+
+ALWAYS_INLINE PartitionTag* PartitionTagPointer(void*) {
+  return &g_checked_ptr_single_tag.partition_tag;
+}
+
+ALWAYS_INLINE void* PartitionTagPointerAdjustSubtract(void* ptr) {
+  return ptr;
+}
+
+ALWAYS_INLINE void* PartitionTagPointerAdjustAdd(void* ptr) {
+  return ptr;
+}
+
+ALWAYS_INLINE void PartitionTagSetValue(void*, size_t, PartitionTag) {}
+
+ALWAYS_INLINE PartitionTag PartitionTagGetValue(void* ptr) {
+  return *PartitionTagPointer(ptr);
+}
+
+ALWAYS_INLINE void PartitionTagClearValue(void* ptr, size_t) {}
+
+#else  // !ENABLE_TAG_FOR_CHECKED_PTR2 && !ENABLE_TAG_FOR_MTE_CHECKED_PTR &&
+       // !ENABLE_TAG_FOR_SINGLE_TAG_CHECKED_PTR
 
 using PartitionTag = uint16_t;
 
@@ -186,7 +233,8 @@ ALWAYS_INLINE PartitionTag PartitionTagGetValue(void*) {
 
 ALWAYS_INLINE void PartitionTagClearValue(void* ptr, size_t) {}
 
-#endif  // !ENABLE_TAG_FOR_CHECKED_PTR2
+#endif  // !ENABLE_TAG_FOR_CHECKED_PTR2 && !ENABLE_TAG_FOR_MTE_CHECKED_PTR &&
+        // !ENABLE_TAG_FOR_SINGLE_TAG_CHECKED_PTR
 
 }  // namespace internal
 }  // namespace base

base/memory/checked_ptr.cc

@@ -35,7 +35,10 @@ BASE_EXPORT bool CheckedPtr2OrMTEImplPartitionAllocSupport::EnabledForPtr(
   // the same result regardless, anyway.
   // TODO(bartekn): Figure out the thread-safety mismatch.
   return IsManagedByPartitionAllocNormalBuckets(ptr)
-#if ENABLE_TAG_FOR_CHECKED_PTR2
+  // Checking offset is not needed for ENABLE_TAG_FOR_SINGLE_TAG_CHECKED_PTR,
+  // but call it anyway for apples-to-apples comparison with
+  // ENABLE_TAG_FOR_CHECKED_PTR2.
+#if ENABLE_TAG_FOR_CHECKED_PTR2 || ENABLE_TAG_FOR_SINGLE_TAG_CHECKED_PTR
          && PartitionAllocGetSlotOffset<ThreadSafe>(ptr) == 0
 #endif
       ;

base/memory/checked_ptr.h

@@ -20,7 +20,8 @@
 
 #define ENABLE_CHECKED_PTR2_OR_MTE_IMPL 0
 #if ENABLE_CHECKED_PTR2_OR_MTE_IMPL
-static_assert(ENABLE_TAG_FOR_CHECKED_PTR2 || ENABLE_TAG_FOR_MTE_CHECKED_PTR,
+static_assert(ENABLE_TAG_FOR_CHECKED_PTR2 || ENABLE_TAG_FOR_MTE_CHECKED_PTR ||
+                  ENABLE_TAG_FOR_SINGLE_TAG_CHECKED_PTR,
               "CheckedPtr2OrMTEImpl can only by used if tags are enabled");
 #endif
 
@@ -185,10 +186,10 @@ struct CheckedPtr2OrMTEImpl {
   static ALWAYS_INLINE void* SafelyUnwrapPtrForDereference(
       uintptr_t wrapped_ptr) {
 #if CHECKED_PTR2_AVOID_BRANCH_WHEN_CHECKING_ENABLED
-    // This variant cannot be used with MTECheckedPtr algorithm, because it
+    // This variant can only be used with CheckedPtr2 algorithm, because it
     // relies on the generation to exist at a constant offset before the
     // allocation.
-    static_assert(!ENABLE_TAG_FOR_MTE_CHECKED_PTR, "");
+    static_assert(ENABLE_TAG_FOR_CHECKED_PTR2, "");
 
     // Top bit tells if the protection is enabled. Use it to decide whether to
     // read the word before the allocation, which exists only if the protection