[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

344–346

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...

347

8? That seems like lots? Maybe not though

361–366

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

354

s/V/I/

371

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

371

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

350

Why not use predecessors ?

Actually, why not:

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

368

Why do you need to separately create a SmallSetVector here?

374

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

388

When is this true?

393

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

331

Can't this just be a SmallVector ?

380

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?

627

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

331

Indeed it can. Fixed!

380

Fixed.

627

Oops, wrong spot:)

lib/Transforms/Utils/LoopSimplify.cpp

 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.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,
   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));
 
   // 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)
         LI->changeLoopFor(BB, NewOuter);
       --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);
+    SmallSetVector<BasicBlock *, 8> ExitBlockSet(ExitBlocks.begin(),
+                                                 ExitBlocks.end());
+
+    // Go through all instructions in NewBB and check if they are using operands
+    // from the inner loop. In this case we'll need to insert LCSSA phi-nodes to
+    // preserve LCSSA.
+    for (Instruction &NewI : *NewBB) {
+      PHINode *PN = dyn_cast<PHINode>(&NewI);
+      if (!PN)
+        break;
+      for (unsigned Idx = 0, E = PN->getNumIncomingValues(); Idx != E; ++Idx) {
+        Instruction *I = dyn_cast<Instruction>(PN->getIncomingValue(Idx));
+        if (!I || !L->contains(I))
+          continue;
+
+        // We've found a value from L used in OuterL. Let's insert phi-nodes at
+        // the corresponding exits blocks and run SSAUpdate to use these values.
+        SSAUpdater SSAUpdate;
+        SSAUpdate.Initialize(I->getType(), I->getName());
+        for (BasicBlock *ExitBB : ExitBlockSet) {
+          if (SSAUpdate.HasValueForBlock(ExitBB))
+            continue;
+          if (!DT->dominates(I->getParent(), ExitBB))
+            continue;
+
+          // We've found the exit block where we need to insert the phi-node.
+          // Create it and register in SSAUpdater.
+          PHINode *NewPN = PHINode::Create(
+              PN->getType(), 1, I->getName() + ".lcssa", &ExitBB->front());
+          for (BasicBlock *PredBB : predecessors(ExitBB))
+            NewPN->addIncoming(I, PredBB);
+          SSAUpdate.AddAvailableValue(ExitBB, NewPN);
+        }
+
+        // Rewrite all uses of I in PN with the new values. We do this for all
+        // operands and not only for the current one to avoid creating multiple
+        // LCSSA phi-nodes for the same value.
+        for (unsigned IdxToRewrite = Idx; IdxToRewrite != E; ++IdxToRewrite)
+          if (PN->getIncomingValue(IdxToRewrite) == I)
+            SSAUpdate.RewriteUse(PN->getOperandUse(IdxToRewrite));
+      }
+    }
+  }
+
   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
for (pred_iterator PI = pred_begin(ExitBlock), PE = pred_end(ExitBlock);
       // allowed.
       if (!L->contains(*PI)) {
         if (rewriteLoopExitBlock(L, ExitBlock, DT, LI, PreserveLCSSA)) {
           ++NumInserted;
           Changed = true;
         }
         break;
       }
   }
 
   // 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
     // common backedge instead.

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
+  br i1 true, label %bb45, label %bb54
+
+bb45:
+  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
+}