Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Register VM #7

Merged
merged 6 commits into from
Apr 12, 2019
Merged

Register VM #7

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
7b83311
Register VM
wweic Mar 30, 2019
e55a072
Address review feedback
wweic Apr 5, 2019
a0681d0
Inline stack slot
wweic Apr 10, 2019
dfcd401
Update reviewer comment
wweic Apr 10, 2019
b04456b
Get rid of special register number for return value
wweic Apr 10, 2019
9a9ac43
Shape tensor in Register
wweic Apr 10, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
132 changes: 75 additions & 57 deletions include/tvm/relay/vm/vm.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,9 +19,10 @@ namespace vm {

using namespace tvm::runtime;

using VirtualRegisterNum = size_t;

enum struct Opcode {
Push,
Pop,
Move,
Ret,
Invoke,
InvokeClosure,
Expand All @@ -31,58 +32,86 @@ enum struct Opcode {
AllocClosure,
GetField,
If,
Select,
LoadConst,
Goto,
Move,
Goto
};

struct Instruction {
struct TensorInfo {
int64_t* shape;
VirtualRegisterNum shape_register;
size_t ndim;
DLDataType dtype;
};

Opcode op;

// Destination register that the opcode writes to
VirtualRegisterNum dst;

union {
size_t stack_index;
TensorInfo tensor_info;

// For InvokeClosure
struct {
VirtualRegisterNum closure;
size_t closure_args_num;
VirtualRegisterNum* closure_args;
};
// For Ret
struct {
VirtualRegisterNum result;
};
// For Move
struct {
VirtualRegisterNum from;
};
struct {
size_t packed_index;
size_t arity;
size_t output_size;
VirtualRegisterNum* packed_args;
};
// For Select node
struct {
VirtualRegisterNum select_cond;
VirtualRegisterNum select_op1;
VirtualRegisterNum select_op2;
};
// For If node
struct {
VirtualRegisterNum if_cond;
size_t true_offset;
size_t false_offset;
};
// For Invoke
struct {
size_t func_index;
size_t num_args;
VirtualRegisterNum* invoke_args_registers;
};
struct {
size_t const_index;
};
struct {
size_t pc_offset;
};
// For GetField
struct {
size_t object_offset;
VirtualRegisterNum object;
size_t field_index;
};
// For AllocDatatype
struct {
size_t constructor_tag;
size_t num_fields;
VirtualRegisterNum* datatype_fields;
};
// For AllocClosure
struct {
size_t clo_index;
size_t num_freevar;
};
struct {
size_t pop_count;
};
struct {
size_t source;
size_t dest;
VirtualRegisterNum* free_vars;
};
};

Expand All @@ -94,28 +123,29 @@ struct Instruction {
};

// Helpers to build instructions.
Instruction Push(size_t stack_index);
Instruction Pop(size_t pop_count);
Instruction Ret();
Instruction InvokePacked(size_t packed_index, size_t arity, size_t output_size);
Instruction AllocTensor(const std::vector<int64_t>& shape, DLDataType dtype);
Instruction AllocDatatype(size_t tag, size_t num_fields);
Instruction AllocClosure(size_t func_index, size_t num_freevar);
Instruction GetField(size_t object_offset, size_t field_index);
Instruction If(size_t true_branch, size_t false_branch);
Instruction Select(VirtualRegisterNum cond, VirtualRegisterNum op1, VirtualRegisterNum op2, VirtualRegisterNum dst);
Instruction Ret(VirtualRegisterNum result);
Instruction InvokePacked(size_t packed_index, size_t arity, size_t output_size, const std::vector<VirtualRegisterNum>& args);
Instruction AllocTensor(VirtualRegisterNum shape_register, const std::vector<int64_t>& shape, DLDataType dtype, VirtualRegisterNum dst);
Instruction AllocDatatype(size_t tag, size_t num_fields, const std::vector<VirtualRegisterNum>& fields, VirtualRegisterNum dst);
Instruction AllocClosure(size_t func_index, size_t num_freevar, const std::vector<VirtualRegisterNum>& free_vars, VirtualRegisterNum dst);
Instruction GetField(VirtualRegisterNum object, size_t field_index, VirtualRegisterNum dst);
Instruction If(VirtualRegisterNum cond, size_t true_branch, size_t false_branch);
Instruction Goto(size_t pc_offset);
Instruction Invoke(size_t func_index);
Instruction InvokeClosure();
Instruction LoadConst(size_t const_index);
Instruction Move(size_t source, size_t dest);
Instruction Invoke(size_t func_index, const std::vector<VirtualRegisterNum>& args, VirtualRegisterNum dst);
Instruction InvokeClosure(VirtualRegisterNum closure, const std::vector<VirtualRegisterNum>& args, VirtualRegisterNum dst);
Instruction LoadConst(size_t const_index, VirtualRegisterNum dst);
Instruction Move(VirtualRegisterNum src, VirtualRegisterNum dst);

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

VMFunction(std::string name, size_t params, std::vector<Instruction> instructions)
: name(name), params(params), instructions(instructions) {}
VMFunction(std::string name, size_t params, std::vector<Instruction> instructions, size_t register_file_size)
: name(name), params(params), instructions(instructions), register_file_size(register_file_size)
{}

VMFunction() {}

Expand All @@ -131,65 +161,53 @@ void VMFunctionPrint(const VMFunction& vm_func);
*/
struct VMFrame {
size_t pc;
size_t bp;
size_t sp;
size_t func_index;
size_t args;
const Instruction* code;

VMFrame(size_t pc, size_t bp, size_t sp, size_t func_index, size_t args, const Instruction* code)
: pc(pc), bp(bp), sp(sp), func_index(func_index), args(args), code(code) {}
std::vector<Object> register_file;

VirtualRegisterNum caller_return_register;

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

struct VirtualMachine {
// TODO(@jroesch):
std::vector<PackedFunc> packed_funcs;
std::vector<VMFunction> functions;
std::vector<VMFrame> frames;
std::vector<Object> stack;
std::vector<Object> constants;

// Frame State
size_t func_index;
const Instruction* code;
size_t pc;
size_t bp;
// Special register to save function call return value
Object return_register;
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Put return_register in the VMFrame so it's possible that we can run VMFunction in parallel.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When Ret executed, the frame is poped out. So result will be gone before next instruction gets executed.


std::vector<TVMContext> ctxs;

// Interface debugging.
std::unordered_map<GlobalVar, size_t, NodeHash, NodeEqual> global_map;
std::unordered_map<size_t, Constructor> tag_index_map;

void PushFrame(size_t arg_count, size_t ret_pc, size_t sp, const VMFunction& vm_func);
void PushFrame(size_t arg_count, size_t ret_pc, const VMFunction& vm_func);
size_t PopFrame();
void InvokeGlobal(const VMFunction& func);
void InvokeGlobal(const VMFunction& func, const std::vector<Object>& args);
void Run();

inline void WriteRegister(VirtualRegisterNum r, Object v);
inline Object ReadRegister(VirtualRegisterNum r);

Object Invoke(const VMFunction& func, const std::vector<Object>& args);
Object Invoke(const GlobalVar& global, const std::vector<Object>& args);

// Ignore the method that dumps register info at compile-time if debugging
// mode is not enabled.
template <typename T = EnableRelayDebug>
typename std::enable_if<T::value, void>::type
DumpRegister();

template <typename T = EnableRelayDebug>
typename std::enable_if<!T::value, void>::type
DumpRegister() {}

// Ignore the method that dumps stack info at compile-time if debugging
// mode is not enabled.
template <typename T = EnableRelayDebug>
typename std::enable_if<T::value, void>::type DumpStack();

template <typename T = EnableRelayDebug>
typename std::enable_if<!T::value, void>::type DumpStack() {}

VirtualMachine() :
functions(), frames(), stack(),
func_index(0), code(nullptr), pc(0), bp(0) {}
functions(), frames(),
func_index(0), code(nullptr), pc(0) {}

void Init(const std::vector<TVMContext>& ctxs);

Expand Down
Loading