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