[LoopNest] Extend `LPMUpdater` and adaptor to handle loop-nest passes

[TaWeiTu]

Summary

This is a follow-up patch of D87045.

The patch implements "loop-nest mode" for LPMUpdater and FunctionToLoopPassAdaptor in which only top-level loops are operated.

createFunctionToLoopPassAdaptor decides whether the returned adaptor is in loop-nest mode or not based on the given pass. If the pass is a loop-nest pass or the pass is a LoopPassManager which contains only loop-nest passes, the loop-nest version of adaptor is returned; otherwise, the normal (loop) version of adaptor is returned.

[Whitney]

LGTM.

[TaWeiTu]

Rebased on top of the current master. Add comments about loop/loop-nest mode.

[asbirlea]

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

243

s/will no/will not

258

legibility:
(!LoopNestMode || !L.getParentLoop())
or more detailed
(!LoopNestMode || (LoopNestMode && !L.getParentLoop()))
?

[TaWeiTu]

Address comments.

[TaWeiTu]

Rebased.

[asbirlea]

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

590

Could you clarify this?

[TaWeiTu]

Clarify the specializations of createFunctionToLoopPassAdaptor()

[TaWeiTu]

Rebase

[TaWeiTu]

Since D87045 has been landed, is there any objection against committing this patch as well?
I'll wait for a few days to see whether the reviewers are happy with the patch or not before landing it.
Thanks!

[Whitney]

Since D87045 has been landed, is there any objection against committing this patch as well?

