[MLIR] Consider AffineIfOp when getting the index set of an Op wrapped in nested loops

[kumasento]

Summary

This diff attempts to resolve the TODO in getOpIndexSet (formerly known as getInstIndexSet), which states "Add support to handle IfInsts surronding op".

Major changes in this diff:

1. Overload getIndexSet. The overloaded version considers both AffineForOp and AffineIfOp.
2. The getInstIndexSet is updated accordingly: its name is changed to getOpIndexSet and its implementation is based on a new API getIVs instead of getLoopIVs.
3. Add addAffineIfOpDomain to FlatAffineConstraints, which extracts new constraints from the integer set of AffineIfOp and merges it to the current constraint system.
4. Update how a Value is determined as dim or symbol for ValuePositionMap in buildDimAndSymbolPositionMaps.

[rriddle]

mlir/include/mlir/Analysis/AffineAnalysis.h

47

nit: Please use /// for top-level comments. Same applies to the rest of this commit.

[kumasento]

Fixed top-level comment problems and reverted AffineOps.cpp changes.

[bondhugula]

Thanks very much for completing these! I'll be able to review this tomorrow. You could consider adding a few test cases to exercise these. These changes could be tested by adding test cases with affine.if's for say memref dep analysis or memref-bound-checking or affine data copy generation - whichever appears easier.

mlir/lib/Analysis/AffineAnalysis.cpp

113–114

This can be named getOpIndexSet now - the inst is a remnant of the past when Operations in MLIR used to be called Instructions and other things.

[kumasento]

1. Added unittest for memref-dependence-check with affine.if
2. Fixed issues in getCommonBlock that doesn't consider the existence of affine.if
3. Fixed issues in addAffineIfDomain that doesn't consider local variable
4. Changed getInstIndexSet naming to getOpIndexSet

[kumasento]

In D84698#2180353, @bondhugula wrote:

Thanks very much for completing these! I'll be able to review this tomorrow. You could consider adding a few test cases to exercise these. These changes could be tested by adding test cases with affine.if's for say memref dep analysis or memref-bound-checking or affine data copy generation - whichever appears easier.

Thank you very much Uday for your comments :)

Now I've added some unittests based on test-memref-dependence-check. They are in a separate file but I can put them into memref-dependence-check.mlir if you prefer.

[kumasento]

Squashed recent commits into one.

[bondhugula]

Looks great so far - thanks. Some more comments.

mlir/include/mlir/Analysis/AffineAnalysis.h

36–47

We really don't need two versions - the one you added should be the single version and the right one to use. Please merge these two into one?

mlir/lib/Analysis/AffineAnalysis.cpp

113–114

opInst is outdated - update doc comment.

114–117

Shouldn't you be asserting that everything in ops is either an AffineForOp or an AffineIfOp?

265–269

Can drop trivial braces here.

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

1 ↗

(On Diff #281488)

Can you append these to the existing file on dependence tests?

[kumasento]

Improved API design, comments, and unittests.

[kumasento]

Hi @bondhugula Thank you for your comments. I've updated this diff according to your suggestions.

[kumasento]

mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp

294

BTW I'm not sure whether there is any way to make this ops creation a bit more concise?

[bondhugula]

mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp

292–296

Please reserve space for ops. ops.reserve(input.size()).

293

Please avoid auto here.

294

This is fine.

mlir/lib/Transforms/Utils/LoopUtils.cpp

2338–2342

Likewise. Reserve space please.

2339

Avoid auto.

2426–2427

2 -> 1 ?

[bondhugula]

mlir/include/mlir/Analysis/AffineAnalysis.h

38

Please continue on previous line. Put ops in back ticks.

mlir/include/mlir/Analysis/AffineStructures.h

222

This line may be confusing. An IntegerSet doesn't have any SSA values with itself. Just drop the "Since an IntegerSet ..., " part and just have the second clause to the effect: "The ifOp's operands are bound to the corresponding ... ".

224

If this only returns success now, please remove the LogicalResult. Also, mention about the assertion here about the unhandled cases. (For eg. IntegerSet having semi-affine expresssions.)

mlir/lib/Analysis/AffineAnalysis.cpp

93

Avoid auto.

102

Avoid auto here.

113–114

Triple /// please.

607

Comment here - not obvious.

613

Debug code.

639

Avoid auto.

mlir/lib/Analysis/Utils.cpp

50

You can use variadic isa: (isa<..., ...>).

[kumasento]

Fixed issues pinpointed by review comments.

[kumasento]

Thank you very much Uday for pointing out all these issues! Sorry for costing you much time for doing so. I have fixed all of them and some other similar things. Would you mind taking another round of review to see if it is OK? Thanks!

mlir/include/mlir/Analysis/AffineStructures.h

222

So sorry for the confusion. As you said that an IntegerSet indeed doesn't have any SSA values, however, the columns in the constraint system should be bound to values. Here what I'm trying to say is that I bind ifOp's operands to those columns to workaround this problem.

But I suppose this point might be too trivial. I would just remove this sentence then. Please let me know if you prefer it stays there (revised based on your comments).

[bondhugula]

mlir/include/mlir/Analysis/Utils.h

42

This doc comment isn't clear. There is no IV for an affine.if - please rewrite this doc.

mlir/lib/Analysis/AffineAnalysis.cpp

607

"a list of parent blocks" is confusing. Instead, the chain of ancestor blocks?

608

-> ... when either a FuncOp or endBlock is reached. ?

615

Instead of checking for FuncOp, check for an op with trait AffineScope (using hasTrait). FuncOp has AffineScope, but there are other ops also that define an AffineScope.

630–631

commonForValue -> commonForIV

mlir/lib/Analysis/AffineStructures.cpp

731

mlir::Value -> Value

mlir/lib/Analysis/Utils.cpp

49

You may want to do a ops->clear() in the beginning.

[kumasento]

Improved doc and API design based on reviews.

[kumasento]

Thank you again Uday for your reviews! Please check again if I've missed any other things. :)

mlir/lib/Analysis/AffineAnalysis.cpp

615

Thanks! Can I say that the idea is not to go out of the current AffineScope?

[bondhugula]

mlir/lib/Analysis/AffineAnalysis.cpp

615

That's exactly right.

[kumasento]

Improved comments in buildDimAndSymbolPositonMaps.

[kumasento]

Hi Uday @bondhugula sorry to bother but would you mind giving this diff one more round of review when you've got time? Thanks!

[bondhugula]

mlir/include/mlir/Analysis/AffineStructures.h

219–220

Reflow to use entire width.

mlir/lib/Analysis/AffineAnalysis.cpp

231

In another word -> In other words

know -> do not know ?

332–333

Document isSymbolDetermined please.

mlir/lib/Analysis/Utils.cpp

50

The name here would be incorrect. If you have AffineIfOp's interleaved, your output is no longer IVs (induction variables). This should be named more like getEnclosingAffineForAndIfOps. In fact, all users of this that just expect IVs would behave wrongly when given an input that has both for's and if's surrounding op.

[kumasento]

1. Changed the naming of getIVs to getEnclosingAffineForAndIfOps
2. Improved updateValuePosMap documentation
3. Added comments to isSymbolDetermined

[kumasento]

Thank you very much Uday @bondhugula for all your help! I've just improved the diff based on your reviews, would you mind taking a look again when you have time? Thanks!

mlir/include/mlir/Analysis/AffineStructures.h

219–220

Thanks! Would it be ok that I place the bullet points in separate lines?

mlir/lib/Analysis/AffineAnalysis.cpp

231

Thank you very much for this suggestion! I realized that this version is still not clear, and I try to rephrase it with more details. Hope this version has much less confusion.

mlir/lib/Analysis/Utils.cpp

50

Thank you for pointing this out! I was confused by the naming as well and I should have raised this issue earlier. I think getEnclosingAffienForAndIfOps is a clear name. I just changed getIVs to that :)

