[CodeExtractor] Split PHI nodes with incoming values from outlined region (PR39433)

[kachkov98]

Summary

If a PHI node out of extracted region has multiple incoming values from it, split this PHI on two parts. First PHI has incomings only from region and extracts with it (they are placed to the separate basic block that added to the list of outlined), and incoming values in original PHI are replaced by first PHI. Similar solution is already used in CodeExtractor for PHIs in entry block (severSplitPHINodes method). It covers PR39433 bug.

[vsk]

Thank you (very much!) for working on this.

lib/Transforms/Utils/CodeExtractor.cpp

619

Nit: please capitalize sentences in comments. Also, maybe "incoming values from the outlining region" would be clearer?

625

Why should PHIs with exactly one predecessor from the outlining region be skipped? What if there's a second exit block with another, different predecessor from the outlining region?

Out1   ->  Out2
  \          \ 
   Exit1      Exit2

1301

Why delete this assertion? ISTM that's it's still useful to have around.

[vsk]

lib/Transforms/Utils/CodeExtractor.cpp

635

Why not "for (BasicBlock &PredBB : predecessors(ExitBB))"?

650

When you replace all uses of ExitBB with NewBB for every terminator in a predecessor of ExitBB, every new PHI inserted into NewBB must have *some* incoming value (just undef?) from *each* predecessor of NewBB. Otherwise, the verifier complains (taken from a stage2 build with hot/cold splitting enabled):

PHINode should have one entry for each predecessor of its parent basic block!
  %conv97.lcssa106.ce = phi i64 [ 1, %if.then31.1 ], [ 1, %if.then31.1 ], [ 2, %if.then31.2 ], [ 2, %if.then31.2 ]
fatal error: error in backend: verification of newFunction failed!

[brzycki]

@kachkov98 thank you for helping to solve this bug. I agree with @vsk 's review comments and I'm most concerned about skipping PHIs with only one predecessor. I think it would be good to add code coverage testing in Utils/CodeExtractorTest.cpp explicitly testing your assumptions of the new code in severSplitPHINodesOfExits()

lib/Transforms/Utils/CodeExtractor.cpp

619

+1

Also please end sentences with a period .

625

Agreed, I think this needs a little more thought. At the very least it needs a comment justifying why these PHIs should not be processed.

[kachkov98]

Fix comments style
Return assertion back
Verify function consistency before outlining

[kachkov98]

lib/Transforms/Utils/CodeExtractor.cpp

625

My intention here was to not create unnecessary output parameters when this single value is constant (if we split it to separate phi, findInputsOutputs() detects that PHI value is used outside the region and creates new argument). If this value is not constant, it will be stored right after its definition. In this example one value will be stored in one argument somewhere in out1, another value will be stored in another argument in out2, than they both are reloaded in codeRepl block, but PHI in exit1 (or exit2) uses only one value from region, so it should be determenistic. I completely agree that this assumption requires additional test, I'll try to provide it soon.

635

The reason is that iterators become invalidated when terminator instuction in predecessor is updated

650

This error message is quite ambigious) I think it means here that PHI shouldn't have duplicated incoming blocks (for example, 2 %if.then31.1 blocks are not correct - how to choose value in that case?). Nevertheless, severSplitPHINodesOfExits just copies these incoming blocks from original PHI, and it means that this PHI was incorrect *before* the extraction. I can confirm that already have this failure while compiling test-suite with HotColdSplitting and PGO, and checking that function is consistent before any extraction resolved it. This check can be done in CodeExtractor (lines 1191-1195), but root cause is that some pass creates this PHI before.

[vsk]

lib/Transforms/Utils/CodeExtractor.cpp

635

Actually, as-written, there is an iterator invalidation bug here. This fixes it:

SmallVector<BasicBlock *, 4> Preds(pred_begin(ExitBB), pred_end(ExitBB));
for (BasicBlock *PredBB : Preds) { ... }

This handles the case where one (or more) of the predecessor blocks is terminated by a switch, which can have multiple uses of ExitBB. The current version of the code walks the wrong user list when this happens.

650

The verifier allows duplicate incoming blocks (as long as they provide a unique incoming value). This isn't a bug in another pass -- I hit the same verification failure with your most recent patch, which verifies oldFunction before doing any work.

The suggested iterator invalidation fix in my comment above addresses the problem.

1192

This should be wrapped in LLVM_DEBUG(), as we don't want to pay the compile-time cost for this in Release builds

