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
18namespace mlir {
19class DialectBytecodeWriter;
20class DialectVersion;
21class Operation;
22
23/// A class to interact with the attributes and types printer when emitting MLIR
24/// bytecode.
25template <class T>
26class AttrTypeBytecodeWriter {
27public:
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.
74class BytecodeWriterConfig {
75public:
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
178private:
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.
190LogicalResult writeBytecodeToFile(Operation *op, raw_ostream &os,
191 const BytecodeWriterConfig &config = {});
192
193} // namespace mlir
194
195#endif // MLIR_BYTECODE_BYTECODEWRITER_H
196

source code of mlir/include/mlir/Bytecode/BytecodeWriter.h