OneFlow: 计算数据的来源

前言

上一篇，我们分析了启动 Runtime，其中我们着重分析了线程的启动，线程是计算的抽象。这篇我们来关注存储，Runtime 在启动的时候，会给 RegstMgr 添加一个 Plan，RegstMgr 根据 Plan 申请分配内存。

回顾

上篇的最后，我们看到了最终一个 Kernel 是如何被调用执行的，数据的来源和目的地都存储在 KernelComputeContext 上。Compute 是最终消费数据的地方，这是数据的目的地，那么数据从哪里来呢？

// oneflow/user/kernels/add_n_kernel.cpp: 41
void Compute(user_op::KernelComputeContext* ctx) const override {
  size_t in_num = ctx->inputs().size();

  user_op::Tensor* out = ctx->Tensor4ArgNameAndIndex("out", 0);
  int64_t n = out->shape().elem_cnt();
  T* out_dptr = out->mut_dptr<T>();

  std::vector<const T*> in_dptrs(in_num);
  for (int32_t i = 0; i < in_num; ++i) {
    in_dptrs.at(i) = ctx->Tensor4ArgNameAndIndex("in", i)->dptr<T>();
  }

  cpu_add<T>(n, out_dptr, in_dptrs);
}

add_n 脚本

首先写一个调用了 add_n 脚本，这个脚本适用于 0.5.0 版本的 OneFlow。它首先将两个矩阵相乘，然后将所有的三个矩阵相加。

from oneflow.compatible import single_client as flow
from oneflow.compatible.single_client import linalg as linalg
from oneflow.compatible.single_client.ops import math_ops as math_ops
from oneflow.compatible.single_client import typing as tp
import numpy as np


@flow.global_function()
def matmul(
        x: tp.Numpy.Placeholder((3, 3), dtype=flow.float),
        y: tp.Numpy.Placeholder((3, 3), dtype=flow.float),
) -> tp.Numpy:
    return linalg.matmul(x, y)


@flow.global_function()
def add_n(
        x: tp.Numpy.Placeholder((3, 3), dtype=flow.float),
        y: tp.Numpy.Placeholder((3, 3), dtype=flow.float),
        z: tp.Numpy.Placeholder((3, 3), dtype=flow.float),
) -> tp.Numpy:
    return math_ops.add_n([x, y, z])


if __name__ == '__main__':
    x = np.arange(0, 9).reshape(3, 3).astype(np.float32)
    y = np.arange(9, 18).reshape(3, 3).astype(np.float32)
    z = matmul(x, y)
    print(z)

    a = add_n(x, y, z)
    print(a)

流程分析

我们尝试一种自底向上的视角去分析。

• Compute 函数中最终消费了数据，它调用了 KernelComputeContext 中 Tensor4ArgNameAndIndex 来获取输入和输出。

// oneflow/user/kernels/add_n_kernel.cpp: 42
void Compute(user_op::KernelComputeContext* ctx) const override {
  std::cout << std::this_thread::get_id() << std::endl;
  size_t in_num = ctx->inputs().size();

  user_op::Tensor* out = ctx->Tensor4ArgNameAndIndex("out", 0);
  int64_t n = out->shape().elem_cnt();
  T* out_dptr = out->mut_dptr<T>();

  std::vector<const T*> in_dptrs(in_num);
  for (int32_t i = 0; i < in_num; ++i) {
    in_dptrs.at(i) = ctx->Tensor4ArgNameAndIndex("in", i)->dptr<T>();
  }

  cpu_add<T>(n, out_dptr, in_dptrs);
}

• UserKernelComputeContext 中会初始化那个元组，不过并没有初始化 BlobTensorView，那这块对应的内存应该在哪里、什么时候进行了初始化？

// oneflow/core/kernel/user_kernel.cpp: 457
struct BnTensorPair {
  std::string bn;
  std::unique_ptr<user_op::BlobTensorView> tensor;
};

BnTensorPair MakeBnTensorPair(const std::string& bn) {
  BnTensorPair pair;
  pair.bn = bn;
  return pair;
}