No objection from me.

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

 #include "llvm/Transforms/Utils/LoopSimplify.h"
 #include "llvm/Transforms/Utils/LoopUtils.h"
 
 namespace llvm {
 
 // Forward declarations of an update tracking API used in the pass manager.
 class LPMUpdater;
 
+namespace {
+
+template <typename PassT>
+using HasRunOnLoopT = decltype(std::declval<PassT>().run(
+    std::declval<Loop &>(), std::declval<LoopAnalysisManager &>(),
+    std::declval<LoopStandardAnalysisResults &>(),
+    std::declval<LPMUpdater &>()));
+
+} // namespace
+
 // Explicit specialization and instantiation declarations for the pass manager.
 // See the comments on the definition of the specialization for details on how
 // it differs from the primary template.
 template <>
 class PassManager<Loop, LoopAnalysisManager, LoopStandardAnalysisResults &,
                   LPMUpdater &>
     : public PassInfoMixin<
           PassManager<Loop, LoopAnalysisManager, LoopStandardAnalysisResults &,
                       LPMUpdater &>> {
-private:
-  template <typename PassT>
-  using HasRunOnLoopT = decltype(std::declval<PassT>().run(
-      std::declval<Loop &>(), std::declval<LoopAnalysisManager &>(),
-      std::declval<LoopStandardAnalysisResults &>(),
-      std::declval<LPMUpdater &>()));
-
 public:
   /// Construct a pass manager.
   ///
   /// If \p DebugLogging is true, we'll log our progress to llvm::dbgs().
   explicit PassManager(bool DebugLogging = false)
       : DebugLogging(DebugLogging) {}
 
   // FIXME: These are equivalent to the default move constructor/move
using RepeatedLoopNestPassModelT =
                           LPMUpdater &>;
     IsLoopNestPass.push_back(true);
     LoopNestPasses.emplace_back(
         new RepeatedLoopNestPassModelT(std::move(Pass)));
   }
 
   static bool isRequired() { return true; }
 
+  size_t getNumLoopPasses() const { return LoopPasses.size(); }
+  size_t getNumLoopNestPasses() const { return LoopNestPasses.size(); }
+
 protected:
   using LoopPassConceptT =
       detail::PassConcept<Loop, LoopAnalysisManager,
                           LoopStandardAnalysisResults &, LPMUpdater &>;
   using LoopNestPassConceptT =
       detail::PassConcept<LoopNest, LoopAnalysisManager,
                           LoopStandardAnalysisResults &, LPMUpdater &>;
 

 /// Loop pass, and Loop passes should use it to update LPM infrastructure if
 /// they modify the loop nest structure.
 class LPMUpdater {
 public:
   /// This can be queried by loop passes which run other loop passes (like pass
   /// managers) to know whether the loop needs to be skipped due to updates to
   /// the loop nest.
   ///
   /// If this returns true, the loop object may have been deleted, so passes
   /// should take care not to touch the object.
   bool skipCurrentLoop() const { return SkipCurrentLoop; }
 
   /// Loop passes should use this method to indicate they have deleted a loop
   /// from the nest.
   ///
   /// 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 && L.getParentLoop()) &&
+           "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.");
     if (&L == CurrentL)
       SkipCurrentLoop = true;
   }
 
   /// Loop passes should use this method to indicate they have added new child
   /// loops of the current loop.
   ///
   /// \p NewChildLoops must contain only the immediate children. Any nested
   /// loops within them will be visited in postorder as usual for the loop pass
   /// manager.
   void addChildLoops(ArrayRef<Loop *> NewChildLoops) {
+    assert(!LoopNestMode &&
+           "Child loops should not be pushed in loop-nest mode.");
     // Insert ourselves back into the worklist first, as this loop should be
     // revisited after all the children have been processed.
     Worklist.insert(CurrentL);
 
 #ifndef NDEBUG
     for (Loop *NewL : NewChildLoops)
       assert(NewL->getParentLoop() == CurrentL && "All of the new loops must "
                                                   "be immediate children of "
#endif
   /// manager.
   void addSiblingLoops(ArrayRef<Loop *> NewSibLoops) {
 #ifndef NDEBUG
     for (Loop *NewL : NewSibLoops)
       assert(NewL->getParentLoop() == ParentL &&
              "All of the new loops must be siblings of the current loop!");
 #endif
 
+    if (LoopNestMode)
+      Worklist.insert(NewSibLoops);
+    else
       appendLoopsToWorklist(NewSibLoops, Worklist);
 
     // No need to skip the current loop or revisit it, as sibling loops
     // shouldn't impact anything.
   }
 
   /// Restart the current loop.
   ///
   /// Loop passes should call this method to indicate the current loop has been
private:
   /// The \c FunctionToLoopPassAdaptor's worklist of loops to process.
   SmallPriorityWorklist<Loop *, 4> &Worklist;
 
   /// The analysis manager for use in the current loop nest.
   LoopAnalysisManager &LAM;
 
   Loop *CurrentL;
   bool SkipCurrentLoop;
+  const bool LoopNestMode;
 
 #ifndef NDEBUG
   // In debug builds we also track the parent loop to implement asserts even in
   // the face of loop deletion.
   Loop *ParentL;
 #endif
 
   LPMUpdater(SmallPriorityWorklist<Loop *, 4> &Worklist,
-             LoopAnalysisManager &LAM)
-      : Worklist(Worklist), LAM(LAM) {}
+             LoopAnalysisManager &LAM, bool LoopNestMode = false)
+      : Worklist(Worklist), LAM(LAM), LoopNestMode(LoopNestMode) {}
 };
 
 template <typename IRUnitT, typename PassT>
 Optional<PreservedAnalyses> LoopPassManager::runSinglePass(
     IRUnitT &IR, PassT &Pass, LoopAnalysisManager &AM,
     LoopStandardAnalysisResults &AR, LPMUpdater &U, PassInstrumentation &PI) {
   // Check the PassInstrumentation's BeforePass callbacks before running the
   // pass, skip its execution completely if asked to (callback returns false).

 /// FunctionAnalysisManager it will run the \c LoopAnalysisManagerFunctionProxy
 /// analysis prior to running the loop passes over the function to enable a \c
 /// LoopAnalysisManager to be used within this run safely.
 template <typename LoopPassT>
 class FunctionToLoopPassAdaptor
     : public PassInfoMixin<FunctionToLoopPassAdaptor<LoopPassT>> {
 public:
   explicit FunctionToLoopPassAdaptor(LoopPassT Pass, bool UseMemorySSA = false,
-                                     bool DebugLogging = false)
+                                     bool DebugLogging = false,
+                                     bool LoopNestMode = false)
       : Pass(std::move(Pass)), LoopCanonicalizationFPM(DebugLogging),
-        UseMemorySSA(UseMemorySSA) {
+        UseMemorySSA(UseMemorySSA), LoopNestMode(LoopNestMode) {
     LoopCanonicalizationFPM.addPass(LoopSimplifyPass());
     LoopCanonicalizationFPM.addPass(LCSSAPass());
   }
 
   /// Runs the loop passes across every loop in the function.
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM) {
     // Before we even compute any loop analyses, first run a miniature function
     // pass pipeline to put loops into their canonical form. Note that we can
if (UseMemorySSA)
       LAMFP.markMSSAUsed();
     LoopAnalysisManager &LAM = LAMFP.getManager();
 
     // A postorder worklist of loops to process.
     SmallPriorityWorklist<Loop *, 4> Worklist;
 
     // Register the worklist and loop analysis manager so that loop passes can
     // update them when they mutate the loop nest structure.
-    LPMUpdater Updater(Worklist, LAM);
+    LPMUpdater Updater(Worklist, LAM, LoopNestMode);
 
     // Add the loop nests in the reverse order of LoopInfo. See method
     // declaration.
+    if (!LoopNestMode) {
       appendLoopsToWorklist(LI, Worklist);
+    } else {
+      for (Loop *L : LI)
+        Worklist.insert(L);
+    }
 
 #ifndef NDEBUG
     PI.pushBeforeNonSkippedPassCallback([&LAR, &LI](StringRef PassID, Any IR) {
       if (isSpecialPass(PassID, {"PassManager"}))
         return;
       assert(any_isa<const Loop *>(IR) || any_isa<const LoopNest *>(IR));
       const Loop *L = any_isa<const Loop *>(IR)
                           ? any_cast<const Loop *>(IR)
                           : &any_cast<const LoopNest *>(IR)->getOutermostLoop();
       assert(L && "Loop should be valid for printing");
 
       // Verify the loop structure and LCSSA form before visiting the loop.
       L->verifyLoop();
       assert(L->isRecursivelyLCSSAForm(LAR.DT, LI) &&
              "Loops must remain in LCSSA form!");
     });
 #endif
 
     do {
       Loop *L = Worklist.pop_back_val();
+      assert(!(LoopNestMode && L->getParentLoop()) &&
+             "L should be a top-level loop in loop-nest mode.");
 
       // Reset the update structure for this loop.
       Updater.CurrentL = L;
       Updater.SkipCurrentLoop = false;
 
 #ifndef NDEBUG
       // Save a parent loop pointer for asserts.
       Updater.ParentL = L->getParentLoop();
#endif
   static bool isRequired() { return true; }
 
 private:
   LoopPassT Pass;
 
   FunctionPassManager LoopCanonicalizationFPM;
 
   bool UseMemorySSA = false;
+  const bool LoopNestMode;
 };
 
 /// A function to deduce a loop pass type and wrap it in the templated
 /// adaptor.
 template <typename LoopPassT>
-FunctionToLoopPassAdaptor<LoopPassT>
+inline FunctionToLoopPassAdaptor<LoopPassT>
 createFunctionToLoopPassAdaptor(LoopPassT Pass, bool UseMemorySSA = false,
                                 bool DebugLogging = false) {
-  return FunctionToLoopPassAdaptor<LoopPassT>(std::move(Pass), UseMemorySSA,
-                                              DebugLogging);
+  return FunctionToLoopPassAdaptor<LoopPassT>(
+      std::move(Pass), UseMemorySSA, DebugLogging,
+      is_detected<HasRunOnLoopT, LoopPassT>::value);
+}
+
+template <>
+inline FunctionToLoopPassAdaptor<LoopPassManager>
+createFunctionToLoopPassAdaptor<LoopPassManager>(LoopPassManager LPM,
+                                                 bool UseMemorySSA,
+                                                 bool DebugLogging) {
+  bool LoopNestMode = (LPM.getNumLoopPasses() == 0);
+  return FunctionToLoopPassAdaptor<LoopPassManager>(
+      std::move(LPM), UseMemorySSA, DebugLogging, LoopNestMode);
 }
 
 /// Pass for printing a loop's contents as textual IR.
 class PrintLoopPass : public PassInfoMixin<PrintLoopPass> {
   raw_ostream &OS;
   std::string Banner;
 
 public:
   PrintLoopPass();
   PrintLoopPass(raw_ostream &OS, const std::string &Banner = "");
 
   PreservedAnalyses run(Loop &L, LoopAnalysisManager &,
                         LoopStandardAnalysisResults &, LPMUpdater &);
 };
 }
 
 #endif // LLVM_TRANSFORMS_SCALAR_LOOPPASSMANAGER_H