Skip to content

Commit 8bc0c84

Browse files
skyetensorflower-gardener
authored andcommitted
[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
1 parent d0e6a3b commit 8bc0c84

File tree

5 files changed

+25
-20
lines changed

5 files changed

+25
-20
lines changed

tensorflow/core/tpu/libtftpu.h

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,8 +41,8 @@ limitations under the License.
4141
extern "C" {
4242
#endif
4343

44-
TFTPU_CAPI_EXPORT void TfTpu_Initialize(bool init_library, int argc,
45-
const char** argv);
44+
TFTPU_CAPI_EXPORT void TfTpu_Initialize(bool init_library, int num_args,
45+
const char** args);
4646

4747
#ifdef __cplusplus
4848
}

tensorflow/core/tpu/tpu_api_dlsym_initializer.cc

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -45,17 +45,18 @@ Status InitializeTpuLibrary(void* library_handle) {
4545
Status s = InitializeTpuStructFns(library_handle);
4646

4747
// Retrieve arguments from environment if applicable
48-
std::vector<const char*> argv_ptr = GetLibTpuInitArguments();
48+
std::pair<std::vector<std::string>, std::vector<const char*> > args =
49+
GetLibTpuInitArguments();
4950

5051
// TPU platform registration must only be performed after the library is
5152
// loaded. We do not want to register a TPU platform in XLA without the
5253
// supporting library providing the necessary APIs.
5354
if (s.ok()) {
54-
void (*initialize_fn)(bool init_library, int argc, const char** argv);
55+
void (*initialize_fn)(bool init_library, int num_args, const char** args);
5556
initialize_fn = reinterpret_cast<decltype(initialize_fn)>(
5657
dlsym(library_handle, "TfTpu_Initialize"));
57-
(*initialize_fn)(/*init_library=*/true, /*argc=*/argv_ptr.size() - 1,
58-
/*argv=*/argv_ptr.data());
58+
(*initialize_fn)(/*init_library=*/true, args.second.size(),
59+
args.second.data());
5960

6061
RegisterTpuPlatform();
6162
}

tensorflow/core/tpu/tpu_executor_dlsym_initializer.cc

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -42,17 +42,18 @@ Status InitializeTpuLibrary(void* library_handle) {
4242
Status s = SetExecutorStructFn(library_handle);
4343

4444
// Retrieve arguments from environment if applicable
45-
std::vector<const char*> argv_ptr = GetLibTpuInitArguments();
45+
std::pair<std::vector<std::string>, std::vector<const char*> > args =
46+
GetLibTpuInitArguments();
4647

4748
// TPU platform registration must only be performed after the library is
4849
// loaded. We do not want to register a TPU platform in XLA without the
4950
// supporting library providing the necessary APIs.
5051
if (s.ok()) {
51-
void (*initialize_fn)(bool init_library, int argc, const char** argv);
52+
void (*initialize_fn)(bool init_library, int num_args, const char** args);
5253
initialize_fn = reinterpret_cast<decltype(initialize_fn)>(
5354
dlsym(library_handle, "TfTpu_Initialize"));
54-
(*initialize_fn)(/*init_library=*/true, /*argc=*/argv_ptr.size() - 1,
55-
/*argv=*/argv_ptr.data());
55+
(*initialize_fn)(/*init_library=*/true, args.second.size(),
56+
args.second.data());
5657

5758
RegisterTpuPlatform();
5859
}

tensorflow/core/tpu/tpu_initializer_helper.cc

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -22,17 +22,17 @@ limitations under the License.
2222
namespace tensorflow {
2323
namespace tpu {
2424

25-
std::vector<const char*> GetLibTpuInitArguments() {
25+
std::pair<std::vector<std::string>, std::vector<const char*>>
26+
GetLibTpuInitArguments() {
27+
// We make copies of the arguments returned by getenv because the memory
28+
// returned may be altered or invalidated by further calls to getenv.
29+
std::vector<std::string> argv;
2630
std::vector<const char*> argv_ptr;
27-
std::vector<absl::string_view> argv;
2831

29-
// Retrieve arguments from environment if applicable
32+
// Retrieve arguments from environment if applicable.
3033
char* env = getenv("LIBTPU_INIT_ARGS");
3134
if (env != nullptr) {
3235
// TODO(frankchn): Handles quotes properly if necessary.
33-
// env pointer is already pointing to an allocated memory block.
34-
// absl::StrSplit returns a string_view that returns a vector of pointers
35-
// into that memory block. This means that we don't need to manage memory.
3636
argv = absl::StrSplit(env, ' ');
3737
}
3838

@@ -42,7 +42,7 @@ std::vector<const char*> GetLibTpuInitArguments() {
4242
}
4343
argv_ptr.push_back(nullptr);
4444

45-
return argv_ptr;
45+
return {argv, argv_ptr};
4646
}
4747

4848
} // namespace tpu

tensorflow/core/tpu/tpu_initializer_helper.h

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,14 +16,17 @@ limitations under the License.
1616
#ifndef TENSORFLOW_CORE_TPU_TPU_INITIALIZER_HELPER_H_
1717
#define TENSORFLOW_CORE_TPU_TPU_INITIALIZER_HELPER_H_
1818

19+
#include <string>
1920
#include <vector>
2021

2122
namespace tensorflow {
2223
namespace tpu {
2324

24-
// This returns an extra nullptr at the end (per the C standard), but this
25-
// should not be counted for 'argc'.
26-
std::vector<const char*> GetLibTpuInitArguments();
25+
// Returns arguments (e.g. flags) set in the LIBTPU_INIT_ARGS environment
26+
// variable. The first return value is the arguments, the second return value is
27+
// pointers to the arguments suitable for passing into the C API.
28+
std::pair<std::vector<std::string>, std::vector<const char*>>
29+
GetLibTpuInitArguments();
2730

2831
} // namespace tpu
2932
} // namespace tensorflow

0 commit comments

Comments
 (0)