// oneflow/core/kernel/user_kernel.cpp: 478
explicit UserKernelComputeContext(DeviceCtx* device_ctx, const KernelConf& kernel_conf,
                                  const JobDesc& job_desc)
    : user_op_conf_(kernel_conf.op_attribute().op_conf()),
      device_ctx_(device_ctx),
      base_ctx_(kernel_conf, job_desc) {
  auto InitInOrOut = [&](const PbMap<std::string, UserOpConf::ListString>& arg_map) {
    for (const auto& it : arg_map) {
      const std::string& arg_name = it.first;
      for (int32_t i = 0; i < it.second.s_size(); ++i) {
        arg2bn_tensor_pair_.emplace(std::make_pair(arg_name, i),
                                    MakeBnTensorPair(GenRepeatedBn(arg_name, i)));
      }
    }
  };
  InitInOrOut(kernel_conf.op_attribute().op_conf().user_conf().input());
  InitInOrOut(kernel_conf.op_attribute().op_conf().user_conf().output());
  arg2bn_tensor_pair_.emplace(std::make_pair("tmp_buffer", 0),
                              MakeBnTensorPair(GenRepeatedBn("tmp_buffer", 0)));
}

• UserKernelComputeContext 提供了一个更新的方法，应该是用这个方法进行了初始化、更新等操作。这里面的逻辑是这样的，传入一个函数，用于将 blob name 映射到 blob。对当前 Context 中存储的每一个 pair，从 pair 里拿到 blob name，调用函数去获取 blob。如果 blob 是空指针，那么将当前的 tensor 更新为空指针。如果不是空指针，那么将 blob 更新到 tensor。tensor 如果已经存在，那么更新数据，tensor 如果不存在，那么初始化。

// oneflow/core/kernel/user_kernel.cpp: 513
bool UpdateTensorWithCorrBlob(const std::function<Blob*(const std::string&)>& BnInOp2Blob) {
  bool updated = false;
  for (auto& pair : arg2bn_tensor_pair_) {
    std::unique_ptr<user_op::BlobTensorView>* arg_tensor_ptr = &pair.second.tensor;
    Blob* blob = BnInOp2Blob(pair.second.bn);
    if (blob == nullptr) {
      if (*arg_tensor_ptr) {
        arg_tensor_ptr->reset(nullptr);
        updated = true;
      }
    } else {
      if (*arg_tensor_ptr) {
        if (arg_tensor_ptr->get()->blob() != blob) {
          arg_tensor_ptr->get()->Reset(blob);
          updated = true;
        } else {
          if (blob->blob_desc().is_dynamic()) { updated = true; }
        }
      } else {
        arg_tensor_ptr->reset(new user_op::BlobTensorView(blob));
        updated = true;
      }
    }
  }
  return updated;
}

• UpdateTensorWithCorrBlob 的调用时机呢？它是什么时候被调用的？在 Lazy 模式下面，会在 ForwardUserKernel 的时候调用。

// oneflow/core/kernel/user_kernel.cpp: 635
void UserKernel::ForwardUserKernel(const std::function<Blob*(const std::string&)>& BnInOp2Blob,
                                   user_op::OpKernelState* opkernel_state) const {
  const bool updated = ctx_->UpdateTensorWithCorrBlob(BnInOp2Blob);

#ifdef WITH_CUDA_GRAPHS
  bool capturing = false;
  if (cuda_graph_ctx_) {
    if (!cuda_graph_ctx_->IsCapturing()) {
      if (cuda_graph_ctx_->IsCaptured() && (!updated)) {
        cuda_graph_ctx_->Launch();
        return;
      }
      capturing = true;
      cuda_graph_ctx_->BeginCapture();
    }
  }
#endif  // WITH_CUDA_GRAPHS

  kernel_->Compute(ctx_.get(), opkernel_state);

#ifdef WITH_CUDA_GRAPHS
  if (cuda_graph_ctx_ && capturing) {
    cuda_graph_ctx_->EndCapture();
    cuda_graph_ctx_->Launch();
  }
#endif  // WITH_CUDA_GRAPHS
}

