experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Expose TPU model server initialization.

Browse Source 
PiperOrigin-RevId: 356631873
Change-Id: I488e73e403a8a761ea1e652bfe7e337033462514
This commit is contained in:
Allen Wang 2021-02-09 17:37:59 -08:00 committed by TensorFlower Gardener
parent 4c01f5d82f
commit 52eb80d341
2 changed files with 7 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tensorflow/core/tpu/tpu_library_init_fns.inc
View File

@ -79,6 +79,8 @@ tensorflow::Status SetTpuOpsStructFns(void* library_handle) {
TFTPU_SET_FN(ops_api_fn, TpuProfiler_Stop);
TFTPU_SET_FN(ops_api_fn, TpuProfiler_CollectData);
TFTPU_SET_FN(ops_api_fn, TfTpu_InitializeTpuModelServer);
return tensorflow::Status::OK();
}

5
tensorflow/core/tpu/tpu_ops_c_api.h
View File

@ -426,6 +426,9 @@ void TpuNodeContext_CloseTpuHost(TF_Status* status);
void TpuNodeContext_Initialize(int device_ordinal, TF_Status* status);
// Globally initialize the TPU system for inference.
TFTPU_CAPI_EXPORT void TfTpu_InitializeTpuModelServer();
struct TfTpu_OpsApiFn {
TFTPU_ADD_FN_IN_STRUCT(TpuCompile_CompileAndBuild);
TFTPU_ADD_FN_IN_STRUCT(TpuCompile_XrtCompileAndBuild);
@ -492,6 +495,8 @@ struct TfTpu_OpsApiFn {
TFTPU_ADD_FN_IN_STRUCT(TpuNodeContext_StopChipHeartbeats);
TFTPU_ADD_FN_IN_STRUCT(TpuNodeContext_CloseTpuHost);
TFTPU_ADD_FN_IN_STRUCT(TpuNodeContext_Initialize);
TFTPU_ADD_FN_IN_STRUCT(TfTpu_InitializeTpuModelServer);
};
} // extern "C"