Searched refs:TpuNodeContext (Results 1 – 11 of 11) sorted by relevance
| /external/tensorflow/tensorflow/stream_executor/tpu/ |
| D | tpu_node_context.cc | 29 StatusOr<std::unique_ptr<TpuNodeContext>> TpuNodeContext::Create( in Create()
40 return std::make_unique<TpuNodeContext>(device_ordinal, node_context); in Create()
43 TpuNodeContext::~TpuNodeContext() { in ~TpuNodeContext()
48 Status TpuNodeContext::StopChipHeartbeats() { in StopChipHeartbeats()
55 Status TpuNodeContext::CloseTpuHost() { in CloseTpuHost()
62 Status TpuNodeContext::Initialize(int device_ordinal) { in Initialize()
69 TpuPlatformInterface* TpuNodeContext::platform() { in platform()
73 int TpuNodeContext::device_ordinal() const { return device_ordinal_; } in device_ordinal()
75 xla::Backend* TpuNodeContext::backend() const { in backend()
84 stream_executor::StreamExecutor* TpuNodeContext::stream_executor() const { in stream_executor()
|
| D | tpu_node_context.h | 42 class TpuNodeContext final {
48 static StatusOr<std::unique_ptr<TpuNodeContext>> Create(int device_ordinal);
50 explicit TpuNodeContext(int device_ordinal, XLA_TpuNodeContext* node_context) in TpuNodeContext() function
54 ~TpuNodeContext();
74 TF_DISALLOW_COPY_AND_ASSIGN(TpuNodeContext);
|
| /external/tensorflow/tensorflow/compiler/xrt/kernels/ |
| D | tpu_execute_op.cc | 138 tpu::TpuNodeContext* node_context, se::Stream* stream, in AllocateOutputTuple()
172 tpu::TpuNodeContext* node_context, se::Stream* stream, in AllocateOutputTensors()
188 OpKernelContext* context, tpu::TpuNodeContext* node_context, in RunExecutable()
230 OpKernelContext* context, tpu::TpuNodeContext* node_context, in ExecuteTPUProgram()
276 TF_ASSIGN_OR_RETURN(std::unique_ptr<tpu::TpuNodeContext> node_context, in DoWork()
277 tpu::TpuNodeContext::Create(device_ordinal)); in DoWork()
398 TF_ASSIGN_OR_RETURN(std::unique_ptr<tpu::TpuNodeContext> node_context, in DoWork()
399 tpu::TpuNodeContext::Create(device_ordinal)); in DoWork()
|
| /external/tensorflow/tensorflow/core/tpu/ |
| D | tpu_execute.cc | 58 using ::tensorflow::tpu::TpuNodeContext;
67 explicit HostTransferManager(TpuNodeContext*, xla::Backend*) {} in HostTransferManager() argument
300 Status status = TpuNodeContext::StopChipHeartbeats(); in TPUCancelExecution()
311 Status status = TpuNodeContext::CloseTpuHost(); in TPUCancelExecution()
418 TpuNodeContext* node_context, xla::DeviceAssignment* device_assignment, in TPUExecute()
|
| D | tpu_execute.h | 46 tpu::TpuNodeContext* node_context, xla::DeviceAssignment* device_assignment,
|
| /external/tensorflow/tensorflow/compiler/xrt/ |
| D | xrt_tpu_device.h | 56 std::unique_ptr<tpu::TpuNodeContext> node_context_;
|
| D | xrt_tpu_device.cc | 39 tpu::TpuNodeContext::Create(device_ordinal)); in Acquire()
|
| /external/tensorflow/tensorflow/core/tpu/kernels/ |
| D | tpu_compile_op.cc | 68 Status close_status = TpuNodeContext::CloseTpuHost(); in Compute()
|
| D | tpu_execute_op.cc | 64 using ::tensorflow::tpu::TpuNodeContext;
397 const VariableUpdateMap& variable_updates, TpuNodeContext* node_context, in AllocateOutputTensors()
628 TF_ASSIGN_OR_RETURN(std::unique_ptr<TpuNodeContext> node_context, in DoWork()
629 TpuNodeContext::Create(device_ordinal)); in DoWork()
|
| D | transfer_ops.cc | 100 TF_CHECK_OK(tpu::TpuNodeContext::CloseTpuHost()); in Cancel()
|
| /external/tensorflow/tensorflow/compiler/jit/ |
| D | xla_tpu_device.cc | 365 TF_RETURN_IF_ERROR(tpu::TpuNodeContext::Initialize(i)); in CreateDevices()
389 device->SetHandleDeviceErrorCallback(&tpu::TpuNodeContext::CloseTpuHost); in CreateDevices()
|