From 23b9a98f5efa23581df6421d01bc2a1f4719a850 Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Fri, 8 Jan 2021 12:56:11 +0800
Subject: [PATCH] fix(mge): fix sublnear cuda and mem leak

GitOrigin-RevId: 82091ec9a6d9a00ecffad9505b59d54e3127c783
---
 imperative/python/megengine/jit/tracing.py |  9 +++++++--
 imperative/python/src/tensor.cpp           | 11 +++++++++--
 imperative/python/src/trace.cpp            |  5 ++---
 imperative/python/src/trace_info.h         |  6 +++---
 4 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/imperative/python/megengine/jit/tracing.py b/imperative/python/megengine/jit/tracing.py
index 4f0ec56d0..85acc0dc1 100644
--- a/imperative/python/megengine/jit/tracing.py
+++ b/imperative/python/megengine/jit/tracing.py
@@ -350,12 +350,16 @@ class trace:
             self._lazy_eval_links = ()
 
     def _take_escaped_tensors(self):
-        escaped_tensors = tuple(filter(lambda x: x() is not None, self._active_tensors.values()))
+        escaped_tensors = tuple(
+            filter(lambda x: x() is not None, self._active_tensors.values())
+        )
         self._active_tensors.clear()
         return escaped_tensors
 
     def _lazy_eval(self, lazy_eval_graph, lazy_eval_tensors, lazy_eval_links):
-        lazy_eval_tensors = list(filter(lambda x: x() is not None, lazy_eval_tensors.values()))
+        lazy_eval_tensors = list(
+            filter(lambda x: x() is not None, lazy_eval_tensors.values())
+        )
         readers = [G.OutputNode(x()._varnode).outputs[0] for x in lazy_eval_tensors]
         self._apply_graph_options(lazy_eval_graph)
         # FIXME
@@ -443,6 +447,7 @@ class trace:
                         x()._reset_varnode()
                         x().mixin_handle = -1
                         x().recording = False
+                        x()._trace_mixin_info = None
 
         try:
             do_enter()
diff --git a/imperative/python/src/tensor.cpp b/imperative/python/src/tensor.cpp
index 40746e335..558e6229a 100644
--- a/imperative/python/src/tensor.cpp
+++ b/imperative/python/src/tensor.cpp
@@ -294,8 +294,13 @@ PyObject* TensorWrapper::copied() {
             return m_tensor->m_trace_info.member;                                   \
         }                                                                           \
         void TensorWrapper::set_##member(PyObject* dest) {                          \
-            Py_INCREF(dest);                                                        \
-            m_tensor->m_trace_info.member = dest;                                   \
+            if (dest == Py_None) {                                                  \
+                Py_XDECREF(m_tensor->m_trace_info.member);                          \
+                m_tensor->m_trace_info.member = nullptr;                            \
+            } else {                                                                \
+                Py_INCREF(dest);                                                    \
+                m_tensor->m_trace_info.member = dest;                               \
+            }                                                                       \
         }
 
 REGISTE_TENSORWRAPPER_PYOBJECT_FUNC(compiled_info)
@@ -463,6 +468,8 @@ PyObject* TensorWrapper::_dev_tensor(){
         auto py_dev_tensor = py::reinterpret_borrow<py::object>(dev_tensor);
         auto sh = interpreter_for_py->put(py_dev_tensor.cast<DeviceTensorND>());
         m_tensor->m_handle = std::move(SharedHandle(sh));
+        Py_DECREF(m_tensor->m_trace_info.compiled_info);
+        m_tensor->m_trace_info.compiled_info = nullptr;
 
         return dev_tensor;
     }
diff --git a/imperative/python/src/trace.cpp b/imperative/python/src/trace.cpp
index 73de092bf..09d1c1753 100644
--- a/imperative/python/src/trace.cpp
+++ b/imperative/python/src/trace.cpp
@@ -55,10 +55,9 @@ apply_result_t apply_trace(ApplyContext& ctx) {
 
     auto args = py::tuple(ctx.nargs + 1);
     args[0] = py::cast(ctx.op);
+    py::tuple args(ctx.nargs);
     for (size_t i = 0; i < ctx.nargs; i++) {
-        args[i + 1] = TensorWrapper::make(
-                              std::move(std::shared_ptr<Tensor>(ctx.args[i])))
-                              .release();
+        args[i + 1] = TensorWrapper::make(ctx.args[i]->shared_from_this());
     }
     auto ret = py::reinterpret_steal<py::object>(
             PyObject_Call(pyf, args.ptr(), nullptr));
diff --git a/imperative/python/src/trace_info.h b/imperative/python/src/trace_info.h
index b7c9a0f6a..d61d9024c 100644
--- a/imperative/python/src/trace_info.h
+++ b/imperative/python/src/trace_info.h
@@ -28,10 +28,10 @@ struct TraceInfo {
         mixin_handle = that.mixin_handle;
         recording = that.recording;
 
-        compiled_info = that.compiled_info;
-        Py_XINCREF(compiled_info);
         trace_mixin_info = that.trace_mixin_info;
         Py_XINCREF(trace_mixin_info);
+        compiled_info = that.compiled_info;
+        Py_XINCREF(compiled_info);
 
         copied = true;
         return *this;
@@ -39,7 +39,7 @@ struct TraceInfo {
 
     ~TraceInfo()  {
         Py_XDECREF(trace_mixin_info);
-        // Py_XDECREF(compiled_info);
+        Py_XDECREF(compiled_info);
     }
 
 private:
-- 
GitLab