• ForwardUserKernel 前面已经分析过了，是在 ForwardDataContent 的时候调用，也就是一个 Kernel Forward 的时候触发，即 Launch 的时候调用，或者说 AsyncLaunchKernel 的时候调用。AsyncLaunchKernel 传入的 lambda 函数，正是上面用来寻找 blob name 对应 blob 的函数。分析一下这个函数。
• 输入是 blob name，输出是 Blob。先根据 blob name 寻找 BlobInfo，然后找到 regst_desc_id，然后根据 id 找到 Regst，再从 Regst 中取出 Blob。

// oneflow/core/actor/actor.cpp: 470
void Actor::AsyncLaunchKernel(const KernelCtx& kernel_ctx,
                              std::function<Regst*(int64_t)> Regst4RegstDescId) {
  for (const ExecKernel& ek : exec_kernel_vec_) {
    ek.kernel->Launch(kernel_ctx, [&](const std::string& bn_in_op) -> Blob* {
      const auto blob_info_it = ek.bn_in_op2blob_info.find(bn_in_op);
      if (blob_info_it == ek.bn_in_op2blob_info.cend()) { return nullptr; }
      const BlobInfo& info = blob_info_it->second;
      if (info.regst_desc_id == -1) { return nullptr; }
      Regst* regst;
      if (info.rs != nullptr) {
        regst = info.rs->Front(info.regst_desc_id);
      } else {
        regst = Regst4RegstDescId(info.regst_desc_id);
      }
      if (regst == nullptr) { return nullptr; }
      if (info.ordinal >= 0) {
        return regst->GetBlobByOrdinal(info.ordinal);
      } else {
        return regst->GetBlobByLbi(info.lbi);
      }
    });
  }
}

• Blob 是从 Regst 中取出来的，Regst 的 Blob 又从何而来？有两种可能，Blob 可以来自用户的输入，也可以来自上一个 Actor 运行的结果。Regst 提供了 SetBlobByOrdinal 允许设置指定位置的 Blob。SetBlobByOrdinal 只被 RegstMgr 的 NewBlobsInOneRegst 方法调用，NewBlobsInOneRegst 则是在 NewRegsts 的时候调用，而 NewRegsts 是在 Actor Init 的时候调用。这说明，Blob 其实是在初始化的时候，就已经分配好了内存，只需要上一个 Actor 将数据写好，下一个 Actor 就会去消费。

下面的代码很长，写着上面的方法，被下面的方法调用。

// oneflow/core/register/register.cpp: 52
void Regst::SetBlobByOrdinal(int64_t ordinal, std::unique_ptr<Blob>&& blob) {
  CHECK(!sorted_blob_vec_.at(ordinal));
  sorted_blob_vec_.at(ordinal).swap(blob);
}

// oneflow/core/register/register_manager.cpp: 191
void RegstMgr::NewBlobsInOneRegst(const std::vector<LbiBlobDescPair>& lbis, Regst* regst,
                                  const RtRegstDesc* rt_regst_desc, char* main_mem_ptr,
                                  char* separated_header_mem_ptr) {
  size_t separated_header_mem_size = rt_regst_desc->SeparatedHeaderByteSize4OneRegst();
  char* cur_body_pointer = nullptr;
  char* cur_header_pointer = nullptr;
  if (separated_header_mem_size > 0) {
    MemoryCase host_mem_case;
    host_mem_case.mutable_host_mem();
    if (separated_header_mem_ptr == nullptr) {
      separated_header_mem_ptr =
          Global<MemoryAllocator>::Get()->Allocate(host_mem_case, separated_header_mem_size);
    }
    cur_header_pointer = separated_header_mem_ptr;
    cur_body_pointer = main_mem_ptr;
  } else {
    CHECK(separated_header_mem_ptr == nullptr);
    cur_header_pointer = main_mem_ptr;
    if (main_mem_ptr == nullptr) {
      cur_body_pointer = nullptr;
    } else {
      cur_body_pointer =
          main_mem_ptr + rt_regst_desc->GetSoleBlobDesc()->AlignedByteSizeOfBlobHeader();
    }
  }
  regst->set_main_mem_ptr(main_mem_ptr);
  regst->set_separated_header_mem_ptr(separated_header_mem_ptr);
  rt_regst_desc->ForEachBlobDescOffsetInOnRegst([&](int64_t ordinal, const LogicalBlobId& lbi,
                                                    const BlobDesc* blob_desc, int64_t body_offset,
                                                    int64_t header_offset) {
    std::unique_ptr<Blob> blob_ptr;
    if (cur_body_pointer == nullptr) {
      blob_ptr.reset(new Blob(regst->regst_desc()->mem_case(), blob_desc,
                              cur_header_pointer + header_offset, nullptr));
    } else {
      blob_ptr.reset(new Blob(regst->regst_desc()->mem_case(), blob_desc,
                              cur_header_pointer + header_offset, cur_body_pointer + body_offset));
      InitNonPODTypeBlobIfNeed(Global<MemoryAllocator>::Get(), blob_ptr.get());
    }
    regst->SetBlobByOrdinal(ordinal, std::move(blob_ptr));
    const int64_t regst_desc_id = rt_regst_desc->regst_desc_id();
    const auto& parallel_ctx = regst_desc_id2parallel_ctx_.at(regst_desc_id);
    if (parallel_ctx.has_parallel_id()) {
      const int64_t parallel_id = parallel_ctx.parallel_id();
      {
        std::lock_guard<std::mutex> lock(mutex_);
        lbi2parallel_id2blob_[lbi][parallel_id] = regst->GetBlobByOrdinal(ordinal);
      }
    }
  });
}

// oneflow/core/register/register_manager.cpp: 150
void RegstMgr::NewRegsts(const RegstDescProto& regst_desc_proto,
                         std::function<void(Regst*)> OneRegstDone) {
  const int64_t regst_desc_id = regst_desc_proto.regst_desc_id();
  const RegstDescTypeProto& regst_desc_type = regst_desc_proto.regst_desc_type();
  const RtRegstDesc* rt_regst_desc = regst_desc_id2rt_regst_desc_.at(regst_desc_id).get();
  char* main_mem_ptr = nullptr;
  char* separated_header_mem_ptr = nullptr;
  int64_t mem_block_id = regst_desc_proto.mem_block_id();
  int64_t header_block_id = regst_desc_proto.separated_header_mem_block_id();
  if (mem_block_id != -1 && mem_block_id2ptr_.find(mem_block_id) != mem_block_id2ptr_.end()) {
    main_mem_ptr = mem_block_id2ptr_.at(mem_block_id) + regst_desc_proto.mem_block_offset();
  }
  if (header_block_id != -1 && mem_block_id2ptr_.find(header_block_id) != mem_block_id2ptr_.end()) {
    separated_header_mem_ptr = mem_block_id2ptr_.at(header_block_id);
  }
  std::vector<LbiBlobDescPair> lbi_pairs;
  if (regst_desc_type.has_data_regst_desc()) {
    for (const LbiBlobDescPair& pair : regst_desc_type.data_regst_desc().lbi2blob_desc()) {
      lbi_pairs.push_back(pair);
    }
    std::sort(lbi_pairs.begin(), lbi_pairs.end(), &CompareLbiBlobDescPair);
    CHECK(!lbi_pairs.empty());
  }
  for (int64_t i = 0; i < rt_regst_desc->register_num(); ++i) {
    Regst* regst = new Regst;
    regst->set_regst_desc(rt_regst_desc);
    if (regst_desc_type.has_data_regst_desc()) {
      NewBlobsInOneRegst(lbi_pairs, regst, rt_regst_desc, main_mem_ptr, separated_header_mem_ptr);
      if (main_mem_ptr != nullptr) { main_mem_ptr += rt_regst_desc->MainByteSize4OneRegst(); }
      if (separated_header_mem_ptr != nullptr) {
        separated_header_mem_ptr += rt_regst_desc->SeparatedHeaderByteSize4OneRegst();
      }
    } else if (regst_desc_type.has_ctrl_regst_desc()) {
      // do nothing
    } else {
      UNIMPLEMENTED();
    }
    OneRegstDone(regst);
  }
}

