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/ascend/ascend_launch_atomic_clean.h"
18 #include "abstract/utils.h"
19 #include "backend/session/single_kernel_graph.h"
20 #include "backend/session/anf_runtime_algorithm.h"
21 #include "debug/anf_ir_dump.h"
22
23 namespace mindspore::device::ascend {
FreeDeviceMem(void * addr)24 void AscendLaunchAtomicClean::FreeDeviceMem(void *addr) { AscendLaunchKernel::FreeDeviceMem(addr); }
25
AlignSizeForLaunchKernel(size_t size)26 size_t AscendLaunchAtomicClean::AlignSizeForLaunchKernel(size_t size) {
27 return AscendLaunchKernel::AlignSizeForLaunchKernel(size);
28 }
29
AllocDeviceMem(size_t size)30 uint8_t *AscendLaunchAtomicClean::AllocDeviceMem(size_t size) { return AscendLaunchKernel::AllocDeviceMem(size); }
31
KernelSelect(const std::shared_ptr<session::KernelGraph> & kernel_graph)32 void AscendLaunchAtomicClean::KernelSelect(const std::shared_ptr<session::KernelGraph> &kernel_graph) {
33 AscendLaunchKernel::KernelSelect(kernel_graph);
34 }
35
KernelBuild(const std::shared_ptr<session::KernelGraph> & kernel_graph)36 void AscendLaunchAtomicClean::KernelBuild(const std::shared_ptr<session::KernelGraph> &kernel_graph) {
37 AscendLaunchKernel::KernelBuild(kernel_graph);
38 }
39
LaunchOpKernel()40 void AscendLaunchAtomicClean::LaunchOpKernel() {
41 if (atomic_clean_graph_ == nullptr) {
42 // construct atomic clean kernel graph and set attr
43 ConstructKernelGraphAndSetAttr();
44 // kernel build
45 KernelBuild(atomic_clean_graph_);
46 }
47 // obtain kernel_mod
48 if (atomic_clean_graph_->execution_order().size() != 1) {
49 MS_LOG(ERROR) << "The execution order of the atomic clean graph should have only one node";
50 }
51 kernel_mod_ = AnfAlgo::GetKernelMod(atomic_clean_graph_->execution_order()[0]);
52 MS_EXCEPTION_IF_NULL(kernel_mod_);
53 // obtain kernel inputs
54 std::vector<kernel::AddressPtr> kernel_inputs;
55 auto input = std::make_shared<kernel::Address>();
56 MS_EXCEPTION_IF_NULL(input);
57 input->addr = input_addr_;
58 MS_EXCEPTION_IF_NULL(input->addr);
59 input->size = total_size_;
60 kernel_inputs.push_back(input);
61 // obtain kernel outputs
62 auto kernel_outputs = ObtainKernelOutputs(kernel_mod_->GetOutputSizeList());
63 // obtain kernel workspace
64 auto kernel_workspaces = ObtainKernelWorkspaces(kernel_mod_->GetWorkspaceSizeList());
65 // launch
66 auto ret_status = kernel_mod_->Launch(kernel_inputs, kernel_workspaces, kernel_outputs, stream_);
67 if (!ret_status) {
68 MS_LOG(ERROR) << "Launch single kernel failed.";
69 }
70 }
71
FreeLaunchDeviceMem()72 void AscendLaunchAtomicClean::FreeLaunchDeviceMem() {
73 input_addr_ = nullptr;
74 FreeOutputAndWorkspaceDeviceMem();
75 }
76
ObtainAtomicCleanKernelGraph()77 std::shared_ptr<session::KernelGraph> AscendLaunchAtomicClean::ObtainAtomicCleanKernelGraph() {
78 std::vector<TypeId> input_dtypes = {dtype_};
79 std::vector<TypeId> output_dtypes = {};
80 // obtain input & output shapes
81 size_t dtype_size = abstract::TypeIdSize(dtype_);
82 if (dtype_size == 0) {
83 MS_LOG(EXCEPTION) << "Divide by zero.";
84 }
85 auto shape = total_size_ / dtype_size;
86 std::vector<std::vector<int64_t>> input_shapes = {{static_cast<int64_t>(shape)}};
87 std::vector<std::vector<size_t>> output_shapes = {};
88 auto atomic_clean_graph = session::SingleKernelGraph::ConstructKernelGraphBasedOnSingleOp(
89 kAtomicAddrCleanOpName, input_dtypes, input_shapes, output_dtypes, output_shapes);
90 MS_EXCEPTION_IF_NULL(atomic_clean_graph);
91 return atomic_clean_graph;
92 }
93
ConstructKernelGraphAndSetAttr()94 void AscendLaunchAtomicClean::ConstructKernelGraphAndSetAttr() {
95 // construct atomic clean kernel graph
96 atomic_clean_graph_ = ObtainAtomicCleanKernelGraph();
97 MS_EXCEPTION_IF_NULL(atomic_clean_graph_);
98 // set atomic clean attr
99 if (!atomic_clean_graph_->execution_order().empty()) {
100 auto clean_node = atomic_clean_graph_->execution_order()[0];
101 // set abstract
102 AbstractBasePtr abstract = std::make_shared<abstract::AbstractNone>();
103 MS_EXCEPTION_IF_NULL(clean_node);
104 clean_node->set_abstract(abstract);
105 // set build info
106 auto builder = std::make_shared<kernel::KernelBuildInfo::KernelBuildInfoBuilder>();
107 MS_EXCEPTION_IF_NULL(builder);
108 builder->SetKernelType(KernelType::TBE_KERNEL);
109 AnfAlgo::SetSelectKernelBuildInfo(builder->Build(), clean_node.get());
110 // set attr
111 std::vector<size_t> clean_size = {total_size_};
112 AnfAlgo::SetNodeAttr(kAttrAtomicAddMemSize, MakeValue(clean_size), clean_node);
113 }
114 }
115 } // namespace mindspore::device::ascend
116