gpu dynamic memory pool can not reuse allReduce in multi-stream

5 years ago · 0f0e8fe874
parent 4a3e5cb944
commit 0f0e8fe874
1 changed files with 3 additions and 0 deletions
--- a/mindspore/ccsrc/device/gpu/gpu_kernel_runtime.cc
+++ b/mindspore/ccsrc/device/gpu/gpu_kernel_runtime.cc
@ -322,6 +322,9 @@ void GPUKernelRuntime::FreeKernelDynamicRes(const mindspore::AnfNodePtr &kernel,
  MS_EXCEPTION_IF_NULL(mem_reuse_util_ptr);
  auto cnode = kernel->cast<CNodePtr>();
  MS_EXCEPTION_IF_NULL(cnode);
+  if (AnfAlgo::GetCNodeName(kernel) == kAllReduceOpName) {
+    return;
+  }
  // Free the input of kernel by reference count.
  for (size_t i = 0; i < AnfAlgo::GetInputTensorNum(kernel); ++i) {
    auto kernel_ref_count_ptr = mem_reuse_util_ptr->GetKernelInputRef(cnode, i);