!10552 fix pynative device memory leak

From: @chujinjin
Reviewed-by: 
Signed-off-by:

mindspore/ccsrc/pipeline/pynative/pynative_execute.cc

@@ -1092,6 +1092,25 @@ void PynativeExecutor::CleanPreMemoryInValueNode(const std::string &cell_id) {
       tensor->set_device_address(nullptr);
     }
   }
+
+  if (dynamic_cell_) {
+    std::set<std::string> forward_op_tensor_id;
+    for (const auto &elem : cell_op_index_with_tensor_id_[top_cell_id_]) {
+      const auto &tensor_id_list = elem.second;
+      for (const auto &tensor_id : tensor_id_list) {
+        forward_op_tensor_id.emplace(tensor_id);
+      }
+    }
+    for (auto &tensor : all_value_node_tensors_) {
+      if (tensor->device_address() != nullptr &&
+          forward_op_tensor_id.find(tensor->id()) != forward_op_tensor_id.end()) {
+        tensor->device_address()->ClearDeviceMemory();
+        tensor->set_device_address(nullptr);
+      }
+    }
+    all_value_node_tensors_.clear();
+  }
+
   top_cell_id_ = cell_id;
 }
 
@@ -2086,6 +2105,37 @@ std::string PynativeExecutor::GetGradCellId(bool has_sens, const py::object &cel
   return cell_id;
 }
 
