1/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15
16// This module exports ParseFlagsFromEnv(), which allows other modules to parse
17// flags from an environtment variable, or a file named by the environment
18// variable.
19
20#include <stdio.h>
21#include <stdlib.h>
22#include <string.h>
23#include <vector>
24
25#include "tensorflow/compiler/xla/legacy_flags/parse_flags_from_env.h"
26#include "tensorflow/compiler/xla/types.h"
27#include "tensorflow/core/platform/logging.h"
28#include "tensorflow/core/platform/macros.h"
29#include "tensorflow/core/platform/mutex.h"
30#include "tensorflow/core/platform/types.h"
31#include "tensorflow/core/util/command_line_flags.h"
32
33namespace xla {
34namespace legacy_flags {
35
36static const char kEnvVar[] = "TF_XLA_FLAGS"; // environment variable queried
37static const char kWS[] = " \t\r\n"; // whitespace
38
39// The following struct represents an argv[]-style array, parsed
40// from data gleaned from the environment.
41//
42// As usual, an anonymous namespace is advisable to avoid
43// constructor/destructor collisions with other "private" types
44// in the same named namespace.
45namespace {
46struct EnvArgv {
47 EnvArgv() : initialized(false), argc(0) {}
48 bool initialized; // whether the other fields have been set.
49 int argc; // elements used in argv[]
50 std::vector<char*> argv; // flag arguments parsed from environment string.
51 std::vector<char*> argv_save; // saved values from argv[] to avoid leaks
52};
53} // anonymous namespace
54
55// Append the string s0[0, .., s0len-1] concatenated with s1[0, .., s1len-1] as
56// a newly allocated nul-terminated string to the array *a. If s0==nullptr, a
57// nullptr is appended without increasing a->argc.
58static void AppendToEnvArgv(const char* s0, size_t s0len, const char* s1,
59 size_t s1len, EnvArgv* a) {
60 if (s0 == nullptr) {
61 a->argv.push_back(nullptr);
62 a->argv_save.push_back(nullptr);
63 } else {
64 string s = string(s0, s0len) + string(s1, s1len);
65 char* str = strdup(s.c_str());
66 a->argv.push_back(str);
67 a->argv_save.push_back(str);
68 a->argc++;
69 }
70}
71
72// Like s.find_first_of(x, pos), but return s.size() when find_first_of() would
73// return string::npos. This avoids if-statements elsewhere.
74static size_t FindFirstOf(const string& s, const char* x, size_t pos) {
75 size_t result = s.find_first_of(x, pos);
76 return result == string::npos ? s.size() : result;
77}
78
79// Like s.find_first_not_of(x, pos), but return s.size() when
80// find_first_not_of() would return string::npos. This avoids if-statements
81// elsewhere.
82static size_t FindFirstNotOf(const string& s, const char* x, size_t pos) {
83 size_t result = s.find_first_not_of(x, pos);
84 return result == string::npos ? s.size() : result;
85}
86
87// Given a string containing flags, parse them into the XLA command line flags.
88// The parse is best effort, and gives up on the first syntax error.
89static void ParseArgvFromString(const string& flag_str, EnvArgv* a) {
90 size_t b = FindFirstNotOf(flag_str, kWS, 0);
91 while (b != flag_str.size() && flag_str[b] == '-') {
92 // b is the index of the start of a flag.
93 // Set e to the index just past the end of the flag.
94 size_t e = b;
95 while (e != flag_str.size() && isascii(flag_str[e]) &&
96 (strchr("-_", flag_str[e]) != nullptr || isalnum(flag_str[e]))) {
97 e++;
98 }
99 if (e != flag_str.size() && flag_str[e] == '=' &&
100 e + 1 != flag_str.size() && strchr("'\"", flag_str[e + 1]) != nullptr) {
101 // A flag of the form --flag="something in double or single quotes"
102 int c;
103 e++; // point just past '='
104 size_t eflag = e;
105 char quote = flag_str[e];
106 e++; // point just past quote
107 // Put in value the string with quotes removed.
108 string value;
109 for (; e != flag_str.size() && (c = flag_str[e]) != quote; e++) {
110 if (quote == '"' && c == '\\' && e + 1 != flag_str.size()) {
111 // Handle backslash in double quoted strings. They are literal in
112 // single-quoted strings.
113 e++;
114 c = flag_str[e];
115 }
116 value += c;
117 }
118 if (e != flag_str.size()) { // skip final " or '
119 e++;
120 }
121 AppendToEnvArgv(flag_str.data() + b, eflag - b, value.data(),
122 value.size(), a);
123 } else { // A flag without a quoted value.
124 e = FindFirstOf(flag_str, kWS, e);
125 AppendToEnvArgv(flag_str.data() + b, e - b, "", 0, a);
126 }
127 b = FindFirstNotOf(flag_str, kWS, e);
128 }
129}
130
131// Call ParseArgvFromString(..., a) on a string derived from the setting of an
132// environment variable kEnvVar, or a file it points to.
133static void SetArgvFromEnv(EnvArgv* a) {
134 if (!a->initialized) {
135 static const char kDummyArgv[] = "<argv[0]>";
136 AppendToEnvArgv(kDummyArgv, strlen(kDummyArgv), nullptr, 0,
137 a); // dummy argv[0]
138 const char* env = getenv(kEnvVar);
139 if (env == nullptr || env[0] == '\0') {
140 // nothing
141 } else if (env[strspn(env, kWS)] == '-') { // flags in env var value
142 ParseArgvFromString(env, a);
143 } else { // assume it's a file name
144 FILE* fp = fopen(env, "r");
145 if (fp != nullptr) {
146 string str;
147 char buf[512];
148 int n;
149 while ((n = fread(buf, 1, sizeof(buf), fp)) > 0) {
150 str.append(buf, n);
151 }
152 fclose(fp);
153 ParseArgvFromString(str, a);
154 }
155 }
156 AppendToEnvArgv(nullptr, 0, nullptr, 0, a); // add trailing nullptr to *a.
157 a->initialized = true;
158 }
159}
160
161// The simulated argv[] parsed from the environment.
162static EnvArgv* env_argv;
163
164// Used to protect accesses to env_argv.
165static tensorflow::mutex env_argv_mu(tensorflow::LINKER_INITIALIZED);
166
167// Call Flags::Parse(argc, argv, flag_list) against any as yet unrecognized
168// flags passed in from the environment.
169bool ParseFlagsFromEnv(const std::vector<tensorflow::Flag>& flag_list) {
170 env_argv_mu.lock();
171 if (env_argv == nullptr) {
172 env_argv = new EnvArgv;
173 }
174 SetArgvFromEnv(env_argv); // a no-op if already initialized
175 bool result =
176 tensorflow::Flags::Parse(&env_argv->argc, &env_argv->argv[0], flag_list);
177 env_argv_mu.unlock();
178 return result;
179}
180
181// Testing only.
182// Reset the env_argv struct so that subsequent calls to ParseFlagsFromEnv()
183// will parse the environment variable (or the file it points to) anew, and set
184// *pargc, and *pargv to point to the internal locations of the argc and argv
185// constructed from the environment.
186void ResetFlagsFromEnvForTesting(int** pargc, std::vector<char*>** pargv) {
187 env_argv_mu.lock();
188 if (env_argv == nullptr) {
189 env_argv = new EnvArgv;
190 }
191 if (!env_argv->argv_save.empty()) {
192 for (int i = 0; env_argv->argv_save[i] != nullptr; i++) {
193 free(env_argv->argv_save[i]);
194 }
195 }
196 env_argv->initialized = false;
197 env_argv->argc = 0;
198 env_argv->argv.clear();
199 env_argv->argv_save.clear();
200 env_argv_mu.unlock();
201 *pargc = &env_argv->argc;
202 *pargv = &env_argv->argv;
203}
204
205} // namespace legacy_flags
206} // namespace xla
207