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

[Andrei Golubev]

================
@@ -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
----------------
andrey-golubev wrote:

honestly, I'm not sure. but seems to be intended behavior? from what I understand:
* resources are (almost) "raw" pointers (with optional deleters)
* 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).

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