From e53b41bad8e1a037e843936655bd207726f8d6bc Mon Sep 17 00:00:00 2001 From: kswang Date: Fri, 3 Jul 2020 15:25:23 +0800 Subject: [PATCH] optimize cpu input mem --- mindspore/ccsrc/device/cpu/cpu_simple_mem_plan.cc | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/mindspore/ccsrc/device/cpu/cpu_simple_mem_plan.cc b/mindspore/ccsrc/device/cpu/cpu_simple_mem_plan.cc index 95f0d25f5b..e6cb6ee53a 100644 --- a/mindspore/ccsrc/device/cpu/cpu_simple_mem_plan.cc +++ b/mindspore/ccsrc/device/cpu/cpu_simple_mem_plan.cc @@ -27,7 +27,12 @@ void CPUSimpleMemPlan::MemPlan(const session::KernelGraph *graph) { MS_EXCEPTION_IF_NULL(kernel); size_t input_num = AnfAlgo::GetInputTensorNum(kernel); for (size_t i = 0; i < input_num; ++i) { - auto address = AnfAlgo::GetPrevNodeOutputAddr(kernel, i); + auto kernel_with_index = AnfAlgo::GetPrevNodeOutput(kernel, i); + MS_EXCEPTION_IF_NULL(kernel_with_index.first); + if (kernel_with_index.first->isa()) { + continue; + } + auto address = AnfAlgo::GetOutputAddr(kernel_with_index.first, kernel_with_index.second, true); MS_EXCEPTION_IF_NULL(address); if (address->ptr_ == nullptr) { total_mem_size += address->size_; @@ -73,7 +78,12 @@ void CPUSimpleMemPlan::MemAssign(const session::KernelGraph *graph, uint8_t *bas MS_EXCEPTION_IF_NULL(kernel); size_t input_num = AnfAlgo::GetInputTensorNum(kernel); for (size_t i = 0; i < input_num; ++i) { - auto address = AnfAlgo::GetPrevNodeMutableOutputAddr(kernel, i); + auto kernel_with_index = AnfAlgo::GetPrevNodeOutput(kernel, i); + MS_EXCEPTION_IF_NULL(kernel_with_index.first); + if (kernel_with_index.first->isa()) { + continue; + } + auto address = AnfAlgo::GetMutableOutputAddr(kernel_with_index.first, kernel_with_index.second, true); MS_EXCEPTION_IF_NULL(address); if (address->ptr_ == nullptr) { address->ptr_ = mem_ptr;