add graph partition

mindspore/ccsrc/backend/session/executor.cc

@@ -89,7 +89,8 @@ bool TensorInVector(const VectorRef *outputs) {
 
 void CompileNodesTask::Run() {
   MS_EXCEPTION_IF_NULL(session_);
-  graph_id_ = session_->CompileGraphImpl(nodes_, output_nodes_);
+  MS_EXCEPTION_IF_NULL(segment_);
+  graph_id_ = session_->CompileGraphImpl(segment_->nodes_, output_nodes_);
 }
 
 void CompileGraphTask::Run() {
@@ -226,10 +227,11 @@ void Executor::SyncRunTask(const std::shared_ptr<Task> &task) {
   MsException::GetInstance().CheckException();
 }
 
-GraphId Executor::CompileGraph(const SessionPtr &session, const AnfNodePtrList &lst, const AnfNodePtrList &outputs) {
+GraphId Executor::CompileGraph(const SessionPtr &session, const GraphSegmentPtr &segment,
+                               const AnfNodePtrList &outputs) {
   auto task = std::make_shared<CompileNodesTask>();
   task->session_ = session;
-  task->nodes_ = lst;
+  task->segment_ = segment;
   task->output_nodes_ = outputs;
   SyncRunTask(task);
   return task->graph_id_;

mindspore/ccsrc/backend/session/executor.h

@@ -63,7 +63,7 @@ class CompileNodesTask : public Task {
   CompileNodesTask() { type_ = kCompileNodes; }
   ~CompileNodesTask() override = default;
   void Run() override;
-  AnfNodePtrList nodes_;
+  GraphSegmentPtr segment_;
   AnfNodePtrList output_nodes_;
   GraphId graph_id_{0};
 };
@@ -151,7 +151,7 @@ class Executor {
   ~Executor();
   void WorkerLoop();
   void WorkerJoin();
-  GraphId CompileGraph(const SessionPtr &session, const AnfNodePtrList &lst, const AnfNodePtrList &outputs);
+  GraphId CompileGraph(const SessionPtr &session, const GraphSegmentPtr &segment, const AnfNodePtrList &outputs);
   GraphId CompileGraph(const SessionPtr &session, NotNull<FuncGraphPtr> func_graph);
   void BuildGraph(const SessionPtr &session, GraphId graphId);
   void RunGraph(const SessionPtr &session, const GraphId &graph_id, const std::vector<tensor::TensorPtr> &inputs,

mindspore/ccsrc/backend/session/session_basic.cc

@@ -1388,9 +1388,9 @@ AnfNodePtr SessionBasic::FindPullNode(const AnfNodePtr &push_node, const std::ve
   return nullptr;
 }
 
-GraphId SessionBasic::CompileGraph(const AnfNodePtrList &lst, const AnfNodePtrList &outputs) {
+GraphId SessionBasic::CompileGraph(const GraphSegmentPtr &segment, const AnfNodePtrList &outputs) {
   MS_EXCEPTION_IF_NULL(executor_);
-  return executor_->CompileGraph(shared_from_this(), lst, outputs);
+  return executor_->CompileGraph(shared_from_this(), segment, outputs);
 }
 
 GraphId SessionBasic::CompileGraph(NotNull<FuncGraphPtr> func_graph) {

mindspore/ccsrc/backend/session/session_basic.h

@@ -68,7 +68,7 @@ class SessionBasic : public std::enable_shared_from_this<SessionBasic> {
 
   virtual ~SessionBasic() { summary_callback_ = nullptr; }
 
-  GraphId CompileGraph(const AnfNodePtrList &lst, const AnfNodePtrList &outputs);
+  GraphId CompileGraph(const GraphSegmentPtr &segment, const AnfNodePtrList &outputs);
   GraphId CompileGraph(NotNull<FuncGraphPtr> func_graph);
   void BuildGraph(GraphId graphId);
   void RunGraph(const GraphId &graph_id, const std::vector<tensor::TensorPtr> &inputs, VectorRef *outputs);
@@ -102,6 +102,8 @@ class SessionBasic : public std::enable_shared_from_this<SessionBasic> {
   virtual void GetModelInputsInfo(uint32_t graph_id, std::vector<tensor::TensorPtr> *inputs) const {}
   std::vector<tensor::TensorPtr> GetInputNeedLockTensors(const GraphId &graph_id,
                                                          const std::vector<tensor::TensorPtr> &inputs);
+  // Get graph by graph id, if not exist return null ptr
+  KernelGraphPtr GetGraph(GraphId graph_id) const;
 #ifdef ENABLE_DEBUGGER
   // set debugger
   void SetDebugger() {
@@ -147,8 +149,6 @@ class SessionBasic : public std::enable_shared_from_this<SessionBasic> {
   virtual void RunOpImpl(const OpRunInfo &op_run_info, const GraphInfo &graph_info,
                          const std::vector<tensor::TensorPtr> &input_tensors, VectorRef *outputs) {}
   void RunInfer(NotNull<FuncGraphPtr> func_graph, const std::vector<tensor::TensorPtr> &inputs);
-  // Get graph by graph id ,if not exist return null ptr
-  KernelGraphPtr GetGraph(GraphId graph_id) const;
 
   virtual void SetSummaryNodes(KernelGraph *graph);
 

mindspore/ccsrc/pipeline/jit/action.cc

@@ -354,7 +354,7 @@ bool TaskEmitAction(const ResourcePtr &res) {
   auto context_ptr = MsContext::GetInstance();
   std::string backend = MsContext::GetInstance()->backend_policy();
   MS_EXCEPTION_IF_NULL(context_ptr);
-  if (CompileGraphs::ContainMixedTarget(func_graph)) {
+  if (func_graph->ContainMultiTarget()) {
     bc_ptr->set_is_multi_graph_sink(false);
     context_ptr->set_param<bool>(MS_CTX_IS_MULTI_GRAPH_SINK, false);
     context_ptr->set_param<bool>(MS_CTX_ENABLE_LOOP_SINK, false);

mindspore/ccsrc/pipeline/jit/pipeline.cc

@@ -923,7 +923,8 @@ bool InitExecDatasetVm(const std::string &queue_name, int64_t size, int64_t batc
   MS_EXCEPTION_IF_NULL(convert_fn);
   // Convert CNodeList to LinConvertResult.
   ConfigManager::GetInstance().set_iter_num(1);
-  auto runner = convert_fn({app_init}, "");
+  auto segment = std::make_shared<GraphSegment>(std::vector<AnfNodePtr>{app_init}, false);
+  auto runner = convert_fn(segment, "");
   if (MsContext::GetInstance()->get_param<int>(MS_CTX_EXECUTION_MODE) != kPynativeMode) {
     backend->Link(runner.graph_id);
   }

mindspore/ccsrc/vm/backend.cc

@@ -34,30 +34,34 @@ namespace compile {
 bool Backend::GetCond(const BaseRef &c, bool *const value) { return BaseRefToBool(c, value); }
 bool Backend::GetIndex(const BaseRef &c, int64_t *const value) { return BaseRefToInt(utils::cast<ValuePtr>(c), value); }
 
-LinConvertResult MsBackend::MsConvert(const AnfNodePtrList &lst, const std::string &target) {
+Backend::Backend(const std::string &name) : name_(name) {
+  MS_LOG(DEBUG) << "select backend:" << name;
+  convert_fn_ = MsVmConvert;
+  is_multi_graph_sink_ = false;
+}
+
+LinConvertResult MsBackend::MsConvert(const GraphSegmentPtr &segment, const std::string &target) {
   MS_LOG(DEBUG) << "MsConvert";
+  MS_EXCEPTION_IF_NULL(segment);
   MS_EXCEPTION_IF_NULL(MsContext::GetInstance());
-  auto cached = g_ConvertCache.find(lst);
+  auto cached = g_ConvertCache.find(segment);
   if (cached != g_ConvertCache.end()) {
     return cached->second;
   }
-
   LinConvertResult result;
-
   FuncGraphPtr fg;
   AnfNodePtrList inputs;
   AnfNodePtrList outputs;
-
+  std::tie(fg, inputs, outputs) = TransformSegmentToAnfGraph(segment->nodes_);
-  std::tie(fg, inputs, outputs) = TransformSegmentToAnfGraph(lst);
   result.inputs = inputs;
   result.outputs = outputs;
   result.graph_id = kInvalidGraphId;
   GraphId graph_id = kInvalidGraphId;
   if (target != target_device_ && !target.empty()) {
     CreateOtherSession(target);
-    graph_id = other_sess_->CompileGraph(lst, outputs);
+    graph_id = other_sess_->CompileGraph(segment, outputs);
   } else {
-    graph_id = target_sess_->CompileGraph(lst, outputs);
+    graph_id = target_sess_->CompileGraph(segment, outputs);
   }
 
   if (MsContext::GetInstance()->get_param<bool>(MS_CTX_PRECOMPILE_ONLY)) {
@@ -79,7 +83,7 @@ LinConvertResult MsBackend::MsConvert(const AnfNodePtrList &lst, const std::stri
   result.graph_id = graph_id;
 
   graph_id_map_[graph_id] = result;
-  (void)g_ConvertCache.emplace(lst, result);
+  (void)g_ConvertCache.emplace(segment, result);
   return result;
 }
 
@@ -154,12 +158,6 @@ void MsBackend::Link(GraphId graph_id) {
   target_sess_->BuildGraph(graph_id);
 }
 
-Backend::Backend(const std::string &name) : name_(name) {
-  MS_LOG(DEBUG) << "select backend:" << name;
-  convert_fn_ = backends[name_];
-  is_multi_graph_sink_ = false;
-}
-
 MsBackend::MsBackend(const std::string &name, const std::string &target, uint32_t device_id) : Backend(name) {
   convert_fn_ = std::bind(&MsBackend::MsConvert, this, std::placeholders::_1, std::placeholders::_2);
   target_sess_ = session::SessionFactory::Get().Create(target);
@@ -194,6 +192,5 @@ VectorRef MsBackend::RunGraph(GraphId graph_id, const VectorRef &args) { return
 #ifdef ENABLE_DEBUGGER
 void MsBackend::SetDebugger() { target_sess_->SetDebugger(); }
 #endif
-
 }  // namespace compile
 }  // namespace mindspore

mindspore/ccsrc/vm/backend.h

@@ -25,6 +25,7 @@
 #include "utils/contract.h"
 #include "ir/anf.h"
 #include "vm/segment_runner.h"
+#include "vm/graph_partition.h"
 #include "vm/vm.h"
 #include "backend/session/session_basic.h"
 
@@ -63,7 +64,7 @@ class MsBackend : public Backend {
   MsBackend(const std::string &name, const std::string &target, uint32_t device_id);
   ~MsBackend() override = default;
 
-  LinConvertResult MsConvert(const AnfNodePtrList &lst, const std::string &target = "");
+  LinConvertResult MsConvert(const GraphSegmentPtr &segment, const std::string &target = "");
   VectorRef MsRunGraph(const GraphId &g, const VectorRef &args, const std::string &target = "");
 
   VectorRef MsSimuRunGraph(const GraphId &g, const VectorRef &args);

mindspore/ccsrc/vm/graph_partition.h

@@ -0,0 +1,48 @@
+/**
+ * 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_VM_GRAPH_PARTITION_H_
+#define MINDSPORE_CCSRC_VM_GRAPH_PARTITION_H_
+#include <vector>
+#include <string>
+#include <memory>
+#include "ir/anf.h"
+#include "ir/func_graph.h"
+#include "ir/graph_utils.h"
+#include "base/base_ref.h"
+
+namespace mindspore {
+extern const char kMsVm[];
+extern const char kGeVm[];
+extern const char kMsConvert[];
+
+namespace compile {
+class GraphPartition {
+ public:
+  explicit GraphPartition(const std::vector<PrimitivePtr> &cut_list, const std::string &backend_name);
+  ~GraphPartition() = default;
+  std::vector<GraphSegmentPtr> Partition(const FuncGraphPtr &func_graph);
+
+ private:
+  bool IsCut(const AnfNodePtr &node);
+  std::vector<PrimitivePtr> cut_list_;
+  std::string backend_name_;
+};
+
+using GraphPartitionPtr = std::shared_ptr<GraphPartition>;
+}  // namespace compile
+}  // namespace mindspore
+#endif  // MINDSPORE_CCSRC_VM_GRAPH_PARTITION_H_

mindspore/ccsrc/vm/segment_runner.cc

@@ -34,10 +34,6 @@
 #include "frontend/operator/ops.h"
 
 namespace mindspore {
-const char kMsConvert[] = "ms";
-const char kMsVm[] = "vm";
-const char kGeVm[] = "ge";
-
 namespace compile {
 // cached conversion
 ConvertCache g_ConvertCache;
@@ -194,8 +190,9 @@ std::tuple<FuncGraphPtr, AnfNodePtrList, AnfNodePtrList> TransformSegmentToAnfGr
 //   This implementation will convert the nodes into a subgraph
 //   that will run using the MsVM.
 template <typename T>
-LinConvertResult Convert(const AnfNodePtrList &lst, const std::string &) {
+LinConvertResult Convert(const GraphSegmentPtr &segment, const std::string &) {
-  auto cached = g_ConvertCache.find(lst);
+  MS_EXCEPTION_IF_NULL(segment);
+  auto cached = g_ConvertCache.find(segment);
   if (cached != g_ConvertCache.end()) {
     return cached->second;
   }
@@ -206,7 +203,7 @@ LinConvertResult Convert(const AnfNodePtrList &lst, const std::string &) {
   AnfNodePtrList inputs;
   AnfNodePtrList outputs;
 
-  std::tie(fg, inputs, outputs) = TransformSegmentToAnfGraph(lst);
+  std::tie(fg, inputs, outputs) = TransformSegmentToAnfGraph(segment->nodes_);
 
   // Clone in case g contains subgraphs that have a different manager
   fg = BasicClone(fg);
@@ -219,18 +216,15 @@ LinConvertResult Convert(const AnfNodePtrList &lst, const std::string &) {
   result.outputs = outputs;
   result.graph_id = UINT32_MAX;
 
-  (void)g_ConvertCache.emplace(lst, result);
+  (void)g_ConvertCache.emplace(segment, result);
   return result;
 }
 
 LinkFuncType MsVmConvert = Convert<VM>;
 
-std::unordered_map<std::string, LinkFuncType> backends = {{kMsVm, MsVmConvert}};
-
 std::set<std::string> backend_list = {
   kMsConvert,
   kMsVm,
 };
-
 }  // namespace compile
 }  // namespace mindspore

mindspore/ccsrc/vm/segment_runner.h

@@ -27,14 +27,10 @@
 
 #include "ir/anf.h"
 #include "vm/vmimpl.h"
+#include "vm/graph_partition.h"
 
 namespace mindspore {
-extern const char kMsVm[];
-extern const char kGeVm[];
-extern const char kMsConvert[];
-
 namespace compile {
-
 struct LinConvertResult {
   RunFuncPtr run;
   RunFuncPtr simu_run;
@@ -43,11 +39,9 @@ struct LinConvertResult {
   uint32_t graph_id;
 };
 
-using LinkFuncType = std::function<LinConvertResult(const AnfNodePtrList &, const std::string &)>;
+using LinkFuncType = std::function<LinConvertResult(const GraphSegmentPtr &, const std::string &)>;
-using ConvertCache = std::unordered_map<BaseRef, LinConvertResult, BaseRefHash>;
+using ConvertCache = std::unordered_map<GraphSegmentPtr, LinConvertResult>;
 extern LinkFuncType MsVmConvert;
-extern LinkFuncType GeVmConvert;
-extern std::unordered_map<std::string, LinkFuncType> backends;
 extern ConvertCache g_ConvertCache;
 extern std::set<std::string> backend_list;
 

mindspore/ccsrc/vm/transform.h

@@ -31,6 +31,7 @@
 #include "frontend/operator/ops.h"
 #include "vm/segment_runner.h"
 #include "vm/backend.h"
+#include "vm/graph_partition.h"
 
 // mindspore namespace is the top level namespace of MindSpore project.
 // Other namespace should be a sub namespace of mindspore namespace in the ME project.
@@ -59,7 +60,6 @@ class CompileGraph {
   void Tie(const AnfNodePtr &n1, const AnfNodePtr &n2) { slots_[n2] = slots_[n1]; }
   void Ret(int64_t nargs);
   int64_t Ref(const AnfNodePtr &node);
-  VectorRef SplitNodes(const FuncGraphPtr &func_graph);
 
   void set_height(int64_t h) {
     height_ = h;
@@ -76,10 +76,9 @@ class CompileGraph {
   }
 
  private:
-  VectorRef SplitNodesWithTarget(const std::vector<AnfNodePtr> &input_nodes, const FuncGraphPtr &graph);
   void PushParameters(const FuncGraphPtr &func_graph);
-  bool SplitGraph(const FuncGraphPtr &func_graph);
+  bool Compile(const FuncGraphPtr &func_graph);
-  int64_t LinConvert(const FuncGraphPtr &func_graph, const AnfNodePtrList &node_list, const std::string &target = "");
+  int64_t LinConvert(const FuncGraphPtr &func_graph, const GraphSegmentPtr &segment, const std::string &target = "");
   int64_t InterpretNode(const FuncGraphPtr &func_graph, const CNodePtr &node);
   int64_t AddCall(const FuncGraphPtr &graph, const CNodePtr &node);
   void AddPadStack(int64_t param_height);
@@ -97,11 +96,12 @@ class CompileGraph {
   void AddInst(const Instruction &inst, const VectorRef &args);
 
   BackendPtr backend_;
+  GraphPartitionPtr graph_partition_;
   LinkFuncType lin_convert_;
-  bool is_gevm_convert_;
+
   int64_t height_{0};
   int64_t max_height_{0};
-  std::vector<PrimitivePtr> cut_list_;
+
   std::unordered_map<AnfNodePtr, int64_t> slots_;
   InstSet inst_;
 };
@@ -123,7 +123,6 @@ class CompileGraphs {
   void Compile(const FuncGraphPtr &func_graph);
   FinalVMPtr Link(const FuncGraphPtr &func_graph);
   FinalVMPtr CompileAndLink(const FuncGraphPtr &func_graph);
-  static bool ContainMixedTarget(const FuncGraphPtr &graph);
 
  private:
   InstSet insts_;

mindspore/core/ir/anf.cc

@@ -301,4 +301,20 @@ std::string GetCNodeTarget(const AnfNodePtr &node) {
   }
   return default_target;
 }
+
+bool ContainMultiTarget(const std::vector<AnfNodePtr> &nodes) {
+  auto context_ptr = MsContext::GetInstance();
+  MS_EXCEPTION_IF_NULL(context_ptr);
+  std::string last_target = context_ptr->get_param<std::string>(MS_CTX_DEVICE_TARGET);
+  for (auto &node : nodes) {
+    if (node->isa<CNode>()) {
+      std::string cur_target = GetCNodeTarget(node);
+      if (last_target != cur_target) {
+        return true;
+      }
+      last_target = cur_target;
+    }
+  }
+  return false;
+}
 }  // namespace mindspore

mindspore/core/ir/anf.h

@@ -482,6 +482,13 @@ void reset_id();
 using TaggedNodeMap = std::unordered_map<AnfNodePtr, size_t>;
 using TaggedGraph = std::pair<FuncGraphPtr, TaggedNodeMap>;
 std::string GetCNodeTarget(const AnfNodePtr &node);
+bool ContainMultiTarget(const std::vector<AnfNodePtr> &nodes);
+struct GraphSegment {
+  GraphSegment(const std::vector<AnfNodePtr> &nodes, bool is_cut) : nodes_(nodes), is_cut_(is_cut) {}
+  std::vector<AnfNodePtr> nodes_;
+  bool is_cut_{false};
+};
+using GraphSegmentPtr = std::shared_ptr<GraphSegment>;
 }  // namespace mindspore
 
 #endif  // MINDSPORE_CORE_IR_ANF_H_

mindspore/core/ir/func_graph.cc

@@ -647,6 +647,19 @@ ParameterPtr FuncGraph::add_weight(const tensor::MetaTensorPtr &meta_tensor) {
   return parameter;
 }
 
+bool FuncGraph::ContainMultiTarget() const {
+  auto graph_manager = manager();
+  MS_EXCEPTION_IF_NULL(graph_manager);
+  FuncGraphSet graphs = graph_manager->func_graphs();
+  for (auto &g : graphs) {
+    auto nodes = TopoSort(g->get_return());
+    if (mindspore::ContainMultiTarget(nodes)) {
+      return true;
+    }
+  }
+  return false;
+}
+
 size_t NewFgSeenGeneration() {
   static size_t fg_seen_generation = 0;
   return ++fg_seen_generation;

mindspore/core/ir/func_graph.h

@@ -354,6 +354,7 @@ class FuncGraph : public FuncGraphBase {
   static void set_drawer(Drawer drawer) { drawer_ = drawer; }
   std::shared_ptr<bool> switch_layer_input() const { return switch_layer_input_; }
   void set_switch_layer_input(std::shared_ptr<bool> switch_layer_input) { switch_layer_input_ = switch_layer_input; }
+  bool ContainMultiTarget() const;
 
  private:
   // graph is manipulated by manager and others

tests/ut/cpp/vm/segment_runner_test.cc

@@ -52,21 +52,11 @@ TEST_F(TestCompileSegmentRunner, test_MsVmConvert1) {
   std::shared_ptr<mindspore::FuncGraphManager> manager = mindspore::Manage(g);
 
   BackendPtr b = std::make_shared<Backend>("vm");
-  CompileGraph transform_(b);
+  auto graph_partition = std::make_shared<GraphPartition>(nonlinear_ops, b->name());
-  auto splits = transform_.SplitNodes(g);
+  auto segments = graph_partition->Partition(g);
   VectorRef args({1.0, 2.0});
 
-  std::vector<BaseRef> todos(splits.size());
+  auto convertResult = MsVmConvert(segments[0], "");
-  auto it = std::copy_if(std::begin(splits), std::end(splits), std::begin(todos),
-                         [](const BaseRef &seg) -> bool { return utils::isa<VectorRef>(seg); });
-  todos.resize(std::distance(todos.begin(), it));
-  ASSERT_EQ(todos.size(), 1);
-
-  AnfNodePtrList anf_list;
-  for (auto &item : utils::cast<VectorRef>(todos[0])) {
-    anf_list.push_back(utils::cast<AnfNodePtr>(item));
-  }
-  auto convertResult = MsVmConvert(anf_list, "");
   auto runResult = (*(convertResult.run))(args);
   ASSERT_TRUE(runResult.size() == 1 && py::cast<double>(BaseRefToPyData(runResult[0])) == 3.0);
 }
@@ -76,21 +66,11 @@ TEST_F(TestCompileSegmentRunner, test_MsVmConvert2) {
   std::shared_ptr<mindspore::FuncGraphManager> manager = mindspore::Manage(g);
 
   BackendPtr b = std::make_shared<Backend>("vm");
-  CompileGraph transform_(b);
+  auto graph_partition = std::make_shared<GraphPartition>(nonlinear_ops, b->name());
-  auto splits = transform_.SplitNodes(g);
+  auto segments = graph_partition->Partition(g);
   VectorRef args({1.0, 2.0});
 
-  std::vector<BaseRef> todos(splits.size());
+  auto convertResult = MsVmConvert(segments[0], "");
-  auto it = std::copy_if(std::begin(splits), std::end(splits), std::begin(todos),
-                         [](const BaseRef &seg) -> bool { return utils::isa<VectorRef>(seg); });
-  todos.resize(std::distance(todos.begin(), it));
-  ASSERT_EQ(todos.size(), 1);
-
-  AnfNodePtrList anf_list;
-  for (auto &item : utils::cast<VectorRef>(todos[0])) {
-    anf_list.push_back(utils::cast<AnfNodePtr>(item));
-  }
-  auto convertResult = MsVmConvert(anf_list, "");
   auto runResult = (*(convertResult.run))(args);
   ASSERT_TRUE(runResult.size() == 1 && py::cast<double>(BaseRefToPyData(runResult[0])) == 2.0);
 }
@@ -100,21 +80,11 @@ TEST_F(TestCompileSegmentRunner, test_if) {
   std::shared_ptr<mindspore::FuncGraphManager> manager = mindspore::Manage(g);
 
   BackendPtr b = std::make_shared<Backend>("vm");
-  CompileGraph transform_(b);
+  auto graph_partition = std::make_shared<GraphPartition>(nonlinear_ops, b->name());
-  auto splits = transform_.SplitNodes(g);
+  auto segments = graph_partition->Partition(g);
   VectorRef args({1.0, 2.0});
 
-  std::vector<BaseRef> todos(splits.size());
+  auto convertResult = MsVmConvert(segments[0], "");
-  auto it = std::copy_if(std::begin(splits), std::end(splits), std::begin(todos),
-                         [](const BaseRef &seg) -> bool { return utils::isa<VectorRef>(seg); });
-  todos.resize(std::distance(todos.begin(), it));
-  ASSERT_EQ(todos.size(), 1);
-
-  AnfNodePtrList anf_list;
-  for (auto &item : utils::cast<VectorRef>(todos[0])) {
-    anf_list.push_back(utils::cast<AnfNodePtr>(item));
-  }
-  auto convertResult = MsVmConvert(anf_list, "");
   auto runResult = (*(convertResult.run))(args);
 
   auto result = py::cast<bool>(BaseRefToPyData(runResult[0]));