[Mlir-commits] [mlir] 5d49511 - Apply the permutation map on each affine nest

[Uday Bondhugula]

Author: Hugo Pompougnac
Date: 2021-12-04T17:48:34+05:30
New Revision: 5d49511b3058cac17eb80d243d3e2ad5cb0cddb9

URL: https://github.com/llvm/llvm-project/commit/5d49511b3058cac17eb80d243d3e2ad5cb0cddb9
DIFF: https://github.com/llvm/llvm-project/commit/5d49511b3058cac17eb80d243d3e2ad5cb0cddb9.diff

LOG: Apply the permutation map on each affine nest

When using -test-loop-permutation="permutation-map=...", applies the
permutation map on each affine nest in the function (and not only the
first one). If the size of the permutation map and the size of a nest
are not consistent, do nothing on this particular nest (instead of
making MLIR crash).

Differential Revision: https://reviews.llvm.org/D112947

Added: 
    

Modified: 
    mlir/test/lib/Dialect/Affine/TestLoopPermutation.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/test/lib/Dialect/Affine/TestLoopPermutation.cpp b/mlir/test/lib/Dialect/Affine/TestLoopPermutation.cpp
index af9d972efa639..5205b25f881ab 100644
--- a/mlir/test/lib/Dialect/Affine/TestLoopPermutation.cpp
+++ b/mlir/test/lib/Dialect/Affine/TestLoopPermutation.cpp
@@ -47,21 +47,22 @@ struct TestLoopPermutation
 } // end anonymous namespace
 
 void TestLoopPermutation::runOnFunction() {
-  // Get the first maximal perfect nest.
-  SmallVector<AffineForOp, 6> nest;
-  for (auto &op : getFunction().front()) {
-    if (auto forOp = dyn_cast<AffineForOp>(op)) {
-      getPerfectlyNestedLoops(nest, forOp);
-      break;
-    }
-  }
-
-  // Nothing to do.
-  if (nest.size() < 2)
-    return;
 
   SmallVector<unsigned, 4> permMap(permList.begin(), permList.end());
-  permuteLoops(nest, permMap);
+
+  SmallVector<AffineForOp, 2> forOps;
+  getFunction().walk([&](AffineForOp forOp) { forOps.push_back(forOp); });
+
+  for (auto forOp : forOps) {
+    SmallVector<AffineForOp, 6> nest;
+    // Get the maximal perfect nest.
+    getPerfectlyNestedLoops(nest, forOp);
+    // Permute if the nest's size is consistent with the specified
+    // permutation.
+    if (nest.size() >= 2 && nest.size() == permMap.size()) {
+      permuteLoops(nest, permMap);
+    }
+  }
 }
 
 namespace mlir {