diff --git a/mindspore/ccsrc/pre_activate/ascend/ascend_backend_optimization.cc b/mindspore/ccsrc/pre_activate/ascend/ascend_backend_optimization.cc index 66ea5ee52626025e7ebf1cb20adf1f9617b4dddf..4294f48e47c944524fa609536bc1fdb5a29b3d96 100644 --- a/mindspore/ccsrc/pre_activate/ascend/ascend_backend_optimization.cc +++ b/mindspore/ccsrc/pre_activate/ascend/ascend_backend_optimization.cc @@ -19,6 +19,7 @@ #include "pre_activate/common/optimizer.h" #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_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" @@ -87,7 +88,6 @@ 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()); ir_fusion_pm->AddPass(std::make_shared()); ir_fusion_pm->AddPass(std::make_shared()); ir_fusion_pm->AddPass(std::make_shared()); @@ -193,8 +193,8 @@ void AscendBackendIRFusionOptimization(const std::shared_ptr(); auto ir_fusion_pm = std::make_shared("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()); ir_fusion_pm->AddPass(std::make_shared()); if (context_ptr->ir_fusion_flag()) { AddAscendBackendOptionalIRFusion(ir_fusion_pm.get()); diff --git a/mindspore/ccsrc/pre_activate/ascend/ir_fusion/fused_batch_norm_fusion.cc b/mindspore/ccsrc/pre_activate/ascend/ir_fusion/fused_batch_norm_fusion.cc index 12f2684b3b6aaa25a76381689b244d9c34c2a4f0..7641772d7a2588eb50ccf33b4faef6d0c42337df 100644 --- a/mindspore/ccsrc/pre_activate/ascend/ir_fusion/fused_batch_norm_fusion.cc +++ b/mindspore/ccsrc/pre_activate/ascend/ir_fusion/fused_batch_norm_fusion.cc @@ -23,6 +23,8 @@ namespace mindspore { namespace opt { namespace { +constexpr size_t kReplaceOutputIndex0 = 3; +constexpr size_t kReplaceOutputIndex1 = 4; bool IsC(const BaseRef &n) { if (utils::isa(n)) { AnfNodePtr in = utils::cast(n); @@ -32,52 +34,6 @@ bool IsC(const BaseRef &n) { return false; } -AnfNodePtr GetBatchNormNode(const AnfNodePtr &node) { - MS_EXCEPTION_IF_NULL(node); - auto depend_cnode = node->cast(); - MS_EXCEPTION_IF_NULL(depend_cnode); - CheckCNodeInputSize(depend_cnode, kDependInputNum); - AnfNodePtr assign_sub = depend_cnode->input(2); - MS_EXCEPTION_IF_NULL(assign_sub); - auto assign_sub_cnode = assign_sub->cast(); - MS_EXCEPTION_IF_NULL(assign_sub_cnode); - CheckCNodeInputSize(assign_sub_cnode, kAssignSubInputNum); - AnfNodePtr mul = assign_sub_cnode->input(2); - MS_EXCEPTION_IF_NULL(mul); - auto mul_cnode = mul->cast(); - MS_EXCEPTION_IF_NULL(mul_cnode); - CheckCNodeInputSize(mul_cnode, kMulInputNum); - AnfNodePtr sub = mul_cnode->input(1); - MS_EXCEPTION_IF_NULL(sub); - auto sub_cnode = sub->cast(); - MS_EXCEPTION_IF_NULL(sub_cnode); - CheckCNodeInputSize(sub_cnode, kSubInputNum); - AnfNodePtr tuple_getitem = sub_cnode->input(2); - MS_EXCEPTION_IF_NULL(tuple_getitem); - auto tuple_getitem_cnode = tuple_getitem->cast(); - MS_EXCEPTION_IF_NULL(tuple_getitem_cnode); - CheckCNodeInputSize(tuple_getitem_cnode, kTupleGetitemInputNum); - return tuple_getitem_cnode->input(1); -} - -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()); -} - void GetBNOutput(const FuncGraphPtr &func_graph, const AnfNodePtr &bn, std::vector *bn_outputs) { MS_EXCEPTION_IF_NULL(func_graph); MS_EXCEPTION_IF_NULL(bn); @@ -92,54 +48,35 @@ void GetBNOutput(const FuncGraphPtr &func_graph, const AnfNodePtr &bn, std::vect MS_EXCEPTION_IF_NULL(output); bn_outputs->push_back(output); } - sort(bn_outputs->begin(), bn_outputs->end(), CompareTupleGetitem); } } // namespace const BaseRef FusedBatchNormFusion::DefinePattern() const { - const auto prim_batch_norm = std::make_shared(kBatchNormOpName); std::shared_ptr Xs = std::make_shared(); VarPtr index0 = std::make_shared(IsC); VarPtr index1 = std::make_shared(IsC); VarPtr index2 = std::make_shared(IsC); - VectorRef batch_norm = VectorRef({prim_batch_norm, data_input_var0_, data_input_var1_, data_input_var2_, Xs}); + VectorRef batch_norm = VectorRef({batch_norm_var_, data_input0_var_, data_input1_var_, data_input2_var_, Xs}); VectorRef tuple_getitem0 = VectorRef({prim::kPrimTupleGetItem, batch_norm, index0}); VectorRef tuple_getitem1 = VectorRef({prim::kPrimTupleGetItem, batch_norm, index1}); VectorRef tuple_getitem2 = VectorRef({prim::kPrimTupleGetItem, batch_norm, index2}); - VectorRef sub0 = VectorRef({prim::kPrimSub, variable_input_var0_, tuple_getitem1}); - VectorRef sub1 = VectorRef({prim::kPrimSub, variable_input_var1_, tuple_getitem2}); - VectorRef mul0 = VectorRef({prim::kPrimMul, sub0, constant_input_var0_}); - VectorRef mul1 = VectorRef({prim::kPrimMul, sub1, constant_input_var1_}); - VectorRef assign_sub0 = VectorRef({prim::kPrimAssignSub, variable_input_var0_, mul0}); - VectorRef assign_sub1 = VectorRef({prim::kPrimAssignSub, variable_input_var1_, mul1}); + VectorRef sub0 = VectorRef({prim::kPrimSub, variable_input0_var_, tuple_getitem1}); + VectorRef sub1 = VectorRef({prim::kPrimSub, variable_input1_var_, tuple_getitem2}); + VectorRef mul0 = VectorRef({prim::kPrimMul, sub0, constant_input0_var_}); + VectorRef mul1 = VectorRef({prim::kPrimMul, sub1, constant_input1_var_}); + VectorRef assign_sub0 = VectorRef({prim::kPrimAssignSub, variable_input0_var_, mul0}); + VectorRef assign_sub1 = VectorRef({prim::kPrimAssignSub, variable_input1_var_, mul1}); VectorRef depend0 = VectorRef({prim::kPrimDepend, tuple_getitem0, assign_sub0}); return VectorRef({prim::kPrimDepend, depend0, assign_sub1}); } -abstract::AbstractTuplePtr FusedBatchNormFusion::CreateAbstractOfFusedBatchNorm(const EquivPtr &equiv, - const AnfNodePtr &bn) const { - MS_EXCEPTION_IF_NULL(equiv); - MS_EXCEPTION_IF_NULL(bn); - auto variable_input0 = utils::cast((*equiv)[variable_input_var0_]); - MS_EXCEPTION_IF_NULL(variable_input0); - auto variable_input1 = utils::cast((*equiv)[variable_input_var1_]); - MS_EXCEPTION_IF_NULL(variable_input1); - auto bn_abstract_tuple = dyn_cast(bn->abstract()); - MS_EXCEPTION_IF_NULL(bn_abstract_tuple); - if (bn_abstract_tuple->elements().size() != kBnOutputNum) { - MS_LOG(EXCEPTION) << "The abstract size of node bn must be " << kBnOutputNum << ", but it is " - << bn_abstract_tuple->elements().size(); - } - AbstractBasePtrList fused_bn_abstract_list{bn_abstract_tuple->elements()[0], variable_input0->abstract(), - variable_input1->abstract(), bn_abstract_tuple->elements()[3], - bn_abstract_tuple->elements()[4]}; - auto abstract_tuple = std::make_shared(fused_bn_abstract_list); - return abstract_tuple; -} - ValuePtr FusedBatchNormFusion::GetFactor(const EquivPtr &equiv) const { MS_EXCEPTION_IF_NULL(equiv); - auto constant_input = utils::cast((*equiv)[constant_input_var0_]); + auto iter_constant_input0 = (*equiv).find(constant_input0_var_); + if (iter_constant_input0 == (*equiv).end()) { + MS_LOG(EXCEPTION) << "The equiv map is expected to contains the constant_input0 var after matched."; + } + auto constant_input = utils::cast(iter_constant_input0->second); MS_EXCEPTION_IF_NULL(constant_input); if (!constant_input->isa()) { return nullptr; @@ -158,53 +95,187 @@ ValuePtr FusedBatchNormFusion::GetFactor(const EquivPtr &equiv) const { return MakeValue(tensor_data[0]); } -const AnfNodePtr FusedBatchNormFusion::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node, - const EquivPtr &equiv) const { +AnfNodePtr FusedBatchNormFusion::CreateBNTrainingReduce(const FuncGraphPtr &func_graph, const AnfNodePtr &node, + const EquivPtr &equiv) const { MS_EXCEPTION_IF_NULL(func_graph); + MS_EXCEPTION_IF_NULL(node); MS_EXCEPTION_IF_NULL(equiv); - // Set inputs - auto data_input0 = utils::cast((*equiv)[data_input_var0_]); - MS_EXCEPTION_IF_NULL(data_input0); - auto data_input1 = utils::cast((*equiv)[data_input_var1_]); + // Set input to create node + auto iter_data_input0 = (*equiv).find(data_input0_var_); + if (iter_data_input0 == (*equiv).end()) { + MS_LOG(EXCEPTION) << "The equiv map is expected to contains the data_input0 var after matched."; + } + std::vector bn_training_reduce_inputs = { + NewValueNode(std::make_shared(kBNTrainingReduceOpName)), + utils::cast(iter_data_input0->second)}; + auto bn_training_reduce = func_graph->NewCNode(bn_training_reduce_inputs); + MS_EXCEPTION_IF_NULL(bn_training_reduce); + bn_training_reduce->set_scope(node->scope()); + // Set abstract + auto iter_data_input1 = (*equiv).find(data_input1_var_); + if (iter_data_input1 == (*equiv).end()) { + MS_LOG(EXCEPTION) << "The equiv map is expected to contains the data_input1 var after matched."; + } + auto data_input1 = utils::cast(iter_data_input1->second); MS_EXCEPTION_IF_NULL(data_input1); - auto data_input2 = utils::cast((*equiv)[data_input_var2_]); + auto iter_data_input2 = (*equiv).find(data_input2_var_); + if (iter_data_input2 == (*equiv).end()) { + MS_LOG(EXCEPTION) << "The equiv map is expected to contains the data_input2 var after matched."; + } + auto data_input2 = utils::cast(iter_data_input2->second); MS_EXCEPTION_IF_NULL(data_input2); - auto variable_input0 = utils::cast((*equiv)[variable_input_var0_]); + AbstractBasePtrList abstract_list{data_input1->abstract(), data_input2->abstract()}; + auto abstract_tuple = std::make_shared(abstract_list); + bn_training_reduce->set_abstract(abstract_tuple); + return bn_training_reduce; +} + +void FusedBatchNormFusion::GetBNTrainingUpdateInputs(const EquivPtr &equiv, + const std::vector &bn_training_reduce_outputs, + std::vector *bn_training_update_inputs) const { + MS_EXCEPTION_IF_NULL(equiv); + MS_EXCEPTION_IF_NULL(bn_training_update_inputs); + auto iter_data_input0 = (*equiv).find(data_input0_var_); + if (iter_data_input0 == (*equiv).end()) { + MS_LOG(EXCEPTION) << "The equiv map is expected to contains the data_input0 var after matched."; + } + auto iter_data_input1 = (*equiv).find(data_input1_var_); + if (iter_data_input1 == (*equiv).end()) { + MS_LOG(EXCEPTION) << "The equiv map is expected to contains the data_input1 var after matched."; + } + auto iter_data_input2 = (*equiv).find(data_input2_var_); + if (iter_data_input2 == (*equiv).end()) { + MS_LOG(EXCEPTION) << "The equiv map is expected to contains the data_input2 var after matched."; + } + auto iter_variable_input0 = (*equiv).find(variable_input0_var_); + if (iter_variable_input0 == (*equiv).end()) { + MS_LOG(EXCEPTION) << "The equiv map is expected to contains the variable_input0 var after matched."; + } + auto iter_variable_input1 = (*equiv).find(variable_input1_var_); + if (iter_variable_input1 == (*equiv).end()) { + MS_LOG(EXCEPTION) << "The equiv map is expected to contains the variable_input1 var after matched."; + } + 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(); + } + *bn_training_update_inputs = { + NewValueNode(std::make_shared(kBNTrainingUpdateOpName)), + utils::cast(iter_data_input0->second), + bn_training_reduce_outputs[0], + bn_training_reduce_outputs[1], + utils::cast(iter_data_input1->second), + utils::cast(iter_data_input2->second), + utils::cast(iter_variable_input0->second), + utils::cast(iter_variable_input1->second), + }; +} + +void FusedBatchNormFusion::GetBNTrainingUpdateAbstractList(const EquivPtr &equiv, const AnfNodePtr &bn, + std::vector *abstract_list) const { + MS_EXCEPTION_IF_NULL(equiv); + MS_EXCEPTION_IF_NULL(bn); + MS_EXCEPTION_IF_NULL(abstract_list); + auto bn_abstract_tuple = dyn_cast(bn->abstract()); + MS_EXCEPTION_IF_NULL(bn_abstract_tuple); + if (bn_abstract_tuple->elements().size() < kBnOutputNum) { + MS_LOG(EXCEPTION) << "The abstract size of node bn must not be less than " << kBnOutputNum << ", but it is " + << bn_abstract_tuple->elements().size(); + } + auto iter_variable_input0 = (*equiv).find(variable_input0_var_); + if (iter_variable_input0 == (*equiv).end()) { + MS_LOG(EXCEPTION) << "The equiv map is expected to contains the variable_input0 var after matched."; + } + auto variable_input0 = utils::cast(iter_variable_input0->second); MS_EXCEPTION_IF_NULL(variable_input0); - auto variable_input1 = utils::cast((*equiv)[variable_input_var1_]); + auto iter_variable_input1 = (*equiv).find(variable_input1_var_); + if (iter_variable_input1 == (*equiv).end()) { + MS_LOG(EXCEPTION) << "The equiv map is expected to contains the variable_input1 var after matched."; + } + auto variable_input1 = utils::cast(iter_variable_input1->second); MS_EXCEPTION_IF_NULL(variable_input1); - std::vector fused_bn_inputs = { - NewValueNode(prim::kPrimFusedBatchNorm), data_input0, data_input1, data_input2, variable_input0, variable_input1}; - auto fused_bn = func_graph->NewCNode(fused_bn_inputs); - fused_bn->set_scope(node->scope()); - MS_EXCEPTION_IF_NULL(fused_bn); + *abstract_list = {bn_abstract_tuple->elements()[0], variable_input0->abstract(), variable_input1->abstract(), + bn_abstract_tuple->elements()[1], bn_abstract_tuple->elements()[2]}; +} + +AnfNodePtr FusedBatchNormFusion::CreateBNTrainingUpdate( + const FuncGraphPtr &func_graph, const AnfNodePtr &node, const EquivPtr &equiv, + const std::vector &bn_training_reduce_outputs) const { + MS_EXCEPTION_IF_NULL(func_graph); + MS_EXCEPTION_IF_NULL(node); + MS_EXCEPTION_IF_NULL(equiv); + // Set input + std::vector bn_training_update_inputs; + GetBNTrainingUpdateInputs(equiv, bn_training_reduce_outputs, &bn_training_update_inputs); + auto bn_training_update = func_graph->NewCNode(bn_training_update_inputs); + MS_EXCEPTION_IF_NULL(bn_training_update); // Set abstract - AnfNodePtr bn = GetBatchNormNode(node); - fused_bn->set_abstract(CreateAbstractOfFusedBatchNorm(equiv, bn)); - // Set attr - AnfAlgo::CopyNodeAttr(kAttrEpsilon, bn, fused_bn); + auto iter_batch_norm = (*equiv).find(batch_norm_var_); + if (iter_batch_norm == (*equiv).end()) { + MS_LOG(EXCEPTION) << "The equiv map is expected to contains the batch_norm var after matched."; + } + AnfNodePtr bn = utils::cast(iter_batch_norm->second); + MS_EXCEPTION_IF_NULL(bn); + AbstractBasePtrList abstract_list; + GetBNTrainingUpdateAbstractList(equiv, bn, &abstract_list); + auto abstract_tuple = std::make_shared(abstract_list); + bn_training_update->set_abstract(abstract_tuple); + AnfAlgo::CopyNodeAttr(kAttrEpsilon, bn, bn_training_update); ValuePtr factor = GetFactor(equiv); if (factor == nullptr) { return nullptr; } - AnfAlgo::SetNodeAttr(kAttrMomentum, factor, fused_bn); - // Replace old nodes with outputs of fused_bn - std::vector fused_bn_outputs; - CreateMultipleOutputsOfAnfNode(func_graph, fused_bn, kBnOutputNum, &fused_bn_outputs); - if (fused_bn_outputs.size() != kBnOutputNum) { - MS_LOG(EXCEPTION) << "The output size of node bn must be " << kBnOutputNum << ", but it is " - << fused_bn_outputs.size(); + AnfAlgo::SetNodeAttr(kAttrFactor, factor, bn_training_update); + AnfAlgo::SetNodeAttr(kAttrIsRef, MakeValue(true), bn_training_update); + bn_training_update->set_scope(node->scope()); + return bn_training_update; +} + +const AnfNodePtr FusedBatchNormFusion::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node, + const EquivPtr &equiv) const { + MS_EXCEPTION_IF_NULL(func_graph); + MS_EXCEPTION_IF_NULL(equiv); + MS_EXCEPTION_IF_NULL(node); + AnfNodePtr bn_training_reduce = CreateBNTrainingReduce(func_graph, node, equiv); + std::vector bn_training_reduce_outputs; + CreateMultipleOutputsOfAnfNode(func_graph, bn_training_reduce, kBNTrainingReduceOutputNum, + &bn_training_reduce_outputs); + AnfNodePtr bn_training_update = CreateBNTrainingUpdate(func_graph, node, equiv, bn_training_reduce_outputs); + if (bn_training_update == nullptr) { + MS_LOG(DEBUG) << "Create BNTrainingUpdate failed for bn node " << node->DebugString(); + return nullptr; + } + std::vector bn_training_update_outputs; + CreateMultipleOutputsOfAnfNode(func_graph, bn_training_update, kBNTrainingUpdateOutputNum, + &bn_training_update_outputs); + if (bn_training_update_outputs.size() < kBNTrainingUpdateOutputNum) { + MS_LOG(EXCEPTION) << "The output size of node bn must be " << kBNTrainingUpdateOutputNum << ", but it is " + << bn_training_update_outputs.size(); + } + // Replace old bn outputs with new outputs + auto iter_batch_norm = (*equiv).find(batch_norm_var_); + if (iter_batch_norm == (*equiv).end()) { + MS_LOG(EXCEPTION) << "The equiv map is expected to contains the batch_norm var after matched."; } + AnfNodePtr bn = utils::cast(iter_batch_norm->second); std::vector bn_outputs; GetBNOutput(func_graph, bn, &bn_outputs); - if (bn_outputs.size() != kBnOutputNum) { - MS_LOG(EXCEPTION) << "The output size of node bn must be " << kBnOutputNum << ", but it is " << bn_outputs.size(); - } auto manager = func_graph->manager(); MS_EXCEPTION_IF_NULL(manager); - (void)manager->Replace(bn_outputs[3], fused_bn_outputs[3]); - (void)manager->Replace(bn_outputs[4], fused_bn_outputs[4]); - return fused_bn_outputs[0]; + for (const auto &output : bn_outputs) { + MS_EXCEPTION_IF_NULL(output); + auto tuple_getitem_cnode = output->cast(); + MS_EXCEPTION_IF_NULL(tuple_getitem_cnode); + AnfNodePtr index_node = tuple_getitem_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 (index == kReplaceOutputIndex0 || index == kReplaceOutputIndex1) { + (void)manager->Replace(output, bn_training_update_outputs[index]); + } + } + return bn_training_update_outputs[0]; } } // namespace opt } // namespace mindspore diff --git a/mindspore/ccsrc/pre_activate/ascend/ir_fusion/fused_batch_norm_fusion.h b/mindspore/ccsrc/pre_activate/ascend/ir_fusion/fused_batch_norm_fusion.h index db25e4f9f537dfbab26a25a9d65f2f91e15ba9bc..e6bf1dda554d7dd29065f8fba60841d4f2b3f2b8 100644 --- a/mindspore/ccsrc/pre_activate/ascend/ir_fusion/fused_batch_norm_fusion.h +++ b/mindspore/ccsrc/pre_activate/ascend/ir_fusion/fused_batch_norm_fusion.h @@ -19,6 +19,7 @@ #include #include #include "pre_activate/common/optimizer.h" +#include "utils/utils.h" namespace mindspore { namespace opt { @@ -26,29 +27,37 @@ class FusedBatchNormFusion : public PatternProcessPass { public: explicit FusedBatchNormFusion(bool multigraph = true) : PatternProcessPass("fused_batch_norm_fusion", multigraph), - data_input_var0_(std::make_shared()), - data_input_var1_(std::make_shared()), - data_input_var2_(std::make_shared()), - variable_input_var0_(std::make_shared()), - variable_input_var1_(std::make_shared()), - constant_input_var0_(std::make_shared()), - constant_input_var1_(std::make_shared()) {} + data_input0_var_(std::make_shared()), + data_input1_var_(std::make_shared()), + data_input2_var_(std::make_shared()), + variable_input0_var_(std::make_shared()), + variable_input1_var_(std::make_shared()), + constant_input0_var_(std::make_shared()), + constant_input1_var_(std::make_shared()), + batch_norm_var_(std::make_shared(std::make_shared(prim::kPrimBatchNorm->name()))) {} ~FusedBatchNormFusion() override = default; const BaseRef DefinePattern() const override; const AnfNodePtr Process(const FuncGraphPtr &, const AnfNodePtr &, const EquivPtr &) const override; private: - abstract::AbstractTuplePtr CreateAbstractOfFusedBatchNorm(const EquivPtr &equiv, const AnfNodePtr &bn) const; - + AnfNodePtr CreateBNTrainingReduce(const FuncGraphPtr &func_graph, const AnfNodePtr &node, + const EquivPtr &equiv) const; + void GetBNTrainingUpdateInputs(const EquivPtr &equiv, const std::vector &bn_training_reduce_outputs, + std::vector *bn_training_update_inputs) const; + void GetBNTrainingUpdateAbstractList(const EquivPtr &equiv, const AnfNodePtr &bn, + std::vector *abstract_list) const; + AnfNodePtr CreateBNTrainingUpdate(const FuncGraphPtr &func_graph, const AnfNodePtr &node, const EquivPtr &equiv, + const std::vector &bn_training_reduce_outputs) const; ValuePtr GetFactor(const EquivPtr &equiv) const; - VarPtr data_input_var0_; - VarPtr data_input_var1_; - VarPtr data_input_var2_; - VarPtr variable_input_var0_; - VarPtr variable_input_var1_; - VarPtr constant_input_var0_; - VarPtr constant_input_var1_; + VarPtr data_input0_var_; + VarPtr data_input1_var_; + VarPtr data_input2_var_; + VarPtr variable_input0_var_; + VarPtr variable_input1_var_; + VarPtr constant_input0_var_; + VarPtr constant_input1_var_; + VarPtr batch_norm_var_; }; } // namespace opt } // namespace mindspore diff --git a/mindspore/nn/layer/normalization.py b/mindspore/nn/layer/normalization.py index 7a102b0bbe9446e6151b0587a3764e66fd783eb4..fd9279cf04c2321b7f05835ddaca6a3c30cc2e33 100644 --- a/mindspore/nn/layer/normalization.py +++ b/mindspore/nn/layer/normalization.py @@ -62,6 +62,7 @@ class _BatchNorm(Cell): self.beta = Parameter(initializer( beta_init, num_features), name="beta", requires_grad=affine) self.group = check_int_positive(device_num_each_group) + self.is_global = False if self.group != 1: self.rank_id = get_rank() self.rank_size = get_group_size() @@ -80,15 +81,18 @@ class _BatchNorm(Cell): self.cast = P.Cast() self.dtype = P.DType() self.reshape = P.Reshape() + self.is_ascend = context.get_context("device_target") == "Ascend" if context.get_context("enable_ge"): self.is_ge_backend = True self.momentum = Tensor(1.0 - momentum, mstype.float32) - self.bn_train = P.BatchNorm(is_training=True, - epsilon=self.eps) else: self.is_ge_backend = False self.momentum = 1.0 - momentum + if self.is_ge_backend or self.is_ascend: + self.bn_train = P.BatchNorm(is_training=True, + epsilon=self.eps) + else: self.bn_train = P.FusedBatchNorm(mode=1, epsilon=self.eps, momentum=self.momentum) @@ -140,24 +144,23 @@ class _BatchNorm(Cell): def construct(self, x): if self.training and self.use_batch_statistics: - if self.is_ge_backend: - if self.is_global: - axes, re_shape = _shape_infer(F.shape(x), self.num_features) - y = self._global_sync(x, axes, re_shape) - else: - y, batch_mean, batch_var, _, _ = \ - self.bn_train(x, - self.gamma, - self.beta, - None, - None) - - mean_sub = self.sub_mean(self.moving_mean, batch_mean) - temp_mean = self.mul_mean(mean_sub, self.momentum) - mean_sub2 = self.sub_var(self.moving_variance, batch_var) - temp_variance = self.mul_var(mean_sub2, self.momentum) - y = F.depend(y, self.assign_sub_mean(self.moving_mean, temp_mean)) - y = F.depend(y, self.assign_sub_var(self.moving_variance, temp_variance)) + if self.is_ge_backend and self.is_global: + axes, re_shape = _shape_infer(F.shape(x), self.num_features) + y = self._global_sync(x, axes, re_shape) + elif self.is_ge_backend or self.is_ascend: + y, batch_mean, batch_var, _, _ = \ + self.bn_train(x, + self.gamma, + self.beta, + None, + None) + + mean_sub = self.sub_mean(self.moving_mean, batch_mean) + temp_mean = self.mul_mean(mean_sub, self.momentum) + mean_sub2 = self.sub_var(self.moving_variance, batch_var) + temp_variance = self.mul_var(mean_sub2, self.momentum) + y = F.depend(y, self.assign_sub_mean(self.moving_mean, temp_mean)) + y = F.depend(y, self.assign_sub_var(self.moving_variance, temp_variance)) else: y = self.bn_train(x, self.gamma, diff --git a/tests/ut/cpp/pre_activate/ascend/ir_fusion/fused_batch_norm_fusion_test.cc b/tests/ut/cpp/pre_activate/ascend/ir_fusion/fused_batch_norm_fusion_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..3d13f4a336e68ff64b41345f0a5d38243d776ee6 --- /dev/null +++ b/tests/ut/cpp/pre_activate/ascend/ir_fusion/fused_batch_norm_fusion_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_fusion/fused_batch_norm_fusion.h" +#include "common/backend_common_test.h" +#include "common/py_func_graph_fetcher.h" + +namespace mindspore { +namespace opt { +class TestHWFusedBatchNormFusion : public BackendCommon { + public: + TestHWFusedBatchNormFusion() : get_py_fun_("gtest_input.pre_activate.fused_batch_norm_fusion_test", true) {} + ~TestHWFusedBatchNormFusion() override = default; + + UT::PyFuncGraphFetcher get_py_fun_; +}; + +TEST_F(TestHWFusedBatchNormFusion, test_fused_batch_norm_fusion) { + FuncGraphPtr g = get_py_fun_.CallAndParseRet("test_fused_batch_norm_fusion", "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 < 6; ++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_fused_batch_norm_fusion", "after"); + EXPECT_TRUE(CheckEqualGraph(g_after, new_graph)); +} +} // namespace opt +} // namespace mindspore \ No newline at end of file diff --git a/tests/ut/cpp/python_input/gtest_input/pre_activate/fused_batch_norm_fusion_test.py b/tests/ut/cpp/python_input/gtest_input/pre_activate/fused_batch_norm_fusion_test.py index 8f4b8b476f2199d7a324ba084cb220de31d925e9..ca93d40443644b9c7fe878cc5d3aa188280d3034 100644 --- a/tests/ut/cpp/python_input/gtest_input/pre_activate/fused_batch_norm_fusion_test.py +++ b/tests/ut/cpp/python_input/gtest_input/pre_activate/fused_batch_norm_fusion_test.py @@ -24,7 +24,8 @@ make_tuple = Primitive('make_tuple') tuple_getitem = Primitive('tuple_getitem') depend = Primitive('depend') BatchNorm = P.BatchNorm() -FusedBatchNorm = P.FusedBatchNorm() +BNTrainingReduce = Primitive('BNTrainingReduce') +BNTrainingUpdate = Primitive('BNTrainingUpdate') constant0 = Tensor(0.1, mstype.float32) constant1 = Tensor(0.1, mstype.float32) @@ -40,7 +41,7 @@ class FnDict: return self.fnDict[name] -def useless_test_fused_batch_norm_fusion(tag): +def test_fused_batch_norm_fusion(tag): fns = FnDict() @fns @@ -60,9 +61,11 @@ def useless_test_fused_batch_norm_fusion(tag): @fns def after(input0, input1, input2, input3, input4, var0, var1): - fused_batch_norm = FusedBatchNorm(input0, input1, input2, var0, var1) - outputs = make_tuple(tuple_getitem(fused_batch_norm, 0), tuple_getitem(fused_batch_norm, 3), - tuple_getitem(fused_batch_norm, 4)) + bn_training_reduce = BNTrainingReduce(input0) + bn_training_update = BNTrainingUpdate(input0, tuple_getitem(bn_training_reduce, 0), + tuple_getitem(bn_training_reduce, 1), input1, input2, var0, var1) + outputs = make_tuple(tuple_getitem(bn_training_update, 0), tuple_getitem(bn_training_update, 3), + tuple_getitem(bn_training_update, 4)) output = tuple_getitem(outputs, 0) return make_tuple(output)