p47068915
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

add right align border for communication op's single output

This commit is contained in:
laiyongqiang 2020-07-18 09:38:12 +08:00
parent a84092e89d
commit b570dec7ab
1 changed files with 8 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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