1 /**
2 * Copyright 2019 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 "backend/kernel_compiler/rts/rt_kernel_info.h"
18 #include <unordered_map>
19 #include <algorithm>
20 #include "utils/convert_utils.h"
21 #include "utils/utils.h"
22 #include "utils/ms_utils.h"
23 #include "backend/session/anf_runtime_algorithm.h"
24
25 namespace mindspore {
26 namespace kernel {
Register(const std::string & name,RtKerDescCreater && fun)27 void RtKerDescFactory::Register(const std::string &name, RtKerDescCreater &&fun) {
28 if (fmap_.find(name) == fmap_.end()) {
29 (void)fmap_.emplace(name, std::move(fun));
30 }
31 }
32
Create(const std::string & name)33 std::shared_ptr<RtKerDesc> RtKerDescFactory::Create(const std::string &name) {
34 const auto &map = Get().fmap_;
35 auto it = map.find(name);
36 if (it != map.end() && it->second) {
37 return (it->second)();
38 }
39 return nullptr;
40 }
41
Get()42 RtKerDescFactory &RtKerDescFactory::Get() {
43 static RtKerDescFactory _this{};
44 return _this;
45 }
46
IsDefaultKernelInfo(const std::string & name)47 static bool IsDefaultKernelInfo(const std::string &name) {
48 static const std::set<std::string> white_list = {kStreamSwitchOpName, kStreamActiveOpName, kLabelSetOpName,
49 kLabelGotoOpName};
50 return white_list.find(name) != white_list.end();
51 }
52
GetRtKelInfo(const CNodePtr & kernel_node,std::vector<std::shared_ptr<kernel::KernelBuildInfo>> * kernel_info_list)53 void GetRtKelInfo(const CNodePtr &kernel_node,
54 std::vector<std::shared_ptr<kernel::KernelBuildInfo>> *kernel_info_list) {
55 MS_EXCEPTION_IF_NULL(kernel_info_list);
56 MS_EXCEPTION_IF_NULL(kernel_node);
57 std::string opNameLower = AnfAlgo::GetCNodeName(kernel_node);
58 (void)std::transform(opNameLower.begin(), opNameLower.end(), opNameLower.begin(), ::tolower);
59
60 auto ker_desc_ptr = RtKerDescFactory::Create(opNameLower);
61 if (ker_desc_ptr != nullptr && !ker_desc_ptr->GetKernelInfo().empty()) {
62 *kernel_info_list = ker_desc_ptr->GetKernelInfo();
63 return;
64 }
65 // if can't find kernel info in kernel info database, use the default kernel info
66 auto node_name = AnfAlgo::GetCNodeName(kernel_node);
67 if (IsDefaultKernelInfo(node_name)) {
68 auto kernel_build_info_builder = std::make_shared<kernel::KernelBuildInfo::KernelBuildInfoBuilder>();
69 // set input infos
70 auto input_num = AnfAlgo::GetInputTensorNum(kernel_node);
71 MS_EXCEPTION_IF_NULL(kernel_build_info_builder);
72 kernel_build_info_builder->SetInputsFormat(std::vector<std::string>(input_num, kOpFormat_DEFAULT));
73 std::vector<TypeId> input_types = {};
74 for (size_t i = 0; i < input_num; i++) {
75 input_types.push_back(AnfAlgo::GetPrevNodeOutputInferDataType(kernel_node, i));
76 }
77 kernel_build_info_builder->SetInputsDeviceType(input_types);
78 // Kernel ops in while-list such as 'LabelSet' always return UMonad.
79 kernel_build_info_builder->SetOutputsFormat({kOpFormat_DEFAULT});
80 kernel_build_info_builder->SetOutputsDeviceType({TypeId::kObjectTypeUMonad});
81 // set other info
82 kernel_build_info_builder->SetFusionType(kernel::FusionType::OPAQUE);
83 kernel_build_info_builder->SetProcessor(kernel::Processor::AICORE);
84 kernel_build_info_builder->SetKernelType(KernelType::RT_KERNEL);
85 kernel_info_list->push_back(kernel_build_info_builder->Build());
86 return;
87 }
88 MS_LOG(DEBUG) << "Rt dose not have op [" << opNameLower << "].";
89 }
90 } // namespace kernel
91 } // namespace mindspore
92