Method CopyCPUTensorToGPU

tensorflow/core/common_runtime/gpu/gpu_util.cc:333–390

static */

/*  static */
void GPUUtil::CopyCPUTensorToGPU(const Tensor* cpu_tensor,
                                 const DeviceContext* device_context,
                                 Device* gpu_device, Tensor* gpu_tensor,
                                 StatusCallback done, bool sync_dst_compute) {
  VLOG(1) << "CopyCPUTensorToGPU";
  const DeviceBase::GpuDeviceInfo* dev_info = nullptr;
  se::Stream* recv_stream = nullptr;
  Status s = PrepareCopy(gpu_device, device_context, *cpu_tensor, gpu_tensor,
                         &dev_info, &recv_stream);
  if (!s.ok()) {
    done(s);
    return;
  }

  se::Stream* recv_host_to_device_stream = nullptr;
  if (MergeComputeAndCopyStream()) {
    recv_host_to_device_stream = recv_stream;
  } else {
    recv_host_to_device_stream =
        static_cast<const GPUDeviceContext*>(device_context)
            ->host_to_device_stream();
    if (recv_host_to_device_stream == nullptr) {
      done(errors::Internal("No send gpu copy-out-stream is available."));
      return;
    }
    // Wait for the recv-stream to make sure the buffer is truly available.
    if (sync_dst_compute) {
      if (recv_host_to_device_stream != recv_stream) {
        recv_host_to_device_stream->ThenWaitFor(recv_stream);
      }
    }
  }

  const int64 total_bytes = cpu_tensor->TotalBytes();
  // Note that 0-size tensors have no backing buffer.
  if (total_bytes > 0) {
    void* src_ptr = GetBase(cpu_tensor);
    void* dst_ptr = GetBase(gpu_tensor);
    DeviceMemoryBase gpu_dst_ptr(dst_ptr, total_bytes);
    recv_host_to_device_stream->ThenMemcpy(&gpu_dst_ptr, src_ptr, total_bytes);
  }

  if (MergeComputeAndCopyStream()) {
    done(Status::OK());
  } else {
    // Use of cpu_tensor may outlive stack scope, so keep a ref.
    TensorReference input_ref(*cpu_tensor);
    dev_info->event_mgr->ThenExecute(
        recv_host_to_device_stream,
        [recv_host_to_device_stream, done, input_ref]() {
          input_ref.Unref();
          if (!recv_host_to_device_stream->ok()) {
            LOG(FATAL) << "CPU->GPU Memcpy failed";
          }
          done(Status::OK());
        });
  }
}