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 src/runtime/vm/bytecode.cc
22 * \brief The bytecode for Relay virtual machine.
23 */
24
25#include <tvm/runtime/logging.h>
26#include <tvm/runtime/ndarray.h>
27#include <tvm/runtime/vm/bytecode.h>
28
29#include <sstream>
30
31namespace tvm {
32namespace runtime {
33namespace vm {
34
35Instruction::Instruction() {}
36
37template <typename T>
38static T* Duplicate(T* src, Index size) {
39 auto dst = new T[size];
40 std::copy(src, src + size, dst);
41 return dst;
42}
43
44Instruction::Instruction(const Instruction& instr) {
45 this->op = instr.op;
46 this->dst = instr.dst;
47
48 switch (instr.op) {
49 case Opcode::Move:
50 this->from = instr.from;
51 return;
52 case Opcode::Fatal:
53 return;
54 case Opcode::Ret:
55 this->result = instr.result;
56 return;
57 case Opcode::AllocTensor:
58 this->alloc_tensor.storage = instr.alloc_tensor.storage;
59 this->alloc_tensor.offset = instr.alloc_tensor.offset;
60 this->alloc_tensor.ndim = instr.alloc_tensor.ndim;
61 this->alloc_tensor.shape =
62 Duplicate<int64_t>(instr.alloc_tensor.shape, instr.alloc_tensor.ndim);
63 this->alloc_tensor.dtype = instr.alloc_tensor.dtype;
64 return;
65 case Opcode::AllocTensorReg:
66 this->alloc_tensor_reg.storage = instr.alloc_tensor_reg.storage;
67 this->alloc_tensor_reg.offset = instr.alloc_tensor_reg.offset;
68 this->alloc_tensor_reg.shape_register = instr.alloc_tensor_reg.shape_register;
69 this->alloc_tensor_reg.dtype = instr.alloc_tensor_reg.dtype;
70 return;
71 case Opcode::AllocADT:
72 this->constructor_tag = instr.constructor_tag;
73 this->num_fields = instr.num_fields;
74 this->datatype_fields = Duplicate<RegName>(instr.datatype_fields, instr.num_fields);
75 return;
76 case Opcode::AllocClosure:
77 this->clo_index = instr.clo_index;
78 this->num_freevar = instr.num_freevar;
79 this->free_vars = Duplicate<RegName>(instr.free_vars, instr.num_freevar);
80 return;
81 case Opcode::InvokePacked:
82 this->packed_index = instr.packed_index;
83 this->arity = instr.arity;
84 this->output_size = instr.output_size;
85 this->packed_args = Duplicate<RegName>(instr.packed_args, instr.arity);
86 return;
87 case Opcode::InvokeClosure:
88 this->closure = instr.closure;
89 this->num_closure_args = instr.num_closure_args;
90 this->closure_args = Duplicate<RegName>(instr.closure_args, instr.num_closure_args);
91 return;
92 case Opcode::Invoke:
93 this->func_index = instr.func_index;
94 this->num_args = instr.num_args;
95 this->invoke_args_registers = Duplicate<RegName>(instr.invoke_args_registers, instr.num_args);
96 return;
97 case Opcode::If:
98 this->if_op = instr.if_op;
99 return;
100 case Opcode::LoadConst:
101 this->const_index = instr.const_index;
102 return;
103 case Opcode::LoadConsti:
104 this->load_consti = instr.load_consti;
105 return;
106 case Opcode::GetField:
107 this->object = instr.object;
108 this->field_index = instr.field_index;
109 return;
110 case Opcode::GetTag:
111 this->get_tag = instr.get_tag;
112 return;
113 case Opcode::Goto:
114 this->pc_offset = instr.pc_offset;
115 return;
116 case Opcode::AllocStorage:
117 this->alloc_storage = instr.alloc_storage;
118 return;
119 case Opcode::ShapeOf:
120 this->shape_of.tensor = instr.shape_of.tensor;
121 return;
122 case Opcode::ReshapeTensor:
123 this->reshape_tensor = instr.reshape_tensor;
124 return;
125 case Opcode::DeviceCopy:
126 this->device_copy = instr.device_copy;
127 return;
128 case Opcode::KillRegister:
129 return;
130 default:
131 std::ostringstream out;
132 out << "Invalid instruction " << static_cast<int>(instr.op);
133 throw std::runtime_error(out.str());
134 }
135}
136
137template <typename T>
138static inline void FreeIf(T* t) {
139 if (t != nullptr) {
140 delete t;
141 }
142}
143
144Instruction& Instruction::operator=(const Instruction& instr) {
145 this->op = instr.op;
146 this->dst = instr.dst;
147
148 switch (instr.op) {
149 case Opcode::Move:
150 this->from = instr.from;
151 return *this;
152 case Opcode::Fatal:
153 return *this;
154 case Opcode::LoadConsti:
155 this->load_consti = instr.load_consti;
156 return *this;
157 case Opcode::Ret:
158 this->result = instr.result;
159 return *this;
160 case Opcode::AllocTensor:
161 this->alloc_tensor.storage = this->alloc_tensor.storage;
162 this->alloc_tensor.offset = instr.alloc_tensor.offset;
163 this->alloc_tensor.ndim = instr.alloc_tensor.ndim;
164 this->alloc_tensor.shape =
165 Duplicate<int64_t>(instr.alloc_tensor.shape, instr.alloc_tensor.ndim);
166 this->alloc_tensor.dtype = instr.alloc_tensor.dtype;
167 return *this;
168 case Opcode::AllocTensorReg:
169 this->alloc_tensor_reg.storage = instr.alloc_tensor_reg.storage;
170 this->alloc_tensor_reg.offset = instr.alloc_tensor_reg.offset;
171 this->alloc_tensor_reg.shape_register = instr.alloc_tensor_reg.shape_register;
172 this->alloc_tensor_reg.dtype = instr.alloc_tensor_reg.dtype;
173 return *this;
174 case Opcode::AllocADT:
175 this->constructor_tag = instr.constructor_tag;
176 this->num_fields = instr.num_fields;
177 FreeIf(this->datatype_fields);
178 this->datatype_fields = Duplicate<RegName>(instr.datatype_fields, instr.num_fields);
179 return *this;
180 case Opcode::AllocClosure:
181 this->clo_index = instr.clo_index;
182 this->num_freevar = instr.num_freevar;
183 FreeIf(this->free_vars);
184 this->free_vars = Duplicate<RegName>(instr.free_vars, instr.num_freevar);
185 return *this;
186 case Opcode::InvokePacked:
187 this->packed_index = instr.packed_index;
188 this->arity = instr.arity;
189 this->output_size = instr.output_size;
190 FreeIf(this->packed_args);
191 this->packed_args = Duplicate<RegName>(instr.packed_args, instr.arity);
192 return *this;
193 case Opcode::InvokeClosure:
194 this->closure = instr.closure;
195 this->num_closure_args = instr.num_closure_args;
196 FreeIf(this->closure_args);
197 this->closure_args = Duplicate<RegName>(instr.closure_args, instr.num_closure_args);
198 return *this;
199 case Opcode::Invoke:
200 this->func_index = instr.func_index;
201 this->num_args = instr.num_args;
202 FreeIf(this->invoke_args_registers);
203 this->invoke_args_registers = Duplicate<RegName>(instr.invoke_args_registers, instr.num_args);
204 return *this;
205 case Opcode::If:
206 this->if_op = instr.if_op;
207 return *this;
208 case Opcode::LoadConst:
209 this->const_index = instr.const_index;
210 return *this;
211 case Opcode::GetField:
212 this->object = instr.object;
213 this->field_index = instr.field_index;
214 return *this;
215 case Opcode::GetTag:
216 this->get_tag = instr.get_tag;
217 return *this;
218 case Opcode::Goto:
219 this->pc_offset = instr.pc_offset;
220 return *this;
221 case Opcode::AllocStorage:
222 this->alloc_storage = instr.alloc_storage;
223 return *this;
224 case Opcode::ShapeOf:
225 this->shape_of.tensor = instr.shape_of.tensor;
226 return *this;
227 case Opcode::ReshapeTensor:
228 this->reshape_tensor = instr.reshape_tensor;
229 return *this;
230 case Opcode::DeviceCopy:
231 this->device_copy = instr.device_copy;
232 return *this;
233 case Opcode::KillRegister:
234 return *this;
235 default:
236 std::ostringstream out;
237 out << "Invalid instruction " << static_cast<int>(instr.op);
238 throw std::runtime_error(out.str());
239 }
240}
241
242Instruction::~Instruction() {
243 switch (this->op) {
244 case Opcode::Move:
245 case Opcode::Ret:
246 case Opcode::AllocTensorReg:
247 case Opcode::If:
248 case Opcode::LoadConst:
249 case Opcode::GetField:
250 case Opcode::GetTag:
251 case Opcode::Goto:
252 case Opcode::LoadConsti:
253 case Opcode::AllocStorage:
254 case Opcode::ShapeOf:
255 case Opcode::ReshapeTensor:
256 case Opcode::DeviceCopy:
257 case Opcode::Fatal:
258 case Opcode::KillRegister:
259 return;
260 case Opcode::AllocTensor:
261 delete[] this->alloc_tensor.shape;
262 return;
263 case Opcode::AllocADT:
264 delete[] this->datatype_fields;
265 return;
266 case Opcode::AllocClosure:
267 delete[] this->free_vars;
268 return;
269 case Opcode::InvokePacked:
270 delete[] this->packed_args;
271 return;
272 case Opcode::InvokeClosure:
273 delete[] this->closure_args;
274 return;
275 case Opcode::Invoke:
276 delete[] this->invoke_args_registers;
277 return;
278 default:
279 std::ostringstream out;
280 LOG(FATAL) << "Invalid instruction " << static_cast<int>(this->op);
281 }
282}
283
284Instruction Instruction::Ret(RegName result) {
285 Instruction instr;
286 instr.op = Opcode::Ret;
287 instr.result = result;
288 return instr;
289}
290
291Instruction Instruction::Fatal() {
292 Instruction instr;
293 instr.op = Opcode::Fatal;
294 return instr;
295}
296
297Instruction Instruction::InvokePacked(Index packed_index, Index arity, Index output_size,
298 const std::vector<RegName>& args) {
299 Instruction instr;
300 instr.op = Opcode::InvokePacked;
301 instr.packed_index = packed_index;
302 instr.arity = arity;
303 instr.output_size = output_size;
304 instr.packed_args = new RegName[arity];
305 for (Index i = 0; i < arity; ++i) {
306 instr.packed_args[i] = args[i];
307 }
308 return instr;
309}
310
311Instruction Instruction::AllocTensor(RegName storage, RegName offset,
312 const std::vector<int64_t>& shape, DLDataType dtype,
313 RegName dst) {
314 Instruction instr;
315 instr.op = Opcode::AllocTensor;
316 instr.dst = dst;
317 instr.alloc_tensor.storage = storage;
318 instr.alloc_tensor.offset = offset;
319 instr.alloc_tensor.ndim = shape.size();
320 instr.alloc_tensor.shape = new int64_t[shape.size()];
321 for (size_t i = 0; i < shape.size(); ++i) {
322 instr.alloc_tensor.shape[i] = shape[i];
323 }
324 instr.alloc_tensor.dtype = dtype;
325 return instr;
326}
327
328Instruction Instruction::AllocTensorReg(RegName storage, RegName offset, RegName shape_register,
329 DLDataType dtype, RegName dst) {
330 Instruction instr;
331 instr.op = Opcode::AllocTensorReg;
332 instr.dst = dst;
333 instr.alloc_tensor_reg.storage = storage;
334 instr.alloc_tensor_reg.offset = offset;
335 instr.alloc_tensor_reg.shape_register = shape_register;
336 instr.alloc_tensor_reg.dtype = dtype;
337 return instr;
338}
339
340Instruction Instruction::AllocStorage(RegName size, Index alignment, DLDataType dtype_hint,
341 Index device_index, RegName dst) {
342 Instruction instr;
343 instr.op = Opcode::AllocStorage;
344 instr.dst = dst;
345 instr.alloc_storage.allocation_size = size;
346 instr.alloc_storage.alignment = alignment;
347 instr.alloc_storage.dtype_hint = dtype_hint;
348 instr.alloc_storage.device_index = device_index;
349 return instr;
350}
351
352Instruction Instruction::ShapeOf(RegName tensor, RegName dst) {
353 Instruction instr;
354 instr.op = Opcode::ShapeOf;
355 instr.dst = dst;
356 instr.shape_of.tensor = tensor;
357 return instr;
358}
359
360Instruction Instruction::ReshapeTensor(RegName tensor, RegName newshape, RegName dst) {
361 Instruction instr;
362 instr.op = Opcode::ReshapeTensor;
363 instr.dst = dst;
364 instr.reshape_tensor.tensor = tensor;
365 instr.reshape_tensor.newshape = newshape;
366 return instr;
367}
368
369Instruction Instruction::DeviceCopy(RegName src, Index src_device_index, Index dst_device_index,
370 RegName dst) {
371 Instruction instr;
372 instr.op = Opcode::DeviceCopy;
373 instr.dst = dst;
374 instr.device_copy.src = src;
375 instr.device_copy.src_device_index = src_device_index;
376 instr.device_copy.dst_device_index = dst_device_index;
377 return instr;
378}
379
380Instruction Instruction::KillRegister(RegName dst) {
381 Instruction instr;
382 instr.op = Opcode::KillRegister;
383 instr.dst = dst;
384 return instr;
385}
386
387Instruction Instruction::AllocADT(Index tag, Index num_fields,
388 const std::vector<RegName>& datatype_fields, RegName dst) {
389 Instruction instr;
390 instr.op = Opcode::AllocADT;
391 instr.dst = dst;
392 instr.constructor_tag = tag;
393 instr.num_fields = num_fields;
394 instr.datatype_fields = new RegName[num_fields];
395 for (Index i = 0; i < num_fields; ++i) {
396 instr.datatype_fields[i] = datatype_fields[i];
397 }
398 return instr;
399}
400
401Instruction Instruction::AllocClosure(Index func_index, Index free_vars,
402 const std::vector<RegName>& free_var_register, RegName dst) {
403 Instruction instr;
404 instr.op = Opcode::AllocClosure;
405 instr.dst = dst;
406 instr.clo_index = func_index;
407 instr.num_freevar = free_vars;
408 instr.free_vars = new RegName[instr.num_freevar];
409 for (Index i = 0; i < instr.num_freevar; ++i) {
410 instr.free_vars[i] = free_var_register[i];
411 }
412 return instr;
413}
414
415Instruction Instruction::GetField(RegName object, Index field_index, RegName dst) {
416 Instruction instr;
417 instr.op = Opcode::GetField;
418 instr.dst = dst;
419 instr.object = object;
420 instr.field_index = field_index;
421 return instr;
422}
423
424Instruction Instruction::GetTag(RegName object, RegName dst) {
425 Instruction instr;
426 instr.op = Opcode::GetTag;
427 instr.dst = dst;
428 instr.get_tag.object = object;
429 return instr;
430}
431
432Instruction Instruction::If(RegName test, RegName target, Index true_branch, Index false_branch) {
433 Instruction instr;
434 instr.op = Opcode::If;
435 instr.if_op.test = test;
436 instr.if_op.target = target;
437 instr.if_op.true_offset = true_branch;
438 instr.if_op.false_offset = false_branch;
439 return instr;
440}
441
442Instruction Instruction::Goto(Index pc_offset) {
443 Instruction instr;
444 instr.op = Opcode::Goto;
445 instr.pc_offset = pc_offset;
446 return instr;
447}
448
449Instruction Instruction::Invoke(Index func_index, const std::vector<RegName>& args_registers,
450 RegName dst) {
451 Instruction instr;
452 instr.op = Opcode::Invoke;
453 instr.dst = dst;
454 instr.func_index = func_index;
455 instr.num_args = args_registers.size();
456 instr.invoke_args_registers = new RegName[instr.num_args];
457 for (Index i = 0; i < instr.num_args; ++i) {
458 instr.invoke_args_registers[i] = args_registers[i];
459 }
460 return instr;
461}
462
463Instruction Instruction::InvokeClosure(RegName closure, const std::vector<RegName>& args,
464 RegName dst) {
465 Instruction instr;
466 instr.op = Opcode::InvokeClosure;
467 instr.dst = dst;
468 instr.closure = closure;
469 instr.num_closure_args = args.size();
470 instr.closure_args = new RegName[args.size()];
471 for (size_t i = 0; i < args.size(); ++i) {
472 instr.closure_args[i] = args[i];
473 }
474 return instr;
475}
476
477Instruction Instruction::LoadConst(Index const_index, RegName dst) {
478 Instruction instr;
479 instr.op = Opcode::LoadConst;
480 instr.dst = dst;
481 instr.const_index = const_index;
482 return instr;
483}
484
485Instruction Instruction::LoadConsti(Index val, RegName dst) {
486 Instruction instr;
487 instr.op = Opcode::LoadConsti;
488 instr.dst = dst;
489 instr.load_consti.val = val;
490 return instr;
491}
492
493Instruction Instruction::Move(RegName src, RegName dst) {
494 Instruction instr;
495 instr.op = Opcode::Move;
496 instr.dst = dst;
497 instr.from = src;
498 return instr;
499}
500
501void DLDatatypePrint(std::ostream& os, const DLDataType& dtype) {
502 switch (dtype.code) {
503 case kDLInt:
504 os << "int";
505 break;
506 case kDLUInt:
507 os << "uint";
508 break;
509 case kDLFloat:
510 os << "float";
511 break;
512 case kDLBfloat:
513 os << "bfloat";
514 break;
515 }
516
517 os << int(dtype.bits);
518 if (dtype.lanes != 1) {
519 os << "x" << dtype.lanes;
520 }
521}
522
523template <typename T>
524std::string StrJoin(T* items, int offset, int cnt, std::string delim = ", ") {
525 if (cnt == 0) {
526 return "";
527 }
528 std::ostringstream oss;
529 oss << items[offset];
530 for (int i = 1; i < cnt; ++i) {
531 oss << delim << items[offset + i];
532 }
533 return oss.str();
534}
535
536void InstructionPrint(std::ostream& os, const Instruction& instr) {
537 switch (instr.op) {
538 case Opcode::Move: {
539 os << "move $" << instr.dst << " $" << instr.from;
540 break;
541 }
542 case Opcode::Ret: {
543 os << "ret $" << instr.result;
544 break;
545 }
546 case Opcode::Fatal: {
547 os << "fatal";
548 break;
549 }
550 case Opcode::InvokePacked: {
551 os << "invoke_packed PackedFunc[" << instr.packed_index << "] (in: $"
552 << StrJoin<RegName>(instr.packed_args, 0, instr.arity - instr.output_size, ", $")
553 << ", out: $"
554 << StrJoin<RegName>(instr.packed_args, instr.arity - instr.output_size, instr.output_size,
555 ", $")
556 << ")";
557 break;
558 }
559 case Opcode::AllocTensor: {
560 os << "alloc_tensor $" << instr.dst << " $" << instr.alloc_tensor.storage << " $"
561 << instr.alloc_tensor.offset << " ["
562 << StrJoin<int64_t>(instr.alloc_tensor.shape, 0, instr.alloc_tensor.ndim) << "] ";
563 DLDatatypePrint(os, instr.alloc_tensor.dtype);
564 break;
565 }
566 case Opcode::AllocTensorReg: {
567 os << "alloc_tensor_reg $" << instr.dst << " $" << instr.alloc_tensor_reg.storage << " $"
568 << instr.alloc_tensor_reg.offset << " $" << instr.alloc_tensor_reg.shape_register << " ";
569 DLDatatypePrint(os, instr.alloc_tensor_reg.dtype);
570 break;
571 }
572 case Opcode::AllocADT: {
573 os << "alloc_data $" << instr.dst << " tag(" << instr.constructor_tag << ") [$"
574 << StrJoin<RegName>(instr.datatype_fields, 0, instr.num_fields, ",$") << "]";
575 break;
576 }
577 case Opcode::AllocClosure: {
578 os << "alloc_closure $" << instr.dst << " VMFunc[" << instr.clo_index << "]($"
579 << StrJoin<RegName>(instr.free_vars, 0, instr.num_freevar, ",$") << ")";
580 break;
581 }
582 case Opcode::If: {
583 os << "if "
584 << "$" << instr.if_op.test << " $" << instr.if_op.target << " " << instr.if_op.true_offset
585 << " " << instr.if_op.false_offset;
586 break;
587 }
588 case Opcode::Invoke: {
589 os << "invoke $" << instr.dst << " VMFunc[" << instr.func_index << "]($"
590 << StrJoin<RegName>(instr.invoke_args_registers, 0, instr.num_args, ",$") << ")";
591 break;
592 }
593 case Opcode::InvokeClosure: {
594 os << "invoke_closure $" << instr.dst << " $" << instr.closure << "($"
595 << StrJoin<RegName>(instr.closure_args, 0, instr.num_closure_args, ",$") << ")";
596 break;
597 }
598 case Opcode::LoadConst: {
599 os << "load_const $" << instr.dst << " Const[" << instr.const_index << "]";
600 break;
601 }
602 case Opcode::LoadConsti: {
603 os << "load_consti $" << instr.dst << " " << instr.load_consti.val;
604 break;
605 }
606 case Opcode::GetField: {
607 os << "get_field $" << instr.dst << " $" << instr.object << "[" << instr.field_index << "]";
608 break;
609 }
610 case Opcode::GetTag: {
611 os << "get_tag $" << instr.dst << " $" << instr.get_tag.object;
612 break;
613 }
614 case Opcode::Goto: {
615 os << "goto " << instr.pc_offset;
616 break;
617 }
618 case Opcode::AllocStorage: {
619 os << "alloc_storage $" << instr.dst << " $" << instr.alloc_storage.allocation_size << " "
620 << instr.alloc_storage.alignment << " "
621 << DLDataType2String(instr.alloc_storage.dtype_hint) << " "
622 << instr.alloc_storage.device_index;
623 break;
624 }
625 case Opcode::ShapeOf: {
626 os << "shape_of $" << instr.dst << " $" << instr.shape_of.tensor;
627 break;
628 }
629 case Opcode::ReshapeTensor: {
630 os << "reshape_tensor $" << instr.dst << " $" << instr.reshape_tensor.tensor << " $"
631 << instr.reshape_tensor.newshape;
632 break;
633 }
634 case Opcode::DeviceCopy: {
635 os << "device_copy $" << instr.dst << " $" << instr.device_copy.src << " "
636 << instr.device_copy.dst_device_index << " " << instr.device_copy.src_device_index;
637 break;
638 }
639 case Opcode::KillRegister: {
640 os << "kill_register $" << instr.dst;
641 break;
642 }
643 default:
644 LOG(FATAL) << "should never hit this case" << static_cast<int>(instr.op);
645 break;
646 }
647}
648
649std::ostream& operator<<(std::ostream& os, const Instruction& instr) {
650 InstructionPrint(os, instr);
651 return os;
652}
653
654} // namespace vm
655} // namespace runtime
656} // namespace tvm
657