[Mlir-commits] [mlir] [MLIR][SCF] Support permutation-based parallel loop fusion (PR #203207)

[Dmitriy Smirnov]

================
@@ -773,27 +930,32 @@ static void fuseIfLegal(ParallelOp firstPloop, ParallelOp &secondPloop,
 
   if (!isFusionLegal(firstPloop, secondPloop, firstToSecondPloopIndices,
                      mayAlias, builder)) {
-    // If second parallel loop consists of two loops of same iteration space
-    // then exchange these loops and re-asses the possibility of fusion.
-    if (secondPloop.getNumLoops() == 2 &&
-        secondPloop.getUpperBound()[0] == secondPloop.getUpperBound()[1] &&
-        secondPloop.getLowerBound()[0] == secondPloop.getLowerBound()[1] &&
-        secondPloop.getStep()[0] == secondPloop.getStep()[1]) {
+    // If iteration space of the second parallel loop is a permutation of the
+    // first one then interchange iteration space of the second parallel loop
+    // and re-asses possibility of fusion.
+    for (auto &perms :
+         computeCandidateInterchangePermutations(firstPloop, secondPloop)) {
+      OpBuilder::InsertionGuard guard(builder);
+      auto newLoop = interchangeLoops(builder, secondPloop, perms);
       firstToSecondPloopIndices.clear();
       firstToSecondPloopIndices.map(block1->getArguments(),
-                                    llvm::reverse(block2->getArguments()));
-      if (!isFusionLegal(firstPloop, secondPloop, firstToSecondPloopIndices,
-                         mayAlias, builder))
-        return;
-      auto newLoop = interchangeLoops(builder, secondPloop);
+                                    newLoop->getBody()->getArguments());
+      if (!isFusionLegal(firstPloop, *newLoop, firstToSecondPloopIndices,
+                         mayAlias, builder)) {
+        newLoop->erase();
+        continue;
+      }
+
+      secondPloop.replaceAllUsesWith(newLoop->getResults());
       secondPloop->erase();
       secondPloop = *newLoop;
       block2 = secondPloop.getBody();
-    } else {
-      return;
+      goto fuseLabel;
----------------
d-smirnov wrote:

Extracted to separate func

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