There are build bots which test llvm using -verify-each (inserting the verifier after each pass), so (at least theoretically) we should safely be able to assume the input IR is correct.

[vsk]

lib/Transforms/Utils/CodeExtractor.cpp

625

I think I follow your explanation. Here's a test:

entry:
  br i1 undef, label %extract-me-1, label %exit-1

; Extracted blocks
extract-me-1:
  br i1 undef, label %exit-1, label %extract-me-2

extract-me-2:
  br label %exit-2

; Blocks that are not extracted
exit-1:
  %p1 = phi [ i8 0, entry ], [ i8 1, extract-me-1 ]
  br label %exit-2

exit-2:
  %p2 = phi [ i8 2, entry ], [ i8 1, exit-1 ]
  ret void

I suppose there is no need to split %p1 or %p2, because there's only one possible incoming value from the codeRepl block to either exit-1 or exit-2.

650

By the way, the in-tree version of the hot/cold splitting pass has several known issues. If you're interested in that pass, I have some patches up to address them: https://reviews.llvm.org/D53887, https://reviews.llvm.org/D54189, https://reviews.llvm.org/D54244. I'd certainly appreciate your feedback!

[kachkov98]

Fix issue when terminator instruction in predecessor is switch
Add unit test for PHIs with one incoming value from region

[kachkov98]

lib/Transforms/Utils/CodeExtractor.cpp

635

Seems that it's fixed now, thank you for explaining the problem!

650

Currently I'm facing the problem that hotcoldsplitting adds to ColdRegion one basic block 2 times (first time when visiting predecessors of SinkBB, second time when visitting successors) - it causes CodeExtractor check fail. Not sure that these patches don't solve the problem, but if it is interesting, I can provide test case for this (what is the best way to do it?)

[vsk]

Thank you, LGTM.

I tested this patch by:

1. Building LNT+externals with hot/cold splitting enabled. I forced outlining to occur whenever a block has more than 1 predecessor, so long as it wouldn't result in the entire function being outlined. This resulted in 48,441 cold functions being extracted/outlined. All output validation tests still passed.
2. Running check-llvm in a stage2 build with hot/cold splitting enabled in the same way described above.

lib/Transforms/Utils/CodeExtractor.cpp

650

D54189 fixes that bug, but I'm not sure whether it still applies cleanly. I'll update it soon.

[kachkov98]

Thank you! I don't have commit access, could you please submit it?

include/llvm/Transforms/Utils/CodeExtractor.h

 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_TRANSFORMS_UTILS_CODEEXTRACTOR_H
 #define LLVM_TRANSFORMS_UTILS_CODEEXTRACTOR_H
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include <limits>
 
 namespace llvm {
 
 class BasicBlock;
 class BlockFrequency;
 class BlockFrequencyInfo;
 class BranchProbabilityInfo;
public:
     /// CommonExitBlock is block outside the outline region. It is the common
     /// successor of blocks inside the region. If there exists a single block
     /// inside the region that is the predecessor of CommonExitBlock, that block
     /// will be returned. Otherwise CommonExitBlock will be split and the
     /// original block will be added to the outline region.
     BasicBlock *findOrCreateBlockForHoisting(BasicBlock *CommonExitBlock);
 
   private:
-    void severSplitPHINodes(BasicBlock *&Header);
+    void severSplitPHINodesOfEntry(BasicBlock *&Header);
+    void severSplitPHINodesOfExits(const SmallPtrSetImpl<BasicBlock *> &Exits);
     void splitReturnBlocks();
 
     Function *constructFunction(const ValueSet &inputs,
                                 const ValueSet &outputs,
                                 BasicBlock *header,
                                 BasicBlock *newRootNode, BasicBlock *newHeader,
                                 Function *oldFunction, Module *M);
 

lib/Transforms/Utils/CodeExtractor.cpp

for (Instruction &II : *BB) {
         if (!definedInRegion(Blocks, U)) {
           Outputs.insert(&II);
           break;
         }
     }
   }
 }
 
