static
| 200 | |
| 201 | // static |
| 202 | void CopyTensor::ViaDMA(StringPiece edge_name, DeviceContext* send_dev_context, |
| 203 | DeviceContext* recv_dev_context, Device* src, |
| 204 | Device* dst, const AllocatorAttributes src_alloc_attr, |
| 205 | const AllocatorAttributes dst_alloc_attr, |
| 206 | const Tensor* input, Tensor* output, |
| 207 | int dev_to_dev_stream_index, StatusCallback done, |
| 208 | bool sync_dst_compute) { |
| 209 | tracing::ScopedAnnotation annotation(edge_name); |
| 210 | VLOG(1) << "Copy " << edge_name; |
| 211 | |
| 212 | const DeviceType src_device_type( |
| 213 | src_alloc_attr.on_host() ? DEVICE_CPU : src->attributes().device_type()); |
| 214 | const DeviceType dst_device_type( |
| 215 | dst_alloc_attr.on_host() ? DEVICE_CPU : dst->attributes().device_type()); |
| 216 | const bool non_cpu_src = src_device_type != DeviceType(DEVICE_CPU); |
| 217 | const bool non_cpu_dst = dst_device_type != DeviceType(DEVICE_CPU); |
| 218 | |
| 219 | // TODO(phawkins): choose an allocator optimal for both the src and dst |
| 220 | // devices, not just the src device. |
| 221 | AllocatorAttributes host_alloc_attrs; |
| 222 | host_alloc_attrs.set_gpu_compatible(true); |
| 223 | host_alloc_attrs.set_on_host(true); |
| 224 | Allocator* cpu_allocator = src->GetAllocator(host_alloc_attrs); |
| 225 | Allocator* out_allocator = dst->GetAllocator(dst_alloc_attr); |
| 226 | |
| 227 | // E.g., gpu -> gpu |
| 228 | if (non_cpu_src && non_cpu_dst) { |
| 229 | // Device to device copy. Look through registry for an appropriate |
| 230 | // CopyFunction. |
| 231 | std::vector<RegistrationInfo>* registry = MutableRegistry(); |
| 232 | for (const RegistrationInfo& ri : *registry) { |
| 233 | if (ri.sender_device_type == src_device_type && |
| 234 | ri.receiver_device_type == dst_device_type) { |
| 235 | CopyDeviceToDevice(ri.copy_function, cpu_allocator, out_allocator, |
| 236 | send_dev_context, recv_dev_context, src, dst, |
| 237 | src_alloc_attr, dst_alloc_attr, input, output, |
| 238 | dev_to_dev_stream_index, std::move(done)); |
| 239 | return; |
| 240 | } |
| 241 | } |
| 242 | |
| 243 | // Fall back to copying via the host. |
| 244 | VLOG(1) << "No function registered to copy from devices of type " |
| 245 | << src_device_type.type() << " to devices of type " |
| 246 | << dst_device_type.type() |
| 247 | << ". Falling back to copying via the host."; |
| 248 | |
| 249 | Tensor* cpu_tensor = |
| 250 | new Tensor(cpu_allocator, input->dtype(), input->shape()); |
| 251 | std::function<void(const Status&)> delete_and_done = std::bind( |
| 252 | [cpu_tensor](StatusCallback done_, |
| 253 | // Begin unbound arguments. |
| 254 | const Status& status) { |
| 255 | delete cpu_tensor; |
| 256 | done_(status); |
| 257 | }, |
| 258 | std::move(done), std::placeholders::_1); |
| 259 | std::function<void(const Status&)> then_copy_to_other_device = std::bind( |
nothing calls this directly
no test coverage detected