Separate out the TPU executor API from the rest to minimize build dependencies.
This is useful for downstream clients like JAX that don't need all the TF-specific APIs. PiperOrigin-RevId: 327136708 Change-Id: I207068d0f128334673f2270b7c9a12f06cc9ec24
This commit is contained in:
Skye Wanderman-Milne
TensorFlower Gardener
parent
45bbbbf96b
commit
9b0c592c95
@ -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"],
|
||||
|
||||
@ -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;
|
||||
|
||||
@ -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();
|
||||
|
||||
27
tensorflow/core/tpu/tpu_executor_api.cc
Normal file
27
tensorflow/core/tpu/tpu_executor_api.cc
Normal file
@ -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
|
||||
30
tensorflow/core/tpu/tpu_executor_api.h
Normal file
30
tensorflow/core/tpu/tpu_executor_api.h
Normal file
@ -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_
|
||||
141
tensorflow/core/tpu/tpu_executor_init_fns.inc
Normal file
141
tensorflow/core/tpu/tpu_executor_init_fns.inc
Normal file
@ -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
|
||||
@ -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();
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user