[LoopSimplify] Update LCSSA after separating nested loops.

[mzolotukhin]

Summary

Usually LCSSA survives this transformation, but in some cases (see
attached test) it doesn't: values from the original loop after
separating might be used from the outer loop. Before the transformation
it was the same loop, so LCSSA phis were not required.

We might move this logic, along with the logic for updating LoopInfo, to
BasicBlockUtils (routine SplitBlockPredecessors), but I now see this
pretty independent on the fix in question.

This fixes PR28272.

[chandlerc]

lib/Transforms/Utils/LoopSimplify.cpp

349–351

Is it possible to track these when putting the instructions into this set? Maybe its just not worth it because its just the PHI nodes...

352

8? That seems like lots? Maybe not though

366–371

It's a bit odd to use this linear search for a unique dominating exit block... I wonder, is there some implicit guarantee that makes it make more sense? Do we instead want to perhaps build up a map from incoming BB of the outer PHI to dominating exit BB?

[sanjoy]

lib/Transforms/Utils/LoopSimplify.cpp

359

s/V/I/

376

What if you have

loop:
  %iv = phi [ 0, %entry ] [ %iv, (%loop, %left, %right)],  [ %iv.inc, %outer ]
  %iv.inc = %iv + 1
  switch(condition), %left, %right, %loop

left:
  br condition, %outer, %loop

right:
  br condition, %outer, %loop

outer:
  br %loop

then the separated loop structure will be

outer_header:
  %iv = [ 0, %entry ] , [ %iv.inc, %outer ]
  br loop

loop:
  %iv.inc = %iv + 1
  switch(condition), %left, %right, %loop

left:
  br condition, %outer, %loop

right:
  br condition, %outer, %loop

outer:
  br %loop

Now %iv.inc is live out of the inner loop and none of the inner loop
exits (%left, %right) dominate %outer.

[sanjoy]

lib/Transforms/Utils/LoopSimplify.cpp

376

Nit: the latter code snippet should be

outer_header:
  %iv = [ 0, %entry ] , [ %iv.inc, %outer ]
  br loop

loop:
  %iv.inc = %iv + 1
  switch(condition), %left, %right, %loop

left:
  br condition, %outer, %loop

right:
  br condition, %outer, %loop

outer:
  br %outer_header

[mzolotukhin]

Hi Sanjoy, Chandler,

Thanks for taking a look!
And thanks for the example, Sanjoy, I need to think over it. Probably, we'll need to insert phis if IncomingBB is reachable from the exit block (and then somehow find a merging point to insert another phi), or something like this.

Michael

[mzolotukhin]

• Use SSAUpdate to properly rewrite values.
• Break edges to exit-blocks after separating loop nest if the exit block has a pred from outside the loop.
• Add Sanjoy's test.

[mzolotukhin]

Ping.

[sanjoy]

lib/Transforms/Utils/LoopSimplify.cpp

355

Why not use predecessors ?

Actually, why not:

if (any_of(predecessors(ExitBlock), [](BasicBlock *BB) {
  return !L->contains(BB);
    }) {
rewriteLoopExitBlock(L, ExitBlock, DT, LI, PreserveLCSSA);
}

373

Why do you need to separately create a SmallSetVector here?

379

NewBB is the header of the new loop, right? (if so, a naming change
would be nice). Then I'm not sure if scanning NewBB for LCSSA
violations is sufficient -- what if we have:

loop:
  %iv = phi [ 0, %entry ] [ %iv, (%loop, %left, %right)],  [ %z, %outer ]
  switch(condition), %left, %right, %loop

left:
  %x = def()
  br condition, %outer_l, %loop

outer_l:
  br outer

right:
  %y = def()
  br condition, %outer_r, %loop

outer_r:
  br outer

outer:
  %z = phi(%x, %y)
  br %loop

then the separated structure would be

outer_l:
  %iv = phi [ 0, %entry ], [ %z, outer ]
  br loop

loop:
  switch(condition), %left, %right, %loop

left:
  %x = def()
  br condition, %outer_l, %loop

outer_l:
  br outer

right:
  %y = def()
  br condition, %outer, %loop

outer_r:
  br outer

outer:
  %z = phi(%x, %y)  ;; this PHI needs to be edited as well
  br %loop

393

When is this true?

398

s/the exit block/an exit block/
s/the phi-node/an LCSSA phi node/

[mzolotukhin]

• Modify test to catch more corner cases including a case mentioned by Sanjoy. Also, add an assert.
• Make processInstruction from LCSSA.cpp externally available.
• Use formLCSSAForInstruction routine instead of doing it manually.
• Use any_of when rewriting loop exit blocks.

[mzolotukhin]

• We don't need to include SSAUpdated any more.

[mzolotukhin]

Hi Sanjoy,

Thanks for review! You're right that we need to look not only at NewBB but in other blocks that happen to be outside the inner loop as well. I updated the test to check this case as well.

Also, I decided to make processInstruction from LCSSA externally available, since it greatly simplifies the patch. I'll submit the corresponding patch for a separate review (we'll need to iron out the API we want to have there), but I'll keep it here as well just to show you the complete picture.

