[CodeExtractor] Allow extracting blocks with exception handling

[sdmitriev]

Summary

This is a CodeExtractor improvement which adds support for extracting blocks which have exception handling constructs if that is legal to do. CodeExtractor performs validation checks to ensure that extraction is legal when it finds invoke instructions or EH pads (landingpad, catchswitch, or cleanuppad) in blocks to be extracted.

I have also added an option to allow extraction of blocks with alloca instructions, but no validation is done for allocas. CodeExtractor caller has to validate it himself before allowing alloca instructions to be extracted. By default allocas are still not allowed in extraction blocks.

[rnk]

lib/Transforms/Utils/CodeExtractor.cpp

207–209

Do not attempt to recover from broken invariants. There are three reasonable options:

1. If this condition is possible but undesirable under normal operation, log with LLVM_DEBUG, dbgs(), and return {}.
2. Assert that this condition is impossible.
3. Use report_fatal_error to assert this condition is possible, but don't compile it out from release builds.

217–220

Ditto, this should probably be an assert, with the whole loop inside an NDEBUG block like it was before.

test/Transforms/CodeExtractor/inline_eh.ll

2

Please add a similar test that uses _CxxFrameHandler3 to get coverage of all the code you added to handle cleanupret, catchret, etc.

[sdmitriev]

lib/Transforms/Utils/CodeExtractor.cpp

207–209

Ok, I will log a message with dbgs(), that will preserve existing behavior. buildExtractionBlockSet was returning an empty list when input list contained a block that was not valid to extract.

217–220

This is actually an important part of the validation check which verifies than extraction is legal. So moving it under NDEBUG is undesirable since this code won’t run in release build. I will log a message and return an empty list.

test/Transforms/CodeExtractor/inline_eh.ll

2

Sure. I will add such test.

[rnk]

lgtm

include/llvm/Transforms/Utils/CodeExtractor.h

class Value;
   public:
     /// \brief Create a code extractor for a sequence of blocks.
     ///
     /// Given a sequence of basic blocks where the first block in the sequence
     /// dominates the rest, prepare a code extractor object for pulling this
     /// sequence out into its new function. When a DominatorTree is also given,
     /// extra checking and transformations are enabled. If AllowVarArgs is true,
     /// vararg functions can be extracted. This is safe, if all vararg handling
-    /// code is extracted, including vastart.
+    /// code is extracted, including vastart. If AllowAlloca is true, then
+    /// extraction of blocks containing alloca instructions would be possible,
+    /// however code extractor won't validate whether extraction is legal. 
     CodeExtractor(ArrayRef<BasicBlock *> BBs, DominatorTree *DT = nullptr,
                   bool AggregateArgs = false, BlockFrequencyInfo *BFI = nullptr,
                   BranchProbabilityInfo *BPI = nullptr,
-                  bool AllowVarArgs = false);
+                  bool AllowVarArgs = false, bool AllowAlloca = false);
 
     /// \brief Create a code extractor for a loop body.
     ///
     /// Behaves just like the generic code sequence constructor, but uses the
     /// block sequence of the loop.
     CodeExtractor(DominatorTree &DT, Loop &L, bool AggregateArgs = false,
                   BlockFrequencyInfo *BFI = nullptr,
                   BranchProbabilityInfo *BPI = nullptr);
 
-    /// \brief Check to see if a block is valid for extraction.
-    ///
-    /// Blocks containing EHPads, allocas and invokes are not valid. If
-    /// AllowVarArgs is true, blocks with vastart can be extracted. This is
-    /// safe, if all vararg handling code is extracted, including vastart.
-    static bool isBlockValidForExtraction(const BasicBlock &BB,
-                                          bool AllowVarArgs);
-
     /// \brief Perform the extraction, returning the new function.
     ///
     /// Returns zero when called on a CodeExtractor instance where isEligible
     /// returns false.
     Function *extractCodeRegion();
 
     /// \brief Test whether this code extractor is eligible.
     ///

lib/Transforms/Utils/CodeExtractor.cpp

 // for functions produced by the code extractor. This is useful when converting
 // extracted functions to pthread-based code, as only one argument (void*) can
 // be passed in to pthread_create().
 static cl::opt<bool>
 AggregateArgsOpt("aggregate-extracted-args", cl::Hidden,
                  cl::desc("Aggregate arguments to code-extracted functions"));
 
 /// \brief Test whether a block is valid for extraction.
