static*/
| 131 | |
| 132 | /*static*/ |
| 133 | void GPUUtil::SetProtoFromGPU(const Tensor& tensor, Device* dev, |
| 134 | const DeviceContext* device_context, |
| 135 | TensorProto* proto, bool is_dead, |
| 136 | StatusCallback done) { |
| 137 | VLOG(1) << "SetProtoFromGPU device_context " << device_context; |
| 138 | const DeviceBase::GpuDeviceInfo* dev_info = nullptr; |
| 139 | se::Stream* send_stream = nullptr; |
| 140 | Status s = PrepareCopy(dev, device_context, tensor, nullptr, &dev_info, |
| 141 | &send_stream); |
| 142 | if (!s.ok()) { |
| 143 | done(s); |
| 144 | return; |
| 145 | } |
| 146 | |
| 147 | auto send_device_to_host_stream = |
| 148 | static_cast<const GPUDeviceContext*>(device_context) |
| 149 | ->device_to_host_stream(); |
| 150 | if (send_device_to_host_stream == nullptr) { |
| 151 | done(errors::Internal("No send gpu copy-out-stream is available.")); |
| 152 | return; |
| 153 | } |
| 154 | // Wait for the sender's main stream to make sure the data are available. |
| 155 | if (send_device_to_host_stream != send_stream) { |
| 156 | send_device_to_host_stream->ThenWaitFor(send_stream); |
| 157 | } |
| 158 | |
| 159 | // Tensor values need to be copied from GPU to CPU ram so that |
| 160 | // we can build the protobuf response for a RecvTensor RPC. |
| 161 | // "device context" identifies the stream where the _Send op executed. |
| 162 | proto->set_dtype(tensor.dtype()); |
| 163 | tensor.shape().AsProto(proto->mutable_tensor_shape()); |
| 164 | |
| 165 | // Prepare a proto with the right data buf size, and DMA the data |
| 166 | // over from the GPU buffer. Note that 0-size tensors do not have a |
| 167 | // backing buffer. |
| 168 | Allocator* alloc = nullptr; |
| 169 | char* buf = nullptr; |
| 170 | const int64 total_bytes = is_dead ? 0 : tensor.TotalBytes(); |
| 171 | if (total_bytes > 0) { |
| 172 | tracing::ScopedAnnotation annotation("SetProtoFromGPU"); |
| 173 | alloc = GPUProcessState::singleton()->GetGpuHostAllocator(0); |
| 174 | buf = static_cast<char*>( |
| 175 | alloc->AllocateRaw(Allocator::kAllocatorAlignment, total_bytes)); |
| 176 | if (LogMemory::IsEnabled()) { |
| 177 | LogMemory::RecordRawAllocation("SetProtoFromGPU", |
| 178 | LogMemory::PROTO_BUFFER_STEP_ID, |
| 179 | total_bytes, buf, alloc); |
| 180 | } |
| 181 | void* src_ptr = GetBase(&tensor); |
| 182 | DeviceMemoryBase gpu_src_ptr(src_ptr, total_bytes); |
| 183 | send_device_to_host_stream->ThenMemcpy(buf, gpu_src_ptr, total_bytes); |
| 184 | } |
| 185 | // Use of tensor may outlive stack scope, so keep a ref. |
| 186 | TensorReference tensor_ref(tensor); |
| 187 | dev_info->event_mgr->ThenExecute( |
| 188 | send_device_to_host_stream, [send_device_to_host_stream, done, proto, buf, |
| 189 | total_bytes, alloc, tensor_ref]() { |
| 190 | if (!send_device_to_host_stream->ok()) { |
nothing calls this directly
no test coverage detected