[CodeExtractor] Ensuring ordering of exitStub creation

[AndrewLitteken]

Summary

Previously the CodeExtractor created exit stubs, and the subsequent return value of the outlined function based on the order of out-of-region blocks after splitting any phi nodes, and collecting the blocks to be outlined. This could cause differences in order if there was a difference of exit block phi nodes between the two regions. This patch moves the collection of the output target blocks to be before this occurs, so that the assignment of target block to output value will be the same, regardless of the contents of the output block.

[jroelofs]

llvm/lib/Transforms/Utils/CodeExtractor.cpp

1268

Should there be an assert:

assert(NumExitBlocks < 0xffff && "too many exit blocks for switch");

or do you think that's too unlikely to matter?

1283

Doesn't the switch need the same type as the cases, or am I misunderstanding what's going on here?

1290–1291

Code following this doesn't assign into the variable, so we may as well drop the & on this one for clarity.

[AndrewLitteken]

llvm/lib/Transforms/Utils/CodeExtractor.cpp

1268

I think it's pretty unlikely that there will ever by that many exits out of a function.

1283

They do, but if there are only two exit blocks, there are only two options, a 0 and 1, returned out of the outlined function. If that's the case, the CodeExtractor simplifies this to a branch instruction anyway.

[jroelofs]

llvm/lib/Transforms/Utils/CodeExtractor.cpp

1268

👍

[paquette]

llvm/lib/Transforms/Utils/CodeExtractor.cpp

1268

I think it would be good to assert anyway if it's possible that it will happen. In the event that it does, I'd prefer to debug an assert than a later crash.

[AndrewLitteken]

Updating with assert.

[paquette]

llvm/lib/Transforms/Utils/CodeExtractor.cpp

1290–1291

Reduce indentation?

 if (Blocks.count(TI->getSuccessor(i)))
   continue;
...

1658

reduce indentation here too?

[paquette]

LGTM

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

class CodeExtractor {
     // If true, varargs functions can be extracted.
     bool AllowVarArgs;
 
     // Bits of intermediate state computed at various phases of extraction.
     SetVector<BasicBlock *> Blocks;
     unsigned NumExitBlocks = std::numeric_limits<unsigned>::max();
     Type *RetTy;
 
+    // Mapping from the original exit blocks, to the new blocks inside
+    // the function.
+    SmallVector<BasicBlock *, 4> OldTargets;
+
     // Suffix to use when creating extracted function (appended to the original
     // function name + "."). If empty, the default is to use the entry block
     // label, if non-empty, otherwise "extracted".
     std::string Suffix;
 
   public:
     /// Create a code extractor for a sequence of blocks.
     ///

llvm/lib/Transforms/Utils/CodeExtractor.cpp

#endif
   for (BasicBlock *Pred :
        llvm::make_early_inc_range(predecessors(CommonExitBlock))) {
     if (Blocks.count(Pred))
       continue;
     Pred->getTerminator()->replaceUsesOfWith(CommonExitBlock, NewExitBlock);
   }
   // Now add the old exit block to the outline region.
   Blocks.insert(CommonExitBlock);
+  OldTargets.push_back(NewExitBlock);
   return CommonExitBlock;
 }
 
 // Find the pair of life time markers for address 'Addr' that are either
 // defined inside the outline region or can legally be shrinkwrapped into the
 // outline region. If there are not other untracked uses of the address, return
 // the pair of markers if found; otherwise return a pair of nullptr.
 CodeExtractor::LifetimeMarkerInfo