-/// severSplitPHINodes - If a PHI node has multiple inputs from outside of the
-/// region, we need to split the entry block of the region so that the PHI node
-/// is easier to deal with.
-void CodeExtractor::severSplitPHINodes(BasicBlock *&Header) {
+/// severSplitPHINodesOfEntry - If a PHI node has multiple inputs from outside
+/// of the region, we need to split the entry block of the region so that the
+/// PHI node is easier to deal with.
+void CodeExtractor::severSplitPHINodesOfEntry(BasicBlock *&Header) {
   unsigned NumPredsFromRegion = 0;
   unsigned NumPredsOutsideRegion = 0;
 
   if (Header != &Header->getParent()->getEntryBlock()) {
     PHINode *PN = dyn_cast<PHINode>(Header->begin());
     if (!PN) return;  // No PHI nodes.
 
     // If the header node contains any PHI nodes, check to see if there is more
for (AfterPHIs = OldPred->begin(); isa<PHINode>(AfterPHIs); ++AfterPHIs) {
           PN->removeIncomingValue(i);
           --i;
         }
       }
     }
   }
 }
 
+/// severSplitPHINodesOfExits - if PHI nodes in exit blocks have inputs from
+/// outlined region, we split these PHIs on two: one with inputs from region
+/// and other with remaining incoming blocks; then first PHIs are placed in
+/// outlined region.
+void CodeExtractor::severSplitPHINodesOfExits(
+    const SmallPtrSetImpl<BasicBlock *> &Exits) {
+  for (BasicBlock *ExitBB : Exits) {
+    BasicBlock *NewBB = nullptr;
+
+    for (PHINode &PN : ExitBB->phis()) {
+      // find all incoming blocks from region
+      SmallVector<unsigned, 2> IncomingVals;
+      for (unsigned i = 0; i < PN.getNumIncomingValues(); ++i)
+        if (Blocks.count(PN.getIncomingBlock(i)))
+          IncomingVals.push_back(i);
+
+      // do not process PHI if there is one (or fewer) predecessor from region
+      if (IncomingVals.size() <= 1)
+        continue;
+
+      // create block for new PHIs and add it to the list of outlined if it
+      // wasn't done before
+      if (!NewBB) {
+        NewBB = BasicBlock::Create(ExitBB->getContext(),
+                                   ExitBB->getName() + ".split",
+                                   ExitBB->getParent(), ExitBB);
+        for (pred_iterator I = pred_begin(ExitBB), E = pred_end(ExitBB);
+             I != E;) {
+          BasicBlock *PredBB = *I++;
+          if (Blocks.count(PredBB))
+            PredBB->getTerminator()->replaceUsesOfWith(ExitBB, NewBB);
+        }
+        BranchInst::Create(ExitBB, NewBB);
+        Blocks.insert(NewBB);
+      }
+
+      // split this PHI
+      PHINode *NewPN =
+          PHINode::Create(PN.getType(), IncomingVals.size(),
+                          PN.getName() + ".ce", NewBB->getFirstNonPHI());
+      for (unsigned i : IncomingVals)
+        NewPN->addIncoming(PN.getIncomingValue(i), PN.getIncomingBlock(i));
+      for (unsigned i : reverse(IncomingVals))
+        PN.removeIncomingValue(i, false);
+      PN.addIncoming(NewPN, NewBB);
+    }
+  }
+}
+
 void CodeExtractor::splitReturnBlocks() {
   for (BasicBlock *Block : Blocks)
     if (ReturnInst *RI = dyn_cast<ReturnInst>(Block->getTerminator())) {
       BasicBlock *New =
           Block->splitBasicBlock(RI->getIterator(), Block->getName() + ".ret");
       if (DT) {
         // Old dominates New. New node dominates all other nodes dominated
         // by Old.
if (!isEligible())
     return nullptr;
 
   // Assumption: this is a single-entry code region, and the header is the first
   // block in the region.
   BasicBlock *header = *Blocks.begin();
   Function *oldFunction = header->getParent();
 
   // For functions with varargs, check that varargs handling is only done in the
   // outlined function, i.e vastart and vaend are only used in outlined blocks.
   if (AllowVarArgs && oldFunction->getFunctionType()->isVarArg()) {
     auto containsVarArgIntrinsic = [](Instruction &I) {
       if (const CallInst *CI = dyn_cast<CallInst>(&I))
         if (const Function *F = CI->getCalledFunction())
           return F->getIntrinsicID() == Intrinsic::vastart ||
                  F->getIntrinsicID() == Intrinsic::vaend;
       return false;
     };
if (BFI) {
     for (BasicBlock *Pred : predecessors(header)) {
       if (Blocks.count(Pred))
         continue;
       EntryFreq +=
           BFI->getBlockFreq(Pred) * BPI->getEdgeProbability(Pred, header);
     }
   }
 
-  // If we have to split PHI nodes or the entry block, do so now.
-  severSplitPHINodes(header);
-
   // If we have any return instructions in the region, split those blocks so
   // that the return is not in the region.
   splitReturnBlocks();
 
+  // 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) {
+      if (!Blocks.count(*SI)) {
+        // Update the branch weight for this successor.
+        if (BFI) {
+          BlockFrequency &BF = ExitWeights[*SI];
+          BF += BFI->getBlockFreq(Block) * BPI->getEdgeProbability(Block, *SI);
+        }
+        ExitBlocks.insert(*SI);
+      }
+    }
+  }
+  NumExitBlocks = ExitBlocks.size();
+
+  // 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);
 
   // 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(),
Function *CodeExtractor::extractCodeRegion() {
 
   if (!HoistingCands.empty()) {
     auto *HoistToBlock = findOrCreateBlockForHoisting(CommonExit);
     Instruction *TI = HoistToBlock->getTerminator();
     for (auto *II : HoistingCands)
       cast<Instruction>(II)->moveBefore(TI);
   }
 
-  // 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) {
-      if (!Blocks.count(*SI)) {
-        // Update the branch weight for this successor.
-        if (BFI) {
-          BlockFrequency &BF = ExitWeights[*SI];
-          BF += BFI->getBlockFreq(Block) * BPI->getEdgeProbability(Block, *SI);
-        }
-        ExitBlocks.insert(*SI);
-      }
-    }
-  }
-  NumExitBlocks = ExitBlocks.size();
-
   // Construct new function based on inputs/outputs & add allocas for all defs.
   Function *newFunction = constructFunction(inputs, outputs, header,
                                             newFuncRoot,
                                             codeReplacer, oldFunction,
                                             oldFunction->getParent());
 
   // Update the entry count of the function.
   if (BFI) {
Function *CodeExtractor::extractCodeRegion() {
   // 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.
+  // Loop over all of the PHI nodes in the header and exit blocks, 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);
     for (unsigned i = 0, e = PN->getNumIncomingValues(); i != e; ++i)
       if (!Blocks.count(PN->getIncomingBlock(i)))
         PN->setIncomingBlock(i, newFuncRoot);
   }
 
-  // Look at all successors of the codeReplacer block.  If any of these blocks
-  // had PHI nodes in them, we need to update the "from" block to be the code
-  // replacer, not the original block in the extracted region.
-  for (BasicBlock *SuccBB : successors(codeReplacer)) {
-    for (PHINode &PN : SuccBB->phis()) {
-      Value *IncomingCodeReplacerVal = nullptr;
-      SmallVector<unsigned, 2> IncomingValsToRemove;
-      for (unsigned I = 0, E = PN.getNumIncomingValues(); I != E; ++I) {
-        BasicBlock *IncomingBB = PN.getIncomingBlock(I);
-
-        // Ignore incoming values from outside of the extracted region.
-        if (!Blocks.count(IncomingBB))
-          continue;
-
-        // Ensure that there is only one incoming value from codeReplacer.
-        if (!IncomingCodeReplacerVal) {
-          PN.setIncomingBlock(I, codeReplacer);
-          IncomingCodeReplacerVal = PN.getIncomingValue(I);
-        } else {
-          assert(IncomingCodeReplacerVal == PN.getIncomingValue(I) &&
-                 "PHI has two incompatbile incoming values from codeRepl");
-          IncomingValsToRemove.push_back(I);
-        }
-      }
-
-      for (unsigned I : reverse(IncomingValsToRemove))
-        PN.removeIncomingValue(I, /*DeletePHIIfEmpty=*/false);
-    }
-  }
+  for (BasicBlock *ExitBB : ExitBlocks)
+    for (PHINode &PN : ExitBB->phis())
+      for (unsigned i = 0, e = PN.getNumIncomingValues(); i != e; ++i)
+        if (Blocks.count(PN.getIncomingBlock(i)))
+          PN.setIncomingBlock(i, codeReplacer);
 
   // Erase debug info intrinsics. Variable updates within the new function are
   // invisible to debuggers. This could be improved by defining a DISubprogram
   // for the new function.
   for (BasicBlock &BB : *newFunction) {
     auto BlockIt = BB.begin();
     // Remove debug info intrinsics from the new function.
     while (BlockIt != BB.end()) {
Function *CodeExtractor::extractCodeRegion() {
   // Mark the new function `noreturn` if applicable.
   bool doesNotReturn = none_of(*newFunction, [](const BasicBlock &BB) {
     return isa<ReturnInst>(BB.getTerminator());
   });
   if (doesNotReturn)
     newFunction->setDoesNotReturn();
 
   LLVM_DEBUG(if (verifyFunction(*newFunction))
-                 report_fatal_error("verifyFunction failed!"));
+             report_fatal_error("verification of newFunction failed!"));
+  LLVM_DEBUG(if (verifyFunction(*oldFunction))
+             report_fatal_error("verification of oldFunction failed!"));
   return newFunction;
 }

test/Transforms/HotColdSplit/duplicate-phi-preds-crash.ll

 declare void @free(i8* %ptr)
 
 declare void @sink() cold
 
 ; CHECK-LABEL: define {{.*}}@realloc2(
 ; CHECK: call {{.*}}@sideeffect(
 ; CHECK: call {{.*}}@realloc(
 ; CHECK-LABEL: codeRepl:
-; CHECK-NEXT: call {{.*}}@realloc2.cold.1(i64 %size, i8* %ptr)
+; CHECK-NEXT: call {{.*}}@realloc2.cold.1(i64 %size, i8* %ptr, i8** %retval.0.ce.loc)
 ; CHECK-LABEL: cleanup:
-; CHECK-NEXT: phi i8* [ null, %if.then ], [ null, %codeRepl ], [ %call, %if.end ]
+; CHECK-NEXT: phi i8* [ null, %if.then ], [ %call, %if.end ], [ %retval.0.ce.reload, %codeRepl ]
 define i8* @realloc2(i8* %ptr, i64 %size) {
 entry:
   %0 = add i64 %size, -1
   %1 = icmp ugt i64 %0, 184549375
   br i1 %1, label %if.then, label %if.end
 
 if.then:                                          ; preds = %entry
   call void @sideeffect(i64 %size)

unittests/Transforms/Utils/CodeExtractorTest.cpp

 #include "llvm/IR/Verifier.h"
 #include "llvm/IRReader/IRReader.h"
 #include "llvm/Support/SourceMgr.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
 
 namespace {
-TEST(CodeExtractor, DISABLED_ExitStub) {
+TEST(CodeExtractor, ExitStub) {
   LLVMContext Ctx;
   SMDiagnostic Err;
   std::unique_ptr<Module> M(parseAssemblyString(R"invalid(
     define i32 @foo(i32 %x, i32 %y, i32 %z) {
     header:
       %0 = icmp ugt i32 %x, %y
       br i1 %0, label %body1, label %body2
 
     body1:
       %1 = add i32 %z, 2
       br label %notExtracted
 
     body2:
       %2 = mul i32 %z, 7
       br label %notExtracted
 
     notExtracted:
       %3 = phi i32 [ %1, %body1 ], [ %2, %body2 ]
       %4 = add i32 %3, %x
       ret i32 %4
     }
   )invalid",
                                                 Err, Ctx));
 
-  // CodeExtractor miscompiles this function. There appear to be some issues
-  // with the handling of outlined regions with live output values.
-  //
-  // In the original function, CE adds two reloads in the codeReplacer block:
-  //
-  //   codeRepl:                                         ; preds = %header
-  //     call void @foo_header.split(i32 %z, i32 %x, i32 %y, i32* %.loc, i32* %.loc1)
-  //     %.reload = load i32, i32* %.loc
-  //     %.reload2 = load i32, i32* %.loc1
-  //     br label %notExtracted
-  //
-  // These reloads must flow into the notExtracted block:
-  //
-  //   notExtracted:                                     ; preds = %codeRepl
-  //     %0 = phi i32 [ %.reload, %codeRepl ], [ %.reload2, %body2 ]
-  //
-  // The problem is that the PHI node in notExtracted now has an incoming
-  // value from a BasicBlock that's in a different function.
-
   Function *Func = M->getFunction("foo");
   SmallVector<BasicBlock *, 3> Candidates;
   for (auto &BB : *Func) {
     if (BB.getName() == "body1")
       Candidates.push_back(&BB);
     if (BB.getName() == "body2")
       Candidates.push_back(&BB);
   }
   // CodeExtractor requires the first basic block
   // to dominate all the other ones.
   Candidates.insert(Candidates.begin(), &Func->getEntryBlock());
 
   DominatorTree DT(*Func);
   CodeExtractor CE(Candidates, &DT);
   EXPECT_TRUE(CE.isEligible());
 
   Function *Outlined = CE.extractCodeRegion();
   EXPECT_TRUE(Outlined);
   EXPECT_FALSE(verifyFunction(*Outlined));
+  EXPECT_FALSE(verifyFunction(*Func));
 }
 } // end anonymous namespace