[LoopFlatten] Mark inner loop as deleted

[nikic]

Summary

If a loop is flattened, the inner loop is removed and the LPM should be informed of this fact, so it can invalidate associated analyses. To support this, we relax an assertion in LPMUpdater to allow invalidating non-top-level loops when running in LoopNestMode, as the pass does not know how exactly it will get scheduled.

[aeubanks]

if we have a loop nest pass in one pass manager delete a loop without invalidating its analyses, then another loop pass in a different pass manager allocate a loop at the same address I think we could still see invalid cached analyses

LoopAnalysisManagerFunctionProxy::Result::invalidate() may walk through LoopInfo to get a list of loops to clear analyses, and if we delete a loop, we won't clear the corresponding cached analyses. We can always blow away the entire LoopAnalysisManagerFunctionProxy, but always properly invalidating seems cleaner and there's less chance of things silently going wrong.

[nikic]

Always invalidate, independently of LoopNestMode.

llvm/include/llvm/Transforms/Scalar/LoopPassManager.h

public:
   /// Note that this loop must either be the current loop or a subloop of the
   /// current loop. This routine must be called prior to removing the loop from
   /// the loop nest.
   ///
   /// If this is called for the current loop, in addition to clearing any
   /// state, this routine will mark that the current loop should be skipped by
   /// the rest of the pass management infrastructure.
   void markLoopAsDeleted(Loop &L, llvm::StringRef Name) {
-    assert((!LoopNestMode || CurrentL == &L) &&
-           "L should be a top-level loop in loop-nest mode.");
-    LAM.clear(L, Name);
     assert((&L == CurrentL || CurrentL->contains(&L)) &&
            "Cannot delete a loop outside of the "
            "subloop tree currently being processed.");
+
+    // In LoopNestMode, ignore deletion of child loops, as these aren't visited
+    // anyway.
+    if (LoopNestMode && CurrentL != &L)
+      return;
+
+    LAM.clear(L, Name);
     if (&L == CurrentL)
       SkipCurrentLoop = true;
   }
 
   void setParentLoop(Loop *L) {
 #ifdef LLVM_ENABLE_ABI_BREAKING_CHECKS
     ParentL = L;
 #endif

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

if (!checkIVUsers(FI))
     return false;
 
   LLVM_DEBUG(dbgs() << "CanFlattenLoopPair: OK\n");
   return true;
 }
 
 static bool DoFlattenLoopPair(FlattenInfo &FI, DominatorTree *DT, LoopInfo *LI,
                               ScalarEvolution *SE, AssumptionCache *AC,
-                              const TargetTransformInfo *TTI) {
+                              const TargetTransformInfo *TTI, LPMUpdater *U) {
   Function *F = FI.OuterLoop->getHeader()->getParent();
   LLVM_DEBUG(dbgs() << "Checks all passed, doing the transformation\n");
   {
     using namespace ore;
     OptimizationRemark Remark(DEBUG_TYPE, "Flattened", FI.InnerLoop->getStartLoc(),
                               FI.InnerLoop->getHeader());
     OptimizationRemarkEmitter ORE(F);
     Remark << "Flattened into outer loop";
LLVM_DEBUG(dbgs() << "Replacing: "; V->dump();
                dbgs() << "with:      "; OuterValue->dump());
     V->replaceAllUsesWith(OuterValue);
   }
 
   // Tell LoopInfo, SCEV and the pass manager that the inner loop has been
   // deleted, and any information that have about the outer loop invalidated.
   SE->forgetLoop(FI.OuterLoop);
   SE->forgetLoop(FI.InnerLoop);
+  if (U)
+    U->markLoopAsDeleted(*FI.InnerLoop, FI.InnerLoop->getName());
   LI->erase(FI.InnerLoop);
 
   // Increment statistic value.
   NumFlattened++;
 
   return true;
 }
 
