1 //===- DLTI.cpp - Data Layout And Target Info MLIR Dialect 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/Dialect/DLTI/DLTI.h"
10 #include "mlir/IR/Builders.h"
11 #include "mlir/IR/BuiltinDialect.h"
12 #include "mlir/IR/Dialect.h"
13 #include "mlir/IR/DialectImplementation.h"
14 #include "llvm/ADT/TypeSwitch.h"
15 
16 using namespace mlir;
17 
18 //===----------------------------------------------------------------------===//
19 // DataLayoutEntryAttr
20 //===----------------------------------------------------------------------===//
21 //
22 constexpr const StringLiteral mlir::DataLayoutEntryAttr::kAttrKeyword;
23 
24 namespace mlir {
25 namespace impl {
26 class DataLayoutEntryStorage : public AttributeStorage {
27 public:
28   using KeyTy = std::pair<DataLayoutEntryKey, Attribute>;
29 
30   DataLayoutEntryStorage(DataLayoutEntryKey entryKey, Attribute value)
31       : entryKey(entryKey), value(value) {}
32 
33   static DataLayoutEntryStorage *construct(AttributeStorageAllocator &allocator,
34                                            const KeyTy &key) {
35     return new (allocator.allocate<DataLayoutEntryStorage>())
36         DataLayoutEntryStorage(key.first, key.second);
37   }
38 
39   bool operator==(const KeyTy &other) const {
40     return other.first == entryKey && other.second == value;
41   }
42 
43   DataLayoutEntryKey entryKey;
44   Attribute value;
45 };
46 } // namespace impl
47 } // namespace mlir
48 
49 DataLayoutEntryAttr DataLayoutEntryAttr::get(Identifier key, Attribute value) {
50   return Base::get(key.getContext(), key, value);
51 }
52 
53 DataLayoutEntryAttr DataLayoutEntryAttr::get(Type key, Attribute value) {
54   return Base::get(key.getContext(), key, value);
55 }
56 
57 DataLayoutEntryKey DataLayoutEntryAttr::getKey() const {
58   return getImpl()->entryKey;
59 }
60 
61 Attribute DataLayoutEntryAttr::getValue() const { return getImpl()->value; }
62 
63 /// Parses an attribute with syntax:
64 ///   attr ::= `#target.` `dl_entry` `<` (type | quoted-string) `,` attr `>`
65 DataLayoutEntryAttr DataLayoutEntryAttr::parse(DialectAsmParser &parser) {
66   if (failed(parser.parseLess()))
67     return {};
68 
69   Type type = nullptr;
70   StringRef identifier;
71   llvm::SMLoc idLoc = parser.getCurrentLocation();
72   OptionalParseResult parsedType = parser.parseOptionalType(type);
73   if (parsedType.hasValue() && failed(parsedType.getValue()))
74     return {};
75   if (!parsedType.hasValue()) {
76     OptionalParseResult parsedString = parser.parseOptionalString(&identifier);
77     if (!parsedString.hasValue() || failed(parsedString.getValue())) {
78       parser.emitError(idLoc) << "expected a type or a quoted string";
79       return {};
80     }
81   }
82 
83   Attribute value;
84   if (failed(parser.parseComma()) || failed(parser.parseAttribute(value)) ||
85       failed(parser.parseGreater()))
86     return {};
87 
88   return type ? get(type, value)
89               : get(parser.getBuilder().getIdentifier(identifier), value);
90 }
91 
92 void DataLayoutEntryAttr::print(DialectAsmPrinter &os) const {
93   os << DataLayoutEntryAttr::kAttrKeyword << "<";
94   if (auto type = getKey().dyn_cast<Type>())
95     os << type;
96   else
97     os << "\"" << getKey().get<Identifier>().strref() << "\"";
98   os << ", " << getValue() << ">";
99 }
100 
101 //===----------------------------------------------------------------------===//
102 // DataLayoutSpecAttr
103 //===----------------------------------------------------------------------===//
104 //
105 constexpr const StringLiteral mlir::DataLayoutSpecAttr::kAttrKeyword;
106 
107 namespace mlir {
108 namespace impl {
109 class DataLayoutSpecStorage : public AttributeStorage {
110 public:
111   using KeyTy = ArrayRef<DataLayoutEntryInterface>;
112 
113   DataLayoutSpecStorage(ArrayRef<DataLayoutEntryInterface> entries)
114       : entries(entries) {}
115 
116   bool operator==(const KeyTy &key) const { return key == entries; }
117 
118   static DataLayoutSpecStorage *construct(AttributeStorageAllocator &allocator,
119                                           const KeyTy &key) {
120     return new (allocator.allocate<DataLayoutSpecStorage>())
121         DataLayoutSpecStorage(allocator.copyInto(key));
122   }
123 
124   ArrayRef<DataLayoutEntryInterface> entries;
125 };
126 } // namespace impl
127 } // namespace mlir
128 
129 DataLayoutSpecAttr
130 DataLayoutSpecAttr::get(MLIRContext *ctx,
131                         ArrayRef<DataLayoutEntryInterface> entries) {
132   return Base::get(ctx, entries);
133 }
134 
135 DataLayoutSpecAttr
136 DataLayoutSpecAttr::getChecked(function_ref<InFlightDiagnostic()> emitError,
137                                MLIRContext *context,
138                                ArrayRef<DataLayoutEntryInterface> entries) {
139   return Base::getChecked(emitError, context, entries);
140 }
141 
142 LogicalResult
143 DataLayoutSpecAttr::verify(function_ref<InFlightDiagnostic()> emitError,
144                            ArrayRef<DataLayoutEntryInterface> entries) {
145   DenseSet<Type> types;
146   DenseSet<Identifier> ids;
147   for (DataLayoutEntryInterface entry : entries) {
148     if (auto type = entry.getKey().dyn_cast<Type>()) {
149       if (!types.insert(type).second)
150         return emitError() << "repeated layout entry key: " << type;
151     } else {
152       auto id = entry.getKey().get<Identifier>();
153       if (!ids.insert(id).second)
154         return emitError() << "repeated layout entry key: " << id;
155     }
156   }
157   return success();
158 }
159 
160 /// Given a list of old and a list of new entries, overwrites old entries with
161 /// new ones if they have matching keys, appends new entries to the old entry
162 /// list otherwise.
163 static void
164 overwriteDuplicateEntries(SmallVectorImpl<DataLayoutEntryInterface> &oldEntries,
165                           ArrayRef<DataLayoutEntryInterface> newEntries) {
166   unsigned oldEntriesSize = oldEntries.size();
167   for (DataLayoutEntryInterface entry : newEntries) {
168     // We expect a small (dozens) number of entries, so it is practically
169     // cheaper to iterate over the list linearly rather than to create an
170     // auxiliary hashmap to avoid duplication. Also note that we never need to
171     // check for duplicate keys the values that were added from `newEntries`.
172     bool replaced = false;
173     for (unsigned i = 0; i < oldEntriesSize; ++i) {
174       if (oldEntries[i].getKey() == entry.getKey()) {
175         oldEntries[i] = entry;
176         replaced = true;
177         break;
178       }
179     }
180     if (!replaced)
181       oldEntries.push_back(entry);
182   }
183 }
184 
185 /// Combines a data layout spec into the given lists of entries organized by
186 /// type class and identifier, overwriting them if necessary. Fails to combine
187 /// if the two entries with identical keys are not compatible.
188 static LogicalResult
189 combineOneSpec(DataLayoutSpecInterface spec,
190                DenseMap<TypeID, DataLayoutEntryList> &entriesForType,
191                DenseMap<Identifier, DataLayoutEntryInterface> &entriesForID) {
192   // A missing spec should be fine.
193   if (!spec)
194     return success();
195 
196   DenseMap<TypeID, DataLayoutEntryList> newEntriesForType;
197   DenseMap<Identifier, DataLayoutEntryInterface> newEntriesForID;
198   spec.bucketEntriesByType(newEntriesForType, newEntriesForID);
199 
200   // Try overwriting the old entries with the new ones.
201   for (const auto &kvp : newEntriesForType) {
202     if (!entriesForType.count(kvp.first)) {
203       entriesForType[kvp.first] = std::move(kvp.second);
204       continue;
205     }
206 
207     Type typeSample = kvp.second.front().getKey().get<Type>();
208     assert(&typeSample.getDialect() !=
209                typeSample.getContext()->getLoadedDialect<BuiltinDialect>() &&
210            "unexpected data layout entry for built-in type");
211 
212     auto interface = typeSample.cast<DataLayoutTypeInterface>();
213     if (!interface.areCompatible(entriesForType.lookup(kvp.first), kvp.second))
214       return failure();
215 
216     overwriteDuplicateEntries(entriesForType[kvp.first], kvp.second);
217   }
218 
219   for (const auto &kvp : newEntriesForID) {
220     Identifier id = kvp.second.getKey().get<Identifier>();
221     Dialect *dialect = id.getDialect();
222     if (!entriesForID.count(id)) {
223       entriesForID[id] = kvp.second;
224       continue;
225     }
226 
227     // Attempt to combine the enties using the dialect interface. If the
228     // dialect is not loaded for some reason, use the default combinator
229     // that conservatively accepts identical entries only.
230     entriesForID[id] =
231         dialect ? dialect->getRegisteredInterface<DataLayoutDialectInterface>()
232                       ->combine(entriesForID[id], kvp.second)
233                 : DataLayoutDialectInterface::defaultCombine(entriesForID[id],
234                                                              kvp.second);
235     if (!entriesForID[id])
236       return failure();
237   }
238 
239   return success();
240 }
241 
242 DataLayoutSpecAttr
243 DataLayoutSpecAttr::combineWith(ArrayRef<DataLayoutSpecInterface> specs) const {
244   // Only combine with attributes of the same kind.
245   // TODO: reconsider this when the need arises.
246   if (llvm::any_of(specs, [](DataLayoutSpecInterface spec) {
247         return !spec.isa<DataLayoutSpecAttr>();
248       }))
249     return {};
250 
251   // Combine all specs in order, with `this` being the last one.
252   DenseMap<TypeID, DataLayoutEntryList> entriesForType;
253   DenseMap<Identifier, DataLayoutEntryInterface> entriesForID;
254   for (DataLayoutSpecInterface spec : specs)
255     if (failed(combineOneSpec(spec, entriesForType, entriesForID)))
256       return nullptr;
257   if (failed(combineOneSpec(*this, entriesForType, entriesForID)))
258     return nullptr;
259 
260   // Rebuild the linear list of entries.
261   SmallVector<DataLayoutEntryInterface> entries;
262   llvm::append_range(entries, llvm::make_second_range(entriesForID));
263   for (const auto &kvp : entriesForType)
264     llvm::append_range(entries, kvp.getSecond());
265 
266   return DataLayoutSpecAttr::get(getContext(), entries);
267 }
268 
269 DataLayoutEntryListRef DataLayoutSpecAttr::getEntries() const {
270   return getImpl()->entries;
271 }
272 
273 /// Parses an attribute with syntax
274 ///   attr ::= `#target.` `dl_spec` `<` attr-list? `>`
275 ///   attr-list ::= attr
276 ///               | attr `,` attr-list
277 DataLayoutSpecAttr DataLayoutSpecAttr::parse(DialectAsmParser &parser) {
278   if (failed(parser.parseLess()))
279     return {};
280 
281   // Empty spec.
282   if (succeeded(parser.parseOptionalGreater()))
283     return get(parser.getBuilder().getContext(), {});
284 
285   SmallVector<DataLayoutEntryInterface> entries;
286   do {
287     entries.emplace_back();
288     if (failed(parser.parseAttribute(entries.back())))
289       return {};
290   } while (succeeded(parser.parseOptionalComma()));
291 
292   if (failed(parser.parseGreater()))
293     return {};
294   return getChecked([&] { return parser.emitError(parser.getNameLoc()); },
295                     parser.getBuilder().getContext(), entries);
296 }
297 
298 void DataLayoutSpecAttr::print(DialectAsmPrinter &os) const {
299   os << DataLayoutSpecAttr::kAttrKeyword << "<";
300   llvm::interleaveComma(getEntries(), os);
301   os << ">";
302 }
303 
304 //===----------------------------------------------------------------------===//
305 // DLTIDialect
306 //===----------------------------------------------------------------------===//
307 
308 constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutAttrName;
309 constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutEndiannessKey;
310 constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutEndiannessBig;
311 constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutEndiannessLittle;
312 
313 namespace {
314 class TargetDataLayoutInterface : public DataLayoutDialectInterface {
315 public:
316   using DataLayoutDialectInterface::DataLayoutDialectInterface;
317 
318   LogicalResult verifyEntry(DataLayoutEntryInterface entry,
319                             Location loc) const final {
320     StringRef entryName = entry.getKey().get<Identifier>().strref();
321     if (entryName == DLTIDialect::kDataLayoutEndiannessKey) {
322       auto value = entry.getValue().dyn_cast<StringAttr>();
323       if (value &&
324           (value.getValue() == DLTIDialect::kDataLayoutEndiannessBig ||
325            value.getValue() == DLTIDialect::kDataLayoutEndiannessLittle))
326         return success();
327       return emitError(loc) << "'" << entryName
328                             << "' data layout entry is expected to be either '"
329                             << DLTIDialect::kDataLayoutEndiannessBig << "' or '"
330                             << DLTIDialect::kDataLayoutEndiannessLittle << "'";
331     }
332     return emitError(loc) << "unknown data layout entry name: " << entryName;
333   }
334 };
335 } // namespace
336 
337 void DLTIDialect::initialize() {
338   addAttributes<DataLayoutEntryAttr, DataLayoutSpecAttr>();
339   addInterfaces<TargetDataLayoutInterface>();
340 }
341 
342 Attribute DLTIDialect::parseAttribute(DialectAsmParser &parser,
343                                       Type type) const {
344   StringRef attrKind;
345   if (parser.parseKeyword(&attrKind))
346     return {};
347 
348   if (attrKind == DataLayoutEntryAttr::kAttrKeyword)
349     return DataLayoutEntryAttr::parse(parser);
350   if (attrKind == DataLayoutSpecAttr::kAttrKeyword)
351     return DataLayoutSpecAttr::parse(parser);
352 
353   parser.emitError(parser.getNameLoc(), "unknown attrribute type: ")
354       << attrKind;
355   return {};
356 }
357 
358 void DLTIDialect::printAttribute(Attribute attr, DialectAsmPrinter &os) const {
359   llvm::TypeSwitch<Attribute>(attr)
360       .Case<DataLayoutEntryAttr, DataLayoutSpecAttr>(
361           [&](auto a) { a.print(os); })
362       .Default([](Attribute) { llvm_unreachable("unknown attribute kind"); });
363 }
364 
365 LogicalResult DLTIDialect::verifyOperationAttribute(Operation *op,
366                                                     NamedAttribute attr) {
367   if (attr.first == DLTIDialect::kDataLayoutAttrName) {
368     if (!attr.second.isa<DataLayoutSpecAttr>()) {
369       return op->emitError() << "'" << DLTIDialect::kDataLayoutAttrName
370                              << "' is expected to be a #dlti.dl_spec attribute";
371     }
372     return success();
373   }
374 
375   return op->emitError() << "attribute '" << attr.first
376                          << "' not supported by dialect";
377 }
378