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

Author: qaco

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 {