Function CopyDeviceToDevice

tensorflow/core/common_runtime/copy_tensor.cc:121–197

void CopyDeviceToDevice(CopyTensor::CopyFunction copy_function,
                        Allocator* cpu_allocator, Allocator* out_allocator,
                        DeviceContext* send_dev_context,
                        DeviceContext* recv_dev_context, Device* src,
                        Device* dst, const AllocatorAttributes src_alloc_attr,
                        const AllocatorAttributes dst_alloc_attr,
                        const Tensor* input, Tensor* output,
                        int dev_to_dev_stream_index, StatusCallback done) {
  if (input->dtype() == DT_VARIANT) {
    Tensor copy(cpu_allocator, DT_VARIANT, input->shape());
    auto* status_cb = new ReffedStatusCallback(std::move(done));
    core::ScopedUnref status_cb_unref(status_cb);

    auto wrapped_done = [status_cb](const Status& s) {
      status_cb->UpdateStatus(s);
      status_cb->Unref();
    };
    auto copier = std::bind(
        [copy_function, cpu_allocator, src, dst, src_alloc_attr, dst_alloc_attr,
         recv_dev_context, send_dev_context, out_allocator, status_cb,
         dev_to_dev_stream_index](StatusCallback wrapped_done_,
                                  // Begin unbound arguments
                                  const Tensor& from, Tensor* to) {
          if (from.dtype() == DT_VARIANT) {
            status_cb->Ref();
            CopyDeviceToDevice(copy_function, cpu_allocator, out_allocator,
                               send_dev_context, recv_dev_context, src, dst,
                               src_alloc_attr, dst_alloc_attr, &from, to,
                               dev_to_dev_stream_index, wrapped_done_);
            return Status::OK();
          } else {
            if (!DMAHelper::CanUseDMA(&from)) {
              Status err = errors::InvalidArgument(
                  "During Variant Device->Device Copy: ", src->name(), " to ",
                  dst->name(), " non-DMA-copy attempted of tensor type: ",
                  DataTypeString(from.dtype()));
              status_cb->UpdateStatus(err);
              return err;
            }
            if (status_cb->ok()) {
              status_cb->Ref();
              *to = Tensor(out_allocator, from.dtype(), from.shape());
              copy_function(send_dev_context, recv_dev_context, src, dst,
                            src_alloc_attr, dst_alloc_attr, &from, to,
                            dev_to_dev_stream_index, std::move(wrapped_done_));
              return Status::OK();
            } else {
              return status_cb->status();
            }
          }
        },
        std::move(wrapped_done), std::placeholders::_1, std::placeholders::_2);

    const Variant* v = input->flat<Variant>().data();
    Variant* v_out = copy.flat<Variant>().data();
    Status s_copy_init;
    for (int64 i = 0; i < input->NumElements(); ++i) {
      s_copy_init =