[mlir][affine] Support affine.parallel in the index set analysis

[Lewuathe]

Summary

Support affine.parallel in the index set analysis. It allows us to do dependence analysis containing affine.parallel in addition to affine.for and affine.if. This change only supports the constant lower/upper bound in affine.parallel. Other complicated affine map bounds will be supported in further commits.

See https://github.com/llvm/llvm-project/issues/57327

[Groverkss]

Some general code-style comments to start the review.

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

246–248

Please update this message.

254–255

Any reason for not using else if?

271–273

nit: remove braces here.

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

654

Is there a reason for using auto here and not the actual type? I think generally the actual type is prefered.

655–660

nit: remove braces here.

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

2330–2331

nit: there is no need for a newline before the if.

2332–2334

nit: remove braces.

[Groverkss]

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

650–652

I don't think this ordering is valid for affine.parallel since it does not define any lexicographic order. For example:

affine.parallel (%i) = (0) to (10) {
  %1 = affine.load %m[9 - %i]
  affine.store %1, [%i]
}

There can be a dependence from %i = 9 to %i = 0, since no lexicographic order is defined.

[Lewuathe]

Apply format

[Lewuathe]

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

650–652

I see. We may need to skip this constraint in the case of affine.parallel in the middle.

[Lewuathe]

Fix format issues

[Lewuathe]

Add test for checking lexicographical order

[Lewuathe]

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

650–652

