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/hccl/hccl_kernel_metadata.h"
18 #include <memory>
19 #include <algorithm>
20 #include <set>
21 #include "common/trans.h"
22 #include "utils/utils.h"
23 #include "backend/kernel_compiler/hccl/hcom_util.h"
24 #include "backend/session/anf_runtime_algorithm.h"
25 #include "frontend/parallel/context.h"
26
27 namespace mindspore {
28 namespace kernel {
29 namespace {
30 constexpr size_t N_nchw = 0;
31 constexpr size_t C_nchw = 1;
GetKernelFormat(const CNodePtr & kernel_node,size_t index)32 std::string GetKernelFormat(const CNodePtr &kernel_node, size_t index) {
33 static const std::set<std::string> kReduceNoSupportedSet = {kOpFormat_FRAC_Z, kOpFormat_FRACTAL_Z_C04,
34 kOpFormat_C1HWNCoC0};
35 MS_EXCEPTION_IF_NULL(kernel_node);
36 auto op_name = AnfAlgo::GetCNodeName(kernel_node);
37 auto parallel_context_instance = parallel::ParallelContext::GetInstance();
38 MS_EXCEPTION_IF_NULL(parallel_context_instance);
39 if (parallel_context_instance->enable_parallel_optimizer() && op_name == kBroadcast) {
40 return kOpFormat_DEFAULT;
41 }
42 if (op_name == kReceive || op_name == kHcomSend || op_name == kAllToAllv) {
43 return kOpFormat_DEFAULT;
44 }
45 auto format = AnfAlgo::GetPrevNodeOutputFormat(kernel_node, index);
46 if (op_name != kReduceScatter && op_name != kAllGatherOpName) {
47 return format;
48 }
49 auto input_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, index);
50 if (op_name == kAllGatherOpName && input_shape.size() <= kShape4dDims) {
51 auto pad_shape = trans::PaddingShapeTo4dDefault(input_shape);
52 if (pad_shape[N_nchw] % kCubeSize != 0 || pad_shape[C_nchw] % kCubeSize != 0) {
53 return kOpFormat_DEFAULT;
54 }
55 }
56 if (format == kOpFormat_FRAC_NZ && input_shape.size() <= kShape2dDims) {
57 return kOpFormat_DEFAULT;
58 }
59 if (kReduceNoSupportedSet.find(format) != kReduceNoSupportedSet.end()) {
60 return kOpFormat_DEFAULT;
61 }
62 return format;
63 }
64 } // namespace
HcclMetadataInfo(const CNodePtr & kernel_node,std::vector<std::shared_ptr<KernelBuildInfo>> * kernel_info_list)65 void HcclMetadataInfo(const CNodePtr &kernel_node, std::vector<std::shared_ptr<KernelBuildInfo>> *kernel_info_list) {
66 static const std::vector<TypeId> kHcclSupportTypes = {kNumberTypeInt8, kNumberTypeInt32, kNumberTypeFloat16,
67 kNumberTypeFloat32, kNumberTypeInt16};
68 MS_EXCEPTION_IF_NULL(kernel_info_list);
69 MS_EXCEPTION_IF_NULL(kernel_node);
70 std::string op_name = AnfAlgo::GetCNodeName(kernel_node);
71 if (op_name != kAllGather && op_name != kAllReduce && op_name != kBroadcast && op_name != kReduceScatter &&
72 op_name != kHcomSend && op_name != kReceive && op_name != kAllToAllv) {
73 MS_LOG(DEBUG) << "Hccl does not have op [" << op_name << "]";
74 return;
75 }
76 TypeId recv_type;
77 if (op_name == kReceive) {
78 if (!HcomUtil::GetHcomReceiveType(kernel_node, &recv_type)) {
79 MS_LOG(EXCEPTION) << "GetHcomReceiveType fail!";
80 }
81 auto res = std::find(kHcclSupportTypes.begin(), kHcclSupportTypes.end(), recv_type);
82 if (res == kHcclSupportTypes.end()) {
83 MS_LOG(EXCEPTION) << "HcclReceive cannot support data type: " << TypeIdToType(recv_type);
84 }
85 }
86 for (const auto &type : kHcclSupportTypes) {
87 std::vector<std::string> inputs_format{};
88 std::vector<TypeId> inputs_type{};
89 size_t input_num = AnfAlgo::GetInputTensorNum(kernel_node);
90 for (size_t input_index = 0; input_index < input_num; ++input_index) {
91 (void)inputs_format.emplace_back(GetKernelFormat(kernel_node, input_index));
92 inputs_type.push_back(type);
93 }
94 std::vector<std::string> outputs_format;
95 std::vector<TypeId> outputs_type;
96 size_t output_num = AnfAlgo::GetOutputTensorNum(kernel_node);
97 for (size_t output_index = 0; output_index < output_num; ++output_index) {
98 outputs_format.emplace_back(GetKernelFormat(kernel_node, output_index));
99 if (op_name == kReceive) {
100 outputs_type.push_back(recv_type);
101 } else {
102 outputs_type.push_back(type);
103 }
104 }
105 auto builder = KernelBuildInfo::KernelBuildInfoBuilder();
106 builder.SetInputsFormat(inputs_format);
107 builder.SetInputsDeviceType(inputs_type);
108 builder.SetOutputsFormat(outputs_format);
109 builder.SetOutputsDeviceType(outputs_type);
110 builder.SetKernelType(HCCL_KERNEL);
111 kernel_info_list->push_back(builder.Build());
112 }
113 }
114 } // namespace kernel
115 } // namespace mindspore
116