[mlir] Allow accessing DialectResourceBlobManager::blobMap #142352
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
@llvm/pr-subscribers-mlir-core Author: Andrei Golubev (andrey-golubev) ChangesAdd 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:
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
|
|
@llvm/pr-subscribers-mlir Author: Andrei Golubev (andrey-golubev) ChangesAdd 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:
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
|
|
cc @joker-eph @River707 (I think you're best reviewers here!) |
andrey-golubev
commented
Jun 2, 2025
| /// the blobs are guaranteed to remain unchanged. | ||
| void access(llvm::function_ref<void(const llvm::StringMap<BlobEntry> &)> | ||
| accessor) const; | ||
|
|
There was a problem hiding this comment.
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?
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.
CoTinker
requested review from
michaelmaitland,
joker-eph,
River707 and
matthias-springer
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
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.