From 0de0aee54e6632f400a18e85dc4c25d6ce1beda4 Mon Sep 17 00:00:00 2001 From: huanghui Date: Tue, 28 Apr 2020 21:37:02 +0800 Subject: [PATCH] Add BatchNormGrad split pass --- .../ir_fission/batch_norm_grad_split.cc | 132 ++++++++++++++++++ .../ascend/ir_fission/batch_norm_grad_split.h | 33 +++++ mindspore/ccsrc/utils/utils.h | 2 + .../ir_fission/batch_norm_grad_split_test.cc | 59 ++++++++ .../pre_activate/batch_norm_grad_split.py | 61 ++++++++ 5 files changed, 287 insertions(+) create mode 100644 mindspore/ccsrc/pre_activate/ascend/ir_fission/batch_norm_grad_split.cc create mode 100644 mindspore/ccsrc/pre_activate/ascend/ir_fission/batch_norm_grad_split.h create mode 100644 tests/ut/cpp/pre_activate/ascend/ir_fission/batch_norm_grad_split_test.cc create mode 100644 tests/ut/cpp/python_input/gtest_input/pre_activate/batch_norm_grad_split.py diff --git a/mindspore/ccsrc/pre_activate/ascend/ir_fission/batch_norm_grad_split.cc b/mindspore/ccsrc/pre_activate/ascend/ir_fission/batch_norm_grad_split.cc new file mode 100644 index 0000000000..cb8670dd00 --- /dev/null +++ b/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 +#include +#include + +#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 *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 bn_update_grad_inputs = { + NewValueNode(std::make_shared(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()); + 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 &bn_update_grad_outputs, + std::vector *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 bn_reduce_grad_inputs = { + NewValueNode(std::make_shared(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()); + 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(); + auto prim = std::make_shared(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(); + 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(cnode, kAttrIsTraining)) { + MS_LOG(INFO) << "is_training must be true"; + return nullptr; + } + + std::vector 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 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 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 diff --git a/mindspore/ccsrc/pre_activate/ascend/ir_fission/batch_norm_grad_split.h b/mindspore/ccsrc/pre_activate/ascend/ir_fission/batch_norm_grad_split.h new file mode 100644 index 0000000000..e539fdb27c --- /dev/null +++ b/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_ diff --git a/mindspore/ccsrc/utils/utils.h b/mindspore/ccsrc/utils/utils.h index 59d7f27c11..5b9ab0a1a0 100644 --- a/mindspore/ccsrc/utils/utils.h +++ b/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"; diff --git a/tests/ut/cpp/pre_activate/ascend/ir_fission/batch_norm_grad_split_test.cc b/tests/ut/cpp/pre_activate/ascend/ir_fission/batch_norm_grad_split_test.cc new file mode 100644 index 0000000000..68c327ade1 --- /dev/null +++ b/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 shp_x{1, 64, 112, 112}; + std::vector shp_b{64}; + auto x_abstract = std::make_shared(kFloat32, shp_x); + auto b_abstract = std::make_shared(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(); + auto pm = std::make_shared(); + auto pass = std::make_shared(); + 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 diff --git a/tests/ut/cpp/python_input/gtest_input/pre_activate/batch_norm_grad_split.py b/tests/ut/cpp/python_input/gtest_input/pre_activate/batch_norm_grad_split.py new file mode 100644 index 0000000000..dc783f1fbd --- /dev/null +++ b/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]