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

[bondhugula]

Summary

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.

[bondhugula]

Update commit summary.

[springerm]

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

42

ivs

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

515–520

I think you can use mlir::getConstantIntValue(symbol) here.

[bondhugula]

Address review comments.

[bondhugula]

Thanks for the review.

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

515–520

Thanks, done.

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

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

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

 #define GET_OP_CLASSES
 #include "mlir/Dialect/Affine/IR/AffineOps.h.inc"
 
 namespace mlir {
 /// Returns true if the provided value is the induction variable of an
 /// 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,
                              SmallVectorImpl<Value> *ivs);
 
 /// Extracts the induction variables from a list of either AffineForOp or
 /// AffineParallelOp and places them in the output argument `ivs`.
 void extractInductionVars(ArrayRef<Operation *> affineOps,

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

 /// innermost.
 static LogicalResult getOpIndexSet(Operation *op,
                                    FlatAffineValueConstraints *indexSet) {
   SmallVector<Operation *, 4> ops;
   getEnclosingAffineOps(*op, &ops);
   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
 getNumCommonLoops(const FlatAffineValueConstraints &srcDomain,
                   const FlatAffineValueConstraints &dstDomain,
                   SmallVectorImpl<AffineForOp> *commonLoops = nullptr) {
   // Find the number of common loops shared by src and dst accesses.
   unsigned minNumLoops =

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

LogicalResult MemRefRegion::compute(Operation *op, unsigned loopDepth,
 
   MemRefAccess access(op);
   memref = access.memref;
   write = access.isStore();
 
   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();
   }
 
   // Build the constraints for this region.
   AffineValueMap accessValueMap;
   access.getAccessMap(&accessValueMap);
   AffineMap accessMap = accessValueMap.getAffineMap();
 
LogicalResult MemRefRegion::compute(Operation *op, unsigned loopDepth,
   // and symbols resp. of cst. This will change below once cst is
   // fully constructed out.
   cst.reset(numDims, numSymbols, 0, operands);
 
   // Add equality constraints.
   // 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)) {
+      Value symbol = operand;
+      if (Operation *op = symbol.getDefiningOp()) {
+        if (auto constOp = dyn_cast<arith::ConstantIndexOp>(op))
           cst.addBound(FlatAffineValueConstraints::EQ, symbol, constOp.value());
-        }
-      }
+    } else {
+      LLVM_DEBUG(llvm::dbgs() << "unknown affine dimensional value");
+      return failure();
     }
   }
 
   // Add lower/upper bounds on loop IVs using bounds from 'sliceState'.
   if (sliceState != nullptr) {
     // Add dim and symbol slice operands.
     for (auto operand : sliceState->lbOperands[0]) {
       cst.addInductionVarOrTerminalSymbol(operand);
LogicalResult MemRefRegion::compute(Operation *op, unsigned loopDepth,
 
   // Set all variables appearing after the first 'rank' variables as
   // symbolic variables - so that the ones corresponding to the memref
   // dimensions are the dimensional variables for the memref region.
   cst.setDimSymbolSeparation(cst.getNumDimAndSymbolVars() - rank);
 
   // 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);
     }
   }
 
   // Project out any local variables (these would have been added for any
   // mod/divs).
   cst.projectOut(cst.getNumDimAndSymbolVars(), cst.getNumLocalVars());
 
bool MemRefAccess::operator==(const MemRefAccess &rhs) const {
   AffineValueMap diff, thisMap, rhsMap;
   getAccessMap(&thisMap);
   rhs.getAccessMap(&rhsMap);
   AffineValueMap::difference(thisMap, rhsMap, &diff);
   return llvm::all_of(diff.getAffineMap().getResults(),
                       [](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',
 /// where each lists loops from outer-most to inner-most in loop nest.
 unsigned mlir::getNumCommonSurroundingLoops(Operation &a, Operation &b) {
   SmallVector<Value, 4> loopsA, loopsB;
   getAffineIVs(a, loopsA);
   getAffineIVs(b, loopsB);

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

 }
 
 /// Returns true if the provided value is the induction variable of a
 /// AffineForOp.
 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())
     return AffineForOp();
   auto *containingInst = ivArg.getOwner()->getParent()->getParentOp();
   if (auto forOp = dyn_cast<AffineForOp>(containingInst))
     // Check to make sure `val` is the induction variable, not an iter_arg.
     return forOp.getInductionVar() == val ? forOp : AffineForOp();
   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,
                                    SmallVectorImpl<Value> *ivs) {
   ivs->reserve(forInsts.size());
   for (auto forInst : forInsts)
     ivs->push_back(forInst.getInductionVar());
 }

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

 // CHECK-LABEL: func @zero_d_memref
 func.func @zero_d_memref() {
   %Z = memref.alloc() : memref<f32>
   affine.for %i = 0 to 100 {
     affine.load %Z[] : memref<f32>
   }
   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

 };
 
 } // namespace
 
 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.
   unsigned innermostLoopIdx = depthToLoops.size() - 1;
   if (depthToLoops[0].size() != 1 || depthToLoops[innermostLoopIdx].size() != 1)
     return;
 
   auto loopNest = depthToLoops[0][0];