Function AutoMemcpy

oneflow/core/kernel/kernel_util.cpp:27–49

namespace oneflow {

void AutoMemcpy(ep::Stream* stream, void* dst, const void* src, size_t sz,
                const MemoryCase& dst_mem_case, const MemoryCase& src_mem_case) {
  ep::primitive::MemcpyKind kind{};
  if (stream->device_type() == DeviceType::kCPU) {
    CHECK(memory::IsHostMem(src_mem_case));
    if (dst_mem_case.device_type() != DeviceType::kMeta) { CHECK(memory::IsHostMem(dst_mem_case)); }
    kind = ep::primitive::MemcpyKind::kDtoD;
  } else {
    if (memory::IsHostMem(src_mem_case)) {
      CHECK(!memory::IsHostMem(dst_mem_case));
      kind = ep::primitive::MemcpyKind::kHtoD;
    } else if (memory::IsHostMem(dst_mem_case)) {
      CHECK(!memory::IsHostMem(src_mem_case));
      kind = ep::primitive::MemcpyKind::kDtoH;
    } else {
      kind = ep::primitive::MemcpyKind::kDtoD;
    }
  }
  std::unique_ptr<ep::primitive::Memcpy> primitive =
      ep::primitive::NewPrimitive<ep::primitive::MemcpyFactory>(stream->device_type(), kind);
  CHECK(primitive);
  primitive->Launch(stream, dst, src, sz);
}

void AutoMemcpy(ep::Stream* stream, Blob* dst, const Blob* src) {
  const size_t body_bytes = src->ByteSizeOfBlobBody();