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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Samarth Narang (snarang181)

<details>
<summary>Changes</summary>

This patch addresses the `O(n)` scan as it relied on `std::distance(block->begin(), op->getIterator())`, which will perform a linear scan of the block each time. 
This patch builds a per-block `Operation* -> index` map, which reduces subsequent lookups in the same block to `O(1)`. 

---
Full diff: https://github.com/llvm/llvm-project/pull/156027.diff


1 Files Affected:

- (modified) mlir/lib/Dialect/Affine/Analysis/Utils.cpp (+16-4) 


``````````diff
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();
   }

``````````

</details>


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