1 | //===- BytecodeWriter.h - MLIR Bytecode Writer ------------------*- 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 header defines interfaces to write MLIR bytecode files/streams. |
10 | // |
11 | //===----------------------------------------------------------------------===// |
12 | |
13 | #ifndef MLIR_BYTECODE_BYTECODEWRITER_H |
14 | #define MLIR_BYTECODE_BYTECODEWRITER_H |
15 | |
16 | #include "mlir/IR/AsmState.h" |
17 | |
18 | namespace mlir { |
19 | class DialectBytecodeWriter; |
20 | class DialectVersion; |
21 | class Operation; |
22 | |
23 | /// A class to interact with the attributes and types printer when emitting MLIR |
24 | /// bytecode. |
25 | template <class T> |
26 | class AttrTypeBytecodeWriter { |
27 | public: |
28 | AttrTypeBytecodeWriter() = default; |
29 | virtual ~AttrTypeBytecodeWriter() = default; |
30 | |
31 | /// Callback writer API used in IRNumbering, where groups are created and |
32 | /// type/attribute components are numbered. At this stage, writer is expected |
33 | /// to be a `NumberingDialectWriter`. |
34 | virtual LogicalResult write(T entry, std::optional<StringRef> &name, |
35 | DialectBytecodeWriter &writer) = 0; |
36 | |
37 | /// Callback writer API used in BytecodeWriter, where groups are created and |
38 | /// type/attribute components are numbered. Here, DialectBytecodeWriter is |
39 | /// expected to be an actual writer. The optional stringref specified by |
40 | /// the user is ignored, since the group was already specified when numbering |
41 | /// the IR. |
42 | LogicalResult write(T entry, DialectBytecodeWriter &writer) { |
43 | std::optional<StringRef> dummy; |
44 | return write(entry, dummy, writer); |
45 | } |
46 | |
47 | /// Return an Attribute/Type printer implemented via the given callable, whose |
48 | /// form should match that of the `write` function above. |
49 | template <typename CallableT, |
50 | std::enable_if_t<std::is_convertible_v< |
51 | CallableT, std::function<LogicalResult( |
52 | T, std::optional<StringRef> &, |
53 | DialectBytecodeWriter &)>>, |
54 | bool> = true> |
55 | static std::unique_ptr<AttrTypeBytecodeWriter<T>> |
56 | fromCallable(CallableT &&writeFn) { |
57 | struct Processor : public AttrTypeBytecodeWriter<T> { |
58 | Processor(CallableT &&writeFn) |
59 | : AttrTypeBytecodeWriter(), writeFn(std::move(writeFn)) {} |
60 | LogicalResult write(T entry, std::optional<StringRef> &name, |
61 | DialectBytecodeWriter &writer) override { |
62 | return writeFn(entry, name, writer); |
63 | } |
64 | |
65 | std::decay_t<CallableT> writeFn; |
66 | }; |
67 | return std::make_unique<Processor>(std::forward<CallableT>(writeFn)); |
68 | } |
69 | }; |
70 | |
71 | /// This class contains the configuration used for the bytecode writer. It |
72 | /// controls various aspects of bytecode generation, and contains all of the |
73 | /// various bytecode writer hooks. |
74 | class BytecodeWriterConfig { |
75 | public: |
76 | /// `producer` is an optional string that can be used to identify the producer |
77 | /// of the bytecode when reading. It has no functional effect on the bytecode |
78 | /// serialization. |
79 | BytecodeWriterConfig(StringRef producer = "MLIR" LLVM_VERSION_STRING); |
80 | /// `map` is a fallback resource map, which when provided will attach resource |
81 | /// printers for the fallback resources within the map. |
82 | BytecodeWriterConfig(FallbackAsmResourceMap &map, |
83 | StringRef producer = "MLIR" LLVM_VERSION_STRING); |
84 | ~BytecodeWriterConfig(); |
85 | |
86 | /// An internal implementation class that contains the state of the |
87 | /// configuration. |
88 | struct Impl; |
89 | |
90 | /// Return an instance of the internal implementation. |
91 | const Impl &getImpl() const { return *impl; } |
92 | |
93 | /// Set the desired bytecode version to emit. This method does not validate |
94 | /// the desired version. The bytecode writer entry point will return failure |
95 | /// if it cannot emit the desired version. |
96 | void setDesiredBytecodeVersion(int64_t bytecodeVersion); |
97 | |
98 | /// Get the set desired bytecode version to emit. |
99 | int64_t getDesiredBytecodeVersion() const; |
100 | |
101 | /// A map containing the dialect versions to emit. |
102 | llvm::StringMap<std::unique_ptr<DialectVersion>> & |
103 | getDialectVersionMap() const; |
104 | |
105 | /// Set a given dialect version to emit on the map. |
106 | template <class T> |
107 | void setDialectVersion(std::unique_ptr<DialectVersion> dialectVersion) const { |
108 | return setDialectVersion(T::getDialectNamespace(), |
109 | std::move(dialectVersion)); |
110 | } |
111 | void setDialectVersion(StringRef dialectName, |
112 | std::unique_ptr<DialectVersion> dialectVersion) const; |
113 | |
114 | //===--------------------------------------------------------------------===// |
115 | // Types and Attributes encoding |
116 | //===--------------------------------------------------------------------===// |
117 | |
118 | /// Retrieve the callbacks. |
119 | ArrayRef<std::unique_ptr<AttrTypeBytecodeWriter<Attribute>>> |
120 | getAttributeWriterCallbacks() const; |
121 | ArrayRef<std::unique_ptr<AttrTypeBytecodeWriter<Type>>> |
122 | getTypeWriterCallbacks() const; |
123 | |
124 | /// Attach a custom bytecode printer callback to the configuration for the |
125 | /// emission of custom type/attributes encodings. |
126 | void attachAttributeCallback( |
127 | std::unique_ptr<AttrTypeBytecodeWriter<Attribute>> callback); |
128 | void |
129 | attachTypeCallback(std::unique_ptr<AttrTypeBytecodeWriter<Type>> callback); |
130 | |
131 | /// Attach a custom bytecode printer callback to the configuration for the |
132 | /// emission of custom type/attributes encodings. |
133 | template <typename CallableT> |
134 | std::enable_if_t<std::is_convertible_v< |
135 | CallableT, |
136 | std::function<LogicalResult(Attribute, std::optional<StringRef> &, |
137 | DialectBytecodeWriter &)>>> |
138 | attachAttributeCallback(CallableT &&emitFn) { |
139 | attachAttributeCallback(AttrTypeBytecodeWriter<Attribute>::fromCallable( |
140 | std::forward<CallableT>(emitFn))); |
141 | } |
142 | template <typename CallableT> |
143 | std::enable_if_t<std::is_convertible_v< |
144 | CallableT, std::function<LogicalResult(Type, std::optional<StringRef> &, |
145 | DialectBytecodeWriter &)>>> |
146 | attachTypeCallback(CallableT &&emitFn) { |
147 | attachTypeCallback(AttrTypeBytecodeWriter<Type>::fromCallable( |
148 | std::forward<CallableT>(emitFn))); |
149 | } |
150 | |
151 | //===--------------------------------------------------------------------===// |
152 | // Resources |
153 | //===--------------------------------------------------------------------===// |
154 | |
155 | /// Set a boolean flag to skip emission of resources into the bytecode file. |
156 | void setElideResourceDataFlag(bool shouldElideResourceData = true); |
157 | |
158 | /// Attach the given resource printer to the writer configuration. |
159 | void attachResourcePrinter(std::unique_ptr<AsmResourcePrinter> printer); |
160 | |
161 | /// Attach an resource printer, in the form of a callable, to the |
162 | /// configuration. |
163 | template <typename CallableT> |
164 | std::enable_if_t<std::is_convertible< |
165 | CallableT, function_ref<void(Operation *, AsmResourceBuilder &)>>::value> |
166 | attachResourcePrinter(StringRef name, CallableT &&printFn) { |
167 | attachResourcePrinter(AsmResourcePrinter::fromCallable( |
168 | name, std::forward<CallableT>(printFn))); |
169 | } |
170 | |
171 | /// Attach resource printers to the AsmState for the fallback resources |
172 | /// in the given map. |
173 | void attachFallbackResourcePrinter(FallbackAsmResourceMap &map) { |
174 | for (auto &printer : map.getPrinters()) |
175 | attachResourcePrinter(printer: std::move(printer)); |
176 | } |
177 | |
178 | private: |
179 | /// A pointer to allocated storage for the impl state. |
180 | std::unique_ptr<Impl> impl; |
181 | }; |
182 | |
183 | //===----------------------------------------------------------------------===// |
184 | // Entry Points |
185 | //===----------------------------------------------------------------------===// |
186 | |
187 | /// Write the bytecode for the given operation to the provided output stream. |
188 | /// For streams where it matters, the given stream should be in "binary" mode. |
189 | /// It only ever fails if setDesiredByteCodeVersion can't be honored. |
190 | LogicalResult writeBytecodeToFile(Operation *op, raw_ostream &os, |
191 | const BytecodeWriterConfig &config = {}); |
192 | |
193 | } // namespace mlir |
194 | |
195 | #endif // MLIR_BYTECODE_BYTECODEWRITER_H |
196 | |