From ea6a1b4491ad7af4d1fc9fec133814713f948ee8 Mon Sep 17 00:00:00 2001 From: "Golubev, Andrey" Date: Mon, 2 Jun 2025 09:21:21 +0000 Subject: [PATCH 1/2] [mlir] Allow accessing DialectResourceBlobManager::blobMap 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. --- .../mlir/IR/DialectResourceBlobManager.h | 7 +- mlir/lib/IR/DialectResourceBlobManager.cpp | 8 ++ mlir/unittests/IR/BlobManagerTest.cpp | 74 +++++++++++++++++++ mlir/unittests/IR/CMakeLists.txt | 1 + 4 files changed, 89 insertions(+), 1 deletion(-) create mode 100644 mlir/unittests/IR/BlobManagerTest.cpp 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 &)> + accessor) const; + private: /// A mutex to protect access to the blob map. - llvm::sys::SmartRWMutex blobMapLock; + mutable llvm::sys::SmartRWMutex 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 &)> accessor) + const { + llvm::sys::SmartScopedReader 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 : tensor<1xi64> } : () -> () + +{-# + dialect_resources: { + builtin: { + blob1: "0x08000000ABCDABCDABCDABCE" + } + } +#-} +)mlir"; + +TEST(DialectResourceBlobManagerTest, Lookup) { + MLIRContext context; + context.loadDialect(); + + OwningOpRef m = parseSourceString(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(); + + OwningOpRef m = parseSourceString(moduleStr, &context); + ASSERT_TRUE(m); + + Block *block = m->getBody(); + auto &op = block->getOperations().front(); + auto resourceAttr = op.getAttrOfType("attr"); + ASSERT_NE(resourceAttr, nullptr); + + const auto &dialectManager = + resourceAttr.getRawHandle().getManagerInterface(&context); + + bool blobsArePresent = false; + dialectManager.getBlobManager().access( + [&](const llvm::StringMap + &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 + &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 From 6417a0262770cd72117086a898e29d12c1dddcac Mon Sep 17 00:00:00 2001 From: "Golubev, Andrey" Date: Fri, 20 Jun 2025 15:11:59 +0000 Subject: [PATCH 2/2] Rename access() to getBlobMap() --- mlir/include/mlir/IR/DialectResourceBlobManager.h | 6 +++--- mlir/lib/IR/DialectResourceBlobManager.cpp | 2 +- mlir/unittests/IR/BlobManagerTest.cpp | 6 +++--- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/mlir/include/mlir/IR/DialectResourceBlobManager.h b/mlir/include/mlir/IR/DialectResourceBlobManager.h index f51b73ef39b5f..6c30efde306e4 100644 --- a/mlir/include/mlir/IR/DialectResourceBlobManager.h +++ b/mlir/include/mlir/IR/DialectResourceBlobManager.h @@ -94,9 +94,9 @@ class DialectResourceBlobManager { } /// Provide access to all the registered blobs via a callable. During access - /// the blobs are guaranteed to remain unchanged. - void access(llvm::function_ref &)> - accessor) const; + /// the blob map is guaranteed to remain unchanged. + void getBlobMap(llvm::function_ref &)> + accessor) const; private: /// A mutex to protect access to the blob map. diff --git a/mlir/lib/IR/DialectResourceBlobManager.cpp b/mlir/lib/IR/DialectResourceBlobManager.cpp index 811e4459d518e..83cc1879241d1 100644 --- a/mlir/lib/IR/DialectResourceBlobManager.cpp +++ b/mlir/lib/IR/DialectResourceBlobManager.cpp @@ -64,7 +64,7 @@ auto DialectResourceBlobManager::insert(StringRef name, } while (true); } -void DialectResourceBlobManager::access( +void DialectResourceBlobManager::getBlobMap( llvm::function_ref &)> accessor) const { llvm::sys::SmartScopedReader reader(blobMapLock); diff --git a/mlir/unittests/IR/BlobManagerTest.cpp b/mlir/unittests/IR/BlobManagerTest.cpp index fe480a2189c5d..d82482ddb7936 100644 --- a/mlir/unittests/IR/BlobManagerTest.cpp +++ b/mlir/unittests/IR/BlobManagerTest.cpp @@ -40,7 +40,7 @@ TEST(DialectResourceBlobManagerTest, Lookup) { ASSERT_NE(dialectManager.getBlobManager().lookup("blob1"), nullptr); } -TEST(DialectResourceBlobManagerTest, Access) { +TEST(DialectResourceBlobManagerTest, GetBlobMap) { MLIRContext context; context.loadDialect(); @@ -56,7 +56,7 @@ TEST(DialectResourceBlobManagerTest, Access) { resourceAttr.getRawHandle().getManagerInterface(&context); bool blobsArePresent = false; - dialectManager.getBlobManager().access( + dialectManager.getBlobManager().getBlobMap( [&](const llvm::StringMap &blobMap) { blobsArePresent = blobMap.contains("blob1"); }); ASSERT_TRUE(blobsArePresent); @@ -65,7 +65,7 @@ TEST(DialectResourceBlobManagerTest, Access) { block->clear(); blobsArePresent = false; - dialectManager.getBlobManager().access( + dialectManager.getBlobManager().getBlobMap( [&](const llvm::StringMap &blobMap) { blobsArePresent = blobMap.contains("blob1"); }); ASSERT_TRUE(blobsArePresent);