[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, 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.
 ///
 /// 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/Scalar/LoopInterchange.cpp

 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/ScalarEvolutionExpressions.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constants.h"
 #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"
 #include "llvm/IR/User.h"
 #include "llvm/IR/Value.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
bool LoopInterchangeTransform::adjustLoopBranches() {
   OuterLoopHeader->replacePhiUsesWith(InnerLoopLatch, OuterLoopLatch);
   InnerLoopHeader->replacePhiUsesWith(OuterLoopPreHeader, InnerLoopPreHeader);
   InnerLoopHeader->replacePhiUsesWith(OuterLoopLatch, InnerLoopLatch);
 
   // Values defined in the outer loop header could be used in the inner loop
   // 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;
 }
 
 bool LoopInterchangeTransform::adjustLoopLinks() {
   // Adjust all branches in the inner and outer loop.
   bool Changed = adjustLoopBranches();
   if (Changed) {

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

if (PN->getIncomingValue(i) == LF.OperandValToReplace) {
                   for (const auto &val : NewPN->incoming_values())
                     if (val == Fixup.OperandValToReplace)
                       Fixup.UserInst = NewPN;
                 }
             }
       }
     }
 
-  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
 /// is called "expanding"), and update the UserInst to reference the newly
 /// expanded value.
 void LSRInstance::Rewrite(const LSRUse &LU, const LSRFixup &LF,
                           const Formula &F,
                           SmallVectorImpl<WeakTrackingVH> &DeadInsts) const {

llvm/lib/Transforms/Utils/LCSSA.cpp

 #include "llvm/Analysis/GlobalsModRef.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/LoopPass.h"
 #include "llvm/Analysis/MemorySSA.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 #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"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Transforms/Utils.h"
 #include "llvm/Transforms/Utils/LoopUtils.h"
static bool isExitBlock(BasicBlock *BB,
   return is_contained(ExitBlocks, BB);
 }
 
 /// 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
   // instructions within the same loops, computing the exit blocks is
   // expensive, and we're not mutating the loop structure.
   SmallDenseMap<Loop*, SmallVector<BasicBlock *,1>> LoopExitBlocks;
 
   while (!Worklist.empty()) {
     UsesToRewrite.clear();
 
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;
 
       // 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());
 
       // Add inputs from inside the loop for this PHI. This is valid
       // because `I` dominates `ExitBB` (checked above).  This implies
       // that every incoming block/edge is dominated by `I` as well,
       // i.e. we can add uses of `I` to those incoming edges/append to the incoming
       // blocks without violating the SSA dominance property.
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);
 
for (Instruction &I : *BB) {
       // is inside the loop and another which is not.
       if (I.getType()->isTokenTy())
         continue;
 
       Worklist.push_back(&I);
     }
   }
 
-  IRBuilder<> Builder(L.getHeader()->getContext());
-  Changed = formLCSSAForInstructions(Worklist, DT, *LI, Builder);
+  Changed = formLCSSAForInstructions(Worklist, DT, *LI);
 
   assert(L.isLCSSAForm(DT));
 
   return Changed;
 }
 
 /// Process a loop nest depth first.
 bool llvm::formLCSSARecursively(Loop &L, const DominatorTree &DT,

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, &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
+}