static bool CanWidenIV(FlattenInfo &FI, DominatorTree *DT, LoopInfo *LI,
   FI.NarrowOuterInductionPHI = FI.OuterInductionPHI;
 
   // After widening, rediscover all the loop components.
   return CanFlattenLoopPair(FI, DT, LI, SE, AC, TTI);
 }
 
 static bool FlattenLoopPair(FlattenInfo &FI, DominatorTree *DT, LoopInfo *LI,
                             ScalarEvolution *SE, AssumptionCache *AC,
-                            const TargetTransformInfo *TTI) {
+                            const TargetTransformInfo *TTI, LPMUpdater *U) {
   LLVM_DEBUG(
       dbgs() << "Loop flattening running on outer loop "
              << FI.OuterLoop->getHeader()->getName() << " and inner loop "
              << FI.InnerLoop->getHeader()->getName() << " in "
              << FI.OuterLoop->getHeader()->getParent()->getName() << "\n");
 
   if (!CanFlattenLoopPair(FI, DT, LI, SE, AC, TTI))
     return false;
static bool FlattenLoopPair(FlattenInfo &FI, DominatorTree *DT, LoopInfo *LI,
   // 'RepeatedInstructionThreshold' is set to only 2, which can probably be
   // relaxed. Because this is making a code change (the IV widening, but not
   // the flattening), we return true here.
   if (FI.Widened && !CanFlatten)
     return true;
 
   // If we have widened and can perform the transformation, do that here.
   if (CanFlatten)
-    return DoFlattenLoopPair(FI, DT, LI, SE, AC, TTI);
+    return DoFlattenLoopPair(FI, DT, LI, SE, AC, TTI, U);
 
   // Otherwise, if we haven't widened the IV, check if the new iteration
   // variable might overflow. In this case, we need to version the loop, and
   // select the original version at runtime if the iteration space is too
   // large.
   // TODO: We currently don't version the loop.
   OverflowResult OR = checkOverflow(FI, DT, AC);
   if (OR == OverflowResult::AlwaysOverflowsHigh ||
       OR == OverflowResult::AlwaysOverflowsLow) {
     LLVM_DEBUG(dbgs() << "Multiply would always overflow, so not profitable\n");
     return false;
   } else if (OR == OverflowResult::MayOverflow) {
     LLVM_DEBUG(dbgs() << "Multiply might overflow, not flattening\n");
     return false;
   }
 
   LLVM_DEBUG(dbgs() << "Multiply cannot overflow, modifying loop in-place\n");
-  return DoFlattenLoopPair(FI, DT, LI, SE, AC, TTI);
+  return DoFlattenLoopPair(FI, DT, LI, SE, AC, TTI, U);
 }
 
 bool Flatten(LoopNest &LN, DominatorTree *DT, LoopInfo *LI, ScalarEvolution *SE,
-             AssumptionCache *AC, TargetTransformInfo *TTI) {
+             AssumptionCache *AC, TargetTransformInfo *TTI, LPMUpdater *U) {
   bool Changed = false;
   for (Loop *InnerLoop : LN.getLoops()) {
     auto *OuterLoop = InnerLoop->getParentLoop();
     if (!OuterLoop)
       continue;
     FlattenInfo FI(OuterLoop, InnerLoop);
-    Changed |= FlattenLoopPair(FI, DT, LI, SE, AC, TTI);
+    Changed |= FlattenLoopPair(FI, DT, LI, SE, AC, TTI, U);
   }
   return Changed;
 }
 
 PreservedAnalyses LoopFlattenPass::run(LoopNest &LN, LoopAnalysisManager &LAM,
                                        LoopStandardAnalysisResults &AR,
                                        LPMUpdater &U) {
 
   bool Changed = false;
 
   // The loop flattening pass requires loops to be
   // in simplified form, and also needs LCSSA. Running
   // this pass will simplify all loops that contain inner loops,
   // regardless of whether anything ends up being flattened.
-  Changed |= Flatten(LN, &AR.DT, &AR.LI, &AR.SE, &AR.AC, &AR.TTI);
+  Changed |= Flatten(LN, &AR.DT, &AR.LI, &AR.SE, &AR.AC, &AR.TTI, &U);
 
   if (!Changed)
     return PreservedAnalyses::all();
 
   return getLoopPassPreservedAnalyses();
 }
 
 namespace {
bool LoopFlattenLegacyPass::runOnFunction(Function &F) {
   auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>();
   DominatorTree *DT = DTWP ? &DTWP->getDomTree() : nullptr;
   auto &TTIP = getAnalysis<TargetTransformInfoWrapperPass>();
   auto *TTI = &TTIP.getTTI(F);
   auto *AC = &getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
   bool Changed = false;
   for (Loop *L : *LI) {
     auto LN = LoopNest::getLoopNest(*L, *SE);
-    Changed |= Flatten(*LN, DT, LI, SE, AC, TTI);
+    Changed |= Flatten(*LN, DT, LI, SE, AC, TTI, nullptr);
   }
   return Changed;
 }