Method CopyGPUTensorToCPU

tensorflow/core/common_runtime/gpu/gpu_util.cc:280–330

static

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

  se::Stream* send_device_to_host_stream = nullptr;
  if (MergeComputeAndCopyStream()) {
    send_device_to_host_stream = send_stream;
  } else {
    send_device_to_host_stream =
        static_cast<const GPUDeviceContext*>(device_context)
            ->device_to_host_stream();
    if (send_device_to_host_stream == nullptr) {
      done(errors::Internal("No send gpu copy-out-stream is available."));
      return;
    }

    // Wait for the sender's main stream to make sure the data are available.
    if (send_device_to_host_stream != send_stream) {
      send_device_to_host_stream->ThenWaitFor(send_stream);
    }
  }

  const int64 total_bytes = gpu_tensor->TotalBytes();
  if (total_bytes > 0) {
    void* src_ptr = GetBase(gpu_tensor);
    DeviceMemoryBase gpu_src_ptr(src_ptr, total_bytes);
    void* dst_ptr = GetBase(cpu_tensor);
    send_device_to_host_stream->ThenMemcpy(dst_ptr, gpu_src_ptr, total_bytes);
  }
  // Use of the input may outlive stack scope, so keep a ref.
  TensorReference input_ref(*gpu_tensor);
  dev_info->event_mgr->ThenExecute(
      send_device_to_host_stream,
      [send_device_to_host_stream, done, input_ref]() {
        if (!send_device_to_host_stream->ok()) {
          LOG(FATAL) << "GPU->CPU Memcpy failed";
        }
        input_ref.Unref();
        done(Status::OK());
      });
}

/*  static */
void GPUUtil::CopyCPUTensorToGPU(const Tensor* cpu_tensor,