diff --git a/tensorflow/core/tpu/BUILD b/tensorflow/core/tpu/BUILD index 0a17ba3d408..d8abbd042b9 100644 --- a/tensorflow/core/tpu/BUILD +++ b/tensorflow/core/tpu/BUILD @@ -124,6 +124,7 @@ cc_library( deps = [ ":libtftpu_header", ":tpu_config_c_api", + ":tpu_executor_api", "//tensorflow/core/tpu/kernels:tpu_compile_c_api_hdrs", "//tensorflow/core/tpu/kernels:tpu_execute_c_api_hdrs", "//tensorflow/core/tpu/kernels:tpu_mesh_state_c_api_hdrs", @@ -133,6 +134,16 @@ cc_library( ], ) +cc_library( + name = "tpu_executor_api", + srcs = ["tpu_executor_api.cc"], + hdrs = ["tpu_executor_api.h"], + deps = [ + ":libtftpu_header", + "//tensorflow/stream_executor/tpu:tpu_executor_c_api_hdrs", + ], +) + cc_library( name = "tpu_api_dlsym_initializer", srcs = if_windows( @@ -167,6 +178,12 @@ cc_library( visibility = ["//visibility:public"], ) +cc_library( + name = "tpu_executor_init_fns", + hdrs = ["tpu_executor_init_fns.inc"], + visibility = ["//visibility:public"], +) + cc_library( name = "tpu_node_device", srcs = ["tpu_node_device.cc"], diff --git a/tensorflow/core/tpu/tpu_api.cc b/tensorflow/core/tpu/tpu_api.cc index cd6ca80e4e7..17520ea6ea4 100644 --- a/tensorflow/core/tpu/tpu_api.cc +++ b/tensorflow/core/tpu/tpu_api.cc @@ -48,11 +48,6 @@ TfTpu_TpuProgramApiFn* TpuProgramApiFn() { return &tpu_program_api_fn; } -TfTpu_ExecutorApiFn* ExecutorApiFn() { - static TfTpu_ExecutorApiFn executor_api_fn; - return &executor_api_fn; -} - TfTpu_NodeContextApiFn* NodeContextApiFn() { static TfTpu_NodeContextApiFn node_context_api_fn; return &node_context_api_fn; diff --git a/tensorflow/core/tpu/tpu_api.h b/tensorflow/core/tpu/tpu_api.h index b6edbfd14bb..a9f7bccfdb4 100644 --- a/tensorflow/core/tpu/tpu_api.h +++ b/tensorflow/core/tpu/tpu_api.h @@ -22,7 +22,7 @@ limitations under the License. #include "tensorflow/core/tpu/kernels/tpu_util_c_api.h" #include "tensorflow/core/tpu/libtftpu.h" #include "tensorflow/core/tpu/tpu_config_c_api.h" -#include "tensorflow/stream_executor/tpu/tpu_executor_c_api.h" +#include "tensorflow/core/tpu/tpu_executor_api.h" #include "tensorflow/stream_executor/tpu/tpu_node_context_c_api.h" namespace tensorflow { @@ -40,8 +40,6 @@ TfTpu_ExecuteApiFn* ExecuteApiFn(); TfTpu_TpuProgramApiFn* TpuProgramApiFn(); -TfTpu_ExecutorApiFn* ExecutorApiFn(); - TfTpu_NodeContextApiFn* NodeContextApiFn(); TfTpu_UtilApiFn* UtilApiFn(); diff --git a/tensorflow/core/tpu/tpu_executor_api.cc b/tensorflow/core/tpu/tpu_executor_api.cc new file mode 100644 index 00000000000..dd02ca27aa4 --- /dev/null +++ b/tensorflow/core/tpu/tpu_executor_api.cc @@ -0,0 +1,27 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/core/tpu/tpu_executor_api.h" + +namespace tensorflow { +namespace tpu { + +TfTpu_ExecutorApiFn* ExecutorApiFn() { + static TfTpu_ExecutorApiFn executor_api_fn; + return &executor_api_fn; +} + +} // namespace tpu +} // namespace tensorflow diff --git a/tensorflow/core/tpu/tpu_executor_api.h b/tensorflow/core/tpu/tpu_executor_api.h new file mode 100644 index 00000000000..ee07dc618a6 --- /dev/null +++ b/tensorflow/core/tpu/tpu_executor_api.h @@ -0,0 +1,30 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef TENSORFLOW_CORE_TPU_TPU_EXECUTOR_API_H_ +#define TENSORFLOW_CORE_TPU_TPU_EXECUTOR_API_H_ + +#include "tensorflow/core/tpu/libtftpu.h" +#include "tensorflow/stream_executor/tpu/tpu_executor_c_api.h" + +namespace tensorflow { +namespace tpu { + +TfTpu_ExecutorApiFn* ExecutorApiFn(); + +} // namespace tpu +} // namespace tensorflow + +#endif // TENSORFLOW_CORE_TPU_TPU_EXECUTOR_API_H_ diff --git a/tensorflow/core/tpu/tpu_executor_init_fns.inc b/tensorflow/core/tpu/tpu_executor_init_fns.inc new file mode 100644 index 00000000000..6299b415a32 --- /dev/null +++ b/tensorflow/core/tpu/tpu_executor_init_fns.inc @@ -0,0 +1,141 @@ +namespace { + +tensorflow::Status SetExecutorStructFn(void* library_handle) { + auto* executor_fn = tensorflow::tpu::ExecutorApiFn(); + + TFTPU_SET_FN(executor_fn, TpuPlatform_New); + TFTPU_SET_FN(executor_fn, TpuPlatform_Free); + TFTPU_SET_FN(executor_fn, TpuPlatform_Initialize); + TFTPU_SET_FN(executor_fn, TpuPlatform_Initialized); + TFTPU_SET_FN(executor_fn, TpuPlatform_GetExecutor); + TFTPU_SET_FN(executor_fn, TpuPlatform_Id); + TFTPU_SET_FN(executor_fn, TpuPlatform_VisibleDeviceCount); + TFTPU_SET_FN(executor_fn, TpuPlatform_TpuMemoryLimit); + TFTPU_SET_FN(executor_fn, TpuPlatform_ShouldRegisterTpuDeviceToDeviceCopy); + TFTPU_SET_FN(executor_fn, TpuPlatform_GetTopologyPtr); + TFTPU_SET_FN(executor_fn, TpuPlatform_GetHostLocation); + + TFTPU_SET_FN(executor_fn, TpuExecutor_Init); + TFTPU_SET_FN(executor_fn, TpuExecutor_Free); + TFTPU_SET_FN(executor_fn, TpuExecutor_PlatformDeviceCount); + TFTPU_SET_FN(executor_fn, TpuExecutor_Allocate); + TFTPU_SET_FN(executor_fn, TpuExecutor_Deallocate); + TFTPU_SET_FN(executor_fn, TpuExecutor_GetAllocatorStats); + TFTPU_SET_FN(executor_fn, TpuExecutor_DeviceMemoryUsage); + TFTPU_SET_FN(executor_fn, TpuExecutor_AllocateStream); + TFTPU_SET_FN(executor_fn, TpuExecutor_DeallocateStream); + TFTPU_SET_FN(executor_fn, TpuExecutor_CreateStreamDependency); + TFTPU_SET_FN(executor_fn, TpuExecutor_GetStatus); + TFTPU_SET_FN(executor_fn, TpuExecutor_GetCoreLocation); + TFTPU_SET_FN(executor_fn, TpuExecutor_AllocateEvent); + TFTPU_SET_FN(executor_fn, TpuExecutor_DeallocateEvent); + TFTPU_SET_FN(executor_fn, TpuExecutor_PollForEventStatus); + TFTPU_SET_FN(executor_fn, TpuExecutor_RecordEvent); + TFTPU_SET_FN(executor_fn, TpuExecutor_WaitForEvent); + TFTPU_SET_FN(executor_fn, TpuExecutor_AllocateTimer); + TFTPU_SET_FN(executor_fn, TpuExecutor_DeallocateTimer); + TFTPU_SET_FN(executor_fn, TpuExecutor_StartTimer); + TFTPU_SET_FN(executor_fn, TpuExecutor_StopTimer); + TFTPU_SET_FN(executor_fn, TpuExecutor_SynchronousMemcpyToHost); + TFTPU_SET_FN(executor_fn, TpuExecutor_SynchronousMemcpyFromHost); + TFTPU_SET_FN(executor_fn, TpuExecutor_MemcpyToHost); + TFTPU_SET_FN(executor_fn, TpuExecutor_MemcpyFromHost); + TFTPU_SET_FN(executor_fn, TpuExecutor_EnqueueInfeed); + TFTPU_SET_FN(executor_fn, TpuExecutor_DequeueOutfeed); + TFTPU_SET_FN(executor_fn, TpuExecutor_WaitForInfeedReady); + TFTPU_SET_FN(executor_fn, TpuExecutor_WaitForOutfeedReady); + TFTPU_SET_FN(executor_fn, TpuExecutor_BlockHostUntilDone); + TFTPU_SET_FN(executor_fn, TpuExecutor_BlockUntilDoneOrFailed); + TFTPU_SET_FN(executor_fn, TpuExecutor_SyncAndForgetFailedStreams); + TFTPU_SET_FN(executor_fn, TpuExecutor_SynchronizeAllActivity); + + TFTPU_SET_FN(executor_fn, TpuStream_New); + TFTPU_SET_FN(executor_fn, TpuStream_Free); + TFTPU_SET_FN(executor_fn, TpuStream_Stream); + TFTPU_SET_FN(executor_fn, TpuStream_Status); + TFTPU_SET_FN(executor_fn, TpuStream_IsSameSharedMemoryLocation); + TFTPU_SET_FN(executor_fn, TpuStream_TpuEnqueueOnDeviceSendRecvLocal); + + TFTPU_SET_FN(executor_fn, TpuEvent_New); + TFTPU_SET_FN(executor_fn, TpuEvent_Free); + + TFTPU_SET_FN(executor_fn, TpuTimer_New); + TFTPU_SET_FN(executor_fn, TpuTimer_Free); + TFTPU_SET_FN(executor_fn, TpuTimer_Nanoseconds); + TFTPU_SET_FN(executor_fn, TpuTimer_Microseconds); + + TFTPU_SET_FN(executor_fn, TpuStatus_New); + TFTPU_SET_FN(executor_fn, TpuStatus_Create); + TFTPU_SET_FN(executor_fn, TpuStatus_Set); + TFTPU_SET_FN(executor_fn, TpuStatus_Free); + TFTPU_SET_FN(executor_fn, TpuStatus_Message); + TFTPU_SET_FN(executor_fn, TpuStatus_Code); + TFTPU_SET_FN(executor_fn, TpuStatus_Ok); + + TFTPU_SET_FN(executor_fn, TpuStreamExecutorConfig_Default); + TFTPU_SET_FN(executor_fn, TpuStreamExecutorConfig_SetOrdinal); + TFTPU_SET_FN(executor_fn, TpuStreamExecutorConfig_Free); + + TFTPU_SET_FN(executor_fn, TpuDeviceDescription_New); + TFTPU_SET_FN(executor_fn, TpuDeviceDescription_Free); + + TFTPU_SET_FN(executor_fn, TpuExecutor_CreateDeviceDescription); + TFTPU_SET_FN(executor_fn, TpuExecutor_NewDeviceOptions); + TFTPU_SET_FN(executor_fn, TpuExecutor_FreeDeviceOptions); + TFTPU_SET_FN(executor_fn, TpuExecutor_HostCallback); + + TFTPU_SET_FN(executor_fn, TpuTransferManager_New); + TFTPU_SET_FN(executor_fn, TpuTransferManager_Free); + TFTPU_SET_FN(executor_fn, TpuTransferManager_PlatformId); + TFTPU_SET_FN(executor_fn, TpuTransferManager_HostShapeToDeviceShape); + TFTPU_SET_FN(executor_fn, TpuTransferManager_TransferLiteralToDeviceAsync); + TFTPU_SET_FN(executor_fn, TpuTransferManager_TransferLiteralFromDevice); + TFTPU_SET_FN(executor_fn, TpuTransferManager_GetByteSizeRequirement); + TFTPU_SET_FN(executor_fn, TpuTransferManager_WriteSingleTupleIndexTable); + TFTPU_SET_FN(executor_fn, TpuTransferManager_GetInfeedLayout); + TFTPU_SET_FN(executor_fn, TpuTransferManager_LinearizeToBuffers); + TFTPU_SET_FN(executor_fn, TpuTransferManager_FreeBuffers); + TFTPU_SET_FN(executor_fn, TpuTransferManager_TransferLiteralToInfeed); + TFTPU_SET_FN(executor_fn, TpuTransferManager_TransferBuffersToInfeed); + TFTPU_SET_FN(executor_fn, TpuTransferManager_TransferLiteralFromOutfeed); + TFTPU_SET_FN(executor_fn, TpuTransferManager_ResetDevices); + + TFTPU_SET_FN(executor_fn, TpuComputationPlacer_New); + TFTPU_SET_FN(executor_fn, TpuComputationPlacer_Free); + + TFTPU_SET_FN(executor_fn, TpuTopology_LogicalDevicesPerHost); + TFTPU_SET_FN(executor_fn, TpuTopology_LogicalDevicesPerChip); + TFTPU_SET_FN(executor_fn, TpuTopology_ChipBounds_X); + TFTPU_SET_FN(executor_fn, TpuTopology_ChipBounds_Y); + TFTPU_SET_FN(executor_fn, TpuTopology_ChipBounds_Z); + TFTPU_SET_FN(executor_fn, TpuTopology_HasChip); + TFTPU_SET_FN(executor_fn, TpuTopology_Core); + TFTPU_SET_FN(executor_fn, TpuTopology_NumCores); + TFTPU_SET_FN(executor_fn, TpuTopology_Cores); + TFTPU_SET_FN(executor_fn, TpuTopology_IdForHost); + TFTPU_SET_FN(executor_fn, TpuTopology_Version); + + TFTPU_SET_FN(executor_fn, TpuCoreLocation_ChipCoordinates); + TFTPU_SET_FN(executor_fn, TpuCoreLocation_HostCoordinates); + TFTPU_SET_FN(executor_fn, TpuCoreLocation_Index); + TFTPU_SET_FN(executor_fn, TpuCoreLocation_Id); + + TFTPU_SET_FN(executor_fn, TpuHostLocation_Id); + + TFTPU_SET_FN(executor_fn, TpuCompiler_New); + TFTPU_SET_FN(executor_fn, TpuCompiler_Free); + + TFTPU_SET_FN(executor_fn, TpuCompiler_RunHloPasses); + TFTPU_SET_FN(executor_fn, TpuCompiler_RunBackend); + TFTPU_SET_FN(executor_fn, TpuCompiler_Compile); + TFTPU_SET_FN(executor_fn, TpuCompiler_ShapeSize); + TFTPU_SET_FN(executor_fn, TpuExecutable_ExecuteAsyncOnStream); + TFTPU_SET_FN(executor_fn, TpuExecutable_Free); + + TFTPU_SET_FN(executor_fn, XlaShapeToTpuShapeRepresentation); + TFTPU_SET_FN(executor_fn, XlaShapeToTpuPaddedShape); + + return tensorflow::Status::OK(); +} + +} // namespace diff --git a/tensorflow/core/tpu/tpu_library_init_fns.inc b/tensorflow/core/tpu/tpu_library_init_fns.inc index cc4b62a2f11..16494d0aa86 100644 --- a/tensorflow/core/tpu/tpu_library_init_fns.inc +++ b/tensorflow/core/tpu/tpu_library_init_fns.inc @@ -1,3 +1,5 @@ +#include "third_party/tensorflow/core/tpu/tpu_executor_init_fns.inc" + namespace { tensorflow::Status SetTpuConfigStructFns(void* library_handle) { @@ -70,144 +72,6 @@ tensorflow::Status SetTpuProgramStructFn(void* library_handle) { return tensorflow::Status::OK(); } -tensorflow::Status SetExecutorStructFn(void* library_handle) { - auto* executor_fn = tensorflow::tpu::ExecutorApiFn(); - - TFTPU_SET_FN(executor_fn, TpuPlatform_New); - TFTPU_SET_FN(executor_fn, TpuPlatform_Free); - TFTPU_SET_FN(executor_fn, TpuPlatform_Initialize); - TFTPU_SET_FN(executor_fn, TpuPlatform_Initialized); - TFTPU_SET_FN(executor_fn, TpuPlatform_GetExecutor); - TFTPU_SET_FN(executor_fn, TpuPlatform_Id); - TFTPU_SET_FN(executor_fn, TpuPlatform_VisibleDeviceCount); - TFTPU_SET_FN(executor_fn, TpuPlatform_TpuMemoryLimit); - TFTPU_SET_FN(executor_fn, TpuPlatform_ShouldRegisterTpuDeviceToDeviceCopy); - TFTPU_SET_FN(executor_fn, TpuPlatform_GetTopologyPtr); - TFTPU_SET_FN(executor_fn, TpuPlatform_GetHostLocation); - - TFTPU_SET_FN(executor_fn, TpuExecutor_Init); - TFTPU_SET_FN(executor_fn, TpuExecutor_Free); - TFTPU_SET_FN(executor_fn, TpuExecutor_PlatformDeviceCount); - TFTPU_SET_FN(executor_fn, TpuExecutor_Allocate); - TFTPU_SET_FN(executor_fn, TpuExecutor_Deallocate); - TFTPU_SET_FN(executor_fn, TpuExecutor_GetAllocatorStats); - TFTPU_SET_FN(executor_fn, TpuExecutor_DeviceMemoryUsage); - TFTPU_SET_FN(executor_fn, TpuExecutor_AllocateStream); - TFTPU_SET_FN(executor_fn, TpuExecutor_DeallocateStream); - TFTPU_SET_FN(executor_fn, TpuExecutor_CreateStreamDependency); - TFTPU_SET_FN(executor_fn, TpuExecutor_GetStatus); - TFTPU_SET_FN(executor_fn, TpuExecutor_GetCoreLocation); - TFTPU_SET_FN(executor_fn, TpuExecutor_AllocateEvent); - TFTPU_SET_FN(executor_fn, TpuExecutor_DeallocateEvent); - TFTPU_SET_FN(executor_fn, TpuExecutor_PollForEventStatus); - TFTPU_SET_FN(executor_fn, TpuExecutor_RecordEvent); - TFTPU_SET_FN(executor_fn, TpuExecutor_WaitForEvent); - TFTPU_SET_FN(executor_fn, TpuExecutor_AllocateTimer); - TFTPU_SET_FN(executor_fn, TpuExecutor_DeallocateTimer); - TFTPU_SET_FN(executor_fn, TpuExecutor_StartTimer); - TFTPU_SET_FN(executor_fn, TpuExecutor_StopTimer); - TFTPU_SET_FN(executor_fn, TpuExecutor_SynchronousMemcpyToHost); - TFTPU_SET_FN(executor_fn, TpuExecutor_SynchronousMemcpyFromHost); - TFTPU_SET_FN(executor_fn, TpuExecutor_MemcpyToHost); - TFTPU_SET_FN(executor_fn, TpuExecutor_MemcpyFromHost); - TFTPU_SET_FN(executor_fn, TpuExecutor_EnqueueInfeed); - TFTPU_SET_FN(executor_fn, TpuExecutor_DequeueOutfeed); - TFTPU_SET_FN(executor_fn, TpuExecutor_WaitForInfeedReady); - TFTPU_SET_FN(executor_fn, TpuExecutor_WaitForOutfeedReady); - TFTPU_SET_FN(executor_fn, TpuExecutor_BlockHostUntilDone); - TFTPU_SET_FN(executor_fn, TpuExecutor_BlockUntilDoneOrFailed); - TFTPU_SET_FN(executor_fn, TpuExecutor_SyncAndForgetFailedStreams); - TFTPU_SET_FN(executor_fn, TpuExecutor_SynchronizeAllActivity); - - TFTPU_SET_FN(executor_fn, TpuStream_New); - TFTPU_SET_FN(executor_fn, TpuStream_Free); - TFTPU_SET_FN(executor_fn, TpuStream_Stream); - TFTPU_SET_FN(executor_fn, TpuStream_Status); - TFTPU_SET_FN(executor_fn, TpuStream_IsSameSharedMemoryLocation); - TFTPU_SET_FN(executor_fn, TpuStream_TpuEnqueueOnDeviceSendRecvLocal); - - TFTPU_SET_FN(executor_fn, TpuEvent_New); - TFTPU_SET_FN(executor_fn, TpuEvent_Free); - - TFTPU_SET_FN(executor_fn, TpuTimer_New); - TFTPU_SET_FN(executor_fn, TpuTimer_Free); - TFTPU_SET_FN(executor_fn, TpuTimer_Nanoseconds); - TFTPU_SET_FN(executor_fn, TpuTimer_Microseconds); - - TFTPU_SET_FN(executor_fn, TpuStatus_New); - TFTPU_SET_FN(executor_fn, TpuStatus_Create); - TFTPU_SET_FN(executor_fn, TpuStatus_Set); - TFTPU_SET_FN(executor_fn, TpuStatus_Free); - TFTPU_SET_FN(executor_fn, TpuStatus_Message); - TFTPU_SET_FN(executor_fn, TpuStatus_Code); - TFTPU_SET_FN(executor_fn, TpuStatus_Ok); - - TFTPU_SET_FN(executor_fn, TpuStreamExecutorConfig_Default); - TFTPU_SET_FN(executor_fn, TpuStreamExecutorConfig_SetOrdinal); - TFTPU_SET_FN(executor_fn, TpuStreamExecutorConfig_Free); - - TFTPU_SET_FN(executor_fn, TpuDeviceDescription_New); - TFTPU_SET_FN(executor_fn, TpuDeviceDescription_Free); - - TFTPU_SET_FN(executor_fn, TpuExecutor_CreateDeviceDescription); - TFTPU_SET_FN(executor_fn, TpuExecutor_NewDeviceOptions); - TFTPU_SET_FN(executor_fn, TpuExecutor_FreeDeviceOptions); - TFTPU_SET_FN(executor_fn, TpuExecutor_HostCallback); - - TFTPU_SET_FN(executor_fn, TpuTransferManager_New); - TFTPU_SET_FN(executor_fn, TpuTransferManager_Free); - TFTPU_SET_FN(executor_fn, TpuTransferManager_PlatformId); - TFTPU_SET_FN(executor_fn, TpuTransferManager_HostShapeToDeviceShape); - TFTPU_SET_FN(executor_fn, TpuTransferManager_TransferLiteralToDeviceAsync); - TFTPU_SET_FN(executor_fn, TpuTransferManager_TransferLiteralFromDevice); - TFTPU_SET_FN(executor_fn, TpuTransferManager_GetByteSizeRequirement); - TFTPU_SET_FN(executor_fn, TpuTransferManager_WriteSingleTupleIndexTable); - TFTPU_SET_FN(executor_fn, TpuTransferManager_GetInfeedLayout); - TFTPU_SET_FN(executor_fn, TpuTransferManager_LinearizeToBuffers); - TFTPU_SET_FN(executor_fn, TpuTransferManager_FreeBuffers); - TFTPU_SET_FN(executor_fn, TpuTransferManager_TransferLiteralToInfeed); - TFTPU_SET_FN(executor_fn, TpuTransferManager_TransferBuffersToInfeed); - TFTPU_SET_FN(executor_fn, TpuTransferManager_TransferLiteralFromOutfeed); - TFTPU_SET_FN(executor_fn, TpuTransferManager_ResetDevices); - - TFTPU_SET_FN(executor_fn, TpuComputationPlacer_New); - TFTPU_SET_FN(executor_fn, TpuComputationPlacer_Free); - - TFTPU_SET_FN(executor_fn, TpuTopology_LogicalDevicesPerHost); - TFTPU_SET_FN(executor_fn, TpuTopology_LogicalDevicesPerChip); - TFTPU_SET_FN(executor_fn, TpuTopology_ChipBounds_X); - TFTPU_SET_FN(executor_fn, TpuTopology_ChipBounds_Y); - TFTPU_SET_FN(executor_fn, TpuTopology_ChipBounds_Z); - TFTPU_SET_FN(executor_fn, TpuTopology_HasChip); - TFTPU_SET_FN(executor_fn, TpuTopology_Core); - TFTPU_SET_FN(executor_fn, TpuTopology_NumCores); - TFTPU_SET_FN(executor_fn, TpuTopology_Cores); - TFTPU_SET_FN(executor_fn, TpuTopology_IdForHost); - TFTPU_SET_FN(executor_fn, TpuTopology_Version); - - TFTPU_SET_FN(executor_fn, TpuCoreLocation_ChipCoordinates); - TFTPU_SET_FN(executor_fn, TpuCoreLocation_HostCoordinates); - TFTPU_SET_FN(executor_fn, TpuCoreLocation_Index); - TFTPU_SET_FN(executor_fn, TpuCoreLocation_Id); - - TFTPU_SET_FN(executor_fn, TpuHostLocation_Id); - - TFTPU_SET_FN(executor_fn, TpuCompiler_New); - TFTPU_SET_FN(executor_fn, TpuCompiler_Free); - - TFTPU_SET_FN(executor_fn, TpuCompiler_RunHloPasses); - TFTPU_SET_FN(executor_fn, TpuCompiler_RunBackend); - TFTPU_SET_FN(executor_fn, TpuCompiler_Compile); - TFTPU_SET_FN(executor_fn, TpuCompiler_ShapeSize); - TFTPU_SET_FN(executor_fn, TpuExecutable_ExecuteAsyncOnStream); - TFTPU_SET_FN(executor_fn, TpuExecutable_Free); - - TFTPU_SET_FN(executor_fn, XlaShapeToTpuShapeRepresentation); - TFTPU_SET_FN(executor_fn, XlaShapeToTpuPaddedShape); - - return tensorflow::Status::OK(); -} - tensorflow::Status SetTpuNodeContextStructFns(void* library_handle) { auto* node_context_fn = tensorflow::tpu::NodeContextApiFn();