[Mlir-commits] [mlir] Amortize cost of block op index lookups (PR #156027)

[Samarth Narang]

https://github.com/snarang181 updated https://github.com/llvm/llvm-project/pull/156027

>From 9535e8b3f8ba50457daab119e197c4aa0b2f9236 Mon Sep 17 00:00:00 2001
From: Samarth Narang <snarang at umass.edu>
Date: Fri, 29 Aug 2025 10:07:31 -0400
Subject: [PATCH] Amortize cost of block op index lookups

---
 mlir/lib/Dialect/Affine/Analysis/Utils.cpp | 20 ++++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index 99ea20bf13b49..c9e1ef9781af7 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -1442,16 +1442,28 @@ template LogicalResult
 mlir::affine::boundCheckLoadOrStoreOp(AffineWriteOpInterface storeOp,
                                       bool emitError);
 
+static inline unsigned getIndexInBlock(
+    Operation *op,
+    llvm::DenseMap<Block *, llvm::DenseMap<Operation *, unsigned>> &cache) {
+  Block *block = op->getBlock();
+  auto &blockMap = cache[block];
+  if (blockMap.empty()) {
+    unsigned idx = 0;
+    for (Operation &it : *block)
+      blockMap[&it] = idx++;
+  }
+  return blockMap.lookup(op);
+}
+
 // Returns in 'positions' the Block positions of 'op' in each ancestor
 // Block from the Block containing operation, stopping at 'limitBlock'.
 static void findInstPosition(Operation *op, Block *limitBlock,
                              SmallVectorImpl<unsigned> *positions) {
+  llvm::DenseMap<Block *, llvm::DenseMap<Operation *, unsigned>> indexCache;
+
   Block *block = op->getBlock();
   while (block != limitBlock) {
-    // FIXME: This algorithm is unnecessarily O(n) and should be improved to not
-    // rely on linear scans.
-    int instPosInBlock = std::distance(block->begin(), op->getIterator());
-    positions->push_back(instPosInBlock);
+    positions->push_back(getIndexInBlock(op, indexCache));
     op = block->getParentOp();
     block = op->getBlock();
   }