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 | |
39 | namespace tvm { |
40 | namespace runtime { |
41 | namespace vm { |
42 | |
43 | struct 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 | */ |
57 | class 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 | |