1//===- AdaptorTest.cpp - Adaptor 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-exception
6//
7//===----------------------------------------------------------------------===//
8
9#include "mlir/Bytecode/BytecodeReader.h"
10#include "mlir/Bytecode/BytecodeWriter.h"
11#include "mlir/IR/AsmState.h"
12#include "mlir/IR/BuiltinAttributes.h"
13#include "mlir/IR/OpImplementation.h"
14#include "mlir/IR/OwningOpRef.h"
15#include "mlir/Parser/Parser.h"
16
17#include "llvm/ADT/StringRef.h"
18#include "llvm/Support/Endian.h"
19#include "llvm/Support/MemoryBufferRef.h"
20#include "gmock/gmock.h"
21#include "gtest/gtest.h"
22
23using namespace llvm;
24using namespace mlir;
25
26StringLiteral irWithResources = R"(
27module @TestDialectResources attributes {
28 bytecode.test = dense_resource<resource> : tensor<4xi32>
29} {}
30{-#
31 dialect_resources: {
32 builtin: {
33 resource: "0x2000000001000000020000000300000004000000",
34 resource_2: "0x2000000001000000020000000300000004000000"
35 }
36 }
37#-}
38)";
39
40TEST(Bytecode, MultiModuleWithResource) {
41 MLIRContext context;
42 Builder builder(&context);
43 ParserConfig parseConfig(&context);
44 OwningOpRef<Operation *> module =
45 parseSourceString<Operation *>(sourceStr: irWithResources, config: parseConfig);
46 ASSERT_TRUE(module);
47
48 // Write the module to bytecode
49 std::string buffer;
50 llvm::raw_string_ostream ostream(buffer);
51 ASSERT_TRUE(succeeded(writeBytecodeToFile(module.get(), ostream)));
52 ostream.flush();
53
54 // Create copy of buffer which is aligned to requested resource alignment.
55 constexpr size_t kAlignment = 0x20;
56 size_t bufferSize = buffer.size();
57 buffer.reserve(res: bufferSize + kAlignment - 1);
58 size_t pad = ~(uintptr_t)buffer.data() + 1 & kAlignment - 1;
59 buffer.insert(pos: 0, n: pad, c: ' ');
60 StringRef alignedBuffer(buffer.data() + pad, bufferSize);
61
62 // Parse it back
63 OwningOpRef<Operation *> roundTripModule =
64 parseSourceString<Operation *>(sourceStr: alignedBuffer, config: parseConfig);
65 ASSERT_TRUE(roundTripModule);
66
67 // FIXME: Parsing external resources does not work on big-endian
68 // platforms currently.
69 if (llvm::endianness::native == llvm::endianness::big)
70 GTEST_SKIP();
71
72 // Try to see if we have a valid resource in the parsed module.
73 auto checkResourceAttribute = [&](Operation *op) {
74 Attribute attr = roundTripModule->getDiscardableAttr(name: "bytecode.test");
75 ASSERT_TRUE(attr);
76 auto denseResourceAttr = dyn_cast<DenseI32ResourceElementsAttr>(attr);
77 ASSERT_TRUE(denseResourceAttr);
78 std::optional<ArrayRef<int32_t>> attrData =
79 denseResourceAttr.tryGetAsArrayRef();
80 ASSERT_TRUE(attrData.has_value());
81 ASSERT_EQ(attrData->size(), static_cast<size_t>(4));
82 EXPECT_EQ((*attrData)[0], 1);
83 EXPECT_EQ((*attrData)[1], 2);
84 EXPECT_EQ((*attrData)[2], 3);
85 EXPECT_EQ((*attrData)[3], 4);
86 };
87
88 checkResourceAttribute(*module);
89 checkResourceAttribute(*roundTripModule);
90}
91
92namespace {
93/// A custom operation for the purpose of showcasing how discardable attributes
94/// are handled in absence of properties.
95class OpWithoutProperties : public Op<OpWithoutProperties> {
96public:
97 // Begin boilerplate.
98 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OpWithoutProperties)
99 using Op::Op;
100 static ArrayRef<StringRef> getAttributeNames() {
101 static StringRef attributeNames[] = {StringRef("inherent_attr")};
102 return ArrayRef(attributeNames);
103 };
104 static StringRef getOperationName() {
105 return "test_op_properties.op_without_properties";
106 }
107 // End boilerplate.
108};
109
110// A trivial supporting dialect to register the above operation.
111class TestOpPropertiesDialect : public Dialect {
112public:
113 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestOpPropertiesDialect)
114 static constexpr StringLiteral getDialectNamespace() {
115 return StringLiteral("test_op_properties");
116 }
117 explicit TestOpPropertiesDialect(MLIRContext *context)
118 : Dialect(getDialectNamespace(), context,
119 TypeID::get<TestOpPropertiesDialect>()) {
120 addOperations<OpWithoutProperties>();
121 }
122};
123} // namespace
124
125constexpr StringLiteral withoutPropertiesAttrsSrc = R"mlir(
126 "test_op_properties.op_without_properties"()
127 {inherent_attr = 42, other_attr = 56} : () -> ()
128)mlir";
129
130TEST(Bytecode, OpWithoutProperties) {
131 MLIRContext context;
132 context.getOrLoadDialect<TestOpPropertiesDialect>();
133 ParserConfig config(&context);
134 OwningOpRef<Operation *> op =
135 parseSourceString(sourceStr: withoutPropertiesAttrsSrc, config);
136
137 std::string bytecode;
138 llvm::raw_string_ostream os(bytecode);
139 ASSERT_TRUE(succeeded(writeBytecodeToFile(op.get(), os)));
140 std::unique_ptr<Block> block = std::make_unique<Block>();
141 ASSERT_TRUE(succeeded(readBytecodeFile(
142 llvm::MemoryBufferRef(os.str(), "string-buffer"), block.get(), config)));
143 Operation *roundtripped = &block->front();
144 EXPECT_EQ(roundtripped->getAttrs().size(), 2u);
145 EXPECT_TRUE(roundtripped->getInherentAttr("inherent_attr") != std::nullopt);
146 EXPECT_TRUE(roundtripped->getDiscardableAttr("other_attr") != Attribute());
147
148 EXPECT_TRUE(OperationEquivalence::computeHash(op.get()) ==
149 OperationEquivalence::computeHash(roundtripped));
150}
151

source code of mlir/unittests/Bytecode/BytecodeTest.cpp