Commit 60da4705 by Tang, Cheng Committed by Tianqi Chen

add tvm external registry entrypoint (#1562)

parent 1ed28aeb
...@@ -83,5 +83,11 @@ struct NodeFactoryReg ...@@ -83,5 +83,11 @@ struct NodeFactoryReg
::dmlc::Registry<::tvm::NodeFactoryReg>::Get()->__REGISTER__(TypeName::_type_key) \ ::dmlc::Registry<::tvm::NodeFactoryReg>::Get()->__REGISTER__(TypeName::_type_key) \
.set_body([]() { return std::make_shared<TypeName>(); }) .set_body([]() { return std::make_shared<TypeName>(); })
TVM_DLL::dmlc::Registry<::tvm::NodeFactoryReg > * GetTVMNodeFactoryRegistry();
#define TVM_EXTERNAL_REGISTER_NODE_TYPE(TypeName) \
static DMLC_ATTRIBUTE_UNUSED ::tvm::NodeFactoryReg & __make_Node ## _ ## TypeName ## __ = \
::tvm::GetTVMNodeFactoryRegistry()->__REGISTER__(TypeName::_type_key) \
.set_body([]() { return std::make_shared<TypeName>(); })
} // namespace tvm } // namespace tvm
#endif // TVM_BASE_H_ #endif // TVM_BASE_H_
...@@ -100,6 +100,6 @@ TVM_REGISTER_GLOBAL("nnvm.compiler._load_param_dict") ...@@ -100,6 +100,6 @@ TVM_REGISTER_GLOBAL("nnvm.compiler._load_param_dict")
*rv = ret; *rv = ret;
}); });
TVM_REGISTER_NODE_TYPE(NDArrayWrapperNode); TVM_EXTERNAL_REGISTER_NODE_TYPE(NDArrayWrapperNode);
} // namespace compiler } // namespace compiler
} // namespace nnvm } // namespace nnvm
...@@ -41,4 +41,8 @@ TVM_REGISTER_API("_TVMSetStream") ...@@ -41,4 +41,8 @@ TVM_REGISTER_API("_TVMSetStream")
.set_body([](TVMArgs args, TVMRetValue *ret) { .set_body([](TVMArgs args, TVMRetValue *ret) {
TVMSetStream(args[0], args[1], args[2]); TVMSetStream(args[0], args[1], args[2]);
}); });
TVM_DLL::dmlc::Registry<::tvm::NodeFactoryReg > * GetTVMNodeFactoryRegistry() {
return ::dmlc::Registry<::tvm::NodeFactoryReg>::Get();
}
} // namespace tvm } // namespace tvm
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment