diff --git a/mindspore/ccsrc/pre_activate/ascend/ascend_backend_optimization.cc b/mindspore/ccsrc/pre_activate/ascend/ascend_backend_optimization.cc index 16a82500b43f94acbd71f73567b7d492e693ae6f..0b2af1258aed4b422ace176511360cad19652ad2 100644 --- a/mindspore/ccsrc/pre_activate/ascend/ascend_backend_optimization.cc +++ b/mindspore/ccsrc/pre_activate/ascend/ascend_backend_optimization.cc @@ -49,6 +49,7 @@ #include "pre_activate/ascend/ir_fusion/matmul_biasadd_fusion.h" #include "pre_activate/ascend/ir_fusion/remove_reshape_pair.h" #include "pre_activate/ascend/ir_fusion/derelu_fusion.h" +#include "pre_activate/ascend/ir_fusion/batchnorm_to_bninfer.h" #include "pre_activate/ascend/format_type/insert_trans_op.h" #include "pre_activate/pass/getitem_tuple.h" #include "pre_activate/pass/optimize_dependence.h" @@ -100,6 +101,7 @@ void AddAscendBackendOptionalIRFusion(PassManager *ir_fusion_pm) { ir_fusion_pm->AddPass(std::make_shared()); ir_fusion_pm->AddPass(std::make_shared()); ir_fusion_pm->AddPass(std::make_shared()); + ir_fusion_pm->AddPass(std::make_shared()); } } // namespace diff --git a/mindspore/ccsrc/pre_activate/ascend/ir_fusion/batchnorm_to_bninfer.cc b/mindspore/ccsrc/pre_activate/ascend/ir_fusion/batchnorm_to_bninfer.cc new file mode 100644 index 0000000000000000000000000000000000000000..90c09c9bf97c6038f10077d6bbfe05dd788d931b --- /dev/null +++ b/mindspore/ccsrc/pre_activate/ascend/ir_fusion/batchnorm_to_bninfer.cc @@ -0,0 +1,126 @@ +/** + * Copyright 2020 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pre_activate/ascend/ir_fusion/batchnorm_to_bninfer.h" +#include +#include +#include "session/anf_runtime_algorithm.h" +#include "ir/primitive.h" +#include "utils/utils.h" +#include "operator/ops.h" +#include "pipeline/static_analysis/abstract_value.h" +#include "pre_activate/common/helper.h" + +namespace mindspore { +namespace opt { +namespace { +CNodePtr CreateBNInfer(const FuncGraphPtr &graph, const CNodePtr &batchnorm, const AnfNodePtr &node) { + MS_EXCEPTION_IF_NULL(graph); + MS_EXCEPTION_IF_NULL(batchnorm); + MS_EXCEPTION_IF_NULL(node); + auto prim = std::make_shared(kBNInferOpName); + std::vector inputs = {NewValueNode(prim)}; + for (size_t i = 1; i < batchnorm->size(); ++i) { + inputs.push_back(batchnorm->input(i)); + } + auto new_node = graph->NewCNode(inputs); + MS_EXCEPTION_IF_NULL(new_node); + new_node->set_scope(batchnorm->scope()); + new_node->set_abstract(node->abstract()); + AnfAlgo::CopyNodeAttr(kAttrIsTraining, batchnorm, new_node); + AnfAlgo::CopyNodeAttr(kAttrEpsilon, batchnorm, new_node); + return new_node; +} + +bool CheckIndex(const AnfNodePtr &index_node) { + MS_EXCEPTION_IF_NULL(index_node); + if (!IsValueNode(index_node)) { + return false; + } + ValueNodePtr value_node = index_node->cast(); + MS_EXCEPTION_IF_NULL(value_node); + int index = GetValue(value_node->value()); + if (index != 0) { + MS_LOG(DEBUG) << "tuple_getitem must be 0th output of BatchNorm"; + return false; + } + return true; +} + +bool CheckBatchNorm(const FuncGraphPtr &graph, const CNodePtr &batchnorm) { + MS_EXCEPTION_IF_NULL(graph); + MS_EXCEPTION_IF_NULL(batchnorm); + if (batchnorm->size() < kBatchNormInputNum + 1) { + MS_LOG(DEBUG) << "BatchNorm's input less than " << kBatchNormInputNum; + return false; + } + if (!AnfAlgo::HasNodeAttr(kAttrIsTraining, batchnorm)) { + return false; + } + auto is_training = AnfAlgo::GetNodeAttr(batchnorm, kAttrIsTraining); + if (is_training) { + MS_LOG(DEBUG) << "is_training is true, no need do fusion"; + return false; + } + + if (IsUsedByOthers(graph, batchnorm)) { + MS_LOG(DEBUG) << "Only the 0th output of BatchNorm is used, then do fusion"; + return false; + } + return true; +} + +bool NeedFusion(const FuncGraphPtr &graph, const AnfNodePtr &node, CNodePtr *batchnorm) { + MS_EXCEPTION_IF_NULL(graph); + MS_EXCEPTION_IF_NULL(node); + auto tuple_getitem = node->cast(); + MS_EXCEPTION_IF_NULL(tuple_getitem); + CheckCNodeInputSize(tuple_getitem, kTupleGetItemInputSize); + AnfNodePtr index_node = tuple_getitem->input(kInputNodeOutputIndexInTupleGetItem); + MS_EXCEPTION_IF_NULL(index_node); + if (!CheckIndex(index_node)) { + return false; + } + + AnfNodePtr batchnorm_anf = tuple_getitem->input(kRealInputNodeIndexInTupleGetItem); + MS_EXCEPTION_IF_NULL(batchnorm_anf); + *batchnorm = batchnorm_anf->cast(); + MS_EXCEPTION_IF_NULL(*batchnorm); + return CheckBatchNorm(graph, *batchnorm); +} +} // namespace + +const BaseRef BatchNorm2BNInfer::DefinePattern() const { + VarPtr Xs = std::make_shared(); + VarPtr Y = std::make_shared(); + MS_EXCEPTION_IF_NULL(Xs); + MS_EXCEPTION_IF_NULL(Y); + VectorRef batchnorm({prim::kPrimBatchNorm, Xs}); + VectorRef pattern({prim::kPrimTupleGetItem, batchnorm, Y}); + return pattern; +} + +const AnfNodePtr BatchNorm2BNInfer::Process(const FuncGraphPtr &graph, const AnfNodePtr &node, const EquivPtr &) const { + MS_EXCEPTION_IF_NULL(graph); + MS_EXCEPTION_IF_NULL(node); + + CNodePtr batchnorm = nullptr; + if (!NeedFusion(graph, node, &batchnorm)) { + return nullptr; + } + return CreateBNInfer(graph, batchnorm, node); +} +} // namespace opt +} // namespace mindspore diff --git a/mindspore/ccsrc/pre_activate/ascend/ir_fusion/batchnorm_to_bninfer.h b/mindspore/ccsrc/pre_activate/ascend/ir_fusion/batchnorm_to_bninfer.h new file mode 100644 index 0000000000000000000000000000000000000000..551fe0f6f905e33e95cb28b99d81906727d582d5 --- /dev/null +++ b/mindspore/ccsrc/pre_activate/ascend/ir_fusion/batchnorm_to_bninfer.h @@ -0,0 +1,33 @@ +/** + * Copyright 2020 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef MINDSPORE_CCSRC_PRE_ACTIVATE_ASCEND_IR_FUSION_BATCHNORM_TO_BNINFER_H_ +#define MINDSPORE_CCSRC_PRE_ACTIVATE_ASCEND_IR_FUSION_BATCHNORM_TO_BNINFER_H_ + +#include +#include "pre_activate/common/optimizer.h" + +namespace mindspore { +namespace opt { +class BatchNorm2BNInfer : public PatternProcessPass { + public: + explicit BatchNorm2BNInfer(bool multigraph = true) : PatternProcessPass("batchnorm_to_bninfer", multigraph) {} + ~BatchNorm2BNInfer() override = default; + const BaseRef DefinePattern() const override; + const AnfNodePtr Process(const FuncGraphPtr &, const AnfNodePtr &, const EquivPtr &) const override; +}; +} // namespace opt +} // namespace mindspore +#endif // MINDSPORE_CCSRC_PRE_ACTIVATE_ASCEND_IR_FUSION_BATCHNORM_TO_BNINFER_H_ diff --git a/mindspore/ccsrc/utils/utils.h b/mindspore/ccsrc/utils/utils.h index 5b8a8b178e3a44132cabce66ea80ecf7e6c83c98..015b371f884ad0de4f74e8c9f83bde0a5d378544 100644 --- a/mindspore/ccsrc/utils/utils.h +++ b/mindspore/ccsrc/utils/utils.h @@ -108,6 +108,7 @@ constexpr auto kLambNextMVOpName = "LambNextMV"; constexpr auto kConfusionTransposeDOpName = "ConfusionTransposeD"; constexpr auto kAdamApplyOneWithDecayOpName = "AdamApplyOneWithDecay"; constexpr auto kBatchNormGradOpName = "BatchNormGrad"; +constexpr auto kBNInferOpName = "BNInfer"; constexpr auto kAdamApplyOneOpName = "AdamApplyOne"; constexpr auto kResizeNearestNeighborGradOpName = "ResizeNearestNeighborGrad"; constexpr auto kFusedMulAddOpName = "FusedMulAdd"; diff --git a/tests/ut/cpp/pre_activate/ascend/ir_fusion/batchnorm_to_bninfer_test.cc b/tests/ut/cpp/pre_activate/ascend/ir_fusion/batchnorm_to_bninfer_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..466cba8e6727240a7afeb97c58838b6be5ee8c59 --- /dev/null +++ b/tests/ut/cpp/pre_activate/ascend/ir_fusion/batchnorm_to_bninfer_test.cc @@ -0,0 +1,72 @@ +/** + * Copyright 2020 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "common/backend_common_test.h" +#include "common/py_func_graph_fetcher.h" +#include "pre_activate/common/optimizer.h" +#include "pre_activate/ascend/ir_fusion/batchnorm_to_bninfer.h" +#include "debug/anf_ir_dump.h" + +namespace mindspore { +namespace opt { +class TestHWOptimizeBatchNorm2BNInfer : public BackendCommon { + public: + TestHWOptimizeBatchNorm2BNInfer() : get_py_fun_("gtest_input.pre_activate.batchnorm_to_bninfer", true) {} + ~TestHWOptimizeBatchNorm2BNInfer() override = default; + + UT::PyFuncGraphFetcher get_py_fun_; +}; + +TEST_F(TestHWOptimizeBatchNorm2BNInfer, test_fusion) { + FuncGraphPtr g = get_py_fun_.CallAndParseRet("test_batchnorm_to_bninfer", "before"); + EXPECT_NE(g, nullptr); + std::vector shp_x{32, 64, 112, 112}; + auto x_abstract = std::make_shared(kFloat32, shp_x); + std::vector shp_y{64}; + auto y_abstract = std::make_shared(kFloat32, shp_y); + AbstractBasePtrList args_spec_list{x_abstract, y_abstract, y_abstract, y_abstract, y_abstract}; + auto fg = GetKernelGraph(g, args_spec_list); + + auto optimizer = std::make_shared(); + auto pm = std::make_shared(); + pm->AddPass(std::make_shared()); + optimizer->AddPassManager(pm); + FuncGraphPtr new_graph = optimizer->Optimize(fg); + + FuncGraphPtr g_after = get_py_fun_.CallAndParseRet("test_batchnorm_to_bninfer", "after"); + EXPECT_TRUE(CheckEqualGraph(g_after, new_graph)); +} + +TEST_F(TestHWOptimizeBatchNorm2BNInfer, test_no_fusion) { + FuncGraphPtr g = get_py_fun_.CallAndParseRet("test_batchnorm_to_bninfer", "no_fusion"); + EXPECT_NE(g, nullptr); + std::vector shp_x{32, 64, 112, 112}; + auto x_abstract = std::make_shared(kFloat32, shp_x); + std::vector shp_y{64}; + auto y_abstract = std::make_shared(kFloat32, shp_y); + AbstractBasePtrList args_spec_list{x_abstract, y_abstract, y_abstract, y_abstract, y_abstract}; + auto fg = GetKernelGraph(g, args_spec_list); + auto origin_graph = std::make_shared(*fg); + + auto optimizer = std::make_shared(); + auto pm = std::make_shared(); + pm->AddPass(std::make_shared()); + optimizer->AddPassManager(pm); + FuncGraphPtr new_graph = optimizer->Optimize(fg); + + EXPECT_TRUE(CheckEqualGraph(origin_graph, new_graph)); +} +} // namespace opt +} // namespace mindspore diff --git a/tests/ut/cpp/python_input/gtest_input/pre_activate/batchnorm_to_bninfer.py b/tests/ut/cpp/python_input/gtest_input/pre_activate/batchnorm_to_bninfer.py new file mode 100644 index 0000000000000000000000000000000000000000..bafe21dddeb9a7cc05ec0a1ced0d4c9d2049355c --- /dev/null +++ b/tests/ut/cpp/python_input/gtest_input/pre_activate/batchnorm_to_bninfer.py @@ -0,0 +1,56 @@ +# Copyright 2020 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ + +from mindspore.ops import operations as P +from mindspore.ops import Primitive + +batch_norm = P.BatchNorm(is_training=False) +bn_infer = Primitive('BNInfer') +make_tuple = Primitive('make_tuple') +tuple_getitem = Primitive('tuple_getitem') + +class FnDict: + def __init__(self): + self.fnDict = {} + + def __call__(self, fn): + self.fnDict[fn.__name__] = fn + + def __getitem__(self, name): + return self.fnDict[name] + +def test_batchnorm_to_bninfer(tag): + fns = FnDict() + + @fns + def before(input0, input1, input2, input3, input4): + res = batch_norm(input0, input1, input2, input3, input4) + res = tuple_getitem(res, 0) + return res + + @fns + def after(input0, input1, input2, input3, input4): + res = bn_infer(input0, input1, input2, input3, input4) + return make_tuple(res) + + @fns + def no_fusion(input0, input1, input2, input3, input4): + res = batch_norm(input0, input1, input2, input3, input4) + item0 = tuple_getitem(res, 0) + item1 = tuple_getitem(res, 1) + item2 = tuple_getitem(res, 2) + return make_tuple(item0, item1, item2) + + return fns[tag]