-bool CodeExtractor::isBlockValidForExtraction(const BasicBlock &BB,
-                                              bool AllowVarArgs) {
-  // Landing pads must be in the function where they were inserted for cleanup.
-  if (BB.isEHPad())
-    return false;
+static bool isBlockValidForExtraction(const BasicBlock &BB,
+                                      const SetVector<BasicBlock *> &Result,
+                                      bool AllowVarArgs, bool AllowAlloca) {
   // taking the address of a basic block moved to another function is illegal
   if (BB.hasAddressTaken())
     return false;
 
   // don't hoist code that uses another basicblock address, as it's likely to
   // lead to unexpected behavior, like cross-function jumps
   SmallPtrSet<User const *, 16> Visited;
   SmallVector<User const *, 16> ToVisit;
if (isa<Instruction>(Curr) && cast<Instruction>(Curr)->getParent() != &BB)
       continue;
 
     for (auto const &U : Curr->operands()) {
       if (auto *UU = dyn_cast<User>(U))
         ToVisit.push_back(UU);
     }
   }
 
-  // Don't hoist code containing allocas or invokes. If explicitly requested,
-  // allow vastart.
+  // If explicitly requested, allow vastart and alloca. For invoke instructions
+  // verify that extraction is valid.
   for (BasicBlock::const_iterator I = BB.begin(), E = BB.end(); I != E; ++I) {
-    if (isa<AllocaInst>(I) || isa<InvokeInst>(I))
+    if (isa<AllocaInst>(I)) {
+       if (!AllowAlloca)
+         return false;
+       continue;
+    }
+
+    if (const auto *II = dyn_cast<InvokeInst>(I)) {
+      // Unwind destination (either a landingpad, catchswitch, or cleanuppad)
+      // must be a part of the subgraph which is being extracted.
+      if (auto *UBB = II->getUnwindDest())
+        if (!Result.count(UBB))
+          return false;
+      continue;
+    }
+
+    // All catch handlers of a catchswitch instruction as well as the unwind
+    // destination must be in the subgraph.
+    if (const auto *CSI = dyn_cast<CatchSwitchInst>(I)) {
+      if (auto *UBB = CSI->getUnwindDest())
+        if (!Result.count(UBB))
           return false;
+      for (auto *HBB : CSI->handlers())
+        if (!Result.count(const_cast<BasicBlock*>(HBB)))
+          return false;
+      continue;
+    }
+
+    // Make sure that entire catch handler is within subgraph. It is sufficient
+    // to check that catch return's block is in the list.
+    if (const auto *CPI = dyn_cast<CatchPadInst>(I)) {
+      for (const auto *U : CPI->users())
+        if (const auto *CRI = dyn_cast<CatchReturnInst>(U))
+          if (!Result.count(const_cast<BasicBlock*>(CRI->getParent())))
+            return false;
+      continue;
+    }
+
+    // And do similar checks for cleanup handler - the entire handler must be
+    // in subgraph which is going to be extracted. For cleanup return should
+    // additionally check that the unwind destination is also in the subgraph.
+    if (const auto *CPI = dyn_cast<CleanupPadInst>(I)) {
+      for (const auto *U : CPI->users())
+        if (const auto *CRI = dyn_cast<CleanupReturnInst>(U))
+          if (!Result.count(const_cast<BasicBlock*>(CRI->getParent())))
+            return false;
+      continue;
+    }
+    if (const auto *CRI = dyn_cast<CleanupReturnInst>(I)) {
+      if (auto *UBB = CRI->getUnwindDest())
+        if (!Result.count(UBB))
+          return false;
+      continue;
+    }
+
     if (const CallInst *CI = dyn_cast<CallInst>(I))
       if (const Function *F = CI->getCalledFunction())
         if (F->getIntrinsicID() == Intrinsic::vastart) {
           if (AllowVarArgs)
             continue;
           else
             return false;
         }
   }
 
   return true;
 }
 
 /// \brief Build a set of blocks to extract if the input blocks are viable.
 static SetVector<BasicBlock *>
 buildExtractionBlockSet(ArrayRef<BasicBlock *> BBs, DominatorTree *DT,
-                        bool AllowVarArgs) {
+                        bool AllowVarArgs, bool AllowAlloca) {
   assert(!BBs.empty() && "The set of blocks to extract must be non-empty");
   SetVector<BasicBlock *> Result;
 
   // Loop over the blocks, adding them to our set-vector, and aborting with an
   // empty set if we encounter invalid blocks.
   for (BasicBlock *BB : BBs) {
     // If this block is dead, don't process it.
     if (DT && !DT->isReachableFromEntry(BB))
       continue;
 
     if (!Result.insert(BB))
       llvm_unreachable("Repeated basic blocks in extraction input");
-    if (!CodeExtractor::isBlockValidForExtraction(*BB, AllowVarArgs)) {
-      Result.clear();
-      return Result;
-    }
+
+  for (auto *BB : Result) {
+    if (!isBlockValidForExtraction(*BB, Result, AllowVarArgs, AllowAlloca))
+      return {};
+
+    // Make sure that the first block is not a landing pad.
+    if (BB == Result.front()) {
+      if (BB->isEHPad()) {
+        DEBUG(dbgs() << "The first block cannot be an unwind block\n");
+        return {};
+      }
+      continue;
     }
 
-#ifndef NDEBUG
-  for (SetVector<BasicBlock *>::iterator I = std::next(Result.begin()),
-                                         E = Result.end();
-       I != E; ++I)
-    for (pred_iterator PI = pred_begin(*I), PE = pred_end(*I);
-         PI != PE; ++PI)
-      assert(Result.count(*PI) &&
-             "No blocks in this region may have entries from outside the region"
-             " except for the first block!");
-#endif
+    // All blocks other than the first must not have predecessors outside of
+    // the subgraph which is being extracted.
+    for (auto *PBB : predecessors(BB))
+      if (!Result.count(PBB)) {
+        DEBUG(dbgs() << "No blocks in this region may have entries from "
+                         "outside the region except for the first block!\n");
+        return {};
+      }
+  }
 
   return Result;
 }
 
 CodeExtractor::CodeExtractor(ArrayRef<BasicBlock *> BBs, DominatorTree *DT,
                              bool AggregateArgs, BlockFrequencyInfo *BFI,
-                             BranchProbabilityInfo *BPI, bool AllowVarArgs)
+                             BranchProbabilityInfo *BPI, bool AllowVarArgs,
+                             bool AllowAlloca)
     : DT(DT), AggregateArgs(AggregateArgs || AggregateArgsOpt), BFI(BFI),
       BPI(BPI), AllowVarArgs(AllowVarArgs),
