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