Browse Source

add right align border for communication op's single output

tags/v0.6.0-beta
laiyongqiang 5 years ago
parent
commit
b570dec7ab
1 changed files with 8 additions and 1 deletions
  1. +8
    -1
      mindspore/ccsrc/backend/optimizer/mem_reuse/mem_reuse_allocator.cc

+ 8
- 1
mindspore/ccsrc/backend/optimizer/mem_reuse/mem_reuse_allocator.cc View File

@@ -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++;
}



Loading…
Cancel
Save