diff --git a/imperative/python/megengine/__init__.py b/imperative/python/megengine/__init__.py index 219fd5ef3bd4e288ecc6546e1aa510228fc94d3f..d56be64e9bbda3ba7fb8a98e8942a6f521c8d176 100644 --- a/imperative/python/megengine/__init__.py +++ b/imperative/python/megengine/__init__.py @@ -71,7 +71,9 @@ if sys.platform == "win32": kernel32.SetErrorMode(old_error_mode) -from .core._imperative_rt.core2 import release_trace_apply_func, sync, full_sync +from .core._imperative_rt.core2 import full_sync as _full_sync +from .core._imperative_rt.core2 import release_trace_apply_func +from .core._imperative_rt.core2 import sync as _sync from .core._imperative_rt.utils import _set_fork_exec_path_for_timed_func from .device import * from .logger import enable_debug_log, get_logger, set_log_file, set_log_level @@ -89,10 +91,9 @@ _set_fork_exec_path_for_timed_func( _persistent_cache_impl_ins = persistent_cache.PersistentCacheOnServer() _persistent_cache_impl_ins.reg() -atexit.register(sync) +atexit.register(_full_sync) atexit.register(release_trace_apply_func) -del sync del release_trace_apply_func del _set_fork_exec_path_for_timed_func del _persistent_cache_impl_ins diff --git a/imperative/src/impl/interpreter_impl.cpp b/imperative/src/impl/interpreter_impl.cpp index 66de1a0fbf6ce9fefc68cd81a4207c3676000f05..9c9c3ce55658214f2314df9aa1f92b2e7df0f54e 100644 --- a/imperative/src/impl/interpreter_impl.cpp +++ b/imperative/src/impl/interpreter_impl.cpp @@ -24,6 +24,7 @@ std::unique_ptr InterpreterImpl::create_channel() { } Interpreter& Interpreter::inst() { + Tensor::_static_init(); static InterpreterImpl inst_; return inst_; } diff --git a/imperative/src/impl/physical_tensor.cpp b/imperative/src/impl/physical_tensor.cpp index 13faadab9a5194f1ba5b80d2c71300045cdbe864..7b09508438cd5393b5237d5d0b24c4fe491855c2 100644 --- a/imperative/src/impl/physical_tensor.cpp +++ b/imperative/src/impl/physical_tensor.cpp @@ -56,6 +56,10 @@ protected: return {}; } + AsyncReleaser() { + EventPool::without_timer(); + } + public: static AsyncReleaser* inst() { static AsyncReleaser releaser; @@ -367,6 +371,10 @@ CompNode::Event* Tensor::get_or_create_event() { return e; } +void Tensor::_static_init() { + EventPool::without_timer(); +} + } // namespace imperative } // namespace mgb diff --git a/imperative/src/include/megbrain/imperative/physical_tensor.h b/imperative/src/include/megbrain/imperative/physical_tensor.h index 86b712929d7d62dee417c37603a438d86e7a628e..7485cae2324dd6b2caffac4dfb24ad667b6bb016 100644 --- a/imperative/src/include/megbrain/imperative/physical_tensor.h +++ b/imperative/src/include/megbrain/imperative/physical_tensor.h @@ -130,6 +130,12 @@ public: void add_release_callback(CompNode cn); CompNode::Event* get_or_create_event(); + + // Make sure all static objects required to destruct a tensor has completed + // construction. All static storage duration object that holds tensors must + // call this method before their constructors completes. + static void _static_init(); + private: TensorLayout m_layout;