1namespace { 2 3tensorflow::Status SetExecutorStructFn(void* library_handle) { 4 auto* executor_fn = tensorflow::tpu::ExecutorApiFn(); 5 6 TFTPU_SET_FN(executor_fn, TpuPlatform_New); 7 TFTPU_SET_FN(executor_fn, TpuPlatform_Free); 8 TFTPU_SET_FN(executor_fn, TpuPlatform_Initialize); 9 TFTPU_SET_FN(executor_fn, TpuPlatform_Initialized); 10 TFTPU_SET_FN(executor_fn, TpuPlatform_GetExecutor); 11 TFTPU_SET_FN(executor_fn, TpuPlatform_Id); 12 TFTPU_SET_FN(executor_fn, TpuPlatform_VisibleDeviceCount); 13 TFTPU_SET_FN(executor_fn, TpuPlatform_TpuMemoryLimit); 14 TFTPU_SET_FN(executor_fn, TpuPlatform_ShouldRegisterTpuDeviceToDeviceCopy); 15 TFTPU_SET_FN(executor_fn, TpuPlatform_GetTopologyPtr); 16 TFTPU_SET_FN(executor_fn, TpuPlatform_GetHostLocation); 17 TFTPU_SET_FN(executor_fn, TpuPlatform_GetRuntimeVersion); 18 19 TFTPU_SET_FN(executor_fn, TpuExecutor_Init); 20 TFTPU_SET_FN(executor_fn, TpuExecutor_Free); 21 TFTPU_SET_FN(executor_fn, TpuExecutor_PlatformDeviceCount); 22 TFTPU_SET_FN(executor_fn, TpuExecutor_Allocate); 23 TFTPU_SET_FN(executor_fn, TpuExecutor_Deallocate); 24 TFTPU_SET_FN(executor_fn, TpuExecutor_GetAllocatorStats); 25 TFTPU_SET_FN(executor_fn, TpuExecutor_DeviceMemoryUsage); 26 TFTPU_SET_FN(executor_fn, TpuExecutor_AllocateStream); 27 TFTPU_SET_FN(executor_fn, TpuExecutor_DeallocateStream); 28 TFTPU_SET_FN(executor_fn, TpuExecutor_CreateStreamDependency); 29 TFTPU_SET_FN(executor_fn, TpuExecutor_GetStatus); 30 TFTPU_SET_FN(executor_fn, TpuExecutor_GetCoreLocation); 31 TFTPU_SET_FN(executor_fn, TpuExecutor_AllocateEvent); 32 TFTPU_SET_FN(executor_fn, TpuExecutor_DeallocateEvent); 33 TFTPU_SET_FN(executor_fn, TpuExecutor_PollForEventStatus); 34 TFTPU_SET_FN(executor_fn, TpuExecutor_RecordEvent); 35 TFTPU_SET_FN(executor_fn, TpuExecutor_WaitForEvent); 36 TFTPU_SET_FN(executor_fn, TpuExecutor_AllocateTimer); 37 TFTPU_SET_FN(executor_fn, TpuExecutor_DeallocateTimer); 38 TFTPU_SET_FN(executor_fn, TpuExecutor_StartTimer); 39 TFTPU_SET_FN(executor_fn, TpuExecutor_StopTimer); 40 TFTPU_SET_FN(executor_fn, TpuExecutor_SynchronousMemcpyToHost); 41 TFTPU_SET_FN(executor_fn, TpuExecutor_SynchronousMemcpyFromHost); 42 TFTPU_SET_FN(executor_fn, TpuExecutor_MemcpyToHost); 43 TFTPU_SET_FN(executor_fn, TpuExecutor_MemcpyFromHost); 44 TFTPU_SET_FN(executor_fn, TpuExecutor_EnqueueInfeed); 45 TFTPU_SET_FN(executor_fn, TpuExecutor_DequeueOutfeed); 46 TFTPU_SET_FN(executor_fn, TpuExecutor_WaitForInfeedReady); 47 TFTPU_SET_FN(executor_fn, TpuExecutor_WaitForOutfeedReady); 48 TFTPU_SET_FN(executor_fn, TpuExecutor_BlockHostUntilDone); 49 TFTPU_SET_FN(executor_fn, TpuExecutor_BlockUntilDoneOrFailed); 50 TFTPU_SET_FN(executor_fn, TpuExecutor_SyncAndForgetFailedStreams); 51 TFTPU_SET_FN(executor_fn, TpuExecutor_SynchronizeAllActivity); 52 TFTPU_SET_FN(executor_fn, TpuExecutor_UnloadAllPrograms); 53 TFTPU_SET_FN(executor_fn, TpuExecutor_EnqueueCompactionOnStreamForHbm); 54 55 TFTPU_SET_FN(executor_fn, TpuStream_New); 56 TFTPU_SET_FN(executor_fn, TpuStream_Free); 57 TFTPU_SET_FN(executor_fn, TpuStream_Stream); 58 TFTPU_SET_FN(executor_fn, TpuStream_Status); 59 TFTPU_SET_FN(executor_fn, TpuStream_IsSameSharedMemoryLocation); 60 TFTPU_SET_FN(executor_fn, TpuStream_EnqueueTransferHostToDevice); 61 TFTPU_SET_FN(executor_fn, TpuStream_EnqueueTransferDeviceToHost); 62 TFTPU_SET_FN(executor_fn, TpuStream_TpuEnqueueOnDeviceSendRecvLocal); 63 64 TFTPU_SET_FN(executor_fn, TpuEvent_New); 65 TFTPU_SET_FN(executor_fn, TpuEvent_Free); 66 67 TFTPU_SET_FN(executor_fn, TpuTimer_New); 68 TFTPU_SET_FN(executor_fn, TpuTimer_Free); 69 TFTPU_SET_FN(executor_fn, TpuTimer_Nanoseconds); 70 TFTPU_SET_FN(executor_fn, TpuTimer_Microseconds); 71 72 TFTPU_SET_FN(executor_fn, TpuStatus_New); 73 TFTPU_SET_FN(executor_fn, TpuStatus_Create); 74 TFTPU_SET_FN(executor_fn, TpuStatus_Set); 75 TFTPU_SET_FN(executor_fn, TpuStatus_Free); 76 TFTPU_SET_FN(executor_fn, TpuStatus_Message); 77 TFTPU_SET_FN(executor_fn, TpuStatus_Code); 78 TFTPU_SET_FN(executor_fn, TpuStatus_Ok); 79 80 TFTPU_SET_FN(executor_fn, TpuStreamExecutorConfig_Default); 81 TFTPU_SET_FN(executor_fn, TpuStreamExecutorConfig_SetOrdinal); 82 TFTPU_SET_FN(executor_fn, TpuStreamExecutorConfig_Free); 83 84 TFTPU_SET_FN(executor_fn, TpuDeviceDescription_New); 85 TFTPU_SET_FN(executor_fn, TpuDeviceDescription_Free); 86 87 TFTPU_SET_FN(executor_fn, TpuExecutor_CreateDeviceDescription); 88 TFTPU_SET_FN(executor_fn, TpuExecutor_NewDeviceOptions); 89 TFTPU_SET_FN(executor_fn, TpuExecutor_FreeDeviceOptions); 90 TFTPU_SET_FN(executor_fn, TpuExecutor_HostCallback); 91 92 TFTPU_SET_FN(executor_fn, TpuTransferManager_New); 93 TFTPU_SET_FN(executor_fn, TpuTransferManager_Free); 94 TFTPU_SET_FN(executor_fn, TpuTransferManager_PlatformId); 95 TFTPU_SET_FN(executor_fn, TpuTransferManager_HostShapeToDeviceShape); 96 TFTPU_SET_FN(executor_fn, TpuTransferManager_TransferLiteralToDeviceAsync); 97 TFTPU_SET_FN(executor_fn, TpuTransferManager_TransferLiteralFromDevice); 98 TFTPU_SET_FN(executor_fn, TpuTransferManager_GetByteSizeRequirement); 99 TFTPU_SET_FN(executor_fn, TpuTransferManager_ChooseCompactLayoutForShape); 100 TFTPU_SET_FN(executor_fn, TpuTransferManager_CanShapedBufferBeAccessedNow); 101 TFTPU_SET_FN(executor_fn, TpuTransferManager_CanBufferBeAccessedNow); 102 TFTPU_SET_FN(executor_fn, TpuTransferManager_WriteSingleTupleIndexTable); 103 TFTPU_SET_FN(executor_fn, TpuTransferManager_GetInfeedLayout); 104 TFTPU_SET_FN(executor_fn, TpuTransferManager_LinearizeToBuffers); 105 TFTPU_SET_FN(executor_fn, TpuTransferManager_FreeBuffers); 106 TFTPU_SET_FN(executor_fn, TpuTransferManager_TransferLiteralToInfeed); 107 TFTPU_SET_FN(executor_fn, TpuTransferManager_TransferBuffersToInfeed); 108 TFTPU_SET_FN(executor_fn, TpuTransferManager_TransferLiteralFromOutfeed); 109 TFTPU_SET_FN(executor_fn, TpuTransferManager_ResetDevices); 110 TFTPU_SET_FN(executor_fn, TpuTransferManager_ReadDynamicShapes); 111 112 TFTPU_SET_FN(executor_fn, TpuComputationPlacer_New); 113 TFTPU_SET_FN(executor_fn, TpuComputationPlacer_Free); 114 TFTPU_SET_FN(executor_fn, TpuComputationPlacer_AssignDevices); 115 TFTPU_SET_FN(executor_fn, TpuComputationPlacer_AssignLocalDevices); 116 117 TFTPU_SET_FN(executor_fn, TpuTopology_LogicalDevicesPerHost); 118 TFTPU_SET_FN(executor_fn, TpuTopology_LogicalDevicesPerChip); 119 TFTPU_SET_FN(executor_fn, TpuTopology_HostCount); 120 TFTPU_SET_FN(executor_fn, TpuTopology_ChipsPerHost); 121 TFTPU_SET_FN(executor_fn, TpuTopology_ChipBounds_X); 122 TFTPU_SET_FN(executor_fn, TpuTopology_ChipBounds_Y); 123 TFTPU_SET_FN(executor_fn, TpuTopology_ChipBounds_Z); 124 TFTPU_SET_FN(executor_fn, TpuTopology_HasChip); 125 TFTPU_SET_FN(executor_fn, TpuTopology_CoreForId); 126 TFTPU_SET_FN(executor_fn, TpuTopology_Core); 127 TFTPU_SET_FN(executor_fn, TpuTopology_NumCores); 128 TFTPU_SET_FN(executor_fn, TpuTopology_Cores); 129 TFTPU_SET_FN(executor_fn, TpuTopology_IdForHost); 130 TFTPU_SET_FN(executor_fn, TpuTopology_Version); 131 132 TFTPU_SET_FN(executor_fn, TpuCoreLocation_ChipCoordinates); 133 TFTPU_SET_FN(executor_fn, TpuCoreLocation_HostCoordinates); 134 TFTPU_SET_FN(executor_fn, TpuCoreLocation_Index); 135 TFTPU_SET_FN(executor_fn, TpuCoreLocation_Id); 136 137 TFTPU_SET_FN(executor_fn, TpuHostLocation_Id); 138 TFTPU_SET_FN(executor_fn, TpuHostLocation_NumCores); 139 TFTPU_SET_FN(executor_fn, TpuHostLocation_Cores); 140 141 TFTPU_SET_FN(executor_fn, TpuCompiler_New); 142 TFTPU_SET_FN(executor_fn, TpuCompiler_Free); 143 144 TFTPU_SET_FN(executor_fn, TpuCompiler_RunHloPasses); 145 TFTPU_SET_FN(executor_fn, TpuCompiler_RunBackend); 146 TFTPU_SET_FN(executor_fn, TpuCompiler_Compile); 147 TFTPU_SET_FN(executor_fn, TpuCompiler_ShapeSize); 148 TFTPU_SET_FN(executor_fn, TpuCompiler_DefaultDeviceShapeRepresentation); 149 TFTPU_SET_FN(executor_fn, TpuExecutable_ExecuteAsyncOnStream); 150 TFTPU_SET_FN(executor_fn, TpuExecutable_FreeXlaShapeIndexArray); 151 TFTPU_SET_FN(executor_fn, TpuExecutable_FreeMaybeOwningDeviceMemoryArray); 152 TFTPU_SET_FN(executor_fn, TpuExecutable_Fingerprint); 153 TFTPU_SET_FN(executor_fn, TpuExecutable_Serialize); 154 TFTPU_SET_FN(executor_fn, TpuExecutableSerialize_GetByteSize); 155 TFTPU_SET_FN(executor_fn, TpuExecutableSerialize_WriteToArray); 156 TFTPU_SET_FN(executor_fn, TpuExecutableSerialize_FreeHandle); 157 TFTPU_SET_FN(executor_fn, TpuExecutable_Deserialize); 158 TFTPU_SET_FN(executor_fn, TpuExecutable_HloModule); 159 TFTPU_SET_FN(executor_fn, TpuExecutable_Free); 160 161 TFTPU_SET_FN(executor_fn, XlaShapeToTpuShapeRepresentation); 162 TFTPU_SET_FN(executor_fn, XlaShapeToTpuPaddedShape); 163 164 TFTPU_SET_FN(executor_fn, TpuAsyncCollectiveOffloadHelper_Init); 165 TFTPU_SET_FN(executor_fn, TpuAsyncCollectiveOffloadHelper_Shutdown); 166 167 return tensorflow::Status::OK(); 168} 169 170} // namespace 171