| 1616 | } |
| 1617 | |
| 1618 | void RdmaTensorRequest::RecvTensorContent() { |
| 1619 | bool can_memcpy = DataTypeCanUseMemcpy(meta_data_->data_type_); |
| 1620 | size_t message_size = |
| 1621 | can_memcpy ? result_tensor_->TotalBytes() : meta_data_->proto_size_; |
| 1622 | RDMA_LOG(1) << "Step 0x" << std::hex << step_id_ << std::dec |
| 1623 | << ": Received tensor content #" << index_ << ": " << key_ |
| 1624 | << " (Size: 0x" << std::hex << message_size << ")"; |
| 1625 | |
| 1626 | #if GOOGLE_CUDA |
| 1627 | if (proxy_tensor_ != nullptr) { |
| 1628 | CountCopies(key_, (void*)DMAHelper::base(proxy_tensor_), |
| 1629 | (void*)DMAHelper::base(result_tensor_), |
| 1630 | result_tensor_->TotalBytes(), false); |
| 1631 | GPUUtil::CopyCPUTensorToGPU( |
| 1632 | proxy_tensor_, recv_args_.device_context, dst_dev_, result_tensor_, |
| 1633 | [this](const Status& s) { |
| 1634 | CHECK(s.ok()) << "copy tensor to gpu sync"; |
| 1635 | Done(s); |
| 1636 | }, |
| 1637 | true /*sync_dst_compute*/); |
| 1638 | return; |
| 1639 | } |
| 1640 | #endif |
| 1641 | |
| 1642 | if (can_memcpy) { |
| 1643 | Done(Status::OK()); |
| 1644 | } else { |
| 1645 | RDMA_LOG(2) << "Decoding proto: " << key_ |
| 1646 | << " (Size: " << meta_data_->proto_size_ << ")"; |
| 1647 | TensorProto proto; |
| 1648 | CHECK(ParseProtoUnlimited(&proto, rdma_addr_, meta_data_->proto_size_)) |
| 1649 | << "fail to parse proto from array"; |
| 1650 | ibv_dereg_mr(mr_); |
| 1651 | free(rdma_addr_); |
| 1652 | Status s = dst_dev_->MakeTensorFromProto(proto, recv_args_.alloc_attrs, |
| 1653 | result_tensor_); |
| 1654 | Done(s); |
| 1655 | } |
| 1656 | } |
| 1657 | |
| 1658 | void RdmaTensorRequest::RecvErrorStatus(const Status& status) { |
| 1659 | if (result_tensor_ == nullptr) { |
no test coverage detected