[llvm] [mlir] Number MLIR blocks; require GraphHasNodeNumbers in llvm::LoopInfoBase (PR #207617)
via llvm-commits
llvm-commits at lists.llvm.org
Sun Jul 5 15:21:55 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-llvm-support
Author: Fangrui Song (MaskRay)
<details>
<summary>Changes</summary>
LoopInfoBase keeps a DenseMap fallback for block types whose GraphTraits
has no getNumber(). MLIR's CFGLoopInfo is the only such in-tree user
(https://reviews.llvm.org/D147323).
Give mlir::Block a stable number, assigned when it is added to a region,
mirroring llvm::BasicBlock/llvm::Function: Block gains getNumber();
Region gains nextBlockNumber/blockNumberEpoch with getMaxBlockNumber(),
getBlockNumberEpoch() and renumberBlocks(); the block ilist traits
assign the number on add/transfer; and
GraphTraits<mlir::Block*>/<mlir::Region*> expose
getNumber/getMaxNumber/getNumberEpoch. This makes
GraphHasNodeNumbers<mlir::Block*> true (and also moves MLIR's dominator
tree, llvm::DominatorTreeBase<mlir::Block>, to the numbered path, since it
keys on the same trait).
With every in-tree user now numbered, require GraphHasNodeNumbers in
LoopInfoBase (static_assert) and drop the DenseMap fallback.
Aided by Claude Opus 4.8
---
Full diff: https://github.com/llvm/llvm-project/pull/207617.diff
6 Files Affected:
- (modified) llvm/include/llvm/Support/GenericLoopInfo.h (+29-49)
- (modified) llvm/include/llvm/Support/GenericLoopInfoImpl.h (+18-33)
- (modified) mlir/include/mlir/IR/Block.h (+18)
- (modified) mlir/include/mlir/IR/Region.h (+23)
- (modified) mlir/include/mlir/IR/RegionGraphTraits.h (+11)
- (modified) mlir/lib/IR/Region.cpp (+15-3)
``````````diff
diff --git a/llvm/include/llvm/Support/GenericLoopInfo.h b/llvm/include/llvm/Support/GenericLoopInfo.h
index 9b3bcd09e289d..faebdcb526e0d 100644
--- a/llvm/include/llvm/Support/GenericLoopInfo.h
+++ b/llvm/include/llvm/Support/GenericLoopInfo.h
@@ -524,10 +524,13 @@ raw_ostream &operator<<(raw_ostream &OS, const LoopBase<BlockT, LoopT> &Loop) {
///
template <class BlockT, class LoopT> class LoopInfoBase {
- // BBMap - Mapping of basic blocks to the inner most loop they occur in
- std::conditional_t<GraphHasNodeNumbers<const BlockT *>, SmallVector<LoopT *>,
- DenseMap<const BlockT *, LoopT *>>
- BBMap;
+ static_assert(GraphHasNodeNumbers<const BlockT *>,
+ "LoopInfo requires GraphTraits<BlockT *>::getNumber (see "
+ "GraphHasNodeNumbers)");
+
+ // Mapping of each block, indexed by its number, to the innermost loop it
+ // occurs in (or null).
+ SmallVector<LoopT *> BBMap;
using ParentT = decltype(std::declval<const BlockT *>()->getParent());
ParentT ParentPtr = nullptr;
@@ -615,25 +618,20 @@ template <class BlockT, class LoopT> class LoopInfoBase {
private:
/// Verify that used block numbers are still valid.
void verifyBlockNumberEpoch(ParentT BBParent) const {
- if constexpr (GraphHasNodeNumbers<BlockT *>) {
- assert(ParentPtr == BBParent &&
- "loop info queried with block of other function");
- assert(BlockNumberEpoch ==
- GraphTraits<ParentT>::getNumberEpoch(ParentPtr) &&
- "loop info used with outdated block numbers");
- }
+ assert(ParentPtr == BBParent &&
+ "loop info queried with block of other function");
+ assert(BlockNumberEpoch ==
+ GraphTraits<ParentT>::getNumberEpoch(ParentPtr) &&
+ "loop info used with outdated block numbers");
}
public:
/// Return the inner most loop that BB lives in. If a basic block is in no
/// loop (for example the entry node), null is returned.
LoopT *getLoopFor(const BlockT *BB) const {
- if constexpr (GraphHasNodeNumbers<const BlockT *>) {
- verifyBlockNumberEpoch(BB->getParent());
- unsigned Number = GraphTraits<const BlockT *>::getNumber(BB);
- return Number < BBMap.size() ? BBMap[Number] : nullptr;
- } else
- return BBMap.lookup(BB);
+ verifyBlockNumberEpoch(BB->getParent());
+ unsigned Number = GraphTraits<const BlockT *>::getNumber(BB);
+ return Number < BBMap.size() ? BBMap[Number] : nullptr;
}
/// Same as getLoopFor.
@@ -684,22 +682,14 @@ template <class BlockT, class LoopT> class LoopInfoBase {
/// This should be used by transformations that restructure the loop hierarchy
/// tree.
void changeLoopFor(const BlockT *BB, LoopT *L) {
- if constexpr (GraphHasNodeNumbers<const BlockT *>) {
- verifyBlockNumberEpoch(BB->getParent());
- unsigned Number = GraphTraits<const BlockT *>::getNumber(BB);
- if (Number >= BBMap.size()) {
- unsigned Max = GraphTraits<decltype(BB->getParent())>::getMaxNumber(
- BB->getParent());
- BBMap.resize(Number >= Max ? Number + 1 : Max);
- }
- BBMap[Number] = L;
- } else {
- if (!L) {
- BBMap.erase(BB);
- return;
- }
- BBMap[BB] = L;
+ verifyBlockNumberEpoch(BB->getParent());
+ unsigned Number = GraphTraits<const BlockT *>::getNumber(BB);
+ if (Number >= BBMap.size()) {
+ unsigned Max =
+ GraphTraits<decltype(BB->getParent())>::getMaxNumber(BB->getParent());
+ BBMap.resize(Number >= Max ? Number + 1 : Max);
}
+ BBMap[Number] = L;
}
/// Replace the specified loop in the top-level loops list with the indicated
@@ -722,24 +712,14 @@ template <class BlockT, class LoopT> class LoopInfoBase {
/// including all of the Loop objects it is nested in and our mapping from
/// BasicBlocks to loops.
void removeBlock(BlockT *BB) {
- if constexpr (GraphHasNodeNumbers<BlockT *>) {
- verifyBlockNumberEpoch(BB->getParent());
- unsigned Number = GraphTraits<BlockT *>::getNumber(BB);
- if (Number >= BBMap.size())
- return;
-
- for (LoopT *L = BBMap[Number]; L; L = L->getParentLoop())
- L->removeBlockFromLoop(BB);
- BBMap[Number] = nullptr;
- } else {
- auto I = BBMap.find(BB);
- if (I != BBMap.end()) {
- for (LoopT *L = I->second; L; L = L->getParentLoop())
- L->removeBlockFromLoop(BB);
+ verifyBlockNumberEpoch(BB->getParent());
+ unsigned Number = GraphTraits<BlockT *>::getNumber(BB);
+ if (Number >= BBMap.size())
+ return;
- BBMap.erase(I);
- }
- }
+ for (LoopT *L = BBMap[Number]; L; L = L->getParentLoop())
+ L->removeBlockFromLoop(BB);
+ BBMap[Number] = nullptr;
}
// Internals
diff --git a/llvm/include/llvm/Support/GenericLoopInfoImpl.h b/llvm/include/llvm/Support/GenericLoopInfoImpl.h
index 4666f4dac9cb6..445dfc05a76f5 100644
--- a/llvm/include/llvm/Support/GenericLoopInfoImpl.h
+++ b/llvm/include/llvm/Support/GenericLoopInfoImpl.h
@@ -578,12 +578,9 @@ template <class BlockT, class LoopT>
void LoopInfoBase<BlockT, LoopT>::analyze(const DomTreeBase<BlockT> &DomTree) {
// Postorder traversal of the dominator tree.
const DomTreeNodeBase<BlockT> *DomRoot = DomTree.getRootNode();
- if constexpr (GraphHasNodeNumbers<const BlockT *>) {
- ParentPtr = DomRoot->getBlock()->getParent();
- BlockNumberEpoch = GraphTraits<ParentT>::getNumberEpoch(ParentPtr);
- unsigned Max = GraphTraits<ParentT>::getMaxNumber(ParentPtr);
- BBMap.resize(Max);
- }
+ ParentPtr = DomRoot->getBlock()->getParent();
+ BlockNumberEpoch = GraphTraits<ParentT>::getNumberEpoch(ParentPtr);
+ BBMap.resize(GraphTraits<ParentT>::getMaxNumber(ParentPtr));
for (auto DomNode : post_order(DomRoot)) {
BlockT *Header = DomNode->getBlock();
@@ -755,33 +752,21 @@ void LoopInfoBase<BlockT, LoopT>::verify(
// Verify that blocks are mapped to valid loops.
#ifndef NDEBUG
- if constexpr (GraphHasNodeNumbers<const BlockT *>) {
- for (auto It : enumerate(BBMap)) {
- LoopT *L = It.value();
- unsigned Number = It.index();
- if (!L)
- continue;
- assert(Loops.count(L) && "orphaned loop");
- // We have no way to map block numbers back to blocks, so find it.
- auto BBIt = find_if(L->Blocks, [&Number](BlockT *BB) {
- return GraphTraits<BlockT *>::getNumber(BB) == Number;
- });
- BlockT *BB = BBIt != L->Blocks.end() ? *BBIt : nullptr;
- assert(BB && "orphaned block");
- for (LoopT *ChildLoop : *L)
- assert(!ChildLoop->contains(BB) &&
- "BBMap should point to the innermost loop containing BB");
- }
- } else {
- for (auto &Entry : BBMap) {
- const BlockT *BB = Entry.first;
- LoopT *L = Entry.second;
- assert(Loops.count(L) && "orphaned loop");
- assert(L->contains(BB) && "orphaned block");
- for (LoopT *ChildLoop : *L)
- assert(!ChildLoop->contains(BB) &&
- "BBMap should point to the innermost loop containing BB");
- }
+ for (auto It : enumerate(BBMap)) {
+ LoopT *L = It.value();
+ unsigned Number = It.index();
+ if (!L)
+ continue;
+ assert(Loops.count(L) && "orphaned loop");
+ // We have no way to map block numbers back to blocks, so find it.
+ auto BBIt = find_if(L->Blocks, [&Number](BlockT *BB) {
+ return GraphTraits<BlockT *>::getNumber(BB) == Number;
+ });
+ BlockT *BB = BBIt != L->Blocks.end() ? *BBIt : nullptr;
+ assert(BB && "orphaned block");
+ for (LoopT *ChildLoop : *L)
+ assert(!ChildLoop->contains(BB) &&
+ "BBMap should point to the innermost loop containing BB");
}
// Recompute LoopInfo to verify loops structure.
diff --git a/mlir/include/mlir/IR/Block.h b/mlir/include/mlir/IR/Block.h
index bb3484af2aaf7..5486b6afd5c88 100644
--- a/mlir/include/mlir/IR/Block.h
+++ b/mlir/include/mlir/IR/Block.h
@@ -52,6 +52,19 @@ class alignas(8) Block : public IRObjectWithUseList<BlockOperand>,
/// model.
Region *getParent() const;
+ /// Return a number uniquely identifying this block within its parent region.
+ /// The number is assigned when the block is added to a region and is stable
+ /// until the region's blocks are renumbered (see Region::renumberBlocks and
+ /// Region::getBlockNumberEpoch). Only valid for a block that is in a region.
+ ///
+ /// Unlike computeBlockNumber(), this is O(1) and stable; it exists so that
+ /// generic graph algorithms (e.g. LoopInfo, DominatorTree) can index blocks
+ /// by number.
+ unsigned getNumber() const {
+ assert(getParent() && "only blocks in a region have a valid number");
+ return blockNumber;
+ }
+
/// Returns the closest surrounding operation that contains this block.
Operation *getParentOp();
@@ -422,6 +435,10 @@ class alignas(8) Block : public IRObjectWithUseList<BlockOperand>,
/// the operations within this block have a valid ordering.
llvm::PointerIntPair<Region *, /*IntBits=*/1, bool> parentValidOpOrderPair;
+ /// Unique number of this block within its parent region, assigned when the
+ /// block is added to a region. See getNumber().
+ unsigned blockNumber = 0;
+
/// This is the list of operations in the block.
OpListType operations;
@@ -432,6 +449,7 @@ class alignas(8) Block : public IRObjectWithUseList<BlockOperand>,
void operator=(Block &) = delete;
friend struct llvm::ilist_traits<Block>;
+ friend class Region;
};
raw_ostream &operator<<(raw_ostream &, Block &);
diff --git a/mlir/include/mlir/IR/Region.h b/mlir/include/mlir/IR/Region.h
index 13b54991832cb..f379ec1d9f3b6 100644
--- a/mlir/include/mlir/IR/Region.h
+++ b/mlir/include/mlir/IR/Region.h
@@ -72,6 +72,22 @@ class Region {
return &Region::blocks;
}
+ //===--------------------------------------------------------------------===//
+ // Block numbering
+ //===--------------------------------------------------------------------===//
+
+ /// One past the largest block number handed out in this region; block numbers
+ /// lie in [0, getMaxBlockNumber()). See Block::getNumber().
+ unsigned getMaxBlockNumber() const { return nextBlockNumber; }
+
+ /// A counter bumped whenever the region's block numbers are reassigned
+ /// (renumberBlocks), so users indexing by block number can detect staleness.
+ unsigned getBlockNumberEpoch() const { return blockNumberEpoch; }
+
+ /// Renumber the blocks in this region 0..N in iteration order, compacting the
+ /// number space and bumping the block-number epoch.
+ void renumberBlocks();
+
//===--------------------------------------------------------------------===//
// Argument Handling
//===--------------------------------------------------------------------===//
@@ -343,6 +359,13 @@ class Region {
/// This is the object we are part of.
Operation *container = nullptr;
+
+ /// Next block number to hand out, and an epoch bumped when the block numbers
+ /// are reassigned by renumberBlocks(). See Block::getNumber().
+ unsigned nextBlockNumber = 0;
+ unsigned blockNumberEpoch = 0;
+
+ friend struct llvm::ilist_traits<Block>;
};
/// This class provides an abstraction over the different types of ranges over
diff --git a/mlir/include/mlir/IR/RegionGraphTraits.h b/mlir/include/mlir/IR/RegionGraphTraits.h
index d7d80d9b35806..570f303d370e4 100644
--- a/mlir/include/mlir/IR/RegionGraphTraits.h
+++ b/mlir/include/mlir/IR/RegionGraphTraits.h
@@ -31,6 +31,8 @@ struct GraphTraits<mlir::Block *> {
return node->succ_begin();
}
static ChildIteratorType child_end(NodeRef node) { return node->succ_end(); }
+
+ static unsigned getNumber(NodeRef node) { return node->getNumber(); }
};
template <>
@@ -63,6 +65,8 @@ struct GraphTraits<const mlir::Block *> {
static ChildIteratorType child_end(NodeRef node) {
return const_cast<mlir::Block *>(node)->succ_end();
}
+
+ static unsigned getNumber(NodeRef node) { return node->getNumber(); }
};
template <>
@@ -97,6 +101,13 @@ struct GraphTraits<mlir::Region *> : public GraphTraits<mlir::Block *> {
static nodes_iterator nodes_end(GraphType fn) {
return nodes_iterator(fn->end());
}
+
+ static unsigned getMaxNumber(GraphType region) {
+ return region->getMaxBlockNumber();
+ }
+ static unsigned getNumberEpoch(GraphType region) {
+ return region->getBlockNumberEpoch();
+ }
};
template <>
diff --git a/mlir/lib/IR/Region.cpp b/mlir/lib/IR/Region.cpp
index 15a941f380225..c8cd0a214204d 100644
--- a/mlir/lib/IR/Region.cpp
+++ b/mlir/lib/IR/Region.cpp
@@ -194,7 +194,9 @@ Region *llvm::ilist_traits<::mlir::Block>::getParentRegion() {
/// We keep the region pointer up to date.
void llvm::ilist_traits<::mlir::Block>::addNodeToList(Block *block) {
assert(!block->getParent() && "already in a region!");
- block->parentValidOpOrderPair.setPointer(getParentRegion());
+ Region *region = getParentRegion();
+ block->parentValidOpOrderPair.setPointer(region);
+ block->blockNumber = region->nextBlockNumber++;
}
/// This is a trait method invoked when an operation is removed from a
@@ -214,9 +216,19 @@ void llvm::ilist_traits<::mlir::Block>::transferNodesFromList(
if (curParent == otherList.getParentRegion())
return;
- // Update the 'parent' member of each Block.
- for (; first != last; ++first)
+ // Update the 'parent' member of each Block and give it a number in its new
+ // region.
+ for (; first != last; ++first) {
first->parentValidOpOrderPair.setPointer(curParent);
+ first->blockNumber = curParent->nextBlockNumber++;
+ }
+}
+
+void Region::renumberBlocks() {
+ nextBlockNumber = 0;
+ for (Block &block : *this)
+ block.blockNumber = nextBlockNumber++;
+ ++blockNumberEpoch;
}
//===----------------------------------------------------------------------===//
``````````
</details>
https://github.com/llvm/llvm-project/pull/207617
More information about the llvm-commits
mailing list