[Mlir-commits] [mlir] 529662a - [mlir] Allow accessing DialectResourceBlobManager::blobMap (#142352)

[llvmlistbot at llvm.org]

Author: Andrei Golubev
Date: 2025-06-23T09:50:28+02:00
New Revision: 529662a6b5f1eec832d35786b6c49da1d21513e8

URL: https://github.com/llvm/llvm-project/commit/529662a6b5f1eec832d35786b6c49da1d21513e8
DIFF: https://github.com/llvm/llvm-project/commit/529662a6b5f1eec832d35786b6c49da1d21513e8.diff

LOG: [mlir] Allow accessing DialectResourceBlobManager::blobMap (#142352)

Add a new API to access all blobs that are stored in the blob manager.
The main purpose (as of now) is to allow users of dialect resources to
iterate over all blobs, especially when the blobs are no longer used in
IR (e.g. the operation that uses the blob is deleted) and thus cannot be
easily accessed without manual tracking of keys.

Added: 
    mlir/unittests/IR/BlobManagerTest.cpp

Modified: 
    mlir/include/mlir/IR/DialectResourceBlobManager.h
    mlir/lib/IR/DialectResourceBlobManager.cpp
    mlir/unittests/IR/CMakeLists.txt

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/DialectResourceBlobManager.h b/mlir/include/mlir/IR/DialectResourceBlobManager.h
index e3f32b7a9ab5f..6c30efde306e4 100644
--- a/mlir/include/mlir/IR/DialectResourceBlobManager.h
+++ b/mlir/include/mlir/IR/DialectResourceBlobManager.h
@@ -93,9 +93,14 @@ class DialectResourceBlobManager {
     return HandleT(&entry, dialect);
   }
 
+  /// Provide access to all the registered blobs via a callable. During access
+  /// the blob map is guaranteed to remain unchanged.
+  void getBlobMap(llvm::function_ref<void(const llvm::StringMap<BlobEntry> &)>
+                      accessor) const;
+
 private:
   /// A mutex to protect access to the blob map.
-  llvm::sys::SmartRWMutex<true> blobMapLock;
+  mutable llvm::sys::SmartRWMutex<true> blobMapLock;
 
   /// The internal map of tracked blobs. StringMap stores entries in distinct
   /// allocations, so we can freely take references to the data without fear of

diff  --git a/mlir/lib/IR/DialectResourceBlobManager.cpp b/mlir/lib/IR/DialectResourceBlobManager.cpp
index b83b31e30ef1f..83cc1879241d1 100644
--- a/mlir/lib/IR/DialectResourceBlobManager.cpp
+++ b/mlir/lib/IR/DialectResourceBlobManager.cpp
@@ -63,3 +63,11 @@ auto DialectResourceBlobManager::insert(StringRef name,
     nameStorage.resize(name.size() + 1);
   } while (true);
 }
+
+void DialectResourceBlobManager::getBlobMap(
+    llvm::function_ref<void(const llvm::StringMap<BlobEntry> &)> accessor)
+    const {
+  llvm::sys::SmartScopedReader<true> reader(blobMapLock);
+
+  accessor(blobMap);
+}

diff  --git a/mlir/unittests/IR/BlobManagerTest.cpp b/mlir/unittests/IR/BlobManagerTest.cpp
new file mode 100644
index 0000000000000..d82482ddb7936
--- /dev/null
+++ b/mlir/unittests/IR/BlobManagerTest.cpp
@@ -0,0 +1,74 @@
+//===- mlir/unittest/IR/BlobManagerTest.cpp - Blob management unit tests --===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "../../test/lib/Dialect/Test/TestDialect.h"
+#include "mlir/IR/DialectResourceBlobManager.h"
+#include "mlir/Parser/Parser.h"
+
+#include "gtest/gtest.h"
+
+using namespace mlir;
+
+namespace {
+
+StringLiteral moduleStr = R"mlir(
+"test.use1"() {attr = dense_resource<blob1> : tensor<1xi64> } : () -> ()
+
+{-#
+    dialect_resources: {
+    builtin: {
+        blob1: "0x08000000ABCDABCDABCDABCE"
+    }
+    }
+#-}
+)mlir";
+
+TEST(DialectResourceBlobManagerTest, Lookup) {
+  MLIRContext context;
+  context.loadDialect<test::TestDialect>();
+
+  OwningOpRef<ModuleOp> m = parseSourceString<ModuleOp>(moduleStr, &context);
+  ASSERT_TRUE(m);
+
+  const auto &dialectManager =
+      mlir::DenseResourceElementsHandle::getManagerInterface(&context);
+  ASSERT_NE(dialectManager.getBlobManager().lookup("blob1"), nullptr);
+}
+
+TEST(DialectResourceBlobManagerTest, GetBlobMap) {
+  MLIRContext context;
+  context.loadDialect<test::TestDialect>();
+
+  OwningOpRef<ModuleOp> m = parseSourceString<ModuleOp>(moduleStr, &context);
+  ASSERT_TRUE(m);
+
+  Block *block = m->getBody();
+  auto &op = block->getOperations().front();
+  auto resourceAttr = op.getAttrOfType<DenseResourceElementsAttr>("attr");
+  ASSERT_NE(resourceAttr, nullptr);
+
+  const auto &dialectManager =
+      resourceAttr.getRawHandle().getManagerInterface(&context);
+
+  bool blobsArePresent = false;
+  dialectManager.getBlobManager().getBlobMap(
+      [&](const llvm::StringMap<DialectResourceBlobManager::BlobEntry>
+              &blobMap) { blobsArePresent = blobMap.contains("blob1"); });
+  ASSERT_TRUE(blobsArePresent);
+
+  // remove operations that use resources - resources must still be accessible
+  block->clear();
+
+  blobsArePresent = false;
+  dialectManager.getBlobManager().getBlobMap(
+      [&](const llvm::StringMap<DialectResourceBlobManager::BlobEntry>
+              &blobMap) { blobsArePresent = blobMap.contains("blob1"); });
+  ASSERT_TRUE(blobsArePresent);
+}
+
+} // end anonymous namespace

diff  --git a/mlir/unittests/IR/CMakeLists.txt b/mlir/unittests/IR/CMakeLists.txt
index 9ab6029c3480d..7700644864570 100644
--- a/mlir/unittests/IR/CMakeLists.txt
+++ b/mlir/unittests/IR/CMakeLists.txt
@@ -18,6 +18,7 @@ add_mlir_unittest(MLIRIRTests
   TypeAttrNamesTest.cpp
   OpPropertiesTest.cpp
   ValueTest.cpp
+  BlobManagerTest.cpp
 
   DEPENDS
   MLIRTestInterfaceIncGen