From adc49de803daf56718b4370f80b0123df373ac8c Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Wed, 6 Jan 2021 16:41:04 +0800 Subject: [PATCH] feat(mge/imperative): expose c++ tensor reference count GitOrigin-RevId: 1940881adc243bd0ec01bf104d50b01f31639bd6 --- imperative/python/src/tensor.cpp | 5 ++++- imperative/python/src/tensor.h | 1 + 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/imperative/python/src/tensor.cpp b/imperative/python/src/tensor.cpp index a6fecc61e..e8484d9a6 100644 --- a/imperative/python/src/tensor.cpp +++ b/imperative/python/src/tensor.cpp @@ -541,6 +541,7 @@ struct TensorWeakRef { } return py::none(); } + int _use_cnt() { return wptr.use_count(); } }; /* ============== convert inputs ============== */ @@ -774,6 +775,7 @@ void init_tensor(py::module m) { .def<&TensorWrapper::_swap_in>("_swap_in") .def<&TensorWrapper::_drop>("_drop") .def<&TensorWrapper::reset_varnode>("_reset_varnode") + .def<&TensorWrapper::_use_cnt>("_use_cnt") .def_getset<&TensorWrapper::varnode>("_varnode") .def_getset<&TensorWrapper::copied>("_copied") .def_getset<&TensorWrapper::mixin_handle, &TensorWrapper::set_mixin_handle>("_mixin_handle") @@ -787,7 +789,8 @@ void init_tensor(py::module m) { py::class_(m, "TensorWeakRef") .def(py::init()) - .def("__call__", &TensorWeakRef::operator()); + .def("__call__", &TensorWeakRef::operator()) + .def("_use_cnt", &TensorWeakRef::_use_cnt); static PyMethodDef method_defs[] = { MGE_PY_INTERFACE(apply, py_apply), diff --git a/imperative/python/src/tensor.h b/imperative/python/src/tensor.h index 2cd721329..84ac2efb7 100644 --- a/imperative/python/src/tensor.h +++ b/imperative/python/src/tensor.h @@ -170,6 +170,7 @@ struct TensorWrapper { void set_compiled_info(PyObject *); PyObject* trace_mixin_info(); void set_trace_mixin_info(PyObject *); + PyObject* _use_cnt() { return PyLong_FromSize_t(m_tensor.use_count()); }; }; -- GitLab