[MLIR][NFC] Refactor affine analysis helpers

[bondhugula]

Summary

Refactor affine analysis helpers: the existing ones were using
unnecessary or improperly named arguments. NFC.

[bondhugula]

Mark commit summary NFC.

[bondhugula]

Fix/improve comments.

[ftynse]

Thanks!

Finding the common block sounds like a general utility. I think I have seen or maybe even written similar code elsewhere, but can't pinpoint where exactly. If you can, you are very welcome to move this to be a generic utility somewhere in lib/IR.

[bondhugula]

In D131557#3712177, @ftynse wrote:

Thanks!

Finding the common block sounds like a general utility. I think I have seen or maybe even written similar code elsewhere, but can't pinpoint where exactly. If you can, you are very welcome to move this to be a generic utility somewhere in lib/IR.

Thanks. I'll looks for such usage and do this in a follow-up PR.

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

if (commonLoops != nullptr)
       commonLoops->push_back(getForInductionVarOwner(srcDomain.getValue(i)));
     ++numCommonLoops;
   }
   if (commonLoops != nullptr)
     assert(commonLoops->size() == numCommonLoops);
   return numCommonLoops;
 }
 
-/// Returns Block common to 'srcAccess.opInst' and 'dstAccess.opInst'.
-static Block *getCommonBlock(const MemRefAccess &srcAccess,
-                             const MemRefAccess &dstAccess,
-                             const FlatAffineValueConstraints &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 &&
+/// Returns the closest surrounding block common to `opA` and `opB`. `opA` and
+/// `opB` should be in the same affine scope and thus such a block is guaranteed
+/// to exist.
+static Block *getCommonBlock(Operation *opA, Operation *opB) {
+  // Get the chain of ancestor blocks for the given `MemRefAccess` instance. The
+  // chain extends up to and includnig an op that starts an affine scope.
+  auto getChainOfAncestorBlocks =
+      [&](Operation *op, SmallVectorImpl<Block *> &ancestorBlocks) {
+        Block *currBlock = op->getBlock();
+        // Loop terminates when the currBlock is nullptr or its parent operation
+        // holds an affine scope.
+        while (currBlock &&
                !currBlock->getParentOp()->hasTrait<OpTrait::AffineScope>()) {
           ancestorBlocks.push_back(currBlock);
           currBlock = currBlock->getParentOp()->getBlock();
         }
+        assert(currBlock &&
+               "parent op starting an affine scope is always expected");
+        ancestorBlocks.push_back(currBlock);
       };
 
-  if (numCommonLoops == 0) {
-    Block *block = srcAccess.opInst->getBlock();
-    while (!block->getParentOp()->hasTrait<OpTrait::AffineScope>())
-      block = block->getParentOp()->getBlock();
-    return block;
-  }
-  Value commonForIV = srcDomain.getValue(numCommonLoops - 1);
-  AffineForOp forOp = getForInductionVarOwner(commonForIV);
-  assert(forOp && "commonForValue was not an induction variable");
-
   // Find the closest common block including those in AffineIf.
   SmallVector<Block *, 4> srcAncestorBlocks, dstAncestorBlocks;
-  getChainOfAncestorBlocks(srcAccess, srcAncestorBlocks, forOp.getBody());
-  getChainOfAncestorBlocks(dstAccess, dstAncestorBlocks, forOp.getBody());
+  getChainOfAncestorBlocks(opA, srcAncestorBlocks);
+  getChainOfAncestorBlocks(opB, dstAncestorBlocks);
 
-  Block *commonBlock = forOp.getBody();
+  Block *commonBlock = nullptr;
   for (int i = srcAncestorBlocks.size() - 1, j = dstAncestorBlocks.size() - 1;
        i >= 0 && j >= 0 && srcAncestorBlocks[i] == dstAncestorBlocks[j];
        i--, j--)
     commonBlock = srcAncestorBlocks[i];
-
+  // This is guaranteed since both ops are from the same affine scope.
+  assert(commonBlock && "ops expected to have a common surrounding block");
   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.
-static bool srcAppearsBeforeDstInAncestralBlock(
-    const MemRefAccess &srcAccess, const MemRefAccess &dstAccess,
-    const FlatAffineValueConstraints &srcDomain, unsigned numCommonLoops) {
+/// Returns true if the ancestor operation of 'srcAccess' appears before the
+/// ancestor operation of 'dstAccess' in their common ancestral block. The
+/// operations for `srcAccess` and `dstAccess` are expected to be in the same
+/// affine scope.
+static bool srcAppearsBeforeDstInAncestralBlock(const MemRefAccess &srcAccess,
+                                                const MemRefAccess &dstAccess) {
   // Get Block common to 'srcAccess.opInst' and 'dstAccess.opInst'.
-  auto *commonBlock =
-      getCommonBlock(srcAccess, dstAccess, srcDomain, numCommonLoops);
+  auto *commonBlock = getCommonBlock(srcAccess.opInst, dstAccess.opInst);
   // Check the dominance relationship between the respective ancestors of the
   // src and dst in the Block of the innermost among the common loops.
   auto *srcInst = commonBlock->findAncestorOpInBlock(*srcAccess.opInst);
-  assert(srcInst != nullptr);
+  assert(srcInst && "src access op must lie in common block");
   auto *dstInst = commonBlock->findAncestorOpInBlock(*dstAccess.opInst);
-  assert(dstInst != nullptr);
+  assert(dstInst && "dest access op must lie in common block");
 
   // Determine whether dstInst comes after srcInst.
   return srcInst->isBeforeInBlock(dstInst);
 }
 
 // Adds ordering constraints to 'dependenceDomain' based on number of loops
 // common to 'src/dstDomain' and requested 'loopDepth'.
 // Note that 'loopDepth' cannot exceed the number of common loops plus one.
DependenceResult mlir::checkMemrefAccessDependence(
   // 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)) {
+      !srcAppearsBeforeDstInAncestralBlock(srcAccess, dstAccess)) {
     return DependenceResult::NoDependence;
   }
 
   // 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();