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_transdata.h"
18
19 #include <algorithm>
20 #include "abstract/utils.h"
21 #include "backend/session/single_kernel_graph.h"
22 #include "backend/session/anf_runtime_algorithm.h"
23
24 namespace mindspore::device::ascend {
FreeDeviceMem(void * addr)25 void AscendLaunchTransData::FreeDeviceMem(void *addr) { AscendLaunchKernel::FreeDeviceMem(addr); }
26
AlignSizeForLaunchKernel(size_t size)27 size_t AscendLaunchTransData::AlignSizeForLaunchKernel(size_t size) {
28 return AscendLaunchKernel::AlignSizeForLaunchKernel(size);
29 }
30
AllocDeviceMem(size_t size)31 uint8_t *AscendLaunchTransData::AllocDeviceMem(size_t size) { return AscendLaunchKernel::AllocDeviceMem(size); }
32
KernelSelect(const std::shared_ptr<session::KernelGraph> & kernel_graph)33 void AscendLaunchTransData::KernelSelect(const std::shared_ptr<session::KernelGraph> &kernel_graph) {
34 AscendLaunchKernel::KernelSelect(kernel_graph);
35 }
36
KernelBuild(const std::shared_ptr<session::KernelGraph> & kernel_graph)37 void AscendLaunchTransData::KernelBuild(const std::shared_ptr<session::KernelGraph> &kernel_graph) {
38 AscendLaunchKernel::KernelBuild(kernel_graph);
39 }
40
LaunchOpKernel()41 void AscendLaunchTransData::LaunchOpKernel() {
42 if (transdata_graph_ == nullptr) {
43 // construct transdata kernel graph and set attr
44 ConstructKernelGraphAndSetAttr();
45 // kernel build
46 KernelBuild(transdata_graph_);
47 }
48 // obtain kernel_mod
49 if (transdata_graph_->execution_order().size() != 1) {
50 MS_LOG(ERROR) << "The execution order of the transdata graph should have only one node";
51 return;
52 }
53 kernel_mod_ = AnfAlgo::GetKernelMod(transdata_graph_->execution_order()[0]);
54 MS_EXCEPTION_IF_NULL(kernel_mod_);
55 // obtain kernel inputs
56 std::vector<kernel::AddressPtr> kernel_inputs;
57 auto input = std::make_shared<kernel::Address>();
58 MS_EXCEPTION_IF_NULL(input);
59 input->addr = input_addr_;
60 MS_EXCEPTION_IF_NULL(input->addr);
61 input->size = total_size_;
62 kernel_inputs.push_back(input);
63 // obtain kernel outputs
64 auto kernel_outputs = ObtainKernelOutputs(kernel_mod_->GetOutputSizeList());
65 // obtain kernel workspaces
66 auto kernel_workspace = ObtainKernelWorkspaces(kernel_mod_->GetWorkspaceSizeList());
67 // launch
68 auto ret_status = kernel_mod_->Launch(kernel_inputs, kernel_workspace, kernel_outputs, stream_);
69 if (!ret_status) {
70 MS_LOG(EXCEPTION) << "Launch transdata single kernel failed";
71 }
72 }
73
FreeLaunchDeviceMem()74 void AscendLaunchTransData::FreeLaunchDeviceMem() {
75 input_addr_ = nullptr;
76 FreeOutputAndWorkspaceDeviceMem();
77 }
78
ObtainTransDataKernelGraph()79 std::shared_ptr<session::KernelGraph> AscendLaunchTransData::ObtainTransDataKernelGraph() {
80 std::vector<TypeId> input_dtypes = {dtype_};
81 std::vector<TypeId> output_dtypes = {dtype_};
82 // obtain input & output shape
83 std::vector<int64_t> input_shape;
84 std::transform(shape_.begin(), shape_.end(), std::back_inserter(input_shape), SizeToLong);
85 std::vector<std::vector<int64_t>> input_shapes = {{input_shape}};
86 std::vector<std::vector<size_t>> output_shapes = {{shape_}};
87 auto transdata_graph = session::SingleKernelGraph::ConstructKernelGraphBasedOnSingleOp(
88 kTransDataOpName, input_dtypes, input_shapes, output_dtypes, output_shapes);
89 MS_EXCEPTION_IF_NULL(transdata_graph);
90 return transdata_graph;
91 }
92
ConstructKernelGraphAndSetAttr()93 void AscendLaunchTransData::ConstructKernelGraphAndSetAttr() {
94 // construct transdata kernel graph
95 transdata_graph_ = ObtainTransDataKernelGraph();
96 MS_EXCEPTION_IF_NULL(transdata_graph_);
97 // set transdata attr
98 if (!transdata_graph_->execution_order().empty()) {
99 auto transdata_node = transdata_graph_->execution_order()[0];
100 // set output infer type and shape
101 AnfAlgo::SetOutputInferTypeAndShape({dtype_}, {shape_}, transdata_node.get());
102 // set build info
103 auto builder = std::make_shared<kernel::KernelBuildInfo::KernelBuildInfoBuilder>();
104 builder->SetKernelType(KernelType::TBE_KERNEL);
105 std::vector<TypeId> device_type = {dtype_};
106 builder->SetInputsDeviceType(device_type);
107 builder->SetOutputsDeviceType(device_type);
108 std::vector<std::string> inputs_format = {src_format_};
109 std::vector<std::string> outputs_format = {dst_format_};
110 builder->SetInputsFormat(inputs_format);
111 builder->SetOutputsFormat(outputs_format);
112 AnfAlgo::SetSelectKernelBuildInfo(builder->Build(), transdata_node.get());
113 // set attr
114 AnfAlgo::SetNodeAttr(kAttrSrcFormat, MakeValue(src_format_), transdata_node);
115 AnfAlgo::SetNodeAttr(kAttrDstFormat, MakeValue(dst_format_), transdata_node);
116 }
117 }
118 } // namespace mindspore::device::ascend
119