[Mlir-commits] [llvm] [mlir] [MLIR][OpenMP] add interchange operation in the OpenMP mlir dialect (PR #186381)

[Ferran Toda]

================
@@ -6785,6 +6783,89 @@ void OpenMPIRBuilder::unrollLoopHeuristic(DebugLoc, CanonicalLoopInfo *Loop) {
             });
 }
 
+std::vector<CanonicalLoopInfo *>
+OpenMPIRBuilder::interchangeLoops(DebugLoc DL,
+                                  ArrayRef<CanonicalLoopInfo *> Loops,
+                                  std::vector<int> Permutation) {
+  assert(Permutation.size() == Loops.size() &&
+         "The permutation must have every input loop");
+  int NumLoops = Loops.size();
+  assert(NumLoops >= 2 && "At least two loops to interchange required");
+
+  CanonicalLoopInfo *OutermostLoop = Loops.front();
+  CanonicalLoopInfo *InnermostLoop = Loops.back();
+  Function *F = OutermostLoop->getBody()->getParent();
+
+  // Loop control blocks that may become orphaned later.
+  SmallVector<BasicBlock *> OldControlBBs;
+  for (CanonicalLoopInfo *Loop : Loops)
+    Loop->collectControlBlocks(OldControlBBs);
+
+  // Create the new loop nest structure
+  std::vector<CanonicalLoopInfo *> Result;
+
+  BasicBlock *Enter = OutermostLoop->getPreheader();
+  BasicBlock *Continue = OutermostLoop->getAfter();
+  BasicBlock *OutroInsertBefore = InnermostLoop->getExit();
+
+  for (int i = 0; i < NumLoops; i++) {
+    int loop_n = Permutation[i] - 1;
+    assert(Loops[loop_n]->isValid() &&
+           "All input loops must be valid canonical loops");
+
+    CanonicalLoopInfo *newLoop = createLoopSkeleton(
+        DL, Loops[loop_n]->getTripCount(), F, InnermostLoop->getBody(),
+        OutroInsertBefore, "interchange" + Twine(i));
+    redirectAllPredecessorsTo(Loops[i]->getPreheader(), newLoop->getPreheader(),
+                              DL);
+    redirectTo(Enter, newLoop->getPreheader(), DL);
+    redirectTo(newLoop->getAfter(), Continue, DL);
+
+    // Setup the position where the next loop connects to this loop.
+    Enter = newLoop->getBody();
+    Continue = newLoop->getLatch();
+    OutroInsertBefore = newLoop->getLatch();
+
+    Result.push_back(newLoop);
+  }
+
+  // Detach the original loops
+  for (int i = 0; i < NumLoops - 1; i++) {
+    BasicBlock *body = Loops[i]->getBody();
+    BasicBlock *region =
+        cast<UncondBrInst>(body->getTerminator())->getSuccessor(0);
+    BasicBlock *cont = Loops[i]->getLatch()->getUniquePredecessor();
----------------
NouTimbaler wrote:

I added this because the transformation leaves an unreachable cycle of blocks of the old loops and removing them would make the resulting emitted code more readable. But, like you pointed out, I can't control the structure of the loop body.
Should I leave the task of removing them to opt?

https://github.com/llvm/llvm-project/pull/186381