提交 b570dec7 编写于 作者: L laiyongqiang

add right align border for communication op's single output

上级 a84092e8
......@@ -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++;
}
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册