Searched refs:TpuDriver (Results 1 – 11 of 11) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/ |
| D | libtpu.h | 32 typedef struct TpuDriver TpuDriver; typedef
94 typedef struct TpuDriver*(PrototypeTpuDriver_Open)(const char* worker);
95 typedef void(PrototypeTpuDriver_Close)(struct TpuDriver* driver);
96 typedef struct TpuStatus*(PrototypeTpuDriver_Reset)(struct TpuDriver* driver);
99 struct TpuDriver* driver);
107 struct TpuDriver* driver, const struct TpuAllocationShape shape);
110 struct TpuDriver* driver, void* dst, const void* src,
114 struct TpuDriver* driver, void* dst, const void* src,
118 struct TpuDriver* driver, const struct HloProto hlo_proto,
122 PrototypeTpuDriver_CompileProgramFromText)(struct TpuDriver* driver,
[all …]
|
| D | tpu_client.h | 93 std::unique_ptr<tpu_driver::TpuDriver> driver,
129 tpu_driver::TpuDriver* driver() { return driver_.get(); } in driver()
135 std::unique_ptr<tpu_driver::TpuDriver> driver_;
153 TpuSharedBuffer(tpu_driver::TpuDriver* driver, in TpuSharedBuffer()
170 tpu_driver::TpuDriver* const driver;
|
| D | libtpu_client.c | 57 struct TpuDriver* driver = driver_fn.TpuDriver_Open("local://"); in main()
|
| D | tpu_client.cc | 90 std::unique_ptr<tpu_driver::TpuDriver> driver, in PyTpuClient()
182 tpu_driver::TpuDriver* driver = client->driver(); in FromLiterals()
404 tpu_driver::TpuDriver* driver = client_->driver(); in CopyToDevice()
|
| /external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/ |
| D | tpu_driver.cc | 27 std::string, std::function<xla::StatusOr<std::unique_ptr<TpuDriver>>(
82 const std::function<xla::StatusOr<std::unique_ptr<TpuDriver>>( in RegisterDriver()
88 /*static*/ xla::StatusOr<std::unique_ptr<TpuDriver>> TpuDriverRegistry::Open( in Open()
|
| D | tpu_driver.h | 150 class TpuDriver {
152 virtual ~TpuDriver() {} in ~TpuDriver()
237 static xla::StatusOr<std::unique_ptr<TpuDriver>> Open(
241 const std::function<xla::StatusOr<std::unique_ptr<TpuDriver>>(
|
| D | direct_tpu_driver.cc | 195 explicit DirectTpuLinearizer(::TpuDriver* driver, ::TpuDriverFn* driver_fn) in DirectTpuLinearizer()
232 ::TpuDriver* driver_;
236 class DirectTpuDriver : public TpuDriver {
472 ::TpuDriver* driver_;
484 xla::StatusOr<std::unique_ptr<TpuDriver>> RegisterDirectTpuDriver( in RegisterDirectTpuDriver()
489 return xla::StatusOr<std::unique_ptr<TpuDriver>>( in RegisterDirectTpuDriver()
496 return xla::StatusOr<std::unique_ptr<TpuDriver>>( in RegisterDirectTpuDriver()
|
| D | grpc_tpu_driver.h | 25 xla::StatusOr<std::unique_ptr<TpuDriver>> CreateGrpcTpuDriver(
|
| D | recording_tpu_driver.cc | 127 class RecordingTpuDriver : public TpuDriver {
129 explicit RecordingTpuDriver(std::unique_ptr<TpuDriver> driver, in RecordingTpuDriver()
470 std::unique_ptr<TpuDriver> driver_;
538 xla::StatusOr<std::unique_ptr<TpuDriver>> RegisterRecordingTpuDriver( in RegisterRecordingTpuDriver()
569 return std::unique_ptr<TpuDriver>( in RegisterRecordingTpuDriver()
|
| D | pod_tpu_driver.cc | 243 class PodTpuDriver : public TpuDriver {
813 absl::flat_hash_map<int32_t, std::unique_ptr<TpuDriver>> drivers_;
815 absl::flat_hash_map<int32_t, TpuDriver*> core_to_driver_;
956 xla::StatusOr<std::unique_ptr<TpuDriver>> CreatePodTpuDriver( in CreatePodTpuDriver()
959 return std::unique_ptr<TpuDriver>(new PodTpuDriver(config, creds)); in CreatePodTpuDriver()
971 -> xla::StatusOr<std::unique_ptr<TpuDriver>> { in __anon6b014a981202()
|
| D | grpc_tpu_driver.cc | 314 class GrpcTpuDriver : public TpuDriver {
1057 xla::StatusOr<std::unique_ptr<TpuDriver>> CreateGrpcTpuDriver( in CreateGrpcTpuDriver()
1080 return std::unique_ptr<TpuDriver>( in CreateGrpcTpuDriver()
1087 -> xla::StatusOr<std::unique_ptr<TpuDriver>> { in __anon4073f9c10402()
|