[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

diff --git a/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h b/mlir/include/mlir/Dialect/Affine/Analysis/Utils.h
index f6fde03..0781748 100644 (file)
--- 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 2ed8392..5c0917d 100644 (file)
--- 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 64ab959..05d55f8 100644 (file)
--- 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();
   }