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 53 TFTPU_SET_FN(executor_fn, TpuStream_New); 54 TFTPU_SET_FN(executor_fn, TpuStream_Free); 55 TFTPU_SET_FN(executor_fn, TpuStream_Stream); 56 TFTPU_SET_FN(executor_fn, TpuStream_Status); 57 TFTPU_SET_FN(executor_fn, TpuStream_IsSameSharedMemoryLocation); 58 TFTPU_SET_FN(executor_fn, TpuStream_EnqueueTransferHostToDevice); 59 TFTPU_SET_FN(executor_fn, TpuStream_EnqueueTransferDeviceToHost); 60 TFTPU_SET_FN(executor_fn, TpuStream_TpuEnqueueOnDeviceSendRecvLocal); 61 62 TFTPU_SET_FN(executor_fn, TpuEvent_New); 63 TFTPU_SET_FN(executor_fn, TpuEvent_Free); 64 65 TFTPU_SET_FN(executor_fn, TpuTimer_New); 66 TFTPU_SET_FN(executor_fn, TpuTimer_Free); 67 TFTPU_SET_FN(executor_fn, TpuTimer_Nanoseconds); 68 TFTPU_SET_FN(executor_fn, TpuTimer_Microseconds); 69 70 TFTPU_SET_FN(executor_fn, TpuStatus_New); 71 TFTPU_SET_FN(executor_fn, TpuStatus_Create); 72 TFTPU_SET_FN(executor_fn, TpuStatus_Set); 73 TFTPU_SET_FN(executor_fn, TpuStatus_Free); 74 TFTPU_SET_FN(executor_fn, TpuStatus_Message); 75 TFTPU_SET_FN(executor_fn, TpuStatus_Code); 76 TFTPU_SET_FN(executor_fn, TpuStatus_Ok); 77 78 TFTPU_SET_FN(executor_fn, TpuStreamExecutorConfig_Default); 79 TFTPU_SET_FN(executor_fn, TpuStreamExecutorConfig_SetOrdinal); 80 TFTPU_SET_FN(executor_fn, TpuStreamExecutorConfig_Free); 81 82 TFTPU_SET_FN(executor_fn, TpuDeviceDescription_New); 83 TFTPU_SET_FN(executor_fn, TpuDeviceDescription_Free); 84 85 TFTPU_SET_FN(executor_fn, TpuExecutor_CreateDeviceDescription); 86 TFTPU_SET_FN(executor_fn, TpuExecutor_NewDeviceOptions); 87 TFTPU_SET_FN(executor_fn, TpuExecutor_FreeDeviceOptions); 88 TFTPU_SET_FN(executor_fn, TpuExecutor_HostCallback); 89 90 TFTPU_SET_FN(executor_fn, TpuTransferManager_New); 91 TFTPU_SET_FN(executor_fn, TpuTransferManager_Free); 92 TFTPU_SET_FN(executor_fn, TpuTransferManager_PlatformId); 93 TFTPU_SET_FN(executor_fn, TpuTransferManager_HostShapeToDeviceShape); 94 TFTPU_SET_FN(executor_fn, TpuTransferManager_TransferLiteralToDeviceAsync); 95 TFTPU_SET_FN(executor_fn, TpuTransferManager_TransferLiteralFromDevice); 96 TFTPU_SET_FN(executor_fn, TpuTransferManager_GetByteSizeRequirement); 97 TFTPU_SET_FN(executor_fn, TpuTransferManager_ChooseCompactLayoutForShape); 98 TFTPU_SET_FN(executor_fn, TpuTransferManager_CanShapedBufferBeAccessedNow); 99 TFTPU_SET_FN(executor_fn, TpuTransferManager_CanBufferBeAccessedNow); 100 TFTPU_SET_FN(executor_fn, TpuTransferManager_WriteSingleTupleIndexTable); 101 TFTPU_SET_FN(executor_fn, TpuTransferManager_GetInfeedLayout); 102 TFTPU_SET_FN(executor_fn, TpuTransferManager_LinearizeToBuffers); 103 TFTPU_SET_FN(executor_fn, TpuTransferManager_FreeBuffers); 104 TFTPU_SET_FN(executor_fn, TpuTransferManager_TransferLiteralToInfeed); 105 TFTPU_SET_FN(executor_fn, TpuTransferManager_TransferBuffersToInfeed); 106 TFTPU_SET_FN(executor_fn, TpuTransferManager_TransferLiteralFromOutfeed); 107 TFTPU_SET_FN(executor_fn, TpuTransferManager_ResetDevices); 108 109 TFTPU_SET_FN(executor_fn, TpuComputationPlacer_New); 110 TFTPU_SET_FN(executor_fn, TpuComputationPlacer_Free); 111 TFTPU_SET_FN(executor_fn, TpuComputationPlacer_AssignDevices); 112 TFTPU_SET_FN(executor_fn, TpuComputationPlacer_AssignLocalDevices); 113 114 TFTPU_SET_FN(executor_fn, TpuTopology_LogicalDevicesPerHost); 115 TFTPU_SET_FN(executor_fn, TpuTopology_LogicalDevicesPerChip); 116 TFTPU_SET_FN(executor_fn, TpuTopology_HostCount); 117 TFTPU_SET_FN(executor_fn, TpuTopology_ChipsPerHost); 118 TFTPU_SET_FN(executor_fn, TpuTopology_ChipBounds_X); 119 TFTPU_SET_FN(executor_fn, TpuTopology_ChipBounds_Y); 120 TFTPU_SET_FN(executor_fn, TpuTopology_ChipBounds_Z); 121 TFTPU_SET_FN(executor_fn, TpuTopology_HasChip); 122 TFTPU_SET_FN(executor_fn, TpuTopology_CoreForId); 123 TFTPU_SET_FN(executor_fn, TpuTopology_Core); 124 TFTPU_SET_FN(executor_fn, TpuTopology_NumCores); 125 TFTPU_SET_FN(executor_fn, TpuTopology_Cores); 126 TFTPU_SET_FN(executor_fn, TpuTopology_IdForHost); 127 TFTPU_SET_FN(executor_fn, TpuTopology_Version); 128 129 TFTPU_SET_FN(executor_fn, TpuCoreLocation_ChipCoordinates); 130 TFTPU_SET_FN(executor_fn, TpuCoreLocation_HostCoordinates); 131 TFTPU_SET_FN(executor_fn, TpuCoreLocation_Index); 132 TFTPU_SET_FN(executor_fn, TpuCoreLocation_Id); 133 134 TFTPU_SET_FN(executor_fn, TpuHostLocation_Id); 135 TFTPU_SET_FN(executor_fn, TpuHostLocation_NumCores); 136 TFTPU_SET_FN(executor_fn, TpuHostLocation_Cores); 137 138 TFTPU_SET_FN(executor_fn, TpuCompiler_New); 139 TFTPU_SET_FN(executor_fn, TpuCompiler_Free); 140 141 TFTPU_SET_FN(executor_fn, TpuCompiler_RunHloPasses); 142 TFTPU_SET_FN(executor_fn, TpuCompiler_RunBackend); 143 TFTPU_SET_FN(executor_fn, TpuCompiler_Compile); 144 TFTPU_SET_FN(executor_fn, TpuCompiler_ShapeSize); 145 TFTPU_SET_FN(executor_fn, TpuExecutable_ExecuteAsyncOnStream); 146 TFTPU_SET_FN(executor_fn, TpuExecutable_FreeXlaShapeIndexArray); 147 TFTPU_SET_FN(executor_fn, TpuExecutable_FreeMaybeOwningDeviceMemoryArray); 148 TFTPU_SET_FN(executor_fn, TpuExecutable_Fingerprint); 149 TFTPU_SET_FN(executor_fn, TpuExecutable_HloModule); 150 TFTPU_SET_FN(executor_fn, TpuExecutable_Free); 151 152 TFTPU_SET_FN(executor_fn, XlaShapeToTpuShapeRepresentation); 153 TFTPU_SET_FN(executor_fn, XlaShapeToTpuPaddedShape); 154 155 return tensorflow::Status::OK(); 156} 157 158} // namespace 159