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