[Mlir-commits] [mlir] 697bd97 - Give each mlir::Block a stable ID within its parent region (#207617)

[llvmlistbot at llvm.org]

Author: Fangrui Song
Date: 2026-07-08T19:34:48-07:00
New Revision: 697bd9704894691d6b0f40b3150d70047a6f181f

URL: https://github.com/llvm/llvm-project/commit/697bd9704894691d6b0f40b3150d70047a6f181f
DIFF: https://github.com/llvm/llvm-project/commit/697bd9704894691d6b0f40b3150d70047a6f181f.diff

LOG: Give each mlir::Block a stable ID within its parent region (#207617)

Assign each mlir::Block a stable ID within its parent region, mirroring
llvm::BasicBlock/llvm::Function (called ID, not number, since a block/op
number denotes position in MLIR): Block gains getBlockID() and reads -1u
while it has no parent region; Region gains nextBlockID with
getMaxBlockID() and getBlockIDEpoch(); the block ilist traits assign the
ID on add/transfer and invalidate it on removal; and
GraphTraits<mlir::Block*>/<mlir::Region*> expose
getNumber/getMaxNumber/getNumberEpoch. This makes
GraphHasNodeNumbers<mlir::Block*> true, moving MLIR's CFGLoopInfo and
dominator tree onto the number-indexed path. MLIR never renumbers
blocks,
so the epoch is a fixed 0.

Prerequisite for requiring GraphHasNodeNumbers in llvm::LoopInfoBase and
dropping its DenseMap fallback.

Aided by Claude Opus 4.8

Added: 
    

Modified: 
    mlir/include/mlir/IR/Block.h
    mlir/include/mlir/IR/Region.h
    mlir/include/mlir/IR/RegionGraphTraits.h
    mlir/lib/IR/Region.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/Block.h b/mlir/include/mlir/IR/Block.h
index bb3484af2aaf7..c84921db220af 100644
--- a/mlir/include/mlir/IR/Block.h
+++ b/mlir/include/mlir/IR/Block.h
@@ -52,6 +52,20 @@ class alignas(8) Block : public IRObjectWithUseList<BlockOperand>,
   /// model.
   Region *getParent() const;
 
+  /// Return an ID uniquely identifying this block within its parent region.
+  /// The ID is assigned when the block joins a region and reassigned when the
+  /// block is moved to a 
diff erent region; it is stable while the block stays
+  /// in a region, and removing a block leaves a hole (IDs are not reused). 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 ID.
+  unsigned getBlockID() const {
+    assert(getParent() && "only blocks in a region have a valid ID");
+    return blockID;
+  }
+
   /// Returns the closest surrounding operation that contains this block.
   Operation *getParentOp();
 
@@ -422,6 +436,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 ID of this block within its parent region, (re)assigned when the
+  /// block joins a region; -1u while the block has no parent. See getBlockID().
+  unsigned blockID = -1u;
+
   /// This is the list of operations in the block.
   OpListType operations;
 
@@ -432,6 +450,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..fa9d8f08d84bb 100644
--- a/mlir/include/mlir/IR/Region.h
+++ b/mlir/include/mlir/IR/Region.h
@@ -72,6 +72,19 @@ class Region {
     return &Region::blocks;
   }
 
+  //===--------------------------------------------------------------------===//
+  // Block numbering
+  //===--------------------------------------------------------------------===//
+
+  /// One past the largest block ID handed out in this region; block IDs lie in
+  /// [0, getMaxBlockID()). See Block::getBlockID().
+  unsigned getMaxBlockID() const { return nextBlockID; }
+
+  /// The block-ID epoch, part of the generic number-indexed graph contract
+  /// (LoopInfo, DominatorTree) for detecting stale IDs. MLIR never renumbers a
+  /// region's blocks, so this is a fixed 0. See Block::getBlockID().
+  unsigned getBlockIDEpoch() const { return 0; }
+
   //===--------------------------------------------------------------------===//
   // Argument Handling
   //===--------------------------------------------------------------------===//
@@ -343,6 +356,11 @@ class Region {
 
   /// This is the object we are part of.
   Operation *container = nullptr;
+
+  /// Next block ID to hand out. See Block::getBlockID().
+  unsigned nextBlockID = 0;
+
+  friend struct llvm::ilist_traits<Block>;
 };
 
 /// This class provides an abstraction over the 
diff erent types of ranges over

diff  --git a/mlir/include/mlir/IR/RegionGraphTraits.h b/mlir/include/mlir/IR/RegionGraphTraits.h
index d7d80d9b35806..bf21d9720169c 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->getBlockID(); }
 };
 
 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->getBlockID(); }
 };
 
 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->getMaxBlockID();
+  }
+  static unsigned getNumberEpoch(GraphType region) {
+    return region->getBlockIDEpoch();
+  }
 };
 
 template <>

diff  --git a/mlir/lib/IR/Region.cpp b/mlir/lib/IR/Region.cpp
index 15a941f380225..775485a3755e0 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->blockID = region->nextBlockID++;
 }
 
 /// This is a trait method invoked when an operation is removed from a
@@ -202,6 +204,8 @@ void llvm::ilist_traits<::mlir::Block>::addNodeToList(Block *block) {
 void llvm::ilist_traits<::mlir::Block>::removeNodeFromList(Block *block) {
   assert(block->getParent() && "not already in a region!");
   block->parentValidOpOrderPair.setPointer(nullptr);
+  // The ID is invalid until the block is added to a region again.
+  block->blockID = -1u;
 }
 
 /// This is a trait method invoked when an operation is moved from one block
@@ -214,9 +218,12 @@ 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 an ID in its new
+  // region.
+  for (; first != last; ++first) {
     first->parentValidOpOrderPair.setPointer(curParent);
+    first->blockID = curParent->nextBlockID++;
+  }
 }
 
 //===----------------------------------------------------------------------===//