vm.cc 30.8 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

/*!
 *  Copyright (c) 2019 by Contributors
 * \file src/runtime/vm/vm.cc
 * \brief The Relay virtual machine.
 */

26
#include <dmlc/memory_io.h>
27 28 29
#include <tvm/logging.h>
#include <tvm/runtime/vm.h>

30
#include <algorithm>
31 32 33 34 35 36
#include <chrono>
#include <iostream>
#include <sstream>
#include <stdexcept>
#include <vector>

37 38
#include "memory_manager.h"
#include "naive_allocator.h"
39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62

using namespace tvm::runtime;

namespace tvm {
namespace runtime {
namespace vm {

Instruction::Instruction() {}

template <typename T>
static T* Duplicate(T* src, Index size) {
  auto dst = new T[size];
  std::copy(src, src + size, dst);
  return dst;
}

Instruction::Instruction(const Instruction& instr) {
  this->op = instr.op;
  this->dst = instr.dst;

  switch (instr.op) {
    case Opcode::Move:
      this->from = instr.from;
      return;
63
    case Opcode::Fatal:
64 65 66 67 68
      return;
    case Opcode::Ret:
      this->result = instr.result;
      return;
    case Opcode::AllocTensor:
69 70 71 72 73 74 75 76
      this->alloc_tensor.ndim = instr.alloc_tensor.ndim;
      this->alloc_tensor.shape = Duplicate<int64_t>(instr.alloc_tensor.shape,
                                                    instr.alloc_tensor.ndim);
      this->alloc_tensor.dtype = instr.alloc_tensor.dtype;
      return;
    case Opcode::AllocTensorReg:
      this->alloc_tensor_reg.shape_register = instr.alloc_tensor_reg.shape_register;
      this->alloc_tensor_reg.dtype = instr.alloc_tensor_reg.dtype;
77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95
      return;
    case Opcode::AllocDatatype:
      this->constructor_tag = instr.constructor_tag;
      this->num_fields = instr.num_fields;
      this->datatype_fields = Duplicate<RegName>(instr.datatype_fields, instr.num_fields);
      return;
    case Opcode::AllocClosure:
      this->clo_index = instr.clo_index;
      this->num_freevar = instr.num_freevar;
      this->free_vars = Duplicate<RegName>(instr.free_vars, instr.num_freevar);
      return;
    case Opcode::InvokePacked:
      this->packed_index = instr.packed_index;
      this->arity = instr.arity;
      this->output_size = instr.output_size;
      this->packed_args = Duplicate<RegName>(instr.packed_args, instr.arity);
      return;
    case Opcode::InvokeClosure:
      this->closure = instr.closure;
96 97
      this->num_closure_args = instr.num_closure_args;
      this->closure_args = Duplicate<RegName>(instr.closure_args, instr.num_closure_args);
98 99 100 101 102 103 104
      return;
    case Opcode::Invoke:
      this->func_index = instr.func_index;
      this->num_args = instr.num_args;
      this->invoke_args_registers = Duplicate<RegName>(instr.invoke_args_registers, instr.num_args);
      return;
    case Opcode::If:
105
      this->if_op = instr.if_op;
106 107 108 109
      return;
    case Opcode::LoadConst:
      this->const_index = instr.const_index;
      return;
110 111 112
    case Opcode::LoadConsti:
      this->load_consti = instr.load_consti;
      return;
113 114 115 116
    case Opcode::GetField:
      this->object = instr.object;
      this->field_index = instr.field_index;
      return;
117 118 119
    case Opcode::GetTag:
      this->get_tag = instr.get_tag;
      return;
120 121 122 123 124 125 126 127 128 129
    case Opcode::Goto:
      this->pc_offset = instr.pc_offset;
      return;
    default:
      std::ostringstream out;
      out << "Invalid instruction " << static_cast<int>(instr.op);
      throw std::runtime_error(out.str());
  }
}

130 131 132 133 134 135 136 137 138 139 140 141 142 143 144
template<typename T>
static inline void FreeIf(T* t) {
  if (t != nullptr) {
    delete t;
  }
}

Instruction& Instruction::operator=(const Instruction& instr) {
  this->op = instr.op;
  this->dst = instr.dst;

  switch (instr.op) {
    case Opcode::Move:
      this->from = instr.from;
      return *this;
145 146 147 148
    case Opcode::Fatal:
      return *this;
    case Opcode::LoadConsti:
      this->load_consti = instr.load_consti;
149 150 151 152 153
      return *this;
    case Opcode::Ret:
      this->result = instr.result;
      return *this;
    case Opcode::AllocTensor:
154 155 156 157 158 159 160 161
      this->alloc_tensor.ndim = instr.alloc_tensor.ndim;
      this->alloc_tensor.shape = Duplicate<int64_t>(instr.alloc_tensor.shape,
                                                    instr.alloc_tensor.ndim);
      this->alloc_tensor.dtype = instr.alloc_tensor.dtype;
      return *this;
    case Opcode::AllocTensorReg:
      this->alloc_tensor_reg.shape_register = instr.alloc_tensor_reg.shape_register;
      this->alloc_tensor_reg.dtype = instr.alloc_tensor_reg.dtype;
162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183
      return *this;
    case Opcode::AllocDatatype:
      this->constructor_tag = instr.constructor_tag;
      this->num_fields = instr.num_fields;
      FreeIf(this->datatype_fields);
      this->datatype_fields = Duplicate<RegName>(instr.datatype_fields, instr.num_fields);
      return *this;
    case Opcode::AllocClosure:
      this->clo_index = instr.clo_index;
      this->num_freevar = instr.num_freevar;
      FreeIf(this->free_vars);
      this->free_vars = Duplicate<RegName>(instr.free_vars, instr.num_freevar);
      return *this;
    case Opcode::InvokePacked:
      this->packed_index = instr.packed_index;
      this->arity = instr.arity;
      this->output_size = instr.output_size;
      FreeIf(this->packed_args);
      this->packed_args = Duplicate<RegName>(instr.packed_args, instr.arity);
      return *this;
    case Opcode::InvokeClosure:
      this->closure = instr.closure;
184
      this->num_closure_args = instr.num_closure_args;
185
      FreeIf(this->closure_args);
186
      this->closure_args = Duplicate<RegName>(instr.closure_args, instr.num_closure_args);
187 188 189 190 191 192 193 194
      return *this;
    case Opcode::Invoke:
      this->func_index = instr.func_index;
      this->num_args = instr.num_args;
      FreeIf(this->invoke_args_registers);
      this->invoke_args_registers = Duplicate<RegName>(instr.invoke_args_registers, instr.num_args);
      return *this;
    case Opcode::If:
195
      this->if_op = instr.if_op;
196 197 198 199 200 201 202 203
      return *this;
    case Opcode::LoadConst:
      this->const_index = instr.const_index;
      return *this;
    case Opcode::GetField:
      this->object = instr.object;
      this->field_index = instr.field_index;
      return *this;
204 205 206
    case Opcode::GetTag:
      this->get_tag = instr.get_tag;
      return *this;
207 208 209 210 211 212 213 214 215 216
    case Opcode::Goto:
      this->pc_offset = instr.pc_offset;
      return *this;
    default:
      std::ostringstream out;
      out << "Invalid instruction " << static_cast<int>(instr.op);
      throw std::runtime_error(out.str());
  }
}

217 218 219 220
Instruction::~Instruction() {
  switch (this->op) {
    case Opcode::Move:
    case Opcode::Ret:
221
    case Opcode::AllocTensorReg:
222 223 224
    case Opcode::If:
    case Opcode::LoadConst:
    case Opcode::GetField:
225
    case Opcode::GetTag:
226
    case Opcode::Goto:
227 228
    case Opcode::LoadConsti:
    case Opcode::Fatal:
229
      return;
230 231 232
    case Opcode::AllocTensor:
      delete this->alloc_tensor.shape;
      return;
233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249
    case Opcode::AllocDatatype:
      delete this->datatype_fields;
      return;
    case Opcode::AllocClosure:
      delete this->free_vars;
      return;
    case Opcode::InvokePacked:
      delete this->packed_args;
      return;
    case Opcode::InvokeClosure:
      delete this->closure_args;
      return;
    case Opcode::Invoke:
      delete this->invoke_args_registers;
      return;
    default:
      std::ostringstream out;
250
      LOG(FATAL) << "Invalid instruction " << static_cast<int>(this->op);
251 252 253 254 255 256 257 258 259 260
  }
}

Instruction Instruction::Ret(RegName result) {
  Instruction instr;
  instr.op = Opcode::Ret;
  instr.result = result;
  return instr;
}

261 262 263 264 265 266
Instruction Instruction::Fatal() {
  Instruction instr;
  instr.op = Opcode::Fatal;
  return instr;
}

267 268 269
Instruction Instruction::InvokePacked(Index packed_index,
                                      Index arity,
                                      Index output_size,
270 271 272 273 274 275 276 277 278 279 280 281 282
                                      const std::vector<RegName>& args) {
  Instruction instr;
  instr.op = Opcode::InvokePacked;
  instr.packed_index = packed_index;
  instr.arity = arity;
  instr.output_size = output_size;
  instr.packed_args = new RegName[arity];
  for (Index i = 0; i < arity; ++i) {
    instr.packed_args[i] = args[i];
  }
  return instr;
}

283
Instruction Instruction::AllocTensor(std::vector<int64_t> shape, DLDataType dtype, Index dst) {
284 285 286
  Instruction instr;
  instr.op = Opcode::AllocTensor;
  instr.dst = dst;
287 288 289 290 291 292 293 294 295 296 297 298 299 300 301
  instr.alloc_tensor.ndim = shape.size();
  instr.alloc_tensor.shape = new int64_t[shape.size()];
  for (size_t i = 0; i < shape.size(); ++i) {
    instr.alloc_tensor.shape[i] = shape[i];
  }
  instr.alloc_tensor.dtype = dtype;
  return instr;
}

Instruction Instruction::AllocTensorReg(RegName shape_register, DLDataType dtype, Index dst) {
  Instruction instr;
  instr.op = Opcode::AllocTensorReg;
  instr.dst = dst;
  instr.alloc_tensor_reg.shape_register = shape_register;
  instr.alloc_tensor_reg.dtype = dtype;
302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341
  return instr;
}

Instruction Instruction::AllocDatatype(Index tag, Index num_fields,
                                       const std::vector<RegName>& datatype_fields, Index dst) {
  Instruction instr;
  instr.op = Opcode::AllocDatatype;
  instr.dst = dst;
  instr.constructor_tag = tag;
  instr.num_fields = num_fields;
  instr.datatype_fields = new RegName[num_fields];
  for (Index i = 0; i < num_fields; ++i) {
    instr.datatype_fields[i] = datatype_fields[i];
  }
  return instr;
}

Instruction Instruction::AllocClosure(Index func_index, Index free_vars,
                                      const std::vector<RegName>& free_var_register, Index dst) {
  Instruction instr;
  instr.op = Opcode::AllocClosure;
  instr.dst = dst;
  instr.clo_index = func_index;
  instr.num_freevar = free_vars;
  instr.free_vars = new RegName[instr.num_freevar];
  for (Index i = 0; i < instr.num_freevar; ++i) {
    instr.free_vars[i] = free_var_register[i];
  }
  return instr;
}

Instruction Instruction::GetField(RegName object, Index field_index, RegName dst) {
  Instruction instr;
  instr.op = Opcode::GetField;
  instr.dst = dst;
  instr.object = object;
  instr.field_index = field_index;
  return instr;
}

342
Instruction Instruction::GetTag(RegName object, RegName dst) {
343
  Instruction instr;
344 345 346
  instr.op = Opcode::GetTag;
  instr.dst = dst;
  instr.get_tag.object = object;
347 348 349
  return instr;
}

350
Instruction Instruction::If(RegName test, RegName target, Index true_branch, Index false_branch) {
351
  Instruction instr;
352 353 354 355 356
  instr.op = Opcode::If;
  instr.if_op.test = test;
  instr.if_op.target = target;
  instr.if_op.true_offset = true_branch;
  instr.if_op.false_offset = false_branch;
357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386
  return instr;
}

Instruction Instruction::Goto(Index pc_offset) {
  Instruction instr;
  instr.op = Opcode::Goto;
  instr.pc_offset = pc_offset;
  return instr;
}

Instruction Instruction::Invoke(Index func_index, const std::vector<RegName>& args_registers,
                                RegName dst) {
  Instruction instr;
  instr.op = Opcode::Invoke;
  instr.dst = dst;
  instr.func_index = func_index;
  instr.num_args = args_registers.size();
  instr.invoke_args_registers = new RegName[instr.num_args];
  for (Index i = 0; i < instr.num_args; ++i) {
    instr.invoke_args_registers[i] = args_registers[i];
  }
  return instr;
}

Instruction Instruction::InvokeClosure(RegName closure, const std::vector<RegName>& args,
                                       RegName dst) {
  Instruction instr;
  instr.op = Opcode::InvokeClosure;
  instr.dst = dst;
  instr.closure = closure;
387
  instr.num_closure_args = args.size();
388 389 390 391 392 393 394 395 396 397 398 399 400 401 402
  instr.closure_args = new RegName[args.size()];
  for (size_t i = 0; i < args.size(); ++i) {
    instr.closure_args[i] = args[i];
  }
  return instr;
}

Instruction Instruction::LoadConst(Index const_index, RegName dst) {
  Instruction instr;
  instr.op = Opcode::LoadConst;
  instr.dst = dst;
  instr.const_index = const_index;
  return instr;
}

403
Instruction Instruction::LoadConsti(Index val, RegName dst) {
404 405 406 407 408 409 410
  Instruction instr;
  instr.op = Opcode::LoadConsti;
  instr.dst = dst;
  instr.load_consti.val = val;
  return instr;
}

411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431
Instruction Instruction::Move(RegName src, RegName dst) {
  Instruction instr;
  instr.op = Opcode::Move;
  instr.dst = dst;
  instr.from = src;
  return instr;
}

void DLDatatypePrint(std::ostream& os, const DLDataType& dtype) {
  switch (dtype.code) {
    case kDLInt:
      os << "int";
      break;
    case kDLUInt:
      os << "uint";
      break;
    case kDLFloat:
      os << "float";
      break;
  }

432 433 434
  os << int(dtype.bits);
  if (dtype.lanes != 1) {
    os << "x" << dtype.lanes;
435 436 437
  }
}

438
template<typename T>
439
std::string StrJoin(T* items, int offset, int cnt, std::string delim = ", ") {
440 441 442 443 444 445 446 447 448 449 450
  if (cnt == 0) {
    return "";
  }
  std::ostringstream oss;
  oss << items[offset];
  for (int i = 1; i < cnt; ++i) {
    oss << delim << items[offset + i];
  }
  return oss.str();
}

451 452 453
void InstructionPrint(std::ostream& os, const Instruction& instr) {
  switch (instr.op) {
    case Opcode::Move: {
454
      os << "move $" << instr.dst << " $" << instr.from;
455 456 457
      break;
    }
    case Opcode::Ret: {
458
      os << "ret $" << instr.result;
459 460
      break;
    }
461 462 463 464
    case Opcode::Fatal: {
      os << "fatal";
      break;
    }
465
    case Opcode::InvokePacked: {
466 467 468
      os << "invoke_packed PackedFunc[" << instr.packed_index << "] (in: $"
         << StrJoin<RegName>(instr.packed_args, 0,
                             instr.arity - instr.output_size, ", $")
469 470
         << ", out: $"
         << StrJoin<RegName>(instr.packed_args, instr.arity - instr.output_size,
471
                             instr.output_size, ", $")
472
         << ")";
473 474 475
      break;
    }
    case Opcode::AllocTensor: {
476
      os << "alloc_tensor $" << instr.dst << " ["
477 478
         << StrJoin<int64_t>(instr.alloc_tensor.shape, 0,
                             instr.alloc_tensor.ndim)
479 480 481 482 483 484 485 486
         << "] ";
      DLDatatypePrint(os, instr.alloc_tensor.dtype);
      break;
    }
    case Opcode::AllocTensorReg: {
      os << "alloc_tensor_reg $" << instr.dst << " $"
         << instr.alloc_tensor_reg.shape_register << " ";
      DLDatatypePrint(os, instr.alloc_tensor_reg.dtype);
487 488 489
      break;
    }
    case Opcode::AllocDatatype: {
490
      os << "alloc_data $" << instr.dst << " tag(" << instr.constructor_tag << ") [$"
491
         << StrJoin<RegName>(instr.datatype_fields, 0, instr.num_fields, ",$") << "]";
492 493 494
      break;
    }
    case Opcode::AllocClosure: {
495 496
      os << "alloc_closure $" << instr.dst << " VMFunc[" << instr.clo_index
         << "]($" << StrJoin<RegName>(instr.free_vars, 0, instr.num_freevar, ",$")
497
         << ")";
498 499 500
      break;
    }
    case Opcode::If: {
501
      os << "if " << "$" << instr.if_op.test << " " << instr.if_op.target << " "
502
         << instr.if_op.true_offset << " " << instr.if_op.false_offset;
503 504 505
      break;
    }
    case Opcode::Invoke: {
506 507
      os << "invoke $" << instr.dst << " VMFunc[" << instr.func_index << "]($"
         << StrJoin<RegName>(instr.invoke_args_registers, 0, instr.num_args, ",$")
508
         << ")";
509 510 511
      break;
    }
    case Opcode::InvokeClosure: {
512
      os << "invoke_closure $" << instr.dst << " $" << instr.closure << "($"
513
         << StrJoin<RegName>(instr.closure_args, 0, instr.num_closure_args, ",$")
514
         << ")";
515 516 517
      break;
    }
    case Opcode::LoadConst: {
518
      os << "load_const $" << instr.dst << " Const[" << instr.const_index << "]";
519 520
      break;
    }
521
    case Opcode::LoadConsti: {
522
      os << "load_consti $" << instr.dst << " Const[" << instr.load_consti.val << "]";
523 524
      break;
    }
525
    case Opcode::GetField: {
526
      os << "get_field $" << instr.dst << " $" << instr.object << "["
527
         << instr.field_index << "]";
528 529
      break;
    }
530
    case Opcode::GetTag: {
531
      os << "get_tag $" << instr.dst << " $" << instr.get_tag.object;
532 533
      break;
    }
534
    case Opcode::Goto: {
535
      os << "goto " << instr.pc_offset;
536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551
      break;
    }
    default:
      LOG(FATAL) << "should never hit this case" << static_cast<int>(instr.op);
      break;
  }
}

std::ostream& operator<<(std::ostream& os, const Instruction& instr) {
  InstructionPrint(os, instr);
  return os;
}

void VMFunctionPrint(std::ostream& os, const VMFunction& vm_func) {
  os << vm_func.name << ": " << std::endl;
  for (size_t i = 0; i < vm_func.instructions.size(); ++i) {
552
    os << i << ": " << vm_func.instructions[i] << ";" << std::endl;
553 554 555 556 557 558 559 560
  }
}

std::ostream& operator<<(std::ostream& os, const VMFunction& vm_func) {
  VMFunctionPrint(os, vm_func);
  return os;
}

561 562 563 564 565 566 567 568 569 570 571 572 573 574
Object CopyTo(Object src, const DLContext& ctx) {
  if (src->tag == ObjectTag::kTensor) {
    auto tensor = ToNDArray(src);
    if (tensor->ctx.device_type != ctx.device_type) {
      auto copy = tensor.CopyTo(ctx);
      return Object::Tensor(copy);
    } else {
      return src;
    }
  } else {
    return src;
  }
}

575 576 577 578 579
PackedFunc VirtualMachine::GetFunction(const std::string& name,
                                       const std::shared_ptr<ModuleNode>& sptr_to_self) {
  if (name == "invoke") {
    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
      std::string func_name = args[0];
580
      auto ctx = this->GetParamsContext();
581 582
      std::vector<Object> func_args;
      for (int i = 1; i < args.size(); ++i) {
583
        Object obj = CopyTo(args[i], ctx);
584 585
        func_args.push_back(obj);
      }
586 587 588 589
      auto it = std::find_if(functions.begin(), functions.end(),
                             [func_name](const VMFunction& func) {
                               return func.name == func_name;
                             });
590

591 592 593 594 595 596 597 598 599 600 601 602 603
      CHECK(it != functions.end()) << "Cannot find function " << func_name << "\n";
      CHECK_EQ(func_args.size() + params_.size(), it->params.size())
          << "The number of provided parameters doesn't match the number of arguments"
          << "\n";
      if (!params_.empty()) {
        for (const auto& p : it->params) {
          const auto& pit = params_.find(p);
          if (pit != params_.end()) {
            func_args.push_back(pit->second);
          }
        }
        CHECK_EQ(func_args.size(), it->params.size());
      }
604 605 606 607 608 609 610 611 612 613 614 615 616 617 618
      *rv = this->Invoke(func_name, func_args);
    });
  } else if (name == "init") {
    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
      CHECK_EQ(args.size() % 2, 0);
      std::vector<TVMContext> contexts;
      for (int i = 0; i < args.size() / 2; ++i) {
        TVMContext ctx;
        int device_type = args[i * 2];
        ctx.device_type = DLDeviceType(device_type);
        ctx.device_id = args[i * 2 + 1];
        contexts.push_back(ctx);
      }
      this->Init(contexts);
    });
619 620 621 622
  } else if (name == "load_params") {
    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
      this->LoadParams(args[0].operator std::string());
    });
