Coldstart method ordering

Summary: Add the logic to read an orderfile of methods that will be placed at the beginning of the code items section.

Reviewed By: NTillmann

Differential Revision: D51071073

fbshipit-source-id: ed661d303a5d1d9211906fe5274b46e9175a101f

libredex/ConfigFiles.cpp

@@ -73,6 +73,8 @@ ConfigFiles::ConfigFiles(const Json::Value& config, const std::string& outdir)
       m_proguard_map(
           new ProguardMap(config.get("proguard_map", "").asString(),
                           config.get("use_new_rename_map", 0).asBool())),
+      m_coldstart_methods_filename(
+          config.get("coldstart_methods_file", "").asString()),
       m_printseeds(config.get("printseeds", "").asString()),
       m_method_profiles(new method_profiles::MethodProfiles()),
       m_secondary_method_profiles(new method_profiles::MethodProfiles()) {
@@ -188,6 +190,7 @@ const std::unordered_set<DexType*>& ConfigFiles::get_do_not_devirt_anon() {
   }
   return m_no_devirtualize_annos;
 }
+
 /**
  * Read an interdex list file and return as a vector of appropriately-formatted
  * classname strings.
@@ -242,6 +245,30 @@ std::vector<std::string> ConfigFiles::load_coldstart_classes() {
   return coldstart_classes;
 }
 
+/**
+ * Read a method ordering file for coldstart and return as a vector of
+ * appropriately-formatted methodname strings.
+ */
+std::vector<std::string> ConfigFiles::load_coldstart_methods() {
+  if (m_coldstart_methods_filename.empty()) {
+    return {};
+  }
+  std::vector<std::string> coldstart_methods;
+  std::ifstream input(m_coldstart_methods_filename);
+  if (!input) {
+    throw RedexException(
+        RedexError::INTERNAL_ERROR,
+        "[error] Can not open <coldstart_meth_ordering> file, path is "s +
+            m_coldstart_methods_filename);
+  }
+
+  std::string method;
+  while (input >> method) {
+    coldstart_methods.emplace_back(method);
+  }
+  return coldstart_methods;
+}
+
 /**
  * Read a map of {list_name : class_list} from json
  */

libredex/ConfigFiles.h

@@ -56,6 +56,13 @@ struct ConfigFiles {
     return m_coldstart_classes;
   }
 
+  const std::vector<std::string>& get_coldstart_methods() {
+    if (m_coldstart_methods.empty()) {
+      m_coldstart_methods = load_coldstart_methods();
+    }
+    return m_coldstart_methods;
+  }
+
   /**
    * NOTE: ONLY use if you know what you are doing!
    */
@@ -211,6 +218,7 @@ struct ConfigFiles {
   std::string outdir;
   GlobalConfig m_global_config;
 
+  std::vector<std::string> load_coldstart_methods();
   std::vector<std::string> load_coldstart_classes();
   std::unordered_map<std::string, std::vector<std::string>> load_class_lists();
   void ensure_agg_method_stats_loaded() const;
@@ -228,7 +236,9 @@ struct ConfigFiles {
   bool m_load_class_lists_attempted{false};
   std::unique_ptr<ProguardMap> m_proguard_map;
   std::string m_coldstart_class_filename;
+  std::string m_coldstart_methods_filename;
   std::vector<std::string> m_coldstart_classes;
+  std::vector<std::string> m_coldstart_methods;
   std::unordered_map<std::string, std::vector<std::string>> m_class_lists;
   bool m_dead_class_list_attempted{false};
   std::string m_printseeds; // Filename to dump computed seeds.

libredex/DexOutput.cpp

@@ -266,6 +266,57 @@ void GatheredTypes::sort_dexmethod_emitlist_cls_order(
                                                       compare_dexmethods));
 }
 
