[MLIR] Add missing asserts in interchangeLoops util, doc comment update

[bondhugula]

Summary

Add missing assert checks for input to mlir::interchangeLoops utility.
Rename interchangeLoops -> permuteLoops; update doc comments to clarify
inputs / return val. Other than the assert checks, this is NFC.

Signed-off-by: Uday Bondhugula <[e-mail]>

[mehdi_amini]

mlir/lib/Transforms/Utils/LoopUtils.cpp

724

Should we also check that the body of these loops contains only a single operation (other than the terminator)?

[bondhugula]

perfect nest check

mlir/lib/Transforms/Utils/LoopUtils.cpp

724

Yes, you are right - we should be. Fixed.

mlir/include/mlir/Transforms/LoopUtils.h

 void interchangeLoops(AffineForOp forOpA, AffineForOp forOpB);
 
 /// Checks if the loop interchange permutation 'loopPermMap', of the perfectly
 /// nested sequence of loops in 'loops', would violate dependences (loop 'i' in
 /// 'loops' is mapped to location 'j = 'loopPermMap[i]' in the interchange).
 bool isValidLoopInterchangePermutation(ArrayRef<AffineForOp> loops,
                                        ArrayRef<unsigned> loopPermMap);
 
-/// Performs a sequence of loop interchanges on perfectly nested 'loops', as
-/// specified by permutation 'loopPermMap' (loop 'i' in 'loops' is mapped to
-/// location 'j = 'loopPermMap[i]' after the loop interchange).
-unsigned interchangeLoops(ArrayRef<AffineForOp> loops,
-                          ArrayRef<unsigned> loopPermMap);
+/// Performs a loop permutation on a perfectly nested loop nest `inputNest`
+/// (where the contained loops appear from outer to inner) as specified by the
+/// permutation `permMap`: loop 'i' in `inputNest` is mapped to location
+/// 'loopPermMap[i]', where positions 0, 1, ... are from the outermost position
+/// to inner. Returns the position in `inputNest` of the AffineForOp that
+/// becomes the new outermost loop of this nest. This method always succeeds,
+/// asserts out on invalid input / specifications.
+unsigned permuteLoops(ArrayRef<AffineForOp> inputNest,
+                      ArrayRef<unsigned> permMap);
 
 // Sinks all sequential loops to the innermost levels (while preserving
 // relative order among them) and moves all parallel loops to the
 // outermost (while again preserving relative order among them).
 // Returns AffineForOp of the root of the new loop nest after loop interchanges.
 AffineForOp sinkSequentialLoops(AffineForOp forOp);
 
 /// Sinks 'forOp' by 'loopDepth' levels by performing a series of loop

mlir/lib/Transforms/Utils/LoopUtils.cpp

bool mlir::isValidLoopInterchangePermutation(ArrayRef<AffineForOp> loops,
   // rooted at 'loops[0]', at loop depths in range [1, maxLoopDepth].
   assert(loopPermMap.size() == loops.size());
   unsigned maxLoopDepth = loops.size();
   std::vector<SmallVector<DependenceComponent, 2>> depCompsVec;
   getDependenceComponents(loops[0], maxLoopDepth, &depCompsVec);
   return checkLoopInterchangeDependences(depCompsVec, loops, loopPermMap);
 }
 
-/// Performs a sequence of loop interchanges of loops in perfectly nested
-/// sequence of loops in 'loops', as specified by permutation in 'loopPermMap'.
-unsigned mlir::interchangeLoops(ArrayRef<AffineForOp> loops,
-                                ArrayRef<unsigned> loopPermMap) {
+/// Return true if `loops` is a perfect nest.
+static bool isPerfectlyNested(ArrayRef<AffineForOp> loops) {
+  auto outerLoop = loops.front();
+  for (auto loop : loops.drop_front()) {
+    auto parentForOp = dyn_cast<AffineForOp>(loop.getParentOp());
+    // parentForOp's body should be just this loop and the terminator.
+    if (parentForOp != outerLoop ||
+        parentForOp.getBody()->getOperations().size() != 2)
+      return false;
+    outerLoop = loop;
+  }
+  return true;
+}
+
+// input[i] should move from position i -> permMap[i]. Returns the position in
+// `input` that becomes the new outermost loop.
+unsigned mlir::permuteLoops(ArrayRef<AffineForOp> input,
+                            ArrayRef<unsigned> permMap) {
+  assert(input.size() == permMap.size() && "invalid permutation map size");
+  // Check whether the permutation spec is valid. This is a small vector - we'll
+  // just sort and check if it's iota.
+  SmallVector<unsigned, 4> checkPermMap(permMap.begin(), permMap.end());
+  llvm::sort(checkPermMap);
+  if (llvm::any_of(llvm::enumerate(checkPermMap),
+                   [](const auto &en) { return en.value() != en.index(); }))
+    assert(false && "invalid permutation map");
+
+  // Nothing to do.
+  if (input.size() < 2)
+    return 0;
+
+  assert(isPerfectlyNested(input) && "input not perfectly nested");
+
   Optional<unsigned> loopNestRootIndex;
-  for (int i = loops.size() - 1; i >= 0; --i) {
-    int permIndex = static_cast<int>(loopPermMap[i]);
+  for (int i = input.size() - 1; i >= 0; --i) {
+    int permIndex = static_cast<int>(permMap[i]);
     // Store the index of the for loop which will be the new loop nest root.
     if (permIndex == 0)
       loopNestRootIndex = i;
     if (permIndex > i) {
       // Sink loop 'i' by 'permIndex - i' levels deeper into the loop nest.
-      sinkLoop(loops[i], permIndex - i);
+      sinkLoop(input[i], permIndex - i);
     }
   }
   assert(loopNestRootIndex.hasValue());
   return loopNestRootIndex.getValue();
 }
 
 // Sinks all sequential loops to the innermost levels (while preserving
 // relative order among them) and moves all parallel loops to the
if (isParallelLoop[i]) {
       loopPermMap[i] = nextSequentialLoop++;
     }
   }
 
   // Check if permutation 'loopPermMap' would violate dependences.
   if (!checkLoopInterchangeDependences(depCompsVec, loops, loopPermMap))
     return forOp;
   // Perform loop interchange according to permutation 'loopPermMap'.
-  unsigned loopNestRootIndex = interchangeLoops(loops, loopPermMap);
+  unsigned loopNestRootIndex = permuteLoops(loops, loopPermMap);
   return loops[loopNestRootIndex];
 }
 
 /// Performs a series of loop interchanges to sink 'forOp' 'loopDepth' levels
 /// deeper in the loop nest.
 void mlir::sinkLoop(AffineForOp forOp, unsigned loopDepth) {
   for (unsigned i = 0; i < loopDepth; ++i) {
     AffineForOp nextForOp = cast<AffineForOp>(forOp.getBody()->front());