-      Blocks(buildExtractionBlockSet(BBs, DT, AllowVarArgs)) {}
+      Blocks(buildExtractionBlockSet(BBs, DT, AllowVarArgs, AllowAlloca)) {}
 
 CodeExtractor::CodeExtractor(DominatorTree &DT, Loop &L, bool AggregateArgs,
                              BlockFrequencyInfo *BFI,
                              BranchProbabilityInfo *BPI)
     : DT(&DT), AggregateArgs(AggregateArgs || AggregateArgsOpt), BFI(BFI),
       BPI(BPI), AllowVarArgs(false),
       Blocks(buildExtractionBlockSet(L.getBlocks(), &DT,
-                                     /* AllowVarArgs */ false)) {}
+                                     /* AllowVarArgs */ false,
+                                     /* AllowAlloca */ false)) {}
 
 /// definedInRegion - Return true if the specified value is defined in the
 /// extracted region.
 static bool definedInRegion(const SetVector<BasicBlock *> &Blocks, Value *V) {
   if (Instruction *I = dyn_cast<Instruction>(V))
     if (Blocks.count(I->getParent()))
       return true;
   return false;
if (Count.hasValue())
           ProfileCount(Count.getValue(), Function::PCT_Real)); // FIXME
     BFI->setBlockFreq(codeReplacer, EntryFreq.getFrequency());
   }
 
   emitCallAndSwitchStatement(newFunction, codeReplacer, inputs, outputs);
 
   moveCodeToFunction(newFunction);
 
