1/*
2 * Licensed to the Apache Software Foundation (ASF) under one
3 * or more contributor license agreements. See the NOTICE file
4 * distributed with this work for additional information
5 * regarding copyright ownership. The ASF licenses this file
6 * to you under the Apache License, Version 2.0 (the
7 * "License"); you may not use this file except in compliance
8 * with the License. You may obtain a copy of the License at
9 *
10 * http://www.apache.org/licenses/LICENSE-2.0
11 *
12 * Unless required by applicable law or agreed to in writing,
13 * software distributed under the License is distributed on an
14 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15 * KIND, either express or implied. See the License for the
16 * specific language governing permissions and limitations
17 * under the License.
18 */
19
20/*!
21 * \file tvm/runtime/vm/executable.h
22 * \brief The Relay virtual machine executable.
23 */
24#ifndef TVM_RUNTIME_VM_EXECUTABLE_H_
25#define TVM_RUNTIME_VM_EXECUTABLE_H_
26
27#include <tvm/runtime/container/map.h>
28#include <tvm/runtime/container/string.h>
29#include <tvm/runtime/module.h>
30#include <tvm/runtime/object.h>
31#include <tvm/runtime/packed_func.h>
32#include <tvm/runtime/vm/bytecode.h>
33
34#include <map>
35#include <string>
36#include <unordered_map>
37#include <vector>
38
39namespace tvm {
40namespace runtime {
41namespace vm {
42
43struct VMFunction;
44
45/*!
46 * \brief The executable emitted by the VM compiler.
47 *
48 * The executable contains information (e.g. data in different memory regions)
49 * to run in a virtual machine.
50 *
51 * - Global section, containing all globals.
52 * - Constant section, storing the constant pool.
53 * - Primitive name section, containing the function name of the primitive ops
54 * used by the virtual machine.
55 * - Code section, handling the VM functions and bytecode.
56 */
57class TVM_DLL Executable : public ModuleNode {
58 public:
59 /*!
60 * \brief Get a PackedFunc from an executable module.
61 *
62 * \param name the name of the function.
63 * \param sptr_to_self The shared_ptr that points to this module node.
64 *
65 * \return PackedFunc or nullptr when it is not available.
66 */
67 PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) final;
68
69 /*!
70 * \brief Write the Executable to the binary stream in serialized form.
71 *
72 * Late-bound constants (if any) must have already been saved by \p
73 * MoveLateBoundConstantsToBinary.
74 *
75 * \param stream The binary stream to save the executable to.
76 */
77 void SaveToBinary(dmlc::Stream* stream) final;
78
79 /*!
80 * \brief Write the Executable to the provided path as a file containing its serialized content.
81 *
82 * Late-bound constants (if any) must have already been saved by \p
83 * MoveLateBoundConstantsToBinary.
84 *
85 * \param path The path to write the serialized data to.
86 * \param format The format of the serialized blob.
87 */
88 void SaveToFile(const std::string& path, const std::string& format) final;
89
90 /*!
91 * \brief Serialize the executable into global section, constant section, and
92 * code section. This object must outlive the returned byte array.
93 *
94 * Late-bound constants (if any) must have already been saved by \p
95 * MoveLateBoundConstantsToBinary.
96 *
97 * \return The binary representation of the VM.
98 */
99 TVMByteArray Save();
100
101 /*!
102 * \brief Load the saved VM executable.
103 *
104 * Late-bound constants (if any) must then be loaded by \p LoadLateBoundConstantsFromBinary.
105 *
106 * \param code The bytecode in string.
107 * \param lib The compiled runtime library.
108 *
109 * \return exe The constructed executable.
110 */
111 static runtime::Module Load(const std::string& code, const runtime::Module lib);
112
113 /*!
114 * \brief Returns the late-bound constants for the executable (if any) as a byte-stream.
115 * Leaves the executable's late-bound constants map empty. Only constants who's byte
116 * tensor size is greater than or equal to \p byte_limit are marked as late-bound. \p byte_limit
117 * may be zero.
118 *
119 * Must be called before \p SaveToBinary and friends if late-bound constants are
120 * desired. Otherwise can be ignore.
121 */
122 void MoveLateBoundConstantsToStream(dmlc::Stream* stream, size_t byte_limit);
123
124 /*!
125 * \brief As for \p MoveLateBoundConstantsToStream, but save to file at \p path.
126 */
127 void MoveLateBoundConstantsToFile(const std::string& path, size_t byte_limit);
128
129 /*!
130 * \brief Get a map of all constants with larger that byte_limit in size.
131 */
132 Map<String, NDArray> GetLateBoundConstants(size_t byte_limit);
133
134 /*!
135 * \brief Restores the late-bound constants for the executable (if any) from given byte-stream.
136 *
137 * Must be called after \p Load but before any other methods if \p MoveLateBoundConstantsToBinary
138 * was used when saving. Otherwise can be ignored.
139 */
140 void LoadLateBoundConstantsFromStream(dmlc::Stream* stream);
141
142 /*!
143 * \brief Restores the late-bound constants for the executable (if any) from given map.
144 *
145 * Must be called after \p Load but before any other methods if \p MoveLateBoundConstantsToBinary
146 * was used when saving. Otherwise can be ignored.
147 */
148 void LoadLateBoundConstantsFromMap(Map<String, NDArray> map);
149
150 /*!
151 * \brief As for \p LoadLateBoundConstantsFromStream, but load from file at \p path.
152 */
153 void LoadLateBoundConstantsFromFile(const std::string& path);
154
155 /*!
156 * \brief Get the serialized form of the `functions`. This is
157 * essentially bytecode serialization.
158 *
159 * \return The serialized vm bytecode.
160 *
161 * \note The bytecode is in the following format:
162 * func_name reg_file_size num_instructions
163 * param1 param2 ... paramM
164 * instruction1
165 * instruction2
166 * ...
167 * instructionN
168 *
169 * Each instruction is printed in the following format:
170 * opcode num_fields field1 ... fieldX # The text format.
171 *
172 * Serializing an `Instruction` requires us to deal with the bytecode. Each line
173 * of the instructions could be serialized as the following format:
174 * hash, opcode, f1, f2, ..., fX, field with variable length
175 * 1. hash: the hash of the instruction. This number will be used to help us
176 * validate if an instruction is well-formed during deserialization.
177 * 2. opcode: the opcode code of the instruction.
178 * 3. f1, f2, ..., fX. These fields together represent the fixed fields in
179 * an instruction, e.g., `from` and `dst` fields of a `Move` instruction. For
180 * example, `DLDataType` will be unpacked into three fields (code, bits, lanes).
181 * 4. The rest of the line indicates the field with variable length, e.g.,
182 * the shape of a tensor, the args used by an `InvokPacked` instruction, etc.
183 *
184 * The field starting from # is only used for debugging. The serialized code
185 * doesn't contain it, therefore the deserializer doens't need to handle it.
186 */
187 std::string GetBytecode() const;
188
189 /*!
190 * \brief Returns a description of all the constants in the executable in human-readable
191 * format. Intended for debugging and diff-testing.
192 */
193 std::string GetConstants() const;
194
195 /*!
196 * \brief Returns a description of all the (virtual) devices in the executable in human-readable
197 * format. Intended for debugging and diff-testing.
198 */
199 std::string GetVirtualDevices() const;
200
201 /*!
202 * \brief Returns a description of all the 'primitive' (ie PackedFuncs) in the executable in
203 * human-readable format. These correspond either to PrimFuncs we've compiled locally, or
204 * functions compiled by a BYOC external codegen. Intended for debugging and diff-testing.
205 */
206 std::string GetPrimitives() const;
207
208 /*!
209 * \brief Print the detailed statistics of the given code, i.e. number of
210 * globls and constants, etc.
211 */
212 std::string Stats() const;
213
214 /*!
215 * \brief Get the `lib` module in an executable. Users have the flexibility to call
216 * `export_library` from the frontend to save the library to disk.
217 *
218 * \return The runtime module that contains the hardware dependent code.
219 */
220 runtime::Module GetLib() const;
221
222 /*!
223 * \brief Set the `lib` module in an executable.
224 *
225 * This allows us to do partial initialization in the case of (de|ser)ialization cases.
226 * This method also ensures correct initialization of library ensuring we only Import a
227 * single library.
228 *
229 * NB: This also provides some abstraction over how libraries are stored as there are plans
230 * to iterate on the way runtime::Module works in the backend of the compiler.
231 */
232 void SetLib(const runtime::Module& lib);
233
234 /*!
235 * \brief Get VMFunction.
236 * \param func_name The function's name.
237 * \return VMFunction.
238 */
239 const VMFunction& GetVMFunctionWithName(const std::string& func_name) const;
240
241 /*!
242 * \brief Get the arity of the VMFunction.
243 * \param func Function name.
244 * \return The number of parameters.
245 */
246 int GetFunctionArity(std::string func) const;
247
248 /*!
249 * \brief Get the parameter name given the function name and parameter index.
250 * \param func Function name.
251 * \param index Parameter index.
252 * \return The parameter name.
253 */
254 std::string GetFunctionParameterName(std::string func, uint32_t index) const;
255
256 virtual ~Executable() {}
257
258 const char* type_key() const final { return "VMExecutable"; }
259
260 /*!
261 * \brief The (compile-time, virtual) devices corresponding to each device index.
262 * Currently we only support at most one device per device type.
263 */
264 std::vector<Device> virtual_devices;
265 /*!
266 * \brief The device index corresponding to the 'host' device. That will hold and evaluate
267 * shape-related data and code.
268 */
269 int host_device_index = -1;
270 /*!
271 * \brief The global constant array.
272 *
273 * LoadConst instructions indexes are w.r.t. this vector. Late-bound constants are removed
274 * from this table after saving late-bound constants.
275 */
276 std::vector<ObjectRef> constants;
277 /*!
278 * \brief For each constant index the name of the late-bound constant, or null if constant is
279 * immediate. Only populated after loading executable but before loading late-bound constants.
280 */
281 std::vector<String> late_bound_constant_names;
282
283 /*! \brief A map from globals (as strings) to their index in the Relay function map. */
284 std::unordered_map<std::string, Index> global_map;
285 /*! \brief A mapping from the packed function's global name (as string) to the index that
286 * corresponds to the position of the `packed_funcs` list in a `VirtualMachine` object.
287 */
288 std::unordered_map<std::string, Index> primitive_map;
289 /*! \brief The structural hashes of the operators in this function. */
290 std::map<Index, Map<String, ObjectRef>> op_attrs;
291 /*! \brief The virtual machine's function table. */
292 std::vector<VMFunction> functions;
293 /*! \brief The index of the device holding each constant. */
294 std::vector<Index> const_device_indexes;
295
296 private:
297 /*!
298 * \brief Save the virtual devices
299 *
300 * /param strm The output stream.
301 */
302 void SaveVirtualDevicesSection(dmlc::Stream* strm);
303
304 /*!
305 * \brief Save the globals.
306 *
307 * \param strm The output stream.
308 */
309 void SaveGlobalSection(dmlc::Stream* strm);
310
311 /*!
312 * \brief Save the constant pool.
313 *
314 * \param stream The output stream.
315 */
316 void SaveConstantSection(dmlc::Stream* stream);
317
318 /*!
319 * \brief Load the constant pool.
320 *
321 * \param stream The input stream.
322 */
323 void LoadConstantSection(dmlc::Stream* stream);
324
325 /*!
326 * \brief Save primitive op names.
327 *
328 * \param strm The output stream.
329 */
330 void SavePrimitiveOpNames(dmlc::Stream* strm);
331
332 /*!
333 * \brief Save the vm functions.
334 *
335 * \param strm The output stream.
336 */
337 void SaveCodeSection(dmlc::Stream* strm);
338
339 /*!
340 * \brief Load the virtual devices
341 *
342 * /param strm The input stream.
343 */
344 void LoadVirtualDevicesSection(dmlc::Stream* strm);
345
346 /*!
347 * \brief Load the globals.
348 *
349 * \param strm The input stream.
350 */
351 void LoadGlobalSection(dmlc::Stream* strm);
352
353 /*!
354 * \brief Load primitive op names.
355 *
356 * \param strm The input stream.
357 */
358 void LoadPrimitiveOpNames(dmlc::Stream* strm);
359
360 /*!
361 * \brief Load the vm functions.
362 *
363 * \param strm The input stream.
364 */
365 void LoadCodeSection(dmlc::Stream* strm);
366
367 /*! \brief The serialized bytecode. */
368 std::string code_;
369};
370
371} // namespace vm
372} // namespace runtime
373} // namespace tvm
374
375#endif // TVM_RUNTIME_VM_EXECUTABLE_H_
376