Function ConcatSplitPartDataContent

oneflow/core/kernel/boxing_kernel.cpp:128–168

static char* GetMutDptr(Blob* blob) { return blob->mut_dptr<char>(); }

void ConcatSplitPartDataContent(ep::Stream* stream, const DataContentDesc& in_desc,
                                const DataContentDesc& out_desc, int32_t part_id,
                                int32_t part_num) {
  size_t one_elem_size = in_desc.OneElemSize();
  BalancedSplitter bs(in_desc.TotalElemNum(), part_num);
  Range range = bs.At(part_id);
  int64_t in_idx = range.begin();
  int64_t in_elem_num = 0;
  const char* in_ptr = nullptr;
  int64_t out_idx = range.begin();
  int64_t out_elem_num = 0;
  char* out_ptr = nullptr;

  while (in_elem_num > 0 || out_elem_num > 0 || in_idx < range.end() || out_idx < range.end()) {
    if (in_elem_num == 0) {
      std::tie(in_elem_num, in_ptr) =
          in_desc.CalcContinuousElemNumStartFrom<const char, GetConstDptr>(in_idx);
      in_elem_num = std::min(in_elem_num, range.end() - in_idx);
      if (in_elem_num == 0) { break; }
      in_idx += in_elem_num;
    }
    if (out_elem_num == 0) {
      std::tie(out_elem_num, out_ptr) =
          out_desc.CalcContinuousElemNumStartFrom<char, GetMutDptr>(out_idx);
      out_elem_num = std::min(out_elem_num, range.end() - out_idx);
      if (out_elem_num == 0) { break; }
      out_idx += out_elem_num;
    }
    int64_t copy_elem_num = std::min(in_elem_num, out_elem_num);
    size_t copy_size = copy_elem_num * one_elem_size;
    Memcpy<DeviceType::kCPU>(stream, out_ptr, in_ptr, copy_size);
    in_elem_num -= copy_elem_num;
    out_elem_num -= copy_elem_num;
    in_ptr += copy_size;
    out_ptr += copy_size;
  }
  CHECK_EQ(in_elem_num, 0);
  CHECK_EQ(out_elem_num, 0);
  CHECK_EQ(in_idx, range.end());
  CHECK_EQ(out_idx, range.end());
}

void ConcatSplitDataContent(ep::Stream* stream,
                            const std::function<Blob*(const std::string&)>& BnInOp2Blob,