1 /**
2 * Copyright 2021 Huawei Technologies Co., Ltd
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 #include "runtime/device/launch_kernel.h"
18
19 namespace mindspore::device {
ObtainKernelAddress(const std::vector<size_t> & list,std::vector<uint8_t * > * addr)20 std::vector<kernel::AddressPtr> LaunchKernel::ObtainKernelAddress(const std::vector<size_t> &list,
21 std::vector<uint8_t *> *addr) {
22 MS_EXCEPTION_IF_NULL(addr);
23 std::vector<kernel::AddressPtr> kernel_address;
24 if (addr->size() < list.size()) {
25 MS_LOG_EXCEPTION << "Error addr size!";
26 }
27 for (size_t i = 0; i < list.size(); ++i) {
28 auto size = AlignSizeForLaunchKernel(list[i]);
29 (*addr)[i] = AllocDeviceMem(size);
30 auto address = std::make_shared<kernel::Address>();
31 MS_EXCEPTION_IF_NULL(address);
32 address->addr = (*addr)[i];
33 MS_EXCEPTION_IF_NULL(address->addr);
34 address->size = size;
35 kernel_address.push_back(address);
36 }
37 return kernel_address;
38 }
39
ObtainKernelInputs(const std::vector<size_t> & inputs_list,const std::vector<uint8_t * > & inputs_addr)40 std::vector<kernel::AddressPtr> LaunchKernel::ObtainKernelInputs(const std::vector<size_t> &inputs_list,
41 const std::vector<uint8_t *> &inputs_addr) {
42 std::vector<kernel::AddressPtr> kernel_inputs;
43 if (inputs_list.size() != inputs_addr.size()) {
44 MS_LOG(ERROR) << "input_list size should equal to input_addr_ size, input_list size: " << inputs_list.size()
45 << ", input_addr_ size: " << inputs_addr.size();
46 }
47 for (size_t i = 0; i < inputs_list.size(); ++i) {
48 auto input_size = AlignSizeForLaunchKernel(inputs_list[i]);
49 auto input = std::make_shared<kernel::Address>();
50 MS_EXCEPTION_IF_NULL(input);
51 input->addr = inputs_addr[i];
52 MS_EXCEPTION_IF_NULL(input->addr);
53 input->size = input_size;
54 kernel_inputs.push_back(input);
55 }
56 return kernel_inputs;
57 }
58
ObtainKernelOutputs(const std::vector<size_t> & outputs_list)59 std::vector<kernel::AddressPtr> LaunchKernel::ObtainKernelOutputs(const std::vector<size_t> &outputs_list) {
60 // init output_addr_
61 outputs_addr_ = std::vector<uint8_t *>(outputs_list.size(), nullptr);
62 auto kernel_outputs = ObtainKernelAddress(outputs_list, &outputs_addr_);
63 return kernel_outputs;
64 }
65
ObtainKernelWorkspaces(const std::vector<size_t> & workspaces_list)66 std::vector<kernel::AddressPtr> LaunchKernel::ObtainKernelWorkspaces(const std::vector<size_t> &workspaces_list) {
67 std::vector<kernel::AddressPtr> kernel_workspace;
68 if (workspaces_list.empty()) {
69 return kernel_workspace;
70 }
71 // init workspace_addr_
72 workspaces_addr_ = std::vector<uint8_t *>(workspaces_list.size(), nullptr);
73 kernel_workspace = ObtainKernelAddress(workspaces_list, &workspaces_addr_);
74 return kernel_workspace;
75 }
76
LaunchSingleKernel(const std::vector<uint8_t * > & inputs_addr)77 void LaunchKernel::LaunchSingleKernel(const std::vector<uint8_t *> &inputs_addr) {
78 MS_EXCEPTION_IF_NULL(kernel_mod_);
79 // obtain kernel inputs
80 auto kernel_inputs = ObtainKernelInputs(kernel_mod_->GetInputSizeList(), inputs_addr);
81 // obtain kernel outputs
82 auto kernel_outputs = ObtainKernelOutputs(kernel_mod_->GetOutputSizeList());
83 // obtain kernel workspace
84 auto kernel_workspaces = ObtainKernelWorkspaces(kernel_mod_->GetWorkspaceSizeList());
85 // launch
86 auto ret_status = kernel_mod_->Launch(kernel_inputs, kernel_workspaces, kernel_outputs, stream_);
87 if (!ret_status) {
88 MS_LOG(ERROR) << "Launch single kernel failed.";
89 }
90 }
91
FreeOutputAndWorkspaceDeviceMem()92 void LaunchKernel::FreeOutputAndWorkspaceDeviceMem() {
93 // free outputs_addr and workspaces_addr_
94 for (size_t i = 0; i < outputs_addr_.size(); ++i) {
95 if (outputs_addr_[i] != nullptr) {
96 FreeDeviceMem(outputs_addr_[i]);
97 outputs_addr_[i] = nullptr;
98 }
99 }
100 for (size_t i = 0; i < workspaces_addr_.size(); ++i) {
101 if (workspaces_addr_[i] != nullptr) {
102 FreeDeviceMem(workspaces_addr_[i]);
103 workspaces_addr_[i] = nullptr;
104 }
105 }
106 outputs_addr_.clear();
107 workspaces_addr_.clear();
108 }
109 } // namespace mindspore::device
110