[mlir] Allow accessing DialectResourceBlobManager::blobMap #142352

andrey-golubev · 2025-06-02T09:35:47Z

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.

llvmbot · 2025-06-02T09:36:24Z

@llvm/pr-subscribers-mlir-core

Author: Andrei Golubev (andrey-golubev)

Changes

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.

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

4 Files Affected:

(modified) mlir/include/mlir/IR/DialectResourceBlobManager.h (+6-1)
(modified) mlir/lib/IR/DialectResourceBlobManager.cpp (+8)
(added) mlir/unittests/IR/BlobManagerTest.cpp (+74)
(modified) mlir/unittests/IR/CMakeLists.txt (+1)

diff --git a/mlir/include/mlir/IR/DialectResourceBlobManager.h b/mlir/include/mlir/IR/DialectResourceBlobManager.h
index e3f32b7a9ab5f..f51b73ef39b5f 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 blobs are guaranteed to remain unchanged.
+  void access(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..811e4459d518e 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::access(
+    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..fe480a2189c5d
--- /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, Access) {
+  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().access(
+      [&](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().access(
+      [&](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

llvmbot · 2025-06-02T09:36:24Z

@llvm/pr-subscribers-mlir

Author: Andrei Golubev (andrey-golubev)

Changes

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.

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

4 Files Affected:

(modified) mlir/include/mlir/IR/DialectResourceBlobManager.h (+6-1)
(modified) mlir/lib/IR/DialectResourceBlobManager.cpp (+8)
(added) mlir/unittests/IR/BlobManagerTest.cpp (+74)
(modified) mlir/unittests/IR/CMakeLists.txt (+1)

diff --git a/mlir/include/mlir/IR/DialectResourceBlobManager.h b/mlir/include/mlir/IR/DialectResourceBlobManager.h
index e3f32b7a9ab5f..f51b73ef39b5f 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 blobs are guaranteed to remain unchanged.
+  void access(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..811e4459d518e 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::access(
+    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..fe480a2189c5d
--- /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, Access) {
+  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().access(
+      [&](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().access(
+      [&](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

andrey-golubev · 2025-06-02T09:36:59Z

cc @joker-eph @River707 (I think you're best reviewers here!)

andrey-golubev · 2025-06-02T09:39:05Z

mlir/include/mlir/IR/DialectResourceBlobManager.h

+  /// the blobs are guaranteed to remain unchanged.
+  void access(llvm::function_ref<void(const llvm::StringMap<BlobEntry> &)>
+                  accessor) const;
+


note: the API itself is a bit weird since we need to keep access protected by a mutex. without a larger redesign of this, it seems to be the only option?

andrey-golubev · 2025-06-12T10:42:28Z

ping @joker-eph @River707

andrey-golubev · 2025-06-19T07:35:49Z

@joker-eph @River707 (gently pinging again)

joker-eph · 2025-06-20T14:20:08Z

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 blobs are guaranteed to remain unchanged.


You can guarantee that the map isn't changed, but the blobs themselves aren't protected right?

hmm, good point. I guess better wording: "the blob map is guaranteed to remain unchanged".

if one is evil, they could hold a BlobEntry for a while and modify it, etc, but this is pretty much out of scope of this (any?) API.

rephrased a bit, thanks!

joker-eph · 2025-06-20T14:22:11Z

mlir/unittests/IR/BlobManagerTest.cpp

+              &blobMap) { blobsArePresent = blobMap.contains("blob1"); });
+  ASSERT_TRUE(blobsArePresent);
+
+  // remove operations that use resources - resources must still be accessible


Is this really a guarantee of the system? Aren't we allowed to somehow have a reference-count mechanism?

honestly, I'm not sure. but seems to be intended behavior? from what I understand:

resources are "raw" (at best unique when with deleter) pointers

blob - underlying data behind resources - are stored into a dialect manager (lifetime == MLIR context's lifetime)

resources are used by operations but as the resources themselves are attributes (mlir::DenseResourceElementsAttr), they are not "owned" by operations

thus, what we effectively get is: dense_resource<> is never really deleted unless one explicitly does something at the dialect (aka blob) manager level. which is fine by me, to be honest.

side-note: i was experimenting with custom deletion on top of dialect resources before (e.g. via a manual ref-count) but I guess ideal solution is to switch to properties and have operation-driven auto-ref-count semantics (might require some resource redesign as well though).

joker-eph · 2025-06-20T14:23:24Z

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 blobs are guaranteed to remain unchanged.
+  void access(llvm::function_ref<void(const llvm::StringMap<BlobEntry> &)>


Suggested change

void access(llvm::function_ref<void(const llvm::StringMap<BlobEntry> &)>

void getBlobMap(llvm::function_ref<void(const llvm::StringMap<BlobEntry> &)>

"access" is a bit unusual, ultimately this is an accessor every if you're using a callback for purposes of mutex handling

right. I was thinking that get is supposed to return something. if you think this is a better name, though, i have nothing against. "access" does sound alien, agree.

changed to getBlobMap

joker-eph

I'm not totally sure about the invariants we want around all this, but this would be a "low-level API / you must know what you're doing" kind of API.

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.

andrey-golubev · 2025-06-23T07:12:43Z

@joker-eph updated according to your suggestions (and also rebased). do you mind merging (i have no rights)?

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.

llvmbot added mlir:core MLIR Core Infrastructure mlir labels Jun 2, 2025

andrey-golubev commented Jun 2, 2025

View reviewed changes

CoTinker requested review from michaelmaitland, joker-eph, River707 and matthias-springer June 6, 2025 02:26

joker-eph reviewed Jun 20, 2025

View reviewed changes

joker-eph approved these changes Jun 20, 2025

View reviewed changes

andrey-golubev added 2 commits June 20, 2025 15:02

Rename access() to getBlobMap()

6417a02

andrey-golubev force-pushed the blob_map_access branch from 2b76715 to 6417a02 Compare June 20, 2025 15:26

joker-eph merged commit 529662a into llvm:main Jun 23, 2025
7 checks passed

andrey-golubev deleted the blob_map_access branch June 23, 2025 07:52

andrey-golubev mentioned this pull request Jun 23, 2025

[mlir] Allow accessing DialectResourceBlobManager::blobMap (cherry-pick) intel/npu-plugin-llvm#142

Merged

	void access(llvm::function_ref<void(const llvm::StringMap<BlobEntry> &)>
	void getBlobMap(llvm::function_ref<void(const llvm::StringMap<BlobEntry> &)>

[mlir] Allow accessing DialectResourceBlobManager::blobMap #142352

[mlir] Allow accessing DialectResourceBlobManager::blobMap #142352

Uh oh!

Conversation

andrey-golubev commented Jun 2, 2025

Uh oh!

llvmbot commented Jun 2, 2025

Uh oh!

llvmbot commented Jun 2, 2025

Uh oh!

andrey-golubev commented Jun 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

andrey-golubev commented Jun 12, 2025

Uh oh!

andrey-golubev commented Jun 19, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

andrey-golubev Jun 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

joker-eph left a comment

Choose a reason for hiding this comment

Uh oh!

andrey-golubev commented Jun 23, 2025

Uh oh!

Uh oh!

Uh oh!

andrey-golubev commented Jun 2, 2025 •

edited

Loading

andrey-golubev Jun 20, 2025 •

edited

Loading