vm.h

Go to the documentation of this file.

/*
 * 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.
 */

#ifndef TVM_RUNTIME_VM_VM_H_
#define TVM_RUNTIME_VM_VM_H_

#include <tvm/runtime/container/closure.h>
#include <tvm/runtime/module.h>
#include <tvm/runtime/object.h>
#include <tvm/runtime/packed_func.h>
#include <tvm/runtime/registry.h>
#include <tvm/runtime/vm/bytecode.h>
#include <tvm/runtime/vm/executable.h>
#include <tvm/runtime/vm/memory_manager.h>

#include <memory>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

namespace tvm {
namespace runtime {
namespace vm {

class VMClosureObj : public ClosureObj {
 public:
  size_t func_index;
  std::vector<ObjectRef> free_vars;

  static constexpr const uint32_t _type_index = TypeIndex::kDynamic;
  static constexpr const char* _type_key = "vm.Closure";
  TVM_DECLARE_FINAL_OBJECT_INFO(VMClosureObj, ClosureObj);
};

class VMClosure : public Closure {
 public:
  VMClosure(size_t func_index, std::vector<ObjectRef> free_vars);
  TVM_DEFINE_OBJECT_REF_METHODS(VMClosure, Closure, VMClosureObj);
};

struct VMFunction {
  std::string name;
  std::vector<std::string> params;
  std::vector<Instruction> instructions;
  Index register_file_size;
  std::vector<Index> params_device_type;

  VMFunction(const std::string& name, std::vector<std::string> params,
             const std::vector<Instruction>& instructions, Index register_file_size,
             const std::vector<Index> params_device_type = {})
      : name(name),
        params(params),
        instructions(instructions),
        register_file_size(register_file_size),
        params_device_type(params_device_type) {}

  VMFunction() {}

  friend std::ostream& operator<<(std::ostream& os, const VMFunction&);
};

struct VMFrame {
  Index pc;
  Index func_index;
  Index args;
  const Instruction* code;

  std::vector<ObjectRef> register_file;

  RegName caller_return_register;

  VMFrame(Index pc, Index func_index, Index args, const Instruction* code, Index register_file_size)
      : pc(pc),
        func_index(func_index),
        args(args),
        code(code),
        register_file(register_file_size),
        caller_return_register(0) {}
};

class VirtualMachine : public runtime::ModuleNode {
 public:
  virtual PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self);

  virtual ~VirtualMachine() {}

  const char* type_key() const final { return "VirtualMachine"; }

  VirtualMachine() : frames_(), func_index_(0), code_(nullptr), pc_(0), exec_(nullptr) {}

  virtual void LoadExecutable(const Executable* exec);

 protected:
  void PushFrame(Index arg_count, Index ret_pc, const VMFunction& vm_func);

  Index PopFrame();

  inline void WriteRegister(RegName reg, const ObjectRef& obj);

  ObjectRef ReadRegister(RegName reg) const;

  int64_t LoadScalarInt(RegName reg) const;

  ObjectRef Invoke(const VMFunction& func, const std::vector<ObjectRef>& args);

  // TODO(@jroesch): I really would like this to be a global variable.
  ObjectRef Invoke(const std::string& name, const std::vector<ObjectRef>& args);

  virtual void InvokePacked(Index packed_index, const PackedFunc& func, Index arg_count,
                            Index output_size, const std::vector<ObjectRef>& args);

  void Init(const std::vector<Device>& devices, const std::vector<AllocatorType>& alloc_types);

  void RunLoop();

  Device GetDevice(Index device_type) const;

  void InvokeGlobal(const VMFunction& func, const std::vector<ObjectRef>& args);

  void SetInput(std::string name, TVMArgs args, int offset);

  virtual void OpStartHook(Instruction instr);

  virtual void OpStopHook();

 protected:
  std::vector<PackedFunc> packed_funcs_;
  std::vector<VMFrame> frames_;
  Index func_index_;
  const Instruction* code_;
  Index pc_;
  ObjectRef return_register_;
  const Executable* exec_;
  std::unordered_map<std::string, std::vector<ObjectRef>> inputs_;
  std::vector<Device> devices_;
  std::vector<Allocator*> allocators_;
  std::vector<ObjectRef> const_pool_;
};

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

#endif  // TVM_RUNTIME_VM_VM_H_