!30988 add Fracz depthwise case in TransDataSplit

Merge pull request !30988 from yuchaojie/op_select
This commit is contained in:
i-robot 2022-03-10 07:19:34 +00:00 committed by Gitee
commit afb33fff62
No known key found for this signature in database
GPG Key ID: 173E9B9CA92EEF8F
1 changed files with 26 additions and 1 deletions

View File

@ -19,11 +19,35 @@
namespace mindspore {
namespace opt {
namespace {
constexpr size_t kNchwDimNum = 4;
constexpr size_t kDimC = 1;
const std::set<std::pair<string, string>> invalid_formats_pair = {
{kOpFormat_C1HWNCoC0, kOpFormat_NCHW}, {kOpFormat_NCHW, kOpFormat_C1HWNCoC0},
{kOpFormat_C1HWNCoC0, kOpFormat_DEFAULT}, {kOpFormat_DEFAULT, kOpFormat_FRACTAL_ZN_LSTM},
{kOpFormat_FRACTAL_ZN_LSTM, kOpFormat_DEFAULT}, {kOpFormat_DEFAULT, kOpFormat_C1HWNCoC0}};
bool IsDepthwiseCase(const CNodePtr &node, const std::string &input_format, const std::string &output_format) {
MS_EXCEPTION_IF_NULL(node);
abstract::BaseShapePtr base_shape;
if (input_format == kOpFormat_FRAC_Z && output_format == kOpFormat_DEFAULT) {
base_shape = common::AnfAlgo::GetPrevNodeOutputDetailShape(node, 0);
} else if (input_format == kOpFormat_DEFAULT && output_format == kOpFormat_FRAC_Z) {
base_shape = common::AnfAlgo::GetOutputDetailShape(node, 0);
} else {
return false;
}
MS_EXCEPTION_IF_NULL(base_shape);
if (base_shape->isa<abstract::Shape>()) {
auto shape_ptr = base_shape->cast<abstract::ShapePtr>();
auto shape_vec = shape_ptr->shape();
return shape_vec.size() == kNchwDimNum && shape_vec[kDimC] == 1;
}
return false;
}
} // namespace
const AnfNodePtr TransDataSplit::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
const EquivPtr &) const {
MS_EXCEPTION_IF_NULL(func_graph);
@ -45,7 +69,8 @@ bool TransDataSplit::IsFormatInvaild(const AnfNodePtr &node) const {
auto output_format = AnfAlgo::GetOutputFormat(node, 0);
auto format_pair = std::make_pair(input_format, output_format);
return invalid_formats_pair.find(format_pair) != invalid_formats_pair.end();
return invalid_formats_pair.find(format_pair) != invalid_formats_pair.end() ||
IsDepthwiseCase(cnode, input_format, output_format);
}
const BaseRef TransDataSplit::DefinePattern() const {