Browse Source

!1720 add reducemean's special kernel fileter rule

Merge pull request !1720 from lianliguang/r0.3
tags/v0.3.0-alpha
mindspore-ci-bot Gitee 6 years ago
parent
commit
9cb129ac99
3 changed files with 18 additions and 2 deletions
  1. +1
    -1
      mindspore/ccsrc/device/ascend/kernel_select_ascend.cc
  2. +13
    -0
      mindspore/ccsrc/kernel/tbe/tbe_kernel_select.cc
  3. +4
    -1
      mindspore/ccsrc/pre_activate/ascend/ascend_helper.cc

+ 1
- 1
mindspore/ccsrc/device/ascend/kernel_select_ascend.cc View File

@@ -178,7 +178,7 @@ void SetTensorDeviceInfo(const kernel::KernelBuildInfo &selected_kernel_info, co
if (AnfAlgo::GetOutputDeviceDataType(real_input_node, 0) == kTypeUnknown || is_ref) {
std::vector<std::string> output_format = {selected_kernel_info.GetInputFormat(input_index)};
builder->SetOutputsFormat(output_format);
std::vector<TypeId> output_type = {AnfAlgo::GetOutputInferDataType(real_input_node, 0)};
std::vector<TypeId> output_type = {selected_kernel_info.GetInputDeviceType(input_index)};
builder->SetOutputsDeviceType(output_type);
AnfAlgo::SetSelectKernelBuildInfo(builder->Build(), real_input_node.get());
}


+ 13
- 0
mindspore/ccsrc/kernel/tbe/tbe_kernel_select.cc View File

@@ -582,6 +582,7 @@ bool IsShapeMatchFormat(const std::vector<size_t> &shape, const std::string &for

bool IsValidKernelInfo(const std::shared_ptr<CNode> &kernel_node, const kernel::KernelBuildInfo &kernel_build_info) {
MS_EXCEPTION_IF_NULL(kernel_node);
auto kernel_name = AnfAlgo::GetCNodeName(kernel_node);
const size_t kCAxis = 1;
for (size_t index = 0; index < kernel_build_info.GetOutputNum(); ++index) {
auto output_shape = AnfAlgo::GetOutputInferShape(kernel_node, index);
@@ -594,6 +595,12 @@ bool IsValidKernelInfo(const std::shared_ptr<CNode> &kernel_node, const kernel::
if (!IsShapeMatchFormat(output_shape, kernel_build_info.GetOutputFormat(index))) {
return false;
}
if (kernel_name == "ReduceMean") {
auto keep_dims = AnfAlgo::GetNodeAttr<bool>(kernel_node, kAttrKeepDims);
if (keep_dims == false && kernel_build_info.GetOutputFormat(index) != kOpFormat_DEFAULT) {
return false;
}
}
}
for (size_t index = 0; index < kernel_build_info.GetInputNum(); ++index) {
auto input_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, index);
@@ -606,6 +613,12 @@ bool IsValidKernelInfo(const std::shared_ptr<CNode> &kernel_node, const kernel::
}
return false;
}
if (kernel_name == "ReduceMean") {
auto keep_dims = AnfAlgo::GetNodeAttr<bool>(kernel_node, kAttrKeepDims);
if (keep_dims == false && kernel_build_info.GetInputFormat(index) != kOpFormat_DEFAULT) {
return false;
}
}
}
if (AnfAlgo::GetCNodeName(kernel_node) == prim::kPrimCast->name()) {
return AnfAlgo::GetOutputInferDataType(kernel_node, 0) == kernel_build_info.GetOutputDeviceType(0) &&


+ 4
- 1
mindspore/ccsrc/pre_activate/ascend/ascend_helper.cc View File

@@ -315,7 +315,10 @@ CNodePtr InsertCastForInput(const FuncGraphPtr &func_graph, const CNodePtr &cnod
auto cur_input = AnfAlgo::GetInputNode(cnode, input_index);
auto kernel_with_index = AnfAlgo::VisitKernel(cur_input, 0);
auto is_weight_boundary = [](const AnfNodePtr &node) -> bool {
if (node->isa<ValueNode>() || node->isa<Parameter>()) {
if (node->isa<ValueNode>()) {
return true;
}
if (node->isa<Parameter>() && AnfAlgo::IsParameterWeight(node->cast<ParameterPtr>())) {
return true;
}
return false;


Loading…
Cancel
Save