modify ConvertSwitchReplacement

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<FloatEnvGetItemSwitch>(), "float_env_getitem_switch", prim::kPrimEnvGetItem);
-  convert_switch_replacement_ =
-    MakeSubstitution(std::make_shared<ConvertSwitchReplacement>(), "convert_switch_replacement", IsCNodeDup);
   exchange_switch_depend_value_ =
     MakeSubstitution(std::make_shared<ExchangeSwitchDependValue>(), "exchange_switch_depend_value", prim::kPrimSwitch);
 

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_;

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<AnfNodePtr> &block_nodes,
   return GenerateMergeNodes(block_nodes, branch_output_abs, func_graph);
 }
 }  // namespace internal
+
+bool ConvertSwitchReplacement::CheckSwitchBranch(const AnfNodePtr &node) {
+  if (!IsValueNode<FuncGraph>(node)) {
+    return false;
+  }
+  // If graph contains FuncGraph, then ignore this node.
+  auto graph = GetValueNode<FuncGraphPtr>(node);
+  for (auto &item : graph->value_nodes()) {
+    auto value_node = item.first;
+    if (IsValueNode<FuncGraph>(value_node)) {
+      return false;
+    }
+  }
+  return true;
+}
+
+bool ConvertSwitchReplacement::CheckSwitchWrapNode(const AnfNodePtr &node) {
+  // {{prim::kPrimSwitch, X, G1, G2}, Xs}.
+  if (node->isa<CNode>()) {
+    auto inp0 = node->cast<CNodePtr>()->input(0);
+    if (IsPrimitiveCNode(inp0, prim::kPrimSwitch)) {
+      auto switch_node = inp0->cast<CNodePtr>();
+      // 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<CNodePtr>();
+  auto switch_cnode = cnode->input(0)->cast<CNodePtr>();
+  auto cond = switch_cnode->input(kCondIndex);
+  auto true_br = switch_cnode->input(kTrueBranchIndex);
+  auto false_br = switch_cnode->input(kFalseBranchIndex);
+
+  auto g1 = GetValueNode<FuncGraphPtr>(true_br);
+  auto g2 = GetValueNode<FuncGraphPtr>(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<AnfNodePtr> 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

mindspore/ccsrc/frontend/optimizer/irpass/branch_culling.h

@@ -95,63 +95,33 @@ AnfNodePtr TransformMergeBranches(const std::vector<AnfNodePtr> &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<CNode>() || 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<AnfNodePtr> all_nodes = DeepScopedGraphSearch(ret);
+
+    bool change = false;
+    for (auto &node : all_nodes) {
+      if (CheckSwitchWrapNode(node)) {
+        TransformSwitchBranchReplace(node);
+        change = true;
+      }
     }
-
-    PatternNode<AnfNodePtr> cond, true_br, false_br;
-
-    auto ConvertSwitchLambda = [&node, &cond, &true_br, &false_br]() -> AnfNodePtr {
-      auto g1_ = GetValueNode<FuncGraphPtr>(true_br.GetNode(node));
-      auto g2_ = GetValueNode<FuncGraphPtr>(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<FuncGraph>(value_node)) {
-          return nullptr;
-        }
-      }
-
-      for (auto &item : g2_->value_nodes()) {
-        auto value_node = item.first;
-        if (IsValueNode<FuncGraph>(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<AnfNodePtr> params;
-      auto cnode = node->cast<CNodePtr>();
-      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<FuncGraph>, node) && false_br.CheckFunc(IsValueNode<FuncGraph>, 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} ->

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));