From bc1fbcac89fb71ec3ba83b742ac39945910fbd69 Mon Sep 17 00:00:00 2001
From: LaiYongqiang <laiyongqiang1@huawei.com>
Date: Wed, 16 Feb 2022 22:13:38 +0800
Subject: [PATCH] AllReduce input and output size aligned by 512 in ascend
 device context

---
 .../ascend/hal/hardware/ascend_device_context.cc   | 14 +++++++++++++-
 .../graph_scheduler/actor/data_prepare_actor.cc    | 13 ++++---------
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/mindspore/ccsrc/plugin/device/ascend/hal/hardware/ascend_device_context.cc b/mindspore/ccsrc/plugin/device/ascend/hal/hardware/ascend_device_context.cc
index 8ecf205af2..65ba19a428 100644
--- a/mindspore/ccsrc/plugin/device/ascend/hal/hardware/ascend_device_context.cc
+++ b/mindspore/ccsrc/plugin/device/ascend/hal/hardware/ascend_device_context.cc
@@ -26,6 +26,7 @@
 #include "utils/utils.h"
 #include "plugin/device/ascend/hal/device/kernel_select_ascend.h"
 #include "runtime/device/kernel_adjust.h"
+#include "runtime/device/memory_manager.h"
 #include "plugin/device/ascend/hal/device/ascend_stream_assign.h"
 #include "plugin/device/ascend/hal/device/kernel_build_ascend.h"
 #include "plugin/device/ascend/hal/hardware/ascend_graph_optimization.h"
@@ -526,8 +527,19 @@ void AscendDeviceContext::FreeMemory(DeviceAddress *const &address) const {
 bool AscendDeviceContext::AllocateContinuousMemory(const std::vector<DeviceAddressPtr> &addr_list, size_t total_size,
                                                    const std::vector<size_t> &size_list) const {
   MS_EXCEPTION_IF_NULL(runtime_instance_);
+  if (addr_list.size() != size_list.size()) {
+    MS_LOG(EXCEPTION) << "AllocateContinuousMemory check failed: input address list size " << addr_list.size()
+                      << " vs input size list size " << size_list.size();
+  }
   runtime_instance_->SetContext();
-  return mem_manager_->MallocContinuousMemFromMemPool(addr_list, total_size, size_list);
+  size_t align_total_size = 0;
+  std::vector<size_t> align_size_list;
+  for (size_t i = 0; i < size_list.size(); i++) {
+    auto align_size = device::MemoryManager::GetCommonAlignSize(size_list[i]);
+    align_size_list.emplace_back(align_size);
+    align_total_size += align_size;
+  }
+  return mem_manager_->MallocContinuousMemFromMemPool(addr_list, align_total_size, align_size_list);
 }
 
 void *AscendDeviceContext::AllocateMemory(size_t size) const {
diff --git a/mindspore/ccsrc/runtime/graph_scheduler/actor/data_prepare_actor.cc b/mindspore/ccsrc/runtime/graph_scheduler/actor/data_prepare_actor.cc
index 9d7279bb0f..f2ebea0b3c 100644
--- a/mindspore/ccsrc/runtime/graph_scheduler/actor/data_prepare_actor.cc
+++ b/mindspore/ccsrc/runtime/graph_scheduler/actor/data_prepare_actor.cc
@@ -25,7 +25,6 @@
 #include "mindrt/include/async/async.h"
 #include "utils/log_adapter.h"
 #include "utils/convert_utils.h"
-#include "runtime/device/memory_manager.h"
 
 namespace mindspore {
 namespace runtime {
@@ -76,10 +75,8 @@ void FetchContinuousMemoryInfo(const CNodePtr &node, std::vector<DeviceTensorPtr
     for (size_t i = 0; i < intput_sizes.size(); ++i) {
       const auto &device_tensor = AnfAlgo::GetPrevNodeMutableOutputAddr(node, i, false);
       MS_EXCEPTION_IF_NULL(device_tensor);
-      auto origin_size = intput_sizes[i];
-      auto align_size = device::MemoryManager::GetCommonAlignSize(origin_size);
-      *total_size += align_size;
-      (void)size_list->emplace_back(align_size);
+      *total_size += intput_sizes[i];
+      (void)size_list->emplace_back(intput_sizes[i]);
       (void)addr_list->emplace_back(device_tensor);
     }
   } else {
@@ -87,10 +84,8 @@ void FetchContinuousMemoryInfo(const CNodePtr &node, std::vector<DeviceTensorPtr
     for (size_t i = 0; i < output_sizes.size(); ++i) {
       const auto &device_tensor = AnfAlgo::GetMutableOutputAddr(node, i, false);
       MS_EXCEPTION_IF_NULL(device_tensor);
-      auto origin_size = output_sizes[i];
-      auto align_size = device::MemoryManager::GetCommonAlignSize(origin_size);
-      *total_size += align_size;
-      (void)size_list->emplace_back(align_size);
+      *total_size += output_sizes[i];
+      (void)size_list->emplace_back(output_sizes[i]);
       (void)addr_list->emplace_back(device_tensor);
     }
   }