template pattern for boxing build exec

oneflow/graph/boxing_task_node.cpp

@@ -55,6 +55,12 @@ OpPair FwBuildBoxingOpModelModel() {
 
 }
 
+void BoxingTaskNode::FwBuildExecGraphAndSetProducedRegisterDescs() {
+  SetOutEdgeRegisterPtr();
+  FwBuildExecGraph();
+  SetProducedRegister();
+}
+
 void BoxingTaskNode::SetOutEdgeRegisterPtr() {
   for (TaskEdge* edge : out_edges()) {
     std::string name = "boxing_out_" + std::to_string(edge->edge_id());

oneflow/graph/boxing_task_node.h

@@ -20,7 +20,6 @@ class BoxingTaskNode : public TaskNode {
       std::pair<const ChainNode*, std::vector<const TaskEdge*>>;
   using Chain2EdgesMap =
       std::unordered_map<const ChainNode*, std::vector<const TaskEdge*>>;
-  void SetOutEdgeRegisterPtr();
   void FwInitChain2SortedEdgesMaps(
       Chain2EdgesMap* chain2sorted_edges,
       const std::unordered_set<TaskEdge*>& (TaskNode::*in_out_edges)() const,
@@ -33,10 +32,14 @@ class BoxingTaskNode : public TaskNode {
   void FwBuildChainSortedEdgesPair(
       const ChainEdgesPair& chain_sorted_in_edges,
       const ChainEdgesPair& chain_sorted_out_edges);
-  void SetProducedRegister();
-  void BpBuildExecGraphAndSetProducedRegisterDescs() override;
+  virtual void FwBuildExecGraph() = 0;
 
  private:
+  void FwBuildExecGraphAndSetProducedRegisterDescs() override;
+  void BpBuildExecGraphAndSetProducedRegisterDescs() override;
+  
+  void SetOutEdgeRegisterPtr();
+  void SetProducedRegister();
   
 };
 

oneflow/graph/in_boxing_task_node.cpp

@@ -5,8 +5,7 @@
 
 namespace oneflow {
 
-void InBoxingTaskNode::FwBuildExecGraphAndSetProducedRegisterDescs() {
-  SetOutEdgeRegisterPtr();
+void InBoxingTaskNode::FwBuildExecGraph() {
   Chain2EdgesMap chain2sorted_in_edges;
   FwInitChain2SortedEdgesMaps(&chain2sorted_in_edges, 
                               &TaskNode::in_edges,
@@ -21,7 +20,6 @@ void InBoxingTaskNode::FwBuildExecGraphAndSetProducedRegisterDescs() {
   for (const ChainEdgesPair& chain_sorted_in_edges : chain2sorted_in_edges) {
     FwBuildChainSortedEdgesPair(chain_sorted_in_edges, chain_sorted_out_edges);
   }
-  SetProducedRegister();
   mut_exec_graph().UpdateSourceAndSink();
 }
 

oneflow/graph/in_boxing_task_node.h

@@ -18,8 +18,7 @@ class InBoxingTaskNode final : public BoxingTaskNode {
   void InitWithFwNode(TaskNode* fw_node) override {
     BoxingTaskNode::InitWithFwNode(fw_node);
   }
-
-  void FwBuildExecGraphAndSetProducedRegisterDescs() override;
+  void FwBuildExecGraph() override;
 
 };
 

oneflow/graph/out_boxing_task_node.cpp

@@ -2,9 +2,7 @@
 
 namespace oneflow {
 
-// In future, we can use template-pattern
-void OutBoxingTaskNode::FwBuildExecGraphAndSetProducedRegisterDescs() {
-  SetOutEdgeRegisterPtr();
+void OutBoxingTaskNode::FwBuildExecGraph() {
   Chain2EdgesMap chain2sorted_out_edges;
   FwInitChain2SortedEdgesMaps(&chain2sorted_out_edges,
                               &TaskNode::out_edges,
@@ -13,13 +11,12 @@ void OutBoxingTaskNode::FwBuildExecGraphAndSetProducedRegisterDescs() {
   ChainEdgesPair chain_sorted_in_edges;
   chain_sorted_in_edges.first = chain_node();
   chain_sorted_in_edges.second.assign(in_edges().begin(), in_edges().end());
-  FwSortEdgesInnerStage(&chain_sorted_in_edges,
+  FwSortEdgesInnerStage(&chain_sorted_in_edges.second,
                         &TaskEdge::src_node,
                         &TaskNode::SoleInEdge);
   for (const ChainEdgesPair& chain_sorted_out_edges : chain2sorted_out_edges) {
     FwBuildChainSortedEdgesPair(chain_sorted_in_edges, chain_sorted_out_edges);
   }
-  SetProducedRegister();
   mut_exec_graph().UpdateSourceAndSink();
 }
 

oneflow/graph/out_boxing_task_node.h

@@ -18,7 +18,7 @@ class OutBoxingTaskNode final : public BoxingTaskNode {
   void InitWithFwNode(TaskNode* fw_node) override {
     BoxingTaskNode::InitWithFwNode(fw_node);
   }
-  void FwBuildExecGraphAndSetProducedRegisterDescs() override;
+  void FwBuildExecGraph() override;
 
 };