Apply the permutation map on each affine nest

[qaco]

Summary

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).

[bondhugula]

Thanks for fixing this. One other comment to make the walk complete.

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

49–50

We could actually just use:

getFunction().walk<AffineForOp>([&](AffineForOp forOp) {
  forOps.push_back(forOp);
});
``

[qaco]

Use getFunction.walk() instead of a for loop on each operation in the function.

[qaco]

In D112947#3123697, @bondhugula wrote:

Thanks for fixing this. One other comment to make the walk complete.

Thanks for your comment. Update done.

[qaco]

@bondhugula I can not commit the patch myself (I have not commit access).

[bondhugula]

In D112947#3124225, @qaco wrote:

@bondhugula I can not commit the patch myself (I have not commit access).

Sorry about the delay - I'm committing this.

[bondhugula]

@qaco The patch wasn't clang-formatted and had white space issues as well. I can fix these before committing. Can you provide your name/email information that you want on the commit? arc patch isn't able to automatically get it.

[qaco]

In D112947#3171215, @bondhugula wrote:

@qaco The patch wasn't clang-formatted and had white space issues as well. I can fix these before committing. Can you provide your name/email information that you want on the commit? arc patch isn't able to automatically get it.

Thank you very much.

My name : Hugo Pompougnac
My email : [e-mail]

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

private:
   /// transformed nest (with i going from outermost to innermost).
   ListOption<unsigned> permList{*this, "permutation-map",
                                 llvm::cl::desc("Specify the loop permutation"),
                                 llvm::cl::OneOrMore, llvm::cl::CommaSeparated};
 };
 
 } // end anonymous namespace
 
 void TestLoopPermutation::runOnFunction() {
-  // Get the first maximal perfect nest.
+
+  SmallVector<unsigned, 4> permMap(permList.begin(), permList.end());
+
+  SmallVector<AffineForOp, 2> forOps;
+  getFunction().walk([&](AffineForOp forOp) { forOps.push_back(forOp); });
+
+  for (auto forOp : forOps) {
     SmallVector<AffineForOp, 6> nest;
-  for (auto &op : getFunction().front()) {
-    if (auto forOp = dyn_cast<AffineForOp>(op)) {
-      getPerfectlyNestedLoops(nest, forOp);
-      break;
+    // Get the maximal perfect nest.
+    // Permute if the nest's size is consistent with the specified
+    // permutation.
+    if (nest.size() >= 2 && nest.size() == permMap.size()) {
+      permuteLoops(nest, permMap);
     }
   }
-
-  // Nothing to do.
-  if (nest.size() < 2)
-    return;
-
-  SmallVector<unsigned, 4> permMap(permList.begin(), permList.end());
-  permuteLoops(nest, permMap);
 }
 
 namespace mlir {
 void registerTestLoopPermutationPass() {
   PassRegistration<TestLoopPermutation>();
 }
 } // namespace mlir