File tree Expand file tree Collapse file tree 5 files changed +25
-20
lines changed Expand file tree Collapse file tree 5 files changed +25
-20
lines changed Original file line number Diff line number Diff line change @@ -41,8 +41,8 @@ limitations under the License.
41
41
extern "C" {
42
42
#endif
43
43
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 );
46
46
47
47
#ifdef __cplusplus
48
48
}
Original file line number Diff line number Diff line change @@ -45,17 +45,18 @@ Status InitializeTpuLibrary(void* library_handle) {
45
45
Status s = InitializeTpuStructFns (library_handle);
46
46
47
47
// 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 ();
49
50
50
51
// TPU platform registration must only be performed after the library is
51
52
// loaded. We do not want to register a TPU platform in XLA without the
52
53
// supporting library providing the necessary APIs.
53
54
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 );
55
56
initialize_fn = reinterpret_cast <decltype (initialize_fn)>(
56
57
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 ());
59
60
60
61
RegisterTpuPlatform ();
61
62
}
Original file line number Diff line number Diff line change @@ -42,17 +42,18 @@ Status InitializeTpuLibrary(void* library_handle) {
42
42
Status s = SetExecutorStructFn (library_handle);
43
43
44
44
// 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 ();
46
47
47
48
// TPU platform registration must only be performed after the library is
48
49
// loaded. We do not want to register a TPU platform in XLA without the
49
50
// supporting library providing the necessary APIs.
50
51
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 );
52
53
initialize_fn = reinterpret_cast <decltype (initialize_fn)>(
53
54
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 ());
56
57
57
58
RegisterTpuPlatform ();
58
59
}
Original file line number Diff line number Diff line change @@ -22,17 +22,17 @@ limitations under the License.
22
22
namespace tensorflow {
23
23
namespace tpu {
24
24
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;
26
30
std::vector<const char *> argv_ptr;
27
- std::vector<absl::string_view> argv;
28
31
29
- // Retrieve arguments from environment if applicable
32
+ // Retrieve arguments from environment if applicable.
30
33
char * env = getenv (" LIBTPU_INIT_ARGS" );
31
34
if (env != nullptr ) {
32
35
// 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.
36
36
argv = absl::StrSplit (env, ' ' );
37
37
}
38
38
@@ -42,7 +42,7 @@ std::vector<const char*> GetLibTpuInitArguments() {
42
42
}
43
43
argv_ptr.push_back (nullptr );
44
44
45
- return argv_ptr;
45
+ return {argv, argv_ptr} ;
46
46
}
47
47
48
48
} // namespace tpu
Original file line number Diff line number Diff line change @@ -16,14 +16,17 @@ limitations under the License.
16
16
#ifndef TENSORFLOW_CORE_TPU_TPU_INITIALIZER_HELPER_H_
17
17
#define TENSORFLOW_CORE_TPU_TPU_INITIALIZER_HELPER_H_
18
18
19
+ #include < string>
19
20
#include < vector>
20
21
21
22
namespace tensorflow {
22
23
namespace tpu {
23
24
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 ();
27
30
28
31
} // namespace tpu
29
32
} // namespace tensorflow
You can’t perform that action at this time.
0 commit comments