[Mlir-commits] [mlir] Cache name collision counter in DialectResourceBlobManager (PR #215771)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-core

Author: Kigyosi Alexandru (akigyosi)

<details>
<summary>Changes</summary>

This change improves repeated-name insertion in DialectResourceBlobManager by caching the next suffix counter per base name, avoiding repeated scans from 1 on every collision. It also adds unit tests that validate naming behavior for repeated inserts and direct-name collisions.

When many items share the same base name, insert() currently restarts suffix probing at 1 each time (blob_1, blob_2, ...), which causes unnecessary repeated lookup work and longer compilation time. Caching the next counter for each base name lets insertion resume from the last known position.

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


3 Files Affected:

- (modified) mlir/include/mlir/IR/DialectResourceBlobManager.h (+6) 
- (modified) mlir/lib/IR/DialectResourceBlobManager.cpp (+3-1) 
- (modified) mlir/unittests/IR/BlobManagerTest.cpp (+25) 


``````````diff
diff --git a/mlir/include/mlir/IR/DialectResourceBlobManager.h b/mlir/include/mlir/IR/DialectResourceBlobManager.h
index 6c30efde306e4..2e268c7c6a2c3 100644
--- a/mlir/include/mlir/IR/DialectResourceBlobManager.h
+++ b/mlir/include/mlir/IR/DialectResourceBlobManager.h
@@ -106,6 +106,12 @@ class DialectResourceBlobManager {
   /// allocations, so we can freely take references to the data without fear of
   /// invalidation during additional insertion/deletion.
   llvm::StringMap<BlobEntry> blobMap;
+
+  /// Cache of the next counter to try for a given base name.
+  /// When a requested name collides, insert() appends "_N".
+  /// Subsequent collisions on the same base name can resume the search
+  /// from the cached counter instead of restarting at 1.
+  llvm::StringMap<size_t> nameCounters;
 };
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/IR/DialectResourceBlobManager.cpp b/mlir/lib/IR/DialectResourceBlobManager.cpp
index 83cc1879241d1..a8af5214eea76 100644
--- a/mlir/lib/IR/DialectResourceBlobManager.cpp
+++ b/mlir/lib/IR/DialectResourceBlobManager.cpp
@@ -53,7 +53,9 @@ auto DialectResourceBlobManager::insert(StringRef name,
   // re-attempt insertion until we find one that is unique.
   llvm::SmallString<32> nameStorage(name);
   nameStorage.push_back('_');
-  size_t nameCounter = 1;
+
+  // Resume the numbering from the last counter used for this base name.
+  size_t &nameCounter = nameCounters.try_emplace(name, 1).first->second;
   do {
     Twine(nameCounter++).toVector(nameStorage);
 
diff --git a/mlir/unittests/IR/BlobManagerTest.cpp b/mlir/unittests/IR/BlobManagerTest.cpp
index d82482ddb7936..ea123b3bb46da 100644
--- a/mlir/unittests/IR/BlobManagerTest.cpp
+++ b/mlir/unittests/IR/BlobManagerTest.cpp
@@ -71,4 +71,29 @@ TEST(DialectResourceBlobManagerTest, GetBlobMap) {
   ASSERT_TRUE(blobsArePresent);
 }
 
+TEST(DialectResourceBlobManagerTest, InsertDisambiguatesNames) {
+  DialectResourceBlobManager manager;
+
+  // Repeated insertion of the same base name appends an increasing counter.
+  EXPECT_EQ(manager.insert("blob").getKey(), "blob");
+  EXPECT_EQ(manager.insert("blob").getKey(), "blob_1");
+  EXPECT_EQ(manager.insert("blob").getKey(), "blob_2");
+  EXPECT_EQ(manager.insert("blob").getKey(), "blob_3");
+
+  // A different base name is numbered independently.
+  EXPECT_EQ(manager.insert("other").getKey(), "other");
+  EXPECT_EQ(manager.insert("other").getKey(), "other_1");
+}
+
+TEST(DialectResourceBlobManagerTest, InsertSkipsNamesTakenDirectly) {
+  DialectResourceBlobManager manager;
+
+  EXPECT_EQ(manager.insert("blob").getKey(), "blob");
+  // Directly claim the name the next collision would otherwise produce.
+  EXPECT_EQ(manager.insert("blob_1").getKey(), "blob_1");
+  // The disambiguation search must skip the taken slot rather than reuse it,
+  // even though the resumed counter points at it first.
+  EXPECT_EQ(manager.insert("blob").getKey(), "blob_2");
+}
+
 } // end anonymous namespace

``````````

</details>


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