[CodeGenPrepare] Fix for using outdated/corrupt LoopInfo

[chill]

Summary

Some transformation in CodeGenPrepare pass may create and/or delete
basic block, but they don't update the LoopInfo, so the LoopInfo may
end up containing dangling pointers and sometimes reused basic blocks,
which leads to "interesting" non-deterministic behaviour.

This patch makes the pass invalidate the LoopInfo whenever the
DominationTree is invalidated and lazily rebuilds it.

[efriedma]

It might be worth trying to teach CodeGenPrepare to update the domtree instead of destroying it, at some point... but that shouldn't block this patch.

llvm/lib/CodeGen/CodeGenPrepare.cpp

602

InitialLI is pointless if you're going to immediately throw it away.

[chill]

llvm/lib/CodeGen/CodeGenPrepare.cpp

602

It's used in the construction of BranchProbabilityInfo and BlockGFrequenceInfo.

[efriedma]

llvm/lib/CodeGen/CodeGenPrepare.cpp

602

I guess, but you don't need to store it in the class for that.

Is the BPI also broken by CFG modifications?

[chill]

llvm/lib/CodeGen/CodeGenPrepare.cpp

602

I guess, but you don't need to store it in the class for that.

Ah, indeed.

Is the BPI also broken by CFG modifications?

Looks like it, yes.

[chill]

llvm/lib/CodeGen/CodeGenPrepare.cpp

602

BFI and BPI use CallbackVH for basic blocks, so they are perhaps fine, but they also keep a raw pointer to the LoopInfo.
Fortunately, LoopInfo is not hard to keep updated, so I opted for this approach in a new patch.