// oneflow/core/actor/actor.cpp: 40
void Actor::Init(const JobDesc* job_desc, const TaskProto& task_proto,
                 const ThreadCtx& thread_ctx) {
  job_desc_ = job_desc;
  actor_id_ = task_proto.task_id();
  thrd_id_ = Global<IDMgr>::Get()->ThrdId4ActorId(actor_id_);
  job_id_ = task_proto.job_id();
  InitDeviceCtx(thread_ctx);
  if (task_proto.has_parallel_ctx()) {
    parallel_ctx_.reset(new ParallelContext(task_proto.parallel_ctx()));
  }
  for (const ExecNodeProto& node : task_proto.exec_sequence().exec_node()) {
    ExecKernel ek;
    ek.kernel = ConstructKernel(job_desc_, node.kernel_conf(), device_ctx_.get());
    exec_kernel_vec_.push_back(std::move(ek));
  }

  is_kernel_launch_synchronized_ =
      std::all_of(exec_kernel_vec_.cbegin(), exec_kernel_vec_.cend(),
                  [](const ExecKernel& ek) { return ek.kernel->IsKernelLaunchSynchronized(); });
  if (!is_kernel_launch_synchronized_) { CHECK_EQ(exec_kernel_vec_.size(), 1); }

  remaining_eord_cnt_ = 0;
  msg_handler_ = nullptr;
  eord_regst_desc_ids_.clear();

  for (const auto& pair : task_proto.produced_regst_desc()) {
    Global<RegstMgr>::Get()->NewRegsts(pair.second, [this](Regst* regst) {
      produced_regsts_[regst->regst_desc_id()].emplace_back(regst);
    });
    int64_t regst_desc_id = pair.second.regst_desc_id();
    CHECK(name2regst_desc_id_.insert({pair.first, {regst_desc_id}}).second);
    if (pair.second.regst_desc_type().has_ctrl_regst_desc()) {
      produced_ctrl_regst_desc_ids_.insert(regst_desc_id);
    }
  }
  for (const auto& pair : produced_regsts_) {
    for (const auto& regst : pair.second) { produced_regst2reading_cnt_[regst.get()] = 0; }
  }

  for (const auto& pair : task_proto.consumed_regst_desc_id()) {
    CHECK(name2regst_desc_id_.find(pair.first) == name2regst_desc_id_.end());
    std::vector<int64_t>& regst_desc_id_vec = name2regst_desc_id_[pair.first];
    for (int64_t regst_desc_id : pair.second.regst_desc_id()) {
      regst_desc_id_vec.push_back(regst_desc_id);
    }
    remaining_eord_cnt_ += pair.second.regst_desc_id_size();
    if (pair.first == "in_ctrl") {
      consumed_ctrl_regst_desc_ids_.insert(regst_desc_id_vec.begin(), regst_desc_id_vec.end());
    }
  }

  total_reading_cnt_ = 0;
  is_inplace_consumed_eord_ = false;
  CheckInplaceRegstDescId(task_proto);
  TakeOverInplaceConsumedAndProduced(task_proto.produced_regst_desc());
  is_naive_consumed_eord_ = false;
  TakeOverNaiveConsumed(task_proto.consumed_regst_desc_id());
  TakeOverNaiveProduced(task_proto.produced_regst_desc());
  InitBnInOp2BlobInfo(task_proto);
  VirtualActorInit(task_proto);
}

• 既然上一个 Actor 将数据输出之后，下一个 Actor 就会自动消费，那么我们其实只需要关注接受数据的 Actor。数据是通过 Push Job 推送来的，并且在 Python 需要提供一个函数，将 numpy 数据复制到 Blob 里面。

下面的过程来自 InferenceSession，通过启动 Push Job 来将数据推送。Python 通过继承一个 C++ 类，实现了 PushBlob 的方法，在 PushBlob 被调用的时候，他会将指针变成 blob，然后调用 push 的回调函数。

# python/oneflow/compatible/single_client/serving/inference_session.py： 427
def _run_push_jobs(self, **kwargs):
    for (
        input_name,
        push_job_name,
    ) in self.inter_user_job_info_.input_or_var_op_name2push_job_name.items():
        if input_name not in kwargs:
            raise ValueError('input "{}" is absent'.format(input_name))
        input_numpy = kwargs[input_name]
        if not isinstance(input_numpy, np.ndarray):
            raise ValueError('input "{}" requires numpy.ndarray'.format(input_name))
        push_fn = input_blob_util._MakePushNdarrayCallback(input_numpy)
        push_job_inst = job_instance_util.MakePushJobInstance(
            push_job_name, input_name, push_fn
        )
        self._run_job(push_job_inst)

