• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2 
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6 
7     http://www.apache.org/licenses/LICENSE-2.0
8 
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 
16 #ifndef TENSORFLOW_COMPILER_XLA_PYTHON_TPU_DRIVER_CLIENT_LIBTPU_H_
17 #define TENSORFLOW_COMPILER_XLA_PYTHON_TPU_DRIVER_CLIENT_LIBTPU_H_
18 
19 #include <stdbool.h>
20 #include <stdint.h>
21 
22 #define TPUDRIVER_CAPI_EXPORT __attribute__((visibility("default")))
23 
24 #ifdef __cplusplus
25 extern "C" {
26 #endif
27 
28 // ------------------- TPU Driver Support -----------------------
29 
30 struct TpuDriverFn;
31 
32 typedef struct TpuDriver TpuDriver;
33 
34 typedef struct TpuEvent TpuEvent;
35 
36 typedef struct TpuBufferHandleInternal TpuBufferHandleInternal;
37 
38 typedef struct TpuCompiledProgramHandleInternal
39     TpuCompiledProgramHandleInternal;
40 
41 typedef struct TpuLoadedProgramHandleInternal TpuLoadedProgramHandleInternal;
42 
43 typedef struct TpuBufferHandle {
44   TpuBufferHandleInternal* internal_handle;
45   TpuEvent* event;
46   int64_t size_in_bytes;
47 } TpuBufferHandle;
48 
49 typedef struct TpuCompiledProgramHandle {
50   TpuCompiledProgramHandleInternal* internal_handle;
51   TpuEvent* event;
52 } TpuCompiledProgramHandle;
53 
54 typedef struct TpuLoadedProgramHandle {
55   TpuLoadedProgramHandleInternal* internal_handle;
56   TpuEvent* event;
57 } TpuLoadedProgramHandle;
58 
59 // HloProto is a serialized xla::HloProto buffer.
60 typedef struct HloProto {
61   void* buffer;
62   int32_t size;
63 } HloProto;
64 
65 // DeviceAssignment is a serialized xla::DeviceAssignmentProto buffer.
66 typedef struct DeviceAssignment {
67   void* bytes;
68   int32_t size;
69 } DeviceAssignment;
70 
71 typedef struct TpuStatus {
72   int32_t code;
73   char* msg;
74 } TpuStatus;
75 
76 typedef struct CompiledProgramShape {
77   struct TpuStatus* status;
78   void* bytes;
79   int32_t size;
80 } CompiledProgramShape;
81 
82 typedef struct TpuAllocationShape {
83   void* bytes;
84   int32_t size;
85 } TpuAllocationShape;
86 
87 typedef struct TpuSystemInfo {
88   void* bytes;
89   int32_t size;
90 } TpuSystemInfo;
91 
92 typedef void(PrototypeTpuDriver_Initialize)(struct TpuDriverFn* driver_fn,
93                                             bool initialize);
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);
97 
98 typedef struct TpuSystemInfo*(PrototypeTpuDriver_QuerySystemInfo)(
99     struct TpuDriver* driver);
100 
101 typedef void(PrototypeTpuDriver_FreeSystemInfo)(struct TpuSystemInfo* info);
102 
103 // TODO(frankchn): Make this not a hard-coded constant.
104 const int32_t MemoryRegion_HBM = 1;
105 
106 typedef int64_t(PrototypeTpuDriver_ComputeLinearizedBytesFromShape)(
107     struct TpuDriver* driver, const struct TpuAllocationShape shape);
108 
109 typedef struct TpuStatus*(PrototypeTpuDriver_LinearizeShape)(
110     struct TpuDriver* driver, void* dst, const void* src,
111     const struct TpuAllocationShape shape);
112 
113 typedef struct TpuStatus*(PrototypeTpuDriver_DelinearizeShape)(
114     struct TpuDriver* driver, void* dst, const void* src,
115     const struct TpuAllocationShape shape);
116 
117 typedef struct TpuCompiledProgramHandle*(PrototypeTpuDriver_CompileProgram)(
118     struct TpuDriver* driver, const struct HloProto hlo_proto,
119     int32_t num_replicas, int32_t eventc, struct TpuEvent** eventv);
120 
121 typedef struct TpuCompiledProgramHandle*(
122     PrototypeTpuDriver_CompileProgramFromText)(struct TpuDriver* driver,
123                                                const char* hlo_text,
124                                                int32_t num_replicas,
125                                                int32_t eventc,
126                                                struct TpuEvent** eventv);
127 
128 /* Note: We are not responsible for freeing the event within the
129  * TpuCompiledProgramHandle. You have to call FreeEvent separately to ensure
130  * that memory does not leak.
131  */
132 typedef void(PrototypeTpuDriver_FreeCompiledProgramHandle)(
133     struct TpuCompiledProgramHandle* handle);
134 
135 typedef struct TpuLoadedProgramHandle*(PrototypeTpuDriver_LoadProgram)(
136     struct TpuDriver* driver, int32_t core_id,
137     const struct TpuCompiledProgramHandle* compiled_program_handle,
138     int32_t eventc, struct TpuEvent** eventv);
139 
140 /* Note: We are not responsible for freeing the event within the
141  * TpuLoadedProgramHandle. You have to call FreeEvent separately to ensure that
142  * memory does not leak.
143  */
144 typedef struct TpuEvent*(PrototypeTpuDriver_UnloadProgram)(
145     struct TpuDriver* driver,
146     struct TpuLoadedProgramHandle* loaded_program_handle, int32_t eventc,
147     struct TpuEvent** eventv);
148 
149 typedef struct TpuEvent*(PrototypeTpuDriver_ExecuteProgram)(
150     struct TpuDriver* driver, struct TpuLoadedProgramHandle* handle,
151     int32_t inputc, struct TpuBufferHandle** input_buffer_handle,
152     int32_t outputc, struct TpuBufferHandle** output_buffer_handle,
153     struct DeviceAssignment device_assignment, int32_t eventc,
154     struct TpuEvent** eventv);
155 
156 typedef struct TpuBufferHandle*(PrototypeTpuDriver_AllocateTuple)(
157     struct TpuDriver* driver, int32_t core_id, int32_t memory_region,
158     int32_t bufferc, struct TpuBufferHandle** buffer_handle, int32_t eventc,
159     struct TpuEvent** eventv);
160 
161 typedef struct TpuBufferHandle*(PrototypeTpuDriver_Allocate)(
162     struct TpuDriver* driver, int32_t core_id, int32_t memory_region,
163     int64_t num_bytes, int32_t eventc, struct TpuEvent** eventv);
164 
165 typedef struct TpuBufferHandle*(PrototypeTpuDriver_AllocateShape)(
166     struct TpuDriver* driver, int32_t core_id, int32_t memory_region,
167     const struct TpuAllocationShape shape, int32_t eventc,
168     struct TpuEvent** eventv);
169 
170 /* Note: We are not responsible for freeing the event within the
171  * TpuBufferHandle. You have to call FreeEvent separately to ensure that memory
172  * does not leak.
173  */
174 typedef struct TpuEvent*(PrototypeTpuDriver_Deallocate)(
175     struct TpuDriver* driver, struct TpuBufferHandle* buffer_handle,
176     int32_t eventc, struct TpuEvent** eventv);
177 
178 typedef struct TpuEvent*(PrototypeTpuDriver_TransferToDevice)(
179     struct TpuDriver* driver, const void* src, struct TpuBufferHandle* dst,
180     int32_t eventc, struct TpuEvent** eventv);
181 
182 typedef struct TpuEvent*(PrototypeTpuDriver_TransferFromDevice)(
183     struct TpuDriver* driver, struct TpuBufferHandle* src, void* dst,
184     int32_t eventc, struct TpuEvent** eventv);
185 
186 typedef struct TpuEvent*(PrototypeTpuDriver_TransferFromDeviceToDevice)(
187     struct TpuDriver* driver, struct TpuBufferHandle* src,
188     struct TpuBufferHandle* dst, int32_t eventc, struct TpuEvent** eventv);
189 
190 typedef struct CompiledProgramShape*(
191     PrototypeTpuDriver_GetCompiledProgramShape)(
192     struct TpuCompiledProgramHandle* handle);
193 
194 typedef void(PrototypeTpuDriver_FreeCompiledProgramShape)(
195     struct CompiledProgramShape* shape);
196 
197 typedef void(PrototypeTpuDriver_EventAddCallback)(
198     struct TpuEvent* event,
199     void (*callback_fn)(struct TpuStatus*, void* additional_info),
200     void* additional_info);
201 
202 typedef struct TpuStatus*(PrototypeTpuDriver_EventAwait)(struct TpuEvent* event,
203                                                          int64_t timeout_in_us);
204 
205 typedef void(PrototypeTpuDriver_FreeEvent)(struct TpuEvent* event);
206 
207 typedef void(PrototypeTpuDriver_FreeStatus)(struct TpuStatus* status);
208 
209 typedef const char*(PrototypeTpuDriver_Version)();
210 
211 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_Initialize TpuDriver_Initialize;
212 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_Open TpuDriver_Open;
213 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_Close TpuDriver_Close;
214 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_Reset TpuDriver_Reset;
215 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_QuerySystemInfo
216     TpuDriver_QuerySystemInfo;
217 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_FreeSystemInfo
218     TpuDriver_FreeSystemInfo;
219 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_ComputeLinearizedBytesFromShape
220     TpuDriver_ComputeLinearizedBytesFromShape;
221 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_LinearizeShape
222     TpuDriver_LinearizeShape;
223 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_DelinearizeShape
224     TpuDriver_DelinearizeShape;
225 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_CompileProgram
226     TpuDriver_CompileProgram;
227 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_CompileProgramFromText
228     TpuDriver_CompileProgramFromText;
229 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_FreeCompiledProgramHandle
230     TpuDriver_FreeCompiledProgramHandle;
231 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_LoadProgram
232     TpuDriver_LoadProgram;
233 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_UnloadProgram
234     TpuDriver_UnloadProgram;
235 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_ExecuteProgram
236     TpuDriver_ExecuteProgram;
237 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_AllocateTuple
238     TpuDriver_AllocateTuple;
239 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_Allocate TpuDriver_Allocate;
240 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_AllocateShape
241     TpuDriver_AllocateShape;
242 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_Deallocate TpuDriver_Deallocate;
243 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_TransferToDevice
244     TpuDriver_TransferToDevice;
245 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_TransferFromDevice
246     TpuDriver_TransferFromDevice;
247 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_TransferFromDeviceToDevice
248     TpuDriver_TransferFromDeviceToDevice;
249 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_GetCompiledProgramShape
250     TpuDriver_GetCompiledProgramShape;
251 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_FreeCompiledProgramShape
252     TpuDriver_FreeCompiledProgramShape;
253 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_EventAddCallback
254     TpuDriver_EventAddCallback;
255 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_EventAwait TpuDriver_EventAwait;
256 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_FreeEvent TpuDriver_FreeEvent;
257 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_FreeStatus TpuDriver_FreeStatus;
258 TPUDRIVER_CAPI_EXPORT extern PrototypeTpuDriver_Version TpuDriver_Version;
259 
260 #ifdef __cplusplus
261 }
262 #endif
263 
264 struct TpuDriverFn {
265   PrototypeTpuDriver_Open* TpuDriver_Open;                          // NOLINT
266   PrototypeTpuDriver_Close* TpuDriver_Close;                        // NOLINT
267   PrototypeTpuDriver_Reset* TpuDriver_Reset;                        // NOLINT
268   PrototypeTpuDriver_ComputeLinearizedBytesFromShape*
269       TpuDriver_ComputeLinearizedBytesFromShape;                    // NOLINT
270   PrototypeTpuDriver_QuerySystemInfo* TpuDriver_QuerySystemInfo;    // NOLINT
271   PrototypeTpuDriver_FreeSystemInfo* TpuDriver_FreeSystemInfo;      // NOLINT
272   PrototypeTpuDriver_LinearizeShape* TpuDriver_LinearizeShape;      // NOLINT
273   PrototypeTpuDriver_DelinearizeShape* TpuDriver_DelinearizeShape;  // NOLINT
274   PrototypeTpuDriver_CompileProgram* TpuDriver_CompileProgram;      // NOLINT
275   PrototypeTpuDriver_CompileProgramFromText*
276       TpuDriver_CompileProgramFromText;                             // NOLINT
277   PrototypeTpuDriver_FreeCompiledProgramHandle*
278       TpuDriver_FreeCompiledProgramHandle;                          // NOLINT
279   PrototypeTpuDriver_LoadProgram* TpuDriver_LoadProgram;            // NOLINT
280   PrototypeTpuDriver_UnloadProgram* TpuDriver_UnloadProgram;        // NOLINT
281   PrototypeTpuDriver_ExecuteProgram* TpuDriver_ExecuteProgram;      // NOLINT
282   PrototypeTpuDriver_AllocateTuple* TpuDriver_AllocateTuple;        // NOLINT
283   PrototypeTpuDriver_Allocate* TpuDriver_Allocate;                  // NOLINT
284   PrototypeTpuDriver_AllocateShape* TpuDriver_AllocateShape;        // NOLINT
285   PrototypeTpuDriver_Deallocate* TpuDriver_Deallocate;              // NOLINT
286   PrototypeTpuDriver_TransferToDevice* TpuDriver_TransferToDevice;  // NOLINT
287   PrototypeTpuDriver_TransferFromDevice*
288       TpuDriver_TransferFromDevice;  // NOLINT
289   PrototypeTpuDriver_TransferFromDeviceToDevice*
290       TpuDriver_TransferFromDeviceToDevice;                         // NOLINT
291   PrototypeTpuDriver_GetCompiledProgramShape*
292       TpuDriver_GetCompiledProgramShape;  // NOLINT
293   PrototypeTpuDriver_FreeCompiledProgramShape*
294       TpuDriver_FreeCompiledProgramShape;                           // NOLINT
295   PrototypeTpuDriver_EventAddCallback* TpuDriver_EventAddCallback;  // NOLINT
296   PrototypeTpuDriver_EventAwait* TpuDriver_EventAwait;              // NOLINT
297   PrototypeTpuDriver_FreeEvent* TpuDriver_FreeEvent;                // NOLINT
298   PrototypeTpuDriver_FreeStatus* TpuDriver_FreeStatus;              // NOLINT
299 
300   PrototypeTpuDriver_Version* TpuDriver_Version;                    // NOLINT
301 };
302 
303 #endif  // TENSORFLOW_COMPILER_XLA_PYTHON_TPU_DRIVER_CLIENT_LIBTPU_H_
304