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 "llvm/Support/raw_ostream.h"
21#include "gmock/gmock.h"
22#include "gtest/gtest.h"
23
24using namespace llvm;
25using namespace mlir;
26
27StringLiteral irWithResources = R"(
28module @TestDialectResources attributes {
29 bytecode.test = dense_resource<resource> : tensor<4xi32>
30} {}
31{-#
32 dialect_resources: {
33 builtin: {
34 resource: "0x2000000001000000020000000300000004000000",
35 resource_2: "0x2000000001000000020000000300000004000000"
36 }
37 }
38#-}
39)";
40
41struct MockOstream final : public raw_ostream {
42 std::unique_ptr<std::byte[]> buffer;
43 size_t size = 0;
44
45 MOCK_METHOD(void, reserveExtraSpace, (uint64_t extraSpace), (override));
46
47 MockOstream() : raw_ostream(true) {}
48 uint64_t current_pos() const override { return pos; }
49
50private:
51 size_t pos = 0;
52
53 void write_impl(const char *ptr, size_t length) override {
54 if (pos + length <= size) {
55 memcpy(dest: (void *)(buffer.get() + pos), src: ptr, n: length);
56 pos += length;
57 } else {
58 report_fatal_error(
59 reason: "Attempted to write past the end of the fixed size buffer.");
60 }
61 }
62};
63
64TEST(Bytecode, MultiModuleWithResource) {
65 MLIRContext context;
66 Builder builder(&context);
67 ParserConfig parseConfig(&context);
68 OwningOpRef<Operation *> module =
69 parseSourceString<Operation *>(sourceStr: irWithResources, config: parseConfig);
70 ASSERT_TRUE(module);
71
72 // Write the module to bytecode.
73 MockOstream ostream;
74 EXPECT_CALL(ostream, reserveExtraSpace).WillOnce(once_action: [&](uint64_t space) {
75 ostream.buffer = std::make_unique<std::byte[]>(num: space);
76 ostream.size = space;
77 });
78 ASSERT_TRUE(succeeded(writeBytecodeToFile(module.get(), ostream)));
79
80 // Create copy of buffer which is aligned to requested resource alignment.
81 std::string buffer((char *)ostream.buffer.get(),
82 (char *)ostream.buffer.get() + ostream.size);
83 constexpr size_t kAlignment = 0x20;
84 size_t bufferSize = buffer.size();
85 buffer.reserve(res: bufferSize + kAlignment - 1);
86 size_t pad = (~(uintptr_t)buffer.data() + 1) & (kAlignment - 1);
87 buffer.insert(pos: 0, n: pad, c: ' ');
88 StringRef alignedBuffer(buffer.data() + pad, bufferSize);
89
90 // Parse it back
91 OwningOpRef<Operation *> roundTripModule =
92 parseSourceString<Operation *>(sourceStr: alignedBuffer, config: parseConfig);
93 ASSERT_TRUE(roundTripModule);
94
95 // FIXME: Parsing external resources does not work on big-endian
96 // platforms currently.
97 if (llvm::endianness::native == llvm::endianness::big)
98 GTEST_SKIP();
99
100 // Try to see if we have a valid resource in the parsed module.
101 auto checkResourceAttribute = [](Operation *parsedModule) {
102 Attribute attr = parsedModule->getDiscardableAttr(name: "bytecode.test");
103 ASSERT_TRUE(attr);
104 auto denseResourceAttr = dyn_cast<DenseI32ResourceElementsAttr>(attr);
105 ASSERT_TRUE(denseResourceAttr);
106 std::optional<ArrayRef<int32_t>> attrData =
107 denseResourceAttr.tryGetAsArrayRef();
108 ASSERT_TRUE(attrData.has_value());
109 ASSERT_EQ(attrData->size(), static_cast<size_t>(4));
110 EXPECT_EQ((*attrData)[0], 1);
111 EXPECT_EQ((*attrData)[1], 2);
112 EXPECT_EQ((*attrData)[2], 3);
113 EXPECT_EQ((*attrData)[3], 4);
114 };
115
116 checkResourceAttribute(*module);
117 checkResourceAttribute(*roundTripModule);
118}
119
120namespace {
121/// A custom operation for the purpose of showcasing how discardable attributes
122/// are handled in absence of properties.
123class OpWithoutProperties : public Op<OpWithoutProperties> {
124public:
125 // Begin boilerplate.
126 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OpWithoutProperties)
127 using Op::Op;
128 static ArrayRef<StringRef> getAttributeNames() {
129 static StringRef attributeNames[] = {StringRef("inherent_attr")};
130 return ArrayRef(attributeNames);
131 };
132 static StringRef getOperationName() {
133 return "test_op_properties.op_without_properties";
134 }
135 // End boilerplate.
136};
137
138// A trivial supporting dialect to register the above operation.
139class TestOpPropertiesDialect : public Dialect {
140public:
141 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestOpPropertiesDialect)
142 static constexpr StringLiteral getDialectNamespace() {
143 return StringLiteral("test_op_properties");
144 }
145 explicit TestOpPropertiesDialect(MLIRContext *context)
146 : Dialect(getDialectNamespace(), context,
147 TypeID::get<TestOpPropertiesDialect>()) {
148 addOperations<OpWithoutProperties>();
149 }
150};
151} // namespace
152
153constexpr StringLiteral withoutPropertiesAttrsSrc = R"mlir(
154 "test_op_properties.op_without_properties"()
155 {inherent_attr = 42, other_attr = 56} : () -> ()
156)mlir";
157
158TEST(Bytecode, OpWithoutProperties) {
159 MLIRContext context;
160 context.getOrLoadDialect<TestOpPropertiesDialect>();
161 ParserConfig config(&context);
162 OwningOpRef<Operation *> op =
163 parseSourceString(sourceStr: withoutPropertiesAttrsSrc, config);
164
165 std::string bytecode;
166 llvm::raw_string_ostream os(bytecode);
167 ASSERT_TRUE(succeeded(writeBytecodeToFile(op.get(), os)));
168 std::unique_ptr<Block> block = std::make_unique<Block>();
169 ASSERT_TRUE(succeeded(readBytecodeFile(
170 llvm::MemoryBufferRef(bytecode, "string-buffer"), block.get(), config)));
171 Operation *roundtripped = &block->front();
172 EXPECT_EQ(roundtripped->getAttrs().size(), 2u);
173 EXPECT_TRUE(roundtripped->getInherentAttr("inherent_attr") != std::nullopt);
174 EXPECT_TRUE(roundtripped->getDiscardableAttr("other_attr") != Attribute());
175
176 EXPECT_TRUE(OperationEquivalence::computeHash(op.get()) ==
177 OperationEquivalence::computeHash(roundtripped));
178}
179

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