CallInst *CodeExtractor::emitCallAndSwitchStatement(Function *newFunction,
 
   // Since there may be multiple exits from the original region, make the new
   // function return an unsigned, switch on that number.  This loop iterates
   // over all of the blocks in the extracted region, updating any terminator
   // instructions in the to-be-extracted region that branch to blocks that are
   // not in the region to be extracted.
   std::map<BasicBlock *, BasicBlock *> ExitBlockMap;
 
+  // Iterate over the previously collected targets, and create new blocks inside
+  // the function to branch to.
   unsigned switchVal = 0;
-  for (BasicBlock *Block : Blocks) {
-    Instruction *TI = Block->getTerminator();
-    for (unsigned i = 0, e = TI->getNumSuccessors(); i != e; ++i)
-      if (!Blocks.count(TI->getSuccessor(i))) {
-        BasicBlock *OldTarget = TI->getSuccessor(i);
-        // add a new basic block which returns the appropriate value
-        BasicBlock *&NewTarget = ExitBlockMap[OldTarget];
-        if (!NewTarget) {
-          // If we don't already have an exit stub for this non-extracted
-          // destination, create one now!
-          NewTarget = BasicBlock::Create(Context,
-                                         OldTarget->getName() + ".exitStub",
-                                         newFunction);
-          unsigned SuccNum = switchVal++;
+  for (BasicBlock *OldTarget : OldTargets) {
+    if (Blocks.count(OldTarget))
+      continue;
+    if (NewTarget)
+      continue;
+
 
-          Value *brVal = nullptr;
-          switch (NumExitBlocks) {
-          case 0:
-          case 1: break;  // No value needed.
-          case 2:         // Conditional branch, return a bool
-            brVal = ConstantInt::get(Type::getInt1Ty(Context), !SuccNum);
-            break;
-          default:
-            brVal = ConstantInt::get(Type::getInt16Ty(Context), SuccNum);
-            break;
-          }
+    assert(NumExitBlocks < 0xffff && "too many exit blocks for switch");
 
-          ReturnInst::Create(Context, brVal, NewTarget);
 
-          // Update the switch instruction.
-          TheSwitch->addCase(ConstantInt::get(Type::getInt16Ty(Context),
-                                              SuccNum),
-                             OldTarget);
-        }
 
-        // rewrite the original branch instruction with this new target
-        TI->setSuccessor(i, NewTarget);
-      }
+  for (BasicBlock *Block : Blocks) {
+    Instruction *TI = Block->getTerminator();
+    for (unsigned i = 0, e = TI->getNumSuccessors(); i != e; ++i) {
+      if (Blocks.count(TI->getSuccessor(i)))
+        continue;
+      BasicBlock *OldTarget = TI->getSuccessor(i);
+      // add a new basic block which returns the appropriate value
+      BasicBlock *NewTarget = ExitBlockMap[OldTarget];
+      assert(NewTarget && "Unknown target block!");
+
   }
 
   // Store the arguments right after the definition of output value.
   // This should be proceeded after creating exit stubs to be ensure that invoke
   // result restore will be placed in the outlined function.
   Function::arg_iterator OAI = OutputArgBegin;
   for (unsigned i = 0, e = outputs.size(); i != e; ++i) {
     auto *OutI = dyn_cast<Instruction>(outputs[i]);
for (BasicBlock *Succ : successors(Block)) {
           BF += BFI->getBlockFreq(Block) * BPI->getEdgeProbability(Block, Succ);
         }
         ExitBlocks.insert(Succ);
       }
     }
   }
   NumExitBlocks = ExitBlocks.size();
 
+  for (BasicBlock *Block : Blocks) {
+    Instruction *TI = Block->getTerminator();
+    for (unsigned i = 0, e = TI->getNumSuccessors(); i != e; ++i) {
+      if (Blocks.count(TI->getSuccessor(i)))
+        continue;
+      BasicBlock *OldTarget = TI->getSuccessor(i);
+      OldTargets.push_back(OldTarget);
+    }
+  }
+
   // If we have to split PHI nodes of the entry or exit blocks, do so now.
   severSplitPHINodesOfEntry(header);
   severSplitPHINodesOfExits(ExitBlocks);
 
   // This takes place of the original loop
   BasicBlock *codeReplacer = BasicBlock::Create(header->getContext(),
                                                 "codeRepl", oldFunction,
                                                 header);

llvm/unittests/Transforms/Utils/CodeExtractorTest.cpp

TEST(CodeExtractor, InputOutputMonitoring) {
   EXPECT_TRUE(Exit && cast<PHINode>(Exit->front()).getNumIncomingValues() == 1);
   // Ensure that there is a PHI in outlined function with 2 incoming values.
   EXPECT_TRUE(ExitSplit &&
               cast<PHINode>(ExitSplit->front()).getNumIncomingValues() == 2);
   EXPECT_FALSE(verifyFunction(*Outlined));
   EXPECT_FALSE(verifyFunction(*Func));
 }
 
+TEST(CodeExtractor, ExitBlockOrderingPhis) {
+  LLVMContext Ctx;
+  SMDiagnostic Err;
+  std::unique_ptr<Module> M(parseAssemblyString(R"invalid(
+    define void @foo(i32 %a, i32 %b) {
+    entry:
+      %0 = alloca i32, align 4
+      br label %test0
+    test0:
+      %c = load i32, i32* %0, align 4
+      br label %test1
+    test1:
+      %e = load i32, i32* %0, align 4
+      br i1 true, label %first, label %test
+    test:
+      %d = load i32, i32* %0, align 4
+      br i1 true, label %first, label %next
+    first:
+      %1 = phi i32 [ %c, %test ], [ %e, %test1 ]
+      ret void
+    next:
+      %2 = add i32 %d, 1
+      %3 = add i32 %e, 1
+      ret void
+    }
+  )invalid",
+                                                Err, Ctx));
+  Function *Func = M->getFunction("foo");
+  SmallVector<BasicBlock *, 3> Candidates{ getBlockByName(Func, "test0"),
+                                           getBlockByName(Func, "test1"),
+                                           getBlockByName(Func, "test") };
+
+  CodeExtractor CE(Candidates);
+  EXPECT_TRUE(CE.isEligible());
+
+  CodeExtractorAnalysisCache CEAC(*Func);
+  Function *Outlined = CE.extractCodeRegion(CEAC);
+  EXPECT_TRUE(Outlined);
+
+  BasicBlock *FirstExitStub = getBlockByName(Outlined, "first.exitStub");
+  BasicBlock *NextExitStub = getBlockByName(Outlined, "next.exitStub");
+
+  Instruction *FirstTerm = FirstExitStub->getTerminator();
+  ReturnInst *FirstReturn = dyn_cast<ReturnInst>(FirstTerm);
+  EXPECT_TRUE(FirstReturn);
+  ConstantInt *CIFirst = dyn_cast<ConstantInt>(FirstReturn->getReturnValue());
+  EXPECT_TRUE(CIFirst->getLimitedValue() == 1u);
+
+  Instruction *NextTerm = NextExitStub->getTerminator();
+  ReturnInst *NextReturn = dyn_cast<ReturnInst>(NextTerm);
+  EXPECT_TRUE(NextReturn);
+  ConstantInt *CINext = dyn_cast<ConstantInt>(NextReturn->getReturnValue());
+  EXPECT_TRUE(CINext->getLimitedValue() == 0u);
+  
+  EXPECT_FALSE(verifyFunction(*Outlined));
+  EXPECT_FALSE(verifyFunction(*Func));
+}
+
+TEST(CodeExtractor, ExitBlockOrdering) {
+  LLVMContext Ctx;
+  SMDiagnostic Err;
+  std::unique_ptr<Module> M(parseAssemblyString(R"invalid(
+    define void @foo(i32 %a, i32 %b) {
+    entry:
+      %0 = alloca i32, align 4
+      br label %test0
+    test0:
+      %c = load i32, i32* %0, align 4
+      br label %test1
+    test1:
+      %e = load i32, i32* %0, align 4
+      br i1 true, label %first, label %test
+    test:
+      %d = load i32, i32* %0, align 4
+      br i1 true, label %first, label %next
+    first:
+      ret void
+    next:
+      %1 = add i32 %d, 1
+      %2 = add i32 %e, 1
+      ret void
+    }
+  )invalid",
+                                                Err, Ctx));
+  Function *Func = M->getFunction("foo");
+  SmallVector<BasicBlock *, 3> Candidates{ getBlockByName(Func, "test0"),
+                                           getBlockByName(Func, "test1"),
+                                           getBlockByName(Func, "test") };
+
+  CodeExtractor CE(Candidates);
+  EXPECT_TRUE(CE.isEligible());
+
+  CodeExtractorAnalysisCache CEAC(*Func);
+  Function *Outlined = CE.extractCodeRegion(CEAC);
+  EXPECT_TRUE(Outlined);
+
+  BasicBlock *FirstExitStub = getBlockByName(Outlined, "first.exitStub");
+  BasicBlock *NextExitStub = getBlockByName(Outlined, "next.exitStub");
+
+  Instruction *FirstTerm = FirstExitStub->getTerminator();
+  ReturnInst *FirstReturn = dyn_cast<ReturnInst>(FirstTerm);
+  EXPECT_TRUE(FirstReturn);
+  ConstantInt *CIFirst = dyn_cast<ConstantInt>(FirstReturn->getReturnValue());
+  EXPECT_TRUE(CIFirst->getLimitedValue() == 1u);
+
+  Instruction *NextTerm = NextExitStub->getTerminator();
+  ReturnInst *NextReturn = dyn_cast<ReturnInst>(NextTerm);
+  EXPECT_TRUE(NextReturn);
+  ConstantInt *CINext = dyn_cast<ConstantInt>(NextReturn->getReturnValue());
+  EXPECT_TRUE(CINext->getLimitedValue() == 0u);
+  
+  EXPECT_FALSE(verifyFunction(*Outlined));
+  EXPECT_FALSE(verifyFunction(*Func));
+}
+
 TEST(CodeExtractor, ExitPHIOnePredFromRegion) {
   LLVMContext Ctx;
   SMDiagnostic Err;
   std::unique_ptr<Module> M(parseAssemblyString(R"invalid(
     define i32 @foo() {
     header:
       br i1 undef, label %extracted1, label %pred