[Mlir-commits] [mlir] 26fedf9 - [MLIR] Generalize/complete getEnclosingAffineForIfOps utility

[Uday Bondhugula]

Author: Uday Bondhugula
Date: 2022-08-25T02:56:31+05:30
New Revision: 26fedf92c3122ff51fb27f434a8c5b5a47da8f94

URL: https://github.com/llvm/llvm-project/commit/26fedf92c3122ff51fb27f434a8c5b5a47da8f94
DIFF: https://github.com/llvm/llvm-project/commit/26fedf92c3122ff51fb27f434a8c5b5a47da8f94.diff

LOG: [MLIR] Generalize/complete getEnclosingAffineForIfOps utility

Rename/generalize getEnclosingAffineForIfOps -> getEnclosingAffineOps.
The utility was originally written only for affine.for ops and then
extended for affine.if as well. It wasn't however updated when
affine.parallel was introduced -- in most cases, analysis has been used
in the presence of affine.for/if but not post parallelization. Extend
utility to also support affine.parallel ops; this allows future changes
to enable affine analysis and opts in the presence of affine.parallel
ops. Fix related stale comments.

This is NFC for all use cases in tree.

Change an associated assert to a utility failure.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D132326

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Affine/Analysis/Utils.h
    mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
    mlir/lib/Dialect/Affine/Analysis/Utils.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h b/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h
index f6fde0383e448..0781748f2f8b2 100644
--- a/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h
+++ b/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h
@@ -38,11 +38,14 @@ class Value;
 //  TODO: handle 'affine.if' ops.
 void getLoopIVs(Operation &op, SmallVectorImpl<AffineForOp> *loops);
 
-/// Populates 'ops' with IVs of the loops surrounding `op`, along with
-/// `affine.if` operations interleaved between these loops, ordered from the
-/// outermost `affine.for` or `affine.if` operation to the innermost one.
-void getEnclosingAffineForAndIfOps(Operation &op,
-                                   SmallVectorImpl<Operation *> *ops);
+/// Populates 'ops' with affine operations enclosing `op` ordered from outermost
+/// to innermost. affine.for, affine.if, or affine.parallel ops comprise such
+/// surrounding affine ops.
+/// TODO: Change this to return a list of enclosing ops up until the op that
+/// starts an `AffineScope`. In such a case, `ops` is guaranteed by design to
+/// have a successive chain of affine parent ops, and this is primarily what is
+/// needed for most analyses.
+void getEnclosingAffineOps(Operation &op, SmallVectorImpl<Operation *> *ops);
 
 /// Returns the nesting depth of this operation, i.e., the number of loops
 /// surrounding this operation.

diff  --git a/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp b/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
index 2ed83929ce8c7..5c0917d865d87 100644
--- a/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
@@ -241,10 +241,12 @@ LogicalResult mlir::getIndexSet(MutableArrayRef<Operation *> ops,
                                 FlatAffineValueConstraints *domain) {
   SmallVector<Value, 4> indices;
   SmallVector<AffineForOp, 8> forOps;
-
   for (Operation *op : ops) {
-    assert((isa<AffineForOp, AffineIfOp>(op)) &&
-           "ops should have either AffineForOp or AffineIfOp");
+    if (!isa<AffineForOp, AffineIfOp>(op)) {
+      // TODO: Support affine.parallel ops.
+      LLVM_DEBUG(llvm::dbgs() << "getIndexSet only handles affine.for/if ops");
+      return failure();
+    }
     if (AffineForOp forOp = dyn_cast<AffineForOp>(op))
       forOps.push_back(forOp);
   }
@@ -271,7 +273,7 @@ LogicalResult mlir::getIndexSet(MutableArrayRef<Operation *> ops,
 static LogicalResult getOpIndexSet(Operation *op,
                                    FlatAffineValueConstraints *indexSet) {
   SmallVector<Operation *, 4> ops;
-  getEnclosingAffineForAndIfOps(*op, &ops);
+  getEnclosingAffineOps(*op, &ops);
   return getIndexSet(ops, indexSet);
 }
 

diff  --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index 64ab959219064..05d55f8514bb6 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -45,18 +45,15 @@ void mlir::getLoopIVs(Operation &op, SmallVectorImpl<AffineForOp> *loops) {
   std::reverse(loops->begin(), loops->end());
 }
 
-/// Populates 'ops' with IVs of the loops surrounding `op`, along with
-/// `affine.if` operations interleaved between these loops, ordered from the
-/// outermost `affine.for` operation to the innermost one.
-void mlir::getEnclosingAffineForAndIfOps(Operation &op,
-                                         SmallVectorImpl<Operation *> *ops) {
+void mlir::getEnclosingAffineOps(Operation &op,
+                                 SmallVectorImpl<Operation *> *ops) {
   ops->clear();
   Operation *currOp = op.getParentOp();
 
-  // Traverse up the hierarchy collecting all `affine.for` and `affine.if`
-  // operations.
+  // Traverse up the hierarchy collecting all `affine.for`, `affine.if`, and
+  // affine.parallel operations.
   while (currOp) {
-    if (isa<AffineIfOp, AffineForOp>(currOp))
+    if (isa<AffineIfOp, AffineForOp, AffineParallelOp>(currOp))
       ops->push_back(currOp);
     currOp = currOp->getParentOp();
   }