+void PynativeExecutor::SaveAllValueNodeTensors(const FuncGraphPtr &graph) {
+  std::unordered_set<tensor::TensorPtr> all_value_node_tensors;
+  auto trace_function = [&all_value_node_tensors](const AnfNodePtr &anf_node) {
+    auto value = GetValueNode(anf_node);
+    if (value) {
+      if (value->isa<tensor::Tensor>()) {
+        auto tensor = value->cast<tensor::TensorPtr>();
+        MS_EXCEPTION_IF_NULL(tensor);
+        if (tensor->device_address()) {
+          all_value_node_tensors.emplace(tensor);
+        }
+      } else if (value->isa<ValueTuple>()) {
+        auto tuple = value->cast<ValueTuplePtr>();
+        MS_EXCEPTION_IF_NULL(tuple);
+        for (size_t i = 0; i < tuple->size(); i++) {
+          if ((*tuple)[i]->isa<tensor::Tensor>()) {
+            auto tensor = (*tuple)[i]->cast<tensor::TensorPtr>();
+            MS_EXCEPTION_IF_NULL(tensor);
+            if (tensor->device_address()) {
+              all_value_node_tensors.emplace(tensor);
+            }
+          }
+        }
+      }
+    }
+    return FOLLOW;
+  };
+  (void)TopoSort(graph->get_return(), SuccDeeperSimple, trace_function);
+  all_value_node_tensors_ = all_value_node_tensors;
+}
+
 void PynativeExecutor::GradNetInner(const GradOperationPtr &grad, const py::object &cell, const py::object &weights,
                                     const py::args &args) {
   auto size = args.size();
@@ -2127,6 +2177,9 @@ void PynativeExecutor::GradNetInner(const GradOperationPtr &grad, const py::obje
   resource->results()[pipeline::kBackend] = compile::CreateBackend();
 
   MS_LOG(INFO) << "Start opt";
+  if (dynamic_cell_) {
+    SaveAllValueNodeTensors(resource->func_graph());
+  }
   PynativeOptimizeAction(resource);
   SaveTensorsInValueNode(resource);
   TaskEmitAction(resource);

mindspore/ccsrc/pipeline/pynative/pynative_execute.h

@@ -198,6 +198,7 @@ class PynativeExecutor : public std::enable_shared_from_this<PynativeExecutor> {
   // Update the abstract and device address info of value node and tensors in bprop graph
   void UpdateAbstractAndDeviceAddress(const OpExecInfoPtr &op_exec_info, const py::object &out_real);
   void SaveTensorsInValueNode(const ResourcePtr &resource);
+  void SaveAllValueNodeTensors(const FuncGraphPtr &graph);
   void CleanPreMemoryInValueNode(const std::string &cell_id);
 
   // Construct grad graph
@@ -304,6 +305,7 @@ class PynativeExecutor : public std::enable_shared_from_this<PynativeExecutor> {
   std::unordered_map<std::string, TensorIdWithTensor> cell_tensor_id_with_tensor_;
   std::unordered_map<std::string, abstract::AbstractBasePtr> node_abs_map_;
   std::unordered_map<std::string, AbstractListMap> prim_abs_list_;
+  std::unordered_set<tensor::TensorPtr> all_value_node_tensors_;
 };
 
 using PynativeExecutorPtr = std::shared_ptr<PynativeExecutor>;

mindspore/ccsrc/runtime/device/ascend/ascend_device_address.cc

@@ -612,7 +612,7 @@ bool AscendDeviceAddress::ConvertFormatAndSyncHostToDevice(const ShapeVector &sh
   return sync_ok;
 }
 
-AscendDeviceAddress::~AscendDeviceAddress() {
+void AscendDeviceAddress::ClearDeviceMemory() {
   if (ptr_ == nullptr) {
     return;
   }
@@ -627,6 +627,8 @@ AscendDeviceAddress::~AscendDeviceAddress() {
   }
 }
 
+AscendDeviceAddress::~AscendDeviceAddress() { ClearDeviceMemory(); }
+
 bool AscendDeviceAddress::DumpMemToFile(bool trans_flag, const std::string &filepath, const std::string &host_fmt,
                                         const ShapeVector &host_shape, TypeId host_type) const {
   bool ret = false;

mindspore/ccsrc/runtime/device/ascend/ascend_device_address.h

@@ -41,6 +41,7 @@ class AscendDeviceAddress : public DeviceAddress {
   ~AscendDeviceAddress() override;
   bool SyncDeviceToHost(const ShapeVector &shape, size_t size, TypeId type, void *host_ptr) const override;
   bool SyncHostToDevice(const ShapeVector &shape, size_t size, TypeId type, const void *host_ptr) const override;
+  void ClearDeviceMemory() override;
   DeviceAddressType DeviceType() const override { return DeviceAddressType::kAscend; }
   bool DumpMemToFile(bool dump_mode, const std::string &filepath, const std::string &host_fmt,
                      const ShapeVector &host_shape, TypeId host_type) const override;

mindspore/ccsrc/runtime/device/cpu/cpu_device_address.h

@@ -35,6 +35,7 @@ class CPUDeviceAddress : public DeviceAddress {
 
   bool SyncDeviceToHost(const ShapeVector &shape, size_t size, TypeId type, void *host_ptr) const override;
   bool SyncHostToDevice(const ShapeVector &shape, size_t size, TypeId type, const void *host_ptr) const override;
+  void ClearDeviceMemory() override {}
   DeviceAddressType DeviceType() const override { return DeviceAddressType::kCPU; }
 };
 }  // namespace cpu

mindspore/ccsrc/runtime/device/gpu/gpu_device_address.cc

@@ -69,7 +69,7 @@ bool GPUDeviceAddress::SyncHostToDevice(const ShapeVector &, size_t size, TypeId
   return GPUDeviceManager::GetInstance().SyncStream(stream);
 }
 
-GPUDeviceAddress::~GPUDeviceAddress() {
+void GPUDeviceAddress::ClearDeviceMemory() {
   if (ptr_ == nullptr) {
     return;
   }
@@ -78,6 +78,8 @@ GPUDeviceAddress::~GPUDeviceAddress() {
     ptr_ = nullptr;
   }
 }
+
+GPUDeviceAddress::~GPUDeviceAddress() { ClearDeviceMemory(); }
 #ifdef ENABLE_DEBUGGER
 bool GPUDeviceAddress::LoadMemToHost(const std::string &tensor_name, int execution_order, const std::string &host_fmt,
                                      const ShapeVector &host_shape, TypeId host_type, size_t slot,

mindspore/ccsrc/runtime/device/gpu/gpu_device_address.h

@@ -38,6 +38,7 @@ class GPUDeviceAddress : public DeviceAddress {
 
   bool SyncDeviceToHost(const ShapeVector &shape, size_t size, TypeId type, void *host_ptr) const override;
   bool SyncHostToDevice(const ShapeVector &shape, size_t size, TypeId type, const void *host_ptr) const override;
+  void ClearDeviceMemory() override;
   void set_status(DeviceAddressStatus status) { status_ = status; }
   DeviceAddressStatus status() const { return status_; }
   DeviceAddressType DeviceType() const override { return DeviceAddressType::kGPU; }

mindspore/core/ir/device_sync.h

@@ -33,6 +33,7 @@ class DeviceSync {
   virtual bool SyncDeviceToHost(const ShapeVector &shape, size_t size, TypeId type, void *host_ptr) const = 0;
   virtual bool SyncHostToDevice(const ShapeVector &shape, size_t size, TypeId type, const void *host_ptr) const = 0;
   virtual void *GetMutablePtr() const = 0;
+  virtual void ClearDeviceMemory() = 0;
 };
 using DeviceSyncPtr = std::shared_ptr<DeviceSync>;
 }  // namespace mindspore