1 | //===- AsmState.h - Assembly State Utilities --------------------*- C++ -*-===// |
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 | // This file defines various classes and utilites for interacting with the MLIR |
10 | // assembly formats. |
11 | // |
12 | //===----------------------------------------------------------------------===// |
13 | |
14 | #ifndef MLIR_IR_ASMSTATE_H_ |
15 | #define MLIR_IR_ASMSTATE_H_ |
16 | |
17 | #include "mlir/Bytecode/BytecodeReaderConfig.h" |
18 | #include "mlir/IR/OperationSupport.h" |
19 | #include "mlir/Support/LLVM.h" |
20 | #include "llvm/ADT/MapVector.h" |
21 | #include "llvm/ADT/StringMap.h" |
22 | |
23 | #include <memory> |
24 | #include <variant> |
25 | |
26 | namespace mlir { |
27 | class AsmResourcePrinter; |
28 | class AsmDialectResourceHandle; |
29 | class Operation; |
30 | |
31 | namespace detail { |
32 | class AsmStateImpl; |
33 | } // namespace detail |
34 | |
35 | //===----------------------------------------------------------------------===// |
36 | // Resources |
37 | //===----------------------------------------------------------------------===// |
38 | |
39 | /// The following classes enable support for parsing and printing resources |
40 | /// within MLIR assembly formats. Resources are a mechanism by which dialects, |
41 | /// and external clients, may attach additional information when parsing or |
42 | /// printing IR without that information being encoded in the IR itself. |
43 | /// Resources are not uniqued within the MLIR context, are not attached directly |
44 | /// to any operation, and are solely intended to live and be processed outside |
45 | /// of the immediate IR. |
46 | /// |
47 | /// Resources are encoded using a key-value pair nested within dictionaries |
48 | /// anchored either on a dialect, or an externally registered entity. |
49 | /// Dictionaries anchored on dialects use the dialect namespace directly, and |
50 | /// dictionaries anchored on external entities use a provided unique identifier. |
51 | /// The resource key is an identifier used to disambiguate the data. The |
52 | /// resource value may be stored in various limited forms, but general encodings |
53 | /// use a string (human readable) or blob format (binary). Within the textual |
54 | /// format, an example may be of the form: |
55 | /// |
56 | /// {-# |
57 | /// // The `dialect_resources` section within the file-level metadata |
58 | /// // dictionary is used to contain any dialect resource entries. |
59 | /// dialect_resources: { |
60 | /// // Here is a dictionary anchored on "foo_dialect", which is a dialect |
61 | /// // namespace. |
62 | /// foo_dialect: { |
63 | /// // `some_dialect_resource` is a key to be interpreted by the dialect, |
64 | /// // and used to initialize/configure/etc. |
65 | /// some_dialect_resource: "Some important resource value" |
66 | /// } |
67 | /// }, |
68 | /// // The `external_resources` section within the file-level metadata |
69 | /// // dictionary is used to contain any non-dialect resource entries. |
70 | /// external_resources: { |
71 | /// // Here is a dictionary anchored on "mlir_reproducer", which is an |
72 | /// // external entity representing MLIR's crash reproducer functionality. |
73 | /// mlir_reproducer: { |
74 | /// // `pipeline` is an entry that holds a crash reproducer pipeline |
75 | /// // resource. |
76 | /// pipeline: "func.func(canonicalize,cse)" |
77 | /// } |
78 | /// } |
79 | /// #-} |
80 | /// |
81 | |
82 | //===----------------------------------------------------------------------===// |
83 | // Resource Entry |
84 | |
85 | /// This class represents a processed binary blob of data. A resource blob is |
86 | /// essentially a collection of data, potentially mutable, with an associated |
87 | /// deleter function (used if the data needs to be destroyed). |
88 | class AsmResourceBlob { |
89 | public: |
90 | /// A deleter function that frees a blob given the data, allocation size, and |
91 | /// allocation aligment. |
92 | using DeleterFn = |
93 | llvm::unique_function<void(void *data, size_t size, size_t align)>; |
94 | |
95 | //===--------------------------------------------------------------------===// |
96 | // Construction |
97 | //===--------------------------------------------------------------------===// |
98 | |
99 | AsmResourceBlob() = default; |
100 | AsmResourceBlob(ArrayRef<char> data, size_t dataAlignment, DeleterFn deleter, |
101 | bool dataIsMutable) |
102 | : data(data), dataAlignment(dataAlignment), deleter(std::move(deleter)), |
103 | dataIsMutable(dataIsMutable) {} |
104 | /// Utility constructor that initializes a blob with a non-char type T. |
105 | template <typename T, typename DelT> |
106 | AsmResourceBlob(ArrayRef<T> data, DelT &&deleteFn, bool dataIsMutable) |
107 | : data((const char *)data.data(), data.size() * sizeof(T)), |
108 | dataAlignment(alignof(T)), |
109 | deleter([deleteFn = std::forward<DelT>(deleteFn)]( |
110 | void *data, size_t size, size_t align) { |
111 | return deleteFn((T *)data, size, align); |
112 | }), |
113 | dataIsMutable(dataIsMutable) {} |
114 | AsmResourceBlob(AsmResourceBlob &&) = default; |
115 | AsmResourceBlob &operator=(AsmResourceBlob &&rhs) { |
116 | // Delete the current blob if necessary. |
117 | if (deleter) |
118 | deleter(const_cast<char *>(data.data()), data.size(), dataAlignment); |
119 | |
120 | // Take the data entries from rhs. |
121 | data = rhs.data; |
122 | dataAlignment = rhs.dataAlignment; |
123 | deleter = std::move(rhs.deleter); |
124 | dataIsMutable = rhs.dataIsMutable; |
125 | return *this; |
126 | } |
127 | AsmResourceBlob(const AsmResourceBlob &) = delete; |
128 | AsmResourceBlob &operator=(const AsmResourceBlob &) = delete; |
129 | ~AsmResourceBlob() { |
130 | if (deleter) |
131 | deleter(const_cast<char *>(data.data()), data.size(), dataAlignment); |
132 | } |
133 | |
134 | //===--------------------------------------------------------------------===// |
135 | // Data Access |
136 | //===--------------------------------------------------------------------===// |
137 | |
138 | /// Return the alignment of the underlying data. |
139 | size_t getDataAlignment() const { return dataAlignment; } |
140 | |
141 | /// Return the raw underlying data of this blob. |
142 | ArrayRef<char> getData() const { return data; } |
143 | |
144 | /// Return the underlying data as an array of the given type. This is an |
145 | /// inherrently unsafe operation, and should only be used when the data is |
146 | /// known to be of the correct type. |
147 | template <typename T> |
148 | ArrayRef<T> getDataAs() const { |
149 | return llvm::ArrayRef<T>((const T *)data.data(), data.size() / sizeof(T)); |
150 | } |
151 | |
152 | /// Return a mutable reference to the raw underlying data of this blob. |
153 | /// Asserts that the blob `isMutable`. |
154 | MutableArrayRef<char> getMutableData() { |
155 | assert(isMutable() && |
156 | "cannot access mutable reference to non-mutable data" ); |
157 | return MutableArrayRef<char>(const_cast<char *>(data.data()), data.size()); |
158 | } |
159 | |
160 | /// Return if the data of this blob is mutable. |
161 | bool isMutable() const { return dataIsMutable; } |
162 | |
163 | /// Return the deleter function of this blob. |
164 | DeleterFn &getDeleter() { return deleter; } |
165 | const DeleterFn &getDeleter() const { return deleter; } |
166 | |
167 | private: |
168 | /// The raw, properly aligned, blob data. |
169 | ArrayRef<char> data; |
170 | |
171 | /// The alignment of the data. |
172 | size_t dataAlignment = 0; |
173 | |
174 | /// An optional deleter function used to deallocate the underlying data when |
175 | /// necessary. |
176 | DeleterFn deleter; |
177 | |
178 | /// Whether the data is mutable. |
179 | bool dataIsMutable; |
180 | }; |
181 | |
182 | /// This class provides a simple utility wrapper for creating heap allocated |
183 | /// AsmResourceBlobs. |
184 | class HeapAsmResourceBlob { |
185 | public: |
186 | /// Create a new heap allocated blob with the given size and alignment. |
187 | /// `dataIsMutable` indicates if the allocated data can be mutated. By |
188 | /// default, we treat heap allocated blobs as mutable. |
189 | static AsmResourceBlob allocate(size_t size, size_t align, |
190 | bool dataIsMutable = true) { |
191 | return AsmResourceBlob( |
192 | ArrayRef<char>((char *)llvm::allocate_buffer(Size: size, Alignment: align), size), align, |
193 | llvm::deallocate_buffer, dataIsMutable); |
194 | } |
195 | /// Create a new heap allocated blob and copy the provided data into it. |
196 | static AsmResourceBlob allocateAndCopyWithAlign(ArrayRef<char> data, |
197 | size_t align, |
198 | bool dataIsMutable = true) { |
199 | AsmResourceBlob blob = allocate(size: data.size(), align, dataIsMutable); |
200 | std::memcpy(dest: blob.getMutableData().data(), src: data.data(), n: data.size()); |
201 | return blob; |
202 | } |
203 | template <typename T> |
204 | static AsmResourceBlob allocateAndCopyInferAlign(ArrayRef<T> data, |
205 | bool dataIsMutable = true) { |
206 | return allocateAndCopyWithAlign( |
207 | data: ArrayRef<char>((const char *)data.data(), data.size() * sizeof(T)), |
208 | align: alignof(T), dataIsMutable); |
209 | } |
210 | }; |
211 | /// This class provides a simple utility wrapper for creating "unmanaged" |
212 | /// AsmResourceBlobs. The lifetime of the data provided to these blobs is |
213 | /// guaranteed to persist beyond the lifetime of this reference. |
214 | class UnmanagedAsmResourceBlob { |
215 | public: |
216 | /// Create a new unmanaged resource directly referencing the provided data. |
217 | /// `dataIsMutable` indicates if the allocated data can be mutated. By |
218 | /// default, we treat unmanaged blobs as immutable. |
219 | static AsmResourceBlob |
220 | allocateWithAlign(ArrayRef<char> data, size_t align, |
221 | AsmResourceBlob::DeleterFn deleter = {}, |
222 | bool dataIsMutable = false) { |
223 | return AsmResourceBlob(data, align, std::move(deleter), dataIsMutable); |
224 | } |
225 | template <typename T> |
226 | static AsmResourceBlob |
227 | allocateInferAlign(ArrayRef<T> data, AsmResourceBlob::DeleterFn deleter = {}, |
228 | bool dataIsMutable = false) { |
229 | return allocateWithAlign( |
230 | data: ArrayRef<char>((const char *)data.data(), data.size() * sizeof(T)), |
231 | align: alignof(T), deleter: std::move(deleter), dataIsMutable); |
232 | } |
233 | }; |
234 | |
235 | /// This class is used to build resource entries for use by the printer. Each |
236 | /// resource entry is represented using a key/value pair. The provided key must |
237 | /// be unique within the current context, which allows for a client to provide |
238 | /// resource entries without worrying about overlap with other clients. |
239 | class AsmResourceBuilder { |
240 | public: |
241 | virtual ~AsmResourceBuilder(); |
242 | |
243 | /// Build a resource entry represented by the given bool. |
244 | virtual void buildBool(StringRef key, bool data) = 0; |
245 | |
246 | /// Build a resource entry represented by the given human-readable string |
247 | /// value. |
248 | virtual void buildString(StringRef key, StringRef data) = 0; |
249 | |
250 | /// Build an resource entry represented by the given binary blob data. |
251 | virtual void buildBlob(StringRef key, ArrayRef<char> data, |
252 | uint32_t dataAlignment) = 0; |
253 | /// Build an resource entry represented by the given binary blob data. This is |
254 | /// a useful overload if the data type is known. Note that this does not |
255 | /// support `char` element types to avoid accidentally not providing the |
256 | /// expected alignment of data in situations that treat blobs generically. |
257 | template <typename T> |
258 | std::enable_if_t<!std::is_same<T, char>::value> buildBlob(StringRef key, |
259 | ArrayRef<T> data) { |
260 | buildBlob( |
261 | key, data: ArrayRef<char>((const char *)data.data(), data.size() * sizeof(T)), |
262 | dataAlignment: alignof(T)); |
263 | } |
264 | /// Build an resource entry represented by the given resource blob. This is |
265 | /// a useful overload if a blob already exists in-memory. |
266 | void buildBlob(StringRef key, const AsmResourceBlob &blob) { |
267 | buildBlob(key, data: blob.getData(), dataAlignment: blob.getDataAlignment()); |
268 | } |
269 | }; |
270 | |
271 | /// This enum represents the different kinds of resource values. |
272 | enum class AsmResourceEntryKind { |
273 | /// A blob of data with an accompanying alignment. |
274 | Blob, |
275 | /// A boolean value. |
276 | Bool, |
277 | /// A string value. |
278 | String, |
279 | }; |
280 | StringRef toString(AsmResourceEntryKind kind); |
281 | |
282 | /// This class represents a single parsed resource entry. |
283 | class AsmParsedResourceEntry { |
284 | public: |
285 | virtual ~AsmParsedResourceEntry(); |
286 | |
287 | /// Return the key of the resource entry. |
288 | virtual StringRef getKey() const = 0; |
289 | |
290 | /// Emit an error at the location of this entry. |
291 | virtual InFlightDiagnostic emitError() const = 0; |
292 | |
293 | /// Return the kind of this value. |
294 | virtual AsmResourceEntryKind getKind() const = 0; |
295 | |
296 | /// Parse the resource entry represented by a boolean. Returns failure if the |
297 | /// entry does not correspond to a bool. |
298 | virtual FailureOr<bool> parseAsBool() const = 0; |
299 | |
300 | /// Parse the resource entry represented by a human-readable string. Returns |
301 | /// failure if the entry does not correspond to a string. |
302 | virtual FailureOr<std::string> parseAsString() const = 0; |
303 | |
304 | /// An allocator function used to allocate memory for a blob when required. |
305 | /// The function is provided a size and alignment, and should return an |
306 | /// aligned allocation buffer. |
307 | using BlobAllocatorFn = |
308 | function_ref<AsmResourceBlob(size_t size, size_t align)>; |
309 | |
310 | /// Parse the resource entry represented by a binary blob. Returns failure if |
311 | /// the entry does not correspond to a blob. If the blob needed to be |
312 | /// allocated, the given allocator function is invoked. |
313 | virtual FailureOr<AsmResourceBlob> |
314 | parseAsBlob(BlobAllocatorFn allocator) const = 0; |
315 | /// Parse the resource entry represented by a binary blob using heap |
316 | /// allocation. |
317 | FailureOr<AsmResourceBlob> parseAsBlob() const { |
318 | return parseAsBlob(allocator: [](size_t size, size_t align) { |
319 | return HeapAsmResourceBlob::allocate(size, align); |
320 | }); |
321 | } |
322 | }; |
323 | |
324 | //===----------------------------------------------------------------------===// |
325 | // Resource Parser/Printer |
326 | |
327 | /// This class represents an instance of a resource parser. This class should be |
328 | /// implemented by non-dialect clients that want to inject additional resources |
329 | /// into MLIR assembly formats. |
330 | class AsmResourceParser { |
331 | public: |
332 | /// Create a new parser with the given identifying name. This name uniquely |
333 | /// identifies the entries of this parser, and differentiates them from other |
334 | /// contexts. |
335 | AsmResourceParser(StringRef name) : name(name.str()) {} |
336 | virtual ~AsmResourceParser(); |
337 | |
338 | /// Return the name of this parser. |
339 | StringRef getName() const { return name; } |
340 | |
341 | /// Parse the given resource entry. Returns failure if the key/data were not |
342 | /// valid, or could otherwise not be processed correctly. Any necessary errors |
343 | /// should be emitted with the provided entry. |
344 | virtual LogicalResult parseResource(AsmParsedResourceEntry &entry) = 0; |
345 | |
346 | /// Return a resource parser implemented via the given callable, whose form |
347 | /// should match that of `parseResource` above. |
348 | template <typename CallableT> |
349 | static std::unique_ptr<AsmResourceParser> fromCallable(StringRef name, |
350 | CallableT &&parseFn) { |
351 | struct Processor : public AsmResourceParser { |
352 | Processor(StringRef name, CallableT &&parseFn) |
353 | : AsmResourceParser(name), parseFn(std::move(parseFn)) {} |
354 | LogicalResult parseResource(AsmParsedResourceEntry &entry) override { |
355 | return parseFn(entry); |
356 | } |
357 | |
358 | std::decay_t<CallableT> parseFn; |
359 | }; |
360 | return std::make_unique<Processor>(name, std::forward<CallableT>(parseFn)); |
361 | } |
362 | |
363 | private: |
364 | std::string name; |
365 | }; |
366 | |
367 | /// This class represents an instance of a resource printer. This class should |
368 | /// be implemented by non-dialect clients that want to inject additional |
369 | /// resources into MLIR assembly formats. |
370 | class AsmResourcePrinter { |
371 | public: |
372 | /// Create a new printer with the given identifying name. This name uniquely |
373 | /// identifies the entries of this printer, and differentiates them from |
374 | /// other contexts. |
375 | AsmResourcePrinter(StringRef name) : name(name.str()) {} |
376 | virtual ~AsmResourcePrinter(); |
377 | |
378 | /// Return the name of this printer. |
379 | StringRef getName() const { return name; } |
380 | |
381 | /// Build any resources to include during printing, utilizing the given |
382 | /// top-level root operation to help determine what information to include. |
383 | /// Provided data should be registered in the form of a key/data pair, to the |
384 | /// given builder. |
385 | virtual void buildResources(Operation *op, |
386 | AsmResourceBuilder &builder) const = 0; |
387 | |
388 | /// Return a resource printer implemented via the given callable, whose form |
389 | /// should match that of `buildResources` above. |
390 | template <typename CallableT> |
391 | static std::unique_ptr<AsmResourcePrinter> fromCallable(StringRef name, |
392 | CallableT &&printFn) { |
393 | struct Printer : public AsmResourcePrinter { |
394 | Printer(StringRef name, CallableT &&printFn) |
395 | : AsmResourcePrinter(name), printFn(std::move(printFn)) {} |
396 | void buildResources(Operation *op, |
397 | AsmResourceBuilder &builder) const override { |
398 | printFn(op, builder); |
399 | } |
400 | |
401 | std::decay_t<CallableT> printFn; |
402 | }; |
403 | return std::make_unique<Printer>(name, std::forward<CallableT>(printFn)); |
404 | } |
405 | |
406 | private: |
407 | std::string name; |
408 | }; |
409 | |
410 | /// A fallback map containing external resources not explicitly handled by |
411 | /// another parser/printer. |
412 | class FallbackAsmResourceMap { |
413 | public: |
414 | /// This class represents an opaque resource. |
415 | struct OpaqueAsmResource { |
416 | OpaqueAsmResource(StringRef key, |
417 | std::variant<AsmResourceBlob, bool, std::string> value) |
418 | : key(key.str()), value(std::move(value)) {} |
419 | |
420 | /// The key identifying the resource. |
421 | std::string key; |
422 | /// An opaque value for the resource, whose variant values align 1-1 with |
423 | /// the kinds defined in AsmResourceEntryKind. |
424 | std::variant<AsmResourceBlob, bool, std::string> value; |
425 | }; |
426 | |
427 | /// Return a parser than can be used for parsing entries for the given |
428 | /// identifier key. |
429 | AsmResourceParser &getParserFor(StringRef key); |
430 | |
431 | /// Build a set of resource printers to print the resources within this map. |
432 | std::vector<std::unique_ptr<AsmResourcePrinter>> getPrinters(); |
433 | |
434 | private: |
435 | struct ResourceCollection : public AsmResourceParser { |
436 | ResourceCollection(StringRef name) : AsmResourceParser(name) {} |
437 | |
438 | /// Parse a resource into this collection. |
439 | LogicalResult parseResource(AsmParsedResourceEntry &entry) final; |
440 | |
441 | /// Build the resources held by this collection. |
442 | void buildResources(Operation *op, AsmResourceBuilder &builder) const; |
443 | |
444 | /// The set of resources parsed into this collection. |
445 | SmallVector<OpaqueAsmResource> resources; |
446 | }; |
447 | |
448 | /// The set of opaque resources. |
449 | llvm::MapVector<std::string, std::unique_ptr<ResourceCollection>, |
450 | llvm::StringMap<unsigned>> |
451 | keyToResources; |
452 | }; |
453 | |
454 | //===----------------------------------------------------------------------===// |
455 | // ParserConfig |
456 | //===----------------------------------------------------------------------===// |
457 | |
458 | /// This class represents a configuration for the MLIR assembly parser. It |
459 | /// contains all of the necessary state to parse a MLIR source file. |
460 | class ParserConfig { |
461 | public: |
462 | /// Construct a parser configuration with the given context. |
463 | /// `verifyAfterParse` indicates if the IR should be verified after parsing. |
464 | /// `fallbackResourceMap` is an optional fallback handler that can be used to |
465 | /// parse external resources not explicitly handled by another parser. |
466 | ParserConfig(MLIRContext *context, bool verifyAfterParse = true, |
467 | FallbackAsmResourceMap *fallbackResourceMap = nullptr) |
468 | : context(context), verifyAfterParse(verifyAfterParse), |
469 | fallbackResourceMap(fallbackResourceMap) { |
470 | assert(context && "expected valid MLIR context" ); |
471 | } |
472 | |
473 | /// Return the MLIRContext to be used when parsing. |
474 | MLIRContext *getContext() const { return context; } |
475 | |
476 | /// Returns if the parser should verify the IR after parsing. |
477 | bool shouldVerifyAfterParse() const { return verifyAfterParse; } |
478 | |
479 | /// Returns the parsing configurations associated to the bytecode read. |
480 | BytecodeReaderConfig &getBytecodeReaderConfig() const { |
481 | return const_cast<BytecodeReaderConfig &>(bytecodeReaderConfig); |
482 | } |
483 | |
484 | /// Return the resource parser registered to the given name, or nullptr if no |
485 | /// parser with `name` is registered. |
486 | AsmResourceParser *getResourceParser(StringRef name) const { |
487 | auto it = resourceParsers.find(Val: name); |
488 | if (it != resourceParsers.end()) |
489 | return it->second.get(); |
490 | if (fallbackResourceMap) |
491 | return &fallbackResourceMap->getParserFor(key: name); |
492 | return nullptr; |
493 | } |
494 | |
495 | /// Attach the given resource parser. |
496 | void attachResourceParser(std::unique_ptr<AsmResourceParser> parser) { |
497 | StringRef name = parser->getName(); |
498 | auto it = resourceParsers.try_emplace(Key: name, Args: std::move(parser)); |
499 | (void)it; |
500 | assert(it.second && |
501 | "resource parser already registered with the given name" ); |
502 | } |
503 | |
504 | /// Attach the given callable resource parser with the given name. |
505 | template <typename CallableT> |
506 | std::enable_if_t<std::is_convertible< |
507 | CallableT, function_ref<LogicalResult(AsmParsedResourceEntry &)>>::value> |
508 | attachResourceParser(StringRef name, CallableT &&parserFn) { |
509 | attachResourceParser(AsmResourceParser::fromCallable( |
510 | name, std::forward<CallableT>(parserFn))); |
511 | } |
512 | |
513 | private: |
514 | MLIRContext *context; |
515 | bool verifyAfterParse; |
516 | DenseMap<StringRef, std::unique_ptr<AsmResourceParser>> resourceParsers; |
517 | FallbackAsmResourceMap *fallbackResourceMap; |
518 | BytecodeReaderConfig bytecodeReaderConfig; |
519 | }; |
520 | |
521 | //===----------------------------------------------------------------------===// |
522 | // AsmState |
523 | //===----------------------------------------------------------------------===// |
524 | |
525 | /// This class provides management for the lifetime of the state used when |
526 | /// printing the IR. It allows for alleviating the cost of recomputing the |
527 | /// internal state of the asm printer. |
528 | /// |
529 | /// The IR should not be mutated in-between invocations using this state, and |
530 | /// the IR being printed must not be an parent of the IR originally used to |
531 | /// initialize this state. This means that if a child operation is provided, a |
532 | /// parent operation cannot reuse this state. |
533 | class AsmState { |
534 | public: |
535 | /// This map represents the raw locations of operations within the output |
536 | /// stream. This maps the original pointer to the operation, to a pair of line |
537 | /// and column in the output stream. |
538 | using LocationMap = DenseMap<Operation *, std::pair<unsigned, unsigned>>; |
539 | |
540 | /// Initialize the asm state at the level of the given operation. A location |
541 | /// map may optionally be provided to be populated when printing. `map` is an |
542 | /// optional fallback resource map, which when provided will attach resource |
543 | /// printers for the fallback resources within the map. |
544 | AsmState(Operation *op, |
545 | const OpPrintingFlags &printerFlags = OpPrintingFlags(), |
546 | LocationMap *locationMap = nullptr, |
547 | FallbackAsmResourceMap *map = nullptr); |
548 | AsmState(MLIRContext *ctx, |
549 | const OpPrintingFlags &printerFlags = OpPrintingFlags(), |
550 | LocationMap *locationMap = nullptr, |
551 | FallbackAsmResourceMap *map = nullptr); |
552 | ~AsmState(); |
553 | |
554 | /// Get the printer flags. |
555 | const OpPrintingFlags &getPrinterFlags() const; |
556 | |
557 | /// Return an instance of the internal implementation. Returns nullptr if the |
558 | /// state has not been initialized. |
559 | detail::AsmStateImpl &getImpl() { return *impl; } |
560 | |
561 | //===--------------------------------------------------------------------===// |
562 | // Resources |
563 | //===--------------------------------------------------------------------===// |
564 | |
565 | /// Attach the given resource printer to the AsmState. |
566 | void attachResourcePrinter(std::unique_ptr<AsmResourcePrinter> printer); |
567 | |
568 | /// Attach an resource printer, in the form of a callable, to the AsmState. |
569 | template <typename CallableT> |
570 | std::enable_if_t<std::is_convertible< |
571 | CallableT, function_ref<void(Operation *, AsmResourceBuilder &)>>::value> |
572 | attachResourcePrinter(StringRef name, CallableT &&printFn) { |
573 | attachResourcePrinter(AsmResourcePrinter::fromCallable( |
574 | name, std::forward<CallableT>(printFn))); |
575 | } |
576 | |
577 | /// Attach resource printers to the AsmState for the fallback resources |
578 | /// in the given map. |
579 | void attachFallbackResourcePrinter(FallbackAsmResourceMap &map) { |
580 | for (auto &printer : map.getPrinters()) |
581 | attachResourcePrinter(printer: std::move(printer)); |
582 | } |
583 | |
584 | /// Returns a map of dialect resources that were referenced when using this |
585 | /// state to print IR. |
586 | DenseMap<Dialect *, SetVector<AsmDialectResourceHandle>> & |
587 | getDialectResources() const; |
588 | |
589 | private: |
590 | AsmState() = delete; |
591 | |
592 | /// A pointer to allocated storage for the impl state. |
593 | std::unique_ptr<detail::AsmStateImpl> impl; |
594 | }; |
595 | |
596 | //===----------------------------------------------------------------------===// |
597 | // AsmPrinter CommandLine Options |
598 | //===----------------------------------------------------------------------===// |
599 | |
600 | /// Register a set of useful command-line options that can be used to configure |
601 | /// various flags within the AsmPrinter. |
602 | void registerAsmPrinterCLOptions(); |
603 | |
604 | } // namespace mlir |
605 | |
606 | #endif // MLIR_IR_ASMSTATE_H_ |
607 | |