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

Add argc and argv arguments to pass to tpu library on startup

Browse Source 
PiperOrigin-RevId: 345777151
Change-Id: I150fac0e0a9559c3968509c2b8e011b20620156b
This commit is contained in:
Frank Chen 2020-12-04 16:11:44 -08:00 committed by TensorFlower Gardener
parent cd39d7a4c8
commit 304b968153
3 changed files with 5 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tensorflow/core/tpu/libtftpu.h
View File

@ -41,7 +41,8 @@ limitations under the License.
extern "C" {
#endif
TFTPU_CAPI_EXPORT void TfTpu_Initialize(bool init_library);
TFTPU_CAPI_EXPORT void TfTpu_Initialize(bool init_library, int argc,
char** argv);
#ifdef __cplusplus
}

2
tensorflow/core/tpu/tpu_api_dlsym_initializer.cc
View File

@ -52,7 +52,7 @@ Status InitializeTpuLibrary(void* library_handle) {
void (*initialize_fn)(bool init_library);
initialize_fn = reinterpret_cast<decltype(initialize_fn)>(
dlsym(library_handle, "TfTpu_Initialize"));
(*initialize_fn)(/*init_library=*/true);
(*initialize_fn)(/*init_library=*/true, /*argc=*/0, /*argv=*/nullptr);
RegisterTpuPlatform();
RegisterTpuSystemDevice();

4
tensorflow/core/tpu/tpu_executor_dlsym_initializer.cc
View File

@ -45,8 +45,8 @@ Status InitializeTpuLibrary(void* library_handle) {
// supporting library providing the necessary APIs.
if (s.ok()) {
void (*initialize_fn)();
initialize_fn = reinterpret_cast<decltype(initialize_fn)>(
dlsym(library_handle, "TfTpu_Initialize"));
initialize_fn = reinterpret_cast<decltype(initialize_fn)>(dlsym(
library_handle, "TfTpu_Initialize", /*argc=*/0, /*argv=*/nullptr));
(*initialize_fn)();
RegisterTpuPlatform();