Remove findInstPositionInBlock from the Block api.

PiperOrigin-RevId: 232704766

mlir/include/mlir/IR/Block.h

@@ -188,11 +188,6 @@ public:
     return const_cast<Block *>(this)->front();
   }
 
-  /// Returns the instructions's position in this block or -1 if the instruction
-  /// is not present.
-  /// TODO: This is needlessly inefficient, and should not be API on Block.
-  int64_t findInstPositionInBlock(const Instruction &inst) const;
-
   /// Returns 'inst' if 'inst' lies in this block, or otherwise finds the
   /// ancestor instruction of 'inst' that lies in this block. Returns nullptr if
   /// the latter fails.

mlir/lib/Analysis/LoopAnalysis.cpp

@@ -320,21 +320,32 @@ bool mlir::isInstwiseShiftValid(ConstOpPointer<AffineForOp> forOp,
                                 ArrayRef<uint64_t> shifts) {
   auto *forBody = forOp->getBody();
   assert(shifts.size() == forBody->getInstructions().size());
-  unsigned s = 0;
-  for (const auto &inst : *forBody) {
+
+  // Work backwards over the body of the block so that we only need to iterator
+  // over the body once.
+  DenseMap<const Instruction *, uint64_t> forBodyShift;
+  for (auto it : llvm::enumerate(llvm::reverse(forBody->getInstructions()))) {
+    const auto &inst = it.value();
+
+    // Get the index of the current instruction, note that we are iterating in
+    // reverse so we need to fix it up.
+    size_t index = shifts.size() - it.index() - 1;
+
+    // Remember the shift of this instruction.
+    uint64_t shift = shifts[index];
+    forBodyShift.try_emplace(&inst, shift);
+
+    // Validate the results of this instruction if it were to be shifted.
     for (unsigned i = 0, e = inst.getNumResults(); i < e; ++i) {
       const Value *result = inst.getResult(i);
       for (const InstOperand &use : result->getUses()) {
         // If an ancestor instruction doesn't lie in the block of forOp,
-        // there is no shift to check. This is a naive way. If performance
-        // becomes an issue, a map can be used to store 'shifts' - to look up
-        // the shift for a instruction in constant time.
+        // there is no shift to check.
         if (auto *ancInst = forBody->findAncestorInstInBlock(*use.getOwner()))
-          if (shifts[s] != shifts[forBody->findInstPositionInBlock(*ancInst)])
+          if (shift != forBodyShift[ancInst])
             return false;
       }
     }
-    s++;
   }
   return true;
 }

mlir/lib/Analysis/Utils.cpp

@@ -318,8 +318,9 @@ static void findInstPosition(const Instruction *inst, Block *limitBlock,
                              SmallVectorImpl<unsigned> *positions) {
   const Block *block = inst->getBlock();
   while (block != limitBlock) {
-    int instPosInBlock = block->findInstPositionInBlock(*inst);
-    assert(instPosInBlock >= 0);
+    // FIXME: This algorithm is unnecessarily O(n) and should be improved to not
+    // rely on lineary scans.
+    int instPosInBlock = std::distance(block->begin(), inst->getIterator());
     positions->push_back(instPosInBlock);
     inst = block->getContainingInst();
     block = inst->getBlock();

mlir/lib/IR/Block.cpp

@@ -90,19 +90,6 @@ Instruction *Block::findAncestorInstInBlock(Instruction *inst) {
   return currInst;
 }
 
-/// Returns the instructions's position in this block or -1 if the instruction
-/// is not present.
-/// TODO: This is needlessly inefficient, and should not be API on Block.
-int64_t Block::findInstPositionInBlock(const Instruction &inst) const {
-  int64_t j = 0;
-  for (const auto &s : instructions) {
-    if (&s == &inst)
-      return j;
-    j++;
-  }
-  return -1;
-}
-
 /// This drops all operand uses from instructions within this block, which is
 /// an essential step in breaking cyclic dependences between references when
 /// they are to be deleted.