[SCEVExpander] Remember phi nodes inserted by LCSSA construction

SCEVExpander keeps track of all instructions it inserted. However,
it currently misses some phi nodes created during LCSSA construction.
Fix this by collecting these into another argument.

This also removes the IRBuilder argument, which was added for
essentially the same purpose, but only handles the root LCSSA nodes,
not those inserted by SSAUpdater.

This was reported as a regression on D149344, but the reduced test
case also reproduces without it.

Differential Revision: https://reviews.llvm.org/D150681

Author: nikic

llvm/include/llvm/Transforms/Utils/LoopUtils.h

@@ -76,10 +76,13 @@ bool formDedicatedExitBlocks(Loop *L, DominatorTree *DT, LoopInfo *LI,
 /// This function may introduce unused PHI nodes. If \p PHIsToRemove is not
 /// nullptr, those are added to it (before removing, the caller has to check if
 /// they still do not have any uses). Otherwise the PHIs are directly removed.
+///
+/// If \p InsertedPHIs is not nullptr, inserted phis will be added to this
+/// vector.
 bool formLCSSAForInstructions(
     SmallVectorImpl<Instruction *> &Worklist, const DominatorTree &DT,
-    const LoopInfo &LI, IRBuilderBase &Builder,
-    SmallVectorImpl<PHINode *> *PHIsToRemove = nullptr);
+    const LoopInfo &LI, SmallVectorImpl<PHINode *> *PHIsToRemove = nullptr,
+    SmallVectorImpl<PHINode *> *InsertedPHIs = nullptr);
 
 /// Put loop into LCSSA form.
 ///

llvm/lib/Transforms/Scalar/LoopInterchange.cpp

@@ -30,7 +30,6 @@
 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
-#include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
@@ -1686,12 +1685,11 @@ bool LoopInterchangeTransform::adjustLoopBranches() {
   // latch. In that case, we need to create LCSSA phis for them, because after
   // interchanging they will be defined in the new inner loop and used in the
   // new outer loop.
-  IRBuilder<> Builder(OuterLoopHeader->getContext());
   SmallVector<Instruction *, 4> MayNeedLCSSAPhis;
   for (Instruction &I :
        make_range(OuterLoopHeader->begin(), std::prev(OuterLoopHeader->end())))
     MayNeedLCSSAPhis.push_back(&I);
-  formLCSSAForInstructions(MayNeedLCSSAPhis, *DT, *LI, Builder);
+  formLCSSAForInstructions(MayNeedLCSSAPhis, *DT, *LI);
 
   return true;
 }

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

@@ -5659,8 +5659,7 @@ void LSRInstance::RewriteForPHI(
       }
     }
 
-  IRBuilder<> Builder(L->getHeader()->getContext());
-  formLCSSAForInstructions(InsertedNonLCSSAInsts, DT, LI, Builder);
+  formLCSSAForInstructions(InsertedNonLCSSAInsts, DT, LI);
 }
 
 /// Emit instructions for the leading candidate expression for this LSRUse (this

llvm/lib/Transforms/Utils/LCSSA.cpp

@@ -40,7 +40,6 @@
 #include "llvm/Analysis/ScalarEvolutionAliasAnalysis.h"
 #include "llvm/IR/DebugInfo.h"
 #include "llvm/IR/Dominators.h"
