Index: mlir/test/lib/Dialect/Affine/TestLoopPermutation.cpp
===================================================================
--- mlir/test/lib/Dialect/Affine/TestLoopPermutation.cpp
+++ mlir/test/lib/Dialect/Affine/TestLoopPermutation.cpp
@@ -47,21 +47,27 @@
 } // end anonymous namespace
 
 void TestLoopPermutation::runOnFunction() {
-  // Get the first maximal perfect nest.
-  SmallVector<AffineForOp, 6> nest;
+
+  SmallVector<unsigned, 4> permMap(permList.begin(), permList.end());
+  
+  SmallVector<AffineForOp, 2> forOps;
   for (auto &op : getFunction().front()) {
     if (auto forOp = dyn_cast<AffineForOp>(op)) {
-      getPerfectlyNestedLoops(nest, forOp);
-      break;
+      forOps.push_back(forOp);
     }
   }
 
-  // Nothing to do.
-  if (nest.size() < 2)
-    return;
-
-  SmallVector<unsigned, 4> permMap(permList.begin(), permList.end());
-  permuteLoops(nest, permMap);
+  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 {