Adding partial inlining option

Summary:
This adds an option to the inliner to do "partial inlining" of hot callsites, where the callee has a hot "prefix", and cold "suffix".

For example (which is similar to the test case), in the following code snippet...
```
   // hot code
   bar(o);
   ...
   void bar(Object o) {
      // hot code
      if (o == nullptr) return;
      // cold code
```
the invocation to "bar" can get partially inlined to the following rewritten code, effectively only inlining the "hot" pure callee code, with a fallback invocation to the actual callee method if the cold code path is ever taken:
```
   // hot code
   if (o == nullptr) goto L1;
   // fallback cold code
   bar(o);
 L1:

   ...
   void bar(Object o) {
      // hot code
      if (o == nullptr) return;
      // cold code
   ...
   void bar(Object o) {
     // unchanged
```

A few more details:
- We consider a block as hot if it was hit in any interaction
- In a callee, we assume that code leading to a `throw` instruction is effectively cold. (To make up for that we don't have frequency information for basic blocks.)
- We want the partially inlined code to be of a certain maximum size
- We only consider for partially inlining "pure" instructions that don't mutate state, so that we can simply invoke the original callee in the fallback-branch if necessary

The new behavior is disabled by default, so this is a behavior-preserving change.

Reviewed By: agampe

Differential Revision: D68050082

fbshipit-source-id: 658185cc6726d664404a5354892bfeefd8013f23

Makefile.am

@@ -311,6 +311,7 @@ libredex_la_SOURCES = \
 	service/method-inliner/LegacyInliner.cpp \
 	service/method-inliner/MethodInliner.cpp \
 	service/method-inliner/ObjectInlinePlugin.cpp \
+	service/method-inliner/PartialInliner.cpp \
 	service/method-inliner/RecursionPruner.cpp \
 	service/method-merger/MethodMerger.cpp \
 	service/method-outliner/OutliningProfileGuidanceImpl.cpp \

libredex/ControlFlow.cpp

@@ -2675,10 +2675,10 @@ Block* ControlFlowGraph::split_block(const cfg::InstructionIterator& it) {
 Block* ControlFlowGraph::split_block_before(Block* old_block,
                                             const IRList::iterator& raw_it) {
   always_assert(editable());
-  // Do not split in front of special move instructions. This would likely end
-  // up being illegal.
-  always_assert(!opcode::is_a_move_result(raw_it->insn->opcode()) &&
-                !opcode::is_a_move_result_pseudo(raw_it->insn->opcode()));
+  // Do not split in front of special move-result instructions. This would
+  // likely end up being illegal.
+  always_assert(raw_it == old_block->end() || raw_it->type != MFLOW_OPCODE ||
+                !opcode::is_move_result_any(raw_it->insn->opcode()));
 
   // new_block will be the predecessor.
   Block* new_block = create_block();

libredex/InlinerConfig.h

@@ -40,6 +40,7 @@ struct InlinerConfig {
   bool debug{false};
   bool check_min_sdk_refs{true};
   bool rewrite_invoke_super{false};
+  bool partial_hot_hot_inline{false};
 
   /*
    * Some versions of ART (5.0.0 - 5.0.2) will fail to verify a method if it

libredex/SourceBlocks.cpp

@@ -1526,7 +1526,7 @@ void insert_synthetic_source_blocks_in_method(
 }
 
 std::unique_ptr<SourceBlock> clone_as_synthetic(SourceBlock* sb,
-                                                DexMethod* ref,
+                                                const DexMethod* ref,
                                                 const SourceBlock::Val& val) {
   std::unique_ptr<SourceBlock> new_sb = std::make_unique<SourceBlock>(*sb);
   new_sb->next.reset();
@@ -1542,7 +1542,7 @@ std::unique_ptr<SourceBlock> clone_as_synthetic(SourceBlock* sb,
 
 std::unique_ptr<SourceBlock> clone_as_synthetic(
     SourceBlock* sb,
-    DexMethod* ref,
+    const DexMethod* ref,
     const std::optional<SourceBlock::Val>& opt_val) {
   std::unique_ptr<SourceBlock> new_sb = std::make_unique<SourceBlock>(*sb);
   new_sb->next.reset();
@@ -1559,7 +1559,9 @@ std::unique_ptr<SourceBlock> clone_as_synthetic(
 }
 
 std::unique_ptr<SourceBlock> clone_as_synthetic(
-    SourceBlock* sb, DexMethod* ref, const std::vector<SourceBlock*>& many) {
+    SourceBlock* sb,
+    const DexMethod* ref,
+    const std::vector<SourceBlock*>& many) {
   std::unique_ptr<SourceBlock> new_sb = std::make_unique<SourceBlock>(*sb);
   new_sb->next.reset();
   new_sb->id = SourceBlock::kSyntheticId;

libredex/SourceBlocksUtils.h

@@ -14,15 +14,17 @@ class DexMethod;
 namespace source_blocks {
 
 std::unique_ptr<SourceBlock> clone_as_synthetic(SourceBlock* sb,
-                                                DexMethod* ref,
+                                                const DexMethod* ref,
                                                 const SourceBlock::Val& val);
 
 std::unique_ptr<SourceBlock> clone_as_synthetic(
     SourceBlock* sb,
-    DexMethod* ref = nullptr,
+    const DexMethod* ref = nullptr,
     const std::optional<SourceBlock::Val>& opt_val = std::nullopt);
 
 std::unique_ptr<SourceBlock> clone_as_synthetic(
-    SourceBlock* sb, DexMethod* ref, const std::vector<SourceBlock*>& many);
+    SourceBlock* sb,
+    const DexMethod* ref,
+    const std::vector<SourceBlock*>& many);
 
 } // namespace source_blocks

opt/methodinline/BridgeSynthInlinePass.cpp

@@ -17,6 +17,7 @@ void BridgeSynthInlinePass::run_pass(DexStoresVector& stores,
                        conf,
                        DEFAULT_COST_CONFIG,
                        /* consider_hot_cold */ false,
+                       /* partial_hot_hot */ false,
                        /* intra_dex */ false,
                        /* inline_for_speed */ nullptr,
                        /* inline_bridge_synth_only */ true);

opt/methodinline/IntraDexInlinePass.cpp

@@ -11,14 +11,22 @@
 
 void IntraDexInlinePass::bind_config() {
   bind("consider_hot_cold", false, m_consider_hot_cold);
+  bind("partial_hot_hot", false, m_partial_hot_hot);
 }
 
 void IntraDexInlinePass::run_pass(DexStoresVector& stores,
                                   ConfigFiles& conf,
                                   PassManager& mgr) {
   inliner::run_inliner(stores, mgr, conf, DEFAULT_COST_CONFIG,
-                       m_consider_hot_cold,
+                       m_consider_hot_cold, m_partial_hot_hot,
                        /* intra_dex */ true);
+  // For partial inlining, we only consider the first time the pass runs, to
+  // avoid repeated partial inlining. (This shouldn't be necessary as the
+  // partial inlining fallback invocation is marked as cold, but just in case
+  // some other Redex optimization disturbs that hotness data.)
+  if (m_partial_hot_hot) {
+    m_partial_hot_hot = false;
+  }
 }
 
 static IntraDexInlinePass s_pass;

opt/methodinline/IntraDexInlinePass.h

@@ -35,4 +35,5 @@ class IntraDexInlinePass : public Pass {
 
  private:
   bool m_consider_hot_cold;
+  bool m_partial_hot_hot;
 };

opt/methodinline/LocalMethodInlinePass.cpp

@@ -32,6 +32,7 @@ void LocalMethodInlinePass::run_pass(DexStoresVector& stores,
                                      PassManager& mgr) {
   inliner::run_inliner(stores, mgr, conf, DEFAULT_COST_CONFIG,
                        /* consider_hot_cold */ false,
+                       /* partial_hot_hot */ false,
                        /* intra_dex */ false,
                        /* inline_for_speed */ nullptr,
                        /* inline_bridge_synth_only */ false,

opt/methodinline/MethodInlinePass.cpp

@@ -103,13 +103,21 @@ void MethodInlinePass::bind_config() {
        DEFAULT_COST_CONFIG.unused_arg_not_top_multiplier,
        m_inliner_cost_config.unused_arg_not_top_multiplier);
   bind("consider_hot_cold", false, m_consider_hot_cold);
+  bind("partial_hot_hot", false, m_partial_hot_hot);
 }
 
 void MethodInlinePass::run_pass(DexStoresVector& stores,
                                 ConfigFiles& conf,
                                 PassManager& mgr) {
   inliner::run_inliner(stores, mgr, conf, m_inliner_cost_config,
-                       m_consider_hot_cold);
+                       m_consider_hot_cold, m_partial_hot_hot);
+  // For partial inlining, we only consider the first time the pass runs, to
+  // avoid repeated partial inlining. (This shouldn't be necessary as the
+  // partial inlining fallback invocation is marked as cold, but just in case
+  // some other Redex optimization disturbs that hotness data.)
+  if (m_partial_hot_hot) {
+    m_partial_hot_hot = false;
+  }
 }
 
 static MethodInlinePass s_pass;

opt/methodinline/MethodInlinePass.h

@@ -35,4 +35,5 @@ class MethodInlinePass : public Pass {
  private:
   InlinerCostConfig m_inliner_cost_config;
   bool m_consider_hot_cold;
+  bool m_partial_hot_hot;
 };

opt/methodinline/PerfMethodInlinePass.cpp

@@ -1136,6 +1136,7 @@ void PerfMethodInlinePass::run_pass(DexStoresVector& stores,
 
   inliner::run_inliner(stores, mgr, conf, DEFAULT_COST_CONFIG,
                        /* consider_hot_cold */ false,
+                       /* partial_hot_hot */ false,
                        /* intra_dex */ true,
                        /* inline_for_speed= */ ifs.get());
 

service/method-inliner/CFGInliner.cpp

@@ -20,6 +20,10 @@
 
 namespace cfg {
 
+const DexString* get_partial_inline_source() {
+  return DexString::make_string("PartiallyInlinedSource");
+}
+
 // TODO:
 //  * should this really be a friend class to ControlFlowGraph, Block, and Edge?
 
@@ -554,13 +558,21 @@ void CFGInliner::add_callee_throws_to_caller(
 
 void CFGInliner::set_dbg_pos_parents(ControlFlowGraph* callee,
                                      DexPosition* callsite_dbg_pos) {
+  auto* partial_inline_source = get_partial_inline_source();
+
   for (auto& entry : callee->m_blocks) {
     Block* b = entry.second;
     for (auto& mie : *b) {
       // Don't overwrite existing parent pointers because those are probably
       // methods that were inlined into callee before
       if (mie.type == MFLOW_POSITION && mie.pos->parent == nullptr) {
-        mie.pos->parent = callsite_dbg_pos;
+        // Deal with specially marked position that represents partially inlined
+        // fallback invocation.
+        if (mie.pos->file == partial_inline_source) {
+          mie.pos = std::make_unique<DexPosition>(*callsite_dbg_pos);
+        } else {
+          mie.pos->parent = callsite_dbg_pos;
+        }
       }
     }
   }

service/method-inliner/CFGInliner.h

@@ -13,6 +13,10 @@
 
 namespace cfg {
 
+// Special marker source file name used by partial inlining to make the position
+// of the fallback code.
+const DexString* get_partial_inline_source();
+
 class CFGInlinerPlugin;
 class CFGInliner {
  public: