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 
17 using namespace mlir;
18 
19 //===----------------------------------------------------------------------===//
20 // Default implementations
21 //===----------------------------------------------------------------------===//
22 
23 /// Reports that the given type is missing the data layout information and
24 /// exits.
25 static LLVM_ATTRIBUTE_NORETURN void reportMissingDataLayout(Type type) {
26   std::string message;
27   llvm::raw_string_ostream os(message);
28   os << "neither the scoping op nor the type class provide data layout "
29         "information for "
30      << type;
31   llvm::report_fatal_error(os.str());
32 }
33 
34 /// Returns the bitwidth of the index type if specified in the param list.
35 /// Assumes 64-bit index otherwise.
36 static unsigned getIndexBitwidth(DataLayoutEntryListRef params) {
37   if (params.empty())
38     return 64;
39   auto attr = params.front().getValue().cast<IntegerAttr>();
40   return attr.getValue().getZExtValue();
41 }
42 
43 unsigned
44 mlir::detail::getDefaultTypeSize(Type type, const DataLayout &dataLayout,
45                                  ArrayRef<DataLayoutEntryInterface> params) {
46   unsigned bits = getDefaultTypeSizeInBits(type, dataLayout, params);
47   return llvm::divideCeil(bits, 8);
48 }
49 
50 unsigned mlir::detail::getDefaultTypeSizeInBits(Type type,
51                                                 const DataLayout &dataLayout,
52                                                 DataLayoutEntryListRef params) {
53   if (type.isa<IntegerType, FloatType>())
54     return type.getIntOrFloatBitWidth();
55 
56   // Index is an integer of some bitwidth.
57   if (type.isa<IndexType>())
58     return dataLayout.getTypeSizeInBits(
59         IntegerType::get(type.getContext(), getIndexBitwidth(params)));
60 
61   // Sizes of vector types are rounded up to those of types with closest
62   // power-of-two number of elements in the innermost dimension. We also assume
63   // there is no bit-packing at the moment element sizes are taken in bytes and
64   // multiplied with 8 bits.
65   // TODO: make this extensible.
66   if (auto vecType = type.dyn_cast<VectorType>())
67     return vecType.getNumElements() / vecType.getShape().back() *
68            llvm::PowerOf2Ceil(vecType.getShape().back()) *
69            dataLayout.getTypeSize(vecType.getElementType()) * 8;
70 
71   if (auto typeInterface = type.dyn_cast<DataLayoutTypeInterface>())
72     return typeInterface.getTypeSizeInBits(dataLayout, params);
73 
74   reportMissingDataLayout(type);
75 }
76 
77 unsigned mlir::detail::getDefaultABIAlignment(
78     Type type, const DataLayout &dataLayout,
79     ArrayRef<DataLayoutEntryInterface> params) {
80   // Natural alignment is the closest power-of-two number above.
81   if (type.isa<FloatType, VectorType>())
82     return llvm::PowerOf2Ceil(dataLayout.getTypeSize(type));
83 
84   // Index is an integer of some bitwidth.
85   if (type.isa<IndexType>())
86     return dataLayout.getTypeABIAlignment(
87         IntegerType::get(type.getContext(), getIndexBitwidth(params)));
88 
89   if (auto intType = type.dyn_cast<IntegerType>()) {
90     return intType.getWidth() < 64
91                ? llvm::PowerOf2Ceil(llvm::divideCeil(intType.getWidth(), 8))
92                : 4;
93   }
94 
95   if (auto typeInterface = type.dyn_cast<DataLayoutTypeInterface>())
96     return typeInterface.getABIAlignment(dataLayout, params);
97 
98   reportMissingDataLayout(type);
99 }
100 
101 unsigned mlir::detail::getDefaultPreferredAlignment(
102     Type type, const DataLayout &dataLayout,
103     ArrayRef<DataLayoutEntryInterface> params) {
104   // Preferred alignment is same as natural for floats and vectors.
105   if (type.isa<FloatType, VectorType>())
106     return dataLayout.getTypeABIAlignment(type);
107 
108   // Preferred alignment is the cloest power-of-two number above for integers
109   // (ABI alignment may be smaller).
110   if (type.isa<IntegerType, IndexType>())
111     return llvm::PowerOf2Ceil(dataLayout.getTypeSize(type));
112 
113   if (auto typeInterface = type.dyn_cast<DataLayoutTypeInterface>())
114     return typeInterface.getPreferredAlignment(dataLayout, params);
115 
116   reportMissingDataLayout(type);
117 }
118 
119 DataLayoutEntryList
120 mlir::detail::filterEntriesForType(DataLayoutEntryListRef entries,
121                                    TypeID typeID) {
122   return llvm::to_vector<4>(llvm::make_filter_range(
123       entries, [typeID](DataLayoutEntryInterface entry) {
124         auto type = entry.getKey().dyn_cast<Type>();
125         return type && type.getTypeID() == typeID;
126       }));
127 }
128 
129 DataLayoutEntryInterface
130 mlir::detail::filterEntryForIdentifier(DataLayoutEntryListRef entries,
131                                        Identifier id) {
132   const auto *it = llvm::find_if(entries, [id](DataLayoutEntryInterface entry) {
133     if (!entry.getKey().is<Identifier>())
134       return false;
135     return entry.getKey().get<Identifier>() == id;
136   });
137   return it == entries.end() ? DataLayoutEntryInterface() : *it;
138 }
139 
140 static DataLayoutSpecInterface getSpec(Operation *operation) {
141   return llvm::TypeSwitch<Operation *, DataLayoutSpecInterface>(operation)
142       .Case<ModuleOp, DataLayoutOpInterface>(
143           [&](auto op) { return op.getDataLayoutSpec(); })
144       .Default([](Operation *) {
145         llvm_unreachable("expected an op with data layout spec");
146         return DataLayoutSpecInterface();
147       });
148 }
149 
150 /// Populates `opsWithLayout` with the list of proper ancestors of `leaf` that
151 /// are either modules or implement the `DataLayoutOpInterface`.
152 static void
153 collectParentLayouts(Operation *leaf,
154                      SmallVectorImpl<DataLayoutSpecInterface> &specs,
155                      SmallVectorImpl<Location> *opLocations = nullptr) {
156   if (!leaf)
157     return;
158 
159   for (Operation *parent = leaf->getParentOp(); parent != nullptr;
160        parent = parent->getParentOp()) {
161     llvm::TypeSwitch<Operation *>(parent)
162         .Case<ModuleOp>([&](ModuleOp op) {
163           // Skip top-level module op unless it has a layout. Top-level module
164           // without layout is most likely the one implicitly added by the
165           // parser and it doesn't have location. Top-level null specification
166           // would have had the same effect as not having a specification at all
167           // (using type defaults).
168           if (!op->getParentOp() && !op.getDataLayoutSpec())
169             return;
170           specs.push_back(op.getDataLayoutSpec());
171           if (opLocations)
172             opLocations->push_back(op.getLoc());
173         })
174         .Case<DataLayoutOpInterface>([&](DataLayoutOpInterface op) {
175           specs.push_back(op.getDataLayoutSpec());
176           if (opLocations)
177             opLocations->push_back(op.getLoc());
178         });
179   }
180 }
181 
182 /// Returns a layout spec that is a combination of the layout specs attached
183 /// to the given operation and all its ancestors.
184 static DataLayoutSpecInterface getCombinedDataLayout(Operation *leaf) {
185   if (!leaf)
186     return {};
187 
188   assert((isa<ModuleOp, DataLayoutOpInterface>(leaf)) &&
189          "expected an op with data layout spec");
190 
191   SmallVector<DataLayoutOpInterface> opsWithLayout;
192   SmallVector<DataLayoutSpecInterface> specs;
193   collectParentLayouts(leaf, specs);
194 
195   // Fast track if there are no ancestors.
196   if (specs.empty())
197     return getSpec(leaf);
198 
199   // Create the list of non-null specs (null/missing specs can be safely
200   // ignored) from the outermost to the innermost.
201   auto nonNullSpecs = llvm::to_vector<2>(llvm::make_filter_range(
202       llvm::reverse(specs),
203       [](DataLayoutSpecInterface iface) { return iface != nullptr; }));
204 
205   // Combine the specs using the innermost as anchor.
206   if (DataLayoutSpecInterface current = getSpec(leaf))
207     return current.combineWith(nonNullSpecs);
208   if (nonNullSpecs.empty())
209     return {};
210   return nonNullSpecs.back().combineWith(
211       llvm::makeArrayRef(nonNullSpecs).drop_back());
212 }
213 
214 LogicalResult mlir::detail::verifyDataLayoutOp(Operation *op) {
215   DataLayoutSpecInterface spec = getSpec(op);
216   // The layout specification may be missing and it's fine.
217   if (!spec)
218     return success();
219 
220   if (failed(spec.verifySpec(op->getLoc())))
221     return failure();
222   if (!getCombinedDataLayout(op)) {
223     InFlightDiagnostic diag =
224         op->emitError()
225         << "data layout does not combine with layouts of enclosing ops";
226     SmallVector<DataLayoutSpecInterface> specs;
227     SmallVector<Location> opLocations;
228     collectParentLayouts(op, specs, &opLocations);
229     for (Location loc : opLocations)
230       diag.attachNote(loc) << "enclosing op with data layout";
231     return diag;
232   }
233   return success();
234 }
235 
236 //===----------------------------------------------------------------------===//
237 // DataLayout
238 //===----------------------------------------------------------------------===//
239 
240 template <typename OpTy>
241 void checkMissingLayout(DataLayoutSpecInterface originalLayout, OpTy op) {
242   if (!originalLayout) {
243     assert((!op || !op.getDataLayoutSpec()) &&
244            "could not compute layout information for an op (failed to "
245            "combine attributes?)");
246   }
247 }
248 
249 mlir::DataLayout::DataLayout() : DataLayout(ModuleOp()) {}
250 
251 mlir::DataLayout::DataLayout(DataLayoutOpInterface op)
252     : originalLayout(getCombinedDataLayout(op)), scope(op) {
253 #ifndef NDEBUG
254   checkMissingLayout(originalLayout, op);
255   collectParentLayouts(op, layoutStack);
256 #endif
257 }
258 
259 mlir::DataLayout::DataLayout(ModuleOp op)
260     : originalLayout(getCombinedDataLayout(op)), scope(op) {
261 #ifndef NDEBUG
262   checkMissingLayout(originalLayout, op);
263   collectParentLayouts(op, layoutStack);
264 #endif
265 }
266 
267 mlir::DataLayout mlir::DataLayout::closest(Operation *op) {
268   // Search the closest parent either being a module operation or implementing
269   // the data layout interface.
270   while (op) {
271     if (auto module = dyn_cast<ModuleOp>(op))
272       return DataLayout(module);
273     if (auto iface = dyn_cast<DataLayoutOpInterface>(op))
274       return DataLayout(iface);
275     op = op->getParentOp();
276   }
277   return DataLayout();
278 }
279 
280 void mlir::DataLayout::checkValid() const {
281 #ifndef NDEBUG
282   SmallVector<DataLayoutSpecInterface> specs;
283   collectParentLayouts(scope, specs);
284   assert(specs.size() == layoutStack.size() &&
285          "data layout object used, but no longer valid due to the change in "
286          "number of nested layouts");
287   for (auto pair : llvm::zip(specs, layoutStack)) {
288     Attribute newLayout = std::get<0>(pair);
289     Attribute origLayout = std::get<1>(pair);
290     assert(newLayout == origLayout &&
291            "data layout object used, but no longer valid "
292            "due to the change in layout attributes");
293   }
294 #endif
295   assert(((!scope && !this->originalLayout) ||
296           (scope && this->originalLayout == getCombinedDataLayout(scope))) &&
297          "data layout object used, but no longer valid due to the change in "
298          "layout spec");
299 }
300 
301 /// Looks up the value for the given type key in the given cache. If there is no
302 /// such value in the cache, compute it using the given callback and put it in
303 /// the cache before returning.
304 static unsigned cachedLookup(Type t, DenseMap<Type, unsigned> &cache,
305                              function_ref<unsigned(Type)> compute) {
306   auto it = cache.find(t);
307   if (it != cache.end())
308     return it->second;
309 
310   auto result = cache.try_emplace(t, compute(t));
311   return result.first->second;
312 }
313 
314 unsigned mlir::DataLayout::getTypeSize(Type t) const {
315   checkValid();
316   return cachedLookup(t, sizes, [&](Type ty) {
317     DataLayoutEntryList list;
318     if (originalLayout)
319       list = originalLayout.getSpecForType(ty.getTypeID());
320     if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
321       return iface.getTypeSize(ty, *this, list);
322     return detail::getDefaultTypeSize(ty, *this, list);
323   });
324 }
325 
326 unsigned mlir::DataLayout::getTypeSizeInBits(Type t) const {
327   checkValid();
328   return cachedLookup(t, bitsizes, [&](Type ty) {
329     DataLayoutEntryList list;
330     if (originalLayout)
331       list = originalLayout.getSpecForType(ty.getTypeID());
332     if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
333       return iface.getTypeSizeInBits(ty, *this, list);
334     return detail::getDefaultTypeSizeInBits(ty, *this, list);
335   });
336 }
337 
338 unsigned mlir::DataLayout::getTypeABIAlignment(Type t) const {
339   checkValid();
340   return cachedLookup(t, abiAlignments, [&](Type ty) {
341     DataLayoutEntryList list;
342     if (originalLayout)
343       list = originalLayout.getSpecForType(ty.getTypeID());
344     if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
345       return iface.getTypeABIAlignment(ty, *this, list);
346     return detail::getDefaultABIAlignment(ty, *this, list);
347   });
348 }
349 
350 unsigned mlir::DataLayout::getTypePreferredAlignment(Type t) const {
351   checkValid();
352   return cachedLookup(t, preferredAlignments, [&](Type ty) {
353     DataLayoutEntryList list;
354     if (originalLayout)
355       list = originalLayout.getSpecForType(ty.getTypeID());
356     if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
357       return iface.getTypePreferredAlignment(ty, *this, list);
358     return detail::getDefaultPreferredAlignment(ty, *this, list);
359   });
360 }
361 
362 //===----------------------------------------------------------------------===//
363 // DataLayoutSpecInterface
364 //===----------------------------------------------------------------------===//
365 
366 void DataLayoutSpecInterface::bucketEntriesByType(
367     DenseMap<TypeID, DataLayoutEntryList> &types,
368     DenseMap<Identifier, DataLayoutEntryInterface> &ids) {
369   for (DataLayoutEntryInterface entry : getEntries()) {
370     if (auto type = entry.getKey().dyn_cast<Type>())
371       types[type.getTypeID()].push_back(entry);
372     else
373       ids[entry.getKey().get<Identifier>()] = entry;
374   }
375 }
376 
377 LogicalResult mlir::detail::verifyDataLayoutSpec(DataLayoutSpecInterface spec,
378                                                  Location loc) {
379   // First, verify individual entries.
380   for (DataLayoutEntryInterface entry : spec.getEntries())
381     if (failed(entry.verifyEntry(loc)))
382       return failure();
383 
384   // Second, dispatch verifications of entry groups to types or dialects they
385   // are are associated with.
386   DenseMap<TypeID, DataLayoutEntryList> types;
387   DenseMap<Identifier, DataLayoutEntryInterface> ids;
388   spec.bucketEntriesByType(types, ids);
389 
390   for (const auto &kvp : types) {
391     auto sampleType = kvp.second.front().getKey().get<Type>();
392     if (sampleType.isa<IndexType>()) {
393       assert(kvp.second.size() == 1 &&
394              "expected one data layout entry for non-parametric 'index' type");
395       if (!kvp.second.front().getValue().isa<IntegerAttr>())
396         return emitError(loc)
397                << "expected integer attribute in the data layout entry for "
398                << sampleType;
399       continue;
400     }
401 
402     if (isa<BuiltinDialect>(&sampleType.getDialect()))
403       return emitError(loc) << "unexpected data layout for a built-in type";
404 
405     auto dlType = sampleType.dyn_cast<DataLayoutTypeInterface>();
406     if (!dlType)
407       return emitError(loc)
408              << "data layout specified for a type that does not support it";
409     if (failed(dlType.verifyEntries(kvp.second, loc)))
410       return failure();
411   }
412 
413   for (const auto &kvp : ids) {
414     Identifier identifier = kvp.second.getKey().get<Identifier>();
415     Dialect *dialect = identifier.getDialect();
416 
417     // Ignore attributes that belong to an unknown dialect, the dialect may
418     // actually implement the relevant interface but we don't know about that.
419     if (!dialect)
420       continue;
421 
422     const auto *iface =
423         dialect->getRegisteredInterface<DataLayoutDialectInterface>();
424     if (failed(iface->verifyEntry(kvp.second, loc)))
425       return failure();
426   }
427 
428   return success();
429 }
430 
431 #include "mlir/Interfaces/DataLayoutAttrInterface.cpp.inc"
432 #include "mlir/Interfaces/DataLayoutOpInterface.cpp.inc"
433 #include "mlir/Interfaces/DataLayoutTypeInterface.cpp.inc"
434