[MLIR] Generalize/complete getEnclosingAffineForIfOps utility

[bondhugula]

Summary

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.

[bondhugula]

Update commit summary.

[Groverkss]

mlir/lib/Dialect/Affine/Analysis/Utils.cpp

52–54

This comment should be updated since it also collects affine.parallel ops.

[bondhugula]

Address review comment.

mlir/include/mlir/Dialect/Affine/Analysis/Utils.h

 class Operation;
 class Value;
 
 /// Populates 'loops' with IVs of the loops surrounding 'op' ordered from
 /// the outermost 'affine.for' operation to the innermost one.
 //  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.
 unsigned getNestingDepth(Operation *op);
 
 /// Returns whether a loop is a parallel loop and contains a reduction loop.
 bool isLoopParallelAndContainsReduction(AffineForOp forOp);
 

mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp

 // TODO: Handle non-unit steps through local variables or stride information in
 // FlatAffineValueConstraints. (For eg., by using iv - lb % step = 0 and/or by
 // introducing a method in FlatAffineValueConstraints
 // setExprStride(ArrayRef<int64_t> expr, int64_t stride)
 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);
   }
   extractForInductionVars(forOps, &indices);
   // Reset while associated Values in 'indices' to the domain.
   domain->reset(forOps.size(), /*numSymbols=*/0, /*numLocals=*/0, indices);
   for (Operation *op : ops) {
     // Add constraints from forOp's bounds.

 /// Computes the iteration domain for 'op' and populates 'indexSet', which
 /// encapsulates the constraints involving loops surrounding 'op' and
 /// potentially involving any Function symbols. The dimensional variables in
 /// 'indexSet' correspond to the loops surrounding 'op' from outermost to
 /// innermost.
 static LogicalResult getOpIndexSet(Operation *op,
                                    FlatAffineValueConstraints *indexSet) {
   SmallVector<Operation *, 4> ops;
-  getEnclosingAffineForAndIfOps(*op, &ops);
+  getEnclosingAffineOps(*op, &ops);
   return getIndexSet(ops, indexSet);
 }
 
 // Returns the number of outer loop common to 'src/dstDomain'.
 // Loops common to 'src/dst' domains are added to 'commonLoops' if non-null.
 static unsigned
 getNumCommonLoops(const FlatAffineValueConstraints &srcDomain,
                   const FlatAffineValueConstraints &dstDomain,

mlir/lib/Dialect/Affine/Analysis/Utils.cpp

void mlir::getLoopIVs(Operation &op, SmallVectorImpl<AffineForOp> *loops) {
   while (currOp) {
     if (AffineForOp currAffineForOp = dyn_cast<AffineForOp>(currOp))
       loops->push_back(currAffineForOp);
     currOp = currOp->getParentOp();
   }
   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,
   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();
   }
   std::reverse(ops->begin(), ops->end());
 }
 
 // Populates 'cst' with FlatAffineValueConstraints which represent original
 // domain of the loop bounds that define 'ivs'.