add BatchNormGrad2BNInferGrad pass

This commit is contained in:
huanghui 2020-05-13 11:23:04 +08:00
parent 42ba885e58
commit 709828a98b
6 changed files with 294 additions and 0 deletions

View File

@ -50,6 +50,7 @@
#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/ir_fusion/batchnormgrad_to_bninfergrad.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"
@ -102,6 +103,7 @@ void AddAscendBackendOptionalIRFusion(PassManager *ir_fusion_pm) {
ir_fusion_pm->AddPass(std::make_shared<TransposeTransDataFusion>());
ir_fusion_pm->AddPass(std::make_shared<GetitemTuple>());
ir_fusion_pm->AddPass(std::make_shared<BatchNorm2BNInfer>());
ir_fusion_pm->AddPass(std::make_shared<BatchNormGrad2BNInferGrad>());
}
} // namespace

View File

@ -96,6 +96,7 @@ bool NeedFusion(const FuncGraphPtr &graph, const AnfNodePtr &node, CNodePtr *bat
AnfNodePtr batchnorm_anf = tuple_getitem->input(kRealInputNodeIndexInTupleGetItem);
MS_EXCEPTION_IF_NULL(batchnorm_anf);
MS_EXCEPTION_IF_NULL(batchnorm);
*batchnorm = batchnorm_anf->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(*batchnorm);
return CheckBatchNorm(graph, *batchnorm);

View File

@ -0,0 +1,127 @@
/**
* 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/batchnormgrad_to_bninfergrad.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 CreateBNInferGrad(const FuncGraphPtr &graph, const CNodePtr &batchnormgrad, const AnfNodePtr &node) {
MS_EXCEPTION_IF_NULL(graph);
MS_EXCEPTION_IF_NULL(batchnormgrad);
auto prim = std::make_shared<Primitive>(kBNInferGradOpName);
std::vector<AnfNodePtr> inputs = {NewValueNode(prim)};
inputs.push_back(batchnormgrad->input(1));
inputs.push_back(batchnormgrad->input(3));
inputs.push_back(batchnormgrad->input(5));
auto new_node = graph->NewCNode(inputs);
MS_EXCEPTION_IF_NULL(new_node);
new_node->set_scope(batchnormgrad->scope());
new_node->set_abstract(node->abstract());
AnfAlgo::CopyNodeAttr(kAttrIsTraining, batchnormgrad, new_node);
AnfAlgo::CopyNodeAttr(kAttrEpsilon, batchnormgrad, 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 BatchNormGrad";
return false;
}
return true;
}
bool CheckBatchNormGrad(const FuncGraphPtr &graph, const CNodePtr &batchnormgrad) {
MS_EXCEPTION_IF_NULL(graph);
MS_EXCEPTION_IF_NULL(batchnormgrad);
if (batchnormgrad->size() < kBatchNormInputNum + 1) {
MS_LOG(DEBUG) << "BatchNormGrad's input less than " << kBatchNormInputNum;
return false;
}
if (!AnfAlgo::HasNodeAttr(kAttrIsTraining, batchnormgrad)) {
return false;
}
auto is_training = AnfAlgo::GetNodeAttr<bool>(batchnormgrad, kAttrIsTraining);
if (is_training) {
MS_LOG(DEBUG) << "is_training is true, no need do fusion";
return false;
}
if (IsUsedByOthers(graph, batchnormgrad)) {
MS_LOG(DEBUG) << "Only the 0th output of BatchNormGrad is used, then do fusion";
return false;
}
return true;
}
bool NeedFusion(const FuncGraphPtr &graph, const AnfNodePtr &node, CNodePtr *batchnormgrad) {
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 batchnormgrad_anf = tuple_getitem->input(kRealInputNodeIndexInTupleGetItem);
MS_EXCEPTION_IF_NULL(batchnormgrad_anf);
MS_EXCEPTION_IF_NULL(batchnormgrad);
*batchnormgrad = batchnormgrad_anf->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(*batchnormgrad);
return CheckBatchNormGrad(graph, *batchnormgrad);
}
} // namespace
const BaseRef BatchNormGrad2BNInferGrad::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 batchnormgrad({prim::kPrimBatchNormGrad, Xs});
VectorRef pattern({prim::kPrimTupleGetItem, batchnormgrad, Y});
return pattern;
}
const AnfNodePtr BatchNormGrad2BNInferGrad::Process(const FuncGraphPtr &graph, const AnfNodePtr &node,
const EquivPtr &) const {
MS_EXCEPTION_IF_NULL(graph);
MS_EXCEPTION_IF_NULL(node);
CNodePtr batchnormgrad = nullptr;
if (!NeedFusion(graph, node, &batchnormgrad)) {
return nullptr;
}
return CreateBNInferGrad(graph, batchnormgrad, node);
}
} // namespace opt
} // namespace mindspore

View File

@ -0,0 +1,34 @@
/**
* 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_BATCHNORMGRAD_TO_BNINFERGRAD_H_
#define MINDSPORE_CCSRC_PRE_ACTIVATE_ASCEND_IR_FUSION_BATCHNORMGRAD_TO_BNINFERGRAD_H_
#include <memory>
#include "pre_activate/common/optimizer.h"
namespace mindspore {
namespace opt {
class BatchNormGrad2BNInferGrad : public PatternProcessPass {
public:
explicit BatchNormGrad2BNInferGrad(bool multigraph = true)
: PatternProcessPass("batchnormgrad_to_bninfergrad", multigraph) {}
~BatchNormGrad2BNInferGrad() 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_BATCHNORMGRAD_TO_BNINFERGRAD_H_

View File

@ -0,0 +1,73 @@
/**
* 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/batchnormgrad_to_bninfergrad.h"
#include "debug/anf_ir_dump.h"
namespace mindspore {
namespace opt {
class TestHWOptimizeBatchNormGrad2BNInferGrad : public BackendCommon {
public:
TestHWOptimizeBatchNormGrad2BNInferGrad()
: get_py_fun_("gtest_input.pre_activate.batchnormgrad_to_bninfergrad", true) {}
~TestHWOptimizeBatchNormGrad2BNInferGrad() override = default;
UT::PyFuncGraphFetcher get_py_fun_;
};
TEST_F(TestHWOptimizeBatchNormGrad2BNInferGrad, test_fusion) {
FuncGraphPtr g = get_py_fun_.CallAndParseRet("test_batchnormgrad_to_bninfergrad", "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, x_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::BatchNormGrad2BNInferGrad>());
optimizer->AddPassManager(pm);
FuncGraphPtr new_graph = optimizer->Optimize(fg);
FuncGraphPtr g_after = get_py_fun_.CallAndParseRet("test_batchnormgrad_to_bninfergrad", "after");
EXPECT_TRUE(CheckEqualGraph(g_after, new_graph));
}
TEST_F(TestHWOptimizeBatchNormGrad2BNInferGrad, test_no_fusion) {
FuncGraphPtr g = get_py_fun_.CallAndParseRet("test_batchnormgrad_to_bninfergrad", "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, x_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::BatchNormGrad2BNInferGrad>());
optimizer->AddPassManager(pm);
FuncGraphPtr new_graph = optimizer->Optimize(fg);
EXPECT_TRUE(CheckEqualGraph(origin_graph, new_graph));
}
} // namespace opt
} // namespace mindspore

View File

@ -0,0 +1,57 @@
# 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=False)
bn_infer_grad = Primitive('BNInferGrad')
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_batchnormgrad_to_bninfergrad(tag):
fns = FnDict()
@fns
def before(input0, input1, input2, input3, input4):
res = batch_norm_grad(input0, input1, input2, input3, input4)
res = tuple_getitem(res, 0)
return res
@fns
def after(input0, input1, input2, input3, input4):
res = bn_infer_grad(input0, input2, input4)
return make_tuple(res)
@fns
def no_fusion(input0, input1, input2, input3, input4):
res = batch_norm_grad(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]