Add batch norm bert fission pass

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<BatchNormBertFission>());
   ir_fusion_pm->AddPass(std::make_shared<SquareSumFusion>());
   ir_fusion_pm->AddPass(std::make_shared<ClipByNormNoDivSquareSumFusion>());
   ir_fusion_pm->AddPass(std::make_shared<LambUpdateWithLRRuleFusion>());

mindspore/ccsrc/pre_activate/ascend/ir_fission/batch_norm_bert_fission.cc

+/**
+ * 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 <vector>
+#include <memory>
+#include <algorithm>
+#include "session/anf_runtime_algorithm.h"
+#include "pre_activate/common/helper.h"
+
+namespace mindspore {
+namespace opt {
+namespace {
+const std::vector<int> 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<CNodePtr>();
+  auto n2_cnode = n2->cast<CNodePtr>();
+  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<ValueNodePtr>();
+  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<ValueNodePtr>();
+  MS_EXCEPTION_IF_NULL(value_node2);
+  return GetValue<int>(value_node1->value()) < GetValue<int>(value_node2->value());
+}
+
+bool GetBatchNormOutputs(const FuncGraphPtr &func_graph, const AnfNodePtr &bn, std::vector<AnfNodePtr> *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<CNodePtr>();
+    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<ValueNodePtr>();
+    MS_EXCEPTION_IF_NULL(value_node);
+    int index = GetValue<int>(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<CNodePtr>();
+  MS_EXCEPTION_IF_NULL(bn_cnode);
+  CheckCNodeInputSize(bn_cnode, kBatchNormInputNum + 1);
+  std::vector<AnfNodePtr> bn_training_reduce_inputs = {
+    NewValueNode(std::make_shared<Primitive>(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::AbstractTuple>(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<AnfNodePtr> &bn_training_reduce_outputs) {
+  MS_EXCEPTION_IF_NULL(func_graph);
+  MS_EXCEPTION_IF_NULL(bn);
+  auto bn_cnode = bn->cast<CNodePtr>();
+  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<AnfNodePtr> bn_training_update_v2_inputs = {
+    NewValueNode(std::make_shared<Primitive>(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<abstract::AbstractTuple>(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<AbstractBasePtr> abstract_list{bn_abstract_tuple->elements()[0], bn_abstract_tuple->elements()[3],
+                                             bn_abstract_tuple->elements()[4]};
+  auto abstract_tuple = std::make_shared<abstract::AbstractTuple>(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<SeqVar>();
+  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<AnfNodePtr> bn_outputs;
+  if (!GetBatchNormOutputs(func_graph, node, &bn_outputs)) {
+    return nullptr;
+  }
+  AnfNodePtr bn_training_reduce = CreateBNTrainingReduce(func_graph, node);
+  std::vector<AnfNodePtr> 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<AnfNodePtr> 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

mindspore/ccsrc/pre_activate/ascend/ir_fission/batch_norm_bert_fission.h

+/**
+ * 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_

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;

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";

tests/ut/cpp/pre_activate/ascend/ir_fission/batch_norm_bert_fission_test.cc

+/**
+ * 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<int> shp_x{32, 64, 112, 112};
+  auto x_abstract = std::make_shared<abstract::AbstractTensor>(kFloat32, shp_x);
+  std::vector<int> shp_y{64};
+  auto y_abstract = std::make_shared<abstract::AbstractTensor>(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<opt::GraphOptimizer>();
+  auto pm = std::make_shared<opt::PassManager>();
+  pm->AddPass(std::make_shared<opt::BatchNormBertFission>());
+  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

tests/ut/cpp/python_input/gtest_input/pre_activate/batch_norm_bert_fission_test.py

+# 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]