Searched refs:transfer_local_device (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/pjrt/ |
| D | pjrt_stream_executor_client.cc | 1295 LocalDeviceState* transfer_local_device, se::Stream* transfer_stream, in CopyToDeviceHelper() argument
1322 TF_RETURN_IF_ERROR(transfer_local_device->ThenMemcpyDeviceToDevice( in CopyToDeviceHelper()
1330 transfer_local_device, std::move(dst_device_buffer), event, in CopyToDeviceHelper()
1335 StallStreamOnError(transfer_local_device, transfer_stream); in CopyToDeviceHelper()
1336 if (transfer_local_device == dst_local_device) { in CopyToDeviceHelper()
1378 LocalDeviceState* transfer_local_device = in CopyToDevice() local
1385 transfer_local_device->allocation_model()); in CopyToDevice()
1388 transfer_local_device->GetDeviceToDeviceStream(); in CopyToDevice()
1405 dst_device, dst_local_device, transfer_local_device, transfer_stream, in CopyToDevice()
1424 transfer_local_device, event, transfer_stream, in CopyToDevice()
|
| D | pjrt_stream_executor_client.h | 583 LocalDeviceState* transfer_local_device,
|