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 9169019839..6d8f21cb17 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" @@ -525,8 +526,19 @@ void AscendDeviceContext::FreeMemory(DeviceAddress *const &address) const { bool AscendDeviceContext::AllocateContinuousMemory(const std::vector &addr_list, size_t total_size, const std::vector &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 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 b8057d529a..453ddea746 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::vectoremplace_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::vectoremplace_back(align_size); + *total_size += output_sizes[i]; + (void)size_list->emplace_back(output_sizes[i]); (void)addr_list->emplace_back(device_tensor); } }