[SimpleLoopUnswitch] Inform pass manager when child loops are deleted

[bjope]

Summary

As part of the nontrivial unswitching we could end up removing child
loops. This patch add a notification to the pass manager when
that happens (using the markLoopAsDeleted callback).

Without this there could be stale LoopAccessAnalysis results cached
in the analysis manager. Those analysis results are cached based on
a Loop* as key. Since the BumpPtrAllocator used to allocate
Loop objects could be resetted between different runs of for
example the loop-distribute pass (running on different functions),
a new Loop object could be created using the same Loop pointer.
And then when requiring the LoopAccessAnalysis for the loop we
got the stale (corrupt) result from the destroyed loop.

Change-Id: Id976debe231ee6595752f501b36e719200834997

[aeubanks]

Do we need to do the same at the end of rebuildLoopAfterUnswitch()? Looks like another loop is getting destroyed there.

[bjope]

In D109257#2983151, @aeubanks wrote:

Do we need to do the same at the end of rebuildLoopAfterUnswitch()? Looks like another loop is getting destroyed there.

I actually added a similar DestroyedLoopCB callback in rebuildLoopAfterUnswitch() in an earlier draft of the patch. But then I realized that the destroy call in that place is for the base loop that is being unswitched. And as far as I could tell that boils down to rebuildLoopAfterUnswitch() returning false, and then unswitchNontrivialInvariants() will pass on information in the UnswitchCB call that the loops has been destroyed, and the markLoopAsDeleted call is made from that call back.
I was thinking that it would be nice to have a code comment explaining that next to the LI.destroy(&L) call. Seems like I forgot about adding such a comment.

[aeubanks]

llvm/test/Transforms/SimpleLoopUnswitch/nontrivial-unswitch-markloopasdeleted.ll

2

maybe to simplify this a bit you could just do require<access-info> instead of running a full optimization pass that uses it. it'd make this test a bit more robust

[aeubanks]

and thanks for debugging/fixing this! these sorts of issues are rough

[bjope]

llvm/test/Transforms/SimpleLoopUnswitch/nontrivial-unswitch-markloopasdeleted.ll

2

Well, I had the same idea. But decided not to do like that in the end.

My original problem was related to loop-distribute, which is a function pass. So in that situation there is a LoopAnalysisManagerFunctionProxy involved. If using require<access-info>, then I got the impression that things might be a bit different. But given that we basically only want to see that there is a "Clearing all analysis results" for the destroyed loop, then maybe it is ugly to use loop-distribute here.

[aeubanks]

llvm/test/Transforms/SimpleLoopUnswitch/nontrivial-unswitch-markloopasdeleted.ll

2

Ah makes sense. Yeah I think checking for the "Clearing all analysis results" is sufficient

[lebedev.ri]

This certainly should go into 13.x, please file a bug if there isn't one already.

[bjope]

In D109257#2983692, @lebedev.ri wrote:

This certainly should go into 13.x, please file a bug if there isn't one already.

Ok, noted. I filed a PR (and added release-13.0.0 as blocks): https://bugs.llvm.org/show_bug.cgi?id=51754

[aeubanks]

I believe the 13.0.0-rc3 release is today, we should probably get this in soon

[aeubanks]

In D109257#2988330, @aeubanks wrote:

I believe the 13.0.0-rc3 release is today, we should probably get this in soon

Sorry, not sure how I missed that this was already submitted.

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

deleteDeadClonedBlocks(Loop &L, ArrayRef<BasicBlock *> ExitBlocks,
   // Drop any remaining references to break cycles.
   for (BasicBlock *BB : DeadBlocks)
     BB->dropAllReferences();
   // Erase them from the IR.
   for (BasicBlock *BB : DeadBlocks)
     BB->eraseFromParent();
 }
 
