From 1389d20888db76bcf6e8decb45b7af88f3c2dcfc Mon Sep 17 00:00:00 2001 From: ziheng Date: Sat, 5 Aug 2017 20:30:40 -0700 Subject: [PATCH] [EXECUTOR] Split graph_executor to header file and (runtime) source file (#300) * [EXECUTOR] Split graph_executor to header file and (runtime) source file * Fix --- apps/graph_executor/src/graph_executor.cc | 210 +++--------------- apps/graph_executor/src/graph_executor.h | 119 ++++++++++ apps/graph_executor/src/graph_executor_ext.cc | 87 ++++++++ apps/graph_executor/tests/test_executor.py | 4 +- 4 files changed, 233 insertions(+), 187 deletions(-) create mode 100644 apps/graph_executor/src/graph_executor.h create mode 100644 apps/graph_executor/src/graph_executor_ext.cc diff --git a/apps/graph_executor/src/graph_executor.cc b/apps/graph_executor/src/graph_executor.cc index e1230f4c..d0a755f3 100644 --- a/apps/graph_executor/src/graph_executor.cc +++ b/apps/graph_executor/src/graph_executor.cc @@ -1,99 +1,23 @@ /*! * Copyright (c) 2017 by Contributors - * \file NNVM Graph executor. + * \file graph_executor.cc */ -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include +#include "./graph_executor.h" namespace tvm { namespace contrib { -using tvm::runtime::TVMArgs; -using tvm::runtime::TVMRetValue; -using tvm::runtime::PackedFunc; -using nnvm::StorageVector; -using nnvm::ShapeVector; -using nnvm::TShape; -using nnvm::NodeAttrs; - -/*! \brief DLPack compatible data types */ -using DLTypeVector = std::vector; - -/*! \brief The executor function */ -using FOpExec = std::function; - -/*! \brief macro to do C API call */ -#define TVM_CCALL(func) \ - { \ - int ret = (func); \ - CHECK_EQ(ret, 0) \ - << TVMGetLastError(); \ - } - -/*! \brief Graph Executor with TVM runtime */ -class GraphExecutor : public runtime::ModuleNode { - public: - const char* type_key() const { - return "GraphExecutor"; - } - PackedFunc GetFunction( - const std::string& name, - const std::shared_ptr& sptr_to_self); - // Destructor - ~GraphExecutor(); - // Setup with a given graph - void Init(const nnvm::Graph& g, TVMContext ctx); - // Copy data to index-th input - void SetInput(int index, DLTensor* data_in); - // Copy index-th output to data_out - void GetOutput(int index, DLTensor* data_out); - // Load parameters from stream - void LoadParams(dmlc::Stream* strm); - // Load parameters from binary file blob - void LoadParamsFromBlob(std::string param_blob); - // Execute the graph. - void Run(); - - private: - // functions - void SetupStorage(); - void SetupOpExecs(); - // Constructor to create TVM op - FOpExec CreateTVMOp(const nnvm::NodeAttrs& attrs, - std::vector inputs, - size_t num_inputs); - // The graph to be executed. - nnvm::Graph graph_; - // The execution context - TVMContext ctx_; - // Common storage pool - std::vector storage_pool_; - // The data shape - std::vector data_shape_; - // The data entry - std::vector data_entry_; - // The operation lambda on each node - std::vector op_execs_; - // The code module. - tvm::runtime::Module module_; -}; - PackedFunc GraphExecutor::GetFunction( const std::string& name, const std::shared_ptr& sptr_to_self) { // return member functions during query. if (name == "set_input") { return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { - this->SetInput(args[0], args[1]); + if (args[0].type_code() == kStr) { + this->SetInput(this->GetIndex(args[0]), args[1]); + } else { + this->SetInput(args[0], args[1]); + } }); } else if (name == "get_output") { return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { @@ -129,10 +53,17 @@ void GraphExecutor::Init(const nnvm::Graph& g, TVMContext ctx) { graph_ = g; ctx_ = ctx; module_ = g.GetAttr("module"); + this->SetupNameIndex(); this->SetupStorage(); this->SetupOpExecs(); } +int GraphExecutor::GetIndex(std::string name) { + CHECK(name_idx_.count(name)) + << name << " is not in the graph."; + return name_idx_.at(name); +} + void GraphExecutor::SetInput(int index, DLTensor* data_in) { const auto& idx = graph_.indexed_graph(); CHECK_LT(static_cast(index), idx.input_nodes().size()); @@ -147,33 +78,6 @@ void GraphExecutor::GetOutput(int index, DLTensor* data_out) { TVM_CCALL(TVMArrayCopyFromTo(&data_entry_[eid], data_out, nullptr)); } - -constexpr uint64_t kTVMNDArrayMagic = 0xDD5E40F096B4A13F; - -bool SaveDLTensor(dmlc::Stream* strm, DLTensor* tensor) { - uint64_t header = kTVMNDArrayMagic, reserved = 0; - strm->Write(&header, sizeof(header)); - strm->Write(&reserved, sizeof(reserved)); - - strm->Write(&tensor->ctx, sizeof(tensor->ctx)); - strm->Write(&tensor->ndim, sizeof(tensor->ndim)); - strm->Write(&tensor->dtype, sizeof(tensor->dtype)); - - int ndim = tensor->ndim; - strm->Write(tensor->shape, sizeof(int64_t) * ndim); - - int type_size = tensor->dtype.bits / 8; - int64_t size = 1; - for (int i = 0; i < ndim; ++i) { - size *= tensor->shape[i]; - } - int64_t data_byte_size = type_size * size; - strm->Write(&data_byte_size, sizeof(data_byte_size)); - strm->Write(tensor->data, data_byte_size); - return true; -} - - bool LoadDLTensor(dmlc::Stream* strm, DLTensor* tensor) { uint64_t header, reserved; CHECK(strm->Read(&header, sizeof(header))) @@ -209,37 +113,6 @@ bool LoadDLTensor(dmlc::Stream* strm, DLTensor* tensor) { return true; } - -constexpr uint64_t kTVMNDArrayListMagic = 0xF7E58D4F05049CB7; - -TVM_REGISTER_GLOBAL("tvm_graph._save_param_dict") -.set_body([](TVMArgs args, TVMRetValue *rv) { - std::string fname = args[0]; - int num_params = args[1]; - std::vector names; - names.reserve(num_params); - std::vector arrays; - arrays.reserve(num_params); - for (int i = 2; i < (2 + 2*num_params); i += 2) { - names.emplace_back(args[i].operator std::string()); - arrays.emplace_back(args[i+1].operator DLTensor*()); - } - - std::unique_ptr fo(dmlc::Stream::Create(fname.c_str(), "w")); - uint64_t header = kTVMNDArrayListMagic, reserved = 0; - fo->Write(&header, sizeof(header)); - fo->Write(&reserved, sizeof(reserved)); - - fo->Write(names); - { - uint64_t sz = static_cast(arrays.size()); - fo->Write(&sz, sizeof(sz)); - for (size_t i = 0; i < sz; ++i) { - SaveDLTensor(fo.get(), arrays[i]); - } - } - }); - void GraphExecutor::LoadParams(dmlc::Stream *strm) { uint64_t header, reserved; CHECK(strm->Read(&header)) @@ -277,6 +150,15 @@ void GraphExecutor::LoadParamsFromBlob(std::string param_blob) { this->LoadParams(&strm); } +void GraphExecutor::SetupNameIndex() { + nnvm::Symbol s; + s.outputs = graph_.outputs; + std::vector input_names = s.ListInputNames(nnvm::Symbol::kAll); + for (size_t i = 0; i < input_names.size(); ++i) { + name_idx_[input_names[i]] = i; + } +} + void GraphExecutor::SetupStorage() { const auto& idx = graph_.indexed_graph(); // Grab saved optimization plan from graph. @@ -399,23 +281,6 @@ FOpExec GraphExecutor::CreateTVMOp(const nnvm::NodeAttrs& attrs, return fexec; } -struct TVMOpParam : public dmlc::Parameter { - std::string func_name; - uint32_t num_inputs; - uint32_t num_outputs; - bool flatten_data; - DMLC_DECLARE_PARAMETER(TVMOpParam) { - DMLC_DECLARE_FIELD(func_name); - DMLC_DECLARE_FIELD(num_inputs) - .set_default(1); - DMLC_DECLARE_FIELD(num_outputs) - .set_default(1); - DMLC_DECLARE_FIELD(flatten_data) - .set_default(false); - } -}; -DMLC_REGISTER_PARAMETER(TVMOpParam); - /*! \brief Parse keyword arguments as PType arguments and save to parsed */ template inline void ParamParser(nnvm::NodeAttrs* attrs) { @@ -436,6 +301,8 @@ inline void ParamParser(nnvm::NodeAttrs* attrs) { attrs->parsed = std::move(param); } +DMLC_REGISTER_PARAMETER(TVMOpParam); + // ewise tvm op NNVM_REGISTER_OP(tvm_op) .set_attr_parser(ParamParser) @@ -448,33 +315,6 @@ NNVM_REGISTER_OP(tvm_op) return param.num_outputs; }); -// Create executor -tvm::runtime::Module CreateExecutor(nnvm::Graph g, TVMContext ctx) { - std::shared_ptr exec = - std::make_shared(); - exec->Init(g, ctx); - return tvm::runtime::Module(exec); -} - -TVM_REGISTER_GLOBAL("tvm_graph._create_executor") -.set_body([](TVMArgs args, TVMRetValue *rv) { - void* graph_handle = args[0]; - int device_type = args[1]; - int device_id = args[2]; - TVMContext ctx{static_cast(device_type), device_id}; - nnvm::Graph g = static_cast(graph_handle)[0]; - *rv = CreateExecutor(g, ctx); - }); - - -TVM_REGISTER_GLOBAL("tvm_graph._get_module_from_graph") -.set_body([](TVMArgs args, TVMRetValue *rv) { - void* graph_handle = args[0]; - nnvm::Graph* g = static_cast(graph_handle); - *rv = g->MoveCopyAttr("module"); - }); - - TVM_REGISTER_GLOBAL("tvm_graph._load_executor") .set_body([](TVMArgs args, TVMRetValue *rv) { std::string sym_json = args[0]; diff --git a/apps/graph_executor/src/graph_executor.h b/apps/graph_executor/src/graph_executor.h new file mode 100644 index 00000000..3953646c --- /dev/null +++ b/apps/graph_executor/src/graph_executor.h @@ -0,0 +1,119 @@ +/*! + * Copyright (c) 2017 by Contributors + * \file graph_executor.h + */ +#ifndef TVM_GRAPH_EXECUTOR_H_ +#define TVM_GRAPH_EXECUTOR_H_ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace tvm { +namespace contrib { + +using tvm::runtime::TVMArgs; +using tvm::runtime::TVMRetValue; +using tvm::runtime::PackedFunc; +using nnvm::StorageVector; +using nnvm::ShapeVector; +using nnvm::TShape; +using nnvm::NodeAttrs; + +/*! \brief DLPack compatible data types */ +using DLTypeVector = std::vector; + +/*! \brief The executor function */ +using FOpExec = std::function; + +/*! \brief macro to do C API call */ +#define TVM_CCALL(func) \ + { \ + int ret = (func); \ + CHECK_EQ(ret, 0) \ + << TVMGetLastError(); \ + } + +constexpr uint64_t kTVMNDArrayMagic = 0xDD5E40F096B4A13F; +constexpr uint64_t kTVMNDArrayListMagic = 0xF7E58D4F05049CB7; + +/*! \brief Graph Executor with TVM runtime */ +class GraphExecutor : public runtime::ModuleNode { + public: + const char* type_key() const { + return "GraphExecutor"; + } + PackedFunc GetFunction( + const std::string& name, + const std::shared_ptr& sptr_to_self); + // Destructor + ~GraphExecutor(); + // Setup with a given graph + void Init(const nnvm::Graph& g, TVMContext ctx); + // Get index of variable + int GetIndex(std::string name); + // Copy data to index-th input + void SetInput(int index, DLTensor* data_in); + // Copy index-th output to data_out + void GetOutput(int index, DLTensor* data_out); + // Load parameters from stream + void LoadParams(dmlc::Stream* strm); + // Load parameters from binary file blob + void LoadParamsFromBlob(std::string param_blob); + // Execute the graph. + void Run(); + + private: + // functions + void SetupNameIndex(); + void SetupStorage(); + void SetupOpExecs(); + // Constructor to create TVM op + FOpExec CreateTVMOp(const nnvm::NodeAttrs& attrs, + std::vector inputs, + size_t num_inputs); + // The graph to be executed. + nnvm::Graph graph_; + // The execution context + TVMContext ctx_; + // Common storage pool + std::vector storage_pool_; + // The data shape + std::vector data_shape_; + // The data entry + std::vector data_entry_; + // The operation lambda on each node + std::vector op_execs_; + // The code module. + tvm::runtime::Module module_; + std::unordered_map name_idx_; +}; + + +struct TVMOpParam : public dmlc::Parameter { + std::string func_name; + uint32_t num_inputs; + uint32_t num_outputs; + bool flatten_data; + DMLC_DECLARE_PARAMETER(TVMOpParam) { + DMLC_DECLARE_FIELD(func_name); + DMLC_DECLARE_FIELD(num_inputs) + .set_default(1); + DMLC_DECLARE_FIELD(num_outputs) + .set_default(1); + DMLC_DECLARE_FIELD(flatten_data) + .set_default(false); + } +}; +} // namespace contrib +} // namespace tvm + +#endif // TVM_GRAPH_EXECUTOR_H_ diff --git a/apps/graph_executor/src/graph_executor_ext.cc b/apps/graph_executor/src/graph_executor_ext.cc new file mode 100644 index 00000000..2ab2a835 --- /dev/null +++ b/apps/graph_executor/src/graph_executor_ext.cc @@ -0,0 +1,87 @@ +/*! + * Copyright (c) 2017 by Contributors + * \file graph_executor_ext.cc + */ +#include "./graph_executor.h" + +namespace tvm { +namespace contrib { + +bool SaveDLTensor(dmlc::Stream* strm, DLTensor* tensor) { + uint64_t header = kTVMNDArrayMagic, reserved = 0; + strm->Write(&header, sizeof(header)); + strm->Write(&reserved, sizeof(reserved)); + + strm->Write(&tensor->ctx, sizeof(tensor->ctx)); + strm->Write(&tensor->ndim, sizeof(tensor->ndim)); + strm->Write(&tensor->dtype, sizeof(tensor->dtype)); + + int ndim = tensor->ndim; + strm->Write(tensor->shape, sizeof(int64_t) * ndim); + + int type_size = tensor->dtype.bits / 8; + int64_t size = 1; + for (int i = 0; i < ndim; ++i) { + size *= tensor->shape[i]; + } + int64_t data_byte_size = type_size * size; + strm->Write(&data_byte_size, sizeof(data_byte_size)); + strm->Write(tensor->data, data_byte_size); + return true; +} + +TVM_REGISTER_GLOBAL("tvm_graph._save_param_dict") +.set_body([](TVMArgs args, TVMRetValue *rv) { + std::string fname = args[0]; + int num_params = args[1]; + std::vector names; + names.reserve(num_params); + std::vector arrays; + arrays.reserve(num_params); + for (int i = 2; i < (2 + 2*num_params); i += 2) { + names.emplace_back(args[i].operator std::string()); + arrays.emplace_back(args[i+1].operator DLTensor*()); + } + + std::unique_ptr fo(dmlc::Stream::Create(fname.c_str(), "w")); + uint64_t header = kTVMNDArrayListMagic, reserved = 0; + fo->Write(&header, sizeof(header)); + fo->Write(&reserved, sizeof(reserved)); + + fo->Write(names); + { + uint64_t sz = static_cast(arrays.size()); + fo->Write(&sz, sizeof(sz)); + for (size_t i = 0; i < sz; ++i) { + SaveDLTensor(fo.get(), arrays[i]); + } + } + }); + +// Create executor +tvm::runtime::Module CreateExecutor(nnvm::Graph g, TVMContext ctx) { + std::shared_ptr exec = + std::make_shared(); + exec->Init(g, ctx); + return tvm::runtime::Module(exec); +} + +TVM_REGISTER_GLOBAL("tvm_graph._create_executor") +.set_body([](TVMArgs args, TVMRetValue *rv) { + void* graph_handle = args[0]; + int device_type = args[1]; + int device_id = args[2]; + TVMContext ctx{static_cast(device_type), device_id}; + nnvm::Graph g = static_cast(graph_handle)[0]; + *rv = CreateExecutor(g, ctx); + }); + + +TVM_REGISTER_GLOBAL("tvm_graph._get_module_from_graph") +.set_body([](TVMArgs args, TVMRetValue *rv) { + void* graph_handle = args[0]; + nnvm::Graph* g = static_cast(graph_handle); + *rv = g->MoveCopyAttr("module"); + }); +} // namespace contrib +} // namespace tvm diff --git a/apps/graph_executor/tests/test_executor.py b/apps/graph_executor/tests/test_executor.py index 3ec60f66..3d89f26c 100644 --- a/apps/graph_executor/tests/test_executor.py +++ b/apps/graph_executor/tests/test_executor.py @@ -17,8 +17,8 @@ def test_compile(): na = tvm.nd.array(np.ones(shape).astype(dtype)) nb = tvm.nd.array(np.ones(shape).astype(dtype)) # set inputs - set_input(0, na) - set_input(1, nb) + set_input('x', na) + set_input('y', nb) # execute run() # get outputs