From 8bc0c84e45f1e172bfa81e9cc5fd4302a7d362e9 Mon Sep 17 00:00:00 2001 From: Skye Wanderman-Milne Date: Thu, 28 Jan 2021 14:35:03 -0800 Subject: [PATCH] [libtpu] Make copies of data returned by getenv. According to http://www.cplusplus.com/reference/cstdlib/getenv/: The pointer returned points to an internal memory block, whose content or validity may be altered by further calls to getenv (but not by other library functions). This appears to have been working without the defensive copies, but we should follow the spec just in case. This change also refactors TfTpu_Initialize to not use the argc/argv convention for simplicity (the implementation doesn't actually require a nullptr at the end of argv). PiperOrigin-RevId: 354396268 Change-Id: Ie14f13d4a67e1012dc13af63df6a9a4fcb22feeb --- tensorflow/core/tpu/libtftpu.h | 4 ++-- tensorflow/core/tpu/tpu_api_dlsym_initializer.cc | 9 +++++---- .../core/tpu/tpu_executor_dlsym_initializer.cc | 9 +++++---- tensorflow/core/tpu/tpu_initializer_helper.cc | 14 +++++++------- tensorflow/core/tpu/tpu_initializer_helper.h | 9 ++++++--- 5 files changed, 25 insertions(+), 20 deletions(-) diff --git a/tensorflow/core/tpu/libtftpu.h b/tensorflow/core/tpu/libtftpu.h index 6c5cd3c08b7..921409dc522 100644 --- a/tensorflow/core/tpu/libtftpu.h +++ b/tensorflow/core/tpu/libtftpu.h @@ -41,8 +41,8 @@ limitations under the License. extern "C" { #endif -TFTPU_CAPI_EXPORT void TfTpu_Initialize(bool init_library, int argc, - const char** argv); +TFTPU_CAPI_EXPORT void TfTpu_Initialize(bool init_library, int num_args, + const char** args); #ifdef __cplusplus } diff --git a/tensorflow/core/tpu/tpu_api_dlsym_initializer.cc b/tensorflow/core/tpu/tpu_api_dlsym_initializer.cc index 9389528934b..4c67d596809 100644 --- a/tensorflow/core/tpu/tpu_api_dlsym_initializer.cc +++ b/tensorflow/core/tpu/tpu_api_dlsym_initializer.cc @@ -45,17 +45,18 @@ Status InitializeTpuLibrary(void* library_handle) { Status s = InitializeTpuStructFns(library_handle); // Retrieve arguments from environment if applicable - std::vector argv_ptr = GetLibTpuInitArguments(); + std::pair, std::vector > args = + GetLibTpuInitArguments(); // TPU platform registration must only be performed after the library is // loaded. We do not want to register a TPU platform in XLA without the // supporting library providing the necessary APIs. if (s.ok()) { - void (*initialize_fn)(bool init_library, int argc, const char** argv); + void (*initialize_fn)(bool init_library, int num_args, const char** args); initialize_fn = reinterpret_cast( dlsym(library_handle, "TfTpu_Initialize")); - (*initialize_fn)(/*init_library=*/true, /*argc=*/argv_ptr.size() - 1, - /*argv=*/argv_ptr.data()); + (*initialize_fn)(/*init_library=*/true, args.second.size(), + args.second.data()); RegisterTpuPlatform(); } diff --git a/tensorflow/core/tpu/tpu_executor_dlsym_initializer.cc b/tensorflow/core/tpu/tpu_executor_dlsym_initializer.cc index 5da9a33fdc0..8c2ae85733f 100644 --- a/tensorflow/core/tpu/tpu_executor_dlsym_initializer.cc +++ b/tensorflow/core/tpu/tpu_executor_dlsym_initializer.cc @@ -42,17 +42,18 @@ Status InitializeTpuLibrary(void* library_handle) { Status s = SetExecutorStructFn(library_handle); // Retrieve arguments from environment if applicable - std::vector argv_ptr = GetLibTpuInitArguments(); + std::pair, std::vector > args = + GetLibTpuInitArguments(); // TPU platform registration must only be performed after the library is // loaded. We do not want to register a TPU platform in XLA without the // supporting library providing the necessary APIs. if (s.ok()) { - void (*initialize_fn)(bool init_library, int argc, const char** argv); + void (*initialize_fn)(bool init_library, int num_args, const char** args); initialize_fn = reinterpret_cast( dlsym(library_handle, "TfTpu_Initialize")); - (*initialize_fn)(/*init_library=*/true, /*argc=*/argv_ptr.size() - 1, - /*argv=*/argv_ptr.data()); + (*initialize_fn)(/*init_library=*/true, args.second.size(), + args.second.data()); RegisterTpuPlatform(); } diff --git a/tensorflow/core/tpu/tpu_initializer_helper.cc b/tensorflow/core/tpu/tpu_initializer_helper.cc index 8648fd8ab28..0856957d39c 100644 --- a/tensorflow/core/tpu/tpu_initializer_helper.cc +++ b/tensorflow/core/tpu/tpu_initializer_helper.cc @@ -22,17 +22,17 @@ limitations under the License. namespace tensorflow { namespace tpu { -std::vector GetLibTpuInitArguments() { +std::pair, std::vector> +GetLibTpuInitArguments() { + // We make copies of the arguments returned by getenv because the memory + // returned may be altered or invalidated by further calls to getenv. + std::vector argv; std::vector argv_ptr; - std::vector argv; - // Retrieve arguments from environment if applicable + // Retrieve arguments from environment if applicable. char* env = getenv("LIBTPU_INIT_ARGS"); if (env != nullptr) { // TODO(frankchn): Handles quotes properly if necessary. - // env pointer is already pointing to an allocated memory block. - // absl::StrSplit returns a string_view that returns a vector of pointers - // into that memory block. This means that we don't need to manage memory. argv = absl::StrSplit(env, ' '); } @@ -42,7 +42,7 @@ std::vector GetLibTpuInitArguments() { } argv_ptr.push_back(nullptr); - return argv_ptr; + return {argv, argv_ptr}; } } // namespace tpu diff --git a/tensorflow/core/tpu/tpu_initializer_helper.h b/tensorflow/core/tpu/tpu_initializer_helper.h index 8d3ec4e324b..cd9b419648d 100644 --- a/tensorflow/core/tpu/tpu_initializer_helper.h +++ b/tensorflow/core/tpu/tpu_initializer_helper.h @@ -16,14 +16,17 @@ limitations under the License. #ifndef TENSORFLOW_CORE_TPU_TPU_INITIALIZER_HELPER_H_ #define TENSORFLOW_CORE_TPU_TPU_INITIALIZER_HELPER_H_ +#include #include namespace tensorflow { namespace tpu { -// This returns an extra nullptr at the end (per the C standard), but this -// should not be counted for 'argc'. -std::vector GetLibTpuInitArguments(); +// Returns arguments (e.g. flags) set in the LIBTPU_INIT_ARGS environment +// variable. The first return value is the arguments, the second return value is +// pointers to the arguments suitable for passing into the C API. +std::pair, std::vector> +GetLibTpuInitArguments(); } // namespace tpu } // namespace tensorflow