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