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 | |
23 | using namespace llvm; |
24 | using namespace mlir; |
25 | |
26 | StringLiteral irWithResources = R"( |
27 | module @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 | |
40 | TEST(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 | |
92 | namespace { |
93 | /// A custom operation for the purpose of showcasing how discardable attributes |
94 | /// are handled in absence of properties. |
95 | class OpWithoutProperties : public Op<OpWithoutProperties> { |
96 | public: |
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. |
111 | class TestOpPropertiesDialect : public Dialect { |
112 | public: |
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 | |
125 | constexpr StringLiteral = R"mlir( |
126 | "test_op_properties.op_without_properties"() |
127 | {inherent_attr = 42, other_attr = 56} : () -> () |
128 | )mlir" ; |
129 | |
130 | TEST(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 | |