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

[Uday Bondhugula via Phabricator via llvm-commits]

bondhugula created this revision.
bondhugula added reviewers: andydavis1, dcaballe.
Herald added subscribers: llvm-commits, Joonsoo, liufengdb, aartbik, lucyrfox, mgester, arpith-jacob, nicolasvasilache, antiagainst, shauheen, burmako, jpienaar, rriddle, mehdi_amini, mgrang.
Herald added a project: LLVM.
bondhugula added a child revision: D77004: [MLIR] Fix permuteLoops utility.

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 <uday at polymagelabs.com>


Repository:
  rG LLVM Github Monorepo

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,40 @@
   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) {
+// 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;
+
+  // Check for perfect nesting.
+  auto innerLoop = input.back();
+  for (auto loop : llvm::reverse(input.drop_back())) {
+    assert(dyn_cast<AffineForOp>(*loop.getBody()->begin()) == innerLoop &&
+           "input not perfectly nested");
+    innerLoop = loop;
+  }
+
   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 +791,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];
 }
 
@@ -779,7 +800,7 @@
 void mlir::sinkLoop(AffineForOp forOp, unsigned loopDepth) {
   for (unsigned i = 0; i < loopDepth; ++i) {
     AffineForOp nextForOp = cast<AffineForOp>(forOp.getBody()->front());
-    interchangeLoops(forOp, nextForOp);
+    permuteLoops(forOp, nextForOp);
   }
 }
 
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.253390.patch
Type: text/x-patch
Size: 4400 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200329/11f3803c/attachment-0001.bin>