+  // Propagate personality info to the new function if there is one.
+  if (oldFunction->hasPersonalityFn())
+    newFunction->setPersonalityFn(oldFunction->getPersonalityFn());
+
   // Update the branch weights for the exit block.
   if (BFI && NumExitBlocks > 1)
     calculateNewCallTerminatorWeights(codeReplacer, ExitWeights, BPI);
 
   // Loop over all of the PHI nodes in the header block, and change any
   // references to the old incoming edge to be the new incoming edge.
   for (BasicBlock::iterator I = header->begin(); isa<PHINode>(I); ++I) {
     PHINode *PN = cast<PHINode>(I);

test/Transforms/CodeExtractor/inline_eh.ll

+; RUN: opt < %s -skip-partial-inlining-cost-analysis -partial-inliner -S  | FileCheck %s
+; RUN: opt < %s -skip-partial-inlining-cost-analysis -passes=partial-inliner -S  | FileCheck %s
+
+declare void @bar()
+declare i32 @__gxx_personality_v0(...)
+declare i8* @__cxa_begin_catch(i8*)
+declare void @__cxa_end_catch()
+
+define internal void @callee(i1 %cond) personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) {
+entry:
+  br i1 %cond, label %if.then, label %if.end
+
+if.then:
+  invoke void @bar()
+          to label %invoke.cont unwind label %lpad
+
+invoke.cont:
+  br label %try.cont
+
+lpad:
+  %0 = landingpad { i8*, i32 }
+          catch i8* null
+  %1 = extractvalue { i8*, i32 } %0, 0
+  %2 = extractvalue { i8*, i32 } %0, 1
+  br label %catch
+
+catch:
+  %3 = call i8* @__cxa_begin_catch(i8* %1)
+  call void @__cxa_end_catch()
+  br label %try.cont
+
+try.cont:
+  br label %if.end
+
+if.end:
+  ret void
+}
+
+define internal void @caller(i1 %cond) {
+; CHECK-LABEL: define {{.*}} @caller
+entry:
+; CHECK: entry:
+; CHECK-NEXT: br i1
+; CHECK: codeRepl.i:
+; CHECK-NEXT: call void @callee.1_{{.*}}()
+  call void @callee(i1 %cond)
+  ret void
+}
+
+; CHECK-LABEL: define {{.*}} @callee.1_{{.*}}() personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*)
+; CHECK: invoke void @bar()
+; CHECK: landingpad

test/Transforms/CodeExtractor/inline_eh_1.ll

+; RUN: opt < %s -skip-partial-inlining-cost-analysis -partial-inliner -S  | FileCheck %s
+; RUN: opt < %s -skip-partial-inlining-cost-analysis -passes=partial-inliner -S  | FileCheck %s
+
+declare dso_local void @bar()
+declare dso_local i32 @__CxxFrameHandler3(...)
+
+define internal void @callee(i1 %cond) personality i8* bitcast (i32 (...)* @__CxxFrameHandler3 to i8*) {
+entry:
+  br i1 %cond, label %if.then, label %if.end
+
+if.then:
+  invoke void @bar()
+          to label %invoke.cont unwind label %ehcleanup
+
+invoke.cont:
+  br label %try.cont
+
+ehcleanup:
+  %0 = cleanuppad within none []
+  cleanupret from %0 unwind label %catch.dispatch
+
+catch.dispatch:
+  %1 = catchswitch within none [label %catch] unwind to caller
+
+catch:
+  %2 = catchpad within %1 [i8* null, i32 64, i8* null]
+  catchret from %2 to label %catchret.dest
+
+catchret.dest:
+  br label %try.cont
+
+try.cont:
+  br label %if.end
+
+if.end:
+  ret void
+}
+
+define internal void @caller(i1 %cond) {
+; CHECK-LABEL: define {{.*}} @caller
+entry:
+; CHECK: entry:
+; CHECK-NEXT: br i1
+; CHECK: codeRepl.i:
+; CHECK-NEXT: call void @callee.1_{{.*}}()
+  call void @callee(i1 %cond)
+  ret void
+}
+
+; CHECK-LABEL: define {{.*}} @callee.1_{{.*}}() personality i8* bitcast (i32 (...)* @__CxxFrameHandler3 to i8*)
+; CHECK: invoke void @bar()
+; CHECK: cleanuppad
+; CHECK-NEXT: cleanupret
+; CHECK: catchswitch
+; CHECK: catchpad
+; CHECK-NEXT: catchret