623 624 625 626 627 628
  } else {
    LOG(FATAL) << "Unknown packed function: " << name;
    return PackedFunc([sptr_to_self, name](TVMArgs args, TVMRetValue* rv) {});
  }
}

629 630 631 632 633 634 635 636 637 638 639 640
TVMContext VirtualMachine::GetParamsContext() const {
  // Use the fallback device if no device index is available.
  int fallback_device_type = static_cast<int>(ctxs[0].device_type);
  // TODO(wweic): For heterogeneous execution, get device information from byte

  const auto& cit =
    std::find_if(ctxs.begin(), ctxs.end(), [&fallback_device_type](const TVMContext& c) {
      return fallback_device_type == static_cast<int>(c.device_type);
    });
  return (cit == ctxs.end() ? ctxs[0] : *cit);
}

641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656
void VirtualMachine::LoadParams(const std::string& params) {
  dmlc::MemoryStringStream mss(const_cast<std::string*>(&params));
  dmlc::Stream* strm = &mss;
  uint64_t header, reserved;
  CHECK(strm->Read(&header)) << "Invalid parameter file";
  CHECK(header == kTVMNDArrayListMagic) << "Invalid parameter file";
  CHECK(strm->Read(&reserved)) << "Invalid parameter file";

  std::vector<std::string> names;
  CHECK(strm->Read(&names)) << "Invalid parameter file";

  uint64_t sz;
  strm->Read(&sz);
  size_t size = static_cast<size_t>(sz);
  CHECK(size == names.size()) << "Invalid parameter file";

657
  auto ctx = GetParamsContext();
658 659 660 661
  for (size_t i = 0; i < size; i++) {
    NDArray arr;
    CHECK(arr.Load(strm)) << "Invalid parameter file";
    runtime::Object obj = runtime::Object::Tensor(arr);
662 663
    auto copy = CopyTo(obj, ctx);
    params_.emplace(std::make_pair(names[i], copy));
664 665 666
  }
}

