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

[Uday Bondhugula via Phabricator via llvm-commits]

This revision was automatically updated to reflect the committed changes.
Closed by commit rG4e4ea2cde44d: [MLIR] Add missing asserts in interchangeLoops util, doc comment update (authored by bondhugula).

Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D77003/new/

https://reviews.llvm.org/D77003

Files:
  mlir/include/mlir/Transforms/LoopUtils.h
  mlir/lib/Transforms/Utils/LoopUtils.cpp


Index: mlir/lib/Transforms/Utils/LoopUtils.cpp
===================================================================
--- mlir/lib/Transforms/Utils/LoopUtils.cpp
+++ mlir/lib/Transforms/Utils/LoopUtils.cpp
@@ -700,19 +700,48 @@
   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());
@@ -770,7 +799,7 @@
   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];
 }
 
Index: mlir/include/mlir/Transforms/LoopUtils.h
===================================================================
--- mlir/include/mlir/Transforms/LoopUtils.h
+++ mlir/include/mlir/Transforms/LoopUtils.h
@@ -108,11 +108,15 @@
 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


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D77003.253442.patch
Type: text/x-patch
Size: 4402 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200329/d8bb0919/attachment-0001.bin>