Home
last modified time | relevance | path

Searched refs:LocalDeviceState (Results 1 – 13 of 13) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/pjrt/
Dlocal_device_state.cc30 LocalDeviceState::LocalDeviceState(se::StreamExecutor* executor, in LocalDeviceState() function in xla::LocalDeviceState
66 LocalDeviceState::~LocalDeviceState() { in ~LocalDeviceState()
73 Status LocalDeviceState::SynchronizeAllActivity() { in SynchronizeAllActivity()
89 Status LocalDeviceState::ThenMemcpyDeviceToDevice( in ThenMemcpyDeviceToDevice()
99 void LocalDeviceState::ThenExecuteOnCallbackThread( in ThenExecuteOnCallbackThread()
106 se::Stream* LocalDeviceState::GetDeviceToHostStream() { in GetDeviceToHostStream()
114 se::Stream* LocalDeviceState::GetDeviceToDeviceStream() { in GetDeviceToDeviceStream()
122 std::unique_ptr<se::Stream> LocalDeviceState::BorrowStreamFromPool() { in BorrowStreamFromPool()
140 void LocalDeviceState::ReturnStreamToPool(std::unique_ptr<se::Stream> stream) { in ReturnStreamToPool()
150 int LocalDeviceState::GetNewPrngSeed() { in GetNewPrngSeed()
Dgpu_device.cc75 StatusOr<std::vector<std::unique_ptr<LocalDeviceState>>> BuildLocalDeviceStates( in BuildLocalDeviceStates()
77 std::vector<std::unique_ptr<LocalDeviceState>> addressable_devices; in BuildLocalDeviceStates()
81 addressable_devices.push_back(absl::make_unique<LocalDeviceState>( in BuildLocalDeviceStates()
82 executor, xla_client, LocalDeviceState::kComputeSynchronized, in BuildLocalDeviceStates()
91 absl::Span<std::unique_ptr<LocalDeviceState> const> addressable_devices, in CreateBFCAllocator()
148 absl::Span<std::unique_ptr<LocalDeviceState> const> addressable_devices) { in GetGpuDeviceAllocator()
230 std::vector<std::unique_ptr<LocalDeviceState>> local_device_states) { in BuildLocalDevices()
245 std::vector<std::unique_ptr<LocalDeviceState>> local_device_states, in BuildDistributedDevices()
274 std::unique_ptr<LocalDeviceState> local_device; in BuildDistributedDevices()
308 std::unique_ptr<LocalDeviceState> local_device_state, in GpuDevice()
[all …]
Dinterpreter_device.cc28 int id, std::unique_ptr<LocalDeviceState> local_device_state) in InterpreterDevice()
47 auto device_state = absl::make_unique<LocalDeviceState>( in GetInterpreterClient()
48 executor, client, LocalDeviceState::kSynchronous, /*asynchronous=*/false, in GetInterpreterClient()
Dtpu_client.cc49 class TpuDeviceState : public LocalDeviceState {
62 : LocalDeviceState(executor, client, LocalDeviceState::kAsynchronous, in TpuDeviceState()
138 std::vector<std::unique_ptr<LocalDeviceState>> local_device_states) { in GetTpuDevices()
161 std::unique_ptr<LocalDeviceState> local_device_state; in GetTpuDevices()
206 std::vector<std::unique_ptr<LocalDeviceState>> local_device_states; in GetTpuClient()
Dpjrt_stream_executor_client.cc127 StatusOr<LocalDeviceState*> PjRtStreamExecutorDevice::GetLocalDeviceState() in GetLocalDeviceState()
271 void StallStreamOnError(LocalDeviceState* local_device, se::Stream* stream) { in StallStreamOnError()
273 case LocalDeviceState::kAsynchronous: in StallStreamOnError()
279 case LocalDeviceState::kComputeSynchronized: in StallStreamOnError()
287 case LocalDeviceState::kSynchronous: in StallStreamOnError()
329 LocalDeviceState* buffer_local_device, in RecordUsage()
330 LocalDeviceState* stream_local_device, in RecordUsage()
339 LocalDeviceState::kSynchronous) || in RecordUsage()
342 LocalDeviceState::kComputeSynchronized && in RecordUsage()
365 LocalDeviceState* local_device, se::Stream* copy_stream, in AllocateDestinationBuffer()
[all …]
Dcpu_device.cc28 std::unique_ptr<LocalDeviceState> local_device_state) in CpuDevice()
53 auto device_state = absl::make_unique<LocalDeviceState>( in GetCpuClient()
54 executor, client, LocalDeviceState::kSynchronous, asynchronous, in GetCpuClient()
Dlocal_device_state.h36 class LocalDeviceState {
91 LocalDeviceState(se::StreamExecutor* executor, LocalClient* client,
94 virtual ~LocalDeviceState();
Dpjrt_stream_executor_client.h60 int id, std::unique_ptr<LocalDeviceState> local_device_state,
95 LocalDeviceState* local_device_state() const { in local_device_state()
102 StatusOr<LocalDeviceState*> GetLocalDeviceState() const;
115 const std::unique_ptr<LocalDeviceState> local_device_state_;
214 LocalDeviceState& device_state(int device_ordinal) const { in device_state()
582 CopyToDeviceHelper(PjRtDevice* dst_device, LocalDeviceState* dst_local_device,
583 LocalDeviceState* transfer_local_device,
Dinterpreter_device.h29 std::unique_ptr<LocalDeviceState> local_device_state);
Dcpu_device.h28 CpuDevice(int id, std::unique_ptr<LocalDeviceState> local_device_state);
Dgpu_device.h30 GpuDevice(int id, std::unique_ptr<LocalDeviceState> local_device_state,
Dtpu_client.h32 std::unique_ptr<LocalDeviceState> local_device_state, in PjRtTpuDevice()
/external/tensorflow/tensorflow/compiler/xla/python/
Doutfeed_receiver_test.cc92 auto device_state = absl::make_unique<LocalDeviceState>( in GetCpuClientWithNonLocalDevice()
93 executor, client, LocalDeviceState::kSynchronous, /*asynchronous=*/true, in GetCpuClientWithNonLocalDevice()