From 87e69c6fde8cb93839cf25722ea2ba09fefa36e9 Mon Sep 17 00:00:00 2001 From: huangbingjian Date: Tue, 20 Jul 2021 15:28:34 +0800 Subject: [PATCH] modify ConvertSwitchReplacement --- mindspore/ccsrc/frontend/optimizer/irpass.cc | 2 - mindspore/ccsrc/frontend/optimizer/irpass.h | 1 - .../optimizer/irpass/branch_culling.cc | 63 +++++++++++++++ .../optimizer/irpass/branch_culling.h | 78 ++++++------------- mindspore/ccsrc/pipeline/jit/pass.cc | 5 +- 5 files changed, 90 insertions(+), 59 deletions(-) diff --git a/mindspore/ccsrc/frontend/optimizer/irpass.cc b/mindspore/ccsrc/frontend/optimizer/irpass.cc index d3ff88e572f..b50c6c2080e 100644 --- a/mindspore/ccsrc/frontend/optimizer/irpass.cc +++ b/mindspore/ccsrc/frontend/optimizer/irpass.cc @@ -157,8 +157,6 @@ OptimizeIRPassLib::OptimizeIRPassLib() { "float_tuple_getitem_switch", prim::kPrimTupleGetItem); float_env_getitem_switch_ = MakeSubstitution(std::make_shared(), "float_env_getitem_switch", prim::kPrimEnvGetItem); - convert_switch_replacement_ = - MakeSubstitution(std::make_shared(), "convert_switch_replacement", IsCNodeDup); exchange_switch_depend_value_ = MakeSubstitution(std::make_shared(), "exchange_switch_depend_value", prim::kPrimSwitch); diff --git a/mindspore/ccsrc/frontend/optimizer/irpass.h b/mindspore/ccsrc/frontend/optimizer/irpass.h index 7b9f20fb7f3..ca0e76c7815 100644 --- a/mindspore/ccsrc/frontend/optimizer/irpass.h +++ b/mindspore/ccsrc/frontend/optimizer/irpass.h @@ -83,7 +83,6 @@ class OptimizeIRPassLib { SubstitutionPtr switch_simplify_; SubstitutionPtr float_tuple_getitem_switch_; SubstitutionPtr float_env_getitem_switch_; - SubstitutionPtr convert_switch_replacement_; SubstitutionPtr exchange_switch_depend_value_; SubstitutionPtr switch_partial_eliminater_; diff --git a/mindspore/ccsrc/frontend/optimizer/irpass/branch_culling.cc b/mindspore/ccsrc/frontend/optimizer/irpass/branch_culling.cc index e4160495c94..8598dd3094b 100644 --- a/mindspore/ccsrc/frontend/optimizer/irpass/branch_culling.cc +++ b/mindspore/ccsrc/frontend/optimizer/irpass/branch_culling.cc @@ -26,6 +26,9 @@ namespace mindspore { namespace opt { namespace irpass { +constexpr size_t kCondIndex = 1; +constexpr size_t kTrueBranchIndex = 2; +constexpr size_t kFalseBranchIndex = 3; namespace internal { AnfNodePtr GenerateSwitchNode(const FuncGraphPtr &graph, const AnfNodePtr &cond, const AnfNodePtr &data, int64_t switch_idx) { @@ -505,6 +508,66 @@ AnfNodePtr TransformMergeBranches(const std::vector &block_nodes, return GenerateMergeNodes(block_nodes, branch_output_abs, func_graph); } } // namespace internal + +bool ConvertSwitchReplacement::CheckSwitchBranch(const AnfNodePtr &node) { + if (!IsValueNode(node)) { + return false; + } + // If graph contains FuncGraph, then ignore this node. + auto graph = GetValueNode(node); + for (auto &item : graph->value_nodes()) { + auto value_node = item.first; + if (IsValueNode(value_node)) { + return false; + } + } + return true; +} + +bool ConvertSwitchReplacement::CheckSwitchWrapNode(const AnfNodePtr &node) { + // {{prim::kPrimSwitch, X, G1, G2}, Xs}. + if (node->isa()) { + auto inp0 = node->cast()->input(0); + if (IsPrimitiveCNode(inp0, prim::kPrimSwitch)) { + auto switch_node = inp0->cast(); + // for switch replace method, only graphs without graph inside can be replaced + if (CheckSwitchBranch(switch_node->input(kTrueBranchIndex)) && + CheckSwitchBranch(switch_node->input(kFalseBranchIndex))) { + return true; + } + } + } + return false; +} + +void ConvertSwitchReplacement::TransformSwitchBranchReplace(const AnfNodePtr &node) { + auto cnode = node->cast(); + auto switch_cnode = cnode->input(0)->cast(); + auto cond = switch_cnode->input(kCondIndex); + auto true_br = switch_cnode->input(kTrueBranchIndex); + auto false_br = switch_cnode->input(kFalseBranchIndex); + + auto g1 = GetValueNode(true_br); + auto g2 = GetValueNode(false_br); + auto true_output = g1->output()->abstract(); + auto false_output = g2->output()->abstract(); + auto trans_g1 = internal::TransformGraphCondTrueBranchNodes(g1, cond); + auto trans_g2 = internal::TransformGraphCondFalseBranchNodes(g2, cond); + + std::vector params; + if (cnode && cnode->size() > 1) { + // There are arguments for the call of switch result, + // usually these are monad states added by auto-monad. + for (size_t i = 1; i < cnode->size(); ++i) { + params.push_back(cnode->inputs().at(i)); + } + } + auto fg = node->func_graph(); + auto cloned_g1 = InlineClone(trans_g1, fg, params); + auto cloned_g2 = InlineClone(trans_g2, fg, params); + auto new_node = internal::TransformMergeBranches({cond, cloned_g1, cloned_g2}, {true_output, false_output}, fg); + fg->manager()->Replace(node, new_node); +} } // namespace irpass } // namespace opt } // namespace mindspore diff --git a/mindspore/ccsrc/frontend/optimizer/irpass/branch_culling.h b/mindspore/ccsrc/frontend/optimizer/irpass/branch_culling.h index adf9aa94f36..5864edd0e21 100644 --- a/mindspore/ccsrc/frontend/optimizer/irpass/branch_culling.h +++ b/mindspore/ccsrc/frontend/optimizer/irpass/branch_culling.h @@ -95,63 +95,33 @@ AnfNodePtr TransformMergeBranches(const std::vector &block_nodes, } // namespace internal // {{prim::kPrimSwitch, X, G1, G2}, Xs} -class ConvertSwitchReplacement : public OptimizerCaller { +class ConvertSwitchReplacement { public: - AnfNodePtr operator()(const OptimizerPtr &, const AnfNodePtr &node) override { - if (!node->isa() || node->func_graph() == nullptr) { - return nullptr; + ConvertSwitchReplacement() = default; + virtual ~ConvertSwitchReplacement() = default; + + bool operator()(const FuncGraphPtr &root, const OptimizerPtr &optimizer) { + AnfNodePtr ret = root->get_return(); + MS_EXCEPTION_IF_NULL(ret); + std::vector all_nodes = DeepScopedGraphSearch(ret); + + bool change = false; + for (auto &node : all_nodes) { + if (CheckSwitchWrapNode(node)) { + TransformSwitchBranchReplace(node); + change = true; + } } - - PatternNode cond, true_br, false_br; - - auto ConvertSwitchLambda = [&node, &cond, &true_br, &false_br]() -> AnfNodePtr { - auto g1_ = GetValueNode(true_br.GetNode(node)); - auto g2_ = GetValueNode(false_br.GetNode(node)); - auto x_ = cond.GetNode(node); - - // for switch replace method, only graphs without graph inside can be replaced - for (auto &item : g1_->value_nodes()) { - auto value_node = item.first; - if (IsValueNode(value_node)) { - return nullptr; - } - } - - for (auto &item : g2_->value_nodes()) { - auto value_node = item.first; - if (IsValueNode(value_node)) { - return nullptr; - } - } - - auto true_output = g1_->output()->abstract(); - auto false_output = g2_->output()->abstract(); - auto trans_g1 = internal::TransformGraphCondTrueBranchNodes(g1_, x_); - auto trans_g2 = internal::TransformGraphCondFalseBranchNodes(g2_, x_); - - std::vector params; - auto cnode = node->cast(); - if (cnode && cnode->size() > 1) { - // There are arguments for the call of switch result, - // usually these are monad states added by auto-monad. - for (size_t i = 1; i < cnode->size(); ++i) { - params.push_back(cnode->inputs().at(i)); - } - } - auto fg = node->func_graph(); - auto cloned_g1 = InlineClone(trans_g1, fg, params); - auto cloned_g2 = InlineClone(trans_g2, fg, params); - auto nnode = internal::TransformMergeBranches({x_, cloned_g1, cloned_g2}, {true_output, false_output}, fg); - - return nnode; - }; - - MATCH_REPLACE_LAMBDA_IF( - node, PCNode(PPrimitive(prim::kPrimSwitch, cond, true_br, false_br)).MinExtraNodes(0), ConvertSwitchLambda, - true_br.CheckFunc(IsValueNode, node) && false_br.CheckFunc(IsValueNode, node)); - - return nullptr; + return change; } + + private: + // Determine whether there are graphs inside the branch graph. + bool CheckSwitchBranch(const AnfNodePtr &node); + // Determine whether node matches {{prim::kPrimSwitch, X, G1, G2}, Xs}. + bool CheckSwitchWrapNode(const AnfNodePtr &node); + // Replace switch branch. + void TransformSwitchBranchReplace(const AnfNodePtr &node); }; // {prim::kPrimSwitch, {prim::kPrimDepend, ValueNode, X}, G1, G2} -> diff --git a/mindspore/ccsrc/pipeline/jit/pass.cc b/mindspore/ccsrc/pipeline/jit/pass.cc index f4e325b91b6..883e9dd717a 100644 --- a/mindspore/ccsrc/pipeline/jit/pass.cc +++ b/mindspore/ccsrc/pipeline/jit/pass.cc @@ -44,6 +44,7 @@ #include "pipeline/jit/pipeline_split.h" #include "pipeline/pynative/pynative_execute.h" #include "pipeline/jit/static_analysis/auto_monad.h" +#include "frontend/optimizer/irpass/branch_culling.h" #include "frontend/optimizer/irpass/gradient_eliminate.h" #include "frontend/optimizer/irpass/parameter_eliminate.h" #include "frontend/optimizer/irpass/updatestate_eliminate.h" @@ -442,7 +443,7 @@ OptPassGroupMap GetOptPassesC(const opt::irpass::OptimizeIRPassLib &) { } OptPassGroupMap GetControlPhases(const opt::irpass::OptimizeIRPassLib &irpass) { - opt::OptPassConfig control_group = opt::OptPassConfig({irpass.convert_switch_replacement_}, true); + opt::OptPassConfig control_group = opt::OptPassConfig(opt::irpass::ConvertSwitchReplacement()); OptPassGroupMap map({ {"control_group", control_group}, {"renormalize", opt::OptPassConfig::Renormalize()}, @@ -498,7 +499,7 @@ void InitOpt(const ResourcePtr &res) { g_pass_opts["opt_trans_graph"] = Optimizer::MakeOptimizer("opt_trans_graph", res, GetOptPassesTransformGraph(irpass), true, true); g_pass_opts["renormal"] = Optimizer::MakeOptimizer("renormal", res, GetOptPassesC(irpass)); - g_pass_opts["opt_control"] = Optimizer::MakeOptimizer("opt_control", res, GetControlPhases(irpass), false, true); + g_pass_opts["opt_control"] = Optimizer::MakeOptimizer("opt_control", res, GetControlPhases(irpass), true, true); g_pass_opts["opt_grad_epilogue"] = Optimizer::MakeOptimizer("opt_grad_epilogue", res, GetOptPynativeGradEpiloguePhases(irpass), true, false); g_pass_opts["opt_prepare"] = Optimizer::MakeOptimizer("opt_prepare", res, GetPreparePhases(irpass));