mlir/lib/Interfaces/DataLayoutInterfaces.cpp

   1 //===- DataLayoutInterfaces.cpp - Data Layout Interface Implementation ----===//
   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/IR/BuiltinDialect.h"
  11 #include "mlir/IR/BuiltinOps.h"
  12 #include "mlir/IR/BuiltinTypes.h"
  13 #include "mlir/IR/Operation.h"
  14
  15 #include "llvm/ADT/TypeSwitch.h"
  16 #include "llvm/Support/MathExtras.h"
  17
  18 using namespace mlir;
  19
  20 //===----------------------------------------------------------------------===//
  21 // Default implementations
  22 //===----------------------------------------------------------------------===//
  23
  24 /// Reports that the given type is missing the data layout information and
  25 /// exits.
  26 [[noreturn]] static void reportMissingDataLayout(Type type) {
  27   std::string message;
  28   llvm::raw_string_ostream os(message);
  29   os << "neither the scoping op nor the type class provide data layout "
  30         "information for "
  31      << type;
  32   llvm::report_fatal_error(Twine(message));
  33 }
  34
  35 /// Returns the bitwidth of the index type if specified in the param list.
  36 /// Assumes 64-bit index otherwise.
  37 static uint64_t getIndexBitwidth(DataLayoutEntryListRef params) {
  38   if (params.empty())
  39     return 64;
  40   auto attr = cast<IntegerAttr>(params.front().getValue());
  41   return attr.getValue().getZExtValue();
  42 }
  43
  44 llvm::TypeSize
  45 mlir::detail::getDefaultTypeSize(Type type, const DataLayout &dataLayout,
  46                                  ArrayRef<DataLayoutEntryInterface> params) {
  47   llvm::TypeSize bits = getDefaultTypeSizeInBits(type, dataLayout, params);
  48   return divideCeil(bits, 8);
  49 }
  50
  51 llvm::TypeSize
  52 mlir::detail::getDefaultTypeSizeInBits(Type type, const DataLayout &dataLayout,
  53                                        DataLayoutEntryListRef params) {
  54   if (isa<IntegerType, FloatType>(type))
  55     return llvm::TypeSize::getFixed(type.getIntOrFloatBitWidth());
  56
  57   if (auto ctype = dyn_cast<ComplexType>(type)) {
  58     Type et = ctype.getElementType();
  59     uint64_t innerAlignment =
  60         getDefaultPreferredAlignment(et, dataLayout, params) * 8;
  61     llvm::TypeSize innerSize = getDefaultTypeSizeInBits(et, dataLayout, params);
  62
  63     // Include padding required to align the imaginary value in the complex
  64     // type.
  65     return llvm::alignTo(innerSize, innerAlignment) + innerSize;
  66   }
  67
  68   // Index is an integer of some bitwidth.
  69   if (isa<IndexType>(type))
  70     return dataLayout.getTypeSizeInBits(
  71         IntegerType::get(type.getContext(), getIndexBitwidth(params)));
  72
  73   // Sizes of vector types are rounded up to those of types with closest
  74   // power-of-two number of elements in the innermost dimension. We also assume
  75   // there is no bit-packing at the moment element sizes are taken in bytes and
  76   // multiplied with 8 bits.
  77   // TODO: make this extensible.
  78   if (auto vecType = dyn_cast<VectorType>(type)) {
  79     uint64_t baseSize = vecType.getNumElements() / vecType.getShape().back() *
  80                         llvm::PowerOf2Ceil(vecType.getShape().back()) *
  81                         dataLayout.getTypeSize(vecType.getElementType()) * 8;
  82     return llvm::TypeSize::get(baseSize, vecType.isScalable());
  83   }
  84
  85   if (auto typeInterface = dyn_cast<DataLayoutTypeInterface>(type))
  86     return typeInterface.getTypeSizeInBits(dataLayout, params);
  87
  88   reportMissingDataLayout(type);
  89 }
  90
  91 static DataLayoutEntryInterface
  92 findEntryForIntegerType(IntegerType intType,
  93                         ArrayRef<DataLayoutEntryInterface> params) {
  94   assert(!params.empty() && "expected non-empty parameter list");
  95   std::map<unsigned, DataLayoutEntryInterface> sortedParams;
  96   for (DataLayoutEntryInterface entry : params) {
  97     sortedParams.insert(std::make_pair(
  98         entry.getKey().get<Type>().getIntOrFloatBitWidth(), entry));
  99   }
 100   auto iter = sortedParams.lower_bound(intType.getWidth());
 101   if (iter == sortedParams.end())
 102     iter = std::prev(iter);
 103
 104   return iter->second;
 105 }
 106
 107 constexpr const static uint64_t kDefaultBitsInByte = 8u;
 108
 109 static uint64_t extractABIAlignment(DataLayoutEntryInterface entry) {
 110   auto values =
 111       cast<DenseIntElementsAttr>(entry.getValue()).getValues<uint64_t>();
 112   return static_cast<uint64_t>(*values.begin()) / kDefaultBitsInByte;
 113 }
 114
 115 static uint64_t
 116 getIntegerTypeABIAlignment(IntegerType intType,
 117                            ArrayRef<DataLayoutEntryInterface> params) {
 118   constexpr uint64_t kDefaultSmallIntAlignment = 4u;
 119   constexpr unsigned kSmallIntSize = 64;
 120   if (params.empty()) {
 121     return intType.getWidth() < kSmallIntSize
 122                ? llvm::PowerOf2Ceil(
 123                      llvm::divideCeil(intType.getWidth(), kDefaultBitsInByte))
 124                : kDefaultSmallIntAlignment;
 125   }
 126
 127   return extractABIAlignment(findEntryForIntegerType(intType, params));
 128 }
 129
 130 static uint64_t
 131 getFloatTypeABIAlignment(FloatType fltType, const DataLayout &dataLayout,
 132                          ArrayRef<DataLayoutEntryInterface> params) {
 133   assert(params.size() <= 1 && "at most one data layout entry is expected for "
 134                                "the singleton floating-point type");
 135   if (params.empty())
 136     return llvm::PowerOf2Ceil(dataLayout.getTypeSize(fltType).getFixedValue());
 137   return extractABIAlignment(params[0]);
 138 }
 139
 140 uint64_t mlir::detail::getDefaultABIAlignment(
 141     Type type, const DataLayout &dataLayout,
 142     ArrayRef<DataLayoutEntryInterface> params) {
 143   // Natural alignment is the closest power-of-two number above. For scalable
 144   // vectors, aligning them to the same as the base vector is sufficient.
 145   if (isa<VectorType>(type))
 146     return llvm::PowerOf2Ceil(dataLayout.getTypeSize(type).getKnownMinValue());
 147
 148   if (auto fltType = dyn_cast<FloatType>(type))
 149     return getFloatTypeABIAlignment(fltType, dataLayout, params);
 150
 151   // Index is an integer of some bitwidth.
 152   if (isa<IndexType>(type))
 153     return dataLayout.getTypeABIAlignment(
 154         IntegerType::get(type.getContext(), getIndexBitwidth(params)));
 155
 156   if (auto intType = dyn_cast<IntegerType>(type))
 157     return getIntegerTypeABIAlignment(intType, params);
 158
 159   if (auto ctype = dyn_cast<ComplexType>(type))
 160     return getDefaultABIAlignment(ctype.getElementType(), dataLayout, params);
 161
 162   if (auto typeInterface = dyn_cast<DataLayoutTypeInterface>(type))
 163     return typeInterface.getABIAlignment(dataLayout, params);
 164
 165   reportMissingDataLayout(type);
 166 }
 167
 168 static uint64_t extractPreferredAlignment(DataLayoutEntryInterface entry) {
 169   auto values =
 170       cast<DenseIntElementsAttr>(entry.getValue()).getValues<uint64_t>();
 171   return *std::next(values.begin(), values.size() - 1) / kDefaultBitsInByte;
 172 }
 173
 174 static uint64_t
 175 getIntegerTypePreferredAlignment(IntegerType intType,
 176                                  const DataLayout &dataLayout,
 177                                  ArrayRef<DataLayoutEntryInterface> params) {
 178   if (params.empty())
 179     return llvm::PowerOf2Ceil(dataLayout.getTypeSize(intType).getFixedValue());
 180
 181   return extractPreferredAlignment(findEntryForIntegerType(intType, params));
 182 }
 183
 184 static uint64_t
 185 getFloatTypePreferredAlignment(FloatType fltType, const DataLayout &dataLayout,
 186                                ArrayRef<DataLayoutEntryInterface> params) {
 187   assert(params.size() <= 1 && "at most one data layout entry is expected for "
 188                                "the singleton floating-point type");
 189   if (params.empty())
 190     return dataLayout.getTypeABIAlignment(fltType);
 191   return extractPreferredAlignment(params[0]);
 192 }
 193
 194 uint64_t mlir::detail::getDefaultPreferredAlignment(
 195     Type type, const DataLayout &dataLayout,
 196     ArrayRef<DataLayoutEntryInterface> params) {
 197   // Preferred alignment is same as natural for floats and vectors.
 198   if (isa<VectorType>(type))
 199     return dataLayout.getTypeABIAlignment(type);
 200
 201   if (auto fltType = dyn_cast<FloatType>(type))
 202     return getFloatTypePreferredAlignment(fltType, dataLayout, params);
 203
 204   // Preferred alignment is the closest power-of-two number above for integers
 205   // (ABI alignment may be smaller).
 206   if (auto intType = dyn_cast<IntegerType>(type))
 207     return getIntegerTypePreferredAlignment(intType, dataLayout, params);
 208
 209   if (isa<IndexType>(type)) {
 210     return dataLayout.getTypePreferredAlignment(
 211         IntegerType::get(type.getContext(), getIndexBitwidth(params)));
 212   }
 213
 214   if (auto ctype = dyn_cast<ComplexType>(type))
 215     return getDefaultPreferredAlignment(ctype.getElementType(), dataLayout,
 216                                         params);
 217
 218   if (auto typeInterface = dyn_cast<DataLayoutTypeInterface>(type))
 219     return typeInterface.getPreferredAlignment(dataLayout, params);
 220
 221   reportMissingDataLayout(type);
 222 }
 223
 224 std::optional<uint64_t> mlir::detail::getDefaultIndexBitwidth(
 225     Type type, const DataLayout &dataLayout,
 226     ArrayRef<DataLayoutEntryInterface> params) {
 227   if (isa<IndexType>(type))
 228     return getIndexBitwidth(params);
 229
 230   if (auto typeInterface = dyn_cast<DataLayoutTypeInterface>(type))
 231     if (std::optional<uint64_t> indexBitwidth =
 232             typeInterface.getIndexBitwidth(dataLayout, params))
 233       return *indexBitwidth;
 234
 235   // Return std::nullopt for all other types, which are assumed to be non
 236   // pointer-like types.
 237   return std::nullopt;
 238 }
 239
 240 // Returns the endianness if specified in the given entry. If the entry is empty
 241 // the default endianness represented by an empty attribute is returned.
 242 Attribute mlir::detail::getDefaultEndianness(DataLayoutEntryInterface entry) {
 243   if (entry == DataLayoutEntryInterface())
 244     return Attribute();
 245
 246   return entry.getValue();
 247 }
 248
 249 // Returns the memory space used for alloca operations if specified in the
 250 // given entry. If the entry is empty the default memory space represented by
 251 // an empty attribute is returned.
 252 Attribute
 253 mlir::detail::getDefaultAllocaMemorySpace(DataLayoutEntryInterface entry) {
 254   if (entry == DataLayoutEntryInterface()) {
 255     return Attribute();
 256   }
 257
 258   return entry.getValue();
 259 }
 260
 261 // Returns the memory space used for the program memory space.  if
 262 // specified in the given entry. If the entry is empty the default
 263 // memory space represented by an empty attribute is returned.
 264 Attribute
 265 mlir::detail::getDefaultProgramMemorySpace(DataLayoutEntryInterface entry) {
 266   if (entry == DataLayoutEntryInterface()) {
 267     return Attribute();
 268   }
 269
 270   return entry.getValue();
 271 }
 272
 273 // Returns the memory space used for global the global memory space. if
 274 // specified in the given entry. If the entry is empty the default memory
 275 // space represented by an empty attribute is returned.
 276 Attribute
 277 mlir::detail::getDefaultGlobalMemorySpace(DataLayoutEntryInterface entry) {
 278   if (entry == DataLayoutEntryInterface()) {
 279     return Attribute();
 280   }
 281
 282   return entry.getValue();
 283 }
 284
 285 // Returns the stack alignment if specified in the given entry. If the entry is
 286 // empty the default alignment zero is returned.
 287 uint64_t
 288 mlir::detail::getDefaultStackAlignment(DataLayoutEntryInterface entry) {
 289   if (entry == DataLayoutEntryInterface())
 290     return 0;
 291
 292   auto value = cast<IntegerAttr>(entry.getValue());
 293   return value.getValue().getZExtValue();
 294 }
 295
 296 std::optional<Attribute>
 297 mlir::detail::getDevicePropertyValue(DataLayoutEntryInterface entry) {
 298   if (entry == DataLayoutEntryInterface())
 299     return std::nullopt;
 300
 301   return entry.getValue();
 302 }
 303
 304 DataLayoutEntryList
 305 mlir::detail::filterEntriesForType(DataLayoutEntryListRef entries,
 306                                    TypeID typeID) {
 307   return llvm::to_vector<4>(llvm::make_filter_range(
 308       entries, [typeID](DataLayoutEntryInterface entry) {
 309         auto type = llvm::dyn_cast_if_present<Type>(entry.getKey());
 310         return type && type.getTypeID() == typeID;
 311       }));
 312 }
 313
 314 DataLayoutEntryInterface
 315 mlir::detail::filterEntryForIdentifier(DataLayoutEntryListRef entries,
 316                                        StringAttr id) {
 317   const auto *it = llvm::find_if(entries, [id](DataLayoutEntryInterface entry) {
 318     if (!entry.getKey().is<StringAttr>())
 319       return false;
 320     return entry.getKey().get<StringAttr>() == id;
 321   });
 322   return it == entries.end() ? DataLayoutEntryInterface() : *it;
 323 }
 324
 325 static DataLayoutSpecInterface getSpec(Operation *operation) {
 326   return llvm::TypeSwitch<Operation *, DataLayoutSpecInterface>(operation)
 327       .Case<ModuleOp, DataLayoutOpInterface>(
 328           [&](auto op) { return op.getDataLayoutSpec(); })
 329       .Default([](Operation *) {
 330         llvm_unreachable("expected an op with data layout spec");
 331         return DataLayoutSpecInterface();
 332       });
 333 }
 334
 335 static TargetSystemSpecInterface getTargetSystemSpec(Operation *operation) {
 336   if (operation) {
 337     ModuleOp moduleOp = dyn_cast<ModuleOp>(operation);
 338     if (!moduleOp)
 339       moduleOp = operation->getParentOfType<ModuleOp>();
 340     return moduleOp.getTargetSystemSpec();
 341   }
 342   return TargetSystemSpecInterface();
 343 }
 344
 345 /// Populates `opsWithLayout` with the list of proper ancestors of `leaf` that
 346 /// are either modules or implement the `DataLayoutOpInterface`.
 347 static void
 348 collectParentLayouts(Operation *leaf,
 349                      SmallVectorImpl<DataLayoutSpecInterface> &specs,
 350                      SmallVectorImpl<Location> *opLocations = nullptr) {
 351   if (!leaf)
 352     return;
 353
 354   for (Operation *parent = leaf->getParentOp(); parent != nullptr;
 355        parent = parent->getParentOp()) {
 356     llvm::TypeSwitch<Operation *>(parent)
 357         .Case<ModuleOp>([&](ModuleOp op) {
 358           // Skip top-level module op unless it has a layout. Top-level module
 359           // without layout is most likely the one implicitly added by the
 360           // parser and it doesn't have location. Top-level null specification
 361           // would have had the same effect as not having a specification at all
 362           // (using type defaults).
 363           if (!op->getParentOp() && !op.getDataLayoutSpec())
 364             return;
 365           specs.push_back(op.getDataLayoutSpec());
 366           if (opLocations)
 367             opLocations->push_back(op.getLoc());
 368         })
 369         .Case<DataLayoutOpInterface>([&](DataLayoutOpInterface op) {
 370           specs.push_back(op.getDataLayoutSpec());
 371           if (opLocations)
 372             opLocations->push_back(op.getLoc());
 373         });
 374   }
 375 }
 376
 377 /// Returns a layout spec that is a combination of the layout specs attached
 378 /// to the given operation and all its ancestors.
 379 static DataLayoutSpecInterface getCombinedDataLayout(Operation *leaf) {
 380   if (!leaf)
 381     return {};
 382
 383   assert((isa<ModuleOp, DataLayoutOpInterface>(leaf)) &&
 384          "expected an op with data layout spec");
 385
 386   SmallVector<DataLayoutOpInterface> opsWithLayout;
 387   SmallVector<DataLayoutSpecInterface> specs;
 388   collectParentLayouts(leaf, specs);
 389
 390   // Fast track if there are no ancestors.
 391   if (specs.empty())
 392     return getSpec(leaf);
 393
 394   // Create the list of non-null specs (null/missing specs can be safely
 395   // ignored) from the outermost to the innermost.
 396   auto nonNullSpecs = llvm::to_vector<2>(llvm::make_filter_range(
 397       llvm::reverse(specs),
 398       [](DataLayoutSpecInterface iface) { return iface != nullptr; }));
 399
 400   // Combine the specs using the innermost as anchor.
 401   if (DataLayoutSpecInterface current = getSpec(leaf))
 402     return current.combineWith(nonNullSpecs);
 403   if (nonNullSpecs.empty())
 404     return {};
 405   return nonNullSpecs.back().combineWith(
 406       llvm::ArrayRef(nonNullSpecs).drop_back());
 407 }
 408
 409 LogicalResult mlir::detail::verifyDataLayoutOp(Operation *op) {
 410   DataLayoutSpecInterface spec = getSpec(op);
 411   // The layout specification may be missing and it's fine.
 412   if (!spec)
 413     return success();
 414
 415   if (failed(spec.verifySpec(op->getLoc())))
 416     return failure();
 417   if (!getCombinedDataLayout(op)) {
 418     InFlightDiagnostic diag =
 419         op->emitError()
 420         << "data layout does not combine with layouts of enclosing ops";
 421     SmallVector<DataLayoutSpecInterface> specs;
 422     SmallVector<Location> opLocations;
 423     collectParentLayouts(op, specs, &opLocations);
 424     for (Location loc : opLocations)
 425       diag.attachNote(loc) << "enclosing op with data layout";
 426     return diag;
 427   }
 428   return success();
 429 }
 430
 431 llvm::TypeSize mlir::detail::divideCeil(llvm::TypeSize numerator,
 432                                         uint64_t denominator) {
 433   uint64_t divided =
 434       llvm::divideCeil(numerator.getKnownMinValue(), denominator);
 435   return llvm::TypeSize::get(divided, numerator.isScalable());
 436 }
 437
 438 //===----------------------------------------------------------------------===//
 439 // DataLayout
 440 //===----------------------------------------------------------------------===//
 441
 442 template <typename OpTy>
 443 void checkMissingLayout(DataLayoutSpecInterface originalLayout, OpTy op) {
 444   if (!originalLayout) {
 445     assert((!op || !op.getDataLayoutSpec()) &&
 446            "could not compute layout information for an op (failed to "
 447            "combine attributes?)");
 448   }
 449 }
 450
 451 mlir::DataLayout::DataLayout() : DataLayout(ModuleOp()) {}
 452
 453 mlir::DataLayout::DataLayout(DataLayoutOpInterface op)
 454     : originalLayout(getCombinedDataLayout(op)),
 455       originalTargetSystemDesc(getTargetSystemSpec(op)), scope(op),
 456       allocaMemorySpace(std::nullopt), programMemorySpace(std::nullopt),
 457       globalMemorySpace(std::nullopt), stackAlignment(std::nullopt) {
 458 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
 459   checkMissingLayout(originalLayout, op);
 460   collectParentLayouts(op, layoutStack);
 461 #endif
 462 }
 463
 464 mlir::DataLayout::DataLayout(ModuleOp op)
 465     : originalLayout(getCombinedDataLayout(op)),
 466       originalTargetSystemDesc(getTargetSystemSpec(op)), scope(op),
 467       allocaMemorySpace(std::nullopt), programMemorySpace(std::nullopt),
 468       globalMemorySpace(std::nullopt), stackAlignment(std::nullopt) {
 469 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
 470   checkMissingLayout(originalLayout, op);
 471   collectParentLayouts(op, layoutStack);
 472 #endif
 473 }
 474
 475 mlir::DataLayout mlir::DataLayout::closest(Operation *op) {
 476   // Search the closest parent either being a module operation or implementing
 477   // the data layout interface.
 478   while (op) {
 479     if (auto module = dyn_cast<ModuleOp>(op))
 480       return DataLayout(module);
 481     if (auto iface = dyn_cast<DataLayoutOpInterface>(op))
 482       return DataLayout(iface);
 483     op = op->getParentOp();
 484   }
 485   return DataLayout();
 486 }
 487
 488 void mlir::DataLayout::checkValid() const {
 489 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
 490   SmallVector<DataLayoutSpecInterface> specs;
 491   collectParentLayouts(scope, specs);
 492   assert(specs.size() == layoutStack.size() &&
 493          "data layout object used, but no longer valid due to the change in "
 494          "number of nested layouts");
 495   for (auto pair : llvm::zip(specs, layoutStack)) {
 496     Attribute newLayout = std::get<0>(pair);
 497     Attribute origLayout = std::get<1>(pair);
 498     assert(newLayout == origLayout &&
 499            "data layout object used, but no longer valid "
 500            "due to the change in layout attributes");
 501   }
 502 #endif
 503   assert(((!scope && !this->originalLayout) ||
 504           (scope && this->originalLayout == getCombinedDataLayout(scope))) &&
 505          "data layout object used, but no longer valid due to the change in "
 506          "layout spec");
 507 }
 508
 509 /// Looks up the value for the given type key in the given cache. If there is no
 510 /// such value in the cache, compute it using the given callback and put it in
 511 /// the cache before returning.
 512 template <typename T>
 513 static T cachedLookup(Type t, DenseMap<Type, T> &cache,
 514                       function_ref<T(Type)> compute) {
 515   auto it = cache.find(t);
 516   if (it != cache.end())
 517     return it->second;
 518
 519   auto result = cache.try_emplace(t, compute(t));
 520   return result.first->second;
 521 }
 522
 523 llvm::TypeSize mlir::DataLayout::getTypeSize(Type t) const {
 524   checkValid();
 525   return cachedLookup<llvm::TypeSize>(t, sizes, [&](Type ty) {
 526     DataLayoutEntryList list;
 527     if (originalLayout)
 528       list = originalLayout.getSpecForType(ty.getTypeID());
 529     if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
 530       return iface.getTypeSize(ty, *this, list);
 531     return detail::getDefaultTypeSize(ty, *this, list);
 532   });
 533 }
 534
 535 llvm::TypeSize mlir::DataLayout::getTypeSizeInBits(Type t) const {
 536   checkValid();
 537   return cachedLookup<llvm::TypeSize>(t, bitsizes, [&](Type ty) {
 538     DataLayoutEntryList list;
 539     if (originalLayout)
 540       list = originalLayout.getSpecForType(ty.getTypeID());
 541     if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
 542       return iface.getTypeSizeInBits(ty, *this, list);
 543     return detail::getDefaultTypeSizeInBits(ty, *this, list);
 544   });
 545 }
 546
 547 uint64_t mlir::DataLayout::getTypeABIAlignment(Type t) const {
 548   checkValid();
 549   return cachedLookup<uint64_t>(t, abiAlignments, [&](Type ty) {
 550     DataLayoutEntryList list;
 551     if (originalLayout)
 552       list = originalLayout.getSpecForType(ty.getTypeID());
 553     if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
 554       return iface.getTypeABIAlignment(ty, *this, list);
 555     return detail::getDefaultABIAlignment(ty, *this, list);
 556   });
 557 }
 558
 559 uint64_t mlir::DataLayout::getTypePreferredAlignment(Type t) const {
 560   checkValid();
 561   return cachedLookup<uint64_t>(t, preferredAlignments, [&](Type ty) {
 562     DataLayoutEntryList list;
 563     if (originalLayout)
 564       list = originalLayout.getSpecForType(ty.getTypeID());
 565     if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
 566       return iface.getTypePreferredAlignment(ty, *this, list);
 567     return detail::getDefaultPreferredAlignment(ty, *this, list);
 568   });
 569 }
 570
 571 std::optional<uint64_t> mlir::DataLayout::getTypeIndexBitwidth(Type t) const {
 572   checkValid();
 573   return cachedLookup<std::optional<uint64_t>>(t, indexBitwidths, [&](Type ty) {
 574     DataLayoutEntryList list;
 575     if (originalLayout)
 576       list = originalLayout.getSpecForType(ty.getTypeID());
 577     if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
 578       return iface.getIndexBitwidth(ty, *this, list);
 579     return detail::getDefaultIndexBitwidth(ty, *this, list);
 580   });
 581 }
 582
 583 mlir::Attribute mlir::DataLayout::getEndianness() const {
 584   checkValid();
 585   if (endianness)
 586     return *endianness;
 587   DataLayoutEntryInterface entry;
 588   if (originalLayout)
 589     entry = originalLayout.getSpecForIdentifier(
 590         originalLayout.getEndiannessIdentifier(originalLayout.getContext()));
 591
 592   if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
 593     endianness = iface.getEndianness(entry);
 594   else
 595     endianness = detail::getDefaultEndianness(entry);
 596   return *endianness;
 597 }
 598
 599 mlir::Attribute mlir::DataLayout::getAllocaMemorySpace() const {
 600   checkValid();
 601   if (allocaMemorySpace)
 602     return *allocaMemorySpace;
 603   DataLayoutEntryInterface entry;
 604   if (originalLayout)
 605     entry = originalLayout.getSpecForIdentifier(
 606         originalLayout.getAllocaMemorySpaceIdentifier(
 607             originalLayout.getContext()));
 608   if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
 609     allocaMemorySpace = iface.getAllocaMemorySpace(entry);
 610   else
 611     allocaMemorySpace = detail::getDefaultAllocaMemorySpace(entry);
 612   return *allocaMemorySpace;
 613 }
 614
 615 mlir::Attribute mlir::DataLayout::getProgramMemorySpace() const {
 616   checkValid();
 617   if (programMemorySpace)
 618     return *programMemorySpace;
 619   DataLayoutEntryInterface entry;
 620   if (originalLayout)
 621     entry = originalLayout.getSpecForIdentifier(
 622         originalLayout.getProgramMemorySpaceIdentifier(
 623             originalLayout.getContext()));
 624   if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
 625     programMemorySpace = iface.getProgramMemorySpace(entry);
 626   else
 627     programMemorySpace = detail::getDefaultProgramMemorySpace(entry);
 628   return *programMemorySpace;
 629 }
 630
 631 mlir::Attribute mlir::DataLayout::getGlobalMemorySpace() const {
 632   checkValid();
 633   if (globalMemorySpace)
 634     return *globalMemorySpace;
 635   DataLayoutEntryInterface entry;
 636   if (originalLayout)
 637     entry = originalLayout.getSpecForIdentifier(
 638         originalLayout.getGlobalMemorySpaceIdentifier(
 639             originalLayout.getContext()));
 640   if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
 641     globalMemorySpace = iface.getGlobalMemorySpace(entry);
 642   else
 643     globalMemorySpace = detail::getDefaultGlobalMemorySpace(entry);
 644   return *globalMemorySpace;
 645 }
 646
 647 uint64_t mlir::DataLayout::getStackAlignment() const {
 648   checkValid();
 649   if (stackAlignment)
 650     return *stackAlignment;
 651   DataLayoutEntryInterface entry;
 652   if (originalLayout)
 653     entry = originalLayout.getSpecForIdentifier(
 654         originalLayout.getStackAlignmentIdentifier(
 655             originalLayout.getContext()));
 656   if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
 657     stackAlignment = iface.getStackAlignment(entry);
 658   else
 659     stackAlignment = detail::getDefaultStackAlignment(entry);
 660   return *stackAlignment;
 661 }
 662
 663 std::optional<Attribute> mlir::DataLayout::getDevicePropertyValue(
 664     TargetSystemSpecInterface::DeviceID deviceID,
 665     StringAttr propertyName) const {
 666   checkValid();
 667   DataLayoutEntryInterface entry;
 668   if (originalTargetSystemDesc) {
 669     if (std::optional<TargetDeviceSpecInterface> device =
 670             originalTargetSystemDesc.getDeviceSpecForDeviceID(deviceID))
 671       entry = device->getSpecForIdentifier(propertyName);
 672   }
 673   // Currently I am not caching the results because we do not return
 674   // default values of these properties. Instead if the property is
 675   // missing, we return std::nullopt so that the users can resort to
 676   // the default value however they want.
 677   if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
 678     return iface.getDevicePropertyValue(entry);
 679   else
 680     return detail::getDevicePropertyValue(entry);
 681 }
 682
 683 //===----------------------------------------------------------------------===//
 684 // DataLayoutSpecInterface
 685 //===----------------------------------------------------------------------===//
 686
 687 void DataLayoutSpecInterface::bucketEntriesByType(
 688     DenseMap<TypeID, DataLayoutEntryList> &types,
 689     DenseMap<StringAttr, DataLayoutEntryInterface> &ids) {
 690   for (DataLayoutEntryInterface entry : getEntries()) {
 691     if (auto type = llvm::dyn_cast_if_present<Type>(entry.getKey()))
 692       types[type.getTypeID()].push_back(entry);
 693     else
 694       ids[entry.getKey().get<StringAttr>()] = entry;
 695   }
 696 }
 697
 698 LogicalResult mlir::detail::verifyDataLayoutSpec(DataLayoutSpecInterface spec,
 699                                                  Location loc) {
 700   // First, verify individual entries.
 701   for (DataLayoutEntryInterface entry : spec.getEntries())
 702     if (failed(entry.verifyEntry(loc)))
 703       return failure();
 704
 705   // Second, dispatch verifications of entry groups to types or dialects they
 706   // are associated with.
 707   DenseMap<TypeID, DataLayoutEntryList> types;
 708   DenseMap<StringAttr, DataLayoutEntryInterface> ids;
 709   spec.bucketEntriesByType(types, ids);
 710
 711   for (const auto &kvp : types) {
 712     auto sampleType = kvp.second.front().getKey().get<Type>();
 713     if (isa<IndexType>(sampleType)) {
 714       assert(kvp.second.size() == 1 &&
 715              "expected one data layout entry for non-parametric 'index' type");
 716       if (!isa<IntegerAttr>(kvp.second.front().getValue()))
 717         return emitError(loc)
 718                << "expected integer attribute in the data layout entry for "
 719                << sampleType;
 720       continue;
 721     }
 722
 723     if (isa<IntegerType, FloatType>(sampleType)) {
 724       for (DataLayoutEntryInterface entry : kvp.second) {
 725         auto value = dyn_cast<DenseIntElementsAttr>(entry.getValue());
 726         if (!value || !value.getElementType().isSignlessInteger(64)) {
 727           emitError(loc) << "expected a dense i64 elements attribute in the "
 728                             "data layout entry "
 729                          << entry;
 730           return failure();
 731         }
 732
 733         auto elements = llvm::to_vector<2>(value.getValues<uint64_t>());
 734         unsigned numElements = elements.size();
 735         if (numElements < 1 || numElements > 2) {
 736           emitError(loc) << "expected 1 or 2 elements in the data layout entry "
 737                          << entry;
 738           return failure();
 739         }
 740
 741         uint64_t abi = elements[0];
 742         uint64_t preferred = numElements == 2 ? elements[1] : abi;
 743         if (preferred < abi) {
 744           emitError(loc)
 745               << "preferred alignment is expected to be greater than or equal "
 746                  "to the abi alignment in data layout entry "
 747               << entry;
 748           return failure();
 749         }
 750       }
 751       continue;
 752     }
 753
 754     if (isa<BuiltinDialect>(&sampleType.getDialect()))
 755       return emitError(loc) << "unexpected data layout for a built-in type";
 756
 757     auto dlType = dyn_cast<DataLayoutTypeInterface>(sampleType);
 758     if (!dlType)
 759       return emitError(loc)
 760              << "data layout specified for a type that does not support it";
 761     if (failed(dlType.verifyEntries(kvp.second, loc)))
 762       return failure();
 763   }
 764
 765   for (const auto &kvp : ids) {
 766     StringAttr identifier = kvp.second.getKey().get<StringAttr>();
 767     Dialect *dialect = identifier.getReferencedDialect();
 768
 769     // Ignore attributes that belong to an unknown dialect, the dialect may
 770     // actually implement the relevant interface but we don't know about that.
 771     if (!dialect)
 772       continue;
 773
 774     const auto *iface = dyn_cast<DataLayoutDialectInterface>(dialect);
 775     if (!iface) {
 776       return emitError(loc)
 777              << "the '" << dialect->getNamespace()
 778              << "' dialect does not support identifier data layout entries";
 779     }
 780     if (failed(iface->verifyEntry(kvp.second, loc)))
 781       return failure();
 782   }
 783
 784   return success();
 785 }
 786
 787 LogicalResult
 788 mlir::detail::verifyTargetSystemSpec(TargetSystemSpecInterface spec,
 789                                      Location loc) {
 790   DenseMap<StringAttr, DataLayoutEntryInterface> deviceDescKeys;
 791   DenseSet<TargetSystemSpecInterface::DeviceID> deviceIDs;
 792   for (const auto &entry : spec.getEntries()) {
 793     auto targetDeviceSpec =
 794         dyn_cast<TargetDeviceSpecInterface>(entry.getValue());
 795
 796     if (!targetDeviceSpec)
 797       return failure();
 798
 799     // First, verify individual target device desc specs.
 800     if (failed(targetDeviceSpec.verifyEntry(loc)))
 801       return failure();
 802
 803     // Check that device IDs are unique across all entries.
 804     auto deviceID =
 805         llvm::dyn_cast<TargetSystemSpecInterface::DeviceID>(entry.getKey());
 806     if (!deviceID)
 807       return failure();
 808
 809     if (!deviceIDs.insert(deviceID).second) {
 810       return failure();
 811     }
 812
 813     // collect all the keys used by all the target device specs.
 814     for (DataLayoutEntryInterface entry : targetDeviceSpec.getEntries()) {
 815       if (auto type = llvm::dyn_cast_if_present<Type>(entry.getKey())) {
 816         // targetDeviceSpec does not support Type as a key.
 817         return failure();
 818       } else {
 819         deviceDescKeys[entry.getKey().get<StringAttr>()] = entry;
 820       }
 821     }
 822   }
 823
 824   for (const auto &[keyName, keyVal] : deviceDescKeys) {
 825     Dialect *dialect = keyName.getReferencedDialect();
 826
 827     // Ignore attributes that belong to an unknown dialect, the dialect may
 828     // actually implement the relevant interface but we don't know about that.
 829     if (!dialect)
 830       return failure();
 831
 832     const auto *iface = dyn_cast<DataLayoutDialectInterface>(dialect);
 833     if (!iface) {
 834       return emitError(loc)
 835              << "the '" << dialect->getNamespace()
 836              << "' dialect does not support identifier data layout entries";
 837     }
 838     if (failed(iface->verifyEntry(keyVal, loc)))
 839       return failure();
 840   }
 841
 842   return success();
 843 }
 844
 845 #include "mlir/Interfaces/DataLayoutAttrInterface.cpp.inc"
 846 #include "mlir/Interfaces/DataLayoutOpInterface.cpp.inc"
 847 #include "mlir/Interfaces/DataLayoutTypeInterface.cpp.inc"