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
21using namespace mlir;
22
23namespace {
24constexpr static llvm::StringLiteral kAttrName = "dltest.layout";
25constexpr static llvm::StringLiteral kEndiannesKeyName = "dltest.endianness";
26constexpr static llvm::StringLiteral kDefaultKeyName =
27 "dltest.default_memory_space";
28constexpr static llvm::StringLiteral kAllocaKeyName =
29 "dltest.alloca_memory_space";
30constexpr static llvm::StringLiteral kManglingModeKeyName =
31 "dltest.mangling_mode";
32constexpr static llvm::StringLiteral kProgramKeyName =
33 "dltest.program_memory_space";
34constexpr static llvm::StringLiteral kGlobalKeyName =
35 "dltest.global_memory_space";
36constexpr static llvm::StringLiteral kStackAlignmentKeyName =
37 "dltest.stack_alignment";
38constexpr static llvm::StringLiteral kFunctionPointerAlignmentKeyName =
39 "dltest.function_pointer_alignment";
40constexpr static llvm::StringLiteral kLegalIntWidthsKeyName =
41 "dltest.legal_int_widths";
42
43constexpr static llvm::StringLiteral kTargetSystemDescAttrName =
44 "dl_target_sys_desc_test.target_system_spec";
45
46/// Trivial array storage for the custom data layout spec attribute, just a list
47/// of entries.
48class DataLayoutSpecStorage : public AttributeStorage {
49public:
50 using KeyTy = ArrayRef<DataLayoutEntryInterface>;
51
52 DataLayoutSpecStorage(ArrayRef<DataLayoutEntryInterface> entries)
53 : entries(entries) {}
54
55 bool operator==(const KeyTy &key) const { return key == entries; }
56
57 static DataLayoutSpecStorage *construct(AttributeStorageAllocator &allocator,
58 const KeyTy &key) {
59 return new (allocator.allocate<DataLayoutSpecStorage>())
60 DataLayoutSpecStorage(allocator.copyInto(elements: key));
61 }
62
63 ArrayRef<DataLayoutEntryInterface> entries;
64};
65
66/// Simple data layout spec containing a list of entries that always verifies
67/// as valid.
68struct CustomDataLayoutSpec
69 : public Attribute::AttrBase<
70 CustomDataLayoutSpec, Attribute, DataLayoutSpecStorage,
71 DLTIQueryInterface::Trait, DataLayoutSpecInterface::Trait> {
72 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CustomDataLayoutSpec)
73
74 using Base::Base;
75
76 static constexpr StringLiteral name = "test.custom_data_layout_spec";
77
78 static CustomDataLayoutSpec get(MLIRContext *ctx,
79 ArrayRef<DataLayoutEntryInterface> entries) {
80 return Base::get(ctx, entries);
81 }
82 CustomDataLayoutSpec
83 combineWith(ArrayRef<DataLayoutSpecInterface> specs) const {
84 return *this;
85 }
86 DataLayoutEntryListRef getEntries() const { return getImpl()->entries; }
87 LogicalResult verifySpec(Location loc) { return success(); }
88 StringAttr getEndiannessIdentifier(MLIRContext *context) const {
89 return Builder(context).getStringAttr(kEndiannesKeyName);
90 }
91 StringAttr getDefaultMemorySpaceIdentifier(MLIRContext *context) const {
92 return Builder(context).getStringAttr(kDefaultKeyName);
93 }
94 StringAttr getAllocaMemorySpaceIdentifier(MLIRContext *context) const {
95 return Builder(context).getStringAttr(kAllocaKeyName);
96 }
97 StringAttr getManglingModeIdentifier(MLIRContext *context) const {
98 return Builder(context).getStringAttr(kManglingModeKeyName);
99 }
100 StringAttr getProgramMemorySpaceIdentifier(MLIRContext *context) const {
101 return Builder(context).getStringAttr(kProgramKeyName);
102 }
103 StringAttr getGlobalMemorySpaceIdentifier(MLIRContext *context) const {
104 return Builder(context).getStringAttr(kGlobalKeyName);
105 }
106 StringAttr getStackAlignmentIdentifier(MLIRContext *context) const {
107 return Builder(context).getStringAttr(kStackAlignmentKeyName);
108 }
109 StringAttr getFunctionPointerAlignmentIdentifier(MLIRContext *context) const {
110 return Builder(context).getStringAttr(kFunctionPointerAlignmentKeyName);
111 }
112 StringAttr getLegalIntWidthsIdentifier(MLIRContext *context) const {
113 return Builder(context).getStringAttr(kLegalIntWidthsKeyName);
114 }
115 FailureOr<Attribute> query(DataLayoutEntryKey key) const {
116 return llvm::cast<mlir::DataLayoutSpecInterface>(Val: *this).queryHelper(key);
117 }
118};
119
120class TargetSystemSpecStorage : public AttributeStorage {
121public:
122 using KeyTy = ArrayRef<DataLayoutEntryInterface>;
123
124 TargetSystemSpecStorage(ArrayRef<DataLayoutEntryInterface> entries)
125 : entries(entries) {}
126
127 bool operator==(const KeyTy &key) const { return key == entries; }
128
129 static TargetSystemSpecStorage *
130 construct(AttributeStorageAllocator &allocator, const KeyTy &key) {
131 return new (allocator.allocate<TargetSystemSpecStorage>())
132 TargetSystemSpecStorage(allocator.copyInto(elements: key));
133 }
134
135 ArrayRef<DataLayoutEntryInterface> entries;
136};
137
138struct CustomTargetSystemSpec
139 : public Attribute::AttrBase<
140 CustomTargetSystemSpec, Attribute, TargetSystemSpecStorage,
141 DLTIQueryInterface::Trait, TargetSystemSpecInterface::Trait> {
142 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CustomDataLayoutSpec)
143
144 using Base::Base;
145
146 static constexpr StringLiteral name = "test.custom_target_system_spec";
147
148 static CustomTargetSystemSpec
149 get(MLIRContext *ctx, ArrayRef<DataLayoutEntryInterface> entries) {
150 return Base::get(ctx, entries);
151 }
152 ArrayRef<DataLayoutEntryInterface> getEntries() const {
153 return getImpl()->entries;
154 }
155 LogicalResult verifySpec(Location loc) { return success(); }
156 std::optional<TargetDeviceSpecInterface>
157 getDeviceSpecForDeviceID(TargetSystemSpecInterface::DeviceID deviceID) {
158 for (const auto &entry : getEntries()) {
159 if (entry.getKey() == DataLayoutEntryKey(deviceID))
160 if (auto deviceSpec =
161 llvm::dyn_cast<TargetDeviceSpecInterface>(entry.getValue()))
162 return deviceSpec;
163 }
164 return std::nullopt;
165 }
166 FailureOr<Attribute> query(DataLayoutEntryKey key) const {
167 return llvm::cast<mlir::TargetSystemSpecInterface>(Val: *this).queryHelper(key);
168 }
169};
170
171/// A type subject to data layout that exits the program if it is queried more
172/// than once. Handy to check if the cache works.
173struct SingleQueryType
174 : public Type::TypeBase<SingleQueryType, Type, TypeStorage,
175 DataLayoutTypeInterface::Trait> {
176 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SingleQueryType)
177
178 using Base::Base;
179
180 static constexpr StringLiteral name = "test.single_query";
181
182 static SingleQueryType get(MLIRContext *ctx) { return Base::get(ctx); }
183
184 llvm::TypeSize getTypeSizeInBits(const DataLayout &layout,
185 DataLayoutEntryListRef params) const {
186 static bool executed = false;
187 if (executed)
188 llvm::report_fatal_error(reason: "repeated call");
189
190 executed = true;
191 return llvm::TypeSize::getFixed(ExactSize: 1);
192 }
193
194 uint64_t getABIAlignment(const DataLayout &layout,
195 DataLayoutEntryListRef params) {
196 static bool executed = false;
197 if (executed)
198 llvm::report_fatal_error(reason: "repeated call");
199
200 executed = true;
201 return 2;
202 }
203
204 uint64_t getPreferredAlignment(const DataLayout &layout,
205 DataLayoutEntryListRef params) {
206 static bool executed = false;
207 if (executed)
208 llvm::report_fatal_error(reason: "repeated call");
209
210 executed = true;
211 return 4;
212 }
213
214 Attribute getEndianness(DataLayoutEntryInterface entry) {
215 static bool executed = false;
216 if (executed)
217 llvm::report_fatal_error(reason: "repeated call");
218
219 executed = true;
220 return Attribute();
221 }
222
223 Attribute getDefaultMemorySpace(DataLayoutEntryInterface entry) {
224 static bool executed = false;
225 if (executed)
226 llvm::report_fatal_error(reason: "repeated call");
227
228 executed = true;
229 return Attribute();
230 }
231
232 Attribute getAllocaMemorySpace(DataLayoutEntryInterface entry) {
233 static bool executed = false;
234 if (executed)
235 llvm::report_fatal_error(reason: "repeated call");
236
237 executed = true;
238 return Attribute();
239 }
240
241 Attribute getProgramMemorySpace(DataLayoutEntryInterface entry) {
242 static bool executed = false;
243 if (executed)
244 llvm::report_fatal_error(reason: "repeated call");
245
246 executed = true;
247 return Attribute();
248 }
249
250 Attribute getGlobalMemorySpace(DataLayoutEntryInterface entry) {
251 static bool executed = false;
252 if (executed)
253 llvm::report_fatal_error(reason: "repeated call");
254
255 executed = true;
256 return Attribute();
257 }
258};
259
260/// A types that is not subject to data layout.
261struct TypeNoLayout : public Type::TypeBase<TypeNoLayout, Type, TypeStorage> {
262 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TypeNoLayout)
263
264 using Base::Base;
265
266 static constexpr StringLiteral name = "test.no_layout";
267
268 static TypeNoLayout get(MLIRContext *ctx) { return Base::get(ctx); }
269};
270
271/// An op that serves as scope for data layout queries with the relevant
272/// attribute attached. This can handle data layout requests for the built-in
273/// types itself.
274struct OpWithLayout : public Op<OpWithLayout, DataLayoutOpInterface::Trait> {
275 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OpWithLayout)
276
277 using Op::Op;
278 static ArrayRef<StringRef> getAttributeNames() { return {}; }
279
280 static StringRef getOperationName() { return "dltest.op_with_layout"; }
281
282 DataLayoutSpecInterface getDataLayoutSpec() {
283 return getOperation()->getAttrOfType<DataLayoutSpecInterface>(kAttrName);
284 }
285
286 TargetSystemSpecInterface getTargetSystemSpec() {
287 return getOperation()->getAttrOfType<TargetSystemSpecInterface>(
288 kTargetSystemDescAttrName);
289 }
290
291 static llvm::TypeSize getTypeSizeInBits(Type type,
292 const DataLayout &dataLayout,
293 DataLayoutEntryListRef params) {
294 // Make a recursive query.
295 if (isa<FloatType>(type))
296 return dataLayout.getTypeSizeInBits(
297 IntegerType::get(type.getContext(), type.getIntOrFloatBitWidth()));
298
299 // Handle built-in types that are not handled by the default process.
300 if (auto iType = dyn_cast<IntegerType>(type)) {
301 for (DataLayoutEntryInterface entry : params)
302 if (llvm::dyn_cast_if_present<Type>(entry.getKey()) == type)
303 return llvm::TypeSize::getFixed(
304 8 *
305 cast<IntegerAttr>(entry.getValue()).getValue().getZExtValue());
306 return llvm::TypeSize::getFixed(ExactSize: 8 * iType.getIntOrFloatBitWidth());
307 }
308
309 // Use the default process for everything else.
310 return detail::getDefaultTypeSize(type, dataLayout, params);
311 }
312
313 static uint64_t getTypeABIAlignment(Type type, const DataLayout &dataLayout,
314 DataLayoutEntryListRef params) {
315 return llvm::PowerOf2Ceil(A: getTypeSize(type, dataLayout, params));
316 }
317
318 static uint64_t getTypePreferredAlignment(Type type,
319 const DataLayout &dataLayout,
320 DataLayoutEntryListRef params) {
321 return 2 * getTypeABIAlignment(type, dataLayout, params);
322 }
323};
324
325struct OpWith7BitByte
326 : public Op<OpWith7BitByte, DataLayoutOpInterface::Trait> {
327 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OpWith7BitByte)
328
329 using Op::Op;
330 static ArrayRef<StringRef> getAttributeNames() { return {}; }
331
332 static StringRef getOperationName() { return "dltest.op_with_7bit_byte"; }
333
334 DataLayoutSpecInterface getDataLayoutSpec() {
335 return getOperation()->getAttrOfType<DataLayoutSpecInterface>(kAttrName);
336 }
337
338 TargetSystemSpecInterface getTargetSystemSpec() {
339 return getOperation()->getAttrOfType<TargetSystemSpecInterface>(
340 kTargetSystemDescAttrName);
341 }
342
343 // Bytes are assumed to be 7-bit here.
344 static llvm::TypeSize getTypeSize(Type type, const DataLayout &dataLayout,
345 DataLayoutEntryListRef params) {
346 return mlir::detail::divideCeil(numerator: dataLayout.getTypeSizeInBits(t: type), denominator: 7);
347 }
348};
349
350/// A dialect putting all the above together.
351struct DLTestDialect : Dialect {
352 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(DLTestDialect)
353
354 explicit DLTestDialect(MLIRContext *ctx)
355 : Dialect(getDialectNamespace(), ctx, TypeID::get<DLTestDialect>()) {
356 ctx->getOrLoadDialect<DLTIDialect>();
357 addAttributes<CustomDataLayoutSpec>();
358 addOperations<OpWithLayout, OpWith7BitByte>();
359 addTypes<SingleQueryType, TypeNoLayout>();
360 }
361 static StringRef getDialectNamespace() { return "dltest"; }
362
363 void printAttribute(Attribute attr,
364 DialectAsmPrinter &printer) const override {
365 printer << "spec<";
366 llvm::interleaveComma(c: cast<CustomDataLayoutSpec>(Val&: attr).getEntries(),
367 os&: printer);
368 printer << ">";
369 }
370
371 Attribute parseAttribute(DialectAsmParser &parser, Type type) const override {
372 bool ok =
373 succeeded(Result: parser.parseKeyword(keyword: "spec")) && succeeded(Result: parser.parseLess());
374 (void)ok;
375 assert(ok);
376 if (succeeded(Result: parser.parseOptionalGreater()))
377 return CustomDataLayoutSpec::get(ctx: parser.getContext(), entries: {});
378
379 SmallVector<DataLayoutEntryInterface> entries;
380 ok = succeeded(Result: parser.parseCommaSeparatedList(parseElementFn: [&]() {
381 entries.emplace_back();
382 ok = succeeded(parser.parseAttribute(entries.back()));
383 assert(ok);
384 return success();
385 }));
386 assert(ok);
387 ok = succeeded(Result: parser.parseGreater());
388 assert(ok);
389 return CustomDataLayoutSpec::get(ctx: parser.getContext(), entries: entries);
390 }
391
392 void printType(Type type, DialectAsmPrinter &printer) const override {
393 if (isa<SingleQueryType>(Val: type))
394 printer << "single_query";
395 else
396 printer << "no_layout";
397 }
398
399 Type parseType(DialectAsmParser &parser) const override {
400 bool ok = succeeded(Result: parser.parseKeyword(keyword: "single_query"));
401 (void)ok;
402 assert(ok);
403 return SingleQueryType::get(ctx: parser.getContext());
404 }
405};
406
407/// A dialect to test DLTI's target system spec and related attributes
408struct DLTargetSystemDescTestDialect : public Dialect {
409 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(DLTargetSystemDescTestDialect)
410
411 explicit DLTargetSystemDescTestDialect(MLIRContext *ctx)
412 : Dialect(getDialectNamespace(), ctx,
413 TypeID::get<DLTargetSystemDescTestDialect>()) {
414 ctx->getOrLoadDialect<DLTIDialect>();
415 addAttributes<CustomTargetSystemSpec>();
416 }
417 static StringRef getDialectNamespace() { return "dl_target_sys_desc_test"; }
418
419 void printAttribute(Attribute attr,
420 DialectAsmPrinter &printer) const override {
421 printer << "target_system_spec<";
422 llvm::interleaveComma(c: cast<CustomTargetSystemSpec>(Val&: attr).getEntries(),
423 os&: printer, each_fn: [&](const auto &it) {
424 printer << dyn_cast<StringAttr>(it.getKey()) << ":"
425 << it.getValue();
426 });
427 printer << ">";
428 }
429
430 Attribute parseAttribute(DialectAsmParser &parser, Type type) const override {
431 bool ok = succeeded(Result: parser.parseKeyword(keyword: "target_system_spec")) &&
432 succeeded(Result: parser.parseLess());
433 (void)ok;
434 assert(ok);
435 if (succeeded(Result: parser.parseOptionalGreater()))
436 return CustomTargetSystemSpec::get(ctx: parser.getContext(), entries: {});
437
438 auto parseTargetDeviceSpecEntry =
439 [&](AsmParser &parser) -> FailureOr<TargetDeviceSpecEntry> {
440 std::string deviceID;
441 if (failed(Result: parser.parseString(string: &deviceID))) {
442 parser.emitError(loc: parser.getCurrentLocation())
443 << "DeviceID is missing, or is not of string type";
444 return failure();
445 }
446 if (failed(Result: parser.parseColon())) {
447 parser.emitError(loc: parser.getCurrentLocation()) << "Missing colon";
448 return failure();
449 }
450
451 TargetDeviceSpecInterface targetDeviceSpec;
452 if (failed(parser.parseAttribute(targetDeviceSpec))) {
453 parser.emitError(loc: parser.getCurrentLocation())
454 << "Error in parsing target device spec";
455 return failure();
456 }
457 return std::make_pair(parser.getBuilder().getStringAttr(deviceID),
458 targetDeviceSpec);
459 };
460
461 SmallVector<DataLayoutEntryInterface> entries;
462 ok = succeeded(Result: parser.parseCommaSeparatedList(parseElementFn: [&]() {
463 auto deviceIDAndTargetDeviceSpecPair = parseTargetDeviceSpecEntry(parser);
464 ok = succeeded(Result: deviceIDAndTargetDeviceSpecPair);
465 assert(ok);
466 auto entry =
467 DataLayoutEntryAttr::get(deviceIDAndTargetDeviceSpecPair->first,
468 deviceIDAndTargetDeviceSpecPair->second);
469 entries.push_back(entry);
470 return success();
471 }));
472 assert(ok);
473 ok = succeeded(Result: parser.parseGreater());
474 assert(ok);
475 return CustomTargetSystemSpec::get(ctx: parser.getContext(), entries: entries);
476 }
477};
478
479} // namespace
480
481TEST(DataLayout, FallbackDefault) {
482 const char *ir = R"MLIR(
483module {}
484 )MLIR";
485
486 DialectRegistry registry;
487 registry.insert<DLTIDialect, DLTestDialect>();
488 MLIRContext ctx(registry);
489
490 OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(ir, &ctx);
491 DataLayout layout(module.get());
492 EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 42)), 6u);
493 EXPECT_EQ(layout.getTypeSize(Float16Type::get(&ctx)), 2u);
494 EXPECT_EQ(layout.getTypeSizeInBits(IntegerType::get(&ctx, 42)), 42u);
495 EXPECT_EQ(layout.getTypeSizeInBits(Float16Type::get(&ctx)), 16u);
496 EXPECT_EQ(layout.getTypeABIAlignment(IntegerType::get(&ctx, 42)), 8u);
497 EXPECT_EQ(layout.getTypeABIAlignment(Float16Type::get(&ctx)), 2u);
498 EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 42)), 8u);
499 EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 2u);
500
501 EXPECT_EQ(layout.getEndianness(), Attribute());
502 EXPECT_EQ(layout.getDefaultMemorySpace(), Attribute());
503 EXPECT_EQ(layout.getAllocaMemorySpace(), Attribute());
504 EXPECT_EQ(layout.getProgramMemorySpace(), Attribute());
505 EXPECT_EQ(layout.getGlobalMemorySpace(), Attribute());
506 EXPECT_EQ(layout.getStackAlignment(), 0u);
507 EXPECT_EQ(layout.getFunctionPointerAlignment(), Attribute());
508 EXPECT_EQ(layout.getLegalIntWidths(), Attribute());
509 EXPECT_EQ(layout.getManglingMode(), Attribute());
510}
511
512TEST(DataLayout, NullSpec) {
513 const char *ir = R"MLIR(
514"dltest.op_with_layout"() : () -> ()
515 )MLIR";
516
517 DialectRegistry registry;
518 registry.insert<DLTIDialect, DLTestDialect>();
519 MLIRContext ctx(registry);
520
521 OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(sourceStr: ir, config: &ctx);
522 auto op =
523 cast<DataLayoutOpInterface>(module->getBody()->getOperations().front());
524 DataLayout layout(op);
525
526 EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 42)), 42u);
527 EXPECT_EQ(layout.getTypeSize(Float16Type::get(&ctx)), 16u);
528 EXPECT_EQ(layout.getTypeSizeInBits(IntegerType::get(&ctx, 42)), 8u * 42u);
529 EXPECT_EQ(layout.getTypeSizeInBits(Float16Type::get(&ctx)), 8u * 16u);
530 EXPECT_EQ(layout.getTypeABIAlignment(IntegerType::get(&ctx, 42)), 64u);
531 EXPECT_EQ(layout.getTypeABIAlignment(Float16Type::get(&ctx)), 16u);
532 EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 42)), 128u);
533 EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 32u);
534 EXPECT_EQ(layout.getTypeIndexBitwidth(Float16Type::get(&ctx)), std::nullopt);
535 EXPECT_EQ(layout.getTypeIndexBitwidth(IndexType::get(&ctx)), 64u);
536
537 EXPECT_EQ(layout.getEndianness(), Attribute());
538 EXPECT_EQ(layout.getDefaultMemorySpace(), Attribute());
539 EXPECT_EQ(layout.getAllocaMemorySpace(), Attribute());
540 EXPECT_EQ(layout.getProgramMemorySpace(), Attribute());
541 EXPECT_EQ(layout.getGlobalMemorySpace(), Attribute());
542 EXPECT_EQ(layout.getStackAlignment(), 0u);
543 EXPECT_EQ(layout.getManglingMode(), Attribute());
544
545 EXPECT_EQ(layout.getDevicePropertyValue(
546 Builder(&ctx).getStringAttr("CPU" /* device ID*/),
547 Builder(&ctx).getStringAttr("L1_cache_size_in_bytes")),
548 std::nullopt);
549 EXPECT_EQ(layout.getDevicePropertyValue(
550 Builder(&ctx).getStringAttr("CPU" /* device ID*/),
551 Builder(&ctx).getStringAttr("max_vector_width")),
552 std::nullopt);
553}
554
555TEST(DataLayout, EmptySpec) {
556 const char *ir = R"MLIR(
557"dltest.op_with_layout"() { dltest.layout = #dltest.spec< > } : () -> ()
558 )MLIR";
559
560 DialectRegistry registry;
561 registry.insert<DLTIDialect, DLTestDialect>();
562 MLIRContext ctx(registry);
563
564 OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(sourceStr: ir, config: &ctx);
565 auto op =
566 cast<DataLayoutOpInterface>(module->getBody()->getOperations().front());
567 DataLayout layout(op);
568 EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 42)), 42u);
569 EXPECT_EQ(layout.getTypeSize(Float16Type::get(&ctx)), 16u);
570 EXPECT_EQ(layout.getTypeSizeInBits(IntegerType::get(&ctx, 42)), 8u * 42u);
571 EXPECT_EQ(layout.getTypeSizeInBits(Float16Type::get(&ctx)), 8u * 16u);
572 EXPECT_EQ(layout.getTypeABIAlignment(IntegerType::get(&ctx, 42)), 64u);
573 EXPECT_EQ(layout.getTypeABIAlignment(Float16Type::get(&ctx)), 16u);
574 EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 42)), 128u);
575 EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 32u);
576 EXPECT_EQ(layout.getTypeIndexBitwidth(Float16Type::get(&ctx)), std::nullopt);
577 EXPECT_EQ(layout.getTypeIndexBitwidth(IndexType::get(&ctx)), 64u);
578
579 EXPECT_EQ(layout.getEndianness(), Attribute());
580 EXPECT_EQ(layout.getDefaultMemorySpace(), Attribute());
581 EXPECT_EQ(layout.getAllocaMemorySpace(), Attribute());
582 EXPECT_EQ(layout.getProgramMemorySpace(), Attribute());
583 EXPECT_EQ(layout.getGlobalMemorySpace(), Attribute());
584 EXPECT_EQ(layout.getStackAlignment(), 0u);
585 EXPECT_EQ(layout.getManglingMode(), Attribute());
586 EXPECT_EQ(layout.getFunctionPointerAlignment(), Attribute());
587 EXPECT_EQ(layout.getLegalIntWidths(), Attribute());
588
589 EXPECT_EQ(layout.getDevicePropertyValue(
590 Builder(&ctx).getStringAttr("CPU" /* device ID*/),
591 Builder(&ctx).getStringAttr("L1_cache_size_in_bytes")),
592 std::nullopt);
593 EXPECT_EQ(layout.getDevicePropertyValue(
594 Builder(&ctx).getStringAttr("CPU" /* device ID*/),
595 Builder(&ctx).getStringAttr("max_vector_width")),
596 std::nullopt);
597}
598
599TEST(DataLayout, SpecWithEntries) {
600 const char *ir = R"MLIR(
601"dltest.op_with_layout"() { dltest.layout = #dltest.spec<
602 #dlti.dl_entry<i42, 5>,
603 #dlti.dl_entry<i16, 6>,
604 #dlti.dl_entry<index, 42>,
605 #dlti.dl_entry<"dltest.endianness", "little">,
606 #dlti.dl_entry<"dltest.default_memory_space", 1 : i32>,
607 #dlti.dl_entry<"dltest.alloca_memory_space", 5 : i32>,
608 #dlti.dl_entry<"dltest.program_memory_space", 3 : i32>,
609 #dlti.dl_entry<"dltest.global_memory_space", 2 : i32>,
610 #dlti.dl_entry<"dltest.stack_alignment", 128 : i32>,
611 #dlti.dl_entry<"dltest.mangling_mode", "o">,
612 #dlti.dl_entry<"dltest.function_pointer_alignment",
613 #dlti.function_pointer_alignment<64, function_dependent = true>>,
614 #dlti.dl_entry<"dltest.legal_int_widths", array<i32: 64>>
615> } : () -> ()
616 )MLIR";
617
618 DialectRegistry registry;
619 registry.insert<DLTIDialect, DLTestDialect>();
620 MLIRContext ctx(registry);
621
622 OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(sourceStr: ir, config: &ctx);
623 auto op =
624 cast<DataLayoutOpInterface>(module->getBody()->getOperations().front());
625 DataLayout layout(op);
626 EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 42)), 5u);
627 EXPECT_EQ(layout.getTypeSize(Float16Type::get(&ctx)), 6u);
628 EXPECT_EQ(layout.getTypeSizeInBits(IntegerType::get(&ctx, 42)), 40u);
629 EXPECT_EQ(layout.getTypeSizeInBits(Float16Type::get(&ctx)), 48u);
630 EXPECT_EQ(layout.getTypeABIAlignment(IntegerType::get(&ctx, 42)), 8u);
631 EXPECT_EQ(layout.getTypeABIAlignment(Float16Type::get(&ctx)), 8u);
632 EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 42)), 16u);
633 EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 16u);
634 EXPECT_EQ(layout.getTypeIndexBitwidth(Float16Type::get(&ctx)), std::nullopt);
635 EXPECT_EQ(layout.getTypeIndexBitwidth(IndexType::get(&ctx)), 42u);
636
637 EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 32)), 32u);
638 EXPECT_EQ(layout.getTypeSize(Float32Type::get(&ctx)), 32u);
639 EXPECT_EQ(layout.getTypeSizeInBits(IntegerType::get(&ctx, 32)), 256u);
640 EXPECT_EQ(layout.getTypeSizeInBits(Float32Type::get(&ctx)), 256u);
641 EXPECT_EQ(layout.getTypeABIAlignment(IntegerType::get(&ctx, 32)), 32u);
642 EXPECT_EQ(layout.getTypeABIAlignment(Float32Type::get(&ctx)), 32u);
643 EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 32)), 64u);
644 EXPECT_EQ(layout.getTypePreferredAlignment(Float32Type::get(&ctx)), 64u);
645
646 EXPECT_EQ(layout.getEndianness(), Builder(&ctx).getStringAttr("little"));
647 EXPECT_EQ(layout.getDefaultMemorySpace(), Builder(&ctx).getI32IntegerAttr(1));
648 EXPECT_EQ(layout.getAllocaMemorySpace(), Builder(&ctx).getI32IntegerAttr(5));
649 EXPECT_EQ(layout.getProgramMemorySpace(), Builder(&ctx).getI32IntegerAttr(3));
650 EXPECT_EQ(layout.getGlobalMemorySpace(), Builder(&ctx).getI32IntegerAttr(2));
651 EXPECT_EQ(layout.getStackAlignment(), 128u);
652 EXPECT_EQ(layout.getManglingMode(), Builder(&ctx).getStringAttr("o"));
653 EXPECT_EQ(
654 layout.getFunctionPointerAlignment(),
655 FunctionPointerAlignmentAttr::get(&ctx, 64, /*function_dependent=*/true));
656 EXPECT_EQ(layout.getLegalIntWidths(),
657 Builder(&ctx).getDenseI32ArrayAttr({64}));
658}
659
660TEST(DataLayout, SpecWithTargetSystemDescEntries) {
661 const char *ir = R"MLIR(
662 module attributes { dl_target_sys_desc_test.target_system_spec =
663 #dl_target_sys_desc_test.target_system_spec<
664 "CPU": #dlti.target_device_spec<
665 #dlti.dl_entry<"L1_cache_size_in_bytes", "4096">,
666 #dlti.dl_entry<"max_vector_op_width", "128">>
667 > } {}
668 )MLIR";
669
670 DialectRegistry registry;
671 registry.insert<DLTIDialect, DLTargetSystemDescTestDialect>();
672 MLIRContext ctx(registry);
673
674 OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(sourceStr: ir, config: &ctx);
675 DataLayout layout(*module);
676 EXPECT_EQ(layout.getDevicePropertyValue(
677 Builder(&ctx).getStringAttr("CPU") /* device ID*/,
678 Builder(&ctx).getStringAttr("L1_cache_size_in_bytes")),
679 std::optional<Attribute>(Builder(&ctx).getStringAttr("4096")));
680 EXPECT_EQ(layout.getDevicePropertyValue(
681 Builder(&ctx).getStringAttr("CPU") /* device ID*/,
682 Builder(&ctx).getStringAttr("max_vector_op_width")),
683 std::optional<Attribute>(Builder(&ctx).getStringAttr("128")));
684}
685
686TEST(DataLayout, Caching) {
687 const char *ir = R"MLIR(
688"dltest.op_with_layout"() { dltest.layout = #dltest.spec<> } : () -> ()
689 )MLIR";
690
691 DialectRegistry registry;
692 registry.insert<DLTIDialect, DLTestDialect>();
693 MLIRContext ctx(registry);
694
695 OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(sourceStr: ir, config: &ctx);
696 auto op =
697 cast<DataLayoutOpInterface>(module->getBody()->getOperations().front());
698 DataLayout layout(op);
699
700 unsigned sum = 0;
701 sum += layout.getTypeSize(SingleQueryType::get(ctx: &ctx));
702 // The second call should hit the cache. If it does not, the function in
703 // SingleQueryType will be called and will abort the process.
704 sum += layout.getTypeSize(SingleQueryType::get(ctx: &ctx));
705 // Make sure the complier doesn't optimize away the query code.
706 EXPECT_EQ(sum, 2u);
707
708 // A fresh data layout has a new cache, so the call to it should be dispatched
709 // down to the type and abort the process.
710 DataLayout second(op);
711 ASSERT_DEATH(second.getTypeSize(SingleQueryType::get(&ctx)), "repeated call");
712}
713
714TEST(DataLayout, CacheInvalidation) {
715 const char *ir = R"MLIR(
716"dltest.op_with_layout"() { dltest.layout = #dltest.spec<
717 #dlti.dl_entry<i42, 5>,
718 #dlti.dl_entry<i16, 6>
719> } : () -> ()
720 )MLIR";
721
722 DialectRegistry registry;
723 registry.insert<DLTIDialect, DLTestDialect>();
724 MLIRContext ctx(registry);
725
726 OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(sourceStr: ir, config: &ctx);
727 auto op =
728 cast<DataLayoutOpInterface>(module->getBody()->getOperations().front());
729 DataLayout layout(op);
730
731 // Normal query is fine.
732 EXPECT_EQ(layout.getTypeSize(Float16Type::get(&ctx)), 6u);
733
734 // Replace the data layout spec with a new, empty spec.
735 op->setAttr(kAttrName, CustomDataLayoutSpec::get(ctx: &ctx, entries: {}));
736
737 // Data layout is no longer valid and should trigger assertion when queried.
738#ifndef NDEBUG
739 ASSERT_DEATH(layout.getTypeSize(Float16Type::get(&ctx)), "no longer valid");
740#endif
741}
742
743TEST(DataLayout, UnimplementedTypeInterface) {
744 const char *ir = R"MLIR(
745"dltest.op_with_layout"() { dltest.layout = #dltest.spec<> } : () -> ()
746 )MLIR";
747
748 DialectRegistry registry;
749 registry.insert<DLTIDialect, DLTestDialect>();
750 MLIRContext ctx(registry);
751
752 OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(sourceStr: ir, config: &ctx);
753 auto op =
754 cast<DataLayoutOpInterface>(module->getBody()->getOperations().front());
755 DataLayout layout(op);
756
757 ASSERT_DEATH(layout.getTypeSize(TypeNoLayout::get(&ctx)),
758 "neither the scoping op nor the type class provide data layout "
759 "information");
760}
761
762TEST(DataLayout, SevenBitByte) {
763 const char *ir = R"MLIR(
764"dltest.op_with_7bit_byte"() { dltest.layout = #dltest.spec<> } : () -> ()
765 )MLIR";
766
767 DialectRegistry registry;
768 registry.insert<DLTIDialect, DLTestDialect>();
769 MLIRContext ctx(registry);
770
771 OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(sourceStr: ir, config: &ctx);
772 auto op =
773 cast<DataLayoutOpInterface>(module->getBody()->getOperations().front());
774 DataLayout layout(op);
775
776 EXPECT_EQ(layout.getTypeSizeInBits(IntegerType::get(&ctx, 42)), 42u);
777 EXPECT_EQ(layout.getTypeSizeInBits(IntegerType::get(&ctx, 32)), 32u);
778 EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 42)), 6u);
779 EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 32)), 5u);
780}
781

source code of mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp