From 725e94d6e986449894002b1ac3bdf203d2e031ac Mon Sep 17 00:00:00 2001
From: Nikolai Tillmann <nikolait@meta.com>
Date: Wed, 1 May 2024 09:50:17 -0700
Subject: [PATCH] derive profiles for synthesized helper methods

Summary:
OEA may clone some methods, expanding some parameter of a non-escaping object type.
This makes it so that the cloned version of the method derives its method-profile from the original method.
This improves PGO data quality.

Reviewed By: wsanville

Differential Revision: D56586847

fbshipit-source-id: 70c230915b4726de07a61cdeeeff34ea303e8560
---
 .../ExpandableMethodParams.cpp                       | 12 ++++++++----
 opt/object-escape-analysis/ExpandableMethodParams.h  |  4 +++-
 opt/object-escape-analysis/ObjectEscapeAnalysis.cpp  |  9 ++++++---
 3 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/opt/object-escape-analysis/ExpandableMethodParams.cpp b/opt/object-escape-analysis/ExpandableMethodParams.cpp
index a502be46eaa..69415c0bf27 100644
--- a/opt/object-escape-analysis/ExpandableMethodParams.cpp
+++ b/opt/object-escape-analysis/ExpandableMethodParams.cpp
@@ -323,7 +323,8 @@ DexMethodRef* ExpandableMethodParams::get_expanded_method_ref(
 
 // Make sure that all newly used expanded methods actually exist as concrete
 // methods.
-size_t ExpandableMethodParams::flush(const Scope& scope) {
+size_t ExpandableMethodParams::flush(
+    const Scope& scope, method_profiles::MethodProfiles* method_profiles) {
   // First, find all expanded_method_ref that made it into the updated code.
   ConcurrentSet<DexMethodRef*> used_expanded_method_refs;
   walk::parallel::opcodes(scope, [&](DexMethod*, IRInstruction* insn) {
@@ -350,9 +351,12 @@ size_t ExpandableMethodParams::flush(const Scope& scope) {
     type_class(expanded_method->get_class())->add_method(expanded_method);
   }
 
-  // Finally, erase the unused method refs.
-  for (auto [method, param_index] : m_candidates) {
-    if (!used_expanded_method_refs.count(method)) {
+  // Finally, derived method-profiles for used candidates, and erase the unused
+  // method refs.
+  for (auto&& [method, p] : m_candidates) {
+    if (used_expanded_method_refs.count(method)) {
+      method_profiles->derive_stats(method->as_def(), {p.first});
+    } else {
       DexMethod::erase_method(method);
       DexMethod::delete_method_DO_NOT_USE(static_cast<DexMethod*>(method));
     }
diff --git a/opt/object-escape-analysis/ExpandableMethodParams.h b/opt/object-escape-analysis/ExpandableMethodParams.h
index c00f87642e9..0b58748c088 100644
--- a/opt/object-escape-analysis/ExpandableMethodParams.h
+++ b/opt/object-escape-analysis/ExpandableMethodParams.h
@@ -12,6 +12,7 @@
 #include "ConcurrentContainers.h"
 #include "DexClass.h"
 #include "IRInstruction.h"
+#include "MethodProfiles.h"
 
 // Helper class to deal with methods that take a
 // (newly created) object, and only use it to read ifields. For those
@@ -78,7 +79,8 @@ class ExpandableMethodParams {
 
   // Make sure that all newly used expanded ctors actually exist as concrete
   // methods.
-  size_t flush(const Scope& scope);
+  size_t flush(const Scope& scope,
+               method_profiles::MethodProfiles* method_profiles);
 
  private:
   mutable InsertOnlyConcurrentMap<MethodKey, MethodInfo, MethodKeyHash>
diff --git a/opt/object-escape-analysis/ObjectEscapeAnalysis.cpp b/opt/object-escape-analysis/ObjectEscapeAnalysis.cpp
index 044cff5d0b4..dfa48dbed12 100644
--- a/opt/object-escape-analysis/ObjectEscapeAnalysis.cpp
+++ b/opt/object-escape-analysis/ObjectEscapeAnalysis.cpp
@@ -1924,7 +1924,8 @@ void reduce(const Scope& scope,
             Stats* stats,
             std::unordered_set<DexMethod*>* inlinable_methods_kept,
             CalleesCache* callees_cache,
-            MethodSummaryCache* method_summary_cache) {
+            MethodSummaryCache* method_summary_cache,
+            method_profiles::MethodProfiles* method_profiles) {
   Timer t("reduce");
 
   // First, we compute all reduced methods
@@ -1975,7 +1976,8 @@ void reduce(const Scope& scope,
       },
       reduced_methods);
 
-  size_t expanded_methods = expandable_method_params.flush(scope);
+  size_t expanded_methods =
+      expandable_method_params.flush(scope, method_profiles);
   stats->expanded_methods += expanded_methods;
 }
 } // namespace
@@ -2080,7 +2082,8 @@ void ObjectEscapeAnalysisPass::run_pass(DexStoresVector& stores,
   Stats stats;
   reduce(scope, m_config, apply_shrinking_plugins, *method_override_graph,
          inliner, method_summaries, excluded_classes, root_methods, &stats,
-         &inlinable_methods_kept, &callees_cache, &method_summary_cache);
+         &inlinable_methods_kept, &callees_cache, &method_summary_cache,
+         &conf.get_method_profiles());
 
   TRACE(OEA, 1, "[object escape analysis] total savings: %zu",
         (size_t)stats.total_savings);