Initiate class merging model in InterDexReshufflePass

Summary:
Initiate the needed class merging model in the `InterDexReshufflePass` and turn on mergeability-aware algo in the first run of the pass.
We currently use the merging model only to get the mapping from mergeable classes to their merging types.
In the future, we might use it for more accurate info on which methods can be deduped, and which methods can be merged into the same method, etc.

Reviewed By: thezhangwei

Differential Revision: D55593486

fbshipit-source-id: a26e443300422a86bba54158837ce4dd41e228c9

opt/class-merging/ModelSpecGenerator.cpp

@@ -6,6 +6,7 @@
  */
 
 #include "ModelSpecGenerator.h"
+#include "ClassMerging.h"
 #include "Model.h"
 #include "PassManager.h"
 #include "ReflectionAnalysis.h"
@@ -217,4 +218,30 @@ void find_all_mergeables_and_roots(const TypeSystem& type_system,
         merging_spec->merging_targets.size(), merging_spec->roots.size());
 }
 
+class_merging::Model construct_global_model(DexClasses& scope,
+                                            PassManager& mgr,
+                                            ConfigFiles& conf,
+                                            DexStoresVector& stores) {
+  class_merging::ModelSpec m_merging_spec;
+  // The specs below should match those used by IntraDexClassMerging
+  m_merging_spec.use_stable_shape_names = true;
+  m_merging_spec.interdex_config.init_type("non-ordered-set");
+  m_merging_spec.interdex_config.init_inferring_mode("class-loads");
+  m_merging_spec.dedup_fill_in_stack_trace = false;
+  size_t global_min_count = 4;
+
+  // The specs below removes dex boundaries and max size of mergers.
+  m_merging_spec.per_dex_grouping = false;
+  m_merging_spec.strategy = class_merging::strategy::BY_CLASS_COUNT;
+  m_merging_spec.min_count = 2;
+  m_merging_spec.max_count = std::numeric_limits<size_t>::max();
+
+  TypeSystem type_system(scope);
+  find_all_mergeables_and_roots(type_system, scope,
+                                /*global_min_count=*/global_min_count, mgr,
+                                &m_merging_spec);
+  return class_merging::construct_model(type_system, scope, conf, mgr, stores,
+                                        m_merging_spec);
+};
+
 } // namespace class_merging

opt/class-merging/ModelSpecGenerator.h

@@ -17,11 +17,15 @@ class DexType;
 class PassManager;
 class DexStore;
 class TypeSystem;
+struct ConfigFiles;
 using DexStoresVector = std::vector<DexStore>;
 using Scope = std::vector<DexClass*>;
+using DexClasses = std::vector<DexClass*>;
 
 namespace class_merging {
 
+class Model;
+
 struct ModelSpec;
 
 /**
@@ -37,4 +41,16 @@ void find_all_mergeables_and_roots(const TypeSystem& type_system,
                                    PassManager& mgr,
                                    ModelSpec* merging_spec);
 
+/**
+ * Construct a merging model that disregards dex boundaries and max
+ * size per merger. It is used by the mergeability-aware version of
+ * the InterDexReshuffle pass, which is run right before
+ * IntraDexClassMerging. As a result, certain merging specs used are
+ * set to match those used by the IntraDexClassMerging pass.
+ */
+class_merging::Model construct_global_model(DexClasses& scope,
+                                            PassManager& mgr,
+                                            ConfigFiles& conf,
+                                            DexStoresVector& stores);
+
 } // namespace class_merging

opt/interdex/InterDexReshufflePass.cpp

@@ -6,13 +6,14 @@
  */
 
 #include "InterDexReshufflePass.h"
-
+#include "ClassMerging.h"
 #include "ConfigFiles.h"
 #include "DedupStrings.h"
 #include "DexClass.h"
 #include "DexStructure.h"
 #include "DexUtil.h"
 #include "InterDexPass.h"
+#include "ModelSpecGenerator.h"
 #include "PassManager.h"
 #include "Show.h"
 #include "StlUtil.h"