I think there should be a dependency from the store to write (e.g. store %m[0] (%i=0) to load %m[0] (%i = 9). But this case is not properly checked.

func.func @no_lexicographic_order() {
  %m = memref.alloc() : memref<10xf32>
  %c7 = arith.constant 7.0 : f32
  affine.parallel (%i) = (0) to (10) {
    %1 = affine.load %m[9 - %i] : memref<10xf32>
    // expected-remark@above {{dependence from 0 to 0 at depth 1 = false}}
    // expected-remark@above {{dependence from 0 to 1 at depth 1 = true}}
    affine.store %1, %m[%i] : memref<10xf32>
    // expected-remark@above {{dependence from 1 to 0 at depth 1 = false}}
    // expected-remark@above {{dependence from 1 to 1 at depth 1 = false}}
  }
  return
}

[bondhugula]

Thanks for contributing this support. This was long overdue. While this is fine, I feel just directly supporting the general (arbitrary bound case) is better -- since the affine.for support already supports that, and we may have to again throw away most of this code when handling the general bound case. Just supporting the constant bound case to start with isn't an incremental step.

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

253

Use comments on top (LLVM/MLIR style), not trailing.

259–260

May be good to fix the comment while on this:
associated -> associating

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

650

Improve the assertion message here: variable expected for the IV value?

655–658

Negate condition and unguard addBound.

661–664

Likewise.

[Lewuathe]

@bondhugula Thank you for the comments. Sorry for taking the time to work on this. I was struggling with how to treat the ordering constraint in the affine.parallel. (Just removing the ordering constraint does not work as expected).

So for now, should we discard this patch completely and restart working on supporting general case?

[bondhugula]

In D136056#3910972, @Lewuathe wrote:

@bondhugula Thank you for the comments. Sorry for taking the time to work on this. I was struggling with how to treat the ordering constraint in the affine.parallel. (Just removing the ordering constraint does not work as expected).

So for now, should we discard this patch completely and restart working on supporting general case?

It would be good to support the generic case. Parts of this patch are still useful towards that. I didn't understand the issue related to "ordering constraint". You would order the dimensions in the order in which they appear in the affine.parallel.

[Groverkss]

In D136056#3911760, @bondhugula wrote:

In D136056#3910972, @Lewuathe wrote:

@bondhugula Thank you for the comments. Sorry for taking the time to work on this. I was struggling with how to treat the ordering constraint in the affine.parallel. (Just removing the ordering constraint does not work as expected).

So for now, should we discard this patch completely and restart working on supporting general case?

It would be good to support the generic case. Parts of this patch are still useful towards that. I didn't understand the issue related to "ordering constraint". You would order the dimensions in the order in which they appear in the affine.parallel.

I think the "ordering constraint" refers to the dependence analysis, since affine.parallel does not define a schedule for iterations, we cannot impose ordering constraints based on lexicographic order as it currently does for every induction variable.

I think the patch is good other than the dependence analysis part since that needs more discussion. For example, checkMemRefDependence takes a loop depth. Currently, loop depth implies the induction variable at that position. What would that mean after introducing affine.paralllel to the dependence analysis? Would it mean after affine.parallel is allowed? Would it mean the loop or the induction variable? I think some things need to be discussed more before we support it.

I would be happy to review the patch if you send the generic case without dependence analysis.

[Lewuathe]

@bondhugula @Groverkss Thanks again. I'd like to work on the general cases for supporting affine.parallel in the getIndexSet.

But I have one question about the test for that case. Where should we write the test if we do not support the dependence analysis for affine.parallel for the time being? Is there any existing test checking the getIndexSet functionality?

[bondhugula]

In D136056#3914145, @Lewuathe wrote:

@bondhugula @Groverkss Thanks again. I'd like to work on the general cases for supporting affine.parallel in the getIndexSet.

But I have one question about the test for that case. Where should we write the test if we do not support the dependence analysis for affine.parallel for the time being? Is there any existing test checking the getIndexSet functionality?

If we don't support it temporarily, you can always return failure() from getIndexSet which will cleanly fail the dependence analysis check. There is a status for "dependence analysis failed" that gets returned already.

[Lewuathe]

@bondhugula Okay, I'll return failure() to unsupport the dependence analysis with affine.parallel for the time being. Thanks!

[Lewuathe]

Support non-constant constraints.

[Lewuathe]

Post review follow-up.

[Lewuathe]

Apply format.

[Groverkss]

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

148

Is this stil lvalid?

149

Is it required that the variable exists in the constraint system? The documentation does not make this clear.

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

616

This statement is misleading. Dependence analysis isn't implemented for affine.parallel. This statement claims it's not possible. Please add a TODO instead explaining why we did not implement it right now, which is affine.parallel does not specify an ordering.

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

654

Is there any reason we specialize for the constant case? If there is, could you write a comment explaining why we specialize? It is not obvious to me.

[Lewuathe]

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

654

The addBound method depends on whether the bound is constant. In the case of non-constant, we need to give operands explicitly. We use this pattern for affine.for and it's also the case for affine.parallel too.

[Lewuathe]

Refine TODO comments.

[bondhugula]

The added code isn't exercised by any test case, right? Is this ready for review?

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

459

What can operations contain? operations -> affineOps?

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

270

Can use auto here since the type is on the RHS.

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

646

The debug message isn't helpful here. Suggest adding more text to it.

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

259–264

This is nothing but checking getParentOpOfType<AffineParallelOp>() is null. We don't need to add a public method!

2325

Pass output argument by reference (MLIR convention).

2329

Use auto here.

2331

Use braces here per new LLVM/MLIR style.

[Lewuathe]

Thank you for taking a look again.

I'll try to add a test case covering getIndexSet somehow. But there seems to be no dedicated test case for getIndexSet. Test cases for dependent analysis only check the logic of getIndexSet. Do you know whether there is already any test code covering getIndexSet usage?

In D136056#3923319, @bondhugula wrote:

In D136056#3914145, @Lewuathe wrote:

@bondhugula @Groverkss Thanks again. I'd like to work on the general cases for supporting affine.parallel in the getIndexSet.

But I have one question about the test for that case. Where should we write the test if we do not support the dependence analysis for affine.parallel for the time being? Is there any existing test checking the getIndexSet functionality?

If we don't support it temporarily, you can always return failure() from getIndexSet which will cleanly fail the dependence analysis check. There is a status for "dependence analysis failed" that gets returned already.

[bondhugula]

In D136056#3948381, @Lewuathe wrote:

Thank you for taking a look again.

I'll try to add a test case covering getIndexSet somehow. But there seems to be no dedicated test case for getIndexSet. Test cases for dependent analysis only check the logic of getIndexSet. Do you know whether there is already any test code covering getIndexSet usage?

I didn't understand - you don't need to add test cases for the utility. Just test dependence analysis in the presence of affine.parallel. This can be done via the affine-scalrep pass or test-memref-dependence-check.

[Lewuathe]

Add test cases and reflect the latest reviews.

[Lewuathe]

@bondhugula

Sorry for the confusion. I just wanted to know how/where to add test cases for the utility functions like this. But I've found a way to test the case including affine.parallel somehow for dependence check and scalrep. I'm so sorry for bothering you many times but could you review this change again when you get a chance?

[bondhugula]

Can you please resolve the other remaining comments and mark them done before a round of review? For eg. all comments in AffineStructures.cpp many of which are trivial are still pending.

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

2332–2334

Drop trivial braces.

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

1078–1080 ↗

(On Diff #477855)

There seems to be something incorrect here. Both these "expected" checks are pointing to the same line (@above from below is the same as +1 here).

mlir/test/lib/Analysis/TestMemRefDependenceCheck.cpp

90 ↗

(On Diff #477855)

Thank you for improving/fixing this.

[Lewuathe]

Validate affine.parallel existence strictly in memref dependence check.

[Lewuathe]

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

1078–1080 ↗

(On Diff #477855)

We should have checked the affine.parallel existence before returning no dependence when there is no AffineWriteOpInterface.

[bondhugula]

This looks great to me - thanks! A comment on some missing test cases below.

mlir/test/Dialect/Affine/scalrep.mlir

796 ↗

(On Diff #478805)

Missing test case for non-constant loop bound case and mixed affine.parallel + affine.for case. (For eg., affine.for surrounded by affine.parallel or vice versa) -- either here or in the dependence check test case file.

[Lewuathe]

Post review follow-up

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

 #include "mlir/IR/OpDefinition.h"
 #include "mlir/Support/LogicalResult.h"
 
 namespace mlir {
 
 class AffineCondition;
 class AffineForOp;
 class AffineIfOp;
+class AffineParallelOp;
 class AffineMap;
 class AffineValueMap;
 class IntegerSet;
 class MLIRContext;
 class Value;
 class MemRefType;
 struct MutableAffineMap;
 
public:
   /// constraint system. Returns failure for the yet unimplemented/unsupported
   /// cases.  Any new variables that are found in the bound operands of the
   /// 'affine.for' operation are added as trailing variables (either
   /// dimensional or symbolic depending on whether the operand is a valid
   /// symbol).
   //  TODO: add support for non-unit strides.
   LogicalResult addAffineForOpDomain(AffineForOp forOp);
 
+  /// Add constraints (lower and upper bounds) for the specified
+  /// 'affine.parallel' operation's Value using IR information stored in its
+  /// bound maps. Returns failure for the yet unimplemented/unsupported cases.
+  /// Asserts if the Value corresponding to the 'affine.parallel' operation
+  /// isn't found in the constraint system.
+  LogicalResult addAffineParallelOpDomain(AffineParallelOp parallelOp);
+
   /// Adds constraints (lower and upper bounds) for each loop in the loop nest
   /// described by the bound maps `lbMaps` and `ubMaps` of a computation slice.
   /// Every pair (`lbMaps[i]`, `ubMaps[i]`) describes the bounds of a loop in
   /// the nest, sorted outer-to-inner. `operands` contains the bound operands
   /// for a single bound map. All the bound maps will use the same bound
   /// operands. Note that some loops described by a computation slice might not
   /// exist yet in the IR so the Value attached to those dimension variables
   /// might be empty. For that reason, this method doesn't perform Value

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

 /// top level of a `AffineScope` region is always a valid symbol for all
 /// uses in that region.
 bool isTopLevelValue(Value value, Region *region);
 
 /// Returns the closest region enclosing `op` that is held by an operation with
 /// trait `AffineScope`; `nullptr` if there is no such region.
 Region *getAffineScope(Operation *op);
 
+/// Returns whether the op is within an affine parallel op which can potentially
+/// preventing the dependence analysis.
+bool isInParallelOp(Operation *op);
+
 /// AffineDmaStartOp starts a non-blocking DMA operation that transfers data
 /// from a source memref to a destination memref. The source and destination
 /// memref need not be of the same dimensionality, but need to have the same
 /// elemental type. The operands include the source and destination memref's
 /// each followed by its indices, size of the data transfer in terms of the
 /// number of elements (of the elemental type of the memref), a tag memref with
 /// its indices, and optionally at the end, a stride and a
 /// number_of_elements_per_stride arguments. The tag location is used by an

 /// not an induction variable, then return nullptr.
 AffineForOp getForInductionVarOwner(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 *> operations,
+                          SmallVectorImpl<Value> *ivs);
+
 /// Builds a perfect nest of affine.for loops, i.e., each loop except the
 /// innermost one contains only another loop and a terminator. The loops iterate
 /// from "lbs" to "ubs" with "steps". The body of the innermost loop is
 /// populated by calling "bodyBuilderFn" and providing it with an OpBuilder, a
 /// Location and a list of loop induction variables.
 void buildAffineLoopNest(OpBuilder &builder, Location loc,
                          ArrayRef<int64_t> lbs, ArrayRef<int64_t> ubs,
                          ArrayRef<int64_t> steps,

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

 // the yet unimplemented cases.
 // 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;
+  SmallVector<Operation *, 8> loopOps;
+  size_t numDims = 0;
   for (Operation *op : ops) {
-    if (!isa<AffineForOp, AffineIfOp>(op)) {
-      // TODO: Support affine.parallel ops.
-      LLVM_DEBUG(llvm::dbgs() << "getIndexSet only handles affine.for/if ops");
+    if (!isa<AffineForOp, AffineIfOp, AffineParallelOp>(op)) {
+      LLVM_DEBUG(llvm::dbgs() << "getIndexSet only handles affine.for/if/"
+                                 "parallel ops");
       return failure();
     }
-    if (AffineForOp forOp = dyn_cast<AffineForOp>(op))
-      forOps.push_back(forOp);
+    if (AffineForOp forOp = dyn_cast<AffineForOp>(op)) {
+      loopOps.push_back(forOp);
+      // An AffineForOp retains only 1 induction variable.
+      numDims += 1;
+    } else if (AffineParallelOp parallelOp = dyn_cast<AffineParallelOp>(op)) {
+      loopOps.push_back(parallelOp);
+      numDims += parallelOp.getNumDims();
+    }
   }
-  extractForInductionVars(forOps, &indices);
-  // Reset while associated Values in 'indices' to the domain.
-  domain->reset(forOps.size(), /*numSymbols=*/0, /*numLocals=*/0, indices);
+  extractInductionVars(loopOps, &indices);
+  // Reset while associating Values in 'indices' to the domain.
+  domain->reset(numDims, /*numSymbols=*/0, /*numLocals=*/0, indices);
   for (Operation *op : ops) {
     // Add constraints from forOp's bounds.
     if (AffineForOp forOp = dyn_cast<AffineForOp>(op)) {
       if (failed(domain->addAffineForOpDomain(forOp)))
         return failure();
     } else if (AffineIfOp ifOp = dyn_cast<AffineIfOp>(op)) {
       domain->addAffineIfOpDomain(ifOp);
-    }
+    } else if (AffineParallelOp parallelOp = dyn_cast<AffineParallelOp>(op))
+      if (failed(domain->addAffineParallelOpDomain(parallelOp)))
+        return failure();
   }
   return success();
 }
 
 /// 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.
DependenceResult mlir::checkMemrefAccessDependence(
   if (!allowRAR && !isa<AffineWriteOpInterface>(srcAccess.opInst) &&
       !isa<AffineWriteOpInterface>(dstAccess.opInst))
     return DependenceResult::NoDependence;
 
   // We can't analyze further if the ops lie in different affine scopes.
   if (getAffineScope(srcAccess.opInst) != getAffineScope(dstAccess.opInst))
     return DependenceResult::Failure;
 
+  // TODO: Support affine.parallel which does not specify the ordering.
+  if (isInParallelOp(srcAccess.opInst) || isInParallelOp(dstAccess.opInst))
+    return DependenceResult::Failure;
+
   // Create access relation from each MemRefAccess.
   FlatAffineRelation srcRel, dstRel;
   if (failed(srcAccess.getAccessRelation(srcRel)))
     return DependenceResult::Failure;
   if (failed(dstAccess.getAccessRelation(dstRel)))
     return DependenceResult::Failure;
 
   FlatAffineValueConstraints srcDomain = srcRel.getDomainSet();
DependenceResult mlir::checkMemrefAccessDependence(
   // Compute the dependence relation by composing `srcRel` with the inverse of
   // `dstRel`. Doing this builds a relation between iteration domain of
   // `srcAccess` to the iteration domain of `dstAccess` which access the same
   // memory locations.
   dstRel.inverse();
   dstRel.compose(srcRel);
   *dependenceConstraints = dstRel;
 
   // Add 'src' happens before 'dst' ordering constraints.
   addOrderingConstraints(srcDomain, dstDomain, loopDepth,
                          dependenceConstraints);
 
   // Return 'NoDependence' if the solution space is empty: no dependence.
   if (dependenceConstraints->isEmpty())
     return DependenceResult::NoDependence;
 
   // Compute dependence direction vector and return true.
   if (dependenceComponents != nullptr)
     computeDirectionVector(srcDomain, dstDomain, loopDepth,

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

if (forOp.hasConstantUpperBound()) {
     addBound(BoundType::UB, pos, forOp.getConstantUpperBound() - 1);
     return success();
   }
   // Non-constant upper bound case.
   return addBound(BoundType::UB, pos, forOp.getUpperBoundMap(),
                   forOp.getUpperBoundOperands());
 }
 
+LogicalResult FlatAffineValueConstraints::addAffineParallelOpDomain(
+    AffineParallelOp parallelOp) {
+  size_t ivPos = 0;
+  for (auto iv : parallelOp.getIVs()) {
+    LLVM_DEBUG(iv.dump(););
+    unsigned pos;
+    if (!findVar(iv, &pos)) {
+      assert(false && "variable expected for the IV value");
+      return failure();
+    }
+
+    AffineMap lowerBound = parallelOp.getLowerBoundMap(ivPos);
+    if (lowerBound.isConstant())
+      addBound(BoundType::LB, pos, lowerBound.getSingleConstantResult());
+    else if (failed(addBound(BoundType::LB, pos, lowerBound,
+                             parallelOp.getLowerBoundsOperands())))
+      return failure();
+
+    auto upperBound = parallelOp.getUpperBoundMap(ivPos);
+    if (upperBound.isConstant())
+      addBound(BoundType::UB, pos, upperBound.getSingleConstantResult());
+    else if (failed(addBound(BoundType::UB, pos, upperBound,
+                             parallelOp.getUpperBoundsOperands())))
+      return failure();
+  }
+  return success();
+}
+
 LogicalResult
 FlatAffineValueConstraints::addDomainFromSliceMaps(ArrayRef<AffineMap> lbMaps,
                                                    ArrayRef<AffineMap> ubMaps,
                                                    ArrayRef<Value> operands) {
   assert(lbMaps.size() == ubMaps.size());
   assert(lbMaps.size() <= getNumDimVars());
 
   for (unsigned i = 0, e = lbMaps.size(); i < e; ++i) {

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

Region *mlir::getAffineScope(Operation *op) {
   while (auto *parentOp = curOp->getParentOp()) {
     if (parentOp->hasTrait<OpTrait::AffineScope>())
       return curOp->getParentRegion();
     curOp = parentOp;
   }
   return nullptr;
 }
 
+bool mlir::isInParallelOp(mlir::Operation *op) {
+  auto *curOp = op;
+  while (auto *parentOp = curOp->getParentOp()) {
+    if (isa<AffineParallelOp>(parentOp))
+      return true;
+    curOp = parentOp;
+  }
+  return false;
+}
+
 // A Value can be used as a dimension id iff it meets one of the following
 // conditions:
 // *) It is valid as a symbol.
 // *) It is an induction variable.
 // *) It is the result of affine apply operation with dimension id arguments.
 bool mlir::isValidDim(Value value) {
   // The value must be an index type.
   if (!value.getType().isIndex())

 /// them.
 void mlir::extractForInductionVars(ArrayRef<AffineForOp> forInsts,
                                    SmallVectorImpl<Value> *ivs) {
   ivs->reserve(forInsts.size());
   for (auto forInst : forInsts)
     ivs->push_back(forInst.getInductionVar());
 }
 
+void mlir::extractInductionVars(ArrayRef<mlir::Operation *> operations,
+                                SmallVectorImpl<mlir::Value> *ivs) {
+  ivs->reserve(operations.size());
+  for (Operation *op : operations) {
+    // Add constraints from forOp's bounds.
+    if (AffineForOp forOp = dyn_cast<AffineForOp>(op))
+      ivs->push_back(forOp.getInductionVar());
+    else if (AffineParallelOp parallelOp = dyn_cast<AffineParallelOp>(op))
+      for (size_t i = 0; i < parallelOp.getBody()->getNumArguments(); i++)
+        ivs->push_back(parallelOp.getBody()->getArgument(i));
+  }
+}
+
 /// Builds an affine loop nest, using "loopCreatorFn" to create individual loop
 /// operations.
 template <typename BoundListTy, typename LoopCreatorTy>
 static void buildAffineLoopNestImpl(
     OpBuilder &builder, Location loc, BoundListTy lbs, BoundListTy ubs,
     ArrayRef<int64_t> steps,
     function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn,
     LoopCreatorTy &&loopCreatorFn) {