667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683
void VirtualMachine::PushFrame(Index arg_count, Index ret_pc, const VMFunction& vm_func) {
  auto frame = VMFrame(ret_pc, func_index, arg_count, code, vm_func.register_file_size);
  frames.push_back(frame);
}

Index VirtualMachine::PopFrame() {
  CHECK_GT(frames.size(), 0);
  const VMFrame& fr = frames.back();
  func_index = fr.func_index;
  code = fr.code;
  pc = fr.pc;
  auto call_stack_size = frames.size();
  frames.pop_back();
  return call_stack_size;
}

void VirtualMachine::InvokeGlobal(const VMFunction& func, const std::vector<Object>& args) {
684
  DLOG(INFO) << "Invoking global " << func.name << " " << args.size();
685

686
  PushFrame(func.params.size(), this->pc + 1, func);
687 688 689
  for (size_t i = 0; i < args.size(); ++i) {
    WriteRegister(i, args[i]);
  }
690
  DLOG(INFO) << "func.params= " << func.params.size();
691 692 693 694 695 696

  code = func.instructions.data();
  pc = 0;
}

Object VirtualMachine::Invoke(const VMFunction& func, const std::vector<Object>& args) {
697
  DLOG(INFO) << "Executing Function: " << std::endl << func;
698 699

  InvokeGlobal(func, args);
700
  RunLoop();
701
  auto alloc = MemoryManager::Global()->GetAllocator(ctxs[0]);
702
  DLOG(INFO) << "Memory used: " << alloc->UsedMemory() << " B";
703 704 705 706
  return return_register;
}

Object VirtualMachine::Invoke(const std::string& name, const std::vector<Object>& args) {
707
  auto func_index = this->global_map[name];
708
  DLOG(INFO) << "Invoke Global " << name << " at index " << func_index;
709 710 711
  return Invoke(this->functions[func_index], args);
}

712 713 714
void VirtualMachine::InvokePacked(Index packed_index, const PackedFunc& func,
                                  Index arg_count, Index output_size,
                                  const std::vector<Object>& args) {
715 716 717 718 719 720 721 722
  size_t arity = 0;
  for (Index i = 0; i < arg_count; i++) {
    if (args[i].ptr_->tag == ObjectTag::kDatatype) {
      arity += args[i].AsDatatype()->fields.size();
    } else {
      ++arity;
    }
  }
723

724 725 726 727
  std::vector<TVMValue> values(arity);
  std::vector<int> codes(arity);
  runtime::TVMArgsSetter setter(values.data(), codes.data());
  int idx = 0;
728
  for (Index i = 0; i < arg_count; i++) {
729 730 731 732 733 734 735 736 737 738
    if (args[i].ptr_->tag == ObjectTag::kDatatype) {
      auto dt_cell = args[i].AsDatatype();
      for (auto obj : dt_cell->fields) {
        NDArray data = ToNDArray(obj);
        setter(idx++, data);
      }
    } else {
      NDArray data = ToNDArray(args[i]);
      setter(idx++, data);
    }
739 740 741
  }

  TVMRetValue rv;
742
  func.CallPacked(TVMArgs(values.data(), codes.data(), arity), &rv);
743 744
}

745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760
void VirtualMachine::Init(const std::vector<TVMContext>& ctxs) {
  this->ctxs = ctxs;

  // Get the list of packed functions.
  CHECK(primitive_map.empty() || lib.operator->())
      << "runtime module should have been built for primitive functions"
      << "\n";
  for (const auto& it : primitive_map) {
    const auto& packed_name = it.first;
    auto packed_index = static_cast<size_t>(it.second);
    if (packed_funcs.size() <= packed_index) {
      packed_funcs.resize(packed_index + 1);
    }
    packed_funcs[packed_index] = lib.GetFunction(packed_name);
  }
}
761 762 763 764 765 766 767 768 769

inline void VirtualMachine::WriteRegister(Index r, const Object& val) {
  frames.back().register_file[r] = val;
}

inline Object VirtualMachine::ReadRegister(Index r) const {
  return frames.back().register_file[r];
}

770 771 772 773 774 775 776 777 778 779 780 781 782 783 784
inline int32_t VirtualMachine::LoadScalarInt(Index r) const {
  int32_t result;
  const auto& obj = ReadRegister(r);
  NDArray array = ToNDArray(obj).CopyTo({kDLCPU, 0});

  if (array->dtype.bits <= 8) {
    result = reinterpret_cast<int8_t*>(array->data)[0];
  } else if (array->dtype.bits <= 16) {
    result = reinterpret_cast<int16_t*>(array->data)[0];
  } else {
    result = reinterpret_cast<int32_t*>(array->data)[0];
  }
  return result;
}