# python/oneflow/compatible/single_client/framework/input_blob_def.py: 249
def _MakePushNdarrayCallback(ndarray):
    copied = np.copy(ndarray, order="C")

    def Copy(ofblob):
        capacity = reduce(lambda x, y: x * y, ofblob.static_shape, 1)
        elem_cnt = reduce(lambda x, y: x * y, copied.shape, 1)
        assert elem_cnt <= capacity, "%s v.s. %s" % (copied.shape, ofblob.static_shape)
        ofblob.CopyFromNdarray(copied)

    return Copy

# python/oneflow/compatible/single_client/framework/job_instance.py: 106
def PushBlob(self, of_blob_ptr):
    try:
        self.push_cb_(ofblob.OfBlob(of_blob_ptr))
    except Exception as e:
        print(traceback.format_exc())
        raise e

• push job instance 什么时候被启动呢？首先从 Python 调用 C++ 的 API，将 job instance 推送过来，接着将这个 job instance 放入 Buffer Manager，等待被取出来。Launch Job 的最后，调用 Buffer Manager 获取 kBufferNameGlobalWaitJobId，向它发送 job id 启动 foreign input job。之后 ForeignInputKernel 被启动，调用了 ForwardDataContent，然后构建一个 blob 传给 job instance 让它去将数据 push 到 blob 里。

// oneflow/api/python/framework/framework.h: 67
inline Maybe<void> LaunchJob(const std::shared_ptr<oneflow::JobInstance>& cb) {
  CHECK_OR_RETURN(GlobalProcessCtx::IsThisProcessMaster());
  CHECK_NOTNULL_OR_RETURN(Global<Oneflow>::Get());
  const auto& job_name = cb->job_name();
  auto* buffer_mgr = Global<BufferMgr<std::shared_ptr<JobInstance>>>::Get();
  int64_t job_id = Global<JobName2JobId>::Get()->at(job_name);
  if (IsPullJob(job_name, *Global<InterUserJobInfo>::Get())) {
    buffer_mgr->Get(GetForeignOutputBufferName(job_name))->Send(cb);
  }
  if (IsPushJob(job_name, *Global<InterUserJobInfo>::Get())) {
    buffer_mgr->Get(GetForeignInputBufferName(job_name))->Send(cb);
  }
  buffer_mgr->Get(GetCallbackNotifierBufferName(job_name))->Send(cb);
  Global<BufferMgr<int64_t>>::Get()->Get(kBufferNameGlobalWaitJobId)->Send(job_id);
  return Maybe<void>::Ok();
}

// oneflow/core/kernel/foreign_input_kernel.cpp: 23
void ForeignInputKernel::ForwardDataContent(
    const KernelCtx& ctx, const std::function<Blob*(const std::string&)>& BnInOp2Blob) const {
  const auto& buffer_name = op_conf().foreign_input_conf().ofblob_buffer_name();
  std::shared_ptr<JobInstance> foreign_job_instance;
  BufferStatus buffer_status = Global<BufferMgr<std::shared_ptr<JobInstance>>>::Get()
                                   ->Get(buffer_name)
                                   ->TryReceive(&foreign_job_instance);
  CHECK_NE(buffer_status, kBufferStatusEmpty);
  if (buffer_status == kBufferStatusSuccess) {
    OfBlob ofblob(ctx.device_ctx, BnInOp2Blob("out"));
    foreign_job_instance->PushBlob(reinterpret_cast<uint64_t>(&ofblob));
  }
}

总结

这篇文章粗略总结了计算需要的数据从哪里来，将到哪儿去。内存是在运行时启动的时候按照 Plan 分配的，当上一个 Actor 输出之后，下一个 Actor 就可以利用上一个 Actor 的输出来进行下一步的计算。当 Python 需要将前端的数据推送的 C++ 这边的时候，使用 JobInstance 来实现推送数据。

不过这里有一个细节，我说的很模糊。在 Launch Job 的时候，向 BufferMgr 发送一个 id，这个 id 是如何被接收的呢？接收了这个 id 之后，如何启动对应的 Job 呢？下一篇将进行分析。