text: /*
   * 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.
   */
  
  /*!
   * \file tvm/runtime/vm/executable.cc
   * \brief The implementation of a virtual machine executable APIs.
   */
  
  #include <dmlc/memory_io.h>
  #include <tvm/runtime/c_runtime_api.h>
  #include <tvm/runtime/debug.h>
  #include <tvm/runtime/registry.h>
  #include <tvm/runtime/vm/executable.h>
  #include <tvm/runtime/vm/vm.h>
  
  #include <algorithm>
  #include <iomanip>
  #include <iostream>
  #include <memory>
  #include <random>
  #include <sstream>
  #include <utility>
  #include <vector>
  
  #include "../file_utils.h"
  #include "../library_module.h"
  #include "serialize_utils.h"
  
  namespace tvm {
  namespace runtime {
  namespace vm {
  
  #define STREAM_CHECK(val, section)                                          \
    ICHECK(val) << "Invalid VM file format in the " << section << " section." \
                << "\n";
  
  // Helper to serialize a vm instruction.
  VMInstructionSerializer SerializeInstruction(const Instruction& instr);
  // Helper to deserialize a serialized vm instruction.
  Instruction DeserializeInstruction(const VMInstructionSerializer& instr);
  
  PackedFunc Executable::GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) {
    if (name == "get_lib") {
      return PackedFunc(
          [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = this->GetLib(); });
    } else if (name == "get_bytecode") {
      return PackedFunc(
          [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = this->GetBytecode(); });
    } else if (name == "get_constants") {
      return PackedFunc([this](TVMArgs args, TVMRetValue* rv) { *rv = this->GetConstants(); });
    } else if (name == "get_virtual_devices") {
      return PackedFunc([this](TVMArgs args, TVMRetValue* rv) { *rv = this->GetVirtualDevices(); });
    } else if (name == "get_primitives") {
      return PackedFunc([this](TVMArgs args, TVMRetValue* rv) { *rv = this->GetPrimitives(); });
    } else if (name == "get_stats") {
      return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = this->Stats(); });
    } else if (name == "save") {
      return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = this->Save(); });
    } else if (name == "get_function_arity") {
      return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
        std::string func_name = args[0];
        *rv = this->GetFunctionArity(func_name);
      });
    } else if (name == "get_function_param_name") {
      return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
        std::string func_name = args[0];
        int index = args[1];
        *rv = this->GetFunctionParameterName(func_name, index);
      });
    } else if (name == "vm_load_executable") {
      return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
        auto vm = make_object<VirtualMachine>();
        ICHECK(sptr_to_self.get() == this);
        vm->LoadExecutable(GetObjectPtr<Executable>(this));
        *rv = Module(vm);
      });
    } else if (name == "move_late_bound_consts") {
      return PackedFunc([this](TVMArgs args, TVMRetValue* rv) {
        CHECK_EQ(args.size(), 2);
        std::string path = args[0];
        uint64_t byte_limit = args[1];
        MoveLateBoundConstantsToFile(path, static_cast<size_t>(byte_limit));
      });
    } else if (name == "load_late_bound_consts") {
      return PackedFunc([this](TVMArgs args, TVMRetValue* rv) {
        CHECK_EQ(args.size(), 1);
        std::string path = args[0];
        LoadLateBoundConstantsFromFile(path);
      });
    } else {
      LOG(FATAL) << "Unknown packed function: " << name;
      return PackedFunc(nullptr);
    }
  }
  
  int Executable::GetFunctionArity(std::string func_name) const {
    auto it = global_map.find(func_name);
    if (it == global_map.end()) {
      LOG(ERROR) << "Cannot find function " << func_name << " in executable";
      return -1;
    }
    const auto& func = functions[it->second];
    return func.params.size();
  }
  
  std::string Executable::GetFunctionParameterName(std::string func_name, uint32_t index) const {
    auto it = global_map.find(func_name);
    if (it == global_map.end()) {
      LOG(ERROR) << "Cannot find function " << func_name << " in executable";
      return "";
    }
    const auto& func = functions[it->second];
    if (index > func.params.size()) {
      LOG(ERROR) << "Invalid parameter index";
      return "";
    }
    return func.params[index];
  }
  
  std::string Executable::GetBytecode() const {
    std::ostringstream oss;
  
    for (size_t i = 0; i < functions.size(); ++i) {
      const auto& func = functions[i];
      // Print the header of the function format.
      oss << "VM Function[" << i << "]: " << func.name << "(";
      bool first = true;
      for (const auto& param : func.params) {
        if (!first) {
          oss << ", ";
        }
        oss << param;
        first = false;
      }
      oss << ")" << std::endl;
      oss << "# reg file size = " << func.register_file_size << std::endl;
      oss << "# instruction count = " << func.instructions.size() << std::endl;
  
      // Print the instructions of a `VMFunction`.
      // The part after ";" is the instruction in text format.
      oss << "opcode, fields # inst(text):" << std::endl;
      for (size_t idx = 0; idx < func.instructions.size(); ++idx) {
        const auto& instr = func.instructions[idx];
        const auto& serialized_instr = SerializeInstruction(instr);
        std::ostringstream line;
        line << std::setw(2) << idx << ": " << serialized_instr.opcode << " ";
        for (auto it : serialized_instr.fields) {
          line << it << " ";
        }
        oss << std::setw(40) << std::setfill(' ') << std::left << line.str();
        oss << "  # " << instr;
        if (oss.str().back() != '\n') oss << std::endl;
      }
      oss << std::endl;
    }
  
    return oss.str();
  }
  
  std::string Executable::GetConstants() const {
    std::ostringstream oss;
    for (size_t i = 0; i < constants.size(); ++i) {
      const auto& constant = constants[i];
      auto ndarray = Downcast<NDArray>(constant);
      oss << "VM Const[" << i
          << "]: " << RuntimeObject2String(ndarray, virtual_devices[host_device_index])
          << " on device index " << const_device_indexes[i] << std::endl;
    }
    return oss.str();
  }
  
  std::string Executable::GetVirtualDevices() const {
    std::ostringstream oss;
    for (size_t i = 0; i < virtual_devices.size(); ++i) {
      const auto& device = virtual_devices[i];
      oss << "VM VirtualDevice[" << i << "]: device type " << device.device_type << " and id "
          << device.device_id << std::endl;
    }
    return oss.str();
  }
  
  std::string Executable::GetPrimitives() const {
    std::ostringstream os;
    std::vector<std::pair<int, std::string>> entries;
    entries.reserve(primitive_map.size());
    for (const auto& kv : primitive_map) {
      entries.emplace_back(kv.second, kv.first);
    }
    std::sort(entries.begin(), entries.end(),
              [](const std::pair<int, std::string>& left, const std::pair<int, std::string>& right) {
                return left.first < right.first;
              });
    for (const auto& entry : entries) {
      os << "VM PackedFunc[" << entry.first << "]: " << entry.second << std::endl;
    }
    return os.str();
  }
  
  std::string Executable::Stats() const {
    std::ostringstream oss;
    oss << "Relay VM executable statistics:" << std::endl;
  
    // Get the number of constants and the shape of each of them.
    oss << "  Constant shapes (# " << constants.size() << "): [";
    for (const auto& it : constants) {
      const auto constant = Downcast<NDArray>(it);
      const auto& shape = constant.Shape();
  
      // Scalar
      if (shape.empty()) {
        oss << "scalar, ";
        continue;
      }
  
      oss << "[";
      for (auto s : shape) {
        oss << s << ", ";
      }
      oss.seekp(-2, oss.cur);
      oss << "], " << std::endl;
    }
    if (!constants.empty()) oss.seekp(-2, oss.cur);
    oss << "]" << std::endl;
  
    // Get the number of globals and the name of each of them.
    oss << "  Globals (#" << global_map.size() << "): [";
    for (const auto& it : global_map) {
      oss << "(\"" << it.first << "\", " << it.second << ")"
          << ", ";
    }
    if (!global_map.empty()) oss.seekp(-2, oss.cur);
    oss << "]" << std::endl;
  
    // Get the number of primitive ops and the name of each of them.
    oss << "  Primitive ops (#" << primitive_map.size() << "): [";
    std::vector<std::string> prim_ops;
    for (const auto& it : primitive_map) {
      auto packed_index = static_cast<size_t>(it.second);
      if (prim_ops.size() <= packed_index) {
        prim_ops.resize(packed_index + 1);
      }
      prim_ops[packed_index] = it.first;
    }
    for (const auto& it : prim_ops) {
      oss << it << ", ";
    }
    if (!prim_ops.empty()) oss.seekp(-2, oss.cur);
    oss << "]" << std::endl;
  
    return oss.str();
  }
  
  void SaveHeader(dmlc::Stream* strm) {
    uint64_t header = kTVMVMBytecodeMagic;
    strm->Write(header);
    std::string version = TVM_VERSION;
    strm->Write(version);
  }
  
  TVMByteArray Executable::Save() {
    // Initialize the stream object.
    code_.clear();
    dmlc::MemoryStringStream strm(&code_);
  
    // Save header
    SaveHeader(&strm);
  
    // Save virtual devices section.
    SaveVirtualDevicesSection(&strm);
  
    // Global section.
    SaveGlobalSection(&strm);
  
    // Constant section.
    SaveConstantSection(&strm);
  
    // Primitive names.
    SavePrimitiveOpNames(&strm);
  
    // Code section.
    SaveCodeSection(&strm);
  
    TVMByteArray arr;
    arr.data = code_.c_str();
    arr.size = code_.length();
    return arr;
  }
  
  void Executable::SaveVirtualDevicesSection(dmlc::Stream* strm) {
    strm->Write(virtual_devices);
    strm->Write(host_device_index);
  }
  
  void Executable::MoveLateBoundConstantsToStream(dmlc::Stream* stream, size_t byte_limit) {
    ICHECK(late_bound_constant_names.empty());
    late_bound_constant_names.reserve(constants.size());
    Map<String, NDArray> map;
    size_t total_late_bound_bytes = 0;
    for (size_t const_index = 0; const_index < constants.size(); ++const_index) {
      const auto ndarray = Downcast<NDArray>(constants[const_index]);
      ICHECK(ndarray.defined()) << "Undefined constant at index " << const_index;
      size_t num_bytes = runtime::GetDataSize(*ndarray.operator->());
      if (num_bytes < byte_limit) {
        // Leave as immediate.
        late_bound_constant_names.emplace_back(nullptr);
        continue;
      }
      total_late_bound_bytes += num_bytes;
      std::ostringstream os;
      os << "const_" << const_index;
      String name = os.str();
      map.Set(name, Downcast<NDArray>(std::move(constants[const_index])));
      late_bound_constant_names.emplace_back(std::move(name));
    }
    VLOG(1) << "moved " << map.size() << " constants of " << total_late_bound_bytes
            << " bytes (out of " << constants.size() << " overall) to be late-bound";
    runtime::SaveParams(stream, map);
  }
  
  void Executable::MoveLateBoundConstantsToFile(const std::string& path, size_t byte_limit) {
    std::string bytes;
    dmlc::MemoryStringStream stream(&bytes);
    MoveLateBoundConstantsToStream(&stream, byte_limit);
    SaveBinaryToFile(path, bytes);
  }
  
  void Executable::LoadLateBoundConstantsFromStream(dmlc::Stream* stream) {
    if (late_bound_constant_names.empty()) {
      VLOG(1) << "Found no late-bound constants to load";
      return;
    }
    ICHECK_EQ(late_bound_constant_names.size(), constants.size());
    Map<String, NDArray> map = runtime::LoadParams(stream);
    VLOG(1) << "loaded " << map.size() << " late-bound constants";
    for (size_t const_index = 0; const_index < constants.size(); ++const_index) {
      if (!late_bound_constant_names[const_index].defined()) {
        ICHECK(constants[const_index].defined())
            << "Undefined immediate constant at index " << const_index;
        continue;
      }
      const String& name = late_bound_constant_names[const_index];
      ICHECK(!constants[const_index].defined()) << "Unexpected constant at index " << const_index;
      auto itr = map.find(name);
      ICHECK(itr != map.end()) << "No binding for late-bound constant at index " << const_index
                               << " with name '" << name << "'";
      constants[const_index] = (*itr).second;
      map.erase(name);
    }
    late_bound_constant_names.clear();
    ICHECK(map.empty()) << "Have " << map.size() << " unused late-bound constants";
  }
  
  void Executable::LoadLateBoundConstantsFromFile(const std::string& path) {
    std::string bytes;
    LoadBinaryFromFile(path, &bytes);
    dmlc::MemoryStringStream stream(&bytes);
    LoadLateBoundConstantsFromStream(&stream);
  }
  
  void Executable::SaveGlobalSection(dmlc::Stream* strm) {
    std::vector<std::pair<std::string, Index>> globals(this->global_map.begin(),
                                                       this->global_map.end());
    auto comp = [](const std::pair<std::string, Index>& a, const std::pair<std::string, Index>& b) {
      return a.second < b.second;
    };
    std::sort(globals.begin(), globals.end(), comp);
  
    std::vector<std::string> glbs;
    for (const auto& it : globals) {
      glbs.push_back(it.first);
    }
    strm->Write(glbs);
  }
  
  namespace {
  // Tags to distinguish immediate vs late-bound constants in constants table bytestream.
  constexpr uint32_t kImmediateConstTag = 0;
  constexpr uint32_t kLateBoundConstTag = 1;
  }  // namespace
  
  void Executable::SaveConstantSection(dmlc::Stream* stream) {
    // Save the overall number of constants.
    stream->Write(static_cast<uint64_t>(constants.size()));
  
    for (size_t const_index = 0; const_index < constants.size(); ++const_index) {
      if (late_bound_constant_names.empty() || !late_bound_constant_names[const_index].defined()) {
        // Tag immediate constants by 0.
        stream->Write(kImmediateConstTag);
        // Write as DLTensor.
        const auto ndarray = Downcast<runtime::NDArray>(constants[const_index]);
        ICHECK(ndarray.defined());
        runtime::SaveDLTensor(stream, ndarray.operator->());
        VLOG(1) << "save " << const_index << " as immediate";
      } else {
        // Tag late-bound constants by 1.
        const String& name = late_bound_constant_names[const_index];
        ICHECK(!constants[const_index].defined());
        stream->Write(kLateBoundConstTag);
        // Write a string.
        stream->Write(std::string(name));
        VLOG(1) << "save " << const_index << " as late-bound";
      }
    }
  
    VLOG(1) << "saved " << constants.size() << " constants";
  
    // Save the const to device index mapping.
    stream->Write(const_device_indexes);
  }
  
  void Executable::LoadConstantSection(dmlc::Stream* stream) {
    uint64_t sz;
    // Load the overall number of constants.
    STREAM_CHECK(stream->Read(&sz, sizeof(sz)), "constants table size");
    size_t size = static_cast<size_t>(sz);
  
    VLOG(1) << "loading " << size << " constants";
  
    constants.resize(size);
    late_bound_constant_names.resize(size);
    bool any_late_bound = false;
  
    // Load each of the constants.
    for (size_t const_index = 0; const_index < size; const_index++) {
      uint32_t tag;
      STREAM_CHECK(stream->Read(&tag, sizeof(tag)), "constant tag");
      if (tag == kImmediateConstTag) {
        // Immediate constants tagged by 0.
        VLOG(1) << "load " << const_index << " as immediate";
        runtime::NDArray ndarray;
        STREAM_CHECK(ndarray.Load(stream), "constant tensor");
        constants[const_index] = std::move(ndarray);
        late_bound_constant_names[const_index] = String(ObjectPtr<StringObj>(nullptr));
      } else if (tag == kLateBoundConstTag) {
        // Late-bound constants tagged by 1.
        VLOG(1) << "load " << const_index << " as late-bound";
        std::string name;
        STREAM_CHECK(stream->Read(&name), "late-bound constant name");
        constants[const_index] = NDArray(nullptr);
        late_bound_constant_names[const_index] = std::move(name);
        any_late_bound = true;
      } else {
        STREAM_CHECK(false, "constant tag");
      }
    }
  
    if (!any_late_bound) {
      late_bound_constant_names.clear();
    }
  
    // Load the const to device index mapping.
    std::vector<Index> indexes;
    indexes.reserve(size);
    STREAM_CHECK(stream->Read(&indexes), "constant devices");
    ICHECK_EQ(size, indexes.size());
    const_device_indexes = std::move(indexes);
  }
  
  void Executable::SavePrimitiveOpNames(dmlc::Stream* strm) {
    std::vector<std::string> primitive_names;
    for (const auto& it : this->primitive_map) {
      auto packed_index = static_cast<size_t>(it.second);
      if (primitive_names.size() <= packed_index) {
        primitive_names.resize(packed_index + 1);
      }
      primitive_names[packed_index] = it.first;
    }
    strm->Write(primitive_names);
    std::map<uint64_t, std::map<std::string, std::string>> primitive_attrs;
    for (const auto& it : this->op_attrs) {
      auto packed_index = static_cast<size_t>(it.first);
      std::map<std::string, std::string> attrs;
      for (const auto& elem : it.second) {
        // TODO(tkonolige): cannot serialize ObjectRefs with dmlc's serializer, so we just serialize
        // strings for now
        if (elem.second.as<StringObj>()) {
          attrs[elem.first] = Downcast<String>(elem.second);
        }
      }
      primitive_attrs[packed_index] = attrs;
    }
    strm->Write(primitive_attrs);
  }
  
  // Serialize a virtual machine instruction. It creates a list that contains the
  // hash, opcode, and all fields of an instruction.
  //
  // For example, the function signature used to create an `AllocTensor`
  // instruction is:
  //   Instruction AllocTensor(std::vector<Index> shape, DLDataType dtype, RegName dst)
  //
  // The serialized form will be:
  //   `hash 5 dtype.code dtype.bits dtype.lanes ndim dst_register val1 val2 ... valn`
  //
  // where hash is the hash of serialized instruction that is computed internally
  // by the `VMInstructionExecutable`. It is used for sanity check before decoding.
  // 5 shows opcode of `AllocTensor`, `(dtype.code dtype.bits dtype.lanes)`
  // represents a `DLDataType`, `ndim` is the number of dimensions, `dst_register`
  // is the destination register, and the rest of it together indicates the shape
  // of the tensor to be allocated.
  VMInstructionSerializer SerializeInstruction(const Instruction& instr) {
    std::vector<Index> fields;
    // Save the opcode.
    VLOG(2) << "Serializing: " << instr << std::endl;
    switch (instr.op) {
      case Opcode::Move: {
        // Number of fields = 2
        fields.assign({instr.from, instr.dst});
        break;
      }
      case Opcode::Ret: {
        // Number of fields = 1
        fields.push_back(instr.result);
        break;
      }
      case Opcode::Fatal: {
        // Number of fields = 0
        break;
      }
      case Opcode::InvokePacked: {
        // Number of fields = 3 + instr.arity
        // Note that arity includes both input arguments and outputs. We will
        // put all the `arity` number of fields in the end for serialization.
        fields.assign({instr.packed_index, instr.arity, instr.output_size});
        // Save the args.
        fields.insert(fields.end(), instr.packed_args, instr.packed_args + instr.arity);
        break;
      }
      case Opcode::AllocTensor: {
        // Number of fields = 7 + instr.alloc_tensor.ndim
        fields.push_back(instr.alloc_tensor.storage);
        fields.push_back(instr.alloc_tensor.offset);
        // Save `DLDataType` and the dst register.
        const auto& dtype = instr.alloc_tensor.dtype;
        fields.push_back(dtype.code);
        fields.push_back(dtype.bits);
        fields.push_back(dtype.lanes);
  
        // The number of dimensions is not needed for constructing an
        // `AllocTensor` instruction as it equals to the length of the `shape`
        // vector. However, we save it to conveniently deserialize the instruction
        // because we will know how many fields are needed by the `shape` argument.
        fields.push_back(instr.alloc_tensor.ndim);
        fields.push_back(instr.dst);
  
        // Save the shape of the tensor.
        // Note that this field is rotated to the end of the list.
        fields.insert(fields.end(), instr.alloc_tensor.shape,
                      instr.alloc_tensor.shape + instr.alloc_tensor.ndim);
        break;
      }
      case Opcode::AllocTensorReg: {
        // Number of fields = 7
        fields.push_back(instr.alloc_tensor_reg.storage);
        fields.push_back(instr.alloc_tensor_reg.offset);
        fields.push_back(instr.alloc_tensor_reg.shape_register);
        // Save `DLDataType` and the dst register.
        const auto& dtype = instr.alloc_tensor_reg.dtype;
        fields.push_back(dtype.code);
        fields.push_back(dtype.bits);
        fields.push_back(dtype.lanes);
        fields.push_back(instr.dst);
        break;
      }
      case Opcode::AllocStorage: {
        fields.push_back(instr.alloc_storage.allocation_size);
        fields.push_back(instr.alloc_storage.alignment);
        // Save `DLDataType` and the dst register.
        const auto& dtype = instr.alloc_storage.dtype_hint;
        fields.push_back(dtype.code);
        fields.push_back(dtype.bits);
        fields.push_back(dtype.lanes);
        fields.push_back(instr.alloc_storage.device_index);
        fields.push_back(instr.dst);
        break;
      }
      case Opcode::AllocADT: {
        // Number of fields = 3 + instr.num_fields
        fields.assign({instr.constructor_tag, instr.num_fields, instr.dst});
  
        // Save the fields.
        fields.insert(fields.end(), instr.datatype_fields, instr.datatype_fields + instr.num_fields);
        break;
      }
      case Opcode::AllocClosure: {
        // Number of fields = 3 + instr.num_freevar
        fields.assign({instr.clo_index, instr.num_freevar, instr.dst});
  
        // Save the free vars.
        fields.insert(fields.end(), instr.free_vars, instr.free_vars + instr.num_freevar);
        break;
      }
      case Opcode::If: {
        // Number of fields = 4
        fields.assign({instr.if_op.test, instr.if_op.target, instr.if_op.true_offset,
                       instr.if_op.false_offset});
        break;
      }
      case Opcode::Invoke: {
        // Number of fields = 3 + instr.num_args
        fields.assign({instr.func_index, instr.num_args, instr.dst});
  
        // Save the args.
        fields.insert(fields.end(), instr.invoke_args_registers,
                      instr.invoke_args_registers + instr.num_args);
        break;
      }
      case Opcode::InvokeClosure: {
        // Number of fields = 3 + instr.num_closure_args
        fields.assign({instr.closure, instr.num_closure_args, instr.dst});
  
        // Save the args.
        fields.insert(fields.end(), instr.closure_args, instr.closure_args + instr.num_closure_args);
        break;
      }
      case Opcode::LoadConst: {
        // Number of fields = 2
        fields.assign({instr.const_index, instr.dst});
        break;
      }
      case Opcode::LoadConsti: {
        // Number of fields = 2
        fields.assign({instr.load_consti.val, instr.dst});
        break;
      }
      case Opcode::GetField: {
        // Number of fields = 3
        fields.assign({instr.object, instr.field_index, instr.dst});
        break;
      }
      case Opcode::GetTag: {
        // Number of fields = 2
        fields.assign({instr.get_tag.object, instr.dst});
        break;
      }
      case Opcode::Goto: {
        // Number of fields = 1
        fields.push_back(instr.pc_offset);
        break;
      }
      case Opcode::ShapeOf: {
        // Number of fields = 2
        fields.assign({instr.shape_of.tensor, instr.dst});
        break;
      }
      case Opcode::ReshapeTensor: {
        // Number of fields = 3
        fields.assign({instr.reshape_tensor.tensor, instr.reshape_tensor.newshape, instr.dst});
        break;
      }
      case Opcode::DeviceCopy: {
        // Number of fields = 4
        fields.assign({instr.device_copy.src, instr.device_copy.src_device_index,
                       instr.device_copy.dst_device_index, instr.dst});
        break;
      }
      default:
        LOG(FATAL) << "Invalid opcode" << static_cast<int>(instr.op);
        break;
    }
  
    return VMInstructionSerializer(static_cast<Index>(instr.op), fields);
  }
  
  void Executable::SaveCodeSection(dmlc::Stream* strm) {
    // Save the number of functions.
    strm->Write(static_cast<uint64_t>(this->functions.size()));
    for (const auto& func : this->functions) {
      // Save the function info.
      VMFunctionSerializer func_format(func.name, func.register_file_size, func.instructions.size(),
                                       func.params, func.param_device_indexes);
      func_format.Save(strm);
  
      // Serialize each instruction.
      for (const auto& instr : func.instructions) {
        const auto& serialized_instr = SerializeInstruction(instr);
        serialized_instr.Save(strm);
      }
    }
  }
  
  void LoadHeader(dmlc::Stream* strm) {
    // Check header.
    uint64_t header;
    STREAM_CHECK(strm->Read(&header), "header");
    STREAM_CHECK(header == kTVMVMBytecodeMagic, "header");
  
    // Check version.
    std::string version;
    STREAM_CHECK(strm->Read(&version), "version");
    STREAM_CHECK(version == TVM_VERSION, "version");
  }
  
  runtime::Module Executable::GetLib() const {
    ICHECK_LE(this->imports_.size(), 1)
        << "The kernel library must be imported as the only module in an Executable";
  
    if (this->imports().size() == 0) {
      return Module(nullptr);
    } else {
      return this->imports_[0];
    }
  }
  
  void Executable::SetLib(const runtime::Module& lib) {
    ICHECK(lib.defined()) << "the provided library can not be null";
  
    ICHECK_EQ(this->imports_.size(), 0)
        << "A VMExecutable should never have more than one import inside an the executable, \n"
        << "the first import should *always* be the library containing"
        << "the platform specific kernel code";
  
    this->Import(lib);
  }
  
  runtime::Module Executable::Load(const std::string& code, const runtime::Module lib) {
    auto exec = make_object<Executable>();
  
    // Support null-initialization of lib, to enable initialization during
    // deserialization before we have deserialized the imports.
    if (lib.defined()) {
      exec->SetLib(lib);
    }
  
    exec->code_ = code;
    dmlc::MemoryStringStream strm(&exec->code_);
  
    // Load header.
    LoadHeader(&strm);
  
    // Virtual devices section
    exec->LoadVirtualDevicesSection(&strm);
  
    // Global section.
    exec->LoadGlobalSection(&strm);
  
    // Constant section.
    exec->LoadConstantSection(&strm);
  
    // Primitive names that will be invoked by `InvokePacked` instructions.
    exec->LoadPrimitiveOpNames(&strm);
  
    // Code section.
    exec->LoadCodeSection(&strm);
  
    return runtime::Module(exec);
  }
  
  void Executable::LoadVirtualDevicesSection(dmlc::Stream* strm) {
    STREAM_CHECK(strm->Read(&virtual_devices), "virtual_device");
    STREAM_CHECK(strm->Read(&host_device_index), "virtual_device");
    ICHECK(host_device_index >= 0 && host_device_index < static_cast<int>(virtual_devices.size()));
  }
  
  void Executable::LoadGlobalSection(dmlc::Stream* strm) {
    std::vector<std::string> globals;
    STREAM_CHECK(strm->Read(&globals), "global");
    for (size_t i = 0; i < globals.size(); i++) {
      this->global_map.insert({globals[i], i});
    }
  }
  
  void Executable::LoadPrimitiveOpNames(dmlc::Stream* strm) {
    std::vector<std::string> primitive_names;
    STREAM_CHECK(strm->Read(&primitive_names), "primitive name");
    for (size_t i = 0; i < primitive_names.size(); i++) {
      this->primitive_map.insert({primitive_names[i], i});
    }
  
    std::map<uint64_t, std::map<std::string, std::string>> primitive_attrs;
    STREAM_CHECK(strm->Read(&primitive_attrs), "primitive attrs");
    for (const auto& fn : primitive_attrs) {
      std::vector<std::pair<String, ObjectRef>> attrs;
      for (const auto& elem : fn.second) {
        attrs.push_back({elem.first, String(elem.second)});
      }
      this->op_attrs[fn.first] = Map<String, ObjectRef>(attrs.begin(), attrs.end());
    }
  }
  
  // Extract the `cnt` number of fields started at `start` from the list
  // `instr_fields`.
  inline std::vector<Index> ExtractFields(const std::vector<Index>& instr_fields, Index start,
                                          Index cnt) {
    ICHECK_LE(static_cast<size_t>(start + cnt), instr_fields.size());
    std::vector<Index> ret;
    for (auto i = start; i < start + cnt; i++) {
      ret.push_back(instr_fields[i]);
    }
    return ret;
  }
  
  Instruction DeserializeInstruction(const VMInstructionSerializer& instr) {
    Opcode opcode = static_cast<Opcode>(instr.opcode);
    switch (opcode) {
      case Opcode::Move: {
        // Number of fields = 2
        DCHECK_EQ(instr.fields.size(), 2U);
        return Instruction::Move(instr.fields[0], instr.fields[1]);
      }
      case Opcode::Ret: {
        // Number of fields = 1
        DCHECK_EQ(instr.fields.size(), 1U);
        return Instruction::Ret(instr.fields[0]);
      }
      case Opcode::Fatal: {
        // Number of fields = 0
        DCHECK(instr.fields.empty());
        return Instruction::Fatal();
      }
      case Opcode::InvokePacked: {
        // Number of fields = 3 + instr.arity
        DCHECK_GE(instr.fields.size(), 3U);
        DCHECK_EQ(instr.fields.size(), 3U + static_cast<size_t>(instr.fields[1]));
  
        Index packed_index = instr.fields[0];
        Index arity = instr.fields[1];
        Index output_size = instr.fields[2];
        std::vector<RegName> args = ExtractFields(instr.fields, 3, arity);
        return Instruction::InvokePacked(packed_index, arity, output_size, args);
      }
      case Opcode::AllocTensor: {
        // Number of fields = 7 + instr.alloc_tensor.ndim
        DCHECK_GE(instr.fields.size(), 7U);
        DCHECK_EQ(instr.fields.size(), 7U + static_cast<size_t>(instr.fields[5]));
  
        RegName storage_reg = instr.fields[0];
        RegName offset = instr.fields[1];
  
        DLDataType dtype;
        dtype.code = instr.fields[2];
        dtype.bits = instr.fields[3];
        dtype.lanes = instr.fields[4];
  
        Index ndim = instr.fields[5];
        RegName dst = instr.fields[6];
  
        std::vector<Index> shape = ExtractFields(instr.fields, 7, ndim);
  
        return Instruction::AllocTensor(storage_reg, offset, shape, dtype, dst);
      }
      case Opcode::AllocTensorReg: {
        // Number of fields = 7
        DCHECK_EQ(instr.fields.size(), 7U);
  
        RegName storage_reg = instr.fields[0];
        RegName offset = instr.fields[1];
        Index shape_register = instr.fields[2];
  
        DLDataType dtype;
        dtype.code = instr.fields[3];
        dtype.bits = instr.fields[4];
        dtype.lanes = instr.fields[5];
  
        RegName dst = instr.fields[6];
  
        return Instruction::AllocTensorReg(storage_reg, offset, shape_register, dtype, dst);
      }
      case Opcode::AllocADT: {
        // Number of fields = 3 + instr.num_fields
        DCHECK_GE(instr.fields.size(), 3U);
        DCHECK_EQ(instr.fields.size(), 3U + static_cast<size_t>(instr.fields[1]));
  
        Index constructor_tag = instr.fields[0];
        Index num_fields = instr.fields[1];
        RegName dst = instr.fields[2];
        std::vector<Index> fields = ExtractFields(instr.fields, 3, num_fields);
  
        return Instruction::AllocADT(constructor_tag, num_fields, fields, dst);
      }
      case Opcode::AllocClosure: {
        // Number of fields = 3 + instr.num_freevar
        DCHECK_GE(instr.fields.size(), 3U);
        DCHECK_EQ(instr.fields.size(), 3U + static_cast<size_t>(instr.fields[1]));
  
        Index clo_index = instr.fields[0];
        Index num_freevar = instr.fields[1];
        RegName dst = instr.fields[2];
        std::vector<Index> free_vars = ExtractFields(instr.fields, 3, num_freevar);
  
        return Instruction::AllocClosure(clo_index, num_freevar, free_vars, dst);
      }
      case Opcode::AllocStorage: {
        // Number of fields = 7
        DCHECK_GE(instr.fields.size(), 7U);
        Index allocation_size = instr.fields[0];
        Index alignment = instr.fields[1];
  
        DLDataType dtype;
        dtype.code = instr.fields[2];
        dtype.bits = instr.fields[3];
        dtype.lanes = instr.fields[4];
  
        Index device_type = instr.fields[5];
        RegName dst = instr.fields[6];
  
        return Instruction::AllocStorage(allocation_size, alignment, dtype, device_type, dst);
      }
      case Opcode::If: {
        // Number of fields = 4
        DCHECK_EQ(instr.fields.size(), 4U);
        Index test = instr.fields[0];
        Index target = instr.fields[1];
        Index true_offset = instr.fields[2];
        Index false_offset = instr.fields[3];
  
        return Instruction::If(test, target, true_offset, false_offset);
      }
      case Opcode::Invoke: {
        // Number of fields = 3 + instr.num_args
        DCHECK_GE(instr.fields.size(), 3U);
        DCHECK_EQ(instr.fields.size(), 3U + static_cast<size_t>(instr.fields[1]));
  
        Index func_index = instr.fields[0];
        Index num_args = instr.fields[1];
        RegName dst = instr.fields[2];
        std::vector<Index> args = ExtractFields(instr.fields, 3, num_args);
  
        return Instruction::Invoke(func_index, args, dst);
      }
      case Opcode::InvokeClosure: {
        // Number of fields = 3 + instr.num_closure_args
        DCHECK_GE(instr.fields.size(), 3U);
        DCHECK_EQ(instr.fields.size(), 3U + static_cast<size_t>(instr.fields[1]));
  
        Index closure = instr.fields[0];
        Index num_closure_args = instr.fields[1];
        RegName dst = instr.fields[2];
        std::vector<Index> args = ExtractFields(instr.fields, 3, num_closure_args);
  
        return Instruction::InvokeClosure(closure, args, dst);
      }
      case Opcode::LoadConst: {
        // Number of fields = 2
        DCHECK_EQ(instr.fields.size(), 2U);
        return Instruction::LoadConst(instr.fields[0], instr.fields[1]);
      }
      case Opcode::LoadConsti: {
        // Number of fields = 2
        DCHECK_EQ(instr.fields.size(), 2U);
        return Instruction::LoadConsti(instr.fields[0], instr.fields[1]);
      }
      case Opcode::GetField: {
        // Number of fields = 3
        DCHECK_EQ(instr.fields.size(), 3U);
        return Instruction::GetField(instr.fields[0], instr.fields[1], instr.fields[2]);
      }
      case Opcode::GetTag: {
        // Number of fields = 2
        DCHECK_EQ(instr.fields.size(), 2U);
        return Instruction::GetTag(instr.fields[0], instr.fields[1]);
      }
      case Opcode::Goto: {
        // Number of fields = 1
        DCHECK_EQ(instr.fields.size(), 1U);
        return Instruction::Goto(instr.fields[0]);
      }
      case Opcode::ShapeOf: {
        // Number of fields = 2
        DCHECK_EQ(instr.fields.size(), 2U);
        return Instruction::ShapeOf(instr.fields[0], instr.fields[1]);
      }
      case Opcode::ReshapeTensor: {
        // Number of fields = 3
        DCHECK_EQ(instr.fields.size(), 3U);
        return Instruction::ReshapeTensor(instr.fields[0], instr.fields[1], instr.fields[2]);
      }
      case Opcode::DeviceCopy: {
        // Number of fields = 4
        DCHECK_EQ(instr.fields.size(), 4U);
        return Instruction::DeviceCopy(instr.fields[0], instr.fields[1], instr.fields[2],
                                       instr.fields[3]);
      }
      default:
        LOG(FATAL) << "Invalid opcode" << instr.opcode;
        return Instruction();
    }
  }
  
  void Executable::LoadCodeSection(dmlc::Stream* strm) {
    // Load the number of functions.
    uint64_t sz;
    STREAM_CHECK(strm->Read(&sz, sizeof(sz)), "code");
  
    size_t num_funcs = static_cast<size_t>(sz);
    this->functions.resize(num_funcs);
    for (size_t i = 0; i < num_funcs; i++) {
      // Load the function info.
      VMFunctionSerializer loaded_func;
      STREAM_CHECK(loaded_func.Load(strm), "code/function");
  
      // Load the instructions.
      std::vector<Instruction> instructions;
      for (size_t j = 0; j < loaded_func.num_instructions; j++) {
        VMInstructionSerializer instr;
        std::vector<Index> instr_fields;
        STREAM_CHECK(instr.Load(strm), "code/instruction");
        instructions.push_back(DeserializeInstruction(instr));
      }
  
      // Create the VM function.
      VMFunction vm_func =
          VMFunction(loaded_func.name, loaded_func.params, instructions,
                     loaded_func.register_file_size, loaded_func.param_device_indexes);
      auto it = this->global_map.find(loaded_func.name);
      ICHECK(it != this->global_map.end());
      ICHECK_LE(it->second, this->global_map.size());
      this->functions[it->second] = vm_func;
    }
  }
  
  void Executable::SaveToBinary(dmlc::Stream* stream) {
    auto code_bytes = this->Save();
    std::string code(code_bytes.data, code_bytes.size);
    stream->Write(code);
  
    ICHECK(this->imports()[0].defined()) << "the library must be imported before serialization";
  }
  
  Module ExecutableLoadBinary(void* strm) {
    dmlc::Stream* stream = static_cast<dmlc::Stream*>(strm);
    std::string code;
    stream->Read(&code);
    auto exec = Executable::Load(code, Module());
    return exec;
  }
  
  void Executable::SaveToFile(const std::string& path, const std::string& format) {
    std::string data;
    dmlc::MemoryStringStream writer(&data);
    dmlc::SeekStream* strm = &writer;
    SaveToBinary(strm);
    SaveBinaryToFile(path, data);
  }
  
  TVM_REGISTER_GLOBAL("runtime.module.loadbinary_VMExecutable").set_body_typed(ExecutableLoadBinary);
  
  // Load module from module.
  Module ExecutableLoadFile(const std::string& file_name, const std::string& format) {
    std::string data;
    LoadBinaryFromFile(file_name, &data);
    dmlc::MemoryStringStream reader(&data);
    dmlc::Stream* strm = &reader;
    auto exec = ExecutableLoadBinary(reinterpret_cast<void*>(strm));
    return exec;
  }
  
  TVM_REGISTER_GLOBAL("runtime.module.loadfile_VMExecutable").set_body_typed(ExecutableLoadFile);
  
  TVM_REGISTER_GLOBAL("runtime.GetNumOfGlobals").set_body([](TVMArgs args, TVMRetValue* rv) {
    runtime::Module mod = args[0];
    const auto* exec = dynamic_cast<Executable*>(mod.operator->());
    ICHECK(exec);
    *rv = static_cast<int>(exec->global_map.size());
  });
  
  TVM_REGISTER_GLOBAL("runtime.GetGlobalFields").set_body([](TVMArgs args, TVMRetValue* rv) {
    runtime::Module mod = args[0];
    const auto* exec = dynamic_cast<Executable*>(mod.operator->());
    ICHECK(exec);
    int idx = args[1];
    std::vector<std::pair<std::string, Index>> globals(exec->global_map.begin(),
                                                       exec->global_map.end());
    auto comp = [](const std::pair<std::string, Index>& a, const std::pair<std::string, Index>& b) {
      return a.second < b.second;
    };
    std::sort(globals.begin(), globals.end(), comp);
    ICHECK_LT(idx, globals.size());
    *rv = globals[idx].first;
  });
  
  TVM_REGISTER_GLOBAL("runtime.GetNumOfPrimitives").set_body([](TVMArgs args, TVMRetValue* rv) {
    runtime::Module mod = args[0];
    const auto* exec = dynamic_cast<Executable*>(mod.operator->());
    ICHECK(exec);
    *rv = static_cast<int>(exec->primitive_map.size());
  });
  
  TVM_REGISTER_GLOBAL("runtime.GetPrimitiveFields").set_body([](TVMArgs args, TVMRetValue* rv) {
    runtime::Module mod = args[0];
    const auto* exec = dynamic_cast<Executable*>(mod.operator->());
    ICHECK(exec);
    int idx = args[1];
    ICHECK_GE(idx, 0);
    ICHECK_LT(idx, exec->primitive_map.size());
  
    for (const auto& it : exec->primitive_map) {
      if (idx == static_cast<int>(it.second)) {
        *rv = it.first;
        break;
      }
    }
  });
  
  TVM_REGISTER_GLOBAL("runtime.Load_Executable")
      .set_body_typed([](std::string code, runtime::Module lib) {
        return Executable::Load(code, lib);
      });
  
  }  // namespace vm
  }  // namespace runtime
  }  // namespace tvm
  