Skip to content

Conversation

snarang181
Copy link
Contributor

@snarang181 snarang181 commented Aug 29, 2025

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).

Copy link

github-actions bot commented Aug 29, 2025

✅ With the latest revision this PR passed the C/C++ code formatter.

@snarang181 snarang181 force-pushed the affine-analysis-opt branch from 27909ee to 9535e8b Compare August 29, 2025 14:12
@snarang181 snarang181 changed the title Amortize cost of block op index lookups [mlir][Affine][NFC] Amortize cost of block op index lookups Aug 29, 2025
@snarang181 snarang181 marked this pull request as ready for review August 29, 2025 14:18
@llvmbot
Copy link
Member

llvmbot commented Aug 29, 2025

@llvm/pr-subscribers-mlir

Author: Samarth Narang (snarang181)

Changes

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

@llvmbot
Copy link
Member

llvmbot commented Aug 29, 2025

@llvm/pr-subscribers-mlir-affine

Author: Samarth Narang (snarang181)

Changes

This patch addresses the O(n) scan as it relied on std::distance(block-&gt;begin(), op-&gt;getIterator()), which will perform a linear scan of the block each time.
This patch builds a per-block Operation* -&gt; 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 --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();
   }

@snarang181
Copy link
Contributor Author

Pinging reviewers to review, thanks.

Copy link
Member

@ftynse ftynse left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think this cache is effective. The cache is re-created every time we enter the findInstPosition function. The function traverses ancestor blocks for the given op once and then exists. There doesn't seem to be any reuse. FWIW, this is worse because the logic always goes until the end of each block and creates data structures, as opposed to going only until the (ancestor of the) op in question.

I suspect the comment may have either referred to using the cached position of the operation (

mutable unsigned orderIndex = 0;
) that is used by isBeforeInBlock, or a larger overhaul of the algorithm that doesn't need numeric positions of an operation in the block.

@snarang181 snarang181 closed this Oct 5, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants