提交 95c378e4 编写于 作者: M mindspore-ci-bot 提交者: Gitee

!2844 optimize cpu input mem alloc

Merge pull request !2844 from kisnwang/optimize-cpu-mem
...@@ -27,7 +27,12 @@ void CPUSimpleMemPlan::MemPlan(const session::KernelGraph *graph) { ...@@ -27,7 +27,12 @@ void CPUSimpleMemPlan::MemPlan(const session::KernelGraph *graph) {
MS_EXCEPTION_IF_NULL(kernel); MS_EXCEPTION_IF_NULL(kernel);
size_t input_num = AnfAlgo::GetInputTensorNum(kernel); size_t input_num = AnfAlgo::GetInputTensorNum(kernel);
for (size_t i = 0; i < input_num; ++i) { 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<Parameter>()) {
continue;
}
auto address = AnfAlgo::GetOutputAddr(kernel_with_index.first, kernel_with_index.second, true);
MS_EXCEPTION_IF_NULL(address); MS_EXCEPTION_IF_NULL(address);
if (address->ptr_ == nullptr) { if (address->ptr_ == nullptr) {
total_mem_size += address->size_; total_mem_size += address->size_;
...@@ -73,7 +78,12 @@ void CPUSimpleMemPlan::MemAssign(const session::KernelGraph *graph, uint8_t *bas ...@@ -73,7 +78,12 @@ void CPUSimpleMemPlan::MemAssign(const session::KernelGraph *graph, uint8_t *bas
MS_EXCEPTION_IF_NULL(kernel); MS_EXCEPTION_IF_NULL(kernel);
size_t input_num = AnfAlgo::GetInputTensorNum(kernel); size_t input_num = AnfAlgo::GetInputTensorNum(kernel);
for (size_t i = 0; i < input_num; ++i) { 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<Parameter>()) {
continue;
}
auto address = AnfAlgo::GetMutableOutputAddr(kernel_with_index.first, kernel_with_index.second, true);
MS_EXCEPTION_IF_NULL(address); MS_EXCEPTION_IF_NULL(address);
if (address->ptr_ == nullptr) { if (address->ptr_ == nullptr) {
address->ptr_ = mem_ptr; address->ptr_ = mem_ptr;
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册