diff --git a/mindspore/ccsrc/pre_activate/ascend/ascend_backend_optimization.cc b/mindspore/ccsrc/pre_activate/ascend/ascend_backend_optimization.cc index f01dd95f060be041f9bfc2ea46c376f7b8b12822..800c862c53f8b2738a31ad59fbbdd8d272d17625 100644 --- a/mindspore/ccsrc/pre_activate/ascend/ascend_backend_optimization.cc +++ b/mindspore/ccsrc/pre_activate/ascend/ascend_backend_optimization.cc @@ -100,6 +100,8 @@ void AddAscendBackendOptionalIRFusion(PassManager *ir_fusion_pm) { ir_fusion_pm->AddPass(std::make_shared()); ir_fusion_pm->AddPass(std::make_shared()); ir_fusion_pm->AddPass(std::make_shared()); + ir_fusion_pm->AddPass(std::make_shared()); + ir_fusion_pm->AddPass(std::make_shared()); ir_fusion_pm->AddPass(std::make_shared()); ir_fusion_pm->AddPass(std::make_shared()); ir_fusion_pm->AddPass(std::make_shared()); diff --git a/mindspore/ccsrc/pre_activate/ascend/ir_fusion/softmax_grad_ext_fusion.cc b/mindspore/ccsrc/pre_activate/ascend/ir_fusion/softmax_grad_ext_fusion.cc index ccb0cbfcb8788afb90291f5691faf4bb68600bb5..429523bd8ba66ae730885838dc531f0134c62034 100644 --- a/mindspore/ccsrc/pre_activate/ascend/ir_fusion/softmax_grad_ext_fusion.cc +++ b/mindspore/ccsrc/pre_activate/ascend/ir_fusion/softmax_grad_ext_fusion.cc @@ -31,6 +31,24 @@ const BaseRef SoftmaxGradExtFusion::DefinePattern() const { return mul_grad; } +const BaseRef SoftmaxGradExtFusionV2::DefinePattern() const { + VectorRef mul({prim::kPrimMul, input1_, input0_}); + VectorRef sum({sum_var_, mul}); + VectorRef sub({prim::kPrimSub, input0_, sum}); + VectorRef mul1({prim::kPrimMul, input1_, sub}); + VectorRef mul_grad({prim::kPrimMul, input2_, mul1}); + return mul_grad; +} + +const BaseRef SoftmaxGradExtFusionV3::DefinePattern() const { + VectorRef mul({prim::kPrimMul, input1_, input0_}); + VectorRef sum({sum_var_, mul}); + VectorRef sub({prim::kPrimSub, input0_, sum}); + VectorRef mul1({prim::kPrimMul, input1_, sub}); + VectorRef mul_grad({prim::kPrimMul, mul1, input2_}); + return mul_grad; +} + const AnfNodePtr SoftmaxGradExtFusion::Process(const FuncGraphPtr &graph, const AnfNodePtr &node, const EquivPtr &equiv) const { MS_EXCEPTION_IF_NULL(graph); @@ -46,7 +64,7 @@ const AnfNodePtr SoftmaxGradExtFusion::Process(const FuncGraphPtr &graph, const MS_EXCEPTION_IF_NULL(fusion_node); fusion_node->set_scope(node->scope()); fusion_node->set_abstract(node->abstract()); - AnfAlgo::CopyNodeAttr(kAttrKeepDims, sum, fusion_node); + AnfAlgo::CopyNodeAttr(kAttrKeepDims, "keepdims", sum, fusion_node); AnfAlgo::CopyNodeAttr(kAttrAxis, sum, fusion_node); return fusion_node; } diff --git a/mindspore/ccsrc/pre_activate/ascend/ir_fusion/softmax_grad_ext_fusion.h b/mindspore/ccsrc/pre_activate/ascend/ir_fusion/softmax_grad_ext_fusion.h index 70c5658e6086f42413ece583261f700ed7818aa7..59032e697335696203bb6dd6ac633a33c7eb3305 100644 --- a/mindspore/ccsrc/pre_activate/ascend/ir_fusion/softmax_grad_ext_fusion.h +++ b/mindspore/ccsrc/pre_activate/ascend/ir_fusion/softmax_grad_ext_fusion.h @@ -17,13 +17,15 @@ #define MINDSPORE_CCSRC_PRE_ACTIVATE_ASCEND_IR_FUSION_SOFTMAX_GRAD_EXT_FUSION_H_ #include +#include #include "pre_activate/common/optimizer.h" namespace mindspore { namespace opt { class SoftmaxGradExtFusion : public PatternProcessPass { public: - explicit SoftmaxGradExtFusion(bool multigraph = true) : PatternProcessPass("softmax_grad_ext_fusion", multigraph) { + explicit SoftmaxGradExtFusion(const std::string &name = "softmax_grad_ext_fusion", bool multigraph = true) + : PatternProcessPass(name, multigraph) { input0_ = std::make_shared(); input1_ = std::make_shared(); input2_ = std::make_shared(); @@ -33,12 +35,28 @@ class SoftmaxGradExtFusion : public PatternProcessPass { const BaseRef DefinePattern() const override; const AnfNodePtr Process(const FuncGraphPtr &, const AnfNodePtr &, const EquivPtr &) const override; - private: + protected: VarPtr input0_; VarPtr input1_; VarPtr input2_; VarPtr sum_var_; }; + +class SoftmaxGradExtFusionV2 : public SoftmaxGradExtFusion { + public: + explicit SoftmaxGradExtFusionV2(bool multigraph = true) + : SoftmaxGradExtFusion("softmax_grad_ext_fusion_v2", multigraph) {} + ~SoftmaxGradExtFusionV2() override = default; + const BaseRef DefinePattern() const override; +}; + +class SoftmaxGradExtFusionV3 : public SoftmaxGradExtFusion { + public: + explicit SoftmaxGradExtFusionV3(bool multigraph = true) + : SoftmaxGradExtFusion("softmax_grad_ext_fusion_v3", multigraph) {} + ~SoftmaxGradExtFusionV3() override = default; + const BaseRef DefinePattern() const override; +}; } // namespace opt } // namespace mindspore #endif // MINDSPORE_CCSRC_PRE_ACTIVATE_ASCEND_IR_FUSION_SOFTMAX_GRAD_EXT_FUSION_H_ diff --git a/tests/ut/cpp/pre_activate/ascend/ir_fusion/softmax_grad_ext_fusion_test.cc b/tests/ut/cpp/pre_activate/ascend/ir_fusion/softmax_grad_ext_fusion_test.cc index 25432336130956788222a1cfbde3107fee62627f..5f02f0e9c178a9573a4ce9bb493284f7d449d4d0 100644 --- a/tests/ut/cpp/pre_activate/ascend/ir_fusion/softmax_grad_ext_fusion_test.cc +++ b/tests/ut/cpp/pre_activate/ascend/ir_fusion/softmax_grad_ext_fusion_test.cc @@ -49,5 +49,47 @@ TEST_F(TestHWOptSoftmaxGradExtFusion, test_fusion) { FuncGraphPtr g_after = get_py_fun_.CallAndParseRet("test_softmax_grad_ext_fusion", "after"); EXPECT_TRUE(CheckEqualGraph(g_after, new_graph)); } + +TEST_F(TestHWOptSoftmaxGradExtFusion, test_fusion_v2) { + FuncGraphPtr g = get_py_fun_.CallAndParseRet("test_softmax_grad_ext_fusion_v2", "before"); + EXPECT_NE(g, nullptr); + std::vector shp{1, 1, 1, 1}; + auto x_abstract = std::make_shared(kFloat32, shp); + AbstractBasePtrList args_spec_list; + for (size_t i = 0; i < 3; ++i) { + args_spec_list.push_back(x_abstract); + } + auto fg = GetKernelGraph(g, args_spec_list); + + auto optimizer = std::make_shared(); + auto pm = std::make_shared(); + pm->AddPass(std::make_shared()); + optimizer->AddPassManager(pm); + FuncGraphPtr new_graph = optimizer->Optimize(fg); + + FuncGraphPtr g_after = get_py_fun_.CallAndParseRet("test_softmax_grad_ext_fusion_v2", "after"); + EXPECT_TRUE(CheckEqualGraph(g_after, new_graph)); +} + +TEST_F(TestHWOptSoftmaxGradExtFusion, test_fusion_v3) { + FuncGraphPtr g = get_py_fun_.CallAndParseRet("test_softmax_grad_ext_fusion_v3", "before"); + EXPECT_NE(g, nullptr); + std::vector shp{1, 1, 1, 1}; + auto x_abstract = std::make_shared(kFloat32, shp); + AbstractBasePtrList args_spec_list; + for (size_t i = 0; i < 3; ++i) { + args_spec_list.push_back(x_abstract); + } + auto fg = GetKernelGraph(g, args_spec_list); + + auto optimizer = std::make_shared(); + auto pm = std::make_shared(); + pm->AddPass(std::make_shared()); + optimizer->AddPassManager(pm); + FuncGraphPtr new_graph = optimizer->Optimize(fg); + + FuncGraphPtr g_after = get_py_fun_.CallAndParseRet("test_softmax_grad_ext_fusion_v3", "after"); + EXPECT_TRUE(CheckEqualGraph(g_after, new_graph)); +} } // namespace opt } // namespace mindspore diff --git a/tests/ut/cpp/python_input/gtest_input/pre_activate/softmax_grad_ext_fusion.py b/tests/ut/cpp/python_input/gtest_input/pre_activate/softmax_grad_ext_fusion.py index fbcc3d7480d738da391e3b940113d4e57379b209..38b8be0493d90d4e15f817aa7ec68ba0985139f6 100644 --- a/tests/ut/cpp/python_input/gtest_input/pre_activate/softmax_grad_ext_fusion.py +++ b/tests/ut/cpp/python_input/gtest_input/pre_activate/softmax_grad_ext_fusion.py @@ -54,3 +54,43 @@ def test_softmax_grad_ext_fusion(tag): return MakeTuple(res) return fns[tag] + + +def test_softmax_grad_ext_fusion_v2(tag): + fns = FnDict() + + @fns + def before(input0, input1, input2): + mul = Mul(input1, input0) + reduce_sum = ReduceSum(mul, axes) + sub = Sub(input0, reduce_sum) + mul1 = Mul(input1, sub) + mul_grad = Mul(input2, mul1) + return mul_grad + + @fns + def after(input0, input1, input2): + res = SoftmaxGradExt(input0, input1, input2) + return MakeTuple(res) + + return fns[tag] + + +def test_softmax_grad_ext_fusion_v3(tag): + fns = FnDict() + + @fns + def before(input0, input1, input2): + mul = Mul(input1, input0) + reduce_sum = ReduceSum(mul, axes) + sub = Sub(input0, reduce_sum) + mul1 = Mul(input1, sub) + mul_grad = Mul(mul1, input2) + return mul_grad + + @fns + def after(input0, input1, input2): + res = SoftmaxGradExt(input0, input1, input2) + return MakeTuple(res) + + return fns[tag]