781 lines · cpp
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-exception6//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 list47/// 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(key));61 }62 63 ArrayRef<DataLayoutEntryInterface> entries;64};65 66/// Simple data layout spec containing a list of entries that always verifies67/// as valid.68struct CustomDataLayoutSpec69 : 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 CustomDataLayoutSpec83 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>(*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(key));133 }134 135 ArrayRef<DataLayoutEntryInterface> entries;136};137 138struct CustomTargetSystemSpec139 : 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 CustomTargetSystemSpec149 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>(*this).queryHelper(key);168 }169};170 171/// A type subject to data layout that exits the program if it is queried more172/// than once. Handy to check if the cache works.173struct SingleQueryType174 : 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("repeated call");189 190 executed = true;191 return llvm::TypeSize::getFixed(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("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("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("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("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("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("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("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 relevant272/// attribute attached. This can handle data layout requests for the built-in273/// 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(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(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 OpWith7BitByte326 : 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(dataLayout.getTypeSizeInBits(type), 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(cast<CustomDataLayoutSpec>(attr).getEntries(),367 printer);368 printer << ">";369 }370 371 Attribute parseAttribute(DialectAsmParser &parser, Type type) const override {372 bool ok =373 succeeded(parser.parseKeyword("spec")) && succeeded(parser.parseLess());374 (void)ok;375 assert(ok);376 if (succeeded(parser.parseOptionalGreater()))377 return CustomDataLayoutSpec::get(parser.getContext(), {});378 379 SmallVector<DataLayoutEntryInterface> entries;380 ok = succeeded(parser.parseCommaSeparatedList([&]() {381 entries.emplace_back();382 ok = succeeded(parser.parseAttribute(entries.back()));383 assert(ok);384 return success();385 }));386 assert(ok);387 ok = succeeded(parser.parseGreater());388 assert(ok);389 return CustomDataLayoutSpec::get(parser.getContext(), entries);390 }391 392 void printType(Type type, DialectAsmPrinter &printer) const override {393 if (isa<SingleQueryType>(type))394 printer << "single_query";395 else396 printer << "no_layout";397 }398 399 Type parseType(DialectAsmParser &parser) const override {400 bool ok = succeeded(parser.parseKeyword("single_query"));401 (void)ok;402 assert(ok);403 return SingleQueryType::get(parser.getContext());404 }405};406 407/// A dialect to test DLTI's target system spec and related attributes408struct 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(cast<CustomTargetSystemSpec>(attr).getEntries(),423 printer, [&](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(parser.parseKeyword("target_system_spec")) &&432 succeeded(parser.parseLess());433 (void)ok;434 assert(ok);435 if (succeeded(parser.parseOptionalGreater()))436 return CustomTargetSystemSpec::get(parser.getContext(), {});437 438 auto parseTargetDeviceSpecEntry =439 [&](AsmParser &parser) -> FailureOr<TargetDeviceSpecEntry> {440 std::string deviceID;441 if (failed(parser.parseString(&deviceID))) {442 parser.emitError(parser.getCurrentLocation())443 << "DeviceID is missing, or is not of string type";444 return failure();445 }446 if (failed(parser.parseColon())) {447 parser.emitError(parser.getCurrentLocation()) << "Missing colon";448 return failure();449 }450 451 TargetDeviceSpecInterface targetDeviceSpec;452 if (failed(parser.parseAttribute(targetDeviceSpec))) {453 parser.emitError(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(parser.parseCommaSeparatedList([&]() {463 auto deviceIDAndTargetDeviceSpecPair = parseTargetDeviceSpecEntry(parser);464 ok = succeeded(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(parser.parseGreater());474 assert(ok);475 return CustomTargetSystemSpec::get(parser.getContext(), entries);476 }477};478 479} // namespace480 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>(ir, &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>(ir, &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>(ir, &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>(ir, &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>(ir, &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));702 // The second call should hit the cache. If it does not, the function in703 // SingleQueryType will be called and will abort the process.704 sum += layout.getTypeSize(SingleQueryType::get(&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 dispatched709 // 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>(ir, &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, {}));736 737 // Data layout is no longer valid and should trigger assertion when queried.738#ifndef NDEBUG739 ASSERT_DEATH(layout.getTypeSize(Float16Type::get(&ctx)), "no longer valid");740#endif741}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>(ir, &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>(ir, &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