From b570dec7abd5feac581fecef10fd410aa927518c Mon Sep 17 00:00:00 2001 From: laiyongqiang Date: Sat, 18 Jul 2020 09:38:12 +0800 Subject: [PATCH] add right align border for communication op's single output --- .../backend/optimizer/mem_reuse/mem_reuse_allocator.cc | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/mindspore/ccsrc/backend/optimizer/mem_reuse/mem_reuse_allocator.cc b/mindspore/ccsrc/backend/optimizer/mem_reuse/mem_reuse_allocator.cc index d10d8cd949..f57a78863a 100644 --- a/mindspore/ccsrc/backend/optimizer/mem_reuse/mem_reuse_allocator.cc +++ b/mindspore/ccsrc/backend/optimizer/mem_reuse/mem_reuse_allocator.cc @@ -194,13 +194,20 @@ void BestFitMemReuse::AssignCommunicationNodeOutputOffset() { // add left align border for the first output and right align border for the last output to alloc align border memory size_t output_index = 0; - for (auto &tensor_idx : current_kernel_->GetOutputRefIndexs()) { + auto output_ref_indexes = current_kernel_->GetOutputRefIndexs(); + for (auto &tensor_idx : output_ref_indexes) { size_t index = GetTensorIndex(tensor_idx); auto tensor_desc = tensor_ptr_list_[index]; MS_EXCEPTION_IF_NULL(tensor_desc); if (output_index == 0 || output_index == output_num - 1) { tensor_desc->size_ += kDefaultMemAlignSize; } + + if ((output_index == 0) && (output_ref_indexes.size() == 1)) { + // add right align border for single output + tensor_desc->size_ += kDefaultMemAlignSize; + } + output_index++; }