diff --git a/mindspore/ccsrc/pre_activate/ascend/ascend_backend_optimization.cc b/mindspore/ccsrc/pre_activate/ascend/ascend_backend_optimization.cc index 4294f48e47..b71faae7ba 100644 --- a/mindspore/ccsrc/pre_activate/ascend/ascend_backend_optimization.cc +++ b/mindspore/ccsrc/pre_activate/ascend/ascend_backend_optimization.cc @@ -20,6 +20,7 @@ #include "pre_activate/ascend/ir_fission/bn_split.h" #include "pre_activate/ascend/ir_fission/bn_grad_split.h" #include "pre_activate/ascend/ir_fission/batch_norm_grad_split.h" +#include "pre_activate/ascend/ir_fission/batch_norm_bert_fission.h" #include "pre_activate/ascend/ir_fusion/fused_batch_norm_fusion.h" #include "pre_activate/ascend/ir_fission/layer_norm_grad_split.h" #include "pre_activate/pass/communication_op_fusion.h" @@ -76,6 +77,7 @@ namespace opt { namespace { void AddAscendBackendOptionalIRFusion(PassManager *ir_fusion_pm) { MS_EXCEPTION_IF_NULL(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()); diff --git a/mindspore/ccsrc/pre_activate/ascend/ir_fission/batch_norm_bert_fission.cc b/mindspore/ccsrc/pre_activate/ascend/ir_fission/batch_norm_bert_fission.cc new file mode 100644 index 0000000000..0b023c3691 --- /dev/null +++ b/mindspore/ccsrc/pre_activate/ascend/ir_fission/batch_norm_bert_fission.cc @@ -0,0 +1,170 @@ +/** + * 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_fission/batch_norm_bert_fission.h" +#include +#include +#include +#include "session/anf_runtime_algorithm.h" +#include "pre_activate/common/helper.h" + +namespace mindspore { +namespace opt { +namespace { +const std::vector kOutputIndex{0, 3, 4, 5}; +constexpr size_t kBatchNormRealOutputNum = 3; + +bool CompareTupleGetitem(const AnfNodePtr &n1, const AnfNodePtr &n2) { + MS_EXCEPTION_IF_NULL(n1); + MS_EXCEPTION_IF_NULL(n2); + auto n1_cnode = n1->cast(); + auto n2_cnode = n2->cast(); + MS_EXCEPTION_IF_NULL(n1_cnode); + MS_EXCEPTION_IF_NULL(n2_cnode); + auto index_input1 = n1_cnode->input(kInputNodeOutputIndexInTupleGetItem); + MS_EXCEPTION_IF_NULL(index_input1); + auto value_node1 = index_input1->cast(); + MS_EXCEPTION_IF_NULL(value_node1); + auto index_input2 = n2_cnode->input(kInputNodeOutputIndexInTupleGetItem); + MS_EXCEPTION_IF_NULL(index_input2); + auto value_node2 = index_input2->cast(); + MS_EXCEPTION_IF_NULL(value_node2); + return GetValue(value_node1->value()) < GetValue(value_node2->value()); +} + +bool GetBatchNormOutputs(const FuncGraphPtr &func_graph, const AnfNodePtr &bn, std::vector *bn_outputs) { + MS_EXCEPTION_IF_NULL(func_graph); + MS_EXCEPTION_IF_NULL(bn_outputs); + auto manager = func_graph->manager(); + MS_EXCEPTION_IF_NULL(manager); + if (manager->node_users().find(bn) == manager->node_users().end()) { + return false; + } + size_t output_num = 0; + for (const auto &node_index : manager->node_users()[bn]) { + AnfNodePtr output = node_index.first; + MS_EXCEPTION_IF_NULL(output); + auto tuple_getiterm_cnode = output->cast(); + MS_EXCEPTION_IF_NULL(tuple_getiterm_cnode); + auto index_node = tuple_getiterm_cnode->input(kInputNodeOutputIndexInTupleGetItem); + MS_EXCEPTION_IF_NULL(index_node); + auto value_node = index_node->cast(); + MS_EXCEPTION_IF_NULL(value_node); + int index = GetValue(value_node->value()); + if (std::find(kOutputIndex.begin(), kOutputIndex.end(), index) == kOutputIndex.end()) { + return false; + } + bn_outputs->push_back(output); + output_num++; + } + return output_num == kBatchNormRealOutputNum; +} + +AnfNodePtr CreateBNTrainingReduce(const FuncGraphPtr &func_graph, const AnfNodePtr &bn) { + MS_EXCEPTION_IF_NULL(func_graph); + MS_EXCEPTION_IF_NULL(bn); + auto bn_cnode = bn->cast(); + MS_EXCEPTION_IF_NULL(bn_cnode); + CheckCNodeInputSize(bn_cnode, kBatchNormInputNum + 1); + std::vector bn_training_reduce_inputs = { + NewValueNode(std::make_shared(kBNTrainingReduceOpName)), bn_cnode->input(1)}; + auto bn_training_reduce = func_graph->NewCNode(bn_training_reduce_inputs); + MS_EXCEPTION_IF_NULL(bn_training_reduce); + auto bn_input1 = bn_cnode->input(2); + MS_EXCEPTION_IF_NULL(bn_input1); + auto bn_input2 = bn_cnode->input(3); + MS_EXCEPTION_IF_NULL(bn_input2); + AbstractBasePtrList abstract_list{bn_input1->abstract(), bn_input2->abstract()}; + auto abstract_tuple = std::make_shared(abstract_list); + bn_training_reduce->set_abstract(abstract_tuple); + bn_training_reduce->set_scope(bn->scope()); + AnfAlgo::CopyNodeAttrs(bn, bn_training_reduce); + return bn_training_reduce; +} + +AnfNodePtr CreateBNTrainingUpdateV2(const FuncGraphPtr &func_graph, const AnfNodePtr &bn, + const std::vector &bn_training_reduce_outputs) { + MS_EXCEPTION_IF_NULL(func_graph); + MS_EXCEPTION_IF_NULL(bn); + auto bn_cnode = bn->cast(); + MS_EXCEPTION_IF_NULL(bn_cnode); + CheckCNodeInputSize(bn_cnode, kBatchNormInputNum + 1); + if (bn_training_reduce_outputs.size() != kBNTrainingReduceOutputNum) { + MS_LOG(EXCEPTION) << "The output size of node bn_training_reduce must be " << kBNTrainingReduceOutputNum + << ", but it is " << bn_training_reduce_outputs.size(); + } + std::vector bn_training_update_v2_inputs = { + NewValueNode(std::make_shared(kBNTrainingUpdateV2OpName)), + bn_cnode->input(1), + bn_training_reduce_outputs[0], + bn_training_reduce_outputs[1], + bn_cnode->input(2), + bn_cnode->input(3)}; + auto bn_training_update_v2 = func_graph->NewCNode(bn_training_update_v2_inputs); + MS_EXCEPTION_IF_NULL(bn_training_update_v2); + + auto bn_abstract_tuple = dyn_cast(bn->abstract()); + MS_EXCEPTION_IF_NULL(bn_abstract_tuple); + if (bn_abstract_tuple->elements().size() != kBatchNormOutputNum) { + MS_LOG(EXCEPTION) << "The abstract size of node bn must be " << kBatchNormOutputNum << ", but it is " + << bn_abstract_tuple->elements().size(); + } + std::vector abstract_list{bn_abstract_tuple->elements()[0], bn_abstract_tuple->elements()[3], + bn_abstract_tuple->elements()[4]}; + auto abstract_tuple = std::make_shared(abstract_list); + bn_training_update_v2->set_abstract(abstract_tuple); + bn_training_update_v2->set_scope(bn->scope()); + AnfAlgo::CopyNodeAttrs(bn, bn_training_update_v2); + return bn_training_update_v2; +} +} // namespace + +const BaseRef BatchNormBertFission::DefinePattern() const { + VarPtr Xs = std::make_shared(); + return VectorRef({prim::kPrimBatchNorm, Xs}); +} + +const AnfNodePtr BatchNormBertFission::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node, + const EquivPtr &) const { + MS_EXCEPTION_IF_NULL(func_graph); + std::vector bn_outputs; + if (!GetBatchNormOutputs(func_graph, node, &bn_outputs)) { + return nullptr; + } + AnfNodePtr bn_training_reduce = CreateBNTrainingReduce(func_graph, node); + std::vector bn_training_reduce_outputs; + CreateMultipleOutputsOfAnfNode(func_graph, bn_training_reduce, kBNTrainingReduceOutputNum, + &bn_training_reduce_outputs); + + AnfNodePtr bn_training_update_v2 = CreateBNTrainingUpdateV2(func_graph, node, bn_training_reduce_outputs); + std::vector bn_training_update_v2_outputs; + CreateMultipleOutputsOfAnfNode(func_graph, bn_training_update_v2, kBNTrainingUpdateV2OutputNum, + &bn_training_update_v2_outputs); + if (bn_training_update_v2_outputs.size() != kBNTrainingUpdateV2OutputNum) { + MS_LOG(EXCEPTION) << "The output size of node bn_training_reduce must be " << kBNTrainingUpdateV2OutputNum + << ", but it is " << bn_training_update_v2_outputs.size(); + } + auto manager = func_graph->manager(); + MS_EXCEPTION_IF_NULL(manager); + sort(bn_outputs.begin(), bn_outputs.end(), CompareTupleGetitem); + size_t output_index = 0; + for (const auto &output : bn_outputs) { + (void)manager->Replace(output, bn_training_update_v2_outputs[output_index]); + output_index++; + } + return nullptr; +} +} // namespace opt +} // namespace mindspore diff --git a/mindspore/ccsrc/pre_activate/ascend/ir_fission/batch_norm_bert_fission.h b/mindspore/ccsrc/pre_activate/ascend/ir_fission/batch_norm_bert_fission.h new file mode 100644 index 0000000000..fc214817fc --- /dev/null +++ b/mindspore/ccsrc/pre_activate/ascend/ir_fission/batch_norm_bert_fission.h @@ -0,0 +1,32 @@ +/** + * 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_FISSION_BATCH_NORM_BERT_FISSION_H_ +#define MINDSPORE_CCSRC_PRE_ACTIVATE_ASCEND_IR_FISSION_BATCH_NORM_BERT_FISSION_H_ + +#include "pre_activate/common/optimizer.h" + +namespace mindspore { +namespace opt { +class BatchNormBertFission : public PatternProcessPass { + public: + explicit BatchNormBertFission(bool multigraph = true) : PatternProcessPass("batch_norm_bert_fission", multigraph) {} + ~BatchNormBertFission() 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_FISSION_BATCH_NORM_BERT_FISSION_H_ diff --git a/mindspore/ccsrc/pre_activate/common/helper.h b/mindspore/ccsrc/pre_activate/common/helper.h index 9ef57d8e7c..1e27db132e 100644 --- a/mindspore/ccsrc/pre_activate/common/helper.h +++ b/mindspore/ccsrc/pre_activate/common/helper.h @@ -47,6 +47,8 @@ constexpr size_t kBn2ReluOutputNum = 4; constexpr size_t kBnInputNum = 6; constexpr size_t kBnOutputNum = 5; +constexpr size_t kBatchNormInputNum = 5; +constexpr size_t kBatchNormOutputNum = 5; constexpr size_t kBN1OutputNum = 2; constexpr size_t kBN2OutputNum = 3; @@ -61,6 +63,7 @@ constexpr size_t kBNGrad3OutputNum = 1; constexpr size_t kBNTrainingReduceOutputNum = 2; constexpr size_t kBNTrainingUpdateOutputNum = 5; +constexpr size_t kBNTrainingUpdateV2OutputNum = 3; constexpr size_t kBNTrainingUpdateGradOutputNum = 2; constexpr size_t kSingleOutputNum = 1; diff --git a/mindspore/ccsrc/utils/utils.h b/mindspore/ccsrc/utils/utils.h index 9fb62a5470..588dfa1405 100644 --- a/mindspore/ccsrc/utils/utils.h +++ b/mindspore/ccsrc/utils/utils.h @@ -52,6 +52,7 @@ constexpr auto kTopKOpName = "TopK"; constexpr auto kExtractImagePatchesOpName = "ExtractImagePatches"; constexpr auto kBNTrainingReduceOpName = "BNTrainingReduce"; constexpr auto kBNTrainingUpdateOpName = "BNTrainingUpdate"; +constexpr auto kBNTrainingUpdateV2OpName = "BNTrainingUpdateV2"; constexpr auto kSimpleMeanGradOpName = "SimpleMeanGrad"; constexpr auto kMeanGradOpName = "MeanGrad"; constexpr auto kSliceOpName = "Slice"; diff --git a/tests/ut/cpp/pre_activate/ascend/ir_fission/batch_norm_bert_fission_test.cc b/tests/ut/cpp/pre_activate/ascend/ir_fission/batch_norm_bert_fission_test.cc new file mode 100644 index 0000000000..e5abf56c2e --- /dev/null +++ b/tests/ut/cpp/pre_activate/ascend/ir_fission/batch_norm_bert_fission_test.cc @@ -0,0 +1,54 @@ +/** + * 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_fission/batch_norm_bert_fission.h" +#include "common/backend_common_test.h" +#include "common/py_func_graph_fetcher.h" + +namespace mindspore { +namespace opt { +class TestHWBatchNormBertFission : public BackendCommon { + public: + TestHWBatchNormBertFission() : get_py_fun_("gtest_input.pre_activate.batch_norm_bert_fission_test", true) {} + ~TestHWBatchNormBertFission() override = default; + + UT::PyFuncGraphFetcher get_py_fun_; +}; + +TEST_F(TestHWBatchNormBertFission, test_fused_batch_norm_fusion) { + FuncGraphPtr g = get_py_fun_.CallAndParseRet("test_batch_norm_bert_fission", "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}; + for (size_t i = 0; i < 4; ++i) { + args_spec_list.push_back(y_abstract); + } + auto kg = 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(kg); + + FuncGraphPtr g_after = get_py_fun_.CallAndParseRet("test_batch_norm_bert_fission", "after"); + EXPECT_TRUE(CheckEqualGraph(g_after, new_graph)); +} +} // namespace opt +} // namespace mindspore diff --git a/tests/ut/cpp/python_input/gtest_input/pre_activate/batch_norm_bert_fission_test.py b/tests/ut/cpp/python_input/gtest_input/pre_activate/batch_norm_bert_fission_test.py new file mode 100644 index 0000000000..21c1409d36 --- /dev/null +++ b/tests/ut/cpp/python_input/gtest_input/pre_activate/batch_norm_bert_fission_test.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 + +make_tuple = Primitive('make_tuple') +tuple_getitem = Primitive('tuple_getitem') +BatchNorm = P.BatchNorm() +BNTrainingReduce = Primitive('BNTrainingReduce') +BNTrainingUpdateV2 = Primitive('BNTrainingUpdateV2') + + +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_batch_norm_bert_fission(tag): + fns = FnDict() + + @fns + def before(input0, input1, input2, input3, input4): + batch_norm = BatchNorm(input0, input1, input2, input3, input4) + outputs = make_tuple(tuple_getitem(batch_norm, 0), tuple_getitem(batch_norm, 3), tuple_getitem(batch_norm, 4)) + output = tuple_getitem(outputs, 0) + return output + + @fns + def after(input0, input1, input2, input3, input4): + bn_training_reduce = BNTrainingReduce(input0) + bn_training_update_v2 = BNTrainingUpdateV2(input0, tuple_getitem(bn_training_reduce, 0), + tuple_getitem(bn_training_reduce, 1), input1, input2) + outputs = make_tuple(tuple_getitem(bn_training_update_v2, 0), tuple_getitem(bn_training_update_v2, 1), + tuple_getitem(bn_training_update_v2, 2)) + output = tuple_getitem(outputs, 0) + return make_tuple(output) + + return fns[tag]