Add BatchNormGrad split pass

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

@@ -0,0 +1,132 @@
+/**
+ * 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_grad_split.h"
+
+#include <vector>
+#include <string>
+#include <memory>
+
+#include "utils/utils.h"
+#include "utils/context/ms_context.h"
+#include "common/utils.h"
+#include "pre_activate/common/helper.h"
+#include "device/kernel_info.h"
+#include "session/anf_runtime_algorithm.h"
+
+namespace mindspore {
+namespace opt {
+namespace {
+void CreateOutputsOfUpdateGrad(const FuncGraphPtr &graph, const CNodePtr &bn_grad_node,
+                               std::vector<AnfNodePtr> *bn_update_grad_outputs) {
+  MS_EXCEPTION_IF_NULL(graph);
+  MS_EXCEPTION_IF_NULL(bn_grad_node);
+  auto bn_grad_inputs = bn_grad_node->inputs();
+  if (bn_grad_inputs.size() < kBNGradInputNum) {
+    MS_LOG(EXCEPTION) << "BNGrad has wrong inputs size";
+  }
+  std::vector<AnfNodePtr> bn_update_grad_inputs = {
+    NewValueNode(std::make_shared<Primitive>(kBNTrainingUpdateGradOpName)), bn_grad_inputs[1], bn_grad_inputs[2],
+    bn_grad_inputs[4], bn_grad_inputs[5]};
+  auto bn_update_grad = graph->NewCNode(bn_update_grad_inputs);
+  MS_EXCEPTION_IF_NULL(bn_update_grad);
+  bn_update_grad->set_kernel_info(std::make_shared<device::KernelInfo>());
+  bn_update_grad->set_scope(bn_grad_node->scope());
+
+  auto types = {AnfAlgo::GetOutputInferDataType(bn_grad_node, 1), AnfAlgo::GetOutputInferDataType(bn_grad_node, 2)};
+  auto shapes = {AnfAlgo::GetOutputInferShape(bn_grad_node, 1), AnfAlgo::GetOutputInferShape(bn_grad_node, 2)};
+  AnfAlgo::SetOutputInferTypeAndShape(types, shapes, bn_update_grad.get());
+
+  AnfAlgo::CopyNodeAttr(kAttrEpsilon, bn_grad_node, bn_update_grad);
+  CreateMultipleOutputsOfAnfNode(graph, bn_update_grad, kBNTrainingUpdateGradOutputNum, bn_update_grad_outputs);
+}
+
+void CreateOutputsOfReduceGrad(const FuncGraphPtr &graph, const CNodePtr &bn_grad_node,
+                               const std::vector<AnfNodePtr> &bn_update_grad_outputs,
+                               std::vector<AnfNodePtr> *bn_reduce_grad_outputs) {
+  MS_EXCEPTION_IF_NULL(graph);
+  MS_EXCEPTION_IF_NULL(bn_grad_node);
+  auto bn_grad_inputs = bn_grad_node->inputs();
+  if (bn_grad_inputs.size() < kBNGradInputNum) {
+    MS_LOG(EXCEPTION) << "BNGrad has wrong inputs size";
+  }
+  if (bn_update_grad_outputs.size() != kBNTrainingUpdateGradOutputNum) {
+    MS_LOG(EXCEPTION) << "BNTrainingReduceGrad_outputs has wrong size";
+  }
+  std::vector<AnfNodePtr> bn_reduce_grad_inputs = {
+    NewValueNode(std::make_shared<Primitive>(kBNTrainingReduceGradOpName)),
+    bn_grad_inputs[1],
+    bn_grad_inputs[2],
+    bn_update_grad_outputs[0],
+    bn_update_grad_outputs[1],
+    bn_grad_inputs[3],
+    bn_grad_inputs[4],
+    bn_grad_inputs[5]};
+  auto bn_reduce_grad = graph->NewCNode(bn_reduce_grad_inputs);
+  MS_EXCEPTION_IF_NULL(bn_reduce_grad);
+  bn_reduce_grad->set_kernel_info(std::make_shared<device::KernelInfo>());
+  bn_reduce_grad->set_scope(bn_grad_node->scope());
+
+  auto types = {AnfAlgo::GetOutputInferDataType(bn_grad_node, 0)};
+  auto shapes = {AnfAlgo::GetOutputInferShape(bn_grad_node, 0)};
+  AnfAlgo::SetOutputInferTypeAndShape(types, shapes, bn_reduce_grad.get());
+
+  AnfAlgo::CopyNodeAttr(kAttrEpsilon, bn_grad_node, bn_reduce_grad);
+  (*bn_reduce_grad_outputs).push_back(bn_reduce_grad);
+}
+
+}  // namespace
+const BaseRef BatchNormGradSplit::DefinePattern() const {
+  VarPtr Xs = std::make_shared<SeqVar>();
+  auto prim = std::make_shared<Primitive>(kBatchNormGradOpName);
+  return VectorRef({prim, Xs});
+}
+
+const AnfNodePtr BatchNormGradSplit::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
+                                             const EquivPtr &) const {
+  MS_EXCEPTION_IF_NULL(node);
+  MS_EXCEPTION_IF_NULL(func_graph);
+  auto cnode = node->cast<CNodePtr>();
+  MS_EXCEPTION_IF_NULL(cnode);
+  auto primitive = AnfAlgo::GetCNodePrimitive(cnode);
+  MS_EXCEPTION_IF_NULL(primitive);
+  if (!primitive->HasAttr(kAttrIsTraining)) {
+    MS_LOG(INFO) << "Op BatchNormGrad must have attrs of is_training";
+    return nullptr;
+  }
+  if (!AnfAlgo::GetNodeAttr<bool>(cnode, kAttrIsTraining)) {
+    MS_LOG(INFO) << "is_training must be true";
+    return nullptr;
+  }
+
+  std::vector<AnfNodePtr> bn_update_grad_outputs;
+  CreateOutputsOfUpdateGrad(func_graph, cnode, &bn_update_grad_outputs);
+  if (bn_update_grad_outputs.size() != kBNTrainingUpdateGradOutputNum) {
+    MS_LOG(EXCEPTION) << "bn_update_grad_outputs has wrong size";
+  }
+
+  std::vector<AnfNodePtr> bn_reduce_grad_outputs;
+  CreateOutputsOfReduceGrad(func_graph, cnode, bn_update_grad_outputs, &bn_reduce_grad_outputs);
+  if (bn_reduce_grad_outputs.size() != kSingleOutputNum) {
+    MS_LOG(EXCEPTION) << "bn_reduce_grad_outputs has wrong size";
+  }
+
+  std::vector<AnfNodePtr> make_tuple_inputs = {NewValueNode(prim::kPrimMakeTuple), bn_reduce_grad_outputs[0],
+                                               bn_update_grad_outputs[0], bn_update_grad_outputs[1]};
+  auto make_tuple = func_graph->NewCNode(make_tuple_inputs);
+  return make_tuple;
+}
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/pre_activate/ascend/ir_fission/batch_norm_grad_split.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_FISSION_BATCH_NORM_GRAD_SPLIT_H_
+#define MINDSPORE_CCSRC_PRE_ACTIVATE_ASCEND_IR_FISSION_BATCH_NORM_GRAD_SPLIT_H_
+
+#include "pre_activate/common/optimizer.h"
+#include "pre_activate/common/helper.h"
+
+namespace mindspore {
+namespace opt {
+class BatchNormGradSplit : public PatternProcessPass {
+ public:
+  explicit BatchNormGradSplit(bool multigraph = true) : PatternProcessPass("batch_norm_grad_split", multigraph) {}
+  ~BatchNormGradSplit() 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_GRAD_SPLIT_H_

mindspore/ccsrc/utils/utils.h

@@ -107,6 +107,7 @@ constexpr auto kLambNextMVOpName = "LambNextMV";
 constexpr auto kConfusionTransposeDOpName = "ConfusionTransposeD";
 constexpr auto kAdamApplyOneWithDecayOpName = "AdamApplyOneWithDecay";
 constexpr auto kBatchNormOpName = "BatchNorm";
+constexpr auto kBatchNormGradOpName = "BatchNormGrad";
 constexpr auto kAdamApplyOneOpName = "AdamApplyOne";
 constexpr auto kDropoutGenMask = "DropoutGenMask";
 constexpr auto kResizeNearestNeighborGrad = "ResizeNearestNeighborGrad";
@@ -162,6 +163,7 @@ constexpr auto kAttrLabelForInsertStreamActive = "label_for_insert_stream_active
 constexpr auto kAttrFusion = "fusion";
 constexpr auto kAttrGroup = "group";
 constexpr auto kAttrOp = "op";
+constexpr auto kAttrIsTraining = "is_training";
 
 // attr value
 constexpr auto kValueTargetSwitch = "target_switch";

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

@@ -0,0 +1,59 @@
+/**
+ * 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 "operator/ops.h"
+#include "ir/meta_tensor.h"
+#include "debug/anf_ir_dump.h"
+#include "utils/utils.h"
+#include "pre_activate/common/optimizer.h"
+#include "pre_activate/ascend/ir_fission/batch_norm_grad_split.h"
+#include "session/anf_runtime_algorithm.h"
+
+namespace mindspore {
+namespace opt {
+class TestHWBatchNormGradSplit : public BackendCommon {
+ public:
+  TestHWBatchNormGradSplit() : get_py_fun_("gtest_input.pre_activate.batch_norm_grad_split", true) {}
+
+ public:
+  UT::PyFuncGraphFetcher get_py_fun_;
+};
+
+TEST_F(TestHWBatchNormGradSplit, test_split) {
+  get_py_fun_.SetDoResolve(true);
+  FuncGraphPtr g = get_py_fun_.CallAndParseRet("test_batch_norm_grad_split", "before");
+  EXPECT_NE(g, nullptr);
+  std::vector<int> shp_x{1, 64, 112, 112};
+  std::vector<int> shp_b{64};
+  auto x_abstract = std::make_shared<abstract::AbstractTensor>(kFloat32, shp_x);
+  auto b_abstract = std::make_shared<abstract::AbstractTensor>(kFloat32, shp_b);
+  AbstractBasePtrList args_spec_list{x_abstract, x_abstract, b_abstract, b_abstract, b_abstract, b_abstract};
+  auto kernel_graph = GetKernelGraph(g, args_spec_list);
+  EXPECT_NE(kernel_graph, nullptr);
+
+  auto optimizer = std::make_shared<opt::GraphOptimizer>();
+  auto pm = std::make_shared<opt::PassManager>();
+  auto pass = std::make_shared<opt::BatchNormGradSplit>();
+  pm->AddPass(pass);
+  optimizer->AddPassManager(pm);
+  auto new_graph = optimizer->Optimize(kernel_graph);
+
+  FuncGraphPtr g_after = get_py_fun_.CallAndParseRet("test_batch_norm_grad_split", "after");
+  EXPECT_TRUE(CheckEqualGraph(g_after, new_graph));
+}
+}  // namespace opt
+}  // namespace mindspore

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

@@ -0,0 +1,61 @@
+# 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.operations import _grad_ops as G
+from mindspore.ops import Primitive
+
+batch_norm_grad = G.BatchNormGrad(is_training=True)
+bn_training_update_grad = Primitive('BNTrainingUpdateGrad')
+bn_training_reduce_grad = Primitive('BNTrainingReduceGrad')
+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_batch_norm_grad_split(tag):
+    fns = FnDict()
+
+    @fns
+    def before(i0, i1, i2, i3, i4, i5):
+        bn_grad_output = batch_norm_grad(i0, i1, i2, i3, i4, i5)
+        item0 = tuple_getitem(bn_grad_output, 0)
+        item1 = tuple_getitem(bn_grad_output, 1)
+        item2 = tuple_getitem(bn_grad_output, 2)
+        output = make_tuple(item0, item1, item2)
+        return output
+
+    @fns
+    def after(i0, i1, i2, i3, i4, i5):
+        bn_update_grad_output = bn_training_update_grad(i0, i1, i3, i4)
+        update_item0 = tuple_getitem(bn_update_grad_output, 0)
+        update_item1 = tuple_getitem(bn_update_grad_output, 1)
+        bn_reduce_grad_output = bn_training_reduce_grad(i0, i1, update_item0, update_item1, i2, i3, i4)
+        output = make_tuple(bn_reduce_grad_output, update_item0, update_item1)
+        item0 = tuple_getitem(output, 0)
+        item1 = tuple_getitem(output, 1)
+        item2 = tuple_getitem(output, 2)
+        output = make_tuple(item0, item1, item2)
+        return make_tuple(output)
+
+    return fns[tag]