[SCEVExpander] Remember phi nodes inserted by LCSSA construction

[nikic]

Summary

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.

[nikic]

@fhahn Maybe it would be better to drop the IRBuilder argument and insert all phis into the vector? You added it in https://github.com/llvm/llvm-project/commit/a9b06a2c14f9a38ba16165f0343faaa9ae713fec specifically for the SCEVExpander use case, but as it is not sufficient, it might make sense to drop it again in this patch.

[fhahn]

In D150681#4367280, @nikic wrote:

@fhahn Maybe it would be better to drop the IRBuilder argument and insert all phis into the vector? You added it in https://github.com/llvm/llvm-project/commit/a9b06a2c14f9a38ba16165f0343faaa9ae713fec specifically for the SCEVExpander use case, but as it is not sufficient, it might make sense to drop it again in this patch.

Sounds good to me, thanks for looking into this!

[nikic]

Drop IRBuilder argument.

[fhahn]

LGTM, thanks!

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

 /// LoopInfo and DominatorTree are required and, since the routine makes no
 /// changes to CFG, preserved.
 ///
 /// Returns true if any modifications are made.
 ///
 /// 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, phis inserted without using Builder (and
+/// thus not executing the insertion callback) will be added to this vector.
 bool formLCSSAForInstructions(
     SmallVectorImpl<Instruction *> &Worklist, const DominatorTree &DT,
     const LoopInfo &LI, IRBuilderBase &Builder,
-    SmallVectorImpl<PHINode *> *PHIsToRemove = nullptr);
+    SmallVectorImpl<PHINode *> *PHIsToRemove = nullptr,
+    SmallVectorImpl<PHINode *> *InsertedPHIs = nullptr);
 
 /// Put loop into LCSSA form.
 ///
 /// Looks at all instructions in the loop which have uses outside of the
 /// current loop. For each, an LCSSA PHI node is inserted and the uses outside
 /// the loop are rewritten to use this node. Sub-loops must be in LCSSA form
 /// already.
 ///

llvm/lib/Transforms/Utils/LCSSA.cpp

 }
 
 /// For every instruction from the worklist, check to see if it has any uses
 /// that are outside the current loop.  If so, insert LCSSA PHI nodes and
 /// 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
while (!Worklist.empty()) {
     if (auto *Inv = dyn_cast<InvokeInst>(I))
       DomBB = Inv->getNormalDest();
 
     const DomTreeNode *DomNode = DT.getNode(DomBB);
 
     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
     // value, and add them to the Phi's map.
     for (BasicBlock *ExitBB : ExitBlocks) {
       if (!DT.dominates(DomNode, DT.getNode(ExitBB)))
         continue;
 
for (auto *DVI : DbgValues) {
       Value *V = AddedPHIs.size() == 1 ? AddedPHIs[0]
                                        : SSAUpdate.FindValueForBlock(UserBB);
       if (V)
         DVI->replaceVariableLocationOp(I, V);
     }
 
     // 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
     // loop and update their exits properly.
     for (auto *PostProcessPN : PostProcessPHIs)
       if (!PostProcessPN->use_empty())
         Worklist.push_back(PostProcessPN);
 

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

Value *SCEVExpander::fixupLCSSAFormFor(Value *V) {
   Instruction *User =
       CastInst::CreateBitOrPointerCast(DefI, ToTy, "tmp.lcssa.user", InsertPt);
   auto RemoveUserOnExit =
       make_scope_exit([User]() { User->eraseFromParent(); });
 
   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, Builder, &PHIsToRemove,
+                           &InsertedPHIs);
+  for (PHINode *PN : InsertedPHIs)
+    rememberInstruction(PN);
   for (PHINode *PN : PHIsToRemove) {
     if (!PN->use_empty())
       continue;
     InsertedValues.erase(PN);
     InsertedPostIncValues.erase(PN);
     PN->eraseFromParent();
   }
 

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

+; 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
+}