1 //===- DataLayoutInterfacesTest.cpp - Unit Tests for Data Layouts ---------===//
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/Interfaces/DataLayoutInterfaces.h"
10 #include "mlir/Dialect/DLTI/DLTI.h"
11 #include "mlir/IR/Builders.h"
12 #include "mlir/IR/BuiltinOps.h"
13 #include "mlir/IR/Dialect.h"
14 #include "mlir/IR/DialectImplementation.h"
15 #include "mlir/IR/OpDefinition.h"
16 #include "mlir/IR/OpImplementation.h"
17 #include "mlir/Parser.h"
18 
19 #include <gtest/gtest.h>
20 
21 using namespace mlir;
22 
23 namespace {
24 constexpr static llvm::StringLiteral kAttrName = "dltest.layout";
25 
26 /// Trivial array storage for the custom data layout spec attribute, just a list
27 /// of entries.
28 class DataLayoutSpecStorage : public AttributeStorage {
29 public:
30   using KeyTy = ArrayRef<DataLayoutEntryInterface>;
31 
32   DataLayoutSpecStorage(ArrayRef<DataLayoutEntryInterface> entries)
33       : entries(entries) {}
34 
35   bool operator==(const KeyTy &key) const { return key == entries; }
36 
37   static DataLayoutSpecStorage *construct(AttributeStorageAllocator &allocator,
38                                           const KeyTy &key) {
39     return new (allocator.allocate<DataLayoutSpecStorage>())
40         DataLayoutSpecStorage(allocator.copyInto(key));
41   }
42 
43   ArrayRef<DataLayoutEntryInterface> entries;
44 };
45 
46 /// Simple data layout spec containing a list of entries that always verifies
47 /// as valid.
48 struct CustomDataLayoutSpec
49     : public Attribute::AttrBase<CustomDataLayoutSpec, Attribute,
50                                  DataLayoutSpecStorage,
51                                  DataLayoutSpecInterface::Trait> {
52   using Base::Base;
53   static CustomDataLayoutSpec get(MLIRContext *ctx,
54                                   ArrayRef<DataLayoutEntryInterface> entries) {
55     return Base::get(ctx, entries);
56   }
57   CustomDataLayoutSpec
58   combineWith(ArrayRef<DataLayoutSpecInterface> specs) const {
59     return *this;
60   }
61   DataLayoutEntryListRef getEntries() const { return getImpl()->entries; }
62   LogicalResult verifySpec(Location loc) { return success(); }
63 };
64 
65 /// A type subject to data layout that exits the program if it is queried more
66 /// than once. Handy to check if the cache works.
67 struct SingleQueryType
68     : public Type::TypeBase<SingleQueryType, Type, TypeStorage,
69                             DataLayoutTypeInterface::Trait> {
70   using Base::Base;
71 
72   static SingleQueryType get(MLIRContext *ctx) { return Base::get(ctx); }
73 
74   unsigned getTypeSize(const DataLayout &layout,
75                        DataLayoutEntryListRef params) {
76     static bool executed = false;
77     if (executed)
78       llvm::report_fatal_error("repeated call");
79 
80     executed = true;
81     return 1;
82   }
83 
84   unsigned getABIAlignment(const DataLayout &layout,
85                            DataLayoutEntryListRef params) {
86     static bool executed = false;
87     if (executed)
88       llvm::report_fatal_error("repeated call");
89 
90     executed = true;
91     return 2;
92   }
93 
94   unsigned getPreferredAlignment(const DataLayout &layout,
95                                  DataLayoutEntryListRef params) {
96     static bool executed = false;
97     if (executed)
98       llvm::report_fatal_error("repeated call");
99 
100     executed = true;
101     return 4;
102   }
103 };
104 
105 /// A types that is not subject to data layout.
106 struct TypeNoLayout : public Type::TypeBase<TypeNoLayout, Type, TypeStorage> {
107   using Base::Base;
108 
109   static TypeNoLayout get(MLIRContext *ctx) { return Base::get(ctx); }
110 };
111 
112 /// An op that serves as scope for data layout queries with the relevant
113 /// attribute attached. This can handle data layout requests for the built-in
114 /// types itself.
115 struct OpWithLayout : public Op<OpWithLayout, DataLayoutOpInterface::Trait> {
116   using Op::Op;
117 
118   static StringRef getOperationName() { return "dltest.op_with_layout"; }
119 
120   DataLayoutSpecInterface getDataLayoutSpec() {
121     return getOperation()->getAttrOfType<DataLayoutSpecInterface>(kAttrName);
122   }
123 
124   static unsigned getTypeSize(Type type, const DataLayout &dataLayout,
125                               DataLayoutEntryListRef params) {
126     // Make a recursive query.
127     if (type.isa<FloatType>())
128       return dataLayout.getTypeSize(
129           IntegerType::get(type.getContext(), type.getIntOrFloatBitWidth()));
130 
131     // Handle built-in types that are not handled by the default process.
132     if (auto iType = type.dyn_cast<IntegerType>()) {
133       for (DataLayoutEntryInterface entry : params)
134         if (entry.getKey().dyn_cast<Type>() == type)
135           return entry.getValue().cast<IntegerAttr>().getValue().getZExtValue();
136       return iType.getIntOrFloatBitWidth();
137     }
138 
139     // Use the default process for everything else.
140     return detail::getDefaultTypeSize(type, dataLayout, params);
141   }
142 
143   static unsigned getTypeABIAlignment(Type type, const DataLayout &dataLayout,
144                                       DataLayoutEntryListRef params) {
145     return llvm::PowerOf2Ceil(getTypeSize(type, dataLayout, params));
146   }
147 
148   static unsigned getTypePreferredAlignment(Type type,
149                                             const DataLayout &dataLayout,
150                                             DataLayoutEntryListRef params) {
151     return 2 * getTypeABIAlignment(type, dataLayout, params);
152   }
153 };
154 
155 /// A dialect putting all the above together.
156 struct DLTestDialect : Dialect {
157   explicit DLTestDialect(MLIRContext *ctx)
158       : Dialect(getDialectNamespace(), ctx, TypeID::get<DLTestDialect>()) {
159     ctx->getOrLoadDialect<DLTIDialect>();
160     addAttributes<CustomDataLayoutSpec>();
161     addOperations<OpWithLayout>();
162     addTypes<SingleQueryType, TypeNoLayout>();
163   }
164   static StringRef getDialectNamespace() { return "dltest"; }
165 
166   void printAttribute(Attribute attr,
167                       DialectAsmPrinter &printer) const override {
168     printer << "spec<";
169     llvm::interleaveComma(attr.cast<CustomDataLayoutSpec>().getEntries(),
170                           printer);
171     printer << ">";
172   }
173 
174   Attribute parseAttribute(DialectAsmParser &parser, Type type) const override {
175     bool ok =
176         succeeded(parser.parseKeyword("spec")) && succeeded(parser.parseLess());
177     (void)ok;
178     assert(ok);
179     if (succeeded(parser.parseOptionalGreater()))
180       return CustomDataLayoutSpec::get(parser.getBuilder().getContext(), {});
181 
182     SmallVector<DataLayoutEntryInterface> entries;
183     do {
184       entries.emplace_back();
185       ok = succeeded(parser.parseAttribute(entries.back()));
186       assert(ok);
187     } while (succeeded(parser.parseOptionalComma()));
188     ok = succeeded(parser.parseGreater());
189     assert(ok);
190     return CustomDataLayoutSpec::get(parser.getBuilder().getContext(), entries);
191   }
192 
193   void printType(Type type, DialectAsmPrinter &printer) const override {
194     if (type.isa<SingleQueryType>())
195       printer << "single_query";
196     else
197       printer << "no_layout";
198   }
199 
200   Type parseType(DialectAsmParser &parser) const override {
201     bool ok = succeeded(parser.parseKeyword("single_query"));
202     (void)ok;
203     assert(ok);
204     return SingleQueryType::get(parser.getBuilder().getContext());
205   }
206 };
207 
208 } // end namespace
209 
210 TEST(DataLayout, FallbackDefault) {
211   const char *ir = R"MLIR(
212 "dltest.op_with_layout"() : () -> ()
213   )MLIR";
214 
215   DialectRegistry registry;
216   registry.insert<DLTIDialect, DLTestDialect>();
217   MLIRContext ctx(registry);
218 
219   OwningModuleRef module = parseSourceString(ir, &ctx);
220   auto op =
221       cast<DataLayoutOpInterface>(module->getBody()->getOperations().front());
222   DataLayout layout(op);
223   EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 42)), 6u);
224   EXPECT_EQ(layout.getTypeSize(Float16Type::get(&ctx)), 2u);
225   EXPECT_EQ(layout.getTypeABIAlignment(IntegerType::get(&ctx, 42)), 8u);
226   EXPECT_EQ(layout.getTypeABIAlignment(Float16Type::get(&ctx)), 2u);
227   EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 42)), 8u);
228   EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 2u);
229 }
230 
231 TEST(DataLayout, EmptySpec) {
232   const char *ir = R"MLIR(
233 "dltest.op_with_layout"() { dltest.layout = #dltest.spec< > } : () -> ()
234   )MLIR";
235 
236   DialectRegistry registry;
237   registry.insert<DLTIDialect, DLTestDialect>();
238   MLIRContext ctx(registry);
239 
240   OwningModuleRef module = parseSourceString(ir, &ctx);
241   auto op =
242       cast<DataLayoutOpInterface>(module->getBody()->getOperations().front());
243   DataLayout layout(op);
244   EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 42)), 42u);
245   EXPECT_EQ(layout.getTypeSize(Float16Type::get(&ctx)), 16u);
246   EXPECT_EQ(layout.getTypeABIAlignment(IntegerType::get(&ctx, 42)), 64u);
247   EXPECT_EQ(layout.getTypeABIAlignment(Float16Type::get(&ctx)), 16u);
248   EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 42)), 128u);
249   EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 32u);
250 }
251 
252 TEST(DataLayout, SpecWithEntries) {
253   const char *ir = R"MLIR(
254 "dltest.op_with_layout"() { dltest.layout = #dltest.spec<
255   #dlti.dl_entry<i42, 5>,
256   #dlti.dl_entry<i16, 6>
257 > } : () -> ()
258   )MLIR";
259 
260   DialectRegistry registry;
261   registry.insert<DLTIDialect, DLTestDialect>();
262   MLIRContext ctx(registry);
263 
264   OwningModuleRef module = parseSourceString(ir, &ctx);
265   auto op =
266       cast<DataLayoutOpInterface>(module->getBody()->getOperations().front());
267   DataLayout layout(op);
268   EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 42)), 5u);
269   EXPECT_EQ(layout.getTypeSize(Float16Type::get(&ctx)), 6u);
270   EXPECT_EQ(layout.getTypeABIAlignment(IntegerType::get(&ctx, 42)), 8u);
271   EXPECT_EQ(layout.getTypeABIAlignment(Float16Type::get(&ctx)), 8u);
272   EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 42)), 16u);
273   EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 16u);
274 
275   EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 32)), 32u);
276   EXPECT_EQ(layout.getTypeSize(Float32Type::get(&ctx)), 32u);
277   EXPECT_EQ(layout.getTypeABIAlignment(IntegerType::get(&ctx, 32)), 32u);
278   EXPECT_EQ(layout.getTypeABIAlignment(Float32Type::get(&ctx)), 32u);
279   EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 32)), 64u);
280   EXPECT_EQ(layout.getTypePreferredAlignment(Float32Type::get(&ctx)), 64u);
281 }
282 
283 TEST(DataLayout, Caching) {
284   const char *ir = R"MLIR(
285 "dltest.op_with_layout"() { dltest.layout = #dltest.spec<> } : () -> ()
286   )MLIR";
287 
288   DialectRegistry registry;
289   registry.insert<DLTIDialect, DLTestDialect>();
290   MLIRContext ctx(registry);
291 
292   OwningModuleRef module = parseSourceString(ir, &ctx);
293   auto op =
294       cast<DataLayoutOpInterface>(module->getBody()->getOperations().front());
295   DataLayout layout(op);
296 
297   unsigned sum = 0;
298   sum += layout.getTypeSize(SingleQueryType::get(&ctx));
299   // The second call should hit the cache. If it does not, the function in
300   // SingleQueryType will be called and will abort the process.
301   sum += layout.getTypeSize(SingleQueryType::get(&ctx));
302   // Make sure the complier doesn't optimize away the query code.
303   EXPECT_EQ(sum, 2u);
304 
305   // A fresh data layout has a new cache, so the call to it should be dispatched
306   // down to the type and abort the proces.
307   DataLayout second(op);
308   ASSERT_DEATH(second.getTypeSize(SingleQueryType::get(&ctx)), "repeated call");
309 }
310 
311 TEST(DataLayout, CacheInvalidation) {
312   const char *ir = R"MLIR(
313 "dltest.op_with_layout"() { dltest.layout = #dltest.spec<
314   #dlti.dl_entry<i42, 5>,
315   #dlti.dl_entry<i16, 6>
316 > } : () -> ()
317   )MLIR";
318 
319   DialectRegistry registry;
320   registry.insert<DLTIDialect, DLTestDialect>();
321   MLIRContext ctx(registry);
322 
323   OwningModuleRef module = parseSourceString(ir, &ctx);
324   auto op =
325       cast<DataLayoutOpInterface>(module->getBody()->getOperations().front());
326   DataLayout layout(op);
327 
328   // Normal query is fine.
329   EXPECT_EQ(layout.getTypeSize(Float16Type::get(&ctx)), 6u);
330 
331   // Replace the data layout spec with a new, empty spec.
332   op->setAttr(kAttrName, CustomDataLayoutSpec::get(&ctx, {}));
333 
334   // Data layout is no longer valid and should trigger assertion when queried.
335 #ifndef NDEBUG
336   ASSERT_DEATH(layout.getTypeSize(Float16Type::get(&ctx)), "no longer valid");
337 #endif
338 }
339 
340 TEST(DataLayout, UnimplementedTypeInterface) {
341   const char *ir = R"MLIR(
342 "dltest.op_with_layout"() { dltest.layout = #dltest.spec<> } : () -> ()
343   )MLIR";
344 
345   DialectRegistry registry;
346   registry.insert<DLTIDialect, DLTestDialect>();
347   MLIRContext ctx(registry);
348 
349   OwningModuleRef module = parseSourceString(ir, &ctx);
350   auto op =
351       cast<DataLayoutOpInterface>(module->getBody()->getOperations().front());
352   DataLayout layout(op);
353 
354   ASSERT_DEATH(layout.getTypeSize(TypeNoLayout::get(&ctx)),
355                "neither the scoping op nor the type class provide data layout "
356                "information");
357 }
358