| 1039 | } |
| 1040 | |
| 1041 | void RdmaTensorResponse::RecvHandler(Rendezvous::ParsedKey parsed, |
| 1042 | const Rendezvous::Args& send_args, |
| 1043 | const Rendezvous::Args& recv_args, |
| 1044 | const Tensor& in, bool is_dead) { |
| 1045 | Status s = PrepareRecvTensor(parsed, &src_dev_); |
| 1046 | if (!s.ok()) { |
| 1047 | SendErrorStatus(s); |
| 1048 | return; |
| 1049 | } |
| 1050 | |
| 1051 | meta_data_changed_ = TensorMetaDataChanged(in, is_dead); |
| 1052 | #ifdef RDMA_DATA_VALIDATION |
| 1053 | // Always send a meta data message with the source checksum |
| 1054 | meta_data_changed_ = rm_.type_ == RDMA_MESSAGE_TENSOR_REQUEST; |
| 1055 | checksum_ = Checksum(src_dev_, send_args.device_context, in); |
| 1056 | #endif |
| 1057 | bool can_memcpy = DataTypeCanUseMemcpy(in.dtype()); |
| 1058 | // string tensor needs to be serialized |
| 1059 | Tensor copy; |
| 1060 | TensorProto proto; |
| 1061 | const bool on_host = send_args.alloc_attrs.on_host(); |
| 1062 | if (src_dev_->tensorflow_gpu_device_info() && !on_host) { |
| 1063 | #if GOOGLE_CUDA |
| 1064 | DeviceContext* send_dev_context = send_args.device_context; |
| 1065 | CHECK(send_dev_context) |
| 1066 | << "send dev name: " << src_dev_->name() |
| 1067 | << " gpu_info: " << src_dev_->tensorflow_gpu_device_info(); |
| 1068 | |
| 1069 | if (can_memcpy) { |
| 1070 | // If the tensor is located on a GDR compatible GPU, there is no need to |
| 1071 | // copy it. We can send directly from the source, just need to make sure |
| 1072 | // we are in sync with the GPU stream. |
| 1073 | // If the tensor's meta-data changed however, we will need to clone it, |
| 1074 | // so anyway we'll have to copy it from GPU to CPU first. If at some |
| 1075 | // point in time Clone() is changed to only save a shallow copy, we can |
| 1076 | // skip the copy here as well. |
| 1077 | if ((in.TotalBytes() > 0) && !meta_data_changed_ && |
| 1078 | (RdmaMemoryMgr::Singleton().FindMemoryRegion( |
| 1079 | (void*)DMAHelper::base(&in), in.TotalBytes()) != nullptr)) { |
| 1080 | StreamGPUOp(src_dev_, send_dev_context, |
| 1081 | [this, in, proto, is_dead](const Status& s) { |
| 1082 | Send(in, proto, is_dead, s); |
| 1083 | }); |
| 1084 | return; |
| 1085 | } |
| 1086 | |
| 1087 | // The tensor must be copied from GPU to CPU, because either: |
| 1088 | // 1. The tensor is located on a non GDR compatible GPU. |
| 1089 | // 2. The tensor's meta-data has changed. |
| 1090 | Allocator* alloc = GPUProcessState::singleton()->GetGpuHostAllocator(0); |
| 1091 | copy = Tensor(alloc, in.dtype(), in.shape()); |
| 1092 | CountCopies(rm_.name_, (void*)DMAHelper::base(&in), |
| 1093 | (void*)DMAHelper::base(©), in.TotalBytes(), true); |
| 1094 | GPUUtil::CopyGPUTensorToCPU( |
| 1095 | src_dev_, send_dev_context, &in, ©, |
| 1096 | [this, copy, proto, is_dead](const Status& s) { |
| 1097 | Send(copy, proto, is_dead, s); |
| 1098 | }); |
nothing calls this directly
no test coverage detected