[SimpleLoopUnswitch] Remove callbacks (#73300)

boomanaiden154 · web-flow · commit 4233176f9bf8 · 2023-11-24T11:36:52.000-08:00
After the removal of the legacyPM version of simple loop unswitch, there
is no longer a need for the callback mechanism to handle PM specific
tasks. This patch removes the callbacks to help simplify the code now
that they're no longer needed.
diff --git a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
@@ -1685,13 +1685,12 @@ deleteDeadClonedBlocks(Loop &L, ArrayRef<BasicBlock *> ExitBlocks,
     BB->eraseFromParent();
 }
 
-static void
-deleteDeadBlocksFromLoop(Loop &L,
-                         SmallVectorImpl<BasicBlock *> &ExitBlocks,
-                         DominatorTree &DT, LoopInfo &LI,
-                         MemorySSAUpdater *MSSAU,
-                         ScalarEvolution *SE,
-                         function_ref<void(Loop &, StringRef)> DestroyLoopCB) {
+static void deleteDeadBlocksFromLoop(Loop &L,
+                                     SmallVectorImpl<BasicBlock *> &ExitBlocks,
+                                     DominatorTree &DT, LoopInfo &LI,
+                                     MemorySSAUpdater *MSSAU,
+                                     ScalarEvolution *SE,
+                                     LPMUpdater &LoopUpdater) {
   // Find all the dead blocks tied to this loop, and remove them from their
   // successors.
   SmallSetVector<BasicBlock *, 8> DeadBlockSet;
@@ -1741,7 +1740,7 @@ deleteDeadBlocksFromLoop(Loop &L,
                         }) &&
            "If the child loop header is dead all blocks in the child loop must "
            "be dead as well!");
-    DestroyLoopCB(*ChildL, ChildL->getName());
+    LoopUpdater.markLoopAsDeleted(*ChildL, ChildL->getName());
     if (SE)
       SE->forgetBlockAndLoopDispositions();
     LI.destroy(ChildL);
@@ -2085,8 +2084,8 @@ static bool rebuildLoopAfterUnswitch(Loop &L, ArrayRef<BasicBlock *> ExitBlocks,
       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).
+    // markLoopAsDeleted for L should be triggered by the caller (it is
+    // typically done within postUnswitch).
     if (SE)
       SE->forgetBlockAndLoopDispositions();
     LI.destroy(&L);
@@ -2123,18 +2122,56 @@ void visitDomSubTree(DominatorTree &DT, BasicBlock *BB, CallableT Callable) {
   } while (!DomWorklist.empty());
 }
 