llvm/lib/CodeGen/CodeGenPrepare.cpp

 class CodeGenPrepare : public FunctionPass {
   const TargetMachine *TM = nullptr;
   const TargetSubtargetInfo *SubtargetInfo = nullptr;
   const TargetLowering *TLI = nullptr;
   const TargetRegisterInfo *TRI = nullptr;
   const TargetTransformInfo *TTI = nullptr;
   const BasicBlockSectionsProfileReader *BBSectionsProfileReader = nullptr;
   const TargetLibraryInfo *TLInfo = nullptr;
-  const LoopInfo *LI = nullptr;
+  const LoopInfo *InitialLI = nullptr;
+  std::unique_ptr<LoopInfo> LI;
   std::unique_ptr<BlockFrequencyInfo> BFI;
   std::unique_ptr<BranchProbabilityInfo> BPI;
   ProfileSummaryInfo *PSI = nullptr;
 
   /// As we scan instructions optimizing them, this is the next instruction
   /// to optimize. Transforms that can invalidate this should update it.
   BasicBlock::iterator CurInstIterator;
 
void resetIteratorIfInvalidatedWhileCalling(BasicBlock *BB, F f) {
     // If the iterator instruction was recursively deleted, start over at the
     // start of the block.
     if (IterHandle != CurValue) {
       CurInstIterator = BB->begin();
       SunkAddrs.clear();
     }
   }
 
+  void resetDT() {
+    DT.reset();
+    LI.reset();
+    InitialLI = nullptr;
+  }
+
   // Get the DominatorTree, building if necessary.
   DominatorTree &getDT(Function &F) {
     if (!DT)
       DT = std::make_unique<DominatorTree>(F);
     return *DT;
   }
 
+  const LoopInfo &getLI(Function &F) {
+    if (InitialLI)
+      return *InitialLI;
+    if (!LI)
+      LI = std::make_unique<LoopInfo>(getDT(F));
+    return *LI;
+  }
+
   void removeAllAssertingVHReferences(Value *V);
   bool eliminateAssumptions(Function &F);
   bool eliminateFallThrough(Function &F);
   bool eliminateMostlyEmptyBlocks(Function &F);
   BasicBlock *findDestBlockOfMergeableEmptyBlock(BasicBlock *BB);
   bool canMergeBlocks(const BasicBlock *BB, const BasicBlock *DestBB) const;
   void eliminateMostlyEmptyBlock(BasicBlock *BB);
   bool isMergingEmptyBlockProfitable(BasicBlock *BB, BasicBlock *DestBB,
bool CodeGenPrepare::runOnFunction(Function &F) {
   FreshBBs.clear();
 
   TM = &getAnalysis<TargetPassConfig>().getTM<TargetMachine>();
   SubtargetInfo = TM->getSubtargetImpl(F);
   TLI = SubtargetInfo->getTargetLowering();
   TRI = SubtargetInfo->getRegisterInfo();
   TLInfo = &getAnalysis<TargetLibraryInfoWrapperPass>().getTLI(F);
   TTI = &getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
-  LI = &getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
-  BPI.reset(new BranchProbabilityInfo(F, *LI));
-  BFI.reset(new BlockFrequencyInfo(F, *BPI, *LI));
+  InitialLI = &getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
+
+
+  BPI.reset(new BranchProbabilityInfo(F, *InitialLI));
+  BFI.reset(new BlockFrequencyInfo(F, *BPI, *InitialLI));
   PSI = &getAnalysis<ProfileSummaryInfoWrapperPass>().getPSI();
   BBSectionsProfileReader =
       getAnalysisIfAvailable<BasicBlockSectionsProfileReader>();
   OptSize = F.hasOptSize();
   // Use the basic-block-sections profile to promote hot functions to .text.hot
   // if requested.
   if (BBSectionsGuidedSectionPrefix && BBSectionsProfileReader &&
       BBSectionsProfileReader->isFunctionHot(F.getName())) {
bool CodeGenPrepare::runOnFunction(Function &F) {
   // If we are optimzing huge function, we need to consider the build time.
   // Because the basic algorithm's complex is near O(N!).
   IsHugeFunc = F.size() > HugeFuncThresholdInCGPP;
 
   bool MadeChange = true;
   bool FuncIterated = false;
   while (MadeChange) {
     MadeChange = false;
-    DT.reset();
+    resetDT();
 
     for (BasicBlock &BB : llvm::make_early_inc_range(F)) {
       if (FuncIterated && !FreshBBs.contains(&BB))
         continue;
 
       ModifyDT ModifiedDTOnIteration = ModifyDT::NotModifyDT;
       bool Changed = optimizeBlock(BB, ModifiedDTOnIteration);
 
for (BasicBlock &BB : llvm::make_early_inc_range(F)) {
         //
         // If the %cmp sink to other BB, the %and will has chance to sink.
         if (Changed)
           FreshBBs.insert(&BB);
         else if (FuncIterated)
           FreshBBs.erase(&BB);
 
         if (ModifiedDTOnIteration == ModifyDT::ModifyBBDT)
-          DT.reset();
+          resetDT();
       } else {
         // For small/normal functions, we restart BB iteration if the dominator
         // tree of the Function was changed.
         if (ModifiedDTOnIteration != ModifyDT::NotModifyDT)
           break;
       }
     }
     // We have iterated all the BB in the (only work for huge) function.

 }
 
 /// Eliminate blocks that contain only PHI nodes, debug info directives, and an
 /// unconditional branch. Passes before isel (e.g. LSR/loopsimplify) often split
 /// edges in ways that are non-optimal for isel. Start by eliminating these
 /// blocks so we can split them the way we want them.
 bool CodeGenPrepare::eliminateMostlyEmptyBlocks(Function &F) {
   SmallPtrSet<BasicBlock *, 16> Preheaders;
-  SmallVector<Loop *, 16> LoopList(LI->begin(), LI->end());
+  SmallVector<Loop *, 16> LoopList(getLI(F).begin(), getLI(F).end());
   while (!LoopList.empty()) {
     Loop *L = LoopList.pop_back_val();
     llvm::append_range(LoopList, *L);
     if (BasicBlock *Preheader = L->getLoopPreheader())
       Preheaders.insert(Preheader);
   }
 
   bool MadeChange = false;
static bool isIVIncrement(const Value *V, const LoopInfo *LI) {
   return false;
 }
 
 bool CodeGenPrepare::replaceMathCmpWithIntrinsic(BinaryOperator *BO,
                                                  Value *Arg0, Value *Arg1,
                                                  CmpInst *Cmp,
                                                  Intrinsic::ID IID) {
   auto IsReplacableIVIncrement = [this, &Cmp](BinaryOperator *BO) {
-    if (!isIVIncrement(BO, LI))
+    Function &F = *BO->getParent()->getParent();
+    const LoopInfo &LI = getLI(F);
+    if (!isIVIncrement(BO, &LI))
       return false;
-    const Loop *L = LI->getLoopFor(BO->getParent());
+    const Loop *L = LI.getLoopFor(BO->getParent());
     assert(L && "L should not be null after isIVIncrement()");
     // Do not risk on moving increment into a child loop.
-    if (LI->getLoopFor(Cmp->getParent()) != L)
+    if (LI.getLoopFor(Cmp->getParent()) != L)
       return false;
 
     // Finally, we need to ensure that the insert point will dominate all
     // existing uses of the increment.
 
-    auto &DT = getDT(*BO->getParent()->getParent());
+    auto &DT = getDT(F);
     if (DT.dominates(Cmp->getParent(), BO->getParent()))
       // If we're moving up the dom tree, all uses are trivially dominated.
       // (This is the common case for code produced by LSR.)
       return true;
 
     // Otherwise, special case the single use in the phi recurrence.
     return BO->hasOneUse() && DT.dominates(Cmp->getParent(), L->getLoopLatch());
   };

 /// A helper class for matching addressing modes.
 ///
 /// This encapsulates the logic for matching the target-legal addressing modes.
 class AddressingModeMatcher {
   SmallVectorImpl<Instruction *> &AddrModeInsts;
   const TargetLowering &TLI;
   const TargetRegisterInfo &TRI;
   const DataLayout &DL;
-  const LoopInfo &LI;
   const std::function<const DominatorTree &()> getDTFn;
+  const std::function<const LoopInfo &()> getLIFn;
 
   /// AccessTy/MemoryInst - This is the type for the access (e.g. double) and
   /// the memory instruction that we're computing this address for.
   Type *AccessTy;
   unsigned AddrSpace;
   Instruction *MemoryInst;
 
   /// This is the addressing mode that we're building up. This is
class AddressingModeMatcher {
   /// True if we are optimizing for size.
   bool OptSize = false;
 
   ProfileSummaryInfo *PSI;
   BlockFrequencyInfo *BFI;
 
   AddressingModeMatcher(
       SmallVectorImpl<Instruction *> &AMI, const TargetLowering &TLI,
-      const TargetRegisterInfo &TRI, const LoopInfo &LI,
-      const std::function<const DominatorTree &()> getDTFn, Type *AT,
-      unsigned AS, Instruction *MI, ExtAddrMode &AM,
-      const SetOfInstrs &InsertedInsts, InstrToOrigTy &PromotedInsts,
-      TypePromotionTransaction &TPT,
+      const TargetRegisterInfo &TRI,
+      const std::function<const DominatorTree &()> getDTFn,
+      const std::function<const LoopInfo &()> getLIFn, Type *AT, unsigned AS,
+      Instruction *MI, ExtAddrMode &AM, const SetOfInstrs &InsertedInsts,
+      InstrToOrigTy &PromotedInsts, TypePromotionTransaction &TPT,
       std::pair<AssertingVH<GetElementPtrInst>, int64_t> &LargeOffsetGEP,
       bool OptSize, ProfileSummaryInfo *PSI, BlockFrequencyInfo *BFI)
       : AddrModeInsts(AMI), TLI(TLI), TRI(TRI),
-        DL(MI->getModule()->getDataLayout()), LI(LI), getDTFn(getDTFn),
-        AccessTy(AT), AddrSpace(AS), MemoryInst(MI), AddrMode(AM),
-        InsertedInsts(InsertedInsts), PromotedInsts(PromotedInsts), TPT(TPT),
-        LargeOffsetGEP(LargeOffsetGEP), OptSize(OptSize), PSI(PSI), BFI(BFI) {
+        DL(MI->getModule()->getDataLayout()), getDTFn(getDTFn),
+        getLIFn(getLIFn), AccessTy(AT), AddrSpace(AS), MemoryInst(MI),
+        AddrMode(AM), InsertedInsts(InsertedInsts),
+        PromotedInsts(PromotedInsts), TPT(TPT), LargeOffsetGEP(LargeOffsetGEP),
+        OptSize(OptSize), PSI(PSI), BFI(BFI) {
     IgnoreProfitability = false;
   }
 
 public:
   /// Find the maximal addressing mode that a load/store of V can fold,
   /// give an access type of AccessTy.  This returns a list of involved
   /// instructions in AddrModeInsts.
   /// \p InsertedInsts The instructions inserted by other CodeGenPrepare
   /// optimizations.
   /// \p PromotedInsts maps the instructions to their type before promotion.
   /// \p The ongoing transaction where every action should be registered.
   static ExtAddrMode
   Match(Value *V, Type *AccessTy, unsigned AS, Instruction *MemoryInst,
         SmallVectorImpl<Instruction *> &AddrModeInsts,
-        const TargetLowering &TLI, const LoopInfo &LI,
+        const TargetLowering &TLI,
         const std::function<const DominatorTree &()> getDTFn,
+        const std::function<const LoopInfo &()> getLIFn,
         const TargetRegisterInfo &TRI, const SetOfInstrs &InsertedInsts,
         InstrToOrigTy &PromotedInsts, TypePromotionTransaction &TPT,
         std::pair<AssertingVH<GetElementPtrInst>, int64_t> &LargeOffsetGEP,
         bool OptSize, ProfileSummaryInfo *PSI, BlockFrequencyInfo *BFI) {
     ExtAddrMode Result;
 
-    bool Success = AddressingModeMatcher(AddrModeInsts, TLI, TRI, LI, getDTFn,
-                                         AccessTy, AS, MemoryInst, Result,
-                                         InsertedInsts, PromotedInsts, TPT,
-                                         LargeOffsetGEP, OptSize, PSI, BFI)
+    bool Success = AddressingModeMatcher(AddrModeInsts, TLI, TRI, getDTFn,
+                                         getLIFn, AccessTy, AS, MemoryInst,
+                                         Result, InsertedInsts, PromotedInsts,
+                                         TPT, LargeOffsetGEP, OptSize, PSI, BFI)
                        .matchAddr(V, 0);
     (void)Success;
     assert(Success && "Couldn't select *anything*?");
     return Result;
   }
 
 private:
   bool matchScaledValue(Value *ScaleReg, int64_t Scale, unsigned Depth);
bool AddressingModeMatcher::matchScaledValue(Value *ScaleReg, int64_t Scale,
   // Okay, we decided that we can add ScaleReg+Scale to AddrMode.  Check now
   // to see if ScaleReg is actually X+C.  If so, we can turn this into adding
   // X*Scale + C*Scale to addr mode. If we found available IV increment, do not
   // go any further: we can reuse it and cannot eliminate it.
   ConstantInt *CI = nullptr;
   Value *AddLHS = nullptr;
   if (isa<Instruction>(ScaleReg) && // not a constant expr.
       match(ScaleReg, m_Add(m_Value(AddLHS), m_ConstantInt(CI))) &&
-      !isIVIncrement(ScaleReg, &LI) && CI->getValue().isSignedIntN(64)) {
+      !isIVIncrement(ScaleReg, &getLIFn()) && CI->getValue().isSignedIntN(64)) {
     TestAddrMode.InBounds = false;
     TestAddrMode.ScaledReg = AddLHS;
     TestAddrMode.BaseOffs += CI->getSExtValue() * TestAddrMode.Scale;
 
     // If this addressing mode is legal, commit it and remember that we folded
     // this instruction.
     if (TLI.isLegalAddressingMode(DL, TestAddrMode, AccessTy, AddrSpace)) {
       AddrModeInsts.push_back(cast<Instruction>(ScaleReg));
       AddrMode = TestAddrMode;
       return true;
     }
     // Restore status quo.
     TestAddrMode = AddrMode;
   }
 
   // If this is an add recurrence with a constant step, return the increment
   // instruction and the canonicalized step.
   auto GetConstantStep =
       [this](const Value *V) -> std::optional<std::pair<Instruction *, APInt>> {
     auto *PN = dyn_cast<PHINode>(V);
     if (!PN)
       return std::nullopt;
-    auto IVInc = getIVIncrement(PN, &LI);
+    auto IVInc = getIVIncrement(PN, &getLIFn());
     if (!IVInc)
       return std::nullopt;
     // TODO: The result of the intrinsics above is two-complement. However when
     // IV inc is expressed as add or sub, iv.next is potentially a poison value.
     // If it has nuw or nsw flags, we need to make sure that these flags are
     // inferrable at the point of memory instruction. Otherwise we are replacing
     // well-defined two-complement computation with poison. Currently, to avoid
     // potentially complex analysis needed to prove this, we reject such cases.
bool AddressingModeMatcher::matchScaledValue(Value *ScaleReg, int64_t Scale,
   //    assignment.
   if (AddrMode.BaseOffs) {
     if (auto IVStep = GetConstantStep(ScaleReg)) {
       Instruction *IVInc = IVStep->first;
       // The following assert is important to ensure a lack of infinite loops.
       // This transforms is (intentionally) the inverse of the one just above.
       // If they don't agree on the definition of an increment, we'd alternate
       // back and forth indefinitely.
-      assert(isIVIncrement(IVInc, &LI) && "implied by GetConstantStep");
+      assert(isIVIncrement(IVInc, &getLIFn()) && "implied by GetConstantStep");
       APInt Step = IVStep->second;
       APInt Offset = Step * AddrMode.Scale;
       if (Offset.isSignedIntN(64)) {
         TestAddrMode.InBounds = false;
         TestAddrMode.ScaledReg = IVInc;
         TestAddrMode.BaseOffs -= Offset.getLimitedValue();
         // If this addressing mode is legal, commit it..
         // (Note that we defer the (expensive) domtree base legality check
for (const std::pair<Use *, Type *> &Pair : MemoryUses) {
     // Do a match against the root of this address, ignoring profitability. This
     // will tell us if the addressing mode for the memory operation will
     // *actually* cover the shared instruction.
     ExtAddrMode Result;
     std::pair<AssertingVH<GetElementPtrInst>, int64_t> LargeOffsetGEP(nullptr,
                                                                       0);
     TypePromotionTransaction::ConstRestorationPt LastKnownGood =
         TPT.getRestorationPoint();
-    AddressingModeMatcher Matcher(MatchedAddrModeInsts, TLI, TRI, LI, getDTFn,
-                                  AddressAccessTy, AS, UserI, Result,
+    AddressingModeMatcher Matcher(MatchedAddrModeInsts, TLI, TRI, getDTFn,
+                                  getLIFn, AddressAccessTy, AS, UserI, Result,
                                   InsertedInsts, PromotedInsts, TPT,
                                   LargeOffsetGEP, OptSize, PSI, BFI);
     Matcher.IgnoreProfitability = true;
     bool Success = Matcher.matchAddr(Address, 0);
     (void)Success;
     assert(Success && "Couldn't select *anything*?");
 
     // The match was to check the profitability, the changes made are not
std::pair<AssertingVH<GetElementPtrInst>, int64_t> LargeOffsetGEP(nullptr,
                                                                       0);
     // Defer the query (and possible computation of) the dom tree to point of
     // actual use.  It's expected that most address matches don't actually need
     // the domtree.
     auto getDTFn = [MemoryInst, this]() -> const DominatorTree & {
       Function *F = MemoryInst->getParent()->getParent();
       return this->getDT(*F);
     };
+    auto getLIFn = [MemoryInst, this]() -> const LoopInfo & {
+      Function *F = MemoryInst->getParent()->getParent();
+      return this->getLI(*F);
+    };
     ExtAddrMode NewAddrMode = AddressingModeMatcher::Match(
-        V, AccessTy, AddrSpace, MemoryInst, AddrModeInsts, *TLI, *LI, getDTFn,
-        *TRI, InsertedInsts, PromotedInsts, TPT, LargeOffsetGEP, OptSize, PSI,
-        BFI.get());
+        V, AccessTy, AddrSpace, MemoryInst, AddrModeInsts, *TLI, getDTFn,
+        getLIFn, *TRI, InsertedInsts, PromotedInsts, TPT, LargeOffsetGEP,
+        OptSize, PSI, BFI.get());
 
     GetElementPtrInst *GEP = LargeOffsetGEP.first;
     if (GEP && !NewGEPBases.count(GEP)) {
       // If splitting the underlying data structure can reduce the offset of a
       // GEP, collect the GEP.  Skip the GEPs that are the new bases of
       // previously split data structures.
       LargeOffsetGEPMap[GEP->getPointerOperand()].push_back(LargeOffsetGEP);
       LargeOffsetGEPID.insert(std::make_pair(GEP, LargeOffsetGEPID.size()));
if (TLI->isSelectSupported(SelectKind) &&
       (!isFormingBranchFromSelectProfitable(TTI, TLI, SI) || OptSize ||
        llvm::shouldOptimizeForSize(SI->getParent(), PSI, BFI.get())))
     return false;
 
   // The DominatorTree needs to be rebuilt by any consumers after this
   // transformation. We simply reset here rather than setting the ModifiedDT
   // flag to avoid restarting the function walk in runOnFunction for each
   // select optimized.
-  DT.reset();
+  resetDT();
 
   // Transform a sequence like this:
   //    start:
   //       %cmp = cmp uge i32 %a, %b
   //       %sel = select i1 %cmp, i32 %c, i32 %d
   //
   // Into:
   //    start:
if (CastInst *CI = dyn_cast<CastInst>(I)) {
     // the address of globals out of a loop).  If this is the case, we don't
     // want to forward-subst the cast.
     if (isa<Constant>(CI->getOperand(0)))
       return false;
 
     if (OptimizeNoopCopyExpression(CI, *TLI, *DL))
       return true;
 
+    const LoopInfo &LI = getLI(*I->getParent()->getParent());
     if ((isa<UIToFPInst>(I) || isa<FPToUIInst>(I) || isa<TruncInst>(I)) &&
         TLI->optimizeExtendOrTruncateConversion(I,
-                                                LI->getLoopFor(I->getParent())))
+                                                LI.getLoopFor(I->getParent())))
       return true;
 
     if (isa<ZExtInst>(I) || isa<SExtInst>(I)) {
       /// Sink a zext or sext into its user blocks if the target type doesn't
       /// fit in one register
       if (TLI->getTypeAction(CI->getContext(),
                              TLI->getValueType(*DL, CI->getType())) ==
           TargetLowering::TypeExpandInteger) {
         return SinkCast(CI);
       } else {
         if (TLI->optimizeExtendOrTruncateConversion(
-                I, LI->getLoopFor(I->getParent())))
+                I, LI.getLoopFor(I->getParent())))
           return true;
 
         bool MadeChange = optimizeExt(I);
         return MadeChange | optimizeExtUses(I);
       }
     }
     return false;
   }
do {
     CurInstIterator = BB.begin();
     ModifiedDT = ModifyDT::NotModifyDT;
     while (CurInstIterator != BB.end()) {
       MadeChange |= optimizeInst(&*CurInstIterator++, ModifiedDT);
       if (ModifiedDT != ModifyDT::NotModifyDT) {
         // For huge function we tend to quickly go though the inner optmization
         // opportunities in the BB. So we go back to the BB head to re-optimize
         // each instruction instead of go back to the function head.
-        if (IsHugeFunc) {
-          DT.reset();
-          getDT(*BB.getParent());
-          break;
-        } else {
+        if (!IsHugeFunc)
           return true;
-        }
+
+        resetDT();
+        break;
       }
     }
   } while (ModifiedDT == ModifyDT::ModifyInstDT);
 
   bool MadeBitReverse = true;
   while (MadeBitReverse) {
     MadeBitReverse = false;
     for (auto &I : reverse(BB)) {