Cluster non cold subtables in aux entries to the top.

PiperOrigin-RevId: 578656159

src/google/protobuf/compiler/cpp/options.h

@@ -54,6 +54,7 @@ struct Options {
   bool opensource_runtime = false;
   bool annotate_accessor = false;
   bool force_split = false;
+  bool profile_driven_cluster_aux_subtable = false;
 #ifdef PROTOBUF_STABLE_EXPERIMENTS
   bool force_eagerly_verified_lazy = true;
   bool force_inline_string = true;

src/google/protobuf/compiler/cpp/parse_function_generator.cc

@@ -116,11 +116,9 @@ ParseFunctionGenerator::ParseFunctionGenerator(
   if (should_generate_tctable()) {
     tc_table_info_.reset(new TailCallTableInfo(
         descriptor_, ordered_fields_,
-        {
-            /* is_lite */ GetOptimizeFor(descriptor->file(), options_) ==
-                FileOptions::LITE_RUNTIME,
-            /* uses_codegen */ true,
-        },
+        {/* is_lite */ GetOptimizeFor(descriptor->file(), options_) ==
+             FileOptions::LITE_RUNTIME,
+         /* uses_codegen */ true, options_.profile_driven_cluster_aux_subtable},
         GeneratedOptionProvider(this), has_bit_indices,
         inlined_string_indices));
   }

src/google/protobuf/generated_message_reflection.cc

@@ -3416,6 +3416,7 @@ const internal::TcParseTableBase* Reflection::CreateTcParseTable() const {
       {
           /* is_lite */ false,
           /* uses_codegen */ false,
+          /* should_profile_driven_cluster_aux_table */ false,
       },
       ReflectionOptionProvider(*this), has_bit_indices, inlined_string_indices);
 

src/google/protobuf/generated_message_tctable_gen.cc

@@ -8,6 +8,7 @@
 #include "google/protobuf/generated_message_tctable_gen.h"
 
 #include <algorithm>
+#include <cstddef>
 #include <cstdint>
 #include <limits>
 #include <string>
@@ -756,6 +757,35 @@ TailCallTableInfo::TailCallTableInfo(
     }
   }
 
+  auto is_non_cold = [](PerFieldOptions options) {
+    return options.presence_probability >= 0.005;
+  };
+  size_t num_non_cold_subtables = 0;
+  if (message_options.should_profile_driven_cluster_aux_subtable) {
+    // We found that clustering non-cold subtables to the top of aux_entries
+    // achieves the best load tests results than other strategies (e.g.,
+    // clustering all non-cold entries).
+    auto is_non_cold_subtable = [&](const FieldDescriptor* field) {
+      auto options = option_provider.GetForField(field);
+      // In the following code where we assign kSubTable to aux entries, only
+      // the following typed fields are supported.
+      return (field->type() == FieldDescriptor::TYPE_MESSAGE ||
+              field->type() == FieldDescriptor::TYPE_GROUP) &&
+             !field->is_map() && !HasLazyRep(field, options) &&
+             !options.is_implicitly_weak && options.use_direct_tcparser_table &&
+             is_non_cold(options);
+    };
+    for (const FieldDescriptor* field : ordered_fields) {
+      if (is_non_cold_subtable(field)) {
+        num_non_cold_subtables++;
+      }
+    }
+  }
+
+  size_t subtable_aux_idx_begin = aux_entries.size();
+  size_t subtable_aux_idx = aux_entries.size();
+  aux_entries.resize(aux_entries.size() + num_non_cold_subtables);
+
   // Fill in mini table entries.
   for (const FieldDescriptor* field : ordered_fields) {
     auto options = option_provider.GetForField(field);
@@ -797,12 +827,18 @@ TailCallTableInfo::TailCallTableInfo(
               TcParseTableBase::FieldEntry::kNoAuxIdx;
         }
       } else {
-        field_entries.back().aux_idx = aux_entries.size();
-        aux_entries.push_back({options.is_implicitly_weak ? kSubMessageWeak
-                               : options.use_direct_tcparser_table
-                                   ? kSubTable
-                                   : kSubMessage,
-                               {field}});
+        AuxType type = options.is_implicitly_weak          ? kSubMessageWeak
+                       : options.use_direct_tcparser_table ? kSubTable
+                                                           : kSubMessage;
+        if (message_options.should_profile_driven_cluster_aux_subtable &&
+            type == kSubTable && is_non_cold(options)) {
+          aux_entries[subtable_aux_idx] = {type, {field}};
+          field_entries.back().aux_idx = subtable_aux_idx;
+          ++subtable_aux_idx;
+        } else {
+          field_entries.back().aux_idx = aux_entries.size();
+          aux_entries.push_back({type, {field}});
+        }
       }
     } else if (field->type() == FieldDescriptor::TYPE_ENUM &&
                !cpp::HasPreservingUnknownEnumSemantics(field)) {
@@ -845,6 +881,8 @@ TailCallTableInfo::TailCallTableInfo(
       entry.inlined_string_idx = idx;
     }
   }
+  ABSL_CHECK_EQ(subtable_aux_idx - subtable_aux_idx_begin,
+                num_non_cold_subtables);
 
   table_size_log2 = 0;  // fallback value
   int num_fast_fields = -1;

src/google/protobuf/generated_message_tctable_gen.h

@@ -38,6 +38,8 @@ struct PROTOBUF_EXPORT TailCallTableInfo {
   struct MessageOptions {
     bool is_lite;
     bool uses_codegen;
+    // TODO: remove this after A/B test is done.
+    bool should_profile_driven_cluster_aux_subtable;
   };
   struct PerFieldOptions {
     // For presence awareness (e.g. PDProto).

src/google/protobuf/message_unittest.inc

@@ -1937,9 +1937,10 @@ TEST(MESSAGE_TEST_NAME, TestRegressionInlinedStringAuxIdxMismatchOnFastParser) {
   if (table != nullptr &&
       table->fast_entry(1)->target() == internal::TcParser::FastSiS1) {
     // optional string str1 = 1;
+    // The aux_idx points to the inlined_string_idx and not the actual aux_idx.
     EXPECT_EQ(table->fast_entry(1)->bits.aux_idx(), 1);
     // optional InlinedStringIdxRegressionProto sub = 2;
-    EXPECT_EQ(table->fast_entry(2)->bits.aux_idx(), 2);
+    EXPECT_EQ(table->fast_entry(2)->bits.aux_idx(), 1);
     // optional string str2 = 3;
     // The aux_idx points to the inlined_string_idx and not the actual aux_idx.
     EXPECT_EQ(table->fast_entry(3)->bits.aux_idx(), 2);