Thanks,
Michael

[sanjoy]

lgtm with nits

lib/Transforms/Utils/LoopSimplify.cpp

332

Can't this just be a SmallVector ?

385

How about avoiding repeated calls to formLCSSAForInstruction on the same instruction using a set?

Given that you'll now send in a worklist of instructions anyway, maybe just unique that worklist before sending it in?

595

I was asking to use any_of instead of the loop you added. :)

If you want to clean this up also, that's fine too of course. :)

[mzolotukhin]

Thanks for reviewing, Sanjoy!

I committed this in r275891 with your remarks addressed.

Michael

lib/Transforms/Utils/LoopSimplify.cpp

332

Indeed it can. Fixed!

385

Fixed.

595

Oops, wrong spot:)

include/llvm/Transforms/Utils/LoopUtils.h

private:
   InductionKind IK;
   /// Step value.
   const SCEV *Step;
 };
 
 BasicBlock *InsertPreheaderForLoop(Loop *L, DominatorTree *DT, LoopInfo *LI,
                                    bool PreserveLCSSA);
 
+/// \brief Ensures LCSSA form for instruction Inst in the scope of loop L.
+///
+/// For the given instruction in the loop which have uses outside of the loop,
+/// an LCSSA PHI node is inserted and the uses outside the loop are rewritten to
+/// use this node.
+///
+/// LoopInfo and DominatorTree are required and preserved.
+///
+/// Returns true if any modifications are made.
+bool formLCSSAForInstruction(Loop &L, Instruction &Inst, DominatorTree &DT,
+                               const SmallVectorImpl<BasicBlock *> &ExitBlocks,
+                               PredIteratorCache &PredCache, LoopInfo *LI);
+
 /// \brief Simplify each loop in a loop nest recursively.
 ///
 /// This takes a potentially un-simplified loop L (and its children) and turns
 /// it into a simplified loop nest with preheaders and single backedges. It will
 /// update \c AliasAnalysis and \c ScalarEvolution analyses if they're non-null.
 bool simplifyLoop(Loop *L, DominatorTree *DT, LoopInfo *LI, ScalarEvolution *SE,
                   AssumptionCache *AC, bool PreserveLCSSA);
 

lib/Transforms/Utils/LCSSA.cpp

 static bool isExitBlock(BasicBlock *BB,
                         const SmallVectorImpl<BasicBlock *> &ExitBlocks) {
   return find(ExitBlocks, BB) != ExitBlocks.end();
 }
 
 /// Given an instruction in the loop, check to see if it has any uses that are
 /// outside the current loop.  If so, insert LCSSA PHI nodes and rewrite the
 /// uses.
