!14294 add the impl of runtime actors
From: @limingqi107 Reviewed-by: Signed-off-by:pull/14294/MERGE
mindspore-ci-bot
4 years ago
committed by
Gitee
commit
d2ee376b7b
@ -0,0 +1,8 @@
|
||||
include_directories(${CMAKE_SOURCE_DIR}/mindspore/core/mindrt/include)
|
||||
include_directories(${CMAKE_SOURCE_DIR}/mindspore/core/mindrt/src)
|
||||
|
||||
file(GLOB_RECURSE FRAMEWORK_SRC_LIST RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} "*.cc")
|
||||
|
||||
set_property(SOURCE ${FRAMEWORK_SRC_LIST}
|
||||
PROPERTY COMPILE_DEFINITIONS SUBMODULE_ID=mindspore::SubModuleId::SM_RUNTIME_FRAMEWORK)
|
||||
add_library(_mindspore_runtime_framework_obj OBJECT ${FRAMEWORK_SRC_LIST})
|
||||
@ -0,0 +1,46 @@
|
||||
/**
|
||||
* Copyright 2021 Huawei Technologies Co., Ltd
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#ifndef MINDSPORE_CCSRC_RUNTIME_FRAMEWORK_ACTOR_ACTOR_COMMON_H_
|
||||
#define MINDSPORE_CCSRC_RUNTIME_FRAMEWORK_ACTOR_ACTOR_COMMON_H_
|
||||
|
||||
#include <utility>
|
||||
#include "mindrt/include/actor/op_actor.h"
|
||||
#include "utils/log_adapter.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace runtime {
|
||||
// The execution result of actor.
|
||||
constexpr int kSuccess = 0;
|
||||
constexpr int kFailure = 1;
|
||||
|
||||
#define SET_OPCONTEXT_FAIL_RET_WITH_ERROR(op_context, message) \
|
||||
{ \
|
||||
MS_LOG(ERROR) << message; \
|
||||
op_context.SetFailed(kFailure); \
|
||||
return; \
|
||||
}
|
||||
|
||||
#define SET_OPCONTEXT_SUCCESS_RET(op_context) \
|
||||
{ \
|
||||
op_context.SetSuccess(kSuccess); \
|
||||
return; \
|
||||
}
|
||||
|
||||
} // namespace runtime
|
||||
} // namespace mindspore
|
||||
|
||||
#endif // MINDSPORE_CCSRC_RUNTIME_FRAMEWORK_ACTOR_ACTOR_COMMON_H_
|
||||
@ -0,0 +1,163 @@
|
||||
/**
|
||||
* Copyright 2021 Huawei Technologies Co., Ltd
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include "runtime/framework/actor/data_source_actor.h"
|
||||
#include "runtime/framework/actor/kernel_actor.h"
|
||||
#include "runtime/framework/actor/memory_manager_actor.h"
|
||||
#include "mindrt/include/async/async.h"
|
||||
#include "common/trans.h"
|
||||
#include "utils/log_adapter.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace runtime {
|
||||
void DataSourceActor::FetchData(OpContext<DeviceTensor> *context) {
|
||||
MS_EXCEPTION_IF_NULL(context);
|
||||
if (buffers_.size() == buffer_capacity_) {
|
||||
// Send output to trigger computing and free memory.
|
||||
SendOutput(context);
|
||||
FreeMemory(context);
|
||||
buffers_.pop();
|
||||
return;
|
||||
}
|
||||
|
||||
// Construct device tensors and fill to the buffers from member nodes.
|
||||
FillDataBuffer();
|
||||
if (buffers_.size() == 0) {
|
||||
SET_OPCONTEXT_FAIL_RET_WITH_ERROR((*context), "The data queue is empty.");
|
||||
}
|
||||
|
||||
// Allocate memory for device tensors.
|
||||
AllocateMemory(context);
|
||||
}
|
||||
|
||||
void DataSourceActor::AllocateMemory(OpContext<DeviceTensor> *context) {
|
||||
auto device_tensors = buffers_.back();
|
||||
Async(memory_manager_aid_, &MemoryManagerActor::AllocateMemory, device_tensors, device_context_, context, GetAID());
|
||||
}
|
||||
|
||||
void DataSourceActor::FreeMemory(OpContext<DeviceTensor> *context) {
|
||||
auto device_tensors = buffers_.front();
|
||||
Async(memory_manager_aid_, &MemoryManagerActor::FreeMemory, device_tensors, device_context_, context);
|
||||
}
|
||||
|
||||
void DataSourceActor::SendOutput(OpContext<DeviceTensor> *context) {
|
||||
MS_EXCEPTION_IF_NULL(context);
|
||||
if (buffers_.size() == 0) {
|
||||
SET_OPCONTEXT_FAIL_RET_WITH_ERROR((*context), "The data queue is empty.");
|
||||
}
|
||||
|
||||
// Send output data.
|
||||
auto output_device_tensors = buffers_.front();
|
||||
for (auto &op_arrow : output_op_arrows_) {
|
||||
MS_EXCEPTION_IF_NULL(op_arrow);
|
||||
if (IntToSize(op_arrow->from_output_index_) >= output_device_tensors.size()) {
|
||||
SET_OPCONTEXT_FAIL_RET_WITH_ERROR((*context), "The output index is of range.");
|
||||
}
|
||||
auto device_address = output_device_tensors[op_arrow->from_output_index_];
|
||||
auto data = std::make_shared<OpData<DeviceTensor>>(op_arrow->to_op_id_, device_address, op_arrow->to_input_index_);
|
||||
Async(op_arrow->to_op_id_, &KernelActor::RunOpData, data, context);
|
||||
}
|
||||
}
|
||||
|
||||
void DeviceQueueDataSourceActor::FillDataBuffer() {
|
||||
// Construct device tensors.
|
||||
std::vector<DeviceTensor *> device_tensors;
|
||||
for (size_t i = 0; i < AnfAlgo::GetOutputTensorNum(data_kernel_); ++i) {
|
||||
auto device_address = AnfAlgo::GetMutableOutputAddr(data_kernel_, i, false);
|
||||
MS_EXCEPTION_IF_NULL(device_address);
|
||||
device_tensors.emplace_back(device_address.get());
|
||||
}
|
||||
|
||||
buffers_.push(device_tensors);
|
||||
}
|
||||
|
||||
void DeviceQueueDataSourceActor::OnMemoryAllocFinish(OpContext<DeviceTensor> *context) {
|
||||
MS_EXCEPTION_IF_NULL(context);
|
||||
MS_EXCEPTION_IF_NULL(device_context_);
|
||||
if (buffers_.size() == 0) {
|
||||
SET_OPCONTEXT_FAIL_RET_WITH_ERROR((*context), "The data queue is empty.");
|
||||
}
|
||||
|
||||
// Construct outputs of data kernel launching.
|
||||
auto device_tensors = buffers_.back();
|
||||
std::vector<AddressPtr> kernel_outputs;
|
||||
for (auto &device_tensor : device_tensors) {
|
||||
MS_EXCEPTION_IF_NULL(device_tensor);
|
||||
kernel_outputs.emplace_back(std::make_shared<Address>(device_tensor->GetMutablePtr(), device_tensor->GetSize()));
|
||||
}
|
||||
|
||||
// Copy data from device queue by data kernel launching.
|
||||
std::vector<AddressPtr> empty_address;
|
||||
auto kernel_mod = AnfAlgo::GetKernelMod(data_kernel_);
|
||||
auto ret = device_context_->LaunchKernel(kernel_mod, empty_address, empty_address, kernel_outputs);
|
||||
if (!ret) {
|
||||
std::string error_info = "Launch kernel failed: " + data_kernel_->ToString();
|
||||
SET_OPCONTEXT_FAIL_RET_WITH_ERROR((*context), error_info);
|
||||
}
|
||||
|
||||
// Send output to trigger computing and free memory.
|
||||
SendOutput(context);
|
||||
FreeMemory(context);
|
||||
buffers_.pop();
|
||||
}
|
||||
|
||||
void HostQueueDataSourceActor::FillDataBuffer() {
|
||||
// Construct device tensors.
|
||||
std::vector<DeviceTensor *> device_tensors;
|
||||
for (auto &data_node : data_nodes_) {
|
||||
auto device_address = AnfAlgo::GetMutableOutputAddr(data_node, 0, false);
|
||||
MS_EXCEPTION_IF_NULL(device_address);
|
||||
device_tensors.emplace_back(device_address.get());
|
||||
}
|
||||
|
||||
buffers_.push(device_tensors);
|
||||
}
|
||||
|
||||
void HostQueueDataSourceActor::OnMemoryAllocFinish(OpContext<DeviceTensor> *context) {
|
||||
MS_EXCEPTION_IF_NULL(context);
|
||||
if (buffers_.size() == 0) {
|
||||
SET_OPCONTEXT_FAIL_RET_WITH_ERROR((*context), "The data queue is empty.");
|
||||
}
|
||||
|
||||
// Get host tensors from host queue and get device tensors from buffers.
|
||||
MS_EXCEPTION_IF_NULL(host_queue_);
|
||||
auto host_tensors = host_queue_->PullData();
|
||||
auto device_tensors = buffers_.back();
|
||||
if (host_tensors.size() != device_tensors.size()) {
|
||||
SET_OPCONTEXT_FAIL_RET_WITH_ERROR((*context),
|
||||
"The length of host tensors is not equal to the length of device tensors.");
|
||||
}
|
||||
|
||||
// Copy data from host tensor to device tensor.
|
||||
for (size_t i = 0; i < host_tensors.size(); ++i) {
|
||||
auto host_tensor = host_tensors[i];
|
||||
auto device_tensor = device_tensors[i];
|
||||
MS_EXCEPTION_IF_NULL(host_tensor);
|
||||
MS_EXCEPTION_IF_NULL(device_tensor);
|
||||
if (!device_tensor->SyncHostToDevice(trans::GetRuntimePaddingShape(data_nodes_[i], 0),
|
||||
LongToSize(host_tensor->data().nbytes()), host_tensor->data_type(),
|
||||
host_tensor->data_c())) {
|
||||
SET_OPCONTEXT_FAIL_RET_WITH_ERROR((*context), "SyncHostToDevice failed.");
|
||||
}
|
||||
}
|
||||
|
||||
// Send output to trigger computing and free memory.
|
||||
SendOutput(context);
|
||||
FreeMemory(context);
|
||||
buffers_.pop();
|
||||
}
|
||||
} // namespace runtime
|
||||
} // namespace mindspore
|
||||
@ -0,0 +1,190 @@
|
||||
/**
|
||||
* Copyright 2021 Huawei Technologies Co., Ltd
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include "runtime/framework/actor/kernel_actor.h"
|
||||
#include "runtime/framework/actor/memory_manager_actor.h"
|
||||
#include "mindrt/include/async/async.h"
|
||||
#include "utils/log_adapter.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace runtime {
|
||||
void KernelActor::RunOpData(OpDataPtr<DeviceTensor> input_data, OpContext<DeviceTensor> *context) {
|
||||
MS_EXCEPTION_IF_NULL(context);
|
||||
auto sequential_num = context->sequential_num_;
|
||||
input_op_datas_[sequential_num].emplace_back(input_data);
|
||||
// When all the input data are collected, then allocate memory and callback launch.
|
||||
if (CheckLaunchCondition(context)) {
|
||||
FetchInputDeviceTensor(context);
|
||||
FetchOutputDeviceTensor();
|
||||
FetchWorkspaceDeviceTensor();
|
||||
AllocateMemory(context);
|
||||
}
|
||||
}
|
||||
|
||||
void KernelActor::RunOpControl(AID *input_control, OpContext<DeviceTensor> *context) {
|
||||
MS_EXCEPTION_IF_NULL(context);
|
||||
auto sequential_num = context->sequential_num_;
|
||||
input_op_controls_[sequential_num].emplace_back(input_control);
|
||||
// When all the input data are collected, then allocate memory and callback launch.
|
||||
if (CheckLaunchCondition(context)) {
|
||||
FetchInputDeviceTensor(context);
|
||||
FetchOutputDeviceTensor();
|
||||
FetchWorkspaceDeviceTensor();
|
||||
AllocateMemory(context);
|
||||
}
|
||||
}
|
||||
|
||||
void KernelActor::AllocateMemory(OpContext<DeviceTensor> *context) {
|
||||
std::vector<DeviceTensor *> alloc_list(output_device_tensors_);
|
||||
alloc_list.insert(alloc_list.end(), workspace_device_tensors_.begin(), workspace_device_tensors_.end());
|
||||
Async(memory_manager_aid_, &MemoryManagerActor::AllocateMemory, alloc_list, device_context_, context, GetAID());
|
||||
}
|
||||
|
||||
void KernelActor::FreeMemory(OpContext<DeviceTensor> *context) {
|
||||
std::vector<DeviceTensor *> free_list(input_device_tensors_);
|
||||
free_list.insert(free_list.end(), output_device_tensors_.begin(), output_device_tensors_.end());
|
||||
free_list.insert(free_list.end(), workspace_device_tensors_.begin(), workspace_device_tensors_.end());
|
||||
Async(memory_manager_aid_, &MemoryManagerActor::FreeMemory, free_list, device_context_, context);
|
||||
}
|
||||
|
||||
void KernelActor::OnMemoryAllocFinish(OpContext<DeviceTensor> *context) {
|
||||
MS_EXCEPTION_IF_NULL(context);
|
||||
MS_EXCEPTION_IF_NULL(kernel_);
|
||||
auto kernel_mod = AnfAlgo::GetKernelMod(kernel_);
|
||||
std::vector<AddressPtr> kernel_inputs;
|
||||
std::vector<AddressPtr> kernel_outputs;
|
||||
std::vector<AddressPtr> kernel_workspaces;
|
||||
FetchLaunchArgs(&kernel_inputs, &kernel_outputs, &kernel_workspaces);
|
||||
MS_EXCEPTION_IF_NULL(device_context_);
|
||||
auto ret = device_context_->LaunchKernel(kernel_mod, kernel_inputs, kernel_workspaces, kernel_outputs);
|
||||
if (!ret) {
|
||||
std::string error_info = "Launch kernel failed: " + kernel_->ToString();
|
||||
SET_OPCONTEXT_FAIL_RET_WITH_ERROR((*context), error_info);
|
||||
}
|
||||
SendOutput(context);
|
||||
FreeMemory(context);
|
||||
}
|
||||
|
||||
bool KernelActor::CheckLaunchCondition(OpContext<DeviceTensor> *context) {
|
||||
MS_EXCEPTION_IF_NULL(context);
|
||||
if (input_datas_num_ != 0) {
|
||||
auto data_iter = input_op_datas_.find(context->sequential_num_);
|
||||
if (data_iter == input_op_datas_.end()) {
|
||||
return false;
|
||||
}
|
||||
if (data_iter->second.size() != input_datas_num_) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
if (input_controls_num_ != 0) {
|
||||
auto control_iter = input_op_controls_.find(context->sequential_num_);
|
||||
if (control_iter == input_op_controls_.end()) {
|
||||
return false;
|
||||
}
|
||||
if (control_iter->second.size() != input_controls_num_) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
void KernelActor::FetchInputDeviceTensor(OpContext<DeviceTensor> *context) {
|
||||
MS_EXCEPTION_IF_NULL(context);
|
||||
auto input_size = input_datas_num_ + device_tensor_store_keys_.size();
|
||||
input_device_tensors_.resize(input_size);
|
||||
|
||||
auto data_iter = input_op_datas_.find(context->sequential_num_);
|
||||
if (data_iter != input_op_datas_.end()) {
|
||||
for (auto &input_data : data_iter->second) {
|
||||
MS_EXCEPTION_IF_NULL(input_data);
|
||||
input_device_tensors_[input_data->index_] = input_data->data_;
|
||||
}
|
||||
}
|
||||
|
||||
for (auto &device_tensor_store_key : device_tensor_store_keys_) {
|
||||
auto device_tensor = DeviceTensorStore::GetInstance().Fetch(device_tensor_store_key.second);
|
||||
input_device_tensors_[device_tensor_store_key.first] = device_tensor.get();
|
||||
}
|
||||
}
|
||||
|
||||
void KernelActor::FetchOutputDeviceTensor() {
|
||||
output_device_tensors_.clear();
|
||||
for (size_t i = 0; i < AnfAlgo::GetOutputTensorNum(kernel_); ++i) {
|
||||
auto device_address = AnfAlgo::GetMutableOutputAddr(kernel_, i, false);
|
||||
MS_EXCEPTION_IF_NULL(device_address);
|
||||
output_device_tensors_.emplace_back(device_address.get());
|
||||
}
|
||||
}
|
||||
|
||||
void KernelActor::FetchWorkspaceDeviceTensor() {
|
||||
workspace_device_tensors_.clear();
|
||||
auto kernel_mod = AnfAlgo::GetKernelMod(kernel_);
|
||||
MS_EXCEPTION_IF_NULL(kernel_mod);
|
||||
auto workspace_sizes = kernel_mod->GetWorkspaceSizeList();
|
||||
for (size_t i = 0; i < workspace_sizes.size(); ++i) {
|
||||
if (workspace_sizes[i] != 0) {
|
||||
auto device_address = AnfAlgo::GetMutableWorkspaceAddr(kernel_, i);
|
||||
MS_EXCEPTION_IF_NULL(device_address);
|
||||
workspace_device_tensors_.emplace_back(device_address.get());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void KernelActor::FetchLaunchArgs(std::vector<AddressPtr> *kernel_inputs, std::vector<AddressPtr> *kernel_outputs,
|
||||
std::vector<AddressPtr> *kernel_workspaces) {
|
||||
MS_EXCEPTION_IF_NULL(kernel_inputs);
|
||||
MS_EXCEPTION_IF_NULL(kernel_outputs);
|
||||
MS_EXCEPTION_IF_NULL(kernel_workspaces);
|
||||
for (auto &input : input_device_tensors_) {
|
||||
MS_EXCEPTION_IF_NULL(input);
|
||||
kernel_inputs->emplace_back(std::make_shared<Address>(input->GetMutablePtr(), input->GetSize()));
|
||||
}
|
||||
|
||||
for (auto &output : output_device_tensors_) {
|
||||
MS_EXCEPTION_IF_NULL(output);
|
||||
kernel_outputs->emplace_back(std::make_shared<Address>(output->GetMutablePtr(), output->GetSize()));
|
||||
}
|
||||
|
||||
for (auto &workspace : workspace_device_tensors_) {
|
||||
MS_EXCEPTION_IF_NULL(workspace);
|
||||
kernel_workspaces->emplace_back(std::make_shared<Address>(workspace->GetMutablePtr(), workspace->GetSize()));
|
||||
}
|
||||
}
|
||||
|
||||
void KernelActor::SendOutput(OpContext<DeviceTensor> *context) {
|
||||
MS_EXCEPTION_IF_NULL(context);
|
||||
// Send output data.
|
||||
for (auto &op_arrow : output_op_arrows_) {
|
||||
MS_EXCEPTION_IF_NULL(op_arrow);
|
||||
if (IntToSize(op_arrow->from_output_index_) >= output_device_tensors_.size()) {
|
||||
std::string error_info = "The output index is out of range: " + kernel_->ToString();
|
||||
SET_OPCONTEXT_FAIL_RET_WITH_ERROR((*context), error_info);
|
||||
}
|
||||
auto device_address = output_device_tensors_[op_arrow->from_output_index_];
|
||||
auto data = std::make_shared<OpData<DeviceTensor>>(op_arrow->to_op_id_, device_address, op_arrow->to_input_index_);
|
||||
Async(op_arrow->to_op_id_, &KernelActor::RunOpData, data, context);
|
||||
}
|
||||
|
||||
// Send output control.
|
||||
auto source_aid = const_cast<AID *>(&GetAID());
|
||||
for (auto &output_control : output_op_controls_) {
|
||||
Async(output_control, &OpActor::RunOpControl, source_aid, context);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace runtime
|
||||
} // namespace mindspore
|
||||
@ -0,0 +1,47 @@
|
||||
/**
|
||||
* Copyright 2021 Huawei Technologies Co., Ltd
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include "runtime/framework/actor/loop_count_actor.h"
|
||||
#include "runtime/framework/actor/data_source_actor.h"
|
||||
#include "runtime/framework/actor/kernel_actor.h"
|
||||
#include "mindrt/include/async/async.h"
|
||||
#include "utils/log_adapter.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace runtime {
|
||||
void LoopCountActor::RunOpControl(AID *input_control, OpContext<DeviceTensor> *context) {
|
||||
MS_EXCEPTION_IF_NULL(context);
|
||||
auto sequential_num = context->sequential_num_;
|
||||
input_op_controls_[sequential_num].emplace_back(input_control);
|
||||
if (input_op_controls_[sequential_num].size() == input_controls_num_) {
|
||||
current_count_++;
|
||||
if (current_count_ == loop_count_) {
|
||||
current_count_ = 0;
|
||||
SET_OPCONTEXT_SUCCESS_RET((*context));
|
||||
}
|
||||
|
||||
// Send output control.
|
||||
for (auto &data_source_aid : data_source_aids_) {
|
||||
Async(data_source_aid, &DataSourceActor::FetchData, context);
|
||||
}
|
||||
auto source_aid = const_cast<AID *>(&GetAID());
|
||||
for (auto &kernel_aid : no_input_kernel_aids_) {
|
||||
Async(kernel_aid, &KernelActor::RunOpControl, source_aid, context);
|
||||
}
|
||||
}
|
||||
}
|
||||
} // namespace runtime
|
||||
} // namespace mindspore
|
||||
@ -0,0 +1,39 @@
|
||||
/**
|
||||
* Copyright 2021 Huawei Technologies Co., Ltd
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#ifndef MINDSPORE_CCSRC_RUNTIME_FRAMEWORK_ACTOR_MEMORY_INTERFACE_ACTOR_H_
|
||||
#define MINDSPORE_CCSRC_RUNTIME_FRAMEWORK_ACTOR_MEMORY_INTERFACE_ACTOR_H_
|
||||
|
||||
#include <utility>
|
||||
#include <string>
|
||||
#include "mindrt/include/actor/op_actor.h"
|
||||
#include "runtime/framework/device_tensor_store.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace runtime {
|
||||
// The actor represents a set of common memory related operations of actor.
|
||||
class MemoryInterfaceActor : public OpActor<DeviceTensor> {
|
||||
public:
|
||||
explicit MemoryInterfaceActor(std::string name) : OpActor(name) {}
|
||||
virtual ~MemoryInterfaceActor() = default;
|
||||
virtual void AllocateMemory(OpContext<DeviceTensor> *context) = 0;
|
||||
virtual void FreeMemory(OpContext<DeviceTensor> *context) = 0;
|
||||
virtual void OnMemoryAllocFinish(OpContext<DeviceTensor> *context) = 0;
|
||||
};
|
||||
} // namespace runtime
|
||||
} // namespace mindspore
|
||||
|
||||
#endif // MINDSPORE_CCSRC_RUNTIME_FRAMEWORK_ACTOR_MEMORY_INTERFACE_ACTOR_H_
|
||||
@ -0,0 +1,62 @@
|
||||
/**
|
||||
* Copyright 2021 Huawei Technologies Co., Ltd
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include "runtime/framework/actor/memory_manager_actor.h"
|
||||
#include "runtime/framework/actor/data_source_actor.h"
|
||||
#include "runtime/framework/actor/kernel_actor.h"
|
||||
#include "mindrt/include/async/async.h"
|
||||
#include "utils/log_adapter.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace runtime {
|
||||
void MemoryManagerActor::AllocateMemory(std::vector<DeviceTensor *> alloc_list, const DeviceContext *device_context,
|
||||
OpContext<DeviceTensor> *op_context, const AID from_aid) {
|
||||
MS_EXCEPTION_IF_NULL(device_context);
|
||||
MS_EXCEPTION_IF_NULL(op_context);
|
||||
|
||||
for (auto &device_tensor : alloc_list) {
|
||||
MS_EXCEPTION_IF_NULL(device_tensor);
|
||||
if (device_tensor->GetPtr() != nullptr) {
|
||||
continue;
|
||||
}
|
||||
// Allocate memory through the device context.
|
||||
if (!device_context->AllocateMemory(device_tensor, device_tensor->GetSize())) {
|
||||
std::string error_info = "Device memory isn't enough and alloc failed, actor name: " + from_aid.Name() +
|
||||
", alloc size: " + std::to_string(device_tensor->GetSize());
|
||||
SET_OPCONTEXT_FAIL_RET_WITH_ERROR((*op_context), error_info);
|
||||
}
|
||||
}
|
||||
|
||||
// Call back to the from actor to process after memory allocation finished.
|
||||
Async(from_aid, &MemoryInterfaceActor::OnMemoryAllocFinish, op_context);
|
||||
}
|
||||
|
||||
void MemoryManagerActor::FreeMemory(std::vector<DeviceTensor *> free_list, const DeviceContext *device_context,
|
||||
OpContext<DeviceTensor> *) {
|
||||
MS_EXCEPTION_IF_NULL(device_context);
|
||||
for (auto &device_tensor : free_list) {
|
||||
MS_EXCEPTION_IF_NULL(device_tensor);
|
||||
// The reference count is decremented to zero to free memory, and reset to the original count.
|
||||
device_tensor->DecreaseRefCountUsed();
|
||||
if (device_tensor->ref_count_dynamic_used() == 0) {
|
||||
// Free memory through the device context.
|
||||
device_context->FreeMemory(device_tensor);
|
||||
device_tensor->ResetRefCountUsed();
|
||||
}
|
||||
}
|
||||
}
|
||||
} // namespace runtime
|
||||
} // namespace mindspore
|
||||
Loading…
Reference in new issue