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/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 | constexpr static llvm::StringLiteral kEndiannesKeyName = "dltest.endianness" ; |
26 | constexpr static llvm::StringLiteral kAllocaKeyName = |
27 | "dltest.alloca_memory_space" ; |
28 | constexpr static llvm::StringLiteral kProgramKeyName = |
29 | "dltest.program_memory_space" ; |
30 | constexpr static llvm::StringLiteral kGlobalKeyName = |
31 | "dltest.global_memory_space" ; |
32 | constexpr static llvm::StringLiteral kStackAlignmentKeyName = |
33 | "dltest.stack_alignment" ; |
34 | |
35 | /// Trivial array storage for the custom data layout spec attribute, just a list |
36 | /// of entries. |
37 | class DataLayoutSpecStorage : public AttributeStorage { |
38 | public: |
39 | using KeyTy = ArrayRef<DataLayoutEntryInterface>; |
40 | |
41 | DataLayoutSpecStorage(ArrayRef<DataLayoutEntryInterface> entries) |
42 | : entries(entries) {} |
43 | |
44 | bool operator==(const KeyTy &key) const { return key == entries; } |
45 | |
46 | static DataLayoutSpecStorage *construct(AttributeStorageAllocator &allocator, |
47 | const KeyTy &key) { |
48 | return new (allocator.allocate<DataLayoutSpecStorage>()) |
49 | DataLayoutSpecStorage(allocator.copyInto(elements: key)); |
50 | } |
51 | |
52 | ArrayRef<DataLayoutEntryInterface> entries; |
53 | }; |
54 | |
55 | /// Simple data layout spec containing a list of entries that always verifies |
56 | /// as valid. |
57 | struct CustomDataLayoutSpec |
58 | : public Attribute::AttrBase<CustomDataLayoutSpec, Attribute, |
59 | DataLayoutSpecStorage, |
60 | DataLayoutSpecInterface::Trait> { |
61 | MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CustomDataLayoutSpec) |
62 | |
63 | using Base::Base; |
64 | |
65 | static constexpr StringLiteral name = "test.custom_data_layout_spec" ; |
66 | |
67 | static CustomDataLayoutSpec get(MLIRContext *ctx, |
68 | ArrayRef<DataLayoutEntryInterface> entries) { |
69 | return Base::get(ctx, entries); |
70 | } |
71 | CustomDataLayoutSpec |
72 | combineWith(ArrayRef<DataLayoutSpecInterface> specs) const { |
73 | return *this; |
74 | } |
75 | DataLayoutEntryListRef getEntries() const { return getImpl()->entries; } |
76 | LogicalResult verifySpec(Location loc) { return success(); } |
77 | StringAttr getEndiannessIdentifier(MLIRContext *context) const { |
78 | return Builder(context).getStringAttr(kEndiannesKeyName); |
79 | } |
80 | StringAttr getAllocaMemorySpaceIdentifier(MLIRContext *context) const { |
81 | return Builder(context).getStringAttr(kAllocaKeyName); |
82 | } |
83 | StringAttr getProgramMemorySpaceIdentifier(MLIRContext *context) const { |
84 | return Builder(context).getStringAttr(kProgramKeyName); |
85 | } |
86 | StringAttr getGlobalMemorySpaceIdentifier(MLIRContext *context) const { |
87 | return Builder(context).getStringAttr(kGlobalKeyName); |
88 | } |
89 | StringAttr getStackAlignmentIdentifier(MLIRContext *context) const { |
90 | return Builder(context).getStringAttr(kStackAlignmentKeyName); |
91 | } |
92 | }; |
93 | |
94 | /// A type subject to data layout that exits the program if it is queried more |
95 | /// than once. Handy to check if the cache works. |
96 | struct SingleQueryType |
97 | : public Type::TypeBase<SingleQueryType, Type, TypeStorage, |
98 | DataLayoutTypeInterface::Trait> { |
99 | MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SingleQueryType) |
100 | |
101 | using Base::Base; |
102 | |
103 | static constexpr StringLiteral name = "test.single_query" ; |
104 | |
105 | static SingleQueryType get(MLIRContext *ctx) { return Base::get(ctx); } |
106 | |
107 | llvm::TypeSize getTypeSizeInBits(const DataLayout &layout, |
108 | DataLayoutEntryListRef params) const { |
109 | static bool executed = false; |
110 | if (executed) |
111 | llvm::report_fatal_error(reason: "repeated call" ); |
112 | |
113 | executed = true; |
114 | return llvm::TypeSize::getFixed(ExactSize: 1); |
115 | } |
116 | |
117 | uint64_t getABIAlignment(const DataLayout &layout, |
118 | DataLayoutEntryListRef params) { |
119 | static bool executed = false; |
120 | if (executed) |
121 | llvm::report_fatal_error(reason: "repeated call" ); |
122 | |
123 | executed = true; |
124 | return 2; |
125 | } |
126 | |
127 | uint64_t getPreferredAlignment(const DataLayout &layout, |
128 | DataLayoutEntryListRef params) { |
129 | static bool executed = false; |
130 | if (executed) |
131 | llvm::report_fatal_error(reason: "repeated call" ); |
132 | |
133 | executed = true; |
134 | return 4; |
135 | } |
136 | |
137 | Attribute getEndianness(DataLayoutEntryInterface entry) { |
138 | static bool executed = false; |
139 | if (executed) |
140 | llvm::report_fatal_error(reason: "repeated call" ); |
141 | |
142 | executed = true; |
143 | return Attribute(); |
144 | } |
145 | |
146 | Attribute getAllocaMemorySpace(DataLayoutEntryInterface entry) { |
147 | static bool executed = false; |
148 | if (executed) |
149 | llvm::report_fatal_error(reason: "repeated call" ); |
150 | |
151 | executed = true; |
152 | return Attribute(); |
153 | } |
154 | |
155 | Attribute getProgramMemorySpace(DataLayoutEntryInterface entry) { |
156 | static bool executed = false; |
157 | if (executed) |
158 | llvm::report_fatal_error(reason: "repeated call" ); |
159 | |
160 | executed = true; |
161 | return Attribute(); |
162 | } |
163 | |
164 | Attribute getGlobalMemorySpace(DataLayoutEntryInterface entry) { |
165 | static bool executed = false; |
166 | if (executed) |
167 | llvm::report_fatal_error(reason: "repeated call" ); |
168 | |
169 | executed = true; |
170 | return Attribute(); |
171 | } |
172 | }; |
173 | |
174 | /// A types that is not subject to data layout. |
175 | struct TypeNoLayout : public Type::TypeBase<TypeNoLayout, Type, TypeStorage> { |
176 | MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TypeNoLayout) |
177 | |
178 | using Base::Base; |
179 | |
180 | static constexpr StringLiteral name = "test.no_layout" ; |
181 | |
182 | static TypeNoLayout get(MLIRContext *ctx) { return Base::get(ctx); } |
183 | }; |
184 | |
185 | /// An op that serves as scope for data layout queries with the relevant |
186 | /// attribute attached. This can handle data layout requests for the built-in |
187 | /// types itself. |
188 | struct OpWithLayout : public Op<OpWithLayout, DataLayoutOpInterface::Trait> { |
189 | MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OpWithLayout) |
190 | |
191 | using Op::Op; |
192 | static ArrayRef<StringRef> getAttributeNames() { return {}; } |
193 | |
194 | static StringRef getOperationName() { return "dltest.op_with_layout" ; } |
195 | |
196 | DataLayoutSpecInterface getDataLayoutSpec() { |
197 | return getOperation()->getAttrOfType<DataLayoutSpecInterface>(kAttrName); |
198 | } |
199 | |
200 | static llvm::TypeSize getTypeSizeInBits(Type type, |
201 | const DataLayout &dataLayout, |
202 | DataLayoutEntryListRef params) { |
203 | // Make a recursive query. |
204 | if (isa<FloatType>(type)) |
205 | return dataLayout.getTypeSizeInBits( |
206 | IntegerType::get(type.getContext(), type.getIntOrFloatBitWidth())); |
207 | |
208 | // Handle built-in types that are not handled by the default process. |
209 | if (auto iType = dyn_cast<IntegerType>(type)) { |
210 | for (DataLayoutEntryInterface entry : params) |
211 | if (llvm::dyn_cast_if_present<Type>(entry.getKey()) == type) |
212 | return llvm::TypeSize::getFixed( |
213 | 8 * |
214 | cast<IntegerAttr>(entry.getValue()).getValue().getZExtValue()); |
215 | return llvm::TypeSize::getFixed(ExactSize: 8 * iType.getIntOrFloatBitWidth()); |
216 | } |
217 | |
218 | // Use the default process for everything else. |
219 | return detail::getDefaultTypeSize(type, dataLayout, params); |
220 | } |
221 | |
222 | static uint64_t getTypeABIAlignment(Type type, const DataLayout &dataLayout, |
223 | DataLayoutEntryListRef params) { |
224 | return llvm::PowerOf2Ceil(A: getTypeSize(type, dataLayout, params)); |
225 | } |
226 | |
227 | static uint64_t getTypePreferredAlignment(Type type, |
228 | const DataLayout &dataLayout, |
229 | DataLayoutEntryListRef params) { |
230 | return 2 * getTypeABIAlignment(type, dataLayout, params); |
231 | } |
232 | }; |
233 | |
234 | struct OpWith7BitByte |
235 | : public Op<OpWith7BitByte, DataLayoutOpInterface::Trait> { |
236 | MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OpWith7BitByte) |
237 | |
238 | using Op::Op; |
239 | static ArrayRef<StringRef> getAttributeNames() { return {}; } |
240 | |
241 | static StringRef getOperationName() { return "dltest.op_with_7bit_byte" ; } |
242 | |
243 | DataLayoutSpecInterface getDataLayoutSpec() { |
244 | return getOperation()->getAttrOfType<DataLayoutSpecInterface>(kAttrName); |
245 | } |
246 | |
247 | // Bytes are assumed to be 7-bit here. |
248 | static llvm::TypeSize getTypeSize(Type type, const DataLayout &dataLayout, |
249 | DataLayoutEntryListRef params) { |
250 | return mlir::detail::divideCeil(numerator: dataLayout.getTypeSizeInBits(t: type), denominator: 7); |
251 | } |
252 | }; |
253 | |
254 | /// A dialect putting all the above together. |
255 | struct DLTestDialect : Dialect { |
256 | MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(DLTestDialect) |
257 | |
258 | explicit DLTestDialect(MLIRContext *ctx) |
259 | : Dialect(getDialectNamespace(), ctx, TypeID::get<DLTestDialect>()) { |
260 | ctx->getOrLoadDialect<DLTIDialect>(); |
261 | addAttributes<CustomDataLayoutSpec>(); |
262 | addOperations<OpWithLayout, OpWith7BitByte>(); |
263 | addTypes<SingleQueryType, TypeNoLayout>(); |
264 | } |
265 | static StringRef getDialectNamespace() { return "dltest" ; } |
266 | |
267 | void printAttribute(Attribute attr, |
268 | DialectAsmPrinter &printer) const override { |
269 | printer << "spec<" ; |
270 | llvm::interleaveComma(c: cast<CustomDataLayoutSpec>(Val&: attr).getEntries(), |
271 | os&: printer); |
272 | printer << ">" ; |
273 | } |
274 | |
275 | Attribute parseAttribute(DialectAsmParser &parser, Type type) const override { |
276 | bool ok = |
277 | succeeded(result: parser.parseKeyword(keyword: "spec" )) && succeeded(result: parser.parseLess()); |
278 | (void)ok; |
279 | assert(ok); |
280 | if (succeeded(result: parser.parseOptionalGreater())) |
281 | return CustomDataLayoutSpec::get(ctx: parser.getContext(), entries: {}); |
282 | |
283 | SmallVector<DataLayoutEntryInterface> entries; |
284 | ok = succeeded(result: parser.parseCommaSeparatedList(parseElementFn: [&]() { |
285 | entries.emplace_back(); |
286 | ok = succeeded(parser.parseAttribute(entries.back())); |
287 | assert(ok); |
288 | return success(); |
289 | })); |
290 | assert(ok); |
291 | ok = succeeded(result: parser.parseGreater()); |
292 | assert(ok); |
293 | return CustomDataLayoutSpec::get(ctx: parser.getContext(), entries: entries); |
294 | } |
295 | |
296 | void printType(Type type, DialectAsmPrinter &printer) const override { |
297 | if (isa<SingleQueryType>(Val: type)) |
298 | printer << "single_query" ; |
299 | else |
300 | printer << "no_layout" ; |
301 | } |
302 | |
303 | Type parseType(DialectAsmParser &parser) const override { |
304 | bool ok = succeeded(result: parser.parseKeyword(keyword: "single_query" )); |
305 | (void)ok; |
306 | assert(ok); |
307 | return SingleQueryType::get(ctx: parser.getContext()); |
308 | } |
309 | }; |
310 | |
311 | } // namespace |
312 | |
313 | TEST(DataLayout, FallbackDefault) { |
314 | const char *ir = R"MLIR( |
315 | module {} |
316 | )MLIR" ; |
317 | |
318 | DialectRegistry registry; |
319 | registry.insert<DLTIDialect, DLTestDialect>(); |
320 | MLIRContext ctx(registry); |
321 | |
322 | OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(ir, &ctx); |
323 | DataLayout layout(module.get()); |
324 | EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 42)), 6u); |
325 | EXPECT_EQ(layout.getTypeSize(Float16Type::get(&ctx)), 2u); |
326 | EXPECT_EQ(layout.getTypeSizeInBits(IntegerType::get(&ctx, 42)), 42u); |
327 | EXPECT_EQ(layout.getTypeSizeInBits(Float16Type::get(&ctx)), 16u); |
328 | EXPECT_EQ(layout.getTypeABIAlignment(IntegerType::get(&ctx, 42)), 8u); |
329 | EXPECT_EQ(layout.getTypeABIAlignment(Float16Type::get(&ctx)), 2u); |
330 | EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 42)), 8u); |
331 | EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 2u); |
332 | |
333 | EXPECT_EQ(layout.getEndianness(), Attribute()); |
334 | EXPECT_EQ(layout.getAllocaMemorySpace(), Attribute()); |
335 | EXPECT_EQ(layout.getProgramMemorySpace(), Attribute()); |
336 | EXPECT_EQ(layout.getGlobalMemorySpace(), Attribute()); |
337 | EXPECT_EQ(layout.getStackAlignment(), 0u); |
338 | } |
339 | |
340 | TEST(DataLayout, NullSpec) { |
341 | const char *ir = R"MLIR( |
342 | "dltest.op_with_layout"() : () -> () |
343 | )MLIR" ; |
344 | |
345 | DialectRegistry registry; |
346 | registry.insert<DLTIDialect, DLTestDialect>(); |
347 | MLIRContext ctx(registry); |
348 | |
349 | OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(sourceStr: ir, config: &ctx); |
350 | auto op = |
351 | cast<DataLayoutOpInterface>(module->getBody()->getOperations().front()); |
352 | DataLayout layout(op); |
353 | |
354 | EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 42)), 42u); |
355 | EXPECT_EQ(layout.getTypeSize(Float16Type::get(&ctx)), 16u); |
356 | EXPECT_EQ(layout.getTypeSizeInBits(IntegerType::get(&ctx, 42)), 8u * 42u); |
357 | EXPECT_EQ(layout.getTypeSizeInBits(Float16Type::get(&ctx)), 8u * 16u); |
358 | EXPECT_EQ(layout.getTypeABIAlignment(IntegerType::get(&ctx, 42)), 64u); |
359 | EXPECT_EQ(layout.getTypeABIAlignment(Float16Type::get(&ctx)), 16u); |
360 | EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 42)), 128u); |
361 | EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 32u); |
362 | EXPECT_EQ(layout.getTypeIndexBitwidth(Float16Type::get(&ctx)), std::nullopt); |
363 | EXPECT_EQ(layout.getTypeIndexBitwidth(IndexType::get(&ctx)), 64u); |
364 | |
365 | EXPECT_EQ(layout.getEndianness(), Attribute()); |
366 | EXPECT_EQ(layout.getAllocaMemorySpace(), Attribute()); |
367 | EXPECT_EQ(layout.getProgramMemorySpace(), Attribute()); |
368 | EXPECT_EQ(layout.getGlobalMemorySpace(), Attribute()); |
369 | EXPECT_EQ(layout.getStackAlignment(), 0u); |
370 | } |
371 | |
372 | TEST(DataLayout, EmptySpec) { |
373 | const char *ir = R"MLIR( |
374 | "dltest.op_with_layout"() { dltest.layout = #dltest.spec< > } : () -> () |
375 | )MLIR" ; |
376 | |
377 | DialectRegistry registry; |
378 | registry.insert<DLTIDialect, DLTestDialect>(); |
379 | MLIRContext ctx(registry); |
380 | |
381 | OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(sourceStr: ir, config: &ctx); |
382 | auto op = |
383 | cast<DataLayoutOpInterface>(module->getBody()->getOperations().front()); |
384 | DataLayout layout(op); |
385 | EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 42)), 42u); |
386 | EXPECT_EQ(layout.getTypeSize(Float16Type::get(&ctx)), 16u); |
387 | EXPECT_EQ(layout.getTypeSizeInBits(IntegerType::get(&ctx, 42)), 8u * 42u); |
388 | EXPECT_EQ(layout.getTypeSizeInBits(Float16Type::get(&ctx)), 8u * 16u); |
389 | EXPECT_EQ(layout.getTypeABIAlignment(IntegerType::get(&ctx, 42)), 64u); |
390 | EXPECT_EQ(layout.getTypeABIAlignment(Float16Type::get(&ctx)), 16u); |
391 | EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 42)), 128u); |
392 | EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 32u); |
393 | EXPECT_EQ(layout.getTypeIndexBitwidth(Float16Type::get(&ctx)), std::nullopt); |
394 | EXPECT_EQ(layout.getTypeIndexBitwidth(IndexType::get(&ctx)), 64u); |
395 | |
396 | EXPECT_EQ(layout.getEndianness(), Attribute()); |
397 | EXPECT_EQ(layout.getAllocaMemorySpace(), Attribute()); |
398 | EXPECT_EQ(layout.getProgramMemorySpace(), Attribute()); |
399 | EXPECT_EQ(layout.getGlobalMemorySpace(), Attribute()); |
400 | EXPECT_EQ(layout.getStackAlignment(), 0u); |
401 | } |
402 | |
403 | TEST(DataLayout, SpecWithEntries) { |
404 | const char *ir = R"MLIR( |
405 | "dltest.op_with_layout"() { dltest.layout = #dltest.spec< |
406 | #dlti.dl_entry<i42, 5>, |
407 | #dlti.dl_entry<i16, 6>, |
408 | #dlti.dl_entry<index, 42>, |
409 | #dlti.dl_entry<"dltest.endianness", "little">, |
410 | #dlti.dl_entry<"dltest.alloca_memory_space", 5 : i32>, |
411 | #dlti.dl_entry<"dltest.program_memory_space", 3 : i32>, |
412 | #dlti.dl_entry<"dltest.global_memory_space", 2 : i32>, |
413 | #dlti.dl_entry<"dltest.stack_alignment", 128 : i32> |
414 | > } : () -> () |
415 | )MLIR" ; |
416 | |
417 | DialectRegistry registry; |
418 | registry.insert<DLTIDialect, DLTestDialect>(); |
419 | MLIRContext ctx(registry); |
420 | |
421 | OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(sourceStr: ir, config: &ctx); |
422 | auto op = |
423 | cast<DataLayoutOpInterface>(module->getBody()->getOperations().front()); |
424 | DataLayout layout(op); |
425 | EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 42)), 5u); |
426 | EXPECT_EQ(layout.getTypeSize(Float16Type::get(&ctx)), 6u); |
427 | EXPECT_EQ(layout.getTypeSizeInBits(IntegerType::get(&ctx, 42)), 40u); |
428 | EXPECT_EQ(layout.getTypeSizeInBits(Float16Type::get(&ctx)), 48u); |
429 | EXPECT_EQ(layout.getTypeABIAlignment(IntegerType::get(&ctx, 42)), 8u); |
430 | EXPECT_EQ(layout.getTypeABIAlignment(Float16Type::get(&ctx)), 8u); |
431 | EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 42)), 16u); |
432 | EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 16u); |
433 | EXPECT_EQ(layout.getTypeIndexBitwidth(Float16Type::get(&ctx)), std::nullopt); |
434 | EXPECT_EQ(layout.getTypeIndexBitwidth(IndexType::get(&ctx)), 42u); |
435 | |
436 | EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 32)), 32u); |
437 | EXPECT_EQ(layout.getTypeSize(Float32Type::get(&ctx)), 32u); |
438 | EXPECT_EQ(layout.getTypeSizeInBits(IntegerType::get(&ctx, 32)), 256u); |
439 | EXPECT_EQ(layout.getTypeSizeInBits(Float32Type::get(&ctx)), 256u); |
440 | EXPECT_EQ(layout.getTypeABIAlignment(IntegerType::get(&ctx, 32)), 32u); |
441 | EXPECT_EQ(layout.getTypeABIAlignment(Float32Type::get(&ctx)), 32u); |
442 | EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 32)), 64u); |
443 | EXPECT_EQ(layout.getTypePreferredAlignment(Float32Type::get(&ctx)), 64u); |
444 | |
445 | EXPECT_EQ(layout.getEndianness(), Builder(&ctx).getStringAttr("little" )); |
446 | EXPECT_EQ(layout.getAllocaMemorySpace(), Builder(&ctx).getI32IntegerAttr(5)); |
447 | EXPECT_EQ(layout.getProgramMemorySpace(), Builder(&ctx).getI32IntegerAttr(3)); |
448 | EXPECT_EQ(layout.getGlobalMemorySpace(), Builder(&ctx).getI32IntegerAttr(2)); |
449 | EXPECT_EQ(layout.getStackAlignment(), 128u); |
450 | } |
451 | |
452 | TEST(DataLayout, Caching) { |
453 | const char *ir = R"MLIR( |
454 | "dltest.op_with_layout"() { dltest.layout = #dltest.spec<> } : () -> () |
455 | )MLIR" ; |
456 | |
457 | DialectRegistry registry; |
458 | registry.insert<DLTIDialect, DLTestDialect>(); |
459 | MLIRContext ctx(registry); |
460 | |
461 | OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(sourceStr: ir, config: &ctx); |
462 | auto op = |
463 | cast<DataLayoutOpInterface>(module->getBody()->getOperations().front()); |
464 | DataLayout layout(op); |
465 | |
466 | unsigned sum = 0; |
467 | sum += layout.getTypeSize(SingleQueryType::get(ctx: &ctx)); |
468 | // The second call should hit the cache. If it does not, the function in |
469 | // SingleQueryType will be called and will abort the process. |
470 | sum += layout.getTypeSize(SingleQueryType::get(ctx: &ctx)); |
471 | // Make sure the complier doesn't optimize away the query code. |
472 | EXPECT_EQ(sum, 2u); |
473 | |
474 | // A fresh data layout has a new cache, so the call to it should be dispatched |
475 | // down to the type and abort the process. |
476 | DataLayout second(op); |
477 | ASSERT_DEATH(second.getTypeSize(SingleQueryType::get(&ctx)), "repeated call" ); |
478 | } |
479 | |
480 | TEST(DataLayout, CacheInvalidation) { |
481 | const char *ir = R"MLIR( |
482 | "dltest.op_with_layout"() { dltest.layout = #dltest.spec< |
483 | #dlti.dl_entry<i42, 5>, |
484 | #dlti.dl_entry<i16, 6> |
485 | > } : () -> () |
486 | )MLIR" ; |
487 | |
488 | DialectRegistry registry; |
489 | registry.insert<DLTIDialect, DLTestDialect>(); |
490 | MLIRContext ctx(registry); |
491 | |
492 | OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(sourceStr: ir, config: &ctx); |
493 | auto op = |
494 | cast<DataLayoutOpInterface>(module->getBody()->getOperations().front()); |
495 | DataLayout layout(op); |
496 | |
497 | // Normal query is fine. |
498 | EXPECT_EQ(layout.getTypeSize(Float16Type::get(&ctx)), 6u); |
499 | |
500 | // Replace the data layout spec with a new, empty spec. |
501 | op->setAttr(kAttrName, CustomDataLayoutSpec::get(ctx: &ctx, entries: {})); |
502 | |
503 | // Data layout is no longer valid and should trigger assertion when queried. |
504 | #ifndef NDEBUG |
505 | ASSERT_DEATH(layout.getTypeSize(Float16Type::get(&ctx)), "no longer valid" ); |
506 | #endif |
507 | } |
508 | |
509 | TEST(DataLayout, UnimplementedTypeInterface) { |
510 | const char *ir = R"MLIR( |
511 | "dltest.op_with_layout"() { dltest.layout = #dltest.spec<> } : () -> () |
512 | )MLIR" ; |
513 | |
514 | DialectRegistry registry; |
515 | registry.insert<DLTIDialect, DLTestDialect>(); |
516 | MLIRContext ctx(registry); |
517 | |
518 | OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(sourceStr: ir, config: &ctx); |
519 | auto op = |
520 | cast<DataLayoutOpInterface>(module->getBody()->getOperations().front()); |
521 | DataLayout layout(op); |
522 | |
523 | ASSERT_DEATH(layout.getTypeSize(TypeNoLayout::get(&ctx)), |
524 | "neither the scoping op nor the type class provide data layout " |
525 | "information" ); |
526 | } |
527 | |
528 | TEST(DataLayout, SevenBitByte) { |
529 | const char *ir = R"MLIR( |
530 | "dltest.op_with_7bit_byte"() { dltest.layout = #dltest.spec<> } : () -> () |
531 | )MLIR" ; |
532 | |
533 | DialectRegistry registry; |
534 | registry.insert<DLTIDialect, DLTestDialect>(); |
535 | MLIRContext ctx(registry); |
536 | |
537 | OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(sourceStr: ir, config: &ctx); |
538 | auto op = |
539 | cast<DataLayoutOpInterface>(module->getBody()->getOperations().front()); |
540 | DataLayout layout(op); |
541 | |
542 | EXPECT_EQ(layout.getTypeSizeInBits(IntegerType::get(&ctx, 42)), 42u); |
543 | EXPECT_EQ(layout.getTypeSizeInBits(IntegerType::get(&ctx, 32)), 32u); |
544 | EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 42)), 6u); |
545 | EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 32)), 5u); |
546 | } |
547 | |