pytorch学习笔记（十三）：backward过程的底层实现解析

// Add 函数 计算方法的实现， 构建反向传导图的关键在 wrap_outputs
auto Add::apply(const variable_list& inputs) -> variable_list {
  check_input_variables("Add", inputs, 2);
  auto& input1 = inputs[0]->data;
  auto& input2 = inputs[1]->data;
  AutoGPU guard(input1->getDevice());

  bool first_sparse = input1->isSparse();
  auto output = first_sparse ? input2->newTensor() : input1->newTensor();
  if (first_sparse) {
    output->cadd(*input2, *input1);
  } else {
    output->cadd(*input1, *input2);
  }

  return wrap_outputs(inputs, as_tensor_list(std::move(output)), [&](FunctionFlags f) {
    return std::make_shared<AddBackward>(std::move(f));
  });
};

variable_list wrap_outputs(const variable_list& inputs, tensor_list&& outputs,
                           function_constructor ctr) {

  // 使用 inputs variables 来计算 反向传导的 Function 的 flag (f.is_executable, f.is_volatile)和 next_functions
  // 这里需要搞清楚的一点是：inputs 是前向传导的 inputs，从它可获得的信息有：当前 函数的 反向传导函数 是否 可执行，

  auto flags = Function::flags(inputs);
  variable_list result;

  // 开始创建 返回的 Variable 了。
  result.reserve(outputs.size());

  if (flags.is_volatile) {  // 如果 is_volatile=true, 那么输出的 Variable 的 is_volatile=true 
    for (auto& output : outputs) {
      if (output.defined()) { // 因为 可能返回 None 嘛，所以这里 check 一下
        result.emplace_back(make_variable(output, false, true)); // requires_grad=false, is_volatile=true
      } else {
        result.emplace_back(); 
      }
    }
  } else {  // 如果 volatile=false， 难道也不管 is_executable 了吗？ 
    // ctr 是一个 lambda 函数， 它返回一个 std::shared_ptr<GradFn>
    // 梯度 使用 Function::flags 计算出来的 flags 其实是给 Backward 用的。
    auto grad_fn = ctr(std::move(flags));  // 用 flags(is_executable, is_volatile) 创建出来一个 Function。
    for (auto& output : outputs) {
      if (output.defined()) {
        result.emplace_back(make_variable(output, grad_fn));
      } else {
        // forward 的输出 变量个数，就是 backward 的输入变量个数。
        ++grad_fn->num_inputs;
        result.emplace_back();
      }
    }
  }
  return result;
}

a = o+e
c = a+b
d = a+e
res = c+d

struct FunctionTask {
  // 每个 FunctionTask 中都维护着一个 base GraphTask
  GraphTask* base;
  std::shared_ptr<Function> fn; // 代表 grad_fn
  InputBuffer inputs; // 累积 grad_fn 的输入

  FunctionTask(GraphTask* base, std::shared_ptr<Function> fn, InputBuffer inputs)
    : base(base)
    , fn(fn)
    , inputs(std::move(inputs)) {}
};

struct ReadyQueue {
  // 用来 存放可被 执行的 FunctionTask
  // queue 是 FunctionTask 的 一个 双端队列
  std::deque<FunctionTask> queue;
  // std::condition_variable 条件变量，同步的时候会用到
  // 用 unique_lock (over mutex) 来进行操作
  std::condition_variable not_empty;
  std::mutex mutex;

  void push_front(FunctionTask item);
  FunctionTask pop_back();
};

struct GraphTask {
  // 记录整个反向计算图的依赖情况 等等。
  std::exception_ptr exception;
  // Indicates if an error occurred while executing any task.  When this is
  // true, it signals all threads to stop executing.
  std::atomic_bool has_error;

  // 剩余 tasks。 在 ReadyQueue 的 push方法 中加一， 在 evaluate_function 中减一操作
  std::atomic<uint64_t> outstanding_tasks;
  bool keep_graph;
  bool has_any_work;
  // 用来 给 notify_all 加锁的
  std::mutex mutex;
  // Notified when a task finishes executing.  Check outstanding_tasks to see
  // if all tasks are done.
  std::condition_variable not_done;
  const Engine::pre_callback_map& pre_callbacks;
  const Engine::post_callback_map& post_callbacks;

  //用来存放 没有 准备好的 FunctionTask
  std::unordered_map<Function*, InputBuffer> not_ready;
  // 记录 所有 Function 节点的 依赖
  std::unordered_map<Function*, int> dependencies;

  // 这个来 表示 GraphTask 是在哪个 device 上创建的
  int owner;

};

struct Engine{
  // 反向传导计算引擎
}