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