diff --git a/mindspore/ccsrc/device/ascend/kernel_select_ascend.cc b/mindspore/ccsrc/device/ascend/kernel_select_ascend.cc index 06af3d2ca7d..dafe9583484 100644 --- a/mindspore/ccsrc/device/ascend/kernel_select_ascend.cc +++ b/mindspore/ccsrc/device/ascend/kernel_select_ascend.cc @@ -31,12 +31,13 @@ namespace mindspore { namespace device { namespace ascend { namespace { +const float kWegihtBaseScore = 1; +const float kFeatureMapBaseScore = 10; enum MatchCountPriority : int { MATCH_COUNT_PRIORITY_BEGIN = 0, MATCH_DTYPE_COUNT = MATCH_COUNT_PRIORITY_BEGIN, MATCH_FORMAT_COUNT, MATCH_SPECIAL_FORMAT_COUNT, - MATCH_5D_FORMAT_COUNT, MATCH_OUTPUT_DTYPE_COUNT, MATCH_COUNT_PRIORITY_END }; @@ -82,13 +83,6 @@ bool IsValidKernelInfo(const std::shared_ptr &kernel_node, const kernel:: } return true; }; - if (AnfAlgo::GetCNodeName(kernel_node) == "Adam") { - auto input_num = AnfAlgo::GetInputTensorNum(kernel_node); - if (AnfAlgo::GetPrevNodeOutputFormat(kernel_node, input_num - 1) != - kernel_build_info.GetInputFormat(input_num - 1)) { - return false; - } - } if (AnfAlgo::GetCNodeName(kernel_node) == prim::kPrimCast->name()) { return AnfAlgo::GetOutputInferDataType(kernel_node, 0) == kernel_build_info.GetOutputDeviceType(0) && AnfAlgo::GetPrevNodeOutputInferDataType(kernel_node, 0) == kernel_build_info.GetInputDeviceType(0); @@ -112,21 +106,7 @@ bool MatchInferOutputDataType(const CNodePtr &cnode, const kernel::KernelBuildIn MS_EXCEPTION_IF_NULL(cnode); // Check input data type for (size_t input_index = 0; input_index < kernel_build_info.GetInputNum(); ++input_index) { - AnfNodePtr cur_input = AnfAlgo::GetInputNode(cnode, input_index); - MS_EXCEPTION_IF_NULL(cur_input); - TypeId input_origin_type; - if (cur_input->isa() && AnfAlgo::IsParameterWeight(cur_input->cast())) { - // weight - input_origin_type = AnfAlgo::GetOutputDeviceDataType(cur_input, 0); - } else if (cur_input->isa()) { - input_origin_type = AnfAlgo::GetOutputDeviceDataType(cur_input, 0); - } else { - // feature map - input_origin_type = AnfAlgo::GetPrevNodeOutputInferDataType(cnode, input_index); - } - if (input_origin_type == kTypeUnknown) { - continue; - } + TypeId input_origin_type = AnfAlgo::GetPrevNodeOutputInferDataType(cnode, input_index); if (kernel_build_info.GetInputDeviceType(input_index) != input_origin_type) { return false; } @@ -140,6 +120,29 @@ bool MatchInferOutputDataType(const CNodePtr &cnode, const kernel::KernelBuildIn return true; } +string GetPriorityMatchFormat(const CNodePtr &cnode) { + string priority_matched_format = kOpFormat_NC1HWC0; + bool is_init = false; + bool need_change_nd = false; + for (size_t index = 0; index < AnfAlgo::GetInputTensorNum(cnode); ++index) { + auto pre_output_format = AnfAlgo::GetPrevNodeOutputFormat(cnode, index); + if (AnfAlgo::IsFeatureMapInput(cnode, index) && + kNeedTransFormatSet.find(pre_output_format) != kNeedTransFormatSet.end()) { + priority_matched_format = !is_init ? priority_matched_format : pre_output_format; + is_init = true; + } + // feature map has two or more special format; + if (priority_matched_format != pre_output_format && pre_output_format != kOpFormat_DEFAULT) { + priority_matched_format = kOpFormat_DEFAULT; + } + auto input_shape_size = AnfAlgo::GetPrevNodeOutputInferShape(cnode, index).size(); + need_change_nd = (need_change_nd || (input_shape_size != 4 && input_shape_size > 1)); + } + if (need_change_nd) { + priority_matched_format = kOpFormat_DEFAULT; + } + return priority_matched_format; +} /** * compare two vector by priority, select a better vector, like compare two num, first compare highest num location, * if equal then next num location @@ -172,34 +175,18 @@ void UpdateCurMatchCounts(const kernel::KernelBuildInfo &kernel_build_info, cons if (cur_kernelinfo_match_counts->size() < MATCH_COUNT_PRIORITY_END) { MS_LOG(EXCEPTION) << "Out of range cur_kernelinfo_match_counts " << MATCH_COUNT_PRIORITY_END; } + auto pri_match_format = GetPriorityMatchFormat(kernel_node); for (size_t input_index = 0; input_index < AnfAlgo::GetInputTensorNum(kernel_node); ++input_index) { - AnfNodePtr input_anf_node = AnfAlgo::GetInputNode(kernel_node, input_index); - MS_EXCEPTION_IF_NULL(input_anf_node); - // if a input parameter is a weight with default format, the input shouldn't participate the judge - if (input_anf_node->isa()) { - auto para = input_anf_node->cast(); - if (AnfAlgo::IsParameterWeight(para) && AnfAlgo::GetOutputDeviceDataType(para, 0) == kTypeUnknown) { - continue; - } - } + auto base_score = AnfAlgo::IsFeatureMapInput(kernel_node, input_index) ? kFeatureMapBaseScore : kWegihtBaseScore; if (kernel_build_info.GetInputFormat(input_index) == AnfAlgo::GetPrevNodeOutputFormat(kernel_node, input_index)) { - if (AnfAlgo::IsFeatureMapInput(kernel_node, input_index) && - kNeedTransFormatSet.find(kernel_build_info.GetInputFormat(input_index)) != kNeedTransFormatSet.end()) { - (*cur_kernelinfo_match_counts)[MATCH_SPECIAL_FORMAT_COUNT]++; - } - (*cur_kernelinfo_match_counts)[MATCH_FORMAT_COUNT]++; + (*cur_kernelinfo_match_counts)[MATCH_FORMAT_COUNT] += base_score; } if (kernel_build_info.GetInputDeviceType(input_index) == AnfAlgo::GetPrevNodeOutputDeviceDataType(kernel_node, input_index)) { - (*cur_kernelinfo_match_counts)[MATCH_DTYPE_COUNT]++; + (*cur_kernelinfo_match_counts)[MATCH_DTYPE_COUNT] += base_score; } - if (kernel_build_info.GetInputFormat(input_index) == kOpFormat_NC1HWC0) { - // input is from a feature map & this input's shape is not 4d - if (AnfAlgo::IsFeatureMapInput(kernel_node, input_index) && - AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, input_index).size() != kShape4dDims) { - continue; - } - (*cur_kernelinfo_match_counts)[MATCH_5D_FORMAT_COUNT]++; + if (kernel_build_info.GetInputFormat(input_index) == pri_match_format) { + (*cur_kernelinfo_match_counts)[MATCH_SPECIAL_FORMAT_COUNT] += base_score; } } @@ -207,7 +194,7 @@ void UpdateCurMatchCounts(const kernel::KernelBuildInfo &kernel_build_info, cons // cal count of same output dtype between abstract and kernel info if (kernel_build_info.GetOutputDeviceType(output_index) == AnfAlgo::GetOutputInferDataType(kernel_node, output_index)) { - (*cur_kernelinfo_match_counts)[MATCH_OUTPUT_DTYPE_COUNT]++; + (*cur_kernelinfo_match_counts)[MATCH_OUTPUT_DTYPE_COUNT] += 1; } } } @@ -517,7 +504,7 @@ void SelectKernelInfo(const CNodePtr &kernel_node) { std::vector> kernel_info_list; MS_EXCEPTION_IF_NULL(kernel_node); kernel::KernelQuery(kernel_node, &kernel_info_list); - std::vector most_match_counts = {-1, -1, -1, -1, -1}; + std::vector most_match_counts = {-1, -1, -1, -1}; int selected_index = -1; auto context_ptr = MsContext::GetInstance(); MS_EXCEPTION_IF_NULL(context_ptr); @@ -527,7 +514,7 @@ void SelectKernelInfo(const CNodePtr &kernel_node) { std::vector node_mix_precision_datatype_index; std::vector node_mix_precision_datatype; for (size_t info_index = 0; info_index < kernel_info_list.size(); ++info_index) { - std::vector cur_kernel_info_match_counts = {0, 0, 0, 0, 0}; + std::vector cur_kernel_info_match_counts = {0, 0, 0, 0}; auto kernel_build_info = *(kernel_info_list[info_index]); if (!IsValidKernelInfo(kernel_node, kernel_build_info)) { continue;