Change redex to output multiple baseline profiles

Summary:
Modifies the baseline profile construction to take into all baseline profile configs and to add manual profile data.

Modifies the ArtProfileWriter pass to output baseline profiles for all baseline profile configs. Also copies the class information from the preprocessed baseline profile list passed into redex.

Reviewed By: jimmycFB

Differential Revision: D67163202

fbshipit-source-id: bf18ed4151adcd6b4c514952d1b99344947afd74

libredex/BaselineProfile.cpp

@@ -10,83 +10,129 @@
 
 namespace baseline_profiles {
 
-BaselineProfile get_baseline_profile(
-    const BaselineProfileConfig& config,
+BaselineProfile get_default_baseline_profile(
+    const std::unordered_map<std::string, BaselineProfileConfig>& configs,
     const method_profiles::MethodProfiles& method_profiles,
+    const bool ingest_baseline_profile_data,
     std::unordered_set<const DexMethodRef*>* method_refs_without_def) {
-  std::unordered_set<const DexType*> classes;
-  std::unordered_set<const DexMethod*> startup_methods;
-  std::unordered_set<const DexMethod*> post_startup_methods;
-  for (auto&& [interaction_id, interaction_config] :
-       config.interaction_configs) {
-    const auto& method_stats = method_profiles.method_stats(interaction_id);
-    for (auto&& [method_ref, stat] : method_stats) {
-      auto method = method_ref->as_def();
-      if (method == nullptr) {
-        if (method_refs_without_def != nullptr) {
-          method_refs_without_def->insert(method_ref);
+  auto baseline_profiles = get_baseline_profiles(configs,
+                                                 method_profiles,
+                                                 ingest_baseline_profile_data,
+                                                 method_refs_without_def);
+  return baseline_profiles.at(DEFAULT_BASELINE_PROFILE_CONFIG_NAME);
+}
+
+std::unordered_map<std::string, BaselineProfile> get_baseline_profiles(
+    const std::unordered_map<std::string, BaselineProfileConfig>& configs,
+    const method_profiles::MethodProfiles& method_profiles,
+    const bool ingest_baseline_profile_data,
+    std::unordered_set<const DexMethodRef*>* method_refs_without_def) {
+  std::unordered_map<std::string, BaselineProfile> baseline_profiles;
+  for (const auto& [config_name, config] : configs) {
+    // If we're not using this as the final pass of baseline profiles, just
+    // continue on all configs that aren't the default
+    if (!ingest_baseline_profile_data &&
+        config_name != DEFAULT_BASELINE_PROFILE_CONFIG_NAME) {
+      continue;
+    }
+    std::unordered_set<const DexType*> classes;
+    std::unordered_set<const DexMethod*> startup_methods;
+    std::unordered_set<const DexMethod*> post_startup_methods;
+    for (auto&& [interaction_id, interaction_config] :
+         config.interaction_configs) {
+      const auto& method_stats =
+          method_profiles.method_stats_for_baseline_config(interaction_id,
+                                                           config_name);
+      for (auto&& [method_ref, stat] : method_stats) {
+        auto method = method_ref->as_def();
+        if (method == nullptr) {
+          if (method_refs_without_def != nullptr) {
+            method_refs_without_def->insert(method_ref);
+          }
+          continue;
         }
-        continue;
-      }
 
-      if (stat.appear_percent < interaction_config.threshold ||
-          stat.call_count < interaction_config.call_threshold) {
-        continue;
-      }
+        if (stat.appear_percent < interaction_config.threshold ||
+            stat.call_count < interaction_config.call_threshold) {
+          continue;
+        }
 
-      if (interaction_config.startup) {
-        startup_methods.insert(method);
-      }
-      if (interaction_config.post_startup) {
-        post_startup_methods.insert(method);
+        if (interaction_config.startup) {
+          startup_methods.insert(method);
+        }
+        if (interaction_config.post_startup) {
+          post_startup_methods.insert(method);
+        }
+        if (interaction_config.classes) {
+          classes.insert(method->get_class());
+        }
       }
-      if (interaction_config.classes) {
-        classes.insert(method->get_class());
+    }
+    static const std::array<std::string, 2> manual_interaction_ids = {
+        "manual_startup", "manual_post_startup"};
+    for (const auto& interaction_id : manual_interaction_ids) {
+      const auto& method_stats =
+          method_profiles.method_stats_for_baseline_config(interaction_id,
+                                                           config_name);
+      for (auto&& [method_ref, stat] : method_stats) {
+        auto method = method_ref->as_def();
+        if (method == nullptr) {
+          if (method_refs_without_def != nullptr) {
+            method_refs_without_def->insert(method_ref);
+          }
+          continue;
+        }
+        if (interaction_id == "manual_startup") {
+          startup_methods.insert(method);
+        }
+        if (interaction_id == "manual_post_startup") {
+          post_startup_methods.insert(method);
+        }
       }
     }
-  }
-
-  // methods = startup_methods | post_startup_methods
-  std::unordered_set<const DexMethod*> methods(startup_methods.begin(),
-                                               startup_methods.end());
-  methods.insert(post_startup_methods.begin(), post_startup_methods.end());
+    // methods = startup_methods | post_startup_methods
+    std::unordered_set<const DexMethod*> methods(startup_methods.begin(),
+                                                 startup_methods.end());
+    methods.insert(post_startup_methods.begin(), post_startup_methods.end());
 
-  // startup_post_startup_methods = startup_methods & post_startup_methods
-  std::unordered_set<const DexMethod*> startup_post_startup_methods(
-      startup_methods.begin(), startup_methods.end());
-  std20::erase_if(startup_post_startup_methods, [&](const DexMethod* m) {
-    return !post_startup_methods.count(m);
-  });
+    // startup_post_startup_methods = startup_methods & post_startup_methods
+    std::unordered_set<const DexMethod*> startup_post_startup_methods(
+        startup_methods.begin(), startup_methods.end());
+    std20::erase_if(startup_post_startup_methods, [&](const DexMethod* m) {
+      return !post_startup_methods.count(m);
+    });
 
-  // startup_methods -= startup_post_startup_methods
-  std20::erase_if(startup_methods, [&](const DexMethod* m) {
-    return startup_post_startup_methods.count(m);
-  });
+    // startup_methods -= startup_post_startup_methods
+    std20::erase_if(startup_methods, [&](const DexMethod* m) {
+      return startup_post_startup_methods.count(m);
+    });
 
-  // post_startup_methods -= startup_post_startup_methods
-  std20::erase_if(post_startup_methods, [&](const DexMethod* m) {
-    return startup_post_startup_methods.count(m);
-  });
+    // post_startup_methods -= startup_post_startup_methods
+    std20::erase_if(post_startup_methods, [&](const DexMethod* m) {
+      return startup_post_startup_methods.count(m);
+    });
 
-  baseline_profiles::BaselineProfile res;
-  for (auto* method : methods) {
-    auto& flags = res.methods[method];
-    if (startup_post_startup_methods.count(method)) {
-      flags.hot = true;
-      flags.startup = true;
-      flags.post_startup = true;
-    } else if (startup_methods.count(method)) {
-      flags.hot = true;
-      flags.startup = true;
-    } else if (post_startup_methods.count(method)) {
-      flags.hot = true;
-      flags.post_startup = true;
+    baseline_profiles::BaselineProfile res;
+    for (auto* method : methods) {
+      auto& flags = res.methods[method];
+      if (startup_post_startup_methods.count(method)) {
+        flags.hot = true;
+        flags.startup = true;
+        flags.post_startup = true;
+      } else if (startup_methods.count(method)) {
+        flags.hot = true;
+        flags.startup = true;
+      } else if (post_startup_methods.count(method)) {
+        flags.hot = true;
+        flags.post_startup = true;
+      }
     }
+    for (auto* type : classes) {
+      res.classes.insert(type_class(type));
+    }
+    baseline_profiles[config_name] = res;
   }
-  for (auto* type : classes) {
-    res.classes.insert(type_class(type));
-  }
-  return res;
+  return baseline_profiles;
 }
 
 } // namespace baseline_profiles

libredex/BaselineProfile.h

@@ -28,9 +28,16 @@ struct BaselineProfile {
   std::unordered_set<const DexClass*> classes;
 };
 
-BaselineProfile get_baseline_profile(
-    const BaselineProfileConfig& config,
+std::unordered_map<std::string, BaselineProfile> get_baseline_profiles(
+    const std::unordered_map<std::string, BaselineProfileConfig>& configs,
     const method_profiles::MethodProfiles& method_profiles,
+    const bool ingest_baseline_profile_data,
+    std::unordered_set<const DexMethodRef*>* method_refs_without_def = nullptr);
+
+BaselineProfile get_default_baseline_profile(
+    const std::unordered_map<std::string, BaselineProfileConfig>& configs,
+    const method_profiles::MethodProfiles& method_profiles,
+    const bool ingest_baseline_profile_data,
     std::unordered_set<const DexMethodRef*>* method_refs_without_def = nullptr);
 
 } // namespace baseline_profiles

libredex/ConfigFiles.cpp

@@ -352,6 +352,12 @@ std::vector<std::string> ConfigFiles::load_coldstart_methods() {
   return coldstart_methods;
 }
 
+std::string ConfigFiles::get_preprocessed_baseline_profile_file(
+    const std::string& config_name) {
+  return m_preprocessed_baseline_profile_directory + "/" + config_name +
+         "/baseline-profile.txt";
+}
+
 /**
  * Read a map of {list_name : class_list} from json
  */

libredex/ConfigFiles.h

@@ -194,6 +194,9 @@ struct ConfigFiles {
 
   bool get_did_use_bzl_baseline_profile_config();
 
+  std::string get_preprocessed_baseline_profile_file(
+      const std::string& config_name);
+
   boost::optional<std::string> get_android_sdk_api_file(int32_t api_level) {
     std::string api_file;
     std::string key = "android_sdk_api_" + std::to_string(api_level) + "_file";

libredex/MethodProfiles.cpp

@@ -52,6 +52,46 @@ bool empty_column(std::string_view sv) { return sv.empty() || sv == "\n"; }
 AccumulatingTimer MethodProfiles::s_process_unresolved_lines_timer(
     "MethodProfiles::process_unresolved_lines");
 
+std::tuple<const StatsMap&, bool> method_stats_for_interaction_id(
+    const std::string& interaction_id, const AllInteractions& interactions) {
+  const auto& search1 = interactions.find(interaction_id);
+  if (search1 != interactions.end()) {
+    return {search1->second, true};
+  }
+  if (interaction_id == COLD_START) {
+    // Originally, the stats file had no interaction_id column and it only
+    // covered coldstart. Search for the default (empty string) for backwards
+    // compatibility when we're searching for coldstart but it's not found.
+    const auto& search2 = interactions.find("");
+    if (search2 != interactions.end()) {
+      return {search2->second, true};
+    }
+  }
+
+  static StatsMap empty_map = {};
+  return {empty_map, false};
+}
+
+const StatsMap& MethodProfiles::method_stats_for_baseline_config(
+    const std::string& interaction_id,
+    const std::string& baseline_config_name) const {
+  if (baseline_config_name !=
+      baseline_profiles::DEFAULT_BASELINE_PROFILE_CONFIG_NAME) {
+    if (m_baseline_manual_interactions.count(baseline_config_name)) {
+      const auto& method_stats =
+          *(m_baseline_manual_interactions.at(baseline_config_name));
+      const auto& [stats, found] =
+          method_stats_for_interaction_id(interaction_id, method_stats);
+      if (found) {
+        return stats;
+      }
+    }
+  }
+  const auto& [stats, _] =
+      method_stats_for_interaction_id(interaction_id, m_method_stats);
+  return stats;
+}
+
 const StatsMap& MethodProfiles::method_stats(
     const std::string& interaction_id) const {
   const auto& search1 = m_method_stats.find(interaction_id);

libredex/MethodProfiles.h

@@ -134,6 +134,10 @@ class MethodProfiles {
   // If no interactions are found by that interaction id, Return an empty map.
   const StatsMap& method_stats(const std::string& interaction_id) const;
 
+  const StatsMap& method_stats_for_baseline_config(
+      const std::string& interaction_id,
+      const std::string& baseline_config_name) const;
+
   const AllInteractions& all_interactions() const { return m_method_stats; }
 
   boost::optional<Stats> get_method_stat(const std::string& interaction_id,