[PartialInlining] Reduce function outlining overhead

[davidxl]

Summary

Allocas whose life range is fully contained inside the outlined region should be pushed to the outlined function. This not only reduces the number of live-in variables to the outline function (thus reduced number of parameters), but can also avoid 'escaping' the address of the locals (if not shrink wrapped in).

[davide]

This is a really important optimization, IMHO (and for my usecase). Thanks for working on this.
Some comments inline.

include/llvm/Transforms/Utils/CodeExtractor.h

99–108 ↗

(On Diff #100537)

Can you update this comment to reflect what's the Allocas argument here?
Also, probably this should be names SinkCandidates or something?

lib/Transforms/Utils/CodeExtractor.cpp

145 ↗

(On Diff #100537)

Now that I look at this more closely, findInputsOutputs is a slightly inaccurate description after you change. Maybe we can consider splitting this in a separate function, as this one is growing quite a bit?

178–188 ↗

(On Diff #100537)

I think this can be expressed much more concisely using std/llvm::any_of.

200–202 ↗

(On Diff #100537)

While you're here, I'd see if you can use a range for instead (and maybe avoid the extra assignment to value, but I'm not sure if it's possible).

817 ↗

(On Diff #100537)

I'm not entirely sure I like this bit.
dyn_cast<> could return nullptr, resulting in trouble when deferencing. If you're always guaranteed to get a function back, I'd use cast<>, (or, alternatively use if (dyn_cast<>)

[davidxl]

lib/Transforms/Utils/CodeExtractor.cpp

200–202 ↗

(On Diff #100537)

This will need begin/end methods in Instruction.

[davidxl]

addressed Davide's feedback.

[wmi]

lib/Transforms/Utils/CodeExtractor.cpp

146–167 ↗

(On Diff #100554)

Here it requires lifetime.start and lifetime.end are in the blocks to be extracted.
I guess there are also many cases for which we have lifetime.start appear in entry block and lifetime.end appear in return block, and the actual uses of the alloca object are all in the blocks to be extracted. For those cases, we can sink the lifetime.start and lifetime.end together into the extracted func. To do that, we may need to start from alloca instructions in non-extracted blocks instead of lifetime.start in extracted blocks. Since the number of non-extracted blocks will be smaller than extracted blocks, so I guess it will be more efficient?

This is the small case I try:

class A {
public:
void memfunc();
};
int cond;

void foo() {
A a;
if (cond)
return;
a.memfunc();
return;
}

void goo() {
foo();
}

The IR of foo:
define void @_Z3foov() local_unnamed_addr #0 {
entry:
%a = alloca %class.A, align 1
%0 = getelementptr inbounds %class.A, %class.A* %a, i64 0, i32 0
call void @llvm.lifetime.start.p0i8(i64 1, i8* nonnull %0) #3
%1 = load i32, i32* @cond, align 4, !tbaa !2
%tobool = icmp eq i32 %1, 0
br i1 %tobool, label %if.end, label %cleanup

if.end:
call void @_ZN1A7memfuncEv(%class.A* nonnull %a) #3
br label %cleanup

cleanup:
call void @llvm.lifetime.end.p0i8(i64 1, i8* nonnull %0) #3
ret void
}

For the testcase, the lifetime.start and lifetime.end are not in the blocks to be extracted, but it is beneficial to have the alloca sinked.

[davidxl]

yes, those cases can be and will be handled but as a follow up. This patch handles cases where locals are declared in the outlined region.

[davidxl]

To be clear, the case handled here is
void foo()
{

...
{
     A a;
     ...
 }

}

[davidxl]

Address Wei's feedback. Find alloca candidates from Entry blocks .

[wmi]

I have no other comments. See whether Davide has further comments.

lib/Transforms/Utils/CodeExtractor.cpp

783 ↗

(On Diff #100581)

They can be moved closer to their uses.

[davide]

LGTM

[davide]

(sent yesterday but forgot to hit submit).

llvm/trunk/include/llvm/Transforms/Utils/CodeExtractor.h

 namespace llvm {
 template <typename T> class ArrayRef;
   class BasicBlock;
   class BlockFrequency;
   class BlockFrequencyInfo;
   class BranchProbabilityInfo;
   class DominatorTree;
   class Function;
+  class Instruction;
   class Loop;
   class Module;
   class RegionNode;
   class Type;
   class Value;
 
   /// \brief Utility class for extracting code into a new function.
   ///
public:
     /// \brief Compute the set of input values and output values for the code.
     ///
     /// These can be used either when performing the extraction or to evaluate
     /// the expected size of a call to the extracted function. Note that this
     /// work cannot be cached between the two as once we decide to extract
     /// a code sequence, that sequence is modified, including changing these
     /// sets, before extraction occurs. These modifications won't have any
     /// significant impact on the cost however.
-    void findInputsOutputs(ValueSet &Inputs, ValueSet &Outputs) const;
+    void findInputsOutputs(ValueSet &Inputs, ValueSet &Outputs,
+                           const ValueSet &Allocas) const;
+    /// Find the set of allocas whose life ranges are contained within the
+    /// outlined region.
+    ///
+    /// Allocas which have life_time markers contained in the outlined region
+    /// should be pushed to the outlined function. The address bitcasts that
+    /// are used by the lifetime markers are also candidates for shrink-
+    /// wrapping. The instructions that need to be sinked are collected in
+    /// 'Allocas'.
+    void findAllocas(ValueSet &Allocas) const;
 
   private:
     void severSplitPHINodes(BasicBlock *&Header);
     void splitReturnBlocks();
 
     Function *constructFunction(const ValueSet &inputs,
                                 const ValueSet &outputs,
                                 BasicBlock *header,

llvm/trunk/lib/Transforms/Utils/CodeExtractor.cpp

 #include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/RegionInfo.h"
 #include "llvm/Analysis/RegionIterator.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Verifier.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/BlockFrequency.h"
 #include "llvm/Support/CommandLine.h"

 static bool definedInCaller(const SetVector<BasicBlock *> &Blocks, Value *V) {
   if (isa<Argument>(V)) return true;
   if (Instruction *I = dyn_cast<Instruction>(V))
     if (!Blocks.count(I->getParent()))
       return true;
   return false;
 }
 
-void CodeExtractor::findInputsOutputs(ValueSet &Inputs,
-                                      ValueSet &Outputs) const {
+void CodeExtractor::findAllocas(ValueSet &SinkCands) const {
+  Function *Func = (*Blocks.begin())->getParent();
+  for (BasicBlock &BB : *Func) {
+    if (Blocks.count(&BB))
+      continue;
+    for (Instruction &II : BB) {
+      auto *AI = dyn_cast<AllocaInst>(&II);
+      if (!AI)
+        continue;
+
+      // Returns true if matching life time markers are found within
+      // the outlined region.
+      auto GetLifeTimeMarkers = [&](Instruction *Addr) {
+        Instruction *LifeStart = nullptr, *LifeEnd = nullptr;
+        for (User *U : Addr->users()) {
+          if (!definedInRegion(Blocks, U))
+            return false;
+
+          IntrinsicInst *IntrInst = dyn_cast<IntrinsicInst>(U);
+          if (IntrInst) {
+            if (IntrInst->getIntrinsicID() == Intrinsic::lifetime_start)
+              LifeStart = IntrInst;
+            if (IntrInst->getIntrinsicID() == Intrinsic::lifetime_end)
+              LifeEnd = IntrInst;
+          }
+        }
+        return LifeStart && LifeEnd;
+      };
+
+      if (GetLifeTimeMarkers(AI)) {
+        SinkCands.insert(AI);
+        continue;
+      }
+
+      // Follow the bitcast:
+      Instruction *MarkerAddr = nullptr;
+      for (User *U : AI->users()) {
+        if (U->stripPointerCasts() == AI) {
+          Instruction *Bitcast = cast<Instruction>(U);
+          if (GetLifeTimeMarkers(Bitcast)) {
+            MarkerAddr = Bitcast;
+            continue;
+          }
+        }
+        if (!definedInRegion(Blocks, U)) {
+          MarkerAddr = nullptr;
+          break;
+        }
+      }
+      if (MarkerAddr) {
+        if (!definedInRegion(Blocks, MarkerAddr))
+          SinkCands.insert(MarkerAddr);
+        SinkCands.insert(AI);
+      }
+    }
+  }
+}
+
+void CodeExtractor::findInputsOutputs(ValueSet &Inputs, ValueSet &Outputs,
+                                      const ValueSet &SinkCands) const {
+
   for (BasicBlock *BB : Blocks) {
     // If a used value is defined outside the region, it's an input.  If an
     // instruction is used outside the region, it's an output.
     for (Instruction &II : *BB) {
       for (User::op_iterator OI = II.op_begin(), OE = II.op_end(); OI != OE;
-           ++OI)
-        if (definedInCaller(Blocks, *OI))
-          Inputs.insert(*OI);
+           ++OI) {
+        Value *V = *OI;
+        if (!SinkCands.count(V) && definedInCaller(Blocks, V))
+          Inputs.insert(V);
+      }
 
       for (User *U : II.users())
         if (!definedInRegion(Blocks, U)) {
           Outputs.insert(&II);
           break;
         }
     }
   }
TI->setMetadata(
       LLVMContext::MD_prof,
       MDBuilder(TI->getContext()).createBranchWeights(BranchWeights));
 }
 
 Function *CodeExtractor::extractCodeRegion() {
   if (!isEligible())
     return nullptr;
 
-  ValueSet inputs, outputs;
+  ValueSet inputs, outputs, SinkingCands;
 
   // Assumption: this is a single-entry code region, and the header is the first
   // block in the region.
   BasicBlock *header = *Blocks.begin();
 
   // Calculate the entry frequency of the new function before we change the root
   //   block.
   BlockFrequency EntryFreq;
BasicBlock *codeReplacer = BasicBlock::Create(header->getContext(),
                                                 header);
 
   // The new function needs a root node because other nodes can branch to the
   // head of the region, but the entry node of a function cannot have preds.
   BasicBlock *newFuncRoot = BasicBlock::Create(header->getContext(), 
                                                "newFuncRoot");
   newFuncRoot->getInstList().push_back(BranchInst::Create(header));
 
+  findAllocas(SinkingCands);
+
   // Find inputs to, outputs from the code region.
-  findInputsOutputs(inputs, outputs);
+  findInputsOutputs(inputs, outputs, SinkingCands);
+
+  // Now sink all instructions which only have non-phi uses inside the region
+  for (auto *II : SinkingCands)
+    cast<Instruction>(II)->moveBefore(*newFuncRoot,
+                                      newFuncRoot->getFirstInsertionPt());
 
   // Calculate the exit blocks for the extracted region and the total exit
   //  weights for each of those blocks.
   DenseMap<BasicBlock *, BlockFrequency> ExitWeights;
   SmallPtrSet<BasicBlock *, 1> ExitBlocks;
   for (BasicBlock *Block : Blocks) {
     for (succ_iterator SI = succ_begin(Block), SE = succ_end(Block); SI != SE;
          ++SI) {

llvm/trunk/test/Transforms/CodeExtractor/PartialInlineAlloca.ll

+
+;  RUN: opt < %s -partial-inliner -skip-partial-inlining-cost-analysis -S | FileCheck  %s
+ ; RUN: opt < %s -passes=partial-inliner -skip-partial-inlining-cost-analysis -S | FileCheck   %s
+
+%"class.base" = type { %"struct.base"* }
+%"struct.base" = type opaque
+
+@g = external local_unnamed_addr global i32, align 4
+
+; Function Attrs: nounwind uwtable
+define i32 @callee_sinkable_bitcast(i32 %arg) local_unnamed_addr #0 {
+; CHECK-LABEL:define{{.*}}@callee_sinkable_bitcast.{{[0-9]}}
+; CHECK: alloca
+; CHECK-NEXT: bitcast
+; CHECK: call void @llvm.lifetime
+bb:
+  %tmp = alloca  %"class.base", align 4
+  %tmp1 = bitcast %"class.base"* %tmp to i8*
+  %tmp2 = load i32, i32* @g, align 4, !tbaa !2
+  %tmp3 = add nsw i32 %tmp2, 1
+  %tmp4 = icmp slt i32 %arg, 0
+  br i1 %tmp4, label %bb6, label %bb5
+
+bb5:                                              ; preds = %bb
+  call void @llvm.lifetime.start.p0i8(i64 4, i8* nonnull %tmp1) #2
+  %tmp11 = bitcast %"class.base"* %tmp to i32*
+  store i32 %tmp3, i32* %tmp11, align 4, !tbaa !2
+  store i32 %tmp3, i32* @g, align 4, !tbaa !2
+  call void @bar(i32* nonnull %tmp11) #2
+  call void @llvm.lifetime.end.p0i8(i64 4, i8* nonnull %tmp1) #2
+  br label %bb6
+
+bb6:                                              ; preds = %bb5, %bb
+  %tmp7 = phi i32 [ 1, %bb5 ], [ 0, %bb ]
+  ret i32 %tmp7
+}
+
+declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture) #1
+
+declare void @bar(i32*) local_unnamed_addr #2
+declare void @bar2(i32*, i32*) local_unnamed_addr #1
+
+
+; Function Attrs: argmemonly nounwind
+declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture) #1
+
+; Function Attrs: nounwind uwtable
+define i32 @caller(i32 %arg) local_unnamed_addr #0 {
+bb:
+  %tmp = tail call i32 @callee_sinkable_bitcast(i32 %arg)
+  ret i32 %tmp
+}
+
+attributes #0 = { nounwind uwtable}
+attributes #1 = { argmemonly nounwind }
+attributes #2 = { nounwind }
+
+!llvm.module.flags = !{!0}
+!llvm.ident = !{!1}
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{!"clang version 5.0.0 (trunk 303574)"}
+!2 = !{!3, !3, i64 0}
+!3 = !{!"int", !4, i64 0}
+!4 = !{!"omnipotent char", !5, i64 0}
+!5 = !{!"Simple C/C++ TBAA"}
+
+

llvm/trunk/test/Transforms/CodeExtractor/PartialInlineAlloca2.ll

+; RUN: opt < %s -partial-inliner -skip-partial-inlining-cost-analysis -S | FileCheck  %s
+; RUN: opt < %s -passes=partial-inliner -skip-partial-inlining-cost-analysis -S | FileCheck   %s
+
+%"class.base" = type { %"struct.base"* }
+%"struct.base" = type opaque
+
+@g = external local_unnamed_addr global i32, align 4
+
+define i32 @callee_no_bitcast(i32 %arg) local_unnamed_addr #0 {
+; CHECK-LABEL:define{{.*}}@callee_no_bitcast.{{[0-9]}}
+; CHECK: alloca
+; CHECK: call void @llvm.lifetime
+bb:
+  %tmp = alloca i8, align 4
+  %tmp2 = load i32, i32* @g, align 4, !tbaa !2
+  %tmp3 = add nsw i32 %tmp2, 1
+  %tmp4 = icmp slt i32 %arg, 0
+  br i1 %tmp4, label %bb6, label %bb5
+
+bb5:                                              ; preds = %bb
+  call void @llvm.lifetime.start.p0i8(i64 4, i8* nonnull %tmp) #2
+  store i32 %tmp3, i32* @g, align 4, !tbaa !2
+  %tmp11 = bitcast i8 * %tmp to i32*
+  call void @bar(i32* nonnull %tmp11) #2
+  call void @llvm.lifetime.end.p0i8(i64 4, i8* nonnull %tmp) #2
+  br label %bb6
+
+bb6:                                              ; preds = %bb5, %bb
+  %tmp7 = phi i32 [ 1, %bb5 ], [ 0, %bb ]
+  ret i32 %tmp7
+}
+
+; Function Attrs: argmemonly nounwind
+declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture) #1
+
+declare void @bar(i32*) local_unnamed_addr #2
+declare void @bar2(i32*, i32*) local_unnamed_addr #1
+
+
+; Function Attrs: argmemonly nounwind
+declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture) #1
+
+; Function Attrs: nounwind uwtable
+define i32 @caller(i32 %arg) local_unnamed_addr #0 {
+bb:
+  %tmp = tail call i32 @callee_no_bitcast(i32 %arg)
+  ret i32 %tmp
+}
+
+attributes #0 = { nounwind uwtable}
+attributes #1 = { argmemonly nounwind }
+attributes #2 = { nounwind }
+
+!llvm.module.flags = !{!0}
+!llvm.ident = !{!1}
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{!"clang version 5.0.0 (trunk 303574)"}
+!2 = !{!3, !3, i64 0}
+!3 = !{!"int", !4, i64 0}
+!4 = !{!"omnipotent char", !5, i64 0}
+!5 = !{!"Simple C/C++ TBAA"}
+
+
+

llvm/trunk/test/Transforms/CodeExtractor/PartialInlineAlloca4.ll

+; RUN: opt < %s -partial-inliner -skip-partial-inlining-cost-analysis -S | FileCheck  %s
+; RUN: opt < %s -passes=partial-inliner -skip-partial-inlining-cost-analysis -S | FileCheck   %s
+
+%"class.base" = type { %"struct.base"* }
+%"struct.base" = type opaque
+
+@g = external local_unnamed_addr global i32, align 4
+
+define i32 @callee_unknown_use1(i32 %arg) local_unnamed_addr #0 {
+; CHECK-LABEL:define{{.*}}@callee_unknown_use1.{{[0-9]}}
+; CHECK-NOT: alloca
+; CHECK: call void @llvm.lifetime
+bb:
+  %tmp = alloca  i8, align 4
+  %tmp2 = load i32, i32* @g, align 4, !tbaa !2
+  %tmp3 = add nsw i32 %tmp2, 1
+  %tmp4 = icmp slt i32 %arg, 0
+  br i1 %tmp4, label %bb6, label %bb5
+
+bb5:                                              ; preds = %bb
+  call void @llvm.lifetime.start.p0i8(i64 4, i8* nonnull %tmp) #2
+  store i32 %tmp3, i32* @g, align 4, !tbaa !2
+  %tmp11 = bitcast i8* %tmp to i32*
+  call void @bar(i32* nonnull %tmp11) #2
+  call void @llvm.lifetime.end.p0i8(i64 4, i8* nonnull %tmp) #2
+  br label %bb6
+
+bb6:                                              ; preds = %bb5, %bb
+  %tmp7 = phi i32 [ 1, %bb5 ], [ 0, %bb ]
+  %tmp1 = bitcast i8* %tmp to i32*
+  ret i32 %tmp7
+}
+
+
+; Function Attrs: argmemonly nounwind
+declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture) #1
+
+declare void @bar(i32*) local_unnamed_addr #2
+declare void @bar2(i32*, i32*) local_unnamed_addr #1
+
+
+; Function Attrs: argmemonly nounwind
+declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture) #1
+
+; Function Attrs: nounwind uwtable
+define i32 @caller(i32 %arg) local_unnamed_addr #0 {
+bb:
+  %tmp = tail call i32 @callee_unknown_use1(i32 %arg)
+  ret i32 %tmp
+}
+
+attributes #0 = { nounwind uwtable}
+attributes #1 = { argmemonly nounwind }
+attributes #2 = { nounwind }
+
+!llvm.module.flags = !{!0}
+!llvm.ident = !{!1}
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{!"clang version 5.0.0 (trunk 303574)"}
+!2 = !{!3, !3, i64 0}
+!3 = !{!"int", !4, i64 0}
+!4 = !{!"omnipotent char", !5, i64 0}
+!5 = !{!"Simple C/C++ TBAA"}
+
+
+

llvm/trunk/test/Transforms/CodeExtractor/PartialInlineAlloca5.ll

+; RUN: opt < %s -partial-inliner -skip-partial-inlining-cost-analysis -S | FileCheck  %s
+; RUN: opt < %s -passes=partial-inliner -skip-partial-inlining-cost-analysis -S | FileCheck   %s
+
+%"class.base" = type { %"struct.base"* }
+%"struct.base" = type opaque
+
+@g = external local_unnamed_addr global i32, align 4
+
+define i32 @callee_unknown_use2(i32 %arg) local_unnamed_addr #0 {
+; CHECK-LABEL:define{{.*}}@callee_unknown_use2.{{[0-9]}}
+; CHECK-NOT: alloca
+; CHECK: call void @llvm.lifetime
+bb:
+  %tmp = alloca i32, align 4
+  %tmp1 = bitcast i32* %tmp to i8*
+  %tmp2 = load i32, i32* @g, align 4, !tbaa !2
+  %tmp3 = add nsw i32 %tmp2, 1
+  %tmp4 = icmp slt i32 %arg, 0
+  br i1 %tmp4, label %bb6, label %bb5
+
+bb5:                                              ; preds = %bb
+  call void @llvm.lifetime.start.p0i8(i64 4, i8* nonnull %tmp1) #2
+  store i32 %tmp3, i32* %tmp, align 4, !tbaa !2
+  store i32 %tmp3, i32* @g, align 4, !tbaa !2
+  call void @bar(i32* nonnull %tmp) #2
+  call void @llvm.lifetime.end.p0i8(i64 4, i8* nonnull %tmp1) #2
+  br label %bb6
+
+bb6:                                              ; preds = %bb5, %bb
+  %tmp7 = phi i32 [ 1, %bb5 ], [ 0, %bb ]
+  %tmp10 = bitcast i8* %tmp1 to i32*
+  ret i32 %tmp7
+}
+
+; Function Attrs: argmemonly nounwind
+declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture) #1
+
+declare void @bar(i32*) local_unnamed_addr #2
+declare void @bar2(i32*, i32*) local_unnamed_addr #1
+
+
+; Function Attrs: argmemonly nounwind
+declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture) #1
+
+; Function Attrs: nounwind uwtable
+define i32 @caller(i32 %arg) local_unnamed_addr #0 {
+bb:
+  %tmp = tail call i32 @callee_unknown_use2(i32 %arg)
+  ret i32 %tmp
+}
+
+attributes #0 = { nounwind uwtable}
+attributes #1 = { argmemonly nounwind }
+attributes #2 = { nounwind }
+
+!llvm.module.flags = !{!0}
+!llvm.ident = !{!1}
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{!"clang version 5.0.0 (trunk 303574)"}
+!2 = !{!3, !3, i64 0}
+!3 = !{!"int", !4, i64 0}
+!4 = !{!"omnipotent char", !5, i64 0}
+!5 = !{!"Simple C/C++ TBAA"}
+
+
+