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