Address post submit review comments for removing Block::findInstPositionInBlock.

PiperOrigin-RevId: 232713514

2 files changed, 6 insertions, 4 deletions

diff --git a/mlir/lib/Analysis/LoopAnalysis.cpp b/mlir/lib/Analysis/LoopAnalysis.cpp
index a4ec58ee6df..545735fd6fd 100644
--- a/mlir/lib/Analysis/LoopAnalysis.cpp
+++ b/mlir/lib/Analysis/LoopAnalysis.cpp
@@ -321,8 +321,8 @@ bool mlir::isInstwiseShiftValid(ConstOpPointer<AffineForOp> forOp,
   auto *forBody = forOp->getBody();
   assert(shifts.size() == forBody->getInstructions().size());
 
-  // Work backwards over the body of the block so that we only need to iterator
-  // over the body once.
+  // Work backwards over the body of the block so that the shift of a use's
+  // ancestor instruction in the block gets recorded before it's looked up.
   DenseMap<const Instruction *, uint64_t> forBodyShift;
   for (auto it : llvm::enumerate(llvm::reverse(forBody->getInstructions()))) {
     const auto &inst = it.value();
@@ -341,9 +341,11 @@ bool mlir::isInstwiseShiftValid(ConstOpPointer<AffineForOp> forOp,
       for (const InstOperand &use : result->getUses()) {
         // If an ancestor instruction doesn't lie in the block of forOp,
         // there is no shift to check.
-        if (auto *ancInst = forBody->findAncestorInstInBlock(*use.getOwner()))
+        if (auto *ancInst = forBody->findAncestorInstInBlock(*use.getOwner())) {
+          assert(forBodyShift.count(ancInst) > 0 && "ancestor expected in map");
           if (shift != forBodyShift[ancInst])
             return false;
+        }
       }
     }
   }
diff --git a/mlir/lib/Analysis/Utils.cpp b/mlir/lib/Analysis/Utils.cpp
index bdc5d19d0be..ae48e644a68 100644
--- a/mlir/lib/Analysis/Utils.cpp
+++ b/mlir/lib/Analysis/Utils.cpp
@@ -388,7 +388,7 @@ static void findInstPosition(const Instruction *inst, Block *limitBlock,
   const Block *block = inst->getBlock();
   while (block != limitBlock) {
     // FIXME: This algorithm is unnecessarily O(n) and should be improved to not
-    // rely on lineary scans.
+    // rely on linear scans.
     int instPosInBlock = std::distance(block->begin(), inst->getIterator());
     positions->push_back(instPosInBlock);
     inst = block->getContainingInst();