75 lines · cpp
1//===- mlir/unittest/IR/BlobManagerTest.cpp - Blob management unit tests --===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "../../test/lib/Dialect/Test/TestDialect.h"10#include "mlir/IR/DialectResourceBlobManager.h"11#include "mlir/Parser/Parser.h"12 13#include "gtest/gtest.h"14 15using namespace mlir;16 17namespace {18 19StringLiteral moduleStr = R"mlir(20"test.use1"() {attr = dense_resource<blob1> : tensor<1xi64> } : () -> ()21 22{-#23 dialect_resources: {24 builtin: {25 blob1: "0x08000000ABCDABCDABCDABCE"26 }27 }28#-}29)mlir";30 31TEST(DialectResourceBlobManagerTest, Lookup) {32 MLIRContext context;33 context.loadDialect<test::TestDialect>();34 35 OwningOpRef<ModuleOp> m = parseSourceString<ModuleOp>(moduleStr, &context);36 ASSERT_TRUE(m);37 38 const auto &dialectManager =39 mlir::DenseResourceElementsHandle::getManagerInterface(&context);40 ASSERT_NE(dialectManager.getBlobManager().lookup("blob1"), nullptr);41}42 43TEST(DialectResourceBlobManagerTest, GetBlobMap) {44 MLIRContext context;45 context.loadDialect<test::TestDialect>();46 47 OwningOpRef<ModuleOp> m = parseSourceString<ModuleOp>(moduleStr, &context);48 ASSERT_TRUE(m);49 50 Block *block = m->getBody();51 auto &op = block->getOperations().front();52 auto resourceAttr = op.getAttrOfType<DenseResourceElementsAttr>("attr");53 ASSERT_NE(resourceAttr, nullptr);54 55 const auto &dialectManager =56 resourceAttr.getRawHandle().getManagerInterface(&context);57 58 bool blobsArePresent = false;59 dialectManager.getBlobManager().getBlobMap(60 [&](const llvm::StringMap<DialectResourceBlobManager::BlobEntry>61 &blobMap) { blobsArePresent = blobMap.contains("blob1"); });62 ASSERT_TRUE(blobsArePresent);63 64 // remove operations that use resources - resources must still be accessible65 block->clear();66 67 blobsArePresent = false;68 dialectManager.getBlobManager().getBlobMap(69 [&](const llvm::StringMap<DialectResourceBlobManager::BlobEntry>70 &blobMap) { blobsArePresent = blobMap.contains("blob1"); });71 ASSERT_TRUE(blobsArePresent);72}73 74} // end anonymous namespace75