785
void VirtualMachine::RunLoop() {
786 787 788 789 790 791
  CHECK(this->code);
  this->pc = 0;
  Index frame_start = frames.size();
  while (true) {
  main_loop:
    auto const& instr = this->code[this->pc];
792
    DLOG(INFO) << "Executing(" << pc << "): " << instr;
793 794 795 796 797 798 799
#if USE_RELAY_DEBUG
    InstructionPrint(std::cout, instr);
#endif  // USE_RELAY_DEBUG

    switch (instr.op) {
      case Opcode::Move: {
        Object from_obj;
800
        from_obj = ReadRegister(instr.from);
801 802 803 804
        WriteRegister(instr.dst, from_obj);
        pc++;
        goto main_loop;
      }
805 806 807
      case Opcode::Fatal: {
        throw std::runtime_error("VM encountered fatal error");
      }
808
      case Opcode::LoadConst: {
809 810 811
        auto constant_obj = this->constants[instr.const_index];
        auto device_obj = CopyTo(constant_obj, ctxs[0]);
        WriteRegister(instr.dst, device_obj);
812 813 814
        pc++;
        goto main_loop;
      }
815
      case Opcode::LoadConsti: {
816 817
        auto tensor = NDArray::Empty({1}, {kDLInt, 64, 1}, {kDLCPU, 0});
        reinterpret_cast<int64_t*>(tensor->data)[0] = instr.load_consti.val;
818 819 820 821
        WriteRegister(instr.dst, Object::Tensor(tensor));
        pc++;
        goto main_loop;
      }
822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837
      case Opcode::Invoke: {
        std::vector<Object> args;
        for (Index i = 0; i < instr.num_args; ++i) {
          args.push_back(ReadRegister(instr.invoke_args_registers[i]));
        }
        InvokeGlobal(this->functions[instr.func_index], args);
        frames.back().caller_return_register = instr.dst;
        goto main_loop;
      }
      case Opcode::InvokePacked: {
        const auto& func = packed_funcs[instr.packed_index];
        const auto& arity = instr.arity;
        std::vector<Object> args;
        for (Index i = 0; i < arity; ++i) {
          args.push_back(ReadRegister(instr.packed_args[i]));
        }
838
        InvokePacked(instr.packed_index, func, arity, instr.output_size, args);
839 840 841 842 843 844 845 846 847 848 849 850 851 852
        for (Index i = 0; i < instr.output_size; ++i) {
          WriteRegister(instr.packed_args[instr.arity - instr.output_size + i],
                        args[instr.arity - instr.output_size + i]);
        }
        pc++;
        goto main_loop;
      }
      case Opcode::InvokeClosure: {
        auto object = ReadRegister(instr.closure);
        const auto& closure = object.AsClosure();
        std::vector<Object> args;
        for (auto free_var : closure->free_vars) {
          args.push_back(free_var);
        }
853
        for (Index i = 0; i < instr.num_closure_args; ++i) {
854 855
          args.push_back(ReadRegister(instr.closure_args[i]));
        }
856 857 858 859 860 861 862 863 864 865 866 867 868 869 870
        InvokeGlobal(this->functions[closure->func_index], args);
        frames.back().caller_return_register = instr.dst;
        goto main_loop;
      }
      case Opcode::GetField: {
        auto object = ReadRegister(instr.object);
        CHECK(object->tag == ObjectTag::kDatatype)
            << "Object is not data type object, register " << instr.object << ", Object tag "
            << static_cast<int>(object->tag);
        const auto& tuple = object.AsDatatype();
        auto field = tuple->fields[instr.field_index];
        WriteRegister(instr.dst, field);
        pc++;
        goto main_loop;
      }
871 872 873 874 875 876 877 878 879 880 881 882 883 884
      case Opcode::GetTag: {
        auto object = ReadRegister(instr.get_tag.object);
        CHECK(object->tag == ObjectTag::kDatatype)
            << "Object is not data type object, register "
            << instr.get_tag.object << ", Object tag "
            << static_cast<int>(object->tag);
        const auto& data = object.AsDatatype();
        auto tag = data->tag;
        auto tag_tensor = NDArray::Empty({1}, {kDLInt, 32, 1}, {kDLCPU, 0});
        reinterpret_cast<int32_t*>(tag_tensor->data)[0] = tag;
        WriteRegister(instr.dst, Object::Tensor(tag_tensor));
        pc++;
        goto main_loop;
      }
885 886 887 888 889
      case Opcode::Goto: {
        pc += instr.pc_offset;
        goto main_loop;
      }
      case Opcode::If: {
890 891
        int32_t test_val = LoadScalarInt(instr.if_op.test);
        int32_t target_val = LoadScalarInt(instr.if_op.target);
892

893 894 895
        if (test_val == target_val) {
          CHECK_NE(instr.if_op.true_offset, 0);
          pc += instr.if_op.true_offset;
896
        } else {
897 898
          CHECK_NE(instr.if_op.false_offset, 0);
          pc += instr.if_op.false_offset;
899 900 901 902 903
        }

        goto main_loop;
      }
      case Opcode::AllocTensor: {
904
        auto shape = std::vector<int64_t>(instr.alloc_tensor.ndim);
Li committed
905
        for (uint32_t i = 0; i < instr.alloc_tensor.ndim; ++i) {
906 907 908 909 910 911 912 913 914 915
          shape[i] = instr.alloc_tensor.shape[i];
        }
        auto allocator = MemoryManager::Global()->GetAllocator(ctxs[0]);
        auto data = allocator->Empty(shape, instr.alloc_tensor.dtype, ctxs[0]);
        auto obj = Object::Tensor(data);
        WriteRegister(instr.dst, obj);
        pc++;
        goto main_loop;
      }
      case Opcode::AllocTensorReg: {
916 917 918 919
        DLContext cpu_ctx;
        cpu_ctx.device_type = kDLCPU;
        cpu_ctx.device_id = 0;

920
        auto shape_tensor_obj = ReadRegister(instr.alloc_tensor_reg.shape_register);
921 922 923 924 925 926 927
        NDArray shape_tensor = ToNDArray(shape_tensor_obj).CopyTo(cpu_ctx);

        int64_t* dims = static_cast<int64_t*>(shape_tensor->data);
        auto num_dims = shape_tensor->shape[0];
        auto shape = std::vector<int64_t>(shape_tensor->shape[0]);
        shape.assign(dims, dims + num_dims);
        auto allocator = MemoryManager::Global()->GetAllocator(ctxs[0]);
928
        auto data = allocator->Empty(shape, instr.alloc_tensor_reg.dtype, ctxs[0]);
929 930 931 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946 947 948 949 950 951 952 953 954 955 956 957 958 959 960 961 962 963 964 965 966 967 968 969 970 971 972 973 974
        auto obj = Object::Tensor(data);
        WriteRegister(instr.dst, obj);
        pc++;
        goto main_loop;
      }
      case Opcode::AllocDatatype: {
        std::vector<Object> fields;
        for (Index i = 0; i < instr.num_fields; ++i) {
          fields.push_back(ReadRegister(instr.datatype_fields[i]));
        }
        Object obj = Object::Datatype(instr.constructor_tag, fields);
        WriteRegister(instr.dst, obj);
        pc++;
        goto main_loop;
      }
      case Opcode::AllocClosure: {
        std::vector<Object> free_vars;
        for (Index i = 0; i < instr.num_freevar; i++) {
          free_vars.push_back(ReadRegister(instr.free_vars[i]));
        }
        WriteRegister(instr.dst, Object::Closure(instr.func_index, free_vars));
        pc++;
        goto main_loop;
      }
      case Opcode::Ret: {
        // If we have hit the point from which we started
        // running, we should return to the caller breaking
        // the dispatch loop.
        return_register = ReadRegister(instr.result);
        auto caller_return_register = frames.back().caller_return_register;

        if (PopFrame() == frame_start) {
          return;
          // Otherwise we are just returning from a local call.
        } else {
          WriteRegister(caller_return_register, return_register);
          goto main_loop;
        }
      }
    }
  }
}

}  // namespace vm
}  // namespace runtime
}  // namespace tvm