[kumasento]

mlir/lib/Analysis/AffineAnalysis.cpp

242

I'm not very confident in this statement. It is based on an assumption that an operand from srcAccessMap or dstAccessMap should not only appear in affine.if. If not, I'm not sure how I could distinguish whether an operand is a dim or a symbol.

[bondhugula]

Looks good, thanks.

mlir/lib/Transforms/Utils/LoopUtils.cpp

2341

You don't need getOperation() here I think.

[kumasento]

Removed redundant getOperation().

[kumasento]

Thank you Uday! @bondhugula Would you mind helping me commit this if it looks OK :)

[bondhugula]

mlir/include/mlir/Analysis/AffineAnalysis.h

36–37

This sentence is no longer accurate: "... of the forOps ...".

[bondhugula]

mlir/include/mlir/Analysis/AffineAnalysis.h

36–37

Please also mention here that constraints from AffineIfOp's index set are also added.

mlir/include/mlir/Analysis/Utils.h

44

This comment is also no longer accurate:

-> "... outermost `affine.for/affine.if operation to ..."

mlir/lib/Analysis/AffineAnalysis.cpp

614

Please add a comment for this while loop here.

[kumasento]

Improved comments.

[kumasento]

Improved comments of getIndexSet.

[kumasento]

Hi Uday @bondhugula Thank you for your latest review :) Please see whether the latest version looks better.

[bondhugula]

Some minor fixes.

mlir/include/mlir/Analysis/AffineAnalysis.h

37–42

values in AffineIfOps' index set -> AffineIfOp's operands

39

-> from the index sets of AffineIfOps ...

[kumasento]

Minor fixes in comments.

[kumasento]

In D84698#2204751, @bondhugula wrote:

Some minor fixes.

Done :) Thank you Uday @bondhugula

mlir/include/mlir/Analysis/AffineAnalysis.h

 class Operation;
 
 /// Returns in `affineApplyOps`, the sequence of those AffineApplyOp
 /// Operations that are reachable via a search starting from `operands` and
 /// ending at those operands that are not the result of an AffineApplyOp.
 void getReachableAffineApplyOps(ArrayRef<Value> operands,
                                 SmallVectorImpl<Operation *> &affineApplyOps);
 
 /// Builds a system of constraints with dimensional identifiers corresponding to
-/// the loop IVs of the forOps appearing in that order. Bounds of the loop are
-/// used to add appropriate inequalities. Any symbols founds in the bound
-/// operands are added as symbols in the system. Returns failure for the yet
-/// unimplemented cases.
+/// the loop IVs of the forOps and AffineIfOp's operands appearing in
+/// that order. Bounds of the loop are used to add appropriate inequalities.
+/// Constraints from the index sets of AffineIfOp are also added. Any symbols
+/// founds in the bound operands are added as symbols in the system. Returns
+/// failure for the yet unimplemented cases. `ops` accepts both AffineForOp and
+/// AffineIfOp.
 //  TODO: handle non-unit strides.
