| 463 | } |
| 464 | |
| 465 | Status TensorHandle::CopyToDevice(EagerContext* ctx, tensorflow::Device* dstd, |
| 466 | tensorflow::Tensor* output) { |
| 467 | tensorflow::Device* srcd = DeviceOrHostCPU(ctx); |
| 468 | const bool dst_cpu = dstd->tensorflow_gpu_device_info() == nullptr; |
| 469 | const bool src_cpu = srcd->tensorflow_gpu_device_info() == nullptr; |
| 470 | bool is_same_device = |
| 471 | (srcd == dstd) || (srcd->name() == dstd->name()) || (dst_cpu && src_cpu); |
| 472 | |
| 473 | const tensorflow::Tensor* src = nullptr; |
| 474 | TF_RETURN_IF_ERROR(Tensor(&src)); |
| 475 | if (is_same_device) { |
| 476 | *output = *src; |
| 477 | return Status::OK(); |
| 478 | } |
| 479 | if (!dst_cpu && (src->dtype() != tensorflow::DT_VARIANT && |
| 480 | !tensorflow::DataTypeCanUseMemcpy(src->dtype()))) { |
| 481 | return tensorflow::errors::InvalidArgument( |
| 482 | "Can't copy Tensor with type ", |
| 483 | tensorflow::DataTypeString(src->dtype()), " to device ", dstd->name(), |
| 484 | "."); |
| 485 | } |
| 486 | tensorflow::AllocatorAttributes attr; |
| 487 | if (src->dtype() == tensorflow::DT_VARIANT) { |
| 488 | attr.set_on_host(true); |
| 489 | } |
| 490 | tensorflow::Tensor dst(dstd->GetAllocator(attr), src->dtype(), src->shape()); |
| 491 | if (src->shape().num_elements() == 0) { |
| 492 | *output = dst; |
| 493 | return Status::OK(); |
| 494 | } |
| 495 | tensorflow::DeviceContext* src_device_context = nullptr; |
| 496 | if (!src_cpu) { |
| 497 | src_device_context = srcd->tensorflow_gpu_device_info()->default_context; |
| 498 | } |
| 499 | tensorflow::DeviceContext* dst_device_context = nullptr; |
| 500 | if (!dst_cpu) { |
| 501 | dst_device_context = dstd->tensorflow_gpu_device_info()->default_context; |
| 502 | } |
| 503 | // TODO(ashankar): The Sync() call below may be more aggressive than |
| 504 | // necessary. It is based on knowledge of implementation details - that |
| 505 | // GPU devices are implemented using 3 streams - one for host->device copies, |
| 506 | // one for device->host copies and one for sending operations to the GPU. |
| 507 | // With that setup, Sync()ing across all 3 streams should be sufficient |
| 508 | // but more than necessary (since it waits for operations that might have |
| 509 | // nothing to do with this tensor to complete). |
| 510 | TF_RETURN_IF_ERROR(srcd->Sync()); |
| 511 | tensorflow::Notification n; |
| 512 | tensorflow::Status status; |
| 513 | tensorflow::CopyTensor::ViaDMA("copy", src_device_context, dst_device_context, |
| 514 | srcd, dstd, tensorflow::AllocatorAttributes(), |
| 515 | tensorflow::AllocatorAttributes(), src, &dst, |
| 516 | 0 /*dev_to_dev_stream_index*/, |
| 517 | [&status, &n](const tensorflow::Status& s) { |
| 518 | status = s; |
| 519 | n.Notify(); |
| 520 | }); |
| 521 | n.WaitForNotification(); |
| 522 | if (status.ok()) { |
no test coverage detected