@@ -69,10 +70,30 @@ void InterDexReshufflePass::run_pass(DexStoresVector& stores,
     }
   }
 
-  InterDexReshuffleImpl impl(conf, mgr, m_config, original_scope, root_dexen,
-                             untouched_dexes);
-  impl.compute_plan();
-  impl.apply_plan();
+  auto has_IDCM_pass = mgr.find_pass("IntraDexClassMergingPass");
+  // The mergeability_aware reshuffle algorithm is only enabled when 1) there
+  // will be a IDCM pass; 2) m_allow_mergeability_aware is set true; and 3) This
+  // is the first run of InterDexReshufflePass.
+  bool enable_mergeability_aware_reshuffle =
+      has_IDCM_pass && m_allow_mergeability_aware && m_run == 0;
+  if (enable_mergeability_aware_reshuffle) {
+    TRACE(PM, 1, "Run regular mergeability-aware InterDexReshuffle");
+    mgr.incr_metric("Mergeability_aware", 1);
+    class_merging::Model merging_model = class_merging::construct_global_model(
+        original_scope, mgr, conf, stores);
+
+    InterDexReshuffleImpl impl(conf, mgr, m_config, original_scope, root_dexen,
+                               untouched_dexes, merging_model);
+    impl.compute_plan();
+    impl.apply_plan();
+  } else {
+    TRACE(PM, 1, "Run regular InterDexReshuffle");
+    mgr.incr_metric("Mergeability_aware", 0);
+    InterDexReshuffleImpl impl(conf, mgr, m_config, original_scope, root_dexen,
+                               untouched_dexes);
+    impl.compute_plan();
+    impl.apply_plan();
+  }
 
   // Sanity check
   std::unordered_set<DexClass*> original_scope_set(original_scope.begin(),
@@ -84,6 +105,8 @@ void InterDexReshufflePass::run_pass(DexStoresVector& stores,
   for (auto cls : original_scope_set) {
     always_assert(new_scope_set.count(cls));
   }
+
+  ++m_run;
 }
 
 static InterDexReshufflePass s_pass;

opt/interdex/InterDexReshufflePass.h

@@ -95,8 +95,14 @@ class InterDexReshufflePass : public Pass {
          "Whether to exclude coldstart classes in between 1pctColdStart and "
          "20pctColdStart marker"
          "from the reshuffle.");
+    bind("allow_mergeability_aware", true, m_allow_mergeability_aware);
   }
 
  private:
   ReshuffleConfig m_config;
+
+  // Which iteration of `run_pass`.
+  size_t m_run{0};
+
+  bool m_allow_mergeability_aware;
 };

service/class-merging/ClassMerging.cpp

@@ -148,4 +148,26 @@ ModelStats merge_model(const TypeSystem& type_system,
   return stats;
 }
 
+Model construct_model(const TypeSystem& type_system,
+                      Scope& scope,
+                      ConfigFiles& conf,
+                      PassManager& mgr,
+                      DexStoresVector& stores,
+                      ModelSpec& spec) {
+  TRACE(CLMG,
+        2,
+        "[ClassMerging] merging %s model merging targets %zu roots %zu",
+        spec.name.c_str(),
+        spec.merging_targets.size(),
+        spec.roots.size());
+  Timer t("erase_model");
+  int32_t min_sdk = mgr.get_redex_options().min_sdk;
+  XStoreRefs xstores(stores);
+  auto refchecker =
+      create_ref_checker(spec.per_dex_grouping, &xstores, conf, min_sdk);
+  auto model =
+      Model::build_model(scope, stores, conf, spec, type_system, *refchecker);
+  return model;
+}
+
 } // namespace class_merging

service/class-merging/ClassMerging.h

@@ -24,4 +24,10 @@ ModelStats merge_model(const TypeSystem&,
                        DexStoresVector& stores,
                        ModelSpec& spec);
 
+Model construct_model(const TypeSystem& type_system,
+                      Scope& scope,
+                      ConfigFiles& conf,
+                      PassManager& mgr,
+                      DexStoresVector& stores,
+                      ModelSpec& spec);
 } // namespace class_merging