-LogicalResult getIndexSet(MutableArrayRef<AffineForOp> forOps,
+LogicalResult getIndexSet(MutableArrayRef<Operation *> ops,
                           FlatAffineConstraints *domain);
 
 /// Encapsulates a memref load or store access information.
 struct MemRefAccess {
   Value memref;
   Operation *opInst;
   SmallVector<Value, 4> indices;
 
   /// Constructs a MemRefAccess from a load or store operation.
   // TODO: add accessors to standard op's load, store, DMA op's to return
   // MemRefAccess, i.e., loadOp->getAccess(), dmaOp->getRead/WriteAccess.

mlir/include/mlir/Analysis/AffineStructures.h

 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/Support/LogicalResult.h"
 
 namespace mlir {
 
 class AffineCondition;
 class AffineForOp;
+class AffineIfOp;
 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 identifiers that are found in the bound operands of the
   /// 'affine.for' operation are added as trailing identifiers (either
   /// dimensional or symbolic depending on whether the operand is a valid
   /// symbol).
   //  TODO: add support for non-unit strides.
   LogicalResult addAffineForOpDomain(AffineForOp forOp);
 
+  /// Adds constraints imposed by the `affine.if` operation. These constraints
+  /// are collected from the IntegerSet attached to the given `affine.if`
+  /// instance argument (`ifOp`). It is asserted that:
+  /// 1) The IntegerSet of the given `affine.if` instance should not contain
+  /// semi-affine expressions,
+  /// 2) The columns of the constraint system created from `ifOp` should match
+  /// the columns in the current one regarding numbers and values.
+  void addAffineIfOpDomain(AffineIfOp ifOp);
+
   /// Adds a lower or an upper bound for the identifier at the specified
   /// position with constraints being drawn from the specified bound map and
   /// operands. If `eq` is true, add a single equality equal to the bound map's
   /// first result expr.
   LogicalResult addLowerOrUpperBound(unsigned pos, AffineMap boundMap,
                                      ValueRange operands, bool eq,
                                      bool lower = true);
 

mlir/include/mlir/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);
+
 /// Returns the nesting depth of this operation, i.e., the number of loops
 /// surrounding this operation.
 unsigned getNestingDepth(Operation *op);
 
 /// Returns in 'sequentialLoops' all sequential loops in loop nest rooted
 /// at 'forOp'.
 void getSequentialLoops(AffineForOp forOp,
                         llvm::SmallDenseSet<Value, 8> *sequentialLoops);

mlir/lib/Analysis/AffineAnalysis.cpp

 // Builds a system of constraints with dimensional identifiers corresponding to
 // the loop IVs of the forOps appearing in that order. Any symbols founds in
 // the bound operands are added as symbols in the system. Returns failure for
 // the yet unimplemented cases.
 // TODO: Handle non-unit steps through local variables or stride information in
 // FlatAffineConstraints. (For eg., by using iv - lb % step = 0 and/or by
 // introducing a method in FlatAffineConstraints setExprStride(ArrayRef<int64_t>
 // expr, int64_t stride)
-LogicalResult mlir::getIndexSet(MutableArrayRef<AffineForOp> forOps,
+LogicalResult mlir::getIndexSet(MutableArrayRef<Operation *> ops,
                                 FlatAffineConstraints *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 (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 (auto forOp : forOps) {
+  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);
+    }
   }
   return success();
 }
 
-// Computes the iteration domain for 'opInst' and populates 'indexSet', which
-// encapsulates the constraints involving loops surrounding 'opInst' and
-// potentially involving any Function symbols. The dimensional identifiers in
-// 'indexSet' correspond to the loops surrounding 'op' from outermost to
-// innermost.
-// TODO: Add support to handle IfInsts surrounding 'op'.
-static LogicalResult getInstIndexSet(Operation *op,
-                                     FlatAffineConstraints *indexSet) {
-  // TODO: Extend this to gather enclosing IfInsts and consider
-  // factoring it out into a utility function.
-  SmallVector<AffineForOp, 4> loops;
-  getLoopIVs(*op, &loops);
-  return getIndexSet(loops, indexSet);
+/// 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 identifiers in
+/// 'indexSet' correspond to the loops surrounding 'op' from outermost to
+/// innermost.
+static LogicalResult getOpIndexSet(Operation *op,
+  SmallVector<Operation *, 4> ops;
+  getEnclosingAffineForAndIfOps(*op, &ops);
+  return getIndexSet(ops, indexSet);
 }
 
 namespace {
 // ValuePositionMap manages the mapping from Values which represent dimension
 // and symbol identifiers from 'src' and 'dst' access functions to positions
 // in new space where some Values are kept separate (using addSrc/DstValue)
 // and some Values are merged (addSymbolValue).
 // Position lookups return the absolute position in the new space which

 // This method populates 'valuePosMap' with mappings from operand Values in
 // 'srcAccessMap'/'dstAccessMap' (as well as those in 'srcDomain'/'dstDomain')
 // to the position of these values in the merged list.
 static void buildDimAndSymbolPositionMaps(
     const FlatAffineConstraints &srcDomain,
     const FlatAffineConstraints &dstDomain, const AffineValueMap &srcAccessMap,
     const AffineValueMap &dstAccessMap, ValuePositionMap *valuePosMap,
     FlatAffineConstraints *dependenceConstraints) {
-  auto updateValuePosMap = [&](ArrayRef<Value> values, bool isSrc) {
+
+  // IsDimState is a tri-state boolean. It is used to distinguish three
+  // different cases of the values passed to updateValuePosMap.
+  // - When it is TRUE, we are certain that all values are dim values.
+  // - When it is FALSE, we are certain that all values are symbol values.
+  // - When it is UNKNOWN, we need to further check whether the value is from a
+  // loop IV to determine its type (dim or symbol).
+
+  // We need this enumeration because sometimes we cannot determine whether a
+  // Value is a symbol or a dim by the information from the Value itself. If a
+  // Value appears in an affine map of a loop, we can determine whether it is a
+  // dim or not by the function `isForInductionVar`. But when a Value is in the
+  // affine set of an if-statement, there is no way to identify its category
+  // (dim/symbol) by itself. Fortunately, the Values to be inserted into
+  // `valuePosMap` come from `srcDomain` and `dstDomain`, and they hold such
+  // information of Value category: `srcDomain` and `dstDomain` organize Values
+  // by their category, such that the position of each Value stored in
+  // `srcDomain` and `dstDomain` marks which category that a Value belongs to.
+  // Therefore, we can separate Values into dim and symbol groups before passing
+  // them to the function `updateValuePosMap`. Specifically, when passing the
+  // dim group, we set IsDimState to TRUE; otherwise, we set it to FALSE.
+  // However, Values from the operands of `srcAccessMap` and `dstAccessMap` are
+  // not explicitly categorized into dim or symbol, and we have to rely on
+  // `isForInductionVar` to make the decision. IsDimState is set to UNKNOWN in
+  // this case.
+  enum IsDimState { TRUE, FALSE, UNKNOWN };
+
+  // This function places each given Value (in `values`) under a respective
+  // category in `valuePosMap`. Specifically, the placement rules are:
+  // 1) If `isDim` is FALSE, then every value in `values` are inserted into
+  // `valuePosMap` as symbols.
+  // 2) If `isDim` is UNKNOWN and the value of the current iteration is NOT an
+  // induction variable of a for-loop, we treat it as symbol as well.
+  // 3) For other cases, we decide whether to add a value to the `src` or the
+  // `dst` section of the dim category simply by the boolean value `isSrc`.
+  auto updateValuePosMap = [&](ArrayRef<Value> values, bool isSrc,
+                               IsDimState isDim) {
     for (unsigned i = 0, e = values.size(); i < e; ++i) {
       auto value = values[i];
-      if (!isForInductionVar(values[i])) {
-        assert(isValidSymbol(values[i]) &&
+      if (isDim == FALSE || (isDim == UNKNOWN && !isForInductionVar(value))) {
+        assert(isValidSymbol(value) &&
                "access operand has to be either a loop IV or a symbol");
         valuePosMap->addSymbolValue(value);
-      } else if (isSrc) {
-        valuePosMap->addSrcValue(value);
       } else {
+        if (isSrc)
+          valuePosMap->addSrcValue(value);
+        else
           valuePosMap->addDstValue(value);
       }
     }
   };
 
-  SmallVector<Value, 4> srcValues, destValues;
-  srcDomain.getIdValues(0, srcDomain.getNumDimAndSymbolIds(), &srcValues);
-  dstDomain.getIdValues(0, dstDomain.getNumDimAndSymbolIds(), &destValues);
-  // Update value position map with identifiers from src iteration domain.
-  updateValuePosMap(srcValues, /*isSrc=*/true);
-  // Update value position map with identifiers from dst iteration domain.
-  updateValuePosMap(destValues, /*isSrc=*/false);
+  // Collect values from the src and dst domains. For each domain, we separate
+  // the collected values into dim and symbol parts.
+  SmallVector<Value, 4> srcDimValues, dstDimValues, srcSymbolValues,
+      dstSymbolValues;
+  srcDomain.getIdValues(0, srcDomain.getNumDimIds(), &srcDimValues);
+  dstDomain.getIdValues(0, dstDomain.getNumDimIds(), &dstDimValues);
+  srcDomain.getIdValues(srcDomain.getNumDimIds(),
+                        srcDomain.getNumDimAndSymbolIds(), &srcSymbolValues);
+  dstDomain.getIdValues(dstDomain.getNumDimIds(),
+                        dstDomain.getNumDimAndSymbolIds(), &dstSymbolValues);
+
+  // Update value position map with dim values from src iteration domain.
+  updateValuePosMap(srcDimValues, /*isSrc=*/true, /*isDim=*/TRUE);
+  // Update value position map with dim values from dst iteration domain.
+  updateValuePosMap(dstDimValues, /*isSrc=*/false, /*isDim=*/TRUE);
+  // Update value position map with symbols from src iteration domain.
+  updateValuePosMap(srcSymbolValues, /*isSrc=*/true, /*isDim=*/FALSE);
+  // Update value position map with symbols from dst iteration domain.
+  updateValuePosMap(dstSymbolValues, /*isSrc=*/false, /*isDim=*/FALSE);
   // Update value position map with identifiers from src access function.
-  updateValuePosMap(srcAccessMap.getOperands(), /*isSrc=*/true);
+  updateValuePosMap(srcAccessMap.getOperands(), /*isSrc=*/true,
+                    /*isDim=*/UNKNOWN);
   // Update value position map with identifiers from dst access function.
-  updateValuePosMap(dstAccessMap.getOperands(), /*isSrc=*/false);
+  updateValuePosMap(dstAccessMap.getOperands(), /*isSrc=*/false,
+                    /*isDim=*/UNKNOWN);
 }
 
 // Sets up dependence constraints columns appropriately, in the format:
 // [src-dim-ids, dst-dim-ids, symbol-ids, local-ids, const_term]
 static void initDependenceConstraints(
     const FlatAffineConstraints &srcDomain,
     const FlatAffineConstraints &dstDomain, const AffineValueMap &srcAccessMap,
     const AffineValueMap &dstAccessMap, const ValuePositionMap &valuePosMap,
static void initDependenceConstraints(
   // Set values corresponding to dependence constraint identifiers.
   SmallVector<Value, 4> srcLoopIVs, dstLoopIVs;
   srcDomain.getIdValues(0, srcDomain.getNumDimIds(), &srcLoopIVs);
   dstDomain.getIdValues(0, dstDomain.getNumDimIds(), &dstLoopIVs);
 
   dependenceConstraints->setIdValues(0, srcLoopIVs.size(), srcLoopIVs);
   dependenceConstraints->setIdValues(
       srcLoopIVs.size(), srcLoopIVs.size() + dstLoopIVs.size(), dstLoopIVs);
 
-  // Set values for the symbolic identifier dimensions.
-  auto setSymbolIds = [&](ArrayRef<Value> values) {
+  // Set values for the symbolic identifier dimensions. `isSymbolDetermined`
+  // indicates whether we are certain that the `values` passed in are all
+  // symbols. If `isSymbolDetermined` is true, then we treat every Value in
+  // `values` as a symbol; otherwise, we let the function `isForInductionVar` to
+  // distinguish whether a Value in `values` is a symbol or not.
+  auto setSymbolIds = [&](ArrayRef<Value> values,
+                          bool isSymbolDetermined = true) {
     for (auto value : values) {
-      if (!isForInductionVar(value)) {
+      if (isSymbolDetermined || !isForInductionVar(value)) {
         assert(isValidSymbol(value) && "expected symbol");
         dependenceConstraints->setIdValue(valuePosMap.getSymPos(value), value);
       }
     }
   };
 
-  setSymbolIds(srcAccessMap.getOperands());
-  setSymbolIds(dstAccessMap.getOperands());
+  // We are uncertain about whether all operands in `srcAccessMap` and
+  // `dstAccessMap` are symbols, so we set `isSymbolDetermined` to false.
+  setSymbolIds(srcAccessMap.getOperands(), /*isSymbolDetermined=*/false);
+  setSymbolIds(dstAccessMap.getOperands(), /*isSymbolDetermined=*/false);
 
   SmallVector<Value, 8> srcSymbolValues, dstSymbolValues;
   srcDomain.getIdValues(srcDomain.getNumDimIds(),
                         srcDomain.getNumDimAndSymbolIds(), &srcSymbolValues);
   dstDomain.getIdValues(dstDomain.getNumDimIds(),
                         dstDomain.getNumDimAndSymbolIds(), &dstSymbolValues);
+  // Since we only take symbol Values out of `srcDomain` and `dstDomain`,
+  // `isSymbolDetermined` is kept to its default value: true.
   setSymbolIds(srcSymbolValues);
   setSymbolIds(dstSymbolValues);
 
   for (unsigned i = 0, e = dependenceConstraints->getNumDimAndSymbolIds();
        i < e; i++)
     assert(dependenceConstraints->getIds()[i].hasValue());
 }
 
if (commonLoops != nullptr)
       commonLoops->push_back(getForInductionVarOwner(srcDomain.getIdValue(i)));
     ++numCommonLoops;
   }
   if (commonLoops != nullptr)
     assert(commonLoops->size() == numCommonLoops);
   return numCommonLoops;
 }
 
-// Returns Block common to 'srcAccess.opInst' and 'dstAccess.opInst'.
+/// Returns Block common to 'srcAccess.opInst' and 'dstAccess.opInst'.
 static Block *getCommonBlock(const MemRefAccess &srcAccess,
                              const MemRefAccess &dstAccess,
                              const FlatAffineConstraints &srcDomain,
                              unsigned numCommonLoops) {
+  // Get the chain of ancestor blocks to the given `MemRefAccess` instance. The
+  // search terminates when either an op with the `AffineScope` trait or
+  // `endBlock` is reached.
+  auto getChainOfAncestorBlocks = [&](const MemRefAccess &access,
+                                      SmallVector<Block *, 4> &ancestorBlocks,
+                                      Block *endBlock = nullptr) {
+    Block *currBlock = access.opInst->getBlock();
+    // Loop terminates when the currBlock is nullptr or equals to the endBlock,
+    // or its parent operation holds an affine scope.
+    while (currBlock && currBlock != endBlock &&
+           !currBlock->getParentOp()->hasTrait<OpTrait::AffineScope>()) {
+      ancestorBlocks.push_back(currBlock);
+      currBlock = currBlock->getParentOp()->getBlock();
+    }
+  };
+
   if (numCommonLoops == 0) {
-    auto *block = srcAccess.opInst->getBlock();
+    Block *block = srcAccess.opInst->getBlock();
     while (!llvm::isa<FuncOp>(block->getParentOp())) {
       block = block->getParentOp()->getBlock();
     }
     return block;
   }
-  auto commonForValue = srcDomain.getIdValue(numCommonLoops - 1);
-  auto forOp = getForInductionVarOwner(commonForValue);
+  Value commonForIV = srcDomain.getIdValue(numCommonLoops - 1);
+  AffineForOp forOp = getForInductionVarOwner(commonForIV);
   assert(forOp && "commonForValue was not an induction variable");
-  return forOp.getBody();
+
+  // Find the closest common block including those in AffineIf.
+  SmallVector<Block *, 4> srcAncestorBlocks, dstAncestorBlocks;
+  getChainOfAncestorBlocks(srcAccess, srcAncestorBlocks, forOp.getBody());
+  getChainOfAncestorBlocks(dstAccess, dstAncestorBlocks, forOp.getBody());
+
+  Block *commonBlock = forOp.getBody();
+  for (int i = srcAncestorBlocks.size() - 1, j = dstAncestorBlocks.size() - 1;
+       i >= 0 && j >= 0 && srcAncestorBlocks[i] == dstAncestorBlocks[j];
+       i--, j--)
+    commonBlock = srcAncestorBlocks[i];
+
+  return commonBlock;
 }
 
 // Returns true if the ancestor operation of 'srcAccess' appears before the
 // ancestor operation of 'dstAccess' in the common ancestral block. Returns
 // false otherwise.
 // Note that because 'srcAccess' or 'dstAccess' may be nested in conditionals,
 // the function is named 'srcAppearsBeforeDstInCommonBlock'. Note that
 // 'numCommonLoops' is the number of contiguous surrounding outer loops.
DependenceResult mlir::checkMemrefAccessDependence(
   srcAccess.getAccessMap(&srcAccessMap);
 
   // Get composed access function for 'dstAccess'.
   AffineValueMap dstAccessMap;
   dstAccess.getAccessMap(&dstAccessMap);
 
   // Get iteration domain for the 'srcAccess' operation.
   FlatAffineConstraints srcDomain;
-  if (failed(getInstIndexSet(srcAccess.opInst, &srcDomain)))
+  if (failed(getOpIndexSet(srcAccess.opInst, &srcDomain)))
     return DependenceResult::Failure;
 
   // Get iteration domain for 'dstAccess' operation.
   FlatAffineConstraints dstDomain;
-  if (failed(getInstIndexSet(dstAccess.opInst, &dstDomain)))
+  if (failed(getOpIndexSet(dstAccess.opInst, &dstDomain)))
     return DependenceResult::Failure;
 
   // Return 'NoDependence' if loopDepth > numCommonLoops and if the ancestor
   // operation of 'srcAccess' does not properly dominate the ancestor
   // operation of 'dstAccess' in the same common operation block.
   // Note: this check is skipped if 'allowRAR' is true, because because RAR
   // deps can exist irrespective of lexicographic ordering b/w src and dst.
   unsigned numCommonLoops = getNumCommonLoops(srcDomain, dstDomain);
   assert(loopDepth <= numCommonLoops + 1);
   if (!allowRAR && loopDepth > numCommonLoops &&
       !srcAppearsBeforeDstInAncestralBlock(srcAccess, dstAccess, srcDomain,
                                            numCommonLoops)) {
     return DependenceResult::NoDependence;
   }
   // Build dim and symbol position maps for each access from access operand
   // Value to position in merged constraint system.
   ValuePositionMap valuePosMap;
   buildDimAndSymbolPositionMaps(srcDomain, dstDomain, srcAccessMap,
                                 dstAccessMap, &valuePosMap,
                                 dependenceConstraints);
-
   initDependenceConstraints(srcDomain, dstDomain, srcAccessMap, dstAccessMap,
                             valuePosMap, dependenceConstraints);
 
   assert(valuePosMap.getNumDims() ==
          srcDomain.getNumDimIds() + dstDomain.getNumDimIds());
 
   // Create memref access constraint by equating src/dst access functions.
   // Note that this check is conservative, and will fail in the future when

mlir/lib/Analysis/AffineStructures.cpp

if (forOp.hasConstantUpperBound()) {
     return success();
   }
   // Non-constant upper bound case.
   return addLowerOrUpperBound(pos, forOp.getUpperBoundMap(),
                               forOp.getUpperBoundOperands(),
                               /*eq=*/false, /*lower=*/false);
 }
 
+void FlatAffineConstraints::addAffineIfOpDomain(AffineIfOp ifOp) {
+  // Create the base constraints from the integer set attached to ifOp.
+  FlatAffineConstraints cst(ifOp.getIntegerSet());
+
+  // Bind ids in the constraints to ifOp operands.
+  SmallVector<Value, 4> operands = ifOp.getOperands();
+  cst.setIdValues(0, cst.getNumDimAndSymbolIds(), operands);
+
+  // Merge the constraints from ifOp to the current domain.
+  mergeAndAlignIdsWithOther(0, &cst);
+}
+
 // Searches for a constraint with a non-zero coefficient at 'colIdx' in
 // equality (isEq=true) or inequality (isEq=false) constraints.
 // Returns true and sets row found in search in 'rowIdx'.
 // Returns false otherwise.
 static bool findConstraintWithNonZeroAt(const FlatAffineConstraints &cst,
                                         unsigned colIdx, bool isEq,
                                         unsigned *rowIdx) {
   assert(colIdx < cst.getNumCols() && "position out of bounds");

mlir/lib/Analysis/Utils.cpp

while (currOp && ((currAffineForOp = dyn_cast<AffineForOp>(currOp)) ||
                     isa<AffineIfOp>(currOp))) {
     if (currAffineForOp)
       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) {
+  ops->clear();
+  Operation *currOp = op.getParentOp();
+
+  // Traverse up the hierarchy collecting all `affine.for` and `affine.if`
+  // operations.
+  while (currOp && (isa<AffineIfOp, AffineForOp>(currOp))) {
+    ops->push_back(currOp);
+    currOp = currOp->getParentOp();
+  }
+  std::reverse(ops->begin(), ops->end());
+}
+
 // Populates 'cst' with FlatAffineConstraints which represent slice bounds.
 LogicalResult
 ComputationSliceState::getAsConstraints(FlatAffineConstraints *cst) {
   assert(!lbOperands.empty());
   // Adds src 'ivs' as dimension identifiers in 'cst'.
   unsigned numDims = ivs.size();
   // Adds operands (dst ivs and symbols) as symbols in 'cst'.
   unsigned numSymbols = lbOperands[0].size();

mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp

mlir::tilePerfectlyNested(MutableArrayRef<AffineForOp> input,
 
   // Move the loop body of the original nest to the new one.
   moveLoopBody(origLoops.back(), innermostPointLoop);
 
   SmallVector<Value, 8> origLoopIVs;
   extractForInductionVars(input, &origLoopIVs);
 
   FlatAffineConstraints cst;
-  getIndexSet(input, &cst);
+  SmallVector<Operation *, 8> ops;
+  ops.reserve(input.size());
+  for (AffineForOp forOp : input)
+    ops.push_back(forOp);
+  getIndexSet(ops, &cst);
   if (!cst.isHyperRectangular(0, width)) {
     rootAffineForOp.emitError("tiled code generation unimplemented for the "
                               "non-hyperrectangular case");
     return failure();
   }
 
   constructTiledIndexSetHyperRect(origLoops, tiledLoops, tileSizes);
 

mlir/lib/Transforms/Utils/LoopUtils.cpp

 static AffineIfOp createSeparationCondition(MutableArrayRef<AffineForOp> loops,
                                             OpBuilder b) {
   if (loops.empty())
     return nullptr;
 
   auto *context = loops[0].getContext();
 
   FlatAffineConstraints cst;
-  getIndexSet(loops, &cst);
+  SmallVector<Operation *, 8> ops;
+  ops.reserve(loops.size());
+  for (AffineForOp forOp : loops)
+    ops.push_back(forOp);
+  getIndexSet(ops, &cst);
 
   // Remove constraints that are independent of these loop IVs.
   cst.removeIndependentConstraints(/*pos=*/0, /*num=*/loops.size());
 
   // Construct the constraint set representing the guard for full tiles. The
   // lower bound (and upper bound) corresponding to the full tile should be
   // larger (and resp. smaller) than any other lower (or upper bound).
   SmallVector<int64_t, 8> fullTileLb, fullTileUb;
createFullTiles(MutableArrayRef<AffineForOp> inputNest,
   FlatAffineConstraints cst;
   for (auto loop : inputNest) {
     // TODO: straightforward to generalize to a non-unit stride.
     if (loop.getStep() != 1) {
       LLVM_DEBUG(llvm::dbgs()
                  << "[tile separation] non-unit stride not implemented\n");
       return failure();
     }
-    getIndexSet({loop}, &cst);
+    SmallVector<Operation *, 1> loopOp{loop.getOperation()};
+    getIndexSet(loopOp, &cst);
     // We will mark everything other than this loop IV as symbol for getting a
     // pair of <lb, ub> with a constant difference.
     cst.setDimSymbolSeparation(cst.getNumDimAndSymbolIds() - 1);
     unsigned lbPos, ubPos;
     if (!cst.getConstantBoundOnDimSize(/*pos=*/0, /*lb=*/nullptr,
                                        /*lbDivisor=*/nullptr, /*ub=*/nullptr,
                                        &lbPos, &ubPos) ||
         lbPos == ubPos) {

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

affine.for %i0 = 0 to 10 {
     %1 = affine.load %0[%a0] : memref<100xf32>
     // expected-remark@above {{dependence from 1 to 0 at depth 1 = false}}
     // expected-remark@above {{dependence from 1 to 0 at depth 2 = false}}
     // expected-remark@above {{dependence from 1 to 1 at depth 1 = false}}
     // expected-remark@above {{dependence from 1 to 1 at depth 2 = false}}
   }
   return
 }
+
+// -----
+
+// Test the case that `affine.if` changes the domain for both load/store simultaneously.
+#set = affine_set<(d0): (d0 - 50 >= 0)>
+
+// CHECK-LABEL: func @test_affine_for_if_same_block() {
+func @test_affine_for_if_same_block() {
+  %0 = alloc() : memref<100xf32>
+  %cf7 = constant 7.0 : f32
+
+  affine.for %i0 = 0 to 100 {
+    affine.if #set(%i0) {
+      %1 = affine.load %0[%i0] : memref<100xf32>
+      // expected-remark@above {{dependence from 0 to 0 at depth 1 = false}}
+      // expected-remark@above {{dependence from 0 to 0 at depth 2 = false}}
+      // expected-remark@above {{dependence from 0 to 1 at depth 1 = false}}
+      // expected-remark@above {{dependence from 0 to 1 at depth 2 = true}}
+      affine.store %cf7, %0[%i0] : memref<100xf32>
+      // expected-remark@above {{dependence from 1 to 0 at depth 1 = false}}
+      // expected-remark@above {{dependence from 1 to 0 at depth 2 = false}}
+      // expected-remark@above {{dependence from 1 to 1 at depth 1 = false}}
+      // expected-remark@above {{dependence from 1 to 1 at depth 2 = false}}
+    }
+  }
+
+  return
+}
+
+// -----
+
+// Test the case that the domain that load/store access is completedly separated by `affine.if`.
+#set = affine_set<(d0): (d0 - 50 >= 0)>
+
+// CHECK-LABEL: func @test_affine_for_if_separated() {
+func @test_affine_for_if_separated() {
+  %0 = alloc() : memref<100xf32>
+  %cf7 = constant 7.0 : f32
+
+  affine.for %i0 = 0 to 10 {
+    affine.if #set(%i0) {
+      %1 = affine.load %0[%i0] : memref<100xf32>
+      // expected-remark@above {{dependence from 0 to 0 at depth 1 = false}}
+      // expected-remark@above {{dependence from 0 to 0 at depth 2 = false}}
+      // expected-remark@above {{dependence from 0 to 1 at depth 1 = false}}
+      // expected-remark@above {{dependence from 0 to 1 at depth 2 = false}}
+    } else {
+      affine.store %cf7, %0[%i0] : memref<100xf32>
+      // expected-remark@above {{dependence from 1 to 0 at depth 1 = false}}
+      // expected-remark@above {{dependence from 1 to 0 at depth 2 = false}}
+      // expected-remark@above {{dependence from 1 to 1 at depth 1 = false}}
+      // expected-remark@above {{dependence from 1 to 1 at depth 2 = false}}
+    }
+  }
+
+  return
+}
+
+// -----
+
+// Test the case that the domain that load/store access has non-empty union set.
+#set1 = affine_set<(d0): (  d0 - 25 >= 0)>
+#set2 = affine_set<(d0): (- d0 + 75 >= 0)>
+
+// CHECK-LABEL: func @test_affine_for_if_partially_joined() {
+func @test_affine_for_if_partially_joined() {
+  %0 = alloc() : memref<100xf32>
+  %cf7 = constant 7.0 : f32
+
+  affine.for %i0 = 0 to 100 {
+    affine.if #set1(%i0) {
+      %1 = affine.load %0[%i0] : memref<100xf32>
+      // expected-remark@above {{dependence from 0 to 0 at depth 1 = false}}
+      // expected-remark@above {{dependence from 0 to 0 at depth 2 = false}}
+      // expected-remark@above {{dependence from 0 to 1 at depth 1 = false}}
+      // expected-remark@above {{dependence from 0 to 1 at depth 2 = true}}
+    }
+    affine.if #set2(%i0) {
+      affine.store %cf7, %0[%i0] : memref<100xf32>
+      // expected-remark@above {{dependence from 1 to 0 at depth 1 = false}}
+      // expected-remark@above {{dependence from 1 to 0 at depth 2 = false}}
+      // expected-remark@above {{dependence from 1 to 1 at depth 1 = false}}
+      // expected-remark@above {{dependence from 1 to 1 at depth 2 = false}}
+    }
+  }
+
+  return
+}
+
+// -----
+
+// Test whether interleaved affine.for/affine.if can be properly handled.
+#set1 = affine_set<(d0): (d0 - 50 >= 0)>
+#set2 = affine_set<(d0, d1): (d0 - 75 >= 0, d1 - 50 >= 0)>
+
+// CHECK-LABEL: func @test_interleaved_affine_for_if() {
+func @test_interleaved_affine_for_if() {
+  %0 = alloc() : memref<100x100xf32>
+  %cf7 = constant 7.0 : f32
+
+  affine.for %i0 = 0 to 100 {
+    affine.if #set1(%i0) {
+      affine.for %i1 = 0 to 100 {
+        %1 = affine.load %0[%i0, %i1] : memref<100x100xf32>
+        // expected-remark@above {{dependence from 0 to 0 at depth 1 = false}}
+        // expected-remark@above {{dependence from 0 to 0 at depth 2 = false}}
+        // expected-remark@above {{dependence from 0 to 0 at depth 3 = false}}
+        // expected-remark@above {{dependence from 0 to 1 at depth 1 = false}}
+        // expected-remark@above {{dependence from 0 to 1 at depth 2 = false}}
+        // expected-remark@above {{dependence from 0 to 1 at depth 3 = true}}
+
+        affine.if #set2(%i0, %i1) {
+          affine.store %cf7, %0[%i0, %i1] : memref<100x100xf32>
+          // expected-remark@above {{dependence from 1 to 0 at depth 1 = false}}
+          // expected-remark@above {{dependence from 1 to 0 at depth 2 = false}}
+          // expected-remark@above {{dependence from 1 to 0 at depth 3 = false}}
+          // expected-remark@above {{dependence from 1 to 1 at depth 1 = false}}
+          // expected-remark@above {{dependence from 1 to 1 at depth 2 = false}}
+          // expected-remark@above {{dependence from 1 to 1 at depth 3 = false}}
+        }
+      }
+    }
+  }
+
+  return
+}
+
+// -----
+
+// Test whether symbols can be handled .
+#set1 = affine_set<(d0)[s0]: (  d0 - s0 floordiv 2 >= 0)>
+#set2 = affine_set<(d0):     (- d0 +            51 >= 0)>
+
+// CHECK-LABEL: func @test_interleaved_affine_for_if() {
+func @test_interleaved_affine_for_if() {
+  %0 = alloc() : memref<101xf32>
+  %c0 = constant 0 : index
+  %N = dim %0, %c0 : memref<101xf32>
+  %cf7 = constant 7.0 : f32
+
+  affine.for %i0 = 0 to 100 {
+    affine.if #set1(%i0)[%N] {
+      %1 = affine.load %0[%i0] : memref<101xf32>
+      // expected-remark@above {{dependence from 0 to 0 at depth 1 = false}}
+      // expected-remark@above {{dependence from 0 to 0 at depth 2 = false}}
+      // expected-remark@above {{dependence from 0 to 1 at depth 1 = false}}
+      // expected-remark@above {{dependence from 0 to 1 at depth 2 = true}}
+    }
+
+    affine.if #set2(%i0) {
+      affine.store %cf7, %0[%i0] : memref<101xf32>
+      // expected-remark@above {{dependence from 1 to 0 at depth 1 = false}}
+      // expected-remark@above {{dependence from 1 to 0 at depth 2 = false}}
+      // expected-remark@above {{dependence from 1 to 1 at depth 1 = false}}
+      // expected-remark@above {{dependence from 1 to 1 at depth 2 = false}}
+    }
+  }
+
+  return
+}