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
26namespace mlir {
27class AsmResourcePrinter;
28class AsmDialectResourceHandle;
29class Operation;
30
31namespace detail {
32class 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).
88class AsmResourceBlob {
89public:
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
167private:
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.
184class HeapAsmResourceBlob {
185public:
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.
214class UnmanagedAsmResourceBlob {
215public:
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.
239class AsmResourceBuilder {
240public:
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.
272enum 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};
280StringRef toString(AsmResourceEntryKind kind);
281
282/// This class represents a single parsed resource entry.
283class AsmParsedResourceEntry {
284public:
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.
330class AsmResourceParser {
331public:
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
363private:
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.
370class AsmResourcePrinter {
371public:
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
406private:
407 std::string name;
408};
409
410/// A fallback map containing external resources not explicitly handled by
411/// another parser/printer.
412class FallbackAsmResourceMap {
413public:
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
434private:
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.
460class ParserConfig {
461public:
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
513private:
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.
533class AsmState {
534public:
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
589private:
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.
602void registerAsmPrinterCLOptions();
603
604} // namespace mlir
605
606#endif // MLIR_IR_ASMSTATE_H_
607

source code of mlir/include/mlir/IR/AsmState.h