-static void deleteDeadBlocksFromLoop(Loop &L,
-                                     SmallVectorImpl<BasicBlock *> &ExitBlocks,
-                                     DominatorTree &DT, LoopInfo &LI,
-                                     MemorySSAUpdater *MSSAU) {
+static void
+deleteDeadBlocksFromLoop(Loop &L,
+                         MemorySSAUpdater *MSSAU,
+                         function_ref<void(Loop &, StringRef)> DestroyLoopCB) {
   // Find all the dead blocks tied to this loop, and remove them from their
   // successors.
   SmallSetVector<BasicBlock *, 8> DeadBlockSet;
 
   // Start with loop/exit blocks and get a transitive closure of reachable dead
   // blocks.
   SmallVector<BasicBlock *, 16> DeathCandidates(ExitBlocks.begin(),
                                                 ExitBlocks.end());
if (!DeadBlockSet.count(ChildL->getHeader()))
       return false;
 
     assert(llvm::all_of(ChildL->blocks(),
                         [&](BasicBlock *ChildBB) {
                           return DeadBlockSet.count(ChildBB);
                         }) &&
            "If the child loop header is dead all blocks in the child loop must "
            "be dead as well!");
+    DestroyLoopCB(*ChildL, ChildL->getName());
     LI.destroy(ChildL);
     return true;
   });
 
   // Remove the loop mappings for the dead blocks and drop all the references
   // from these blocks to others to handle cyclic references as we start
   // deleting the blocks themselves.
   for (auto *BB : DeadBlockSet) {
static bool rebuildLoopAfterUnswitch(Loop &L, ArrayRef<BasicBlock *> ExitBlocks,
   // Actually delete the loop if nothing remained within it.
   if (Blocks.empty()) {
     assert(SubLoops.empty() &&
            "Failed to remove all subloops from the original loop!");
     if (Loop *ParentL = L.getParentLoop())
       ParentL->removeChildLoop(llvm::find(*ParentL, &L));
     else
       LI.removeLoop(llvm::find(LI, &L));
+    // markLoopAsDeleted for L should be triggered by the caller (it is typically
+    // done by using the UnswitchCB callback).
     LI.destroy(&L);
     return false;
   }
 
   return true;
 }
 
 /// Helper to visit a dominator subtree, invoking a callable on each node.
#endif
   } while (!DomWorklist.empty());
 }
 
 static void unswitchNontrivialInvariants(
     Loop &L, Instruction &TI, ArrayRef<Value *> Invariants,
     SmallVectorImpl<BasicBlock *> &ExitBlocks, IVConditionInfo &PartialIVInfo,
     DominatorTree &DT, LoopInfo &LI, AssumptionCache &AC,
     function_ref<void(bool, bool, ArrayRef<Loop *>)> UnswitchCB,
-    ScalarEvolution *SE, MemorySSAUpdater *MSSAU) {
+    ScalarEvolution *SE, MemorySSAUpdater *MSSAU,
+    function_ref<void(Loop &, StringRef)> DestroyLoopCB) {
   auto *ParentBB = TI.getParent();
   BranchInst *BI = dyn_cast<BranchInst>(&TI);
   SwitchInst *SI = BI ? nullptr : cast<SwitchInst>(&TI);
 
   // We can only unswitch switches, conditional branches with an invariant
   // condition, or combining invariant conditions with an instruction or
   // partially invariant instructions.
   assert((SI || (BI && BI->isConditional())) &&
static void unswitchNontrivialInvariants(
   // handles inserting all the cloned blocks into the correct loops.
   SmallVector<Loop *, 4> NonChildClonedLoops;
   for (std::unique_ptr<ValueToValueMapTy> &VMap : VMaps)
     buildClonedLoops(L, ExitBlocks, *VMap, LI, NonChildClonedLoops);
 
   // Now that our cloned loops have been built, we can update the original loop.
   // First we delete the dead blocks from it and then we rebuild the loop
   // structure taking these deletions into account.
-  deleteDeadBlocksFromLoop(L, ExitBlocks, DT, LI, MSSAU);
+  deleteDeadBlocksFromLoop(L, ExitBlocks, DT, LI, MSSAU, DestroyLoopCB);
 
   if (MSSAU && VerifyMemorySSA)
     MSSAU->getMemorySSA()->verifyMemorySSA();
 
   SmallVector<Loop *, 4> HoistedLoops;
   bool IsStillLoop = rebuildLoopAfterUnswitch(L, ExitBlocks, LI, HoistedLoops);
 
   if (MSSAU && VerifyMemorySSA)
LLVM_DEBUG(dbgs() << "  Computed multiplier  " << CostMultiplier
                     << " for unswitch candidate: " << TI << "\n");
   return CostMultiplier;
 }
 
 static bool unswitchBestCondition(
     Loop &L, DominatorTree &DT, LoopInfo &LI, AssumptionCache &AC,
     AAResults &AA, TargetTransformInfo &TTI,
     function_ref<void(bool, bool, ArrayRef<Loop *>)> UnswitchCB,
-    ScalarEvolution *SE, MemorySSAUpdater *MSSAU) {
+    ScalarEvolution *SE, MemorySSAUpdater *MSSAU,
+    function_ref<void(Loop &, StringRef)> DestroyLoopCB) {
   // Collect all invariant conditions within this loop (as opposed to an inner
   // loop which would be handled when visiting that inner loop).
   SmallVector<std::pair<Instruction *, TinyPtrVector<Value *>>, 4>
       UnswitchCandidates;
 
   // Whether or not we should also collect guards in the loop.
   bool CollectGuards = false;
   if (UnswitchGuards) {
if (isGuard(BestUnswitchTI))
     BestUnswitchTI = turnGuardIntoBranch(cast<IntrinsicInst>(BestUnswitchTI), L,
                                          ExitBlocks, DT, LI, MSSAU);
 
   LLVM_DEBUG(dbgs() << "  Unswitching non-trivial (cost = "
                     << BestUnswitchCost << ") terminator: " << *BestUnswitchTI
                     << "\n");
   unswitchNontrivialInvariants(L, *BestUnswitchTI, BestUnswitchInvariants,
                                ExitBlocks, PartialIVInfo, DT, LI, AC,
-                               UnswitchCB, SE, MSSAU);
+                               UnswitchCB, SE, MSSAU, DestroyLoopCB);
   return true;
 }
 
 /// Unswitch control flow predicated on loop invariant conditions.
 ///
 /// This first hoists all branches or switches which are trivial (IE, do not
 /// require duplicating any part of the loop) out of the loop body. It then
 /// looks at other loop invariant control flows and tries to unswitch those as

 ///
 /// If `SE` is non-null, we will update that analysis based on the unswitching
 /// done.
 static bool
 unswitchLoop(Loop &L, DominatorTree &DT, LoopInfo &LI, AssumptionCache &AC,
              AAResults &AA, TargetTransformInfo &TTI, bool Trivial,
              bool NonTrivial,
              function_ref<void(bool, bool, ArrayRef<Loop *>)> UnswitchCB,
-             ScalarEvolution *SE, MemorySSAUpdater *MSSAU) {
+             ScalarEvolution *SE, MemorySSAUpdater *MSSAU,
+             function_ref<void(Loop &, StringRef)> DestroyLoopCB) {
   assert(L.isRecursivelyLCSSAForm(DT, LI) &&
          "Loops must be in LCSSA form before unswitching.");
 
   // Must be in loop simplified form: we need a preheader and dedicated exits.
   if (!L.isLoopSimplifyForm())
     return false;
 
   // Try trivial unswitch first before loop over other basic blocks in the loop.
unswitchLoop(Loop &L, DominatorTree &DT, LoopInfo &LI, AssumptionCache &AC,
   // For non-trivial unswitching, because it often creates new loops, we rely on
   // the pass manager to iterate on the loops rather than trying to immediately
   // reach a fixed point. There is no substantial advantage to iterating
   // internally, and if any of the new loops are simplified enough to contain
   // trivial unswitching we want to prefer those.
 
   // Try to unswitch the best invariant condition. We prefer this full unswitch to
   // a partial unswitch when possible below the threshold.
-  if (unswitchBestCondition(L, DT, LI, AC, AA, TTI, UnswitchCB, SE, MSSAU))
+  if (unswitchBestCondition(L, DT, LI, AC, AA, TTI, UnswitchCB, SE, MSSAU,
+                            DestroyLoopCB))
     return true;
 
   // No other opportunities to unswitch.
   return false;
 }
 
 PreservedAnalyses SimpleLoopUnswitchPass::run(Loop &L, LoopAnalysisManager &AM,
                                               LoopStandardAnalysisResults &AR,
if (CurrentLoopValid) {
             {DisableUnswitchMD});
         L.setLoopID(NewLoopID);
       } else
         U.revisitCurrentLoop();
     } else
       U.markLoopAsDeleted(L, LoopName);
   };
 
+  auto DestroyLoopCB = [&U](Loop &L, StringRef Name) {
+    U.markLoopAsDeleted(L, Name);
+  };
+
   Optional<MemorySSAUpdater> MSSAU;
   if (AR.MSSA) {
     MSSAU = MemorySSAUpdater(AR.MSSA);
     if (VerifyMemorySSA)
       AR.MSSA->verifyMemorySSA();
   }
   if (!unswitchLoop(L, AR.DT, AR.LI, AR.AC, AR.AA, AR.TTI, Trivial, NonTrivial,
                     UnswitchCB, &AR.SE,
-                    MSSAU.hasValue() ? MSSAU.getPointer() : nullptr))
+                    MSSAU.hasValue() ? MSSAU.getPointer() : nullptr,
+                    DestroyLoopCB))
     return PreservedAnalyses::all();
 
   if (AR.MSSA && VerifyMemorySSA)
     AR.MSSA->verifyMemorySSA();
 
   // Historically this pass has had issues with the dominator tree so verify it
   // in asserts builds.
   assert(AR.DT.verify(DominatorTree::VerificationLevel::Fast));
