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