Function InitDataSourceDistributedInfo

oneflow/user/data/distributed_util.h:28–55

namespace data {

inline Maybe<void> InitDataSourceDistributedInfo(user_op::KernelInitContext* ctx,
                                                 size_t& world_size, int64_t& rank) {
  auto nd_sbp_str_vec = ctx->Attr<std::vector<std::string>>("nd_sbp");
  if (nd_sbp_str_vec.empty()) {
    world_size = GlobalProcessCtx::WorldSize();
    rank = GlobalProcessCtx::Rank();
  } else {
    const Shape& hierarchy = *ctx->parallel_desc().hierarchy();
    CHECK_EQ_OR_RETURN(hierarchy.NumAxes(), nd_sbp_str_vec.size());
    rank = 0;
    world_size = 1;

    using index_helper_t = NdIndexOffsetHelper<int64_t, SHAPE_MAX_AXIS_SIZE>;
    index_helper_t index_helper(hierarchy.dim_vec().data(), hierarchy.NumAxes());
    int64_t nd_index[SHAPE_MAX_AXIS_SIZE] = {0};
    index_helper.OffsetToNdIndex(ctx->parallel_ctx().parallel_id(), nd_index);

    for (int i = hierarchy.NumAxes() - 1; i >= 0; --i) {
      SbpParallel sbp;
      CHECK_OR_RETURN(ParseSbpParallelFromString(nd_sbp_str_vec[i], &sbp));
      if (sbp.has_split_parallel()) {
        rank += nd_index[i] * world_size;
        world_size *= hierarchy.At(i);
      }
    }
  }
  return Maybe<void>::Ok();
}

}  // namespace data