-#include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/PredIteratorCache.h"
@@ -77,15 +76,13 @@ static bool isExitBlock(BasicBlock *BB,
 /// rewrite the uses.
 bool llvm::formLCSSAForInstructions(SmallVectorImpl<Instruction *> &Worklist,
                                     const DominatorTree &DT, const LoopInfo &LI,
-                                    IRBuilderBase &Builder,
-                                    SmallVectorImpl<PHINode *> *PHIsToRemove) {
+                                    SmallVectorImpl<PHINode *> *PHIsToRemove,
+                                    SmallVectorImpl<PHINode *> *InsertedPHIs) {
   SmallVector<Use *, 16> UsesToRewrite;
   SmallSetVector<PHINode *, 16> LocalPHIsToRemove;
   PredIteratorCache PredCache;
   bool Changed = false;
 
-  IRBuilderBase::InsertPointGuard InsertPtGuard(Builder);
-
   // Cache the Loop ExitBlocks across this loop.  We expect to get a lot of
   // instructions within the same loops, computing the exit blocks is
   // expensive, and we're not mutating the loop structure.
@@ -146,8 +143,8 @@ bool llvm::formLCSSAForInstructions(SmallVectorImpl<Instruction *> &Worklist,
     SmallVector<PHINode *, 16> AddedPHIs;
     SmallVector<PHINode *, 8> PostProcessPHIs;
 
-    SmallVector<PHINode *, 4> InsertedPHIs;
-    SSAUpdater SSAUpdate(&InsertedPHIs);
+    SmallVector<PHINode *, 4> LocalInsertedPHIs;
+    SSAUpdater SSAUpdate(&LocalInsertedPHIs);
     SSAUpdate.Initialize(I->getType(), I->getName());
 
     // Insert the LCSSA phi's into all of the exit blocks dominated by the
@@ -159,9 +156,10 @@ bool llvm::formLCSSAForInstructions(SmallVectorImpl<Instruction *> &Worklist,
       // If we already inserted something for this BB, don't reprocess it.
       if (SSAUpdate.HasValueForBlock(ExitBB))
         continue;
-      Builder.SetInsertPoint(&ExitBB->front());
-      PHINode *PN = Builder.CreatePHI(I->getType(), PredCache.size(ExitBB),
-                                      I->getName() + ".lcssa");
+      PHINode *PN = PHINode::Create(I->getType(), PredCache.size(ExitBB),
+                                    I->getName() + ".lcssa", &ExitBB->front());
+      if (InsertedPHIs)
+        InsertedPHIs->push_back(PN);
       // Get the debug location from the original instruction.
       PN->setDebugLoc(I->getDebugLoc());
 
@@ -251,10 +249,12 @@ bool llvm::formLCSSAForInstructions(SmallVectorImpl<Instruction *> &Worklist,
 
     // SSAUpdater might have inserted phi-nodes inside other loops. We'll need
     // to post-process them to keep LCSSA form.
-    for (PHINode *InsertedPN : InsertedPHIs) {
+    for (PHINode *InsertedPN : LocalInsertedPHIs) {
       if (auto *OtherLoop = LI.getLoopFor(InsertedPN->getParent()))
         if (!L->contains(OtherLoop))
           PostProcessPHIs.push_back(InsertedPN);
+      if (InsertedPHIs)
+        InsertedPHIs->push_back(InsertedPN);
     }
 
     // Post process PHI instructions that were inserted into another disjoint
@@ -386,8 +386,7 @@ bool llvm::formLCSSA(Loop &L, const DominatorTree &DT, const LoopInfo *LI) {
     }
   }
 
-  IRBuilder<> Builder(L.getHeader()->getContext());
-  Changed = formLCSSAForInstructions(Worklist, DT, *LI, Builder);
+  Changed = formLCSSAForInstructions(Worklist, DT, *LI);
 
   assert(L.isLCSSAForm(DT));
 

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

@@ -2561,7 +2561,11 @@ Value *SCEVExpander::fixupLCSSAFormFor(Value *V) {
   SmallVector<Instruction *, 1> ToUpdate;
   ToUpdate.push_back(DefI);
   SmallVector<PHINode *, 16> PHIsToRemove;
-  formLCSSAForInstructions(ToUpdate, SE.DT, SE.LI, Builder, &PHIsToRemove);
+  SmallVector<PHINode *, 16> InsertedPHIs;
+  formLCSSAForInstructions(ToUpdate, SE.DT, SE.LI, &PHIsToRemove,
+                           &InsertedPHIs);
+  for (PHINode *PN : InsertedPHIs)
+    rememberInstruction(PN);
   for (PHINode *PN : PHIsToRemove) {
     if (!PN->use_empty())
       continue;

llvm/test/Transforms/LoopIdiom/remove-inserted-lcssa.ll

@@ -0,0 +1,76 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 2
+; RUN: opt -S -passes=loop-idiom < %s | FileCheck %s
+
+; Make sure that any inserted LCSSA phi nodes are removed if the transform
+; is aborted.
+
+define void @test() {
+; CHECK-LABEL: define void @test() {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca [64 x i8], align 16
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[PHI:%.*]] = phi i64 [ 0, [[ENTRY:%.*]] ], [ 1, [[LOOP_LATCH:%.*]] ]
+; CHECK-NEXT:    br i1 false, label [[LOOP_EXIT2:%.*]], label [[LOOP_LATCH]]
+; CHECK:       loop.latch:
+; CHECK-NEXT:    [[OR:%.*]] = or i64 [[PHI]], 4
+; CHECK-NEXT:    br i1 false, label [[LOOP_EXIT:%.*]], label [[LOOP]]
+; CHECK:       loop.exit:
+; CHECK-NEXT:    [[OR_LCSSA:%.*]] = phi i64 [ [[OR]], [[LOOP_LATCH]] ]
+; CHECK-NEXT:    br label [[LOOP2_PREHEADER:%.*]]
+; CHECK:       loop.exit2:
+; CHECK-NEXT:    br label [[LOOP2_PREHEADER]]
+; CHECK:       loop2.preheader:
+; CHECK-NEXT:    [[PHI5_PH:%.*]] = phi ptr [ null, [[LOOP_EXIT2]] ], [ [[ALLOCA]], [[LOOP_EXIT]] ]
+; CHECK-NEXT:    [[PHI6_PH:%.*]] = phi i64 [ 0, [[LOOP_EXIT2]] ], [ [[OR_LCSSA]], [[LOOP_EXIT]] ]
+; CHECK-NEXT:    br label [[LOOP2:%.*]]
+; CHECK:       loop2:
+; CHECK-NEXT:    [[PHI5:%.*]] = phi ptr [ [[GETELEMENTPTR7:%.*]], [[LOOP2]] ], [ [[PHI5_PH]], [[LOOP2_PREHEADER]] ]
+; CHECK-NEXT:    [[PHI6:%.*]] = phi i64 [ [[ADD:%.*]], [[LOOP2]] ], [ [[PHI6_PH]], [[LOOP2_PREHEADER]] ]
+; CHECK-NEXT:    [[GETELEMENTPTR:%.*]] = getelementptr i8, ptr [[ALLOCA]], i64 [[PHI6]]
+; CHECK-NEXT:    [[LOAD:%.*]] = load i8, ptr [[GETELEMENTPTR]], align 1
+; CHECK-NEXT:    store i8 [[LOAD]], ptr [[PHI5]], align 1
+; CHECK-NEXT:    [[GETELEMENTPTR7]] = getelementptr i8, ptr [[PHI5]], i64 1
+; CHECK-NEXT:    [[ADD]] = add i64 [[PHI6]], 1
+; CHECK-NEXT:    [[ICMP:%.*]] = icmp eq i64 [[PHI6]], 0
+; CHECK-NEXT:    br i1 [[ICMP]], label [[LOOP2_EXIT:%.*]], label [[LOOP2]]
+; CHECK:       loop2.exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %alloca = alloca [64 x i8], align 16
+  br label %loop
+
+loop:
+  %phi = phi i64 [ 0, %entry ], [ 1, %loop.latch ]
+  br i1 false, label %loop.exit2, label %loop.latch
+
+loop.latch:
+  %or = or i64 %phi, 4
+  br i1 false, label %loop.exit, label %loop
+
+loop.exit:
+  br label %loop2.preheader
+
+loop.exit2:
+  br label %loop2.preheader
+
+loop2.preheader:
+  %phi5.ph = phi ptr [ null, %loop.exit2 ], [ %alloca, %loop.exit ]
+  %phi6.ph = phi i64 [ 0, %loop.exit2 ], [ %or, %loop.exit ]
+  br label %loop2
+
+loop2:
+  %phi5 = phi ptr [ %getelementptr7, %loop2 ], [ %phi5.ph, %loop2.preheader ]
+  %phi6 = phi i64 [ %add, %loop2 ], [ %phi6.ph, %loop2.preheader ]
+  %getelementptr = getelementptr i8, ptr %alloca, i64 %phi6
+  %load = load i8, ptr %getelementptr, align 1
+  store i8 %load, ptr %phi5, align 1
+  %getelementptr7 = getelementptr i8, ptr %phi5, i64 1
+  %add = add i64 %phi6, 1
+  %icmp = icmp eq i64 %phi6, 0
+  br i1 %icmp, label %loop2.exit, label %loop2
+
+loop2.exit:
+  ret void
+}