+void postUnswitch(Loop &L, LPMUpdater &U, StringRef LoopName,
+                  bool CurrentLoopValid, bool PartiallyInvariant,
+                  bool InjectedCondition, ArrayRef<Loop *> NewLoops) {
+  // If we did a non-trivial unswitch, we have added new (cloned) loops.
+  if (!NewLoops.empty())
+    U.addSiblingLoops(NewLoops);
+
+  // If the current loop remains valid, we should revisit it to catch any
+  // other unswitch opportunities. Otherwise, we need to mark it as deleted.
+  if (CurrentLoopValid) {
+    if (PartiallyInvariant) {
+      // Mark the new loop as partially unswitched, to avoid unswitching on
+      // the same condition again.
+      auto &Context = L.getHeader()->getContext();
+      MDNode *DisableUnswitchMD = MDNode::get(
+          Context,
+          MDString::get(Context, "llvm.loop.unswitch.partial.disable"));
+      MDNode *NewLoopID = makePostTransformationMetadata(
+          Context, L.getLoopID(), {"llvm.loop.unswitch.partial"},
+          {DisableUnswitchMD});
+      L.setLoopID(NewLoopID);
+    } else if (InjectedCondition) {
+      // Do the same for injection of invariant conditions.
+      auto &Context = L.getHeader()->getContext();
+      MDNode *DisableUnswitchMD = MDNode::get(
+          Context,
+          MDString::get(Context, "llvm.loop.unswitch.injection.disable"));
+      MDNode *NewLoopID = makePostTransformationMetadata(
+          Context, L.getLoopID(), {"llvm.loop.unswitch.injection"},
+          {DisableUnswitchMD});
+      L.setLoopID(NewLoopID);
+    } else
+      U.revisitCurrentLoop();
+  } else
+    U.markLoopAsDeleted(L, LoopName);
+}
+
 static void unswitchNontrivialInvariants(
     Loop &L, Instruction &TI, ArrayRef<Value *> Invariants,
     IVConditionInfo &PartialIVInfo, DominatorTree &DT, LoopInfo &LI,
-    AssumptionCache &AC,
-    function_ref<void(bool, bool, bool, ArrayRef<Loop *>)> UnswitchCB,
-    ScalarEvolution *SE, MemorySSAUpdater *MSSAU,
-    function_ref<void(Loop &, StringRef)> DestroyLoopCB, bool InsertFreeze,
-    bool InjectedCondition) {
+    AssumptionCache &AC, ScalarEvolution *SE, MemorySSAUpdater *MSSAU,
+    LPMUpdater &LoopUpdater, bool InsertFreeze, bool InjectedCondition) {
   auto *ParentBB = TI.getParent();
   BranchInst *BI = dyn_cast<BranchInst>(&TI);
   SwitchInst *SI = BI ? nullptr : cast<SwitchInst>(&TI);
 
+  // Save the current loop name in a variable so that we can report it even
+  // after it has been deleted.
+  std::string LoopName(L.getName());
+
   // We can only unswitch switches, conditional branches with an invariant
   // condition, or combining invariant conditions with an instruction or
   // partially invariant instructions.
@@ -2447,7 +2484,7 @@ static void unswitchNontrivialInvariants(
   // 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, SE,DestroyLoopCB);
+  deleteDeadBlocksFromLoop(L, ExitBlocks, DT, LI, MSSAU, SE, LoopUpdater);
 
   if (MSSAU && VerifyMemorySSA)
     MSSAU->getMemorySSA()->verifyMemorySSA();
@@ -2583,7 +2620,8 @@ static void unswitchNontrivialInvariants(
   for (Loop *UpdatedL : llvm::concat<Loop *>(NonChildClonedLoops, HoistedLoops))
     if (UpdatedL->getParentLoop() == ParentL)
       SibLoops.push_back(UpdatedL);
-  UnswitchCB(IsStillLoop, PartiallyInvariant, InjectedCondition, SibLoops);
+  postUnswitch(L, LoopUpdater, LoopName, IsStillLoop, PartiallyInvariant,
+               InjectedCondition, SibLoops);
 
   if (MSSAU && VerifyMemorySSA)
     MSSAU->getMemorySSA()->verifyMemorySSA();
@@ -3430,12 +3468,11 @@ static bool shouldInsertFreeze(Loop &L, Instruction &TI, DominatorTree &DT,
       Cond, &AC, L.getLoopPreheader()->getTerminator(), &DT);
 }
 
-static bool unswitchBestCondition(
-    Loop &L, DominatorTree &DT, LoopInfo &LI, AssumptionCache &AC,
-    AAResults &AA, TargetTransformInfo &TTI,
-    function_ref<void(bool, bool, bool, ArrayRef<Loop *>)> UnswitchCB,
-    ScalarEvolution *SE, MemorySSAUpdater *MSSAU,
-    function_ref<void(Loop &, StringRef)> DestroyLoopCB) {
+static bool unswitchBestCondition(Loop &L, DominatorTree &DT, LoopInfo &LI,
+                                  AssumptionCache &AC, AAResults &AA,
+                                  TargetTransformInfo &TTI, ScalarEvolution *SE,
+                                  MemorySSAUpdater *MSSAU,
+                                  LPMUpdater &LoopUpdater) {
   // Collect all invariant conditions within this loop (as opposed to an inner
   // loop which would be handled when visiting that inner loop).
   SmallVector<NonTrivialUnswitchCandidate, 4> UnswitchCandidates;
@@ -3498,8 +3535,8 @@ static bool unswitchBestCondition(
   LLVM_DEBUG(dbgs() << "  Unswitching non-trivial (cost = " << Best.Cost
                     << ") terminator: " << *Best.TI << "\n");
   unswitchNontrivialInvariants(L, *Best.TI, Best.Invariants, PartialIVInfo, DT,
-                               LI, AC, UnswitchCB, SE, MSSAU, DestroyLoopCB,
-                               InsertFreeze, InjectedCondition);
+                               LI, AC, SE, MSSAU, LoopUpdater, InsertFreeze,
+                               InjectedCondition);
   return true;
 }
 
@@ -3518,20 +3555,18 @@ static bool unswitchBestCondition(
 /// true, we will attempt to do non-trivial unswitching as well as trivial
 /// unswitching.
 ///
-/// The `UnswitchCB` callback provided will be run after unswitching is
-/// complete, with the first parameter set to `true` if the provided loop
-/// remains a loop, and a list of new sibling loops created.
+/// The `postUnswitch` function will be run after unswitching is complete
+/// with information on whether or not the provided loop remains a loop and
+/// a list of new sibling loops created.
 ///
 /// 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, bool, ArrayRef<Loop *>)> UnswitchCB,
-             ScalarEvolution *SE, MemorySSAUpdater *MSSAU,
-             ProfileSummaryInfo *PSI, BlockFrequencyInfo *BFI,
-             function_ref<void(Loop &, StringRef)> DestroyLoopCB) {
+static bool unswitchLoop(Loop &L, DominatorTree &DT, LoopInfo &LI,
+                         AssumptionCache &AC, AAResults &AA,
+                         TargetTransformInfo &TTI, bool Trivial,
+                         bool NonTrivial, ScalarEvolution *SE,
+                         MemorySSAUpdater *MSSAU, ProfileSummaryInfo *PSI,
+                         BlockFrequencyInfo *BFI, LPMUpdater &LoopUpdater) {
   assert(L.isRecursivelyLCSSAForm(DT, LI) &&
          "Loops must be in LCSSA form before unswitching.");
 
@@ -3543,8 +3578,9 @@ unswitchLoop(Loop &L, DominatorTree &DT, LoopInfo &LI, AssumptionCache &AC,
   if (Trivial && unswitchAllTrivialConditions(L, DT, LI, SE, MSSAU)) {
     // If we unswitched successfully we will want to clean up the loop before
     // processing it further so just mark it as unswitched and return.
-    UnswitchCB(/*CurrentLoopValid*/ true, /*PartiallyInvariant*/ false,
-               /*InjectedCondition*/ false, {});
+    postUnswitch(L, LoopUpdater, L.getName(),
+                 /*CurrentLoopValid*/ true, /*PartiallyInvariant*/ false,
+                 /*InjectedCondition*/ false, {});
     return true;
   }
 
@@ -3613,8 +3649,7 @@ unswitchLoop(Loop &L, DominatorTree &DT, LoopInfo &LI, AssumptionCache &AC,
 
   // 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,
-                            DestroyLoopCB))
+  if (unswitchBestCondition(L, DT, LI, AC, AA, TTI, SE, MSSAU, LoopUpdater))
     return true;
 
   // No other opportunities to unswitch.
@@ -3634,61 +3669,14 @@ PreservedAnalyses SimpleLoopUnswitchPass::run(Loop &L, LoopAnalysisManager &AM,
   LLVM_DEBUG(dbgs() << "Unswitching loop in " << F.getName() << ": " << L
                     << "\n");
 
-  // Save the current loop name in a variable so that we can report it even
-  // after it has been deleted.
-  std::string LoopName = std::string(L.getName());
-
-  auto UnswitchCB = [&L, &U, &LoopName](bool CurrentLoopValid,
-                                        bool PartiallyInvariant,
-                                        bool InjectedCondition,
-                                        ArrayRef<Loop *> NewLoops) {
-    // If we did a non-trivial unswitch, we have added new (cloned) loops.
-    if (!NewLoops.empty())
-      U.addSiblingLoops(NewLoops);
-
-    // If the current loop remains valid, we should revisit it to catch any
-    // other unswitch opportunities. Otherwise, we need to mark it as deleted.
-    if (CurrentLoopValid) {
-      if (PartiallyInvariant) {
-        // Mark the new loop as partially unswitched, to avoid unswitching on
-        // the same condition again.
-        auto &Context = L.getHeader()->getContext();
-        MDNode *DisableUnswitchMD = MDNode::get(
-            Context,
-            MDString::get(Context, "llvm.loop.unswitch.partial.disable"));
-        MDNode *NewLoopID = makePostTransformationMetadata(
-            Context, L.getLoopID(), {"llvm.loop.unswitch.partial"},
-            {DisableUnswitchMD});
-        L.setLoopID(NewLoopID);
-      } else if (InjectedCondition) {
-        // Do the same for injection of invariant conditions.
-        auto &Context = L.getHeader()->getContext();
-        MDNode *DisableUnswitchMD = MDNode::get(
-            Context,
-            MDString::get(Context, "llvm.loop.unswitch.injection.disable"));
-        MDNode *NewLoopID = makePostTransformationMetadata(
-            Context, L.getLoopID(), {"llvm.loop.unswitch.injection"},
-            {DisableUnswitchMD});
-        L.setLoopID(NewLoopID);
-      } else
-        U.revisitCurrentLoop();
-    } else
-      U.markLoopAsDeleted(L, LoopName);
-  };
-
-  auto DestroyLoopCB = [&U](Loop &L, StringRef Name) {
-    U.markLoopAsDeleted(L, Name);
-  };
-
   std::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 ? &*MSSAU : nullptr, PSI, AR.BFI,
-                    DestroyLoopCB))
+                    &AR.SE, MSSAU ? &*MSSAU : nullptr, PSI, AR.BFI, U))
     return PreservedAnalyses::all();
 
   if (AR.MSSA && VerifyMemorySSA)
