Use LoopNest as the primary object on which LoopInterchange operates

[TaWeiTu]

Summary

Replace the usage of LoopList in LoopInterchange completely with LoopNest.
Due to this change, we no longer need to explicitly pass OuterLoop and InnerLoop to processLoop as in D96650.
Also, make interchangeNestedLoops a member function of LoopNest so that LoopNestAnalysis is preserved.

[Whitney]

llvm/include/llvm/Analysis/LoopNestAnalysis.h

165

Thinking we should have a more generic set function, like setLoops?

llvm/lib/Transforms/Scalar/LoopInterchange.cpp

504

why this change?

[TaWeiTu]

llvm/include/llvm/Analysis/LoopNestAnalysis.h

165

Agreed. But I also think having interchangeNestedLoops here is good in the sense that we can do the asserts below.
Should I add setLoop in this patch?

llvm/lib/Transforms/Scalar/LoopInterchange.cpp

504

Though unrelated to this patch, I think this makes the code a bit clearer since Interchanged is always true when we reach here, and the change seems too trivial for a separate patch.
I can split this into another patch or just push a NFC commit of course, if you think that makes the purpose of this patch clearer.

[Whitney]

llvm/include/llvm/Analysis/LoopNestAnalysis.h

165

Can we have a generic set function with a generic assert (guarded by expensive checks macro)?
like assert Loops[0] is the outermost loop, there are no other loops in between Loops[i] and Loops[i+1], Loops[Loops.size()-1] is the innermost loop, etc...

then we can have interchangeNestedLoops in interchange pass?

llvm/lib/Transforms/Scalar/LoopInterchange.cpp

504

I see it now, let's just keep it here.

[TaWeiTu]

llvm/include/llvm/Analysis/LoopNestAnalysis.h

165

OK. But I think having a separate patch for that is more appropriate because this patch is primarily about removing the usage of LoopList and synchronizing the API between legacy and new pass manager.
Can I leave interchangeNestedLoops here for now and move it back to LoopInterchange after/in the patch of setLoops?

[Whitney]

llvm/include/llvm/Analysis/LoopNestAnalysis.h

165

ok, can you please add a TODO comment?

[TaWeiTu]

• Add TODOs

[TaWeiTu]

llvm/include/llvm/Analysis/LoopNestAnalysis.h

165

Thanks! Added TODO.

llvm/include/llvm/Analysis/LoopNestAnalysis.h

public:
 
   /// Return the function to which the loop-nest belongs.
   Function *getParent() const {
     return Loops.front()->getHeader()->getParent();
   }
 
   StringRef getName() const { return Loops.front()->getName(); }
 
+  /// Interchange two loops. The outer loop is specified by \p OuterLoopId while
+  /// the inner loop is specified by \p InnerLoopId. Interchange is only allowed
+  /// for nested loops pair, that is, the inner loop should be the only subloop
+  /// of the outer loop.
+  //
+  // TODO: Implement a generic setLoops function with generic assertions and
+  // move interchangeNestedLoops into LoopInterchange pass.
+  void interchangeNestedLoops(unsigned OuterLoopId, unsigned InnerLoopId) {
+    assert((Loops[InnerLoopId]->getParentLoop() == Loops[OuterLoopId] &&
+            Loops[OuterLoopId]->getSubLoops().size() == 1) &&
+           "Loops that are interchanged should be nested.");
+    std::swap(Loops[OuterLoopId], Loops[InnerLoopId]);
+  }
+
 protected:
   const unsigned MaxPerfectDepth; // maximum perfect nesting depth level.
   LoopVectorTy Loops; // the loops in the nest (in breadth first order).
 };
 
 raw_ostream &operator<<(raw_ostream &, const LoopNest &);
 
 /// This analysis provides information for a loop nest. The analysis runs on

llvm/lib/Transforms/Scalar/LoopInterchange.cpp

