1#if defined(PLATFORM_GOOGLE) 2#include "third_party/tensorflow/core/tpu/tpu_executor_init_fns.inc" 3#else 4#include "tensorflow/core/tpu/tpu_executor_init_fns.inc" 5#endif 6 7namespace { 8 9tensorflow::Status SetTpuOpsStructFns(void* library_handle) { 10 // Constant cast so that we can initialize the functions. The functions are 11 // mutable here because this is the only place where they are initialized. 12 auto* ops_api_fn = const_cast<TfTpu_OpsApiFn*>(tensorflow::tpu::OpsApiFn()); 13 14 TFTPU_SET_FN(ops_api_fn, ConfigureDistributedTpuOp_DoWork); 15 TFTPU_SET_FN(ops_api_fn, WaitForDistributedTpuOp_DoWork); 16 TFTPU_SET_FN(ops_api_fn, InitializeHostForDistributedTpuOp_DoWork); 17 TFTPU_SET_FN(ops_api_fn, SetGlobalTPUArrayOp_DoWork); 18 TFTPU_SET_FN(ops_api_fn, DisconnectDistributedTpuChipsOp_DoWork); 19 TFTPU_SET_FN(ops_api_fn, TpuConfigurationApi_FreeCharArray); 20 TFTPU_SET_FN(ops_api_fn, TpuConfigurationApi_FreeInt32Array); 21 TFTPU_SET_FN(ops_api_fn, TpuConfigurationApi_HasTPUPodState); 22 TFTPU_SET_FN(ops_api_fn, TpuConfigurationApi_TpusPerHost); 23 TFTPU_SET_FN(ops_api_fn, TpuConfigurationApi_TpuMemoryLimit); 24 TFTPU_SET_FN(ops_api_fn, 25 TpuConfigurationApi_RemoteCompilationCacheSizeInBytes); 26 TFTPU_SET_FN(ops_api_fn, 27 TpuConfigurationApi_CompilationCacheServerAddressFromConfig); 28 TFTPU_SET_FN(ops_api_fn, TpuConfigurationApi_GetServerAddressAndPort); 29 30 TFTPU_SET_FN(ops_api_fn, TpuMeshState_Create); 31 TFTPU_SET_FN(ops_api_fn, TpuMeshState_Free); 32 TFTPU_SET_FN(ops_api_fn, TpuMeshState_MeshCommonState); 33 34 TFTPU_SET_FN(ops_api_fn, TpuCompile_CompileAndBuild); 35 TFTPU_SET_FN(ops_api_fn, TpuCompile_XrtCompileAndBuild); 36 37 TFTPU_SET_FN(ops_api_fn, TpuExecutable_LoadProgramAndEnqueueToStream); 38 TFTPU_SET_FN(ops_api_fn, HardwareLayout_HostShapeToDeviceShape); 39 TFTPU_SET_FN(ops_api_fn, HardwareLayout_ShapeSize); 40 TFTPU_SET_FN(ops_api_fn, HardwareLayout_ShapeSizeCompact); 41 TFTPU_SET_FN(ops_api_fn, HardwareLayout_ShapeSizeCompactRaw); 42 TFTPU_SET_FN(ops_api_fn, TpuExecute_RuntimeInputToPaddedData); 43 44 TFTPU_SET_FN(ops_api_fn, TpuProgram_New); 45 TFTPU_SET_FN(ops_api_fn, TpuProgram_Free); 46 TFTPU_SET_FN(ops_api_fn, TpuProgram_NewArray); 47 TFTPU_SET_FN(ops_api_fn, TpuProgram_FreeArray); 48 TFTPU_SET_FN(ops_api_fn, TpuProgram_UnloadAndDestroy); 49 TFTPU_SET_FN(ops_api_fn, TpuProgram_GetProgramSize); 50 TFTPU_SET_FN(ops_api_fn, TpuProgram_LogProgramMemorySummary); 51 TFTPU_SET_FN(ops_api_fn, TpuProgram_GetExecutableInfo); 52 TFTPU_SET_FN(ops_api_fn, TpuProgram_GetHostTransferInfo); 53 TFTPU_SET_FN(ops_api_fn, TpuProgram_GetHloMetadata); 54 TFTPU_SET_FN(ops_api_fn, TpuProgram_GetMayModifyVariables); 55 TFTPU_SET_FN(ops_api_fn, TpuProgram_HasSharding); 56 TFTPU_SET_FN(ops_api_fn, TpuProgram_GetTpuProgram); 57 TFTPU_SET_FN(ops_api_fn, TpuProgram_SerializeTpuExecutable); 58 TFTPU_SET_FN(ops_api_fn, TpuProgram_SerializeCompilerMetadata); 59 TFTPU_SET_FN(ops_api_fn, 60 TpuProgram_DeserializeFromGetTpuProgramResponseProto); 61 62 TFTPU_SET_FN(ops_api_fn, TpuNodeContext_Create); 63 TFTPU_SET_FN(ops_api_fn, TpuNodeContext_Free); 64 TFTPU_SET_FN(ops_api_fn, TpuNodeContext_Initialize); 65 TFTPU_SET_FN(ops_api_fn, TpuNodeContext_StopChipHeartbeats); 66 TFTPU_SET_FN(ops_api_fn, TpuNodeContext_CloseTpuHost); 67 68 TFTPU_SET_FN(ops_api_fn, TpuTopology_AvailableCoreCount); 69 TFTPU_SET_FN(ops_api_fn, TpuNetUtil_RecycleUnusedPort); 70 TFTPU_SET_FN(ops_api_fn, TpuCompile_IsTpuCompilationEnabled); 71 TFTPU_SET_FN(ops_api_fn, TpuCompile_ShouldTpuCompileOpIgnoreCancellation); 72 TFTPU_SET_FN(ops_api_fn, TpuCompile_CreateCompilationCacheKey); 73 TFTPU_SET_FN(ops_api_fn, TpuCompile_DestroyCompilationCacheKey); 74 TFTPU_SET_FN(ops_api_fn, TpuCompile_CreateGuaranteedConstFingerprint); 75 76 TFTPU_SET_FN(ops_api_fn, TpuProfiler_Create); 77 TFTPU_SET_FN(ops_api_fn, TpuProfiler_Destroy); 78 TFTPU_SET_FN(ops_api_fn, TpuProfiler_Start); 79 TFTPU_SET_FN(ops_api_fn, TpuProfiler_Stop); 80 TFTPU_SET_FN(ops_api_fn, TpuProfiler_CollectData); 81 82 TFTPU_SET_FN(ops_api_fn, TfTpu_InitializeTpuModelServer); 83 84 return tensorflow::Status::OK(); 85} 86 87tensorflow::Status InitializeTpuStructFns(void* library_handle) { 88 TF_RETURN_IF_ERROR(SetTpuOpsStructFns(library_handle)); 89 TF_RETURN_IF_ERROR(SetExecutorStructFn(library_handle)); 90 91 return tensorflow::Status::OK(); 92} 93 94} // namespace 95