From 09af925fb65dd9a714d4e09fc02e0f8b699585f9 Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Tue, 5 Jan 2021 15:56:16 +0800 Subject: [PATCH] fix(mge): fix cpp trace function release GitOrigin-RevId: 924f945c211bc17596710410e616ab4b1e2e612e --- imperative/python/megengine/__init__.py | 3 --- imperative/python/src/tensor.cpp | 22 +++++++-------------- imperative/python/src/tensor.h | 4 ++-- imperative/python/src/trace.cpp | 26 +++++++++++++++---------- 4 files changed, 25 insertions(+), 30 deletions(-) diff --git a/imperative/python/megengine/__init__.py b/imperative/python/megengine/__init__.py index d56be64e9..aac03103a 100644 --- a/imperative/python/megengine/__init__.py +++ b/imperative/python/megengine/__init__.py @@ -72,7 +72,6 @@ if sys.platform == "win32": kernel32.SetErrorMode(old_error_mode) 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 * @@ -92,9 +91,7 @@ _persistent_cache_impl_ins = persistent_cache.PersistentCacheOnServer() _persistent_cache_impl_ins.reg() atexit.register(_full_sync) -atexit.register(release_trace_apply_func) -del release_trace_apply_func del _set_fork_exec_path_for_timed_func del _persistent_cache_impl_ins diff --git a/imperative/python/src/tensor.cpp b/imperative/python/src/tensor.cpp index a93031777..d22ca77e2 100644 --- a/imperative/python/src/tensor.cpp +++ b/imperative/python/src/tensor.cpp @@ -34,22 +34,15 @@ namespace mgb::imperative::python { std::unique_ptr interpreter_for_py; -py::object cpp_apply_with_tracing, cpp_apply_const_with_tracing, - cpp_apply_compiled_mode, cpp_apply_const_compiled_mode; +PyObject *cpp_apply_with_tracing, *cpp_apply_const_with_tracing, + *cpp_apply_compiled_mode, *cpp_apply_const_compiled_mode; -py::object cpp_apply_backward_varnode; +PyObject *cpp_apply_backward_varnode; -void release_trace_apply_func(){ - cpp_apply_with_tracing.release(); - cpp_apply_const_with_tracing.release(); - cpp_apply_compiled_mode.release(); - cpp_apply_const_compiled_mode.release(); - cpp_apply_backward_varnode.release(); -} #define REGISTE_APPLY_FUNC(mode) \ void set_##mode(py::object pyf) { \ - mode = pybind11::reinterpret_steal(pyf); \ + mode = pyf.ptr(); \ } REGISTE_APPLY_FUNC(cpp_apply_with_tracing) @@ -242,14 +235,15 @@ TensorWrapper::TensorWrapper(PyObject* args, PyObject* kwargs) { // const op if (is_const && is_tracing) { - py::object pyf; + PyObject *pyf; if (is_compiled) { pyf = cpp_apply_const_compiled_mode; } else { pyf = cpp_apply_const_with_tracing; } - auto ret = pyf(*tup); + auto ret = py::reinterpret_steal( + PyObject_Call(pyf, tup.ptr(), nullptr)); auto py_ret = py::reinterpret_borrow(ret); if (auto* t = try_cast(py_ret[0].ptr())) { m_tensor = t->m_tensor; @@ -744,8 +738,6 @@ void init_tensor(py::module m) { }, py::call_guard()); - m.def("release_trace_apply_func", &release_trace_apply_func); - py::handle grad_key_type = GradKeyWrapper::wrap_t::type() .def<&GradKeyWrapper::attach>("attach") .def<&GradKeyWrapper::is_attached_to>("is_attached_to") diff --git a/imperative/python/src/tensor.h b/imperative/python/src/tensor.h index c8b5b846c..213d0df08 100644 --- a/imperative/python/src/tensor.h +++ b/imperative/python/src/tensor.h @@ -253,8 +253,8 @@ auto apply(std::shared_ptr op, T&& tensors) void init_tensor(pybind11::module); -extern pybind11::object cpp_apply_with_tracing, cpp_apply_compiled_mode; -extern pybind11::object cpp_apply_backward_varnode; +extern PyObject *cpp_apply_with_tracing, *cpp_apply_compiled_mode; +extern PyObject *cpp_apply_backward_varnode; } // namespace mgb::imperative::python diff --git a/imperative/python/src/trace.cpp b/imperative/python/src/trace.cpp index b21729db3..ff37a6035 100644 --- a/imperative/python/src/trace.cpp +++ b/imperative/python/src/trace.cpp @@ -6,7 +6,8 @@ * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or + * implied. */ #include "./trace.h" @@ -23,12 +24,13 @@ apply_result_t apply_trace(ApplyContext& ctx) { if (ctx.backward) { // reach here when symbolic=True or compiled=True // call megbrain_graph.py apply(BackwardGraph, *args) - auto args = py::tuple(ctx.nargs); + auto args = py::tuple(ctx.nargs + 1); + args[0] = py::cast(ctx.op); for (size_t i = 0; i < ctx.nargs; i++) { - args[i] = py::cast(ctx.args[i]->m_var); + args[i + 1] = py::cast(ctx.args[i]->m_var); } - py::object ret = cpp_apply_backward_varnode(py::cast(ctx.op), *args); - + py::object ret = py::reinterpret_steal( + PyObject_Call(cpp_apply_backward_varnode, args.ptr(), nullptr)); if (!ret) { throw py::value_error("invalid py object call"); } @@ -36,13 +38,13 @@ apply_result_t apply_trace(ApplyContext& ctx) { // assumption: python function always returns PyList auto tup = py::reinterpret_borrow(ret); for (auto i = 0; i < tup.size(); i++) { - auto pitem = tup[i].cast(); + auto pitem = tup[i].cast(); outputs.emplace_back(std::make_shared(pitem)); } return outputs; } - py::object pyf; + PyObject* pyf; if (is_compiled) { // run apply in compiled mode, step 2, 3, etc pyf = cpp_apply_compiled_mode; @@ -51,11 +53,15 @@ apply_result_t apply_trace(ApplyContext& ctx) { pyf = cpp_apply_with_tracing; } - auto args = py::tuple(ctx.nargs); + auto args = py::tuple(ctx.nargs + 1); + args[0] = py::cast(ctx.op); for (size_t i = 0; i < ctx.nargs; i++) { - args[i] = TensorWrapper::make(std::move(std::shared_ptr(ctx.args[i]))).release(); + args[i + 1] = TensorWrapper::make( + std::move(std::shared_ptr(ctx.args[i]))) + .release(); } - auto ret = pyf(py::cast(ctx.op), *args); + auto ret = py::reinterpret_steal( + PyObject_Call(pyf, args.ptr(), nullptr)); // assumption: python function always returns PyList auto tup = py::reinterpret_borrow(ret); -- GitLab