[MLIR] Update/fix memref region computation for affine.parallel ops

When the affine.parallel op was introduced, affine utilities weren't
extended to handle it. Extending these is straightforward and natural
given that addAffineParallelOpDomain has also been added.
Update/complete memref region compute to account for affine.parallel
ops. Handle failure cleanly.

Add and expose utilities missing for affine.parallel to be consistent
with affine.for.

All of these allow various affine passes to work with a combination of
affine.parallel and affine.for ops.

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

Author: bondhugula

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

@@ -39,6 +39,10 @@ class Value;
 /// from the outermost 'affine.for' operation to the innermost one.
 void getAffineForIVs(Operation &op, SmallVectorImpl<AffineForOp> *loops);
 
+/// Populates 'ivs' with IVs of the surrounding affine.for and affine.parallel
+/// ops ordered from the outermost one to the innermost.
+void getAffineIVs(Operation &op, SmallVectorImpl<Value> &ivs);
+
 /// 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.

mlir/include/mlir/Dialect/Affine/IR/AffineOps.h

@@ -441,10 +441,21 @@ namespace mlir {
 /// AffineForOp.
 bool isAffineForInductionVar(Value val);
 
+/// Returns true if `val` is the induction variable of an AffineParallelOp.
+bool isAffineParallelInductionVar(Value val);
+
+/// Returns true if the provided value is the induction variable of an
+/// AffineForOp or AffineParallelOp.
+bool isAffineInductionVar(Value val);
+
 /// Returns the loop parent of an induction variable. If the provided value is
 /// not an induction variable, then return nullptr.
 AffineForOp getForInductionVarOwner(Value val);
 
+/// Returns true if the provided value is among the induction variables of an
+/// AffineParallelOp.
+AffineParallelOp getAffineParallelInductionVarOwner(Value val);
+
 /// Extracts the induction variables from a list of AffineForOps and places them
 /// in the output argument `ivs`.
 void extractForInductionVars(ArrayRef<AffineForOp> forInsts,

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

@@ -287,14 +287,6 @@ static LogicalResult getOpIndexSet(Operation *op,
   return getIndexSet(ops, indexSet);
 }
 
-/// Returns true if `val` is an induction of an affine.parallel op.
-static bool isAffineParallelInductionVar(Value val) {
-  auto ivArg = val.dyn_cast<BlockArgument>();
-  if (!ivArg)
-    return false;
-  return isa<AffineParallelOp>(ivArg.getOwner()->getParentOp());
-}
-
 // 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

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

@@ -454,19 +454,17 @@ LogicalResult MemRefRegion::compute(Operation *op, unsigned loopDepth,
   unsigned rank = access.getRank();
 
   LLVM_DEBUG(llvm::dbgs() << "MemRefRegion::compute: " << *op
-                          << "depth: " << loopDepth << "\n";);
+                          << "\ndepth: " << loopDepth << "\n";);
 
   // 0-d memrefs.
   if (rank == 0) {
-    SmallVector<AffineForOp, 4> ivs;
-    getAffineForIVs(*op, &ivs);
+    SmallVector<Value, 4> ivs;
+    getAffineIVs(*op, ivs);
     assert(loopDepth <= ivs.size() && "invalid 'loopDepth'");
     // The first 'loopDepth' IVs are symbols for this region.
     ivs.resize(loopDepth);
-    SmallVector<Value, 4> regionSymbols;
-    extractForInductionVars(ivs, &regionSymbols);
     // A 0-d memref has a 0-d region.
-    cst.reset(rank, loopDepth, /*numLocals=*/0, regionSymbols);
+    cst.reset(rank, loopDepth, /*numLocals=*/0, ivs);
     return success();
   }
 
@@ -503,23 +501,24 @@ LogicalResult MemRefRegion::compute(Operation *op, unsigned loopDepth,
   // Add inequalities for loop lower/upper bounds.
   for (unsigned i = 0; i < numDims + numSymbols; ++i) {
     auto operand = operands[i];
-    if (auto loop = getForInductionVarOwner(operand)) {
+    if (auto affineFor = getForInductionVarOwner(operand)) {
       // Note that cst can now have more dimensions than accessMap if the
       // bounds expressions involve outer loops or other symbols.
       // TODO: rewrite this to use getInstIndexSet; this way
       // conditionals will be handled when the latter supports it.
-      if (failed(cst.addAffineForOpDomain(loop)))
+      if (failed(cst.addAffineForOpDomain(affineFor)))
         return failure();
-    } else {
-      // Has to be a valid symbol.
-      auto symbol = operand;
-      assert(isValidSymbol(symbol));
+    } else if (auto parallelOp = getAffineParallelInductionVarOwner(operand)) {
+      if (failed(cst.addAffineParallelOpDomain(parallelOp)))
+        return failure();
+    } else if (isValidSymbol(operand)) {
       // Check if the symbol is a constant.
-      if (auto *op = symbol.getDefiningOp()) {
-        if (auto constOp = dyn_cast<arith::ConstantIndexOp>(op)) {
-          cst.addBound(FlatAffineValueConstraints::EQ, symbol, constOp.value());
-        }
-      }
+      Value symbol = operand;
+      if (auto constVal = getConstantIntValue(symbol))
+        cst.addBound(FlatAffineValueConstraints::EQ, symbol, constVal.value());
+    } else {
+      LLVM_DEBUG(llvm::dbgs() << "unknown affine dimensional value");
+      return failure();
     }
   }
 
@@ -552,16 +551,14 @@ LogicalResult MemRefRegion::compute(Operation *op, unsigned loopDepth,
 
   // Eliminate any loop IVs other than the outermost 'loopDepth' IVs, on which
   // this memref region is symbolic.
-  SmallVector<AffineForOp, 4> enclosingIVs;
-  getAffineForIVs(*op, &enclosingIVs);
+  SmallVector<Value, 4> enclosingIVs;
+  getAffineIVs(*op, enclosingIVs);
   assert(loopDepth <= enclosingIVs.size() && "invalid loop depth");
   enclosingIVs.resize(loopDepth);
   SmallVector<Value, 4> vars;
   cst.getValues(cst.getNumDimVars(), cst.getNumDimAndSymbolVars(), &vars);
-  for (auto var : vars) {
-    AffineForOp iv;
-    if ((iv = getForInductionVarOwner(var)) &&
-        !llvm::is_contained(enclosingIVs, iv)) {
+  for (Value var : vars) {
+    if ((isAffineInductionVar(var)) && !llvm::is_contained(enclosingIVs, var)) {
       cst.projectOut(var);
     }
   }
@@ -1264,21 +1261,19 @@ bool MemRefAccess::operator==(const MemRefAccess &rhs) const {
                       [](AffineExpr e) { return e == 0; });
 }
 
-/// Populates 'loops' with IVs of the surrounding affine.for and affine.parallel
-/// ops ordered from the outermost one to the innermost.
-static void getAffineIVs(Operation &op, SmallVectorImpl<Value> &loops) {
+void mlir::getAffineIVs(Operation &op, SmallVectorImpl<Value> &ivs) {
   auto *currOp = op.getParentOp();
   AffineForOp currAffineForOp;
-  // Traverse up the hierarchy collecting all 'affine.for' operation while
-  // skipping over 'affine.if' operations.
+  // Traverse up the hierarchy collecting all 'affine.for' and affine.parallel
+  // operation while skipping over 'affine.if' operations.
   while (currOp) {
     if (AffineForOp currAffineForOp = dyn_cast<AffineForOp>(currOp))
-      loops.push_back(currAffineForOp.getInductionVar());
+      ivs.push_back(currAffineForOp.getInductionVar());
     else if (auto parOp = dyn_cast<AffineParallelOp>(currOp))
-      llvm::append_range(loops, parOp.getIVs());
+      llvm::append_range(ivs, parOp.getIVs());
     currOp = currOp->getParentOp();
   }
-  std::reverse(loops.begin(), loops.end());
+  std::reverse(ivs.begin(), ivs.end());
 }
 
 /// Returns the number of surrounding loops common to 'loopsA' and 'loopsB',

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

@@ -2504,8 +2504,14 @@ bool mlir::isAffineForInductionVar(Value val) {
   return getForInductionVarOwner(val) != AffineForOp();
 }
 
-/// Returns the loop parent of an induction variable. If the provided value is
-/// not an induction variable, then return nullptr.
+bool mlir::isAffineParallelInductionVar(Value val) {
+  return getAffineParallelInductionVarOwner(val) != nullptr;
+}
+
+bool mlir::isAffineInductionVar(Value val) {
+  return isAffineForInductionVar(val) || isAffineParallelInductionVar(val);
+}
+
 AffineForOp mlir::getForInductionVarOwner(Value val) {
   auto ivArg = val.dyn_cast<BlockArgument>();
   if (!ivArg || !ivArg.getOwner())
@@ -2517,6 +2523,17 @@ AffineForOp mlir::getForInductionVarOwner(Value val) {
   return AffineForOp();
 }
 
+AffineParallelOp mlir::getAffineParallelInductionVarOwner(Value val) {
+  auto ivArg = val.dyn_cast<BlockArgument>();
+  if (!ivArg || !ivArg.getOwner())
+    return nullptr;
+  Operation *containingOp = ivArg.getOwner()->getParentOp();
+  auto parallelOp = dyn_cast<AffineParallelOp>(containingOp);
+  if (parallelOp && llvm::is_contained(parallelOp.getIVs(), val))
+    return parallelOp;
+  return nullptr;
+}
+
 /// Extracts the induction variables from a list of AffineForOps and returns
 /// them.
 void mlir::extractForInductionVars(ArrayRef<AffineForOp> forInsts,

mlir/test/Transforms/memref-bound-check.mlir

@@ -293,3 +293,14 @@ func.func @zero_d_memref() {
   }
   return
 }
+
+// CHECK-LABEL: func @affine_parallel
+func.func @affine_parallel(%M: memref<2048x2048xf64>) {
+  affine.parallel (%i) = (0) to (3000) {
+    affine.for %j = 0 to 2048 {
+      affine.load %M[%i, %j] : memref<2048x2048xf64>
+      // expected-error@above {{'affine.load' op memref out of upper bound access along dimension #1}}
+    }
+  }
+  return
+}

mlir/test/lib/Dialect/Affine/TestAffineDataCopy.cpp

@@ -60,7 +60,8 @@ void TestAffineDataCopy::runOnOperation() {
   // Gather all AffineForOps by loop depth.
   std::vector<SmallVector<AffineForOp, 2>> depthToLoops;
   gatherLoops(getOperation(), depthToLoops);
-  assert(!depthToLoops.empty() && "Loop nest not found");
+  if (depthToLoops.empty())
+    return;
 
   // Only support tests with a single loop nest and a single innermost loop
   // for now.