+void GatheredTypes::sort_dexmethod_emitlist_coldstart_order(
+    std::vector<DexMethod*>& lmeth) {
+  const auto& coldstart_methods = m_config->get_coldstart_methods();
+
+  if (coldstart_methods.empty()) {
+    return;
+  }
+
+  std::unordered_map<const DexMethod*, unsigned int> coldstart_method_ordering;
+  for (auto& method : coldstart_methods) {
+    DexMethodRef* method_ref = DexMethod::get_method(method);
+    if (method_ref == nullptr || !method_ref->is_def()) {
+      continue;
+    }
+    auto* meth = method_ref->as_def();
+    always_assert_log(coldstart_method_ordering.count(meth) == 0,
+                      "%s already found at position %u", SHOW(meth),
+                      coldstart_method_ordering.at(meth));
+    coldstart_method_ordering[meth] = coldstart_method_ordering.size();
+  }
+
+  std::unordered_map<const DexMethodRef*, size_t> initial_ordering;
+  for (auto* m : lmeth) {
+    initial_ordering[m] = initial_ordering.size();
+  }
+
+  auto coldstart_cmp = [&initial_ordering, &coldstart_method_ordering](
+                           const DexMethod* a, const DexMethod* b) {
+    const auto end_it = coldstart_method_ordering.end();
+    const auto a_it = coldstart_method_ordering.find(a);
+    const auto b_it = coldstart_method_ordering.find(b);
+    if (a_it == end_it && b_it == end_it) {
+      return initial_ordering.at(a) < initial_ordering.at(b);
+    }
+    if (a_it != end_it && b_it != end_it) {
+      const auto a_idx = a_it->second;
+      const auto b_idx = b_it->second;
+      if (a_idx != b_idx) {
+        return a_idx < b_idx;
+      }
+      return initial_ordering.at(a) < initial_ordering.at(b);
+    }
+    if (a_it != end_it) {
+      return true;
+    }
+    return false;
+  };
+
+  std::stable_sort(lmeth.begin(), lmeth.end(), coldstart_cmp);
+}
+
 void GatheredTypes::sort_dexmethod_emitlist_profiled_order(
     std::vector<DexMethod*>& lmeth) {
   // Use std::ref to avoid comparator copies.
@@ -1044,6 +1095,10 @@ void DexOutput::generate_code_items(const std::vector<SortMode>& mode) {
       TRACE(CUSTOMSORT, 2, "using method profiled order for bytecode sorting");
       m_gtypes->sort_dexmethod_emitlist_profiled_order(lmeth);
       break;
+    case SortMode::METHOD_COLDSTART_ORDER:
+      TRACE(CUSTOMSORT, 2, "using bps ordering for coldstart");
+      m_gtypes->sort_dexmethod_emitlist_coldstart_order(lmeth);
+      break;
     case SortMode::CLINIT_FIRST:
       TRACE(CUSTOMSORT, 2,
             "sorting <clinit> sections before all other bytecode");
@@ -2986,6 +3041,8 @@ static SortMode make_sort_bytecode(const std::string& sort_bytecode) {
     return SortMode::METHOD_PROFILED_ORDER;
   } else if (sort_bytecode == "method_similarity_order") {
     return SortMode::METHOD_SIMILARITY;
+  } else if (sort_bytecode == "method_coldstart_order") {
+    return SortMode::METHOD_COLDSTART_ORDER;
   } else {
     return SortMode::DEFAULT;
   }

libredex/DexOutput.h

@@ -46,6 +46,7 @@ enum class SortMode {
   CLASS_ORDER,
   CLASS_STRINGS,
   CLINIT_FIRST,
+  METHOD_COLDSTART_ORDER,
   METHOD_PROFILED_ORDER,
   METHOD_SIMILARITY,
   DEFAULT
@@ -273,6 +274,7 @@ class GatheredTypes {
 
   void sort_dexmethod_emitlist_method_similarity_order(
       std::vector<DexMethod*>& lmeth);
+  void sort_dexmethod_emitlist_coldstart_order(std::vector<DexMethod*>& lmeth);
   void sort_dexmethod_emitlist_default_order(std::vector<DexMethod*>& lmeth);
   void sort_dexmethod_emitlist_cls_order(std::vector<DexMethod*>& lmeth);
   void sort_dexmethod_emitlist_clinit_order(std::vector<DexMethod*>& lmeth);

libredex/GlobalConfig.cpp

@@ -148,6 +148,7 @@ void GlobalConfig::bind_config() {
   bind("android_sdk_api_29_file", "", string_param);
   bind("bytecode_sort_mode", {}, string_vector_param);
   bind("coldstart_classes", "", string_param);
+  bind("coldstart_methods_file", "", string_param);
   bind("compute_xml_reachability", false, bool_param);
   bind("unused_keep_rule_abort", false, bool_param);
   bind("debug_info_kind", "", string_param);

redex.py

@@ -1011,6 +1011,18 @@ def _handle_profiles(args: argparse.Namespace) -> None:
     else:
         logging.info("No block profiles found in %s", args.packed_profiles)
 
+    coldstart_method_ordering_str = join_str.join(
+        f"{f.path}"
+        for f in os.scandir(directory)
+        if f.is_file() and f.name.startswith("coldstart_method_ordering")
+    )
+    if coldstart_method_ordering_str:
+        logging.debug("Found coldstart ordering: %s", coldstart_method_ordering_str)
+        # Assume there's at most one.
+        args.passthru.append(f"coldstart_methods_file={coldstart_method_ordering_str}")
+    else:
+        logging.info("No coldstart ordering found in %s", args.packed_profiles)
+
 
 def _handle_secondary_method_profiles(args: argparse.Namespace) -> None:
     if not args.secondary_packed_profiles:

test/integ/DexOutputTest.cpp

@@ -28,6 +28,7 @@ class DexOutputTest : public RedexIntegrationTest {
 
   std::vector<DexMethod*> get_ordered_methods(DexOutput& dout) {
     std::vector<DexMethod*> code_items;
+    code_items.reserve(dout.m_code_item_emits.size());
     for (auto& ci : dout.m_code_item_emits) {
       code_items.push_back(ci.method);
     }
@@ -266,3 +267,122 @@ TEST_F(DexOutputTest, TEST_COMPRESSION_ORDERER_PERF_SENSITIVE) {
   EXPECT_TRUE(std::equal(method_names.begin(), method_names.end(),
                          expected_order.begin()));
 }
+
+TEST_F(DexOutputTest, TEST_COLDSTART_ORDER) {
+  std::vector<SortMode> sort_modes{SortMode::METHOD_COLDSTART_ORDER,
+                                   SortMode::METHOD_SIMILARITY};
+  std::string profile_path = std::getenv("coldstart_methods_file");
+
+  Json::Value cfg;
+  std::istringstream temp_json("{\"coldstart_methods_file\":\"" + profile_path +
+                               "\"}");
+  temp_json >> cfg;
+
+  ConfigFiles config_files(cfg);
+  config_files.parse_global_config();
+  std::unique_ptr<PositionMapper> pos_mapper(PositionMapper::make(""));
+
+  (*classes)[1]->set_perf_sensitive(PerfSensitiveGroup::BETAMAP_ORDERED);
+  (*classes)[2]->set_perf_sensitive(PerfSensitiveGroup::BETAMAP_ORDERED);
+
+  DexMethod::make_method("LDexOutputTest2$Class;.someRandomMethodNotInDex:(I)I")
+      ->make_concrete(ACC_PUBLIC, false);
+
+  auto scope = build_class_scope(stores);
+
+  // Lower the code
+  walk::parallel::methods<>(
+      scope, [](DexMethod* m) { instruction_lowering::lower(m, true); });
+
+  auto gtypes = std::make_shared<GatheredTypes>(&*classes);
+  DexOutput dout("", &*classes, std::move(gtypes), nullptr, true, 0, nullptr, 0,
+                 DebugInfoKind::NoCustomSymbolication, nullptr, config_files,
+                 pos_mapper.get(), nullptr, nullptr, DexOutputConfig(), 25);
+
+  dout.prepare(SortMode::DEFAULT, sort_modes, config_files, "dex\n039");
+  auto code_items = get_ordered_methods(dout);
+
+  uint32_t i = 0;
+  std::vector<std::string> method_names(code_items.size());
+  for (auto* m : code_items) {
+    method_names[i++] = show(m);
+  }
+
+  std::vector<std::string> expected_order = {
+      "LDexOutputTest$PerfSensitiveClass;.<init>:(LDexOutputTest;)V",
+      "LDexOutputTest$SecondPerfSensitiveClass;.<init>:(LDexOutputTest;)V",
+      "LDexOutputTest$PerfSensitiveClass;.EjustReturnFive:()I",
+      "LDexOutputTest$SecondPerfSensitiveClass;.EjustReturnFive:()I",
+      "LDexOutputTest$SecondPerfSensitiveClass;.FsomeLogic:(I)I",
+      "LDexOutputTest$PerfSensitiveClass;.FsomeLogic:(I)I",
+      "LDexOutputTest$NonPerfSensitiveClass;.<init>:(LDexOutputTest;)V",
+      "LDexOutputTest;.<init>:()V",
+      "LDexOutputTest$NonPerfSensitiveClass;.EjustReturnFive:()I",
+      "LDexOutputTest;.AjustReturnFive:()I",
+      "LDexOutputTest;.EjustReturnFive:()I",
+      "LDexOutputTest;.DgetSixpublic:()I",
+      "LDexOutputTest$NonPerfSensitiveClass;.FsomeLogic:(I)I",
+      "LDexOutputTest;.CsomeLogic:(I)I",
+      "LDexOutputTest;.FsomeLogic:(I)I",
+      "LDexOutputTest;.HsomeLogic:(I)I",
+      "LDexOutputTest;.BjustCallSixpublic:()I",
+      "LDexOutputTest;.GjustCallSixpublic:()I"};
+
+  EXPECT_TRUE(std::equal(method_names.begin(), method_names.end(),
+                         expected_order.begin()));
+}
+
+TEST_F(DexOutputTest, TEST_COLDSTART_ORDER_EMPTY_FILE) {
+  std::vector<SortMode> sort_modes{SortMode::METHOD_COLDSTART_ORDER,
+                                   SortMode::METHOD_SIMILARITY};
+  Json::Value cfg;
+  ConfigFiles config_files(cfg);
+  config_files.parse_global_config();
+  std::unique_ptr<PositionMapper> pos_mapper(PositionMapper::make(""));
+
+  (*classes)[1]->set_perf_sensitive(PerfSensitiveGroup::BETAMAP_ORDERED);
+  (*classes)[2]->set_perf_sensitive(PerfSensitiveGroup::BETAMAP_ORDERED);
+
+  auto scope = build_class_scope(stores);
+
+  // Lower the code
+  walk::parallel::methods<>(
+      scope, [](DexMethod* m) { instruction_lowering::lower(m, true); });
+
+  auto gtypes = std::make_shared<GatheredTypes>(&*classes);
+  DexOutput dout("", &*classes, std::move(gtypes), nullptr, true, 0, nullptr, 0,
+                 DebugInfoKind::NoCustomSymbolication, nullptr, config_files,
+                 pos_mapper.get(), nullptr, nullptr, DexOutputConfig(), 25);
+
+  dout.prepare(SortMode::DEFAULT, sort_modes, config_files, "dex\n039");
+  auto code_items = get_ordered_methods(dout);
+
+  uint32_t i = 0;
+  std::vector<std::string> method_names(code_items.size());
+  for (auto* m : code_items) {
+    method_names[i++] = show(m);
+  }
+
+  std::vector<std::string> expected_order = {
+      "LDexOutputTest$NonPerfSensitiveClass;.<init>:(LDexOutputTest;)V",
+      "LDexOutputTest$PerfSensitiveClass;.<init>:(LDexOutputTest;)V",
+      "LDexOutputTest$SecondPerfSensitiveClass;.<init>:(LDexOutputTest;)V",
+      "LDexOutputTest;.<init>:()V",
+      "LDexOutputTest$NonPerfSensitiveClass;.EjustReturnFive:()I",
+      "LDexOutputTest$PerfSensitiveClass;.EjustReturnFive:()I",
+      "LDexOutputTest$SecondPerfSensitiveClass;.EjustReturnFive:()I",
+      "LDexOutputTest;.AjustReturnFive:()I",
+      "LDexOutputTest;.EjustReturnFive:()I",
+      "LDexOutputTest;.DgetSixpublic:()I",
+      "LDexOutputTest$NonPerfSensitiveClass;.FsomeLogic:(I)I",
+      "LDexOutputTest$PerfSensitiveClass;.FsomeLogic:(I)I",
+      "LDexOutputTest$SecondPerfSensitiveClass;.FsomeLogic:(I)I",
+      "LDexOutputTest;.CsomeLogic:(I)I",
+      "LDexOutputTest;.FsomeLogic:(I)I",
+      "LDexOutputTest;.HsomeLogic:(I)I",
+      "LDexOutputTest;.BjustCallSixpublic:()I",
+      "LDexOutputTest;.GjustCallSixpublic:()I"};
+
+  EXPECT_TRUE(std::equal(method_names.begin(), method_names.end(),
+                         expected_order.begin()));
+}

test/integ/coldstart_ordering.csv

@@ -0,0 +1,10 @@
+LDexOutputTest$PerfSensitiveClass;.<init>:(LDexOutputTest;)V
+LDexOutputTest$PerfSensitiveClass;.someRandomNonExistantMethod1:(I)I
+LDexOutputTest2$Class;.someRandomMethodNotInDex:(I)I
+LDexOutputTest$SecondPerfSensitiveClass;.<init>:(LDexOutputTest;)V
+LDexOutputTest$PerfSensitiveClass;.EjustReturnFive:()I
+LDexOutputTest$PerfSensitiveClass;.someRandomNonExistantMethod2:(I)I
+LDexOutputTest$SecondPerfSensitiveClass;.EjustReturnFive:()I
+LDexOutputTest$SecondPerfSensitiveClass;.FsomeLogic:(I)I
+LDexOutputTest$PerfSensitiveClass;.FsomeLogic:(I)I
+LDexOutputTest$PerfSensitiveClass;.someRandomNonExistantMethod3:(I)I