if (CurrentLoopValid) {
       // condition, we should not re-add the current loop to avoid unswitching
       // on the same condition again.
       if (!PartiallyInvariant)
         LPM.addLoop(*L);
     } else
       LPM.markLoopAsDeleted(*L);
   };
 
+  auto DestroyLoopCB = [&LPM](Loop &L, StringRef /* Name */) {
+    LPM.markLoopAsDeleted(L);
+  };
+
   if (VerifyMemorySSA)
     MSSA->verifyMemorySSA();
 
   bool Changed = unswitchLoop(*L, DT, LI, AC, AA, TTI, true, NonTrivial,
-                              UnswitchCB, SE, &MSSAU);
+                              UnswitchCB, SE, &MSSAU, DestroyLoopCB);
 
   if (VerifyMemorySSA)
     MSSA->verifyMemorySSA();
 
   // Historically this pass has had issues with the dominator tree so verify it
   // in asserts builds.
   assert(DT.verify(DominatorTree::VerificationLevel::Fast));
 

llvm/test/Transforms/SimpleLoopUnswitch/nontrivial-unswitch-markloopasdeleted.ll

+; RUN: opt < %s -enable-loop-distribute -passes='loop-distribute,loop-mssa(simple-loop-unswitch<nontrivial>),loop-distribute' -o /dev/null -S -debug-pass-manager=verbose 2>&1 | FileCheck %s
+
+
+; Running loop-distribute will result in LoopAccessAnalysis being required and
+; cached in the LoopAnalysisManagerFunctionProxy.
+;
+; CHECK: Running analysis: LoopAccessAnalysis on Loop at depth 2 containing: %loop_a_inner<header><latch><exiting>
+
+
+; Then simple-loop-unswitch is removing/replacing some loops (resulting in
+; Loop objects used as key in the analyses cache is destroyed). So here we
+; want to see that any analysis results cached on the destroyed loop is
+; cleared. A special case here is that loop_a_inner is destroyed when
+; unswitching the parent loop.
+;
+; The bug solved and verified by this test case was related to the
+; SimpleLoopUnswitch not marking the Loop as removed, so we missed clearing
+; the analysis caches.
+;
+; CHECK: Running pass: SimpleLoopUnswitchPass on Loop at depth 1 containing: %loop_begin<header>,%loop_b,%loop_b_inner,%loop_b_inner_exit,%loop_a,%loop_a_inner,%loop_a_inner_exit,%latch<latch><exiting>
+; CHECK-NEXT: Clearing all analysis results for: loop_a_inner
+
+
+; When running loop-distribute the second time we can see that loop_a_inner
+; isn't analysed because the loop no longer exists (instead we find a new loop,
+; loop_a_inner.us). This kind of verifies that it was correct to remove the
+; loop_a_inner related analysis above.
+;
+; CHECK: Running analysis: LoopAccessAnalysis on Loop at depth 2 containing: %loop_a_inner.us<header><latch><exiting>
+
+
+define i32 @test6(i1* %ptr, i1 %cond1, i32* %a.ptr, i32* %b.ptr) {
+entry:
+  br label %loop_begin
+
+loop_begin:
+  %v = load i1, i1* %ptr
+  br i1 %cond1, label %loop_a, label %loop_b
+
+loop_a:
+  br label %loop_a_inner
+
+loop_a_inner:
+  %va = load i1, i1* %ptr
+  %a = load i32, i32* %a.ptr
+  br i1 %va, label %loop_a_inner, label %loop_a_inner_exit
+
+loop_a_inner_exit:
+  %a.lcssa = phi i32 [ %a, %loop_a_inner ]
+  br label %latch
+
+loop_b:
+  br label %loop_b_inner
+
+loop_b_inner:
+  %vb = load i1, i1* %ptr
+  %b = load i32, i32* %b.ptr
+  br i1 %vb, label %loop_b_inner, label %loop_b_inner_exit
+
+loop_b_inner_exit:
+  %b.lcssa = phi i32 [ %b, %loop_b_inner ]
+  br label %latch
+
+latch:
+  %ab.phi = phi i32 [ %a.lcssa, %loop_a_inner_exit ], [ %b.lcssa, %loop_b_inner_exit ]
+  br i1 %v, label %loop_begin, label %loop_exit
+
+loop_exit:
+  %ab.lcssa = phi i32 [ %ab.phi, %latch ]
+  ret i32 %ab.lcssa
+}