[llvm-branch-commits] [mlir] 5a72ea0 - [MLIR] Add method mlir::isPerfectNestRoot

[Uday Bondhugula via llvm-branch-commits]

Author: Uday Bondhugula
Date: 2021-09-23T08:20:06+05:30
New Revision: 5a72ea079cde160b044ec7241928bb9520a5ae75

URL: https://github.com/llvm/llvm-project/commit/5a72ea079cde160b044ec7241928bb9520a5ae75
DIFF: https://github.com/llvm/llvm-project/commit/5a72ea079cde160b044ec7241928bb9520a5ae75.diff

LOG: [MLIR] Add method mlir::isPerfectNestRoot

Add method to check if an affine.for op is the root of a nest that is
nested perfectly all the way.

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Transforms/LoopUtils.h b/mlir/include/mlir/Transforms/LoopUtils.h
index 8c851215e3c08..fe89a86c4d2e8 100644
--- a/mlir/include/mlir/Transforms/LoopUtils.h
+++ b/mlir/include/mlir/Transforms/LoopUtils.h
@@ -57,6 +57,9 @@ LogicalResult loopUnrollUpToFactor(AffineForOp forOp, uint64_t unrollFactor);
 /// in it from outermost to innermost.
 bool LLVM_ATTRIBUTE_UNUSED isPerfectlyNested(ArrayRef<AffineForOp> loops);
 
+/// Returns true if the nest rooted at `loop` is a perfectly nested loop nest.
+bool isPerfectNestRoot(AffineForOp loop);
+
 /// Get perfectly nested sequence of loops starting at root of loop nest
 /// (the first op being another AffineFor, and the second op - a terminator).
 /// A loop is perfectly nested iff: the first op in the loop's body is another

diff  --git a/mlir/lib/Transforms/Utils/LoopUtils.cpp b/mlir/lib/Transforms/Utils/LoopUtils.cpp
index 4427001be2204..41a6aa692ee7c 100644
--- a/mlir/lib/Transforms/Utils/LoopUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopUtils.cpp
@@ -1543,6 +1543,24 @@ mlir::isPerfectlyNested(ArrayRef<AffineForOp> loops) {
   return true;
 }
 
+/// Returns true if no other affine for op's are nested within `forOp`.
+static bool isInnermostAffineForOp(AffineForOp forOp) {
+  // Only for the innermost affine.for op's.
+  bool isInnermost = true;
+  forOp.walk([&](AffineForOp thisForOp) {
+    isInnermost = (thisForOp == forOp);
+    return WalkResult::interrupt();
+  });
+  return isInnermost;
+}
+
+/// Returns true if the nest rooted at `loop` is a perfectly nested loop nest.
+bool mlir::isPerfectNestRoot(AffineForOp loop) {
+  SmallVector<AffineForOp, 4> nest;
+  getPerfectlyNestedLoops(nest, loop);
+  return isInnermostAffineForOp(nest.back());
+}
+
 // input[i] should move from position i -> permMap[i]. Returns the position in
 // `input` that becomes the new outermost loop.
 unsigned mlir::permuteLoops(MutableArrayRef<AffineForOp> input,