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/bytecode.h |
22 | * \brief The bytecode for Relay virtual machine. |
23 | */ |
24 | #ifndef TVM_RUNTIME_VM_BYTECODE_H_ |
25 | #define TVM_RUNTIME_VM_BYTECODE_H_ |
26 | |
27 | #include <tvm/runtime/data_type.h> |
28 | #include <tvm/runtime/logging.h> |
29 | |
30 | #include <iostream> |
31 | #include <vector> |
32 | |
33 | namespace tvm { |
34 | namespace runtime { |
35 | namespace vm { |
36 | |
37 | /*! \brief A register name. */ |
38 | using RegName = int64_t; |
39 | |
40 | /*! \brief An alias for the integer type used ubiquitously |
41 | * in the VM. |
42 | */ |
43 | using Index = int64_t; |
44 | |
45 | /*! \brief An enumeration of Relay's opcodes. |
46 | * |
47 | * The opcode is used to implement instruction |
48 | * as a tagged union. |
49 | */ |
50 | enum class Opcode { |
51 | Move = 0U, |
52 | Ret = 1U, |
53 | Invoke = 2U, |
54 | InvokeClosure = 3U, |
55 | InvokePacked = 4U, |
56 | AllocTensor = 5U, |
57 | AllocTensorReg = 6U, |
58 | AllocADT = 7U, |
59 | AllocClosure = 8U, |
60 | GetField = 9U, |
61 | If = 10U, |
62 | LoadConst = 11U, |
63 | Goto = 12U, |
64 | GetTag = 13U, |
65 | LoadConsti = 14U, |
66 | Fatal = 15U, |
67 | AllocStorage = 16U, |
68 | ShapeOf = 17U, |
69 | ReshapeTensor = 18U, |
70 | DeviceCopy = 19U, |
71 | KillRegister = 20U, |
72 | }; |
73 | |
74 | /*! \brief A single virtual machine instruction. |
75 | * |
76 | * The representation of the instruction is as |
77 | * a tagged union. |
78 | * |
79 | * The first field represents which instruction, |
80 | * and by extension which field of the union |
81 | * is active. |
82 | */ |
83 | struct Instruction { |
84 | /*! \brief The instruction opcode. */ |
85 | Opcode op; |
86 | |
87 | /*! \brief The destination register. */ |
88 | RegName dst; |
89 | |
90 | union { |
91 | struct /* AllocTensor Operands */ { |
92 | /*! \brief The storage to allocate from. */ |
93 | RegName storage; |
94 | /*! \brief The offset into the storage to allocate from. */ |
95 | Index offset; |
96 | /*! \brief The number of dimensions. */ |
97 | uint32_t ndim; |
98 | /*! \brief The shape of tensor. */ |
99 | int64_t* shape; |
100 | /*! \brief The datatype of tensor to be allocated. */ |
101 | DLDataType dtype; |
102 | } alloc_tensor; |
103 | struct /* AllocTensorReg Operands */ { |
104 | /*! \brief The storage to allocate from. */ |
105 | RegName storage; |
106 | /*! \brief The offset into the storage to allocate from. */ |
107 | Index offset; |
108 | /*! \brief The register to read the shape out of. */ |
109 | RegName shape_register; |
110 | /*! \brief The datatype of tensor to be allocated. */ |
111 | DLDataType dtype; |
112 | } alloc_tensor_reg; |
113 | struct /* InvokeClosure Operands */ { |
114 | /*! \brief The register containing the closure. */ |
115 | RegName closure; |
116 | /*! \brief The number of arguments to the closure. */ |
117 | Index num_closure_args; |
118 | /*! \brief The closure arguments as an array. */ |
119 | RegName* closure_args; |
120 | }; |
121 | struct /* Return Operands */ { |
122 | /*! \brief The register to return. */ |
123 | RegName result; |
124 | }; |
125 | struct /* Move Operands */ { |
126 | /*! \brief The source register for a move operation. */ |
127 | RegName from; |
128 | }; |
129 | struct /* InvokePacked Operands */ { |
130 | /*! \brief The index into the packed function table. */ |
131 | Index packed_index; |
132 | /*! \brief The arity of the packed function. */ |
133 | Index arity; |
134 | /*! \brief The number of outputs produced by the packed function. */ |
135 | Index output_size; |
136 | /*! \brief The arguments to pass to the packed function. */ |
137 | RegName* packed_args; |
138 | }; |
139 | struct /* If Operands */ { |
140 | /*! \brief The register containing the test value. */ |
141 | RegName test; |
142 | /*! \brief The register containing the target value. */ |
143 | RegName target; |
144 | /*! \brief The program counter offset for the true branch. */ |
145 | Index true_offset; |
146 | /*! \brief The program counter offset for the false branch. */ |
147 | Index false_offset; |
148 | } if_op; |
149 | struct /* Invoke Operands */ { |
150 | /*! \brief The function to call. */ |
151 | Index func_index; |
152 | /*! \brief The number of arguments to the function. */ |
153 | Index num_args; |
154 | /*! \brief The registers containing the arguments. */ |
155 | RegName* invoke_args_registers; |
156 | }; |
157 | struct /* LoadConst Operands */ { |
158 | /* \brief The index into the constant pool. */ |
159 | Index const_index; |
160 | }; |
161 | struct /* LoadConsti Operands */ { |
162 | /* \brief The index into the constant pool. */ |
163 | Index val; |
164 | } load_consti; |
165 | struct /* Jump Operands */ { |
166 | /*! \brief The jump offset. */ |
167 | Index pc_offset; |
168 | }; |
169 | struct /* Proj Operands */ { |
170 | /*! \brief The register to project from. */ |
171 | RegName object; |
172 | /*! \brief The field to read out. */ |
173 | Index field_index; |
174 | }; |
175 | struct /* GetTag Operands */ { |
176 | /*! \brief The register to project from. */ |
177 | RegName object; |
178 | } get_tag; |
179 | struct /* AllocADT Operands */ { |
180 | // TODO(mbs): Needs a DeviceAndScope. |
181 | /*! \brief The datatype's constructor tag. */ |
182 | Index constructor_tag; |
183 | /*! \brief The number of fields to store in the datatype. */ |
184 | Index num_fields; |
185 | /*! \brief The fields as an array. */ |
186 | RegName* datatype_fields; |
187 | }; |
188 | struct /* AllocClosure Operands */ { |
189 | // TODO(mbs): Needs a DeviceAndScope. |
190 | /*! \brief The index into the function table. */ |
191 | Index clo_index; |
192 | /*! \brief The number of free variables to capture. */ |
193 | Index num_freevar; |
194 | /*! \brief The free variables as an array. */ |
195 | RegName* free_vars; |
196 | }; |
197 | struct /* AllocStorage Operands */ { |
198 | /*! \brief The size of the allocation. */ |
199 | RegName allocation_size; |
200 | /*! \brief The alignment of the allocation. */ |
201 | Index alignment; |
202 | /*! \brief The hint of the dtype. */ |
203 | DLDataType dtype_hint; |
204 | /*! \brief The index of the device on which the allocation will be made. */ |
205 | Index device_index; |
206 | } alloc_storage; |
207 | struct /* ShapeOf Operands */ { |
208 | RegName tensor; |
209 | } shape_of; |
210 | struct /* ReshapeTensor Operands */ { |
211 | RegName tensor; |
212 | RegName newshape; |
213 | } reshape_tensor; |
214 | struct /* DeviceCopy Operands */ { |
215 | RegName src; |
216 | /*! \brief The index of the source device to copy from. */ |
217 | Index src_device_index; |
218 | /*! \brief The index of the destination deviceto copy to. */ |
219 | Index dst_device_index; |
220 | } device_copy; |
221 | }; |
222 | |
223 | /*! |
224 | * \brief Construct a return instruction. |
225 | * \param return_reg The register containing the return value. |
226 | * \return The return instruction. |
227 | */ |
228 | static Instruction Ret(RegName return_reg); |
229 | /*! |
230 | * \brief Construct a fatal instruction. |
231 | * \return The fatal instruction. |
232 | */ |
233 | static Instruction Fatal(); |
234 | /*! |
235 | * \brief Construct a invoke packed instruction. |
236 | * \param packed_index The index of the packed function. |
237 | * \param arity The arity of the function. |
238 | * \param output_size The number of outputs of the packed function. |
239 | * \param args The argument registers. |
240 | * \return The invoke packed instruction. |
241 | */ |
242 | static Instruction InvokePacked(Index packed_index, Index arity, Index output_size, |
243 | const std::vector<RegName>& args); |
244 | /*! |
245 | * \brief Construct an allocate tensor instruction with constant shape. |
246 | * \param storage The storage to allocate out of. |
247 | * \param offset The offset to allocate at. |
248 | * \param shape The shape of the tensor. |
249 | * \param dtype The dtype of the tensor. |
250 | * \param dst The destination register. |
251 | * \return The allocate tensor instruction. |
252 | */ |
253 | static Instruction AllocTensor(RegName storage, Index offset, const std::vector<int64_t>& shape, |
254 | DLDataType dtype, RegName dst); |
255 | /*! |
256 | * \brief Construct an allocate tensor instruction with register. |
257 | * \param storage The storage to allocate out of. |
258 | * \param offset The offset into the storage to allocate from. |
259 | * \param shape_register The register containing the shape. |
260 | * \param dtype The dtype of the tensor. |
261 | * \param dst The destination register. |
262 | * \return The allocate tensor instruction. |
263 | */ |
264 | static Instruction AllocTensorReg(RegName storage, Index offset, RegName shape_register, |
265 | DLDataType dtype, RegName dst); |
266 | /*! |
267 | * \brief Construct an allocate datatype instruction. |
268 | * \param tag The datatype tag. |
269 | * \param num_fields The number of fields for the datatype. |
270 | * \param fields The registers containing the fields. |
271 | * \param dst The register name of the destination. |
272 | * \return The allocate instruction tensor. |
273 | */ |
274 | static Instruction AllocADT(Index tag, Index num_fields, const std::vector<RegName>& fields, |
275 | RegName dst); |
276 | /*! |
277 | * \brief Construct an allocate closure instruction. |
278 | * \param func_index The index of the function table. |
279 | * \param num_freevar The number of free variables. |
280 | * \param free_vars The registers of the free variables. |
281 | * \param dst The destination register. |
282 | * \return The allocate closure instruction. |
283 | */ |
284 | static Instruction AllocClosure(Index func_index, Index num_freevar, |
285 | const std::vector<RegName>& free_vars, RegName dst); |
286 | /*! |
287 | * \brief Construct a get field instruction. |
288 | * \param object_reg The register containing the object to project from. |
289 | * \param field_index The field to read out of the object. |
290 | * \param dst The destination register. |
291 | * \return The get field instruction. |
292 | */ |
293 | static Instruction GetField(RegName object_reg, Index field_index, RegName dst); |
294 | /*! |
295 | * \brief Construct a get_tag instruction. |
296 | * \param object_reg The register containing the object to project from. |
297 | * \param dst The destination register. |
298 | * \return The get_tag instruction. |
299 | */ |
300 | static Instruction GetTag(RegName object_reg, RegName dst); |
301 | /*! |
302 | * \brief Construct an if instruction. |
303 | * \param test The register containing the test value. |
304 | * \param target The register containing the target value. |
305 | * \param true_branch The offset to the true branch. |
306 | * \param false_branch The offset to the false branch. |
307 | * \return The if instruction. |
308 | */ |
309 | static Instruction If(RegName test, RegName target, Index true_branch, Index false_branch); |
310 | /*! |
311 | * \brief Construct a goto instruction. |
312 | * \param pc_offset The offset from the current pc. |
313 | * \return The goto instruction. |
314 | */ |
315 | static Instruction Goto(Index pc_offset); |
316 | /*! |
317 | * \brief Construct an invoke instruction. |
318 | * \param func_index The index of the function to invoke. |
319 | * \param args The registers containing the arguments. |
320 | * \param dst The destination register. |
321 | * \return The invoke instruction. |
322 | */ |
323 | static Instruction Invoke(Index func_index, const std::vector<RegName>& args, RegName dst); |
324 | /*! |
325 | * \brief Construct an invoke closure instruction. |
326 | * \param closure The register of the closure to invoke. |
327 | * \param args The registers containing the arguments. |
328 | * \param dst The destination register. |
329 | * \return The invoke closure instruction. |
330 | */ |
331 | static Instruction InvokeClosure(RegName closure, const std::vector<RegName>& args, RegName dst); |
332 | /*! |
333 | * \brief Construct a load constant instruction. |
334 | * \param const_index The index of the constant. |
335 | * \param dst The destination register. |
336 | * \return The load constant instruction. |
337 | */ |
338 | static Instruction LoadConst(Index const_index, RegName dst); |
339 | /*! |
340 | * \brief Construct a load_constanti instruction. |
341 | * \param val The interger constant value. |
342 | * \param dst The destination register. |
343 | * \return The load_constanti instruction. |
344 | */ |
345 | static Instruction LoadConsti(Index val, RegName dst); |
346 | /*! |
347 | * \brief Construct a move instruction. |
348 | * \param src The source register. |
349 | * \param dst The destination register. |
350 | * \return The move instruction. |
351 | */ |
352 | static Instruction Move(RegName src, RegName dst); |
353 | /*! |
354 | * \brief Allocate a storage block. |
355 | * \param size The size of the allocation. |
356 | * \param alignment The allocation's alignment. |
357 | * \param dtype_hint The data type hint for the allocator. |
358 | * \param device_index The index of the device to allocate on. |
359 | * \param dst The destination to place the storage. |
360 | * \return The alloc storage instruction. |
361 | */ |
362 | static Instruction AllocStorage(RegName size, Index alignment, DLDataType dtype_hint, |
363 | Index device_index, RegName dst); |
364 | /*! |
365 | * \brief Get the shape of an input tensor. |
366 | * \param tensor The input tensor. |
367 | * \param dst The destination to store the shape of the given tensor. |
368 | * \return The shape of instruction. |
369 | */ |
370 | static Instruction ShapeOf(RegName tensor, RegName dst); |
371 | /*! |
372 | * \brief Reshape the tensor given the new shape. |
373 | * \param tensor The input tensor. |
374 | * \param newshape The shape tensor. |
375 | * \param dst The destination to store the output tensor with new shape. |
376 | * \return The reshape tensor instruction. |
377 | */ |
378 | static Instruction ReshapeTensor(RegName tensor, RegName newshape, RegName dst); |
379 | /*! |
380 | * \brief Copy tensor cross different devices. |
381 | * \param src The source register. |
382 | * \param src_device_index The index of the device holding the tensor in the source register. |
383 | * \param dst_device_index The index of the device to hold the tensor in the destination register. |
384 | * \param dst The destination register to store the copied tensor. |
385 | * \return The device copy instruction. |
386 | */ |
387 | static Instruction DeviceCopy(RegName src, Index src_device_index, Index dst_device_index, |
388 | RegName dst); |
389 | |
390 | static Instruction KillRegister(RegName dst); |
391 | |
392 | Instruction(); |
393 | Instruction(const Instruction& instr); |
394 | Instruction& operator=(const Instruction& instr); |
395 | ~Instruction(); |
396 | |
397 | friend std::ostream& operator<<(std::ostream& os, const Instruction&); |
398 | }; |
399 | |
400 | } // namespace vm |
401 | } // namespace runtime |
402 | } // namespace tvm |
403 | |
404 | #endif // TVM_RUNTIME_VM_BYTECODE_H_ |
405 | |