-static bool processInstruction(Loop &L, Instruction &Inst, DominatorTree &DT,
+bool llvm::formLCSSAForInstruction(Loop &L, Instruction &Inst, DominatorTree &DT,
                                const SmallVectorImpl<BasicBlock *> &ExitBlocks,
                                PredIteratorCache &PredCache, LoopInfo *LI) {
   SmallVector<Use *, 16> UsesToRewrite;
 
   // Tokens cannot be used in PHI nodes, so we skip over them.
   // We can run into tokens which are live out of a loop with catchswitch
   // instructions in Windows EH if the catchswitch has one catchpad which
   // is inside the loop and another which is not.
for (auto *I : PostProcessPHIs) {
     SmallVector<BasicBlock *, 8> EBs;
     OtherLoop->getExitBlocks(EBs);
     if (EBs.empty())
       continue;
 
     // Recurse and re-process each PHI instruction. FIXME: we should really
     // convert this entire thing to a worklist approach where we process a
     // vector of instructions...
-    processInstruction(*OtherLoop, *I, DT, EBs, PredCache, LI);
+    formLCSSAForInstruction(*OtherLoop, *I, DT, EBs, PredCache, LI);
   }
 
   // Remove PHI nodes that did not have any uses rewritten.
   for (PHINode *PN : AddedPHIs)
     if (PN->use_empty())
       PN->eraseFromParent();
 
   return true;
for (BasicBlock *BB : L.blocks()) {
     for (Instruction &I : *BB) {
       // Reject two common cases fast: instructions with no uses (like stores)
       // and instructions with one use that is in the same block as this.
       if (I.use_empty() ||
           (I.hasOneUse() && I.user_back()->getParent() == BB &&
            !isa<PHINode>(I.user_back())))
         continue;
 
-      Changed |= processInstruction(L, I, DT, ExitBlocks, PredCache, LI);
+      Changed |= formLCSSAForInstruction(L, I, DT, ExitBlocks, PredCache, LI);
     }
   }
 
   // If we modified the code, remove any caches about the loop from SCEV to
   // avoid dangling entries.
   // FIXME: This is a big hammer, can we clear the cache more selectively?
   if (SE && Changed)
     SE->forgetLoop(&L);

lib/Transforms/Utils/LoopSimplify.cpp

 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
+#include "llvm/IR/PredIteratorCache.h"
 #include "llvm/IR/Type.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/LoopUtils.h"
+#include "llvm/Transforms/Utils/SSAUpdater.h"
 using namespace llvm;
 
 #define DEBUG_TYPE "loop-simplify"
 
 STATISTIC(NumInserted, "Number of pre-header or exit blocks inserted");
 STATISTIC(NumNested  , "Number of nested loops split out");
 
 // If the block isn't already, move the new block to right after some 'outside
static Loop *separateNestedLoop(Loop *L, BasicBlock *Preheader,
   // not part of the inner loop.
   const std::vector<Loop*> &SubLoops = L->getSubLoops();
   for (size_t I = 0; I != SubLoops.size(); )
     if (BlocksInL.count(SubLoops[I]->getHeader()))
       ++I;   // Loop remains in L
     else
       NewOuter->addChildLoop(L->removeChildLoop(SubLoops.begin() + I));
 
+  SmallSetVector<BasicBlock *, 8> OuterLoopBlocks;
+  OuterLoopBlocks.insert(NewBB);
   // Now that we know which blocks are in L and which need to be moved to
   // OuterLoop, move any blocks that need it.
   for (unsigned i = 0; i != L->getBlocks().size(); ++i) {
     BasicBlock *BB = L->getBlocks()[i];
     if (!BlocksInL.count(BB)) {
       // Move this block to the parent, updating the exit blocks sets
       L->removeBlockFromLoop(BB);
-      if ((*LI)[BB] == L)
+      if ((*LI)[BB] == L) {
         LI->changeLoopFor(BB, NewOuter);
+        OuterLoopBlocks.insert(BB);
+      }
       --i;
     }
   }
 
+  // Split edges to exit blocks from the inner loop, if they emerged in the
+  // process of separating the outer one.
+  SmallVector<BasicBlock *, 8> ExitBlocks;
+  L->getExitBlocks(ExitBlocks);
+  SmallSetVector<BasicBlock *, 8> ExitBlockSet(ExitBlocks.begin(),
+                                               ExitBlocks.end());
+  for (BasicBlock *ExitBlock : ExitBlockSet) {
+    for (pred_iterator PI = pred_begin(ExitBlock), PE = pred_end(ExitBlock);
+         PI != PE; ++PI)
+      // Must be exactly this loop: no subloops, parent loops, or non-loop preds
+      // allowed.
+      if (!L->contains(*PI)) {
+        rewriteLoopExitBlock(L, ExitBlock, DT, LI, PreserveLCSSA);
+        break;
+      }
+  }
+
+  if (PreserveLCSSA) {
+    // Fix LCSSA form for L. Some values, which previously were only used inside
+    // L, can now be used in NewOuter loop. We need to insert phi-nodes for them
+    // in corresponding exit blocks.
+    // We've just rewritten the exit blocks, so we have to construct
+    // ExitBlocksSet again.
+    SmallVector<BasicBlock *, 8> ExitBlocks;
+    L->getExitBlocks(ExitBlocks);
+
+    PredIteratorCache PredCache;
+
+    // Go through all instructions in OuterLoopBlocks and check if they are
+    // using operands from the inner loop. In this case we'll need to fix LCSSA
+    // for these instructions.
+    for (BasicBlock *OuterBB: OuterLoopBlocks) {
+      for (Instruction &I : *OuterBB) {
+        for (Value *Op : I.operands()) {
+          Instruction *OpI = dyn_cast<Instruction>(Op);
+          if (!OpI || !L->contains(OpI))
+            continue;
+          formLCSSAForInstruction(*L, *OpI, *DT, ExitBlocks, PredCache, LI);
+        }
+      }
+    }
+    assert(NewOuter->isRecursivelyLCSSAForm(*DT) &&
+           "LCSSA is broken after separating nested loops!");
+  }
+
   return NewOuter;
 }
 
 /// \brief This method is called when the specified loop has more than one
 /// backedge in it.
 ///
 /// If this occurs, revector all of these backedges to target a new basic block
 /// and have that block branch to the loop header.  This ensures that loops
 /// have exactly one backedge.
 static BasicBlock *insertUniqueBackedgeBlock(Loop *L, BasicBlock *Preheader,
                                              DominatorTree *DT, LoopInfo *LI) {
   assert(L->getNumBackEdges() > 1 && "Must have > 1 backedge!");
 
   // Get information about the loop
ReprocessLoop:
   // loop preheader/header will dominate the exit blocks.  If the exit block has
   // predecessors from outside of the loop, split the edge now.
   SmallVector<BasicBlock*, 8> ExitBlocks;
   L->getExitBlocks(ExitBlocks);
 
   SmallSetVector<BasicBlock *, 8> ExitBlockSet(ExitBlocks.begin(),
                                                ExitBlocks.end());
   for (BasicBlock *ExitBlock : ExitBlockSet) {
-    for (pred_iterator PI = pred_begin(ExitBlock), PE = pred_end(ExitBlock);
-         PI != PE; ++PI)
-      // Must be exactly this loop: no subloops, parent loops, or non-loop preds
-      // allowed.
-      if (!L->contains(*PI)) {
-        if (rewriteLoopExitBlock(L, ExitBlock, DT, LI, PreserveLCSSA)) {
-          ++NumInserted;
-          Changed = true;
-        }
-        break;
-      }
+    if (any_of(predecessors(ExitBlock),
+               [L](BasicBlock *BB) { return !L->contains(BB); })) {
+      rewriteLoopExitBlock(L, ExitBlock, DT, LI, PreserveLCSSA);
   }
 
   // If the header has more than two predecessors at this point (from the
   // preheader and from multiple backedges), we must adjust the loop.
   BasicBlock *LoopLatch = L->getLoopLatch();
   if (!LoopLatch) {
     // If this is really a nested loop, rip it out into a child loop.  Don't do
     // this for loops with a giant number of backedges, just factor them into a

test/Transforms/LoopSimplify/pr28272.ll

+; RUN: opt < %s -lcssa -loop-unroll -S | FileCheck %s
+target triple = "x86_64-unknown-linux-gnu"
+
+; Check that we don't crash on this test.
+; LoopSimplify is invoked by LoopUnroll.
+; CHECK-LABEL: @foo
+define void @foo() {
+entry:
+  br label %header
+
+header:
+  br label %loop1
+
+loop1:
+  br i1 true, label %loop1, label %bb43
+
+bb43:
+  %a = phi i32 [ undef, %loop1 ], [ 0, %bb45 ], [ %a, %bb54 ]
+  %b = phi i32 [ 0, %loop1 ], [ 1, %bb54 ], [ %c, %bb45 ]
+  br i1 true, label %bb114, label %header
+
+bb114:
+  %c = add i32 0, 1
+  %d = add i32 0, 1
+  br i1 true, label %bb45, label %bb54
+
+bb45:
+  %x = add i32 %d, 0
+  br label %bb43
+
+bb54:
+  br label %bb43
+}
+
+; CHECK-LABEL: @foo2
+define void @foo2() {
+entry:
+  br label %outer
+
+outer.loopexit:
+  br label %outer
+
+outer:
+  br label %loop1
+
+loop1:
+  br i1 true, label %loop1, label %loop2.preheader
+
+loop2.preheader:
+  %a.ph = phi i32 [ undef, %loop1 ]
+  %b.ph = phi i32 [ 0, %loop1 ]
+  br label %loop2
+
+loop2:
+  %a = phi i32 [ 0, %loop2.if.true ], [ %a, %loop2.if.false ], [ %a.ph, %loop2.preheader ], [0, %bb]
+  %b = phi i32 [ 1, %loop2.if.false ], [ %c, %loop2.if.true ], [ %b.ph, %loop2.preheader ], [%c, %bb]
+  br i1 true, label %loop2.if, label %outer.loopexit
+
+loop2.if:
+  %c = add i32 0, 1
+  switch i32 undef, label %loop2.if.false [i32 0, label %loop2.if.true
+                                       i32 1, label %bb]
+
+loop2.if.true:
+  br i1 undef, label %loop2, label %bb
+
+loop2.if.false:
+  br label %loop2
+
+bb:
+  br label %loop2
+}