Searched refs:local_device_state (Results 1 – 11 of 11) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/pjrt/ |
| D | tpu_client.h | 32 std::unique_ptr<LocalDeviceState> local_device_state, in PjRtTpuDevice() argument
35 : PjRtStreamExecutorDevice(core.Id(), std::move(local_device_state), in PjRtTpuDevice()
|
| D | interpreter_device.cc | 28 int id, std::unique_ptr<LocalDeviceState> local_device_state) in InterpreterDevice() argument
29 : PjRtStreamExecutorDevice(id, std::move(local_device_state), in InterpreterDevice()
|
| D | cpu_device.cc | 28 std::unique_ptr<LocalDeviceState> local_device_state) in CpuDevice() argument
29 : PjRtStreamExecutorDevice(id, std::move(local_device_state), in CpuDevice()
|
| D | BUILD | 80 ":local_device_state",
112 name = "local_device_state",
113 srcs = ["local_device_state.cc"],
114 hdrs = ["local_device_state.h"],
180 ":local_device_state",
238 ":local_device_state",
|
| D | pjrt_stream_executor_client.cc | 1062 LocalDeviceState* local_device_state = in Release() local
1064 ->local_device_state(); in Release()
1073 stream = local_device_state->BorrowStreamFromPool(); in Release()
1080 local_device_state->ReturnStreamToPool(std::move(stream)); in Release()
1083 if (local_device_state->allocation_model() == in Release()
1095 local_device_state->compute_stream()) && in Release()
1098 block_stream = local_device_state->BorrowStreamFromPool(); in Release()
1105 local_device_state->ThenExecuteOnCallbackThread( in Release()
1107 [device_buffer, block_stream_ptr, local_device_state]() { in Release()
1108 local_device_state->ReturnStreamToPool( in Release()
[all …]
|
| D | tpu_client.cc | 161 std::unique_ptr<LocalDeviceState> local_device_state; in GetTpuDevices() local
163 local_device_state = std::move(local_device_states[device_ordinal]); in GetTpuDevices()
166 core, std::move(local_device_state), task_id, coords_array, in GetTpuDevices()
|
| D | pjrt_stream_executor_client.h | 60 int id, std::unique_ptr<LocalDeviceState> local_device_state,
64 local_device_state ? local_device_state->device_ordinal() : -1), in id_()
65 local_device_state_(std::move(local_device_state)), in id_()
95 LocalDeviceState* local_device_state() const { in local_device_state() function
217 ->local_device_state(); in device_state()
|
| D | interpreter_device.h | 29 std::unique_ptr<LocalDeviceState> local_device_state);
|
| D | cpu_device.h | 28 CpuDevice(int id, std::unique_ptr<LocalDeviceState> local_device_state);
|
| D | gpu_device.h | 30 GpuDevice(int id, std::unique_ptr<LocalDeviceState> local_device_state,
|
| D | gpu_device.cc | 308 std::unique_ptr<LocalDeviceState> local_device_state, in GpuDevice() argument
310 : PjRtStreamExecutorDevice(id, std::move(local_device_state), in GpuDevice()
|