for (auto &Row : DepMatrix) {
     for (auto D : Row)
       LLVM_DEBUG(dbgs() << D << " ");
     LLVM_DEBUG(dbgs() << "\n");
   }
 }
 #endif
 
 static bool populateDependencyMatrix(CharMatrix &DepMatrix, unsigned Level,
-                                     Loop *L, DependenceInfo *DI) {
+                                     const Loop *L, DependenceInfo *DI) {
   using ValueVector = SmallVector<Value *, 16>;
 
   ValueVector MemInstr;
 
   // For each block.
   for (BasicBlock *BB : L->blocks()) {
     // Scan the BB and collect legal loads and stores.
     for (Instruction &I : *BB) {
for (unsigned Row = 0; Row < NumRows; ++Row) {
     if (InnerDep == '*' || OuterDep == '*')
       return false;
     if (!validDepInterchange(DepMatrix, Row, OuterLoopId, InnerDep, OuterDep))
       return false;
   }
   return true;
 }
 
-static LoopVector populateWorklist(Loop &L) {
-  LLVM_DEBUG(dbgs() << "Calling populateWorklist on Func: "
-                    << L.getHeader()->getParent()->getName() << " Loop: %"
-                    << L.getHeader()->getName() << '\n');
-  LoopVector LoopList;
-  Loop *CurrentLoop = &L;
-  const std::vector<Loop *> *Vec = &CurrentLoop->getSubLoops();
-  while (!Vec->empty()) {
-    // The current loop has multiple subloops in it hence it is not tightly
-    // nested.
-    // Discard all loops above it added into Worklist.
-    if (Vec->size() != 1)
-      return {};
-
-    LoopList.push_back(CurrentLoop);
-    CurrentLoop = Vec->front();
-    Vec = &CurrentLoop->getSubLoops();
-  }
-  LoopList.push_back(CurrentLoop);
-  return LoopList;
-}
-
 static PHINode *getInductionVariable(Loop *L, ScalarEvolution *SE) {
   PHINode *InnerIndexVar = L->getCanonicalInductionVariable();
   if (InnerIndexVar)
     return InnerIndexVar;
   if (L->getLoopLatch() == nullptr || L->getLoopPredecessor() == nullptr)
     return nullptr;
   for (BasicBlock::iterator I = L->getHeader()->begin(); isa<PHINode>(I); ++I) {
     PHINode *PhiVar = cast<PHINode>(I);
struct LoopInterchange {
 
   /// Interface to emit optimization remarks.
   OptimizationRemarkEmitter *ORE;
 
   LoopInterchange(ScalarEvolution *SE, LoopInfo *LI, DependenceInfo *DI,
                   DominatorTree *DT, OptimizationRemarkEmitter *ORE)
       : SE(SE), LI(LI), DI(DI), DT(DT), ORE(ORE) {}
 
-  bool run(Loop *L) {
-    if (L->getParentLoop())
-      return false;
-
-    return processLoopList(populateWorklist(*L));
-  }
-
   bool run(LoopNest &LN) {
     const auto &LoopList = LN.getLoops();
     for (unsigned I = 1; I < LoopList.size(); ++I)
       if (LoopList[I]->getParentLoop() != LoopList[I - 1])
         return false;
-    return processLoopList(LoopList);
+    return processLoopNest(LN);
   }
 
-  bool isComputableLoopNest(ArrayRef<Loop *> LoopList) {
-    for (Loop *L : LoopList) {
+  bool isComputableLoopNest(LoopNest &LN) {
+    for (Loop *L : LN.getLoops()) {
       const SCEV *ExitCountOuter = SE->getBackedgeTakenCount(L);
       if (isa<SCEVCouldNotCompute>(ExitCountOuter)) {
         LLVM_DEBUG(dbgs() << "Couldn't compute backedge count\n");
         return false;
       }
       if (L->getNumBackEdges() != 1) {
         LLVM_DEBUG(dbgs() << "NumBackEdges is not equal to 1\n");
         return false;
       }
       if (!L->getExitingBlock()) {
         LLVM_DEBUG(dbgs() << "Loop doesn't have unique exit block\n");
         return false;
       }
     }
     return true;
   }
 
-  unsigned selectLoopForInterchange(ArrayRef<Loop *> LoopList) {
+  unsigned selectLoopForInterchange(LoopNest &LN) {
     // TODO: Add a better heuristic to select the loop to be interchanged based
     // on the dependence matrix. Currently we select the innermost loop.
-    return LoopList.size() - 1;
+    return LN.getNestDepth() - 1;
   }
 
-  bool processLoopList(ArrayRef<Loop *> LoopList) {
+  bool processLoopNest(LoopNest &LN) {
     bool Changed = false;
-    unsigned LoopNestDepth = LoopList.size();
+    unsigned LoopNestDepth = LN.getNestDepth();
     if (LoopNestDepth < 2) {
       LLVM_DEBUG(dbgs() << "Loop doesn't contain minimum nesting level.\n");
       return false;
     }
     if (LoopNestDepth > MaxLoopNestDepth) {
       LLVM_DEBUG(dbgs() << "Cannot handle loops of depth greater than "
                         << MaxLoopNestDepth << "\n");
       return false;
     }
-    if (!isComputableLoopNest(LoopList)) {
+    if (!isComputableLoopNest(LN)) {
       LLVM_DEBUG(dbgs() << "Not valid loop candidate for interchange\n");
       return false;
     }
 
     LLVM_DEBUG(dbgs() << "Processing LoopList of size = " << LoopNestDepth
                       << "\n");
 
     CharMatrix DependencyMatrix;
-    Loop *OuterMostLoop = *(LoopList.begin());
+    const Loop &OuterMostLoop = LN.getOutermostLoop();
     if (!populateDependencyMatrix(DependencyMatrix, LoopNestDepth,
-                                  OuterMostLoop, DI)) {
+                                  &OuterMostLoop, DI)) {
       LLVM_DEBUG(dbgs() << "Populating dependency matrix failed\n");
       return false;
     }
 #ifdef DUMP_DEP_MATRICIES
     LLVM_DEBUG(dbgs() << "Dependence before interchange\n");
     printDepMatrix(DependencyMatrix);
 #endif
 
     // Get the Outermost loop exit.
-    BasicBlock *LoopNestExit = OuterMostLoop->getExitBlock();
+    BasicBlock *LoopNestExit = OuterMostLoop.getExitBlock();
     if (!LoopNestExit) {
       LLVM_DEBUG(dbgs() << "OuterMostLoop needs an unique exit block");
       return false;
     }
 
-    unsigned SelecLoopId = selectLoopForInterchange(LoopList);
+    unsigned SelecLoopId = selectLoopForInterchange(LN);
     // Move the selected loop outwards to the best possible position.
-    Loop *LoopToBeInterchanged = LoopList[SelecLoopId];
     for (unsigned i = SelecLoopId; i > 0; i--) {
-      bool Interchanged = processLoop(LoopToBeInterchanged, LoopList[i - 1], i,
-                                      i - 1, LoopNestExit, DependencyMatrix);
+      bool Interchanged =
+          processLoop(LN, i, i - 1, LoopNestExit, DependencyMatrix);
       if (!Interchanged)
         return Changed;
       // Update the DependencyMatrix
       interChangeDependencies(DependencyMatrix, i, i - 1);
 #ifdef DUMP_DEP_MATRICIES
       LLVM_DEBUG(dbgs() << "Dependence after interchange\n");
       printDepMatrix(DependencyMatrix);
 #endif
-      Changed |= Interchanged;
+      Changed = true;
     }
     return Changed;
   }
 
-  bool processLoop(Loop *InnerLoop, Loop *OuterLoop, unsigned InnerLoopId,
-                   unsigned OuterLoopId, BasicBlock *LoopNestExit,
+  bool processLoop(LoopNest &LN, unsigned InnerLoopId, unsigned OuterLoopId,
+                   BasicBlock *LoopNestExit,
                    std::vector<std::vector<char>> &DependencyMatrix) {
     LLVM_DEBUG(dbgs() << "Processing InnerLoopId = " << InnerLoopId
                       << " and OuterLoopId = " << OuterLoopId << "\n");
+    Loop *OuterLoop = LN.getLoop(OuterLoopId);
+    Loop *InnerLoop = LN.getLoop(InnerLoopId);
     LoopInterchangeLegality LIL(OuterLoop, InnerLoop, SE, ORE);
     if (!LIL.canInterchangeLoops(InnerLoopId, OuterLoopId, DependencyMatrix)) {
       LLVM_DEBUG(dbgs() << "Not interchanging loops. Cannot prove legality.\n");
       return false;
     }
     LLVM_DEBUG(dbgs() << "Loops are legal to interchange\n");
     LoopInterchangeProfitability LIP(OuterLoop, InnerLoop, SE, ORE);
     if (!LIP.isProfitable(InnerLoopId, OuterLoopId, DependencyMatrix)) {
       LLVM_DEBUG(dbgs() << "Interchanging loops not profitable.\n");
       return false;
     }
 
     ORE->emit([&]() {
       return OptimizationRemark(DEBUG_TYPE, "Interchanged",
                                 InnerLoop->getStartLoc(),
                                 InnerLoop->getHeader())
              << "Loop interchanged with enclosing loop.";
     });
 
+    LN.interchangeNestedLoops(OuterLoopId, InnerLoopId);
     LoopInterchangeTransform LIT(OuterLoop, InnerLoop, SE, LI, DT, LoopNestExit,
                                  LIL);
     LIT.transform();
     LLVM_DEBUG(dbgs() << "Loops interchanged.\n");
     LoopsInterchanged++;
 
     assert(InnerLoop->isLCSSAForm(*DT) &&
            "Inner loop not left in LCSSA form after loop interchange!");
if (skipLoop(L))
       return false;
 
     auto *SE = &getAnalysis<ScalarEvolutionWrapperPass>().getSE();
     auto *LI = &getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
     auto *DI = &getAnalysis<DependenceAnalysisWrapperPass>().getDI();
     auto *DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
     auto *ORE = &getAnalysis<OptimizationRemarkEmitterWrapperPass>().getORE();
 
-    return LoopInterchange(SE, LI, DI, DT, ORE).run(L);
+    std::unique_ptr<LoopNest> LN = LoopNest::getLoopNest(*L, *SE);
+    return LoopInterchange(SE, LI, DI, DT, ORE).run(*LN);
   }
 };
 
 char LoopInterchangeLegacyPass::ID = 0;
 
 INITIALIZE_PASS_BEGIN(LoopInterchangeLegacyPass, "loop-interchange",
                       "Interchanges loops for cache reuse", false, false)
 INITIALIZE_PASS_DEPENDENCY(LoopPass)
PreservedAnalyses LoopInterchangePass::run(LoopNest &LN,
                                            LoopStandardAnalysisResults &AR,
                                            LPMUpdater &U) {
   Function &F = *LN.getParent();
 
   DependenceInfo DI(&F, &AR.AA, &AR.SE, &AR.LI);
   OptimizationRemarkEmitter ORE(&F);
   if (!LoopInterchange(&AR.SE, &AR.LI, &DI, &AR.DT, &ORE).run(LN))
     return PreservedAnalyses::all();
-  return getLoopPassPreservedAnalyses();
+  auto PA = getLoopPassPreservedAnalyses();
+  PA.preserve<LoopNestAnalysis>();
+  return PA;
 }