From 9f8a49861de03c08eb693d9aa68d2516e09a32f1 Mon Sep 17 00:00:00 2001
From: willzhang4a58 <willzhang4a58@gmail.com>
Date: Fri, 31 Mar 2017 03:28:12 -0400
Subject: [PATCH] template pattern for boxing build exec

---
 oneflow/graph/boxing_task_node.cpp     | 6 ++++++
 oneflow/graph/boxing_task_node.h       | 9 ++++++---
 oneflow/graph/in_boxing_task_node.cpp  | 4 +---
 oneflow/graph/in_boxing_task_node.h    | 3 +--
 oneflow/graph/out_boxing_task_node.cpp | 7 ++-----
 oneflow/graph/out_boxing_task_node.h   | 2 +-
 6 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/oneflow/graph/boxing_task_node.cpp b/oneflow/graph/boxing_task_node.cpp
index 547d6b9b42..6996854e0e 100644
--- a/oneflow/graph/boxing_task_node.cpp
+++ b/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());
diff --git a/oneflow/graph/boxing_task_node.h b/oneflow/graph/boxing_task_node.h
index c1a8b0851b..803dbbff1f 100644
--- a/oneflow/graph/boxing_task_node.h
+++ b/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();
   
 };
 
diff --git a/oneflow/graph/in_boxing_task_node.cpp b/oneflow/graph/in_boxing_task_node.cpp
index a86870fc5a..67ed03da57 100644
--- a/oneflow/graph/in_boxing_task_node.cpp
+++ b/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();
 }
 
diff --git a/oneflow/graph/in_boxing_task_node.h b/oneflow/graph/in_boxing_task_node.h
index 27c6a368ce..7f8db422fe 100644
--- a/oneflow/graph/in_boxing_task_node.h
+++ b/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;
 
 };
 
diff --git a/oneflow/graph/out_boxing_task_node.cpp b/oneflow/graph/out_boxing_task_node.cpp
index eff1ab2204..d29291947d 100644
--- a/oneflow/graph/out_boxing_task_node.cpp
+++ b/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();
 }
 
diff --git a/oneflow/graph/out_boxing_task_node.h b/oneflow/graph/out_boxing_task_node.h
index a2182adb65..e6d970a473 100644
--- a/oneflow/graph/out_boxing_task_node.h
+++ b/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;
 
 };
 
-- 
GitLab