Commit d27c11e0 by tqchen Committed by Tianqi Chen

[DLL] Use local dll, not de-allocate function in shutdown

parent c6c42af0
...@@ -35,7 +35,7 @@ class NNVMError(Exception): ...@@ -35,7 +35,7 @@ class NNVMError(Exception):
def _load_lib(): def _load_lib():
"""Load libary by searching possible path.""" """Load libary by searching possible path."""
lib_path = libinfo.find_lib_path() lib_path = libinfo.find_lib_path()
lib = ctypes.CDLL(lib_path[0], ctypes.RTLD_GLOBAL) lib = ctypes.CDLL(lib_path[0], ctypes.RTLD_LOCAL)
# DMatrix functions # DMatrix functions
lib.NNGetLastError.restype = ctypes.c_char_p lib.NNGetLastError.restype = ctypes.c_char_p
return lib return lib
......
...@@ -71,11 +71,12 @@ GetAttrDict(const NodeAttrs& attrs) { ...@@ -71,11 +71,12 @@ GetAttrDict(const NodeAttrs& attrs) {
TVM_REGISTER_GLOBAL("nnvm._register_compute") TVM_REGISTER_GLOBAL("nnvm._register_compute")
.set_body([](TVMArgs args, TVMRetValue *rv) { .set_body([](TVMArgs args, TVMRetValue *rv) {
PackedFunc f = args[1]; // Intentionally copy and not de-allocate it, to avoid free pyobject during shutdown
PackedFunc* f = new PackedFunc(args[1].operator PackedFunc());
Op& op = ::dmlc::Registry<nnvm::Op>::Get()->__REGISTER_OR_GET__(args[0]); Op& op = ::dmlc::Registry<nnvm::Op>::Get()->__REGISTER_OR_GET__(args[0]);
auto fcompute = [f](const NodeAttrs& attrs, const Array<Tensor>& inputs) auto fcompute = [f](const NodeAttrs& attrs, const Array<Tensor>& inputs)
-> Array<Tensor> { -> Array<Tensor> {
TVMRetValue ret = f(GetAttrDict(attrs), inputs); TVMRetValue ret = (*f)(GetAttrDict(attrs), inputs);
if ((*ret.ptr<std::shared_ptr<tvm::Node> >())->derived_from<tvm::TensorNode>()) { if ((*ret.ptr<std::shared_ptr<tvm::Node> >())->derived_from<tvm::TensorNode>()) {
return {ret.operator Tensor()}; return {ret.operator Tensor()};
} else { } else {
...@@ -87,12 +88,13 @@ TVM_REGISTER_GLOBAL("nnvm._register_compute") ...@@ -87,12 +88,13 @@ TVM_REGISTER_GLOBAL("nnvm._register_compute")
TVM_REGISTER_GLOBAL("nnvm._register_schedule") TVM_REGISTER_GLOBAL("nnvm._register_schedule")
.set_body([](TVMArgs args, TVMRetValue *rv) { .set_body([](TVMArgs args, TVMRetValue *rv) {
PackedFunc f = args[1]; // Intentionally copy and not de-allocate it, to avoid free pyobject during shutdown
PackedFunc* f = new PackedFunc(args[1].operator PackedFunc());
Op& op = ::dmlc::Registry<nnvm::Op>::Get()->__REGISTER_OR_GET__(args[0]); Op& op = ::dmlc::Registry<nnvm::Op>::Get()->__REGISTER_OR_GET__(args[0]);
auto fschedule = [f](const NodeAttrs& attrs, auto fschedule = [f](const NodeAttrs& attrs,
const Array<Tensor>& outs, const Array<Tensor>& outs,
const std::string& target) { const std::string& target) {
return f(GetAttrDict(attrs), outs, target).operator Schedule(); return (*f)(GetAttrDict(attrs), outs, target).operator Schedule();
}; };
op.set_attr<FTVMSchedule>("FTVMSchedule", fschedule, args[2]); op.set_attr<FTVMSchedule>("FTVMSchedule", fschedule, args[2]);
}); });
......
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