Feat graph autograd engine

oneflow/core/autograd/autograd_engine.cpp

@@ -14,6 +14,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
+#include <stack>
+#include <queue>
 #include "oneflow/core/autograd/autograd_engine.h"
 #include "oneflow/core/autograd/autograd_meta.h"
 #include "oneflow/core/framework/tensor.h"
@@ -70,7 +72,7 @@ StackFunctionNode::StackFunctionNode(
   for (int i = 0; i < inputs.size(); ++i) {
     input_meta_datas_.at(i) = inputs.at(i)->mut_autograd_meta();
     if (input_meta_datas_.at(i)->requires_grad()) {
-      next_functions_->emplace_back(inputs.at(i)->grad_fn_node());
+      next_functions_->emplace_back(inputs.at(i)->mut_grad_fn_node());
     }
   }
 
@@ -85,14 +87,14 @@ StackFunctionNode::StackFunctionNode(
   is_in_stack_ = false;
 }
 
-Maybe<void> StackFunctionNode::AccGrad4RetainGradTensor() {
+Maybe<void> FunctionNode::AccGrad4RetainGradTensor() {
   for (const std::shared_ptr<AutogradMeta>& out : output_meta_datas_) {
     if (out->retain_grad()) { JUST(CopyOrAccGrad(out.get(), /*autograd_mode=*/false)); }
   }
   return Maybe<void>::Ok();
 }
 
-Maybe<void> StackFunctionNode::AccGrad4LeafTensor(bool create_graph) {
+Maybe<void> FunctionNode::AccGrad4LeafTensor(bool create_graph) {
   for (const std::shared_ptr<AutogradMeta>& out : output_meta_datas_) {
     if (out->is_leaf() && out->requires_grad()) {
       JUST(CopyOrAccGrad(out.get(), /*autograd_mode=*/false));
@@ -101,19 +103,18 @@ Maybe<void> StackFunctionNode::AccGrad4LeafTensor(bool create_graph) {
   return Maybe<void>::Ok();
 }
 
-void StackFunctionNode::ReleaseOutTensorArgs() {
+void FunctionNode::ReleaseOutTensorArgs() {
   for (const std::shared_ptr<AutogradMeta>& meta_data : output_meta_datas_) {
     meta_data->now_grad_arg()->Release();
   }
 }
 
 void StackFunctionNode::ReleaseData() {
-  // Releases backward function and makes useless tensors release as early as possible
   if (!input_meta_datas_.empty()) { backward_fn_.reset(); }
   is_in_stack_ = false;
 }
 
-Maybe<bool> StackFunctionNode::Apply(bool create_graph) {
+Maybe<bool> FunctionNode::Apply(bool create_graph) {
   CHECK_NOTNULL_OR_RETURN(backward_fn_.get())
       << "This FunctionNode with name `" << GetOpTypeName() << "` has been released.";
   if (!IsReadyToRun(output_meta_datas_)) { return false; }
@@ -191,10 +192,11 @@ Maybe<TensorTuple> StackAutogradEngine::RunBackwardAndReturnInputsTensorGrad(
       if (!retain_graph) { func_node->ReleaseData(); }
     }
   }
+  // Gets input grads and resume retain_grad
   for (int i = 0; i < inputs.size(); ++i) {
     input_now_grads->at(i) = inputs.at(i)->acc_grad();
     if (!ori_retain_grad.at(i)) {
-      inputs.at(i)->mut_acc_grad().reset();
+      inputs.at(i)->set_acc_grad(nullptr);
       inputs.at(i)->set_retain_grad(false);
     }
   }
@@ -230,8 +232,210 @@ std::shared_ptr<FunctionNode> StackAutogradEngine::AddBackwardFuncPtr(
   return func_node;
 }
 
+void GraphFunctionNode::ReleaseData() {
+  if (!input_meta_datas_.empty()) { backward_fn_.reset(); }
+}
+
+GraphFunctionNode::GraphFunctionNode(
+    const std::string& op_type_name,
+    const std::shared_ptr<const std::function<Maybe<void>(const TensorTuple&, TensorTuple*, bool)>>&
+        backward_fn,
+    const TensorTuple& inputs, const TensorTuple& outputs)
+    : FunctionNode(op_type_name) {
+  input_meta_datas_.resize(inputs.size());
+  next_functions_->reserve(inputs.size());
+  for (int i = 0; i < inputs.size(); ++i) {
+    input_meta_datas_.at(i) = inputs.at(i)->mut_autograd_meta();
+    if (input_meta_datas_.at(i)->requires_grad()) {
+      next_functions_->emplace_back(inputs.at(i)->mut_grad_fn_node());
+    }
+  }
+
+  output_meta_datas_.resize(outputs.size());
+  output_tensor_infos_.reserve(outputs.size());
+  for (int i = 0; i < outputs.size(); ++i) {
+    output_meta_datas_.at(i) = outputs.at(i)->mut_autograd_meta();
+    output_tensor_infos_.emplace_back(TensorInfo(*outputs.at(i)));
+  }
+
+  backward_fn_ = backward_fn;
+}
+
+GraphTask::GraphTask(const TensorTuple& outputs, bool retain_graph, bool create_graph)
+    : retain_graph_(retain_graph), create_graph_(create_graph) {
+  roots_.reserve(outputs.size());
+  for (const auto& out_tensor : outputs) {
+    FunctionNode* node = out_tensor->mut_grad_fn_node().get();
+    roots_.push_back(node);
+    dependencies_.insert(std::make_pair(node, 0));
+  }
+}
+
+// Computes the number of dependencies for each FunctionNode
+Maybe<void> GraphTask::ComputeDependencies() {
+  HashSet<FunctionNode*> seen;
+  std::stack<FunctionNode*> stack;
+  for (FunctionNode* node : roots_) { stack.push(node); }
+
+  while (!stack.empty()) {
+    FunctionNode* node = stack.top();
+    stack.pop();
+    if (/*bool has_seen=*/!seen.insert(node).second) { continue; }
+    for (const auto& next_grad_fn : *(node->GetNextFunctions())) {
+      FunctionNode* next_node = next_grad_fn.get();
+      dependencies_[next_node] += 1;
+      if (seen.find(next_node) == seen.end()) { stack.push(next_node); }
+    }
+  }
+  return Maybe<void>::Ok();
+}
+
+// Computes the number of dependencies for each FunctionNode and prunes useless FunctionNode
+// according to input tensors
+Maybe<void> GraphTask::ComputeDependenciesAndPruneNode(const TensorTuple& inputs) {
+  struct NodeFrame {
+    NodeFrame(FunctionNode* node) : node_(node), next_function_idx_(0) {}
+    FunctionNode* node_;
+    size_t next_function_idx_;
+
+    FunctionNode* GetNextFunction() {
+      if (next_function_idx_ < node_->GetNextFunctions()->size()) {
+        next_function_idx_ += 1;
+        return node_->GetNextFunctions()->at(next_function_idx_ - 1).get();
+      } else {
+        return nullptr;
+      }
+    }
+  };
+
+  for (const auto& input : inputs) {
+    CHECK_NOTNULL_OR_RETURN(input->mut_grad_fn_node().get());
+    need_execute_.insert(input->mut_grad_fn_node().get());
+  }
+
+  HashSet<FunctionNode*> seen;
+  std::stack<NodeFrame> stack;
+
+  // Note: dfs to determine each FunctionNode should execute or not.
+  for (const auto& root : roots_) { stack.push(NodeFrame(root)); }
+  while (!stack.empty()) {
+    NodeFrame& frame = stack.top();
+    if (/*bool has_seen=*/seen.find(frame.node_) != seen.end()) {
+      stack.pop();
+      continue;
+    }
+    if (FunctionNode* node = frame.GetNextFunction()) {
+      dependencies_[node] += 1;
+      if (seen.find(node) == seen.end()) {
+        stack.push(NodeFrame(node));
+        continue;  // recurse
+      }
+    } else {
+      bool need_execute = std::any_of(frame.node_->GetNextFunctions()->begin(),
+                                      frame.node_->GetNextFunctions()->end(),
+                                      [&](const std::shared_ptr<FunctionNode>& fn) {
+                                        return need_execute_.find(fn.get()) != need_execute_.end();
+                                      });
+      if (need_execute) { need_execute_.insert(frame.node_); }
+      seen.insert(frame.node_);
+      stack.pop();
+    }
+  }
+  return Maybe<void>::Ok();
+}
+
+Maybe<void> GraphTask::Apply(bool save_grad_for_leaf) {
+  std::queue<FunctionNode*> queue;
+  for (FunctionNode* node : roots_) {
+    if (dependencies_[node] == 0) { queue.push(node); }
+  }
+
+  while (!queue.empty()) {
+    FunctionNode* node = queue.front();
+    queue.pop();
+    if (!need_execute_.empty() && need_execute_.find(node) == need_execute_.end()) {
+      node->ReleaseOutTensorArgs();
+      continue;
+    }
+    if (/*bool not_ready_to_apply=*/!(JUST(node->Apply(create_graph_)))) { continue; }
+    if (save_grad_for_leaf) { JUST(node->AccGrad4LeafTensor(create_graph_)); }
+    JUST(node->AccGrad4RetainGradTensor());
+    node->ReleaseOutTensorArgs();
+    if (!retain_graph_) { node->ReleaseData(); }
+
+    for (const auto& next_grad_fn : *(node->GetNextFunctions())) {
+      FunctionNode* next_node = next_grad_fn.get();
+      dependencies_[next_node] -= 1;
+      if (dependencies_[next_node] == 0) { queue.push(next_node); }
+    }
+  }
+  return Maybe<void>::Ok();
+}
+
+Maybe<void> GraphAutogradEngine::RunBackwardAndSaveGrads4LeafTensor(const TensorTuple& outputs,
+                                                                    const TensorTuple& out_grads,
+                                                                    bool retain_graph,
+                                                                    bool create_graph) {
+  for (int i = 0; i < outputs.size(); ++i) {
+    JUST(outputs.at(i)->now_grad_arg()->PushPartialTensor(out_grads.at(i)));
+  }
+  GraphTask graph_task(outputs, retain_graph, create_graph);
+  JUST(graph_task.ComputeDependencies());
+  JUST(graph_task.Apply(/*save_grad_for_leaf=*/true));
+  return Maybe<void>::Ok();
+}
+
+Maybe<TensorTuple> GraphAutogradEngine::RunBackwardAndReturnInputsTensorGrad(
+    const TensorTuple& outputs, const TensorTuple& inputs, const TensorTuple& out_grads,
+    bool retain_graph, bool create_graph) {
+  std::shared_ptr<TensorTuple> input_now_grads = std::make_shared<TensorTuple>(inputs.size());
+  GraphTask graph_task(outputs, retain_graph, create_graph);
+  std::vector<bool> ori_retain_grad(inputs.size());
+  for (int i = 0; i < inputs.size(); ++i) {
+    ori_retain_grad.at(i) = inputs.at(i)->retain_grad();
+    inputs.at(i)->set_retain_grad(true);
+  }
+  for (int i = 0; i < outputs.size(); ++i) {
+    JUST(outputs.at(i)->now_grad_arg()->PushPartialTensor(out_grads.at(i)));
+  }
+
+  JUST(graph_task.ComputeDependenciesAndPruneNode(inputs));
+  JUST(graph_task.Apply(/*save_grad_for_leaf=*/false));
+
+  // Gets input grads and resume retain_grad
+  for (int i = 0; i < inputs.size(); ++i) {
+    input_now_grads->at(i) = inputs.at(i)->acc_grad();
+    if (!ori_retain_grad.at(i)) {
+      inputs.at(i)->set_acc_grad(nullptr);
+      inputs.at(i)->set_retain_grad(false);
+    }
+  }
+  return input_now_grads;
+}
+
+std::shared_ptr<FunctionNode> GraphAutogradEngine::AddBackwardFuncPtr(
+    const std::string& op_type_name,
+    const std::shared_ptr<const std::function<Maybe<void>(const TensorTuple&, TensorTuple*, bool)>>&
+        backward_fn,
+    const TensorTuple& inputs, TensorTuple* outputs) {
+  // Firstly push function_node of tensor in stack which is leaf and requires_grad
+  for (const std::shared_ptr<Tensor>& in_tensor : inputs) {
+    if (in_tensor->is_leaf() && in_tensor->requires_grad()) {
+      if (!in_tensor->grad_fn_node()) { AddAccumulateFunctionNode(in_tensor); }
+    }
+  }
+
+  std::shared_ptr<FunctionNode> func_node =
+      std::make_shared<GraphFunctionNode>(op_type_name, backward_fn, inputs, *outputs);
+  for (const std::shared_ptr<Tensor>& out_tensor : *outputs) {
+    out_tensor->set_grad_fn_node(func_node);
+  }
+  return func_node;
+}
+
 AutogradEngine* GetThreadLocalAutogradEngine() {
-  thread_local static StackAutogradEngine autograd_engine;
+  // thread_local static StackAutogradEngine autograd_engine;
+  thread_local static GraphAutogradEngine autograd_engine;
   return &autograd_engine;
 }