!1098 Add BatchNorm2BNInfer pass

Merge pull request !1098 from huanghui/add-batchnorm-bninfer-pass

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<DereluFusion>());
   ir_fusion_pm->AddPass(std::make_shared<TransposeTransDataFusion>());
   ir_fusion_pm->AddPass(std::make_shared<GetitemTuple>());
+  ir_fusion_pm->AddPass(std::make_shared<BatchNorm2BNInfer>());
 }
 }  // namespace
 

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 <memory>
+#include <vector>
+#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<Primitive>(kBNInferOpName);
+  std::vector<AnfNodePtr> 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<Int32Imm>(index_node)) {
+    return false;
+  }
+  ValueNodePtr value_node = index_node->cast<ValueNodePtr>();
+  MS_EXCEPTION_IF_NULL(value_node);
+  int index = GetValue<int>(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<bool>(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<CNodePtr>();
+  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<CNodePtr>();
+  MS_EXCEPTION_IF_NULL(*batchnorm);
+  return CheckBatchNorm(graph, *batchnorm);
+}
+}  // namespace
+
+const BaseRef BatchNorm2BNInfer::DefinePattern() const {
+  VarPtr Xs = std::make_shared<SeqVar>();
+  VarPtr Y = std::make_shared<Var>();
+  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

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 <memory>
+#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_

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

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<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, y_abstract, y_abstract, y_abstract, y_abstract};
+  auto fg = 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::BatchNorm2BNInfer>());
+  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<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, y_abstract, y_abstract, y_abstract, y_abstract};
+  auto fg = GetKernelGraph(g, args_spec_list);
+  auto origin_graph = std::make_shared<session::KernelGraph>(*fg);
+
+  auto optimizer = std::make_shared<opt::GraphOptimizer>();
+  auto pm = std::make_shared<opt::PassManager>();
+  pm->AddPass(std::make_shared<opt::BatchNorm2BNInfer>());
+  optimizer->AddPassManager(pm);
+  FuncGraphPtr new_graph = optimizer->Optimize(fg);
+
+  EXPECT_TRUE(CheckEqualGraph(origin_graph, new_graph));
+}
+}  // namespace opt
+}  // namespace mindspore

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]