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
33namespace tvm {
34namespace runtime {
35namespace vm {
36
37/*! \brief A register name. */
38using RegName = int64_t;
39
40/*! \brief An alias for the integer type used ubiquitously
41 * in the VM.
42 */
43using 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 */
50enum 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 */
83struct 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