xref: /aosp_15_r20/external/tensorflow/configure.py (revision b6fb3261f9314811a0f4371741dbb8839866f948)
1*b6fb3261SAndroid Build Coastguard Worker# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2*b6fb3261SAndroid Build Coastguard Worker#
3*b6fb3261SAndroid Build Coastguard Worker# Licensed under the Apache License, Version 2.0 (the "License");
4*b6fb3261SAndroid Build Coastguard Worker# you may not use this file except in compliance with the License.
5*b6fb3261SAndroid Build Coastguard Worker# You may obtain a copy of the License at
6*b6fb3261SAndroid Build Coastguard Worker#
7*b6fb3261SAndroid Build Coastguard Worker#     http://www.apache.org/licenses/LICENSE-2.0
8*b6fb3261SAndroid Build Coastguard Worker#
9*b6fb3261SAndroid Build Coastguard Worker# Unless required by applicable law or agreed to in writing, software
10*b6fb3261SAndroid Build Coastguard Worker# distributed under the License is distributed on an "AS IS" BASIS,
11*b6fb3261SAndroid Build Coastguard Worker# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12*b6fb3261SAndroid Build Coastguard Worker# See the License for the specific language governing permissions and
13*b6fb3261SAndroid Build Coastguard Worker# limitations under the License.
14*b6fb3261SAndroid Build Coastguard Worker# ==============================================================================
15*b6fb3261SAndroid Build Coastguard Worker"""configure script to get build parameters from user."""
16*b6fb3261SAndroid Build Coastguard Worker
17*b6fb3261SAndroid Build Coastguard Workerimport argparse
18*b6fb3261SAndroid Build Coastguard Workerimport errno
19*b6fb3261SAndroid Build Coastguard Workerimport glob
20*b6fb3261SAndroid Build Coastguard Workerimport os
21*b6fb3261SAndroid Build Coastguard Workerimport platform
22*b6fb3261SAndroid Build Coastguard Workerimport re
23*b6fb3261SAndroid Build Coastguard Workerimport subprocess
24*b6fb3261SAndroid Build Coastguard Workerimport sys
25*b6fb3261SAndroid Build Coastguard Worker
26*b6fb3261SAndroid Build Coastguard Worker# pylint: disable=g-import-not-at-top
27*b6fb3261SAndroid Build Coastguard Workertry:
28*b6fb3261SAndroid Build Coastguard Worker  from shutil import which
29*b6fb3261SAndroid Build Coastguard Workerexcept ImportError:
30*b6fb3261SAndroid Build Coastguard Worker  from distutils.spawn import find_executable as which
31*b6fb3261SAndroid Build Coastguard Worker# pylint: enable=g-import-not-at-top
32*b6fb3261SAndroid Build Coastguard Worker
33*b6fb3261SAndroid Build Coastguard Worker_DEFAULT_CUDA_VERSION = '11'
34*b6fb3261SAndroid Build Coastguard Worker_DEFAULT_CUDNN_VERSION = '2'
35*b6fb3261SAndroid Build Coastguard Worker_DEFAULT_TENSORRT_VERSION = '6'
36*b6fb3261SAndroid Build Coastguard Worker_DEFAULT_CUDA_COMPUTE_CAPABILITIES = '3.5,7.0'
37*b6fb3261SAndroid Build Coastguard Worker
38*b6fb3261SAndroid Build Coastguard Worker_SUPPORTED_ANDROID_NDK_VERSIONS = [
39*b6fb3261SAndroid Build Coastguard Worker    10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21
40*b6fb3261SAndroid Build Coastguard Worker]
41*b6fb3261SAndroid Build Coastguard Worker
42*b6fb3261SAndroid Build Coastguard Worker_DEFAULT_PROMPT_ASK_ATTEMPTS = 10
43*b6fb3261SAndroid Build Coastguard Worker
44*b6fb3261SAndroid Build Coastguard Worker_TF_BAZELRC_FILENAME = '.tf_configure.bazelrc'
45*b6fb3261SAndroid Build Coastguard Worker_TF_WORKSPACE_ROOT = ''
46*b6fb3261SAndroid Build Coastguard Worker_TF_BAZELRC = ''
47*b6fb3261SAndroid Build Coastguard Worker_TF_CURRENT_BAZEL_VERSION = None
48*b6fb3261SAndroid Build Coastguard Worker
49*b6fb3261SAndroid Build Coastguard WorkerNCCL_LIB_PATHS = [
50*b6fb3261SAndroid Build Coastguard Worker    'lib64/', 'lib/powerpc64le-linux-gnu/', 'lib/x86_64-linux-gnu/', ''
51*b6fb3261SAndroid Build Coastguard Worker]
52*b6fb3261SAndroid Build Coastguard Worker
53*b6fb3261SAndroid Build Coastguard Worker# List of files to configure when building Bazel on Apple platforms.
54*b6fb3261SAndroid Build Coastguard WorkerAPPLE_BAZEL_FILES = [
55*b6fb3261SAndroid Build Coastguard Worker    'tensorflow/lite/ios/BUILD', 'tensorflow/lite/objc/BUILD',
56*b6fb3261SAndroid Build Coastguard Worker    'tensorflow/lite/swift/BUILD',
57*b6fb3261SAndroid Build Coastguard Worker    'tensorflow/lite/tools/benchmark/experimental/ios/BUILD'
58*b6fb3261SAndroid Build Coastguard Worker]
59*b6fb3261SAndroid Build Coastguard Worker
60*b6fb3261SAndroid Build Coastguard Worker# List of files to move when building for iOS.
61*b6fb3261SAndroid Build Coastguard WorkerIOS_FILES = [
62*b6fb3261SAndroid Build Coastguard Worker    'tensorflow/lite/objc/TensorFlowLiteObjC.podspec',
63*b6fb3261SAndroid Build Coastguard Worker    'tensorflow/lite/swift/TensorFlowLiteSwift.podspec',
64*b6fb3261SAndroid Build Coastguard Worker]
65*b6fb3261SAndroid Build Coastguard Worker
66*b6fb3261SAndroid Build Coastguard Worker
67*b6fb3261SAndroid Build Coastguard Workerclass UserInputError(Exception):
68*b6fb3261SAndroid Build Coastguard Worker  pass
69*b6fb3261SAndroid Build Coastguard Worker
70*b6fb3261SAndroid Build Coastguard Worker
71*b6fb3261SAndroid Build Coastguard Workerdef is_windows():
72*b6fb3261SAndroid Build Coastguard Worker  return platform.system() == 'Windows'
73*b6fb3261SAndroid Build Coastguard Worker
74*b6fb3261SAndroid Build Coastguard Worker
75*b6fb3261SAndroid Build Coastguard Workerdef is_linux():
76*b6fb3261SAndroid Build Coastguard Worker  return platform.system() == 'Linux'
77*b6fb3261SAndroid Build Coastguard Worker
78*b6fb3261SAndroid Build Coastguard Worker
79*b6fb3261SAndroid Build Coastguard Workerdef is_macos():
80*b6fb3261SAndroid Build Coastguard Worker  return platform.system() == 'Darwin'
81*b6fb3261SAndroid Build Coastguard Worker
82*b6fb3261SAndroid Build Coastguard Worker
83*b6fb3261SAndroid Build Coastguard Workerdef is_ppc64le():
84*b6fb3261SAndroid Build Coastguard Worker  return platform.machine() == 'ppc64le'
85*b6fb3261SAndroid Build Coastguard Worker
86*b6fb3261SAndroid Build Coastguard Worker
87*b6fb3261SAndroid Build Coastguard Workerdef is_cygwin():
88*b6fb3261SAndroid Build Coastguard Worker  return platform.system().startswith('CYGWIN_NT')
89*b6fb3261SAndroid Build Coastguard Worker
90*b6fb3261SAndroid Build Coastguard Worker
91*b6fb3261SAndroid Build Coastguard Workerdef get_input(question):
92*b6fb3261SAndroid Build Coastguard Worker  try:
93*b6fb3261SAndroid Build Coastguard Worker    try:
94*b6fb3261SAndroid Build Coastguard Worker      answer = raw_input(question)
95*b6fb3261SAndroid Build Coastguard Worker    except NameError:
96*b6fb3261SAndroid Build Coastguard Worker      answer = input(question)  # pylint: disable=bad-builtin
97*b6fb3261SAndroid Build Coastguard Worker  except EOFError:
98*b6fb3261SAndroid Build Coastguard Worker    answer = ''
99*b6fb3261SAndroid Build Coastguard Worker  return answer
100*b6fb3261SAndroid Build Coastguard Worker
101*b6fb3261SAndroid Build Coastguard Worker
102*b6fb3261SAndroid Build Coastguard Workerdef symlink_force(target, link_name):
103*b6fb3261SAndroid Build Coastguard Worker  """Force symlink, equivalent of 'ln -sf'.
104*b6fb3261SAndroid Build Coastguard Worker
105*b6fb3261SAndroid Build Coastguard Worker  Args:
106*b6fb3261SAndroid Build Coastguard Worker    target: items to link to.
107*b6fb3261SAndroid Build Coastguard Worker    link_name: name of the link.
108*b6fb3261SAndroid Build Coastguard Worker  """
109*b6fb3261SAndroid Build Coastguard Worker  try:
110*b6fb3261SAndroid Build Coastguard Worker    os.symlink(target, link_name)
111*b6fb3261SAndroid Build Coastguard Worker  except OSError as e:
112*b6fb3261SAndroid Build Coastguard Worker    if e.errno == errno.EEXIST:
113*b6fb3261SAndroid Build Coastguard Worker      os.remove(link_name)
114*b6fb3261SAndroid Build Coastguard Worker      os.symlink(target, link_name)
115*b6fb3261SAndroid Build Coastguard Worker    else:
116*b6fb3261SAndroid Build Coastguard Worker      raise e
117*b6fb3261SAndroid Build Coastguard Worker
118*b6fb3261SAndroid Build Coastguard Worker
119*b6fb3261SAndroid Build Coastguard Workerdef write_to_bazelrc(line):
120*b6fb3261SAndroid Build Coastguard Worker  with open(_TF_BAZELRC, 'a') as f:
121*b6fb3261SAndroid Build Coastguard Worker    f.write(line + '\n')
122*b6fb3261SAndroid Build Coastguard Worker
123*b6fb3261SAndroid Build Coastguard Worker
124*b6fb3261SAndroid Build Coastguard Workerdef write_action_env_to_bazelrc(var_name, var):
125*b6fb3261SAndroid Build Coastguard Worker  write_to_bazelrc('build --action_env {}="{}"'.format(var_name, str(var)))
126*b6fb3261SAndroid Build Coastguard Worker
127*b6fb3261SAndroid Build Coastguard Worker
128*b6fb3261SAndroid Build Coastguard Workerdef run_shell(cmd, allow_non_zero=False, stderr=None):
129*b6fb3261SAndroid Build Coastguard Worker  if stderr is None:
130*b6fb3261SAndroid Build Coastguard Worker    stderr = sys.stdout
131*b6fb3261SAndroid Build Coastguard Worker  if allow_non_zero:
132*b6fb3261SAndroid Build Coastguard Worker    try:
133*b6fb3261SAndroid Build Coastguard Worker      output = subprocess.check_output(cmd, stderr=stderr)
134*b6fb3261SAndroid Build Coastguard Worker    except subprocess.CalledProcessError as e:
135*b6fb3261SAndroid Build Coastguard Worker      output = e.output
136*b6fb3261SAndroid Build Coastguard Worker  else:
137*b6fb3261SAndroid Build Coastguard Worker    output = subprocess.check_output(cmd, stderr=stderr)
138*b6fb3261SAndroid Build Coastguard Worker  return output.decode('UTF-8').strip()
139*b6fb3261SAndroid Build Coastguard Worker
140*b6fb3261SAndroid Build Coastguard Worker
141*b6fb3261SAndroid Build Coastguard Workerdef cygpath(path):
142*b6fb3261SAndroid Build Coastguard Worker  """Convert path from posix to windows."""
143*b6fb3261SAndroid Build Coastguard Worker  return os.path.abspath(path).replace('\\', '/')
144*b6fb3261SAndroid Build Coastguard Worker
145*b6fb3261SAndroid Build Coastguard Worker
146*b6fb3261SAndroid Build Coastguard Workerdef get_python_path(environ_cp, python_bin_path):
147*b6fb3261SAndroid Build Coastguard Worker  """Get the python site package paths."""
148*b6fb3261SAndroid Build Coastguard Worker  python_paths = []
149*b6fb3261SAndroid Build Coastguard Worker  if environ_cp.get('PYTHONPATH'):
150*b6fb3261SAndroid Build Coastguard Worker    python_paths = environ_cp.get('PYTHONPATH').split(':')
151*b6fb3261SAndroid Build Coastguard Worker  try:
152*b6fb3261SAndroid Build Coastguard Worker    stderr = open(os.devnull, 'wb')
153*b6fb3261SAndroid Build Coastguard Worker    library_paths = run_shell([
154*b6fb3261SAndroid Build Coastguard Worker        python_bin_path, '-c',
155*b6fb3261SAndroid Build Coastguard Worker        'import site; print("\\n".join(site.getsitepackages()))'
156*b6fb3261SAndroid Build Coastguard Worker    ],
157*b6fb3261SAndroid Build Coastguard Worker                              stderr=stderr).split('\n')
158*b6fb3261SAndroid Build Coastguard Worker  except subprocess.CalledProcessError:
159*b6fb3261SAndroid Build Coastguard Worker    library_paths = [
160*b6fb3261SAndroid Build Coastguard Worker        run_shell([
161*b6fb3261SAndroid Build Coastguard Worker            python_bin_path, '-c',
162*b6fb3261SAndroid Build Coastguard Worker            'from distutils.sysconfig import get_python_lib;'
163*b6fb3261SAndroid Build Coastguard Worker            'print(get_python_lib())'
164*b6fb3261SAndroid Build Coastguard Worker        ])
165*b6fb3261SAndroid Build Coastguard Worker    ]
166*b6fb3261SAndroid Build Coastguard Worker
167*b6fb3261SAndroid Build Coastguard Worker  all_paths = set(python_paths + library_paths)
168*b6fb3261SAndroid Build Coastguard Worker  # Sort set so order is deterministic
169*b6fb3261SAndroid Build Coastguard Worker  all_paths = sorted(all_paths)
170*b6fb3261SAndroid Build Coastguard Worker
171*b6fb3261SAndroid Build Coastguard Worker  paths = []
172*b6fb3261SAndroid Build Coastguard Worker  for path in all_paths:
173*b6fb3261SAndroid Build Coastguard Worker    if os.path.isdir(path):
174*b6fb3261SAndroid Build Coastguard Worker      paths.append(path)
175*b6fb3261SAndroid Build Coastguard Worker  return paths
176*b6fb3261SAndroid Build Coastguard Worker
177*b6fb3261SAndroid Build Coastguard Worker
178*b6fb3261SAndroid Build Coastguard Workerdef get_python_major_version(python_bin_path):
179*b6fb3261SAndroid Build Coastguard Worker  """Get the python major version."""
180*b6fb3261SAndroid Build Coastguard Worker  return run_shell([python_bin_path, '-c', 'import sys; print(sys.version[0])'])
181*b6fb3261SAndroid Build Coastguard Worker
182*b6fb3261SAndroid Build Coastguard Worker
183*b6fb3261SAndroid Build Coastguard Workerdef setup_python(environ_cp):
184*b6fb3261SAndroid Build Coastguard Worker  """Setup python related env variables."""
185*b6fb3261SAndroid Build Coastguard Worker  # Get PYTHON_BIN_PATH, default is the current running python.
186*b6fb3261SAndroid Build Coastguard Worker  default_python_bin_path = sys.executable
187*b6fb3261SAndroid Build Coastguard Worker  ask_python_bin_path = ('Please specify the location of python. [Default is '
188*b6fb3261SAndroid Build Coastguard Worker                         '{}]: ').format(default_python_bin_path)
189*b6fb3261SAndroid Build Coastguard Worker  while True:
190*b6fb3261SAndroid Build Coastguard Worker    python_bin_path = get_from_env_or_user_or_default(environ_cp,
191*b6fb3261SAndroid Build Coastguard Worker                                                      'PYTHON_BIN_PATH',
192*b6fb3261SAndroid Build Coastguard Worker                                                      ask_python_bin_path,
193*b6fb3261SAndroid Build Coastguard Worker                                                      default_python_bin_path)
194*b6fb3261SAndroid Build Coastguard Worker    # Check if the path is valid
195*b6fb3261SAndroid Build Coastguard Worker    if os.path.isfile(python_bin_path) and os.access(python_bin_path, os.X_OK):
196*b6fb3261SAndroid Build Coastguard Worker      break
197*b6fb3261SAndroid Build Coastguard Worker    elif not os.path.exists(python_bin_path):
198*b6fb3261SAndroid Build Coastguard Worker      print('Invalid python path: {} cannot be found.'.format(python_bin_path))
199*b6fb3261SAndroid Build Coastguard Worker    else:
200*b6fb3261SAndroid Build Coastguard Worker      print('{} is not executable.  Is it the python binary?'.format(
201*b6fb3261SAndroid Build Coastguard Worker          python_bin_path))
202*b6fb3261SAndroid Build Coastguard Worker    environ_cp['PYTHON_BIN_PATH'] = ''
203*b6fb3261SAndroid Build Coastguard Worker
204*b6fb3261SAndroid Build Coastguard Worker  # Convert python path to Windows style before checking lib and version
205*b6fb3261SAndroid Build Coastguard Worker  if is_windows() or is_cygwin():
206*b6fb3261SAndroid Build Coastguard Worker    python_bin_path = cygpath(python_bin_path)
207*b6fb3261SAndroid Build Coastguard Worker
208*b6fb3261SAndroid Build Coastguard Worker  # Get PYTHON_LIB_PATH
209*b6fb3261SAndroid Build Coastguard Worker  python_lib_path = environ_cp.get('PYTHON_LIB_PATH')
210*b6fb3261SAndroid Build Coastguard Worker  if not python_lib_path:
211*b6fb3261SAndroid Build Coastguard Worker    python_lib_paths = get_python_path(environ_cp, python_bin_path)
212*b6fb3261SAndroid Build Coastguard Worker    if environ_cp.get('USE_DEFAULT_PYTHON_LIB_PATH') == '1':
213*b6fb3261SAndroid Build Coastguard Worker      python_lib_path = python_lib_paths[0]
214*b6fb3261SAndroid Build Coastguard Worker    else:
215*b6fb3261SAndroid Build Coastguard Worker      print('Found possible Python library paths:\n  %s' %
216*b6fb3261SAndroid Build Coastguard Worker            '\n  '.join(python_lib_paths))
217*b6fb3261SAndroid Build Coastguard Worker      default_python_lib_path = python_lib_paths[0]
218*b6fb3261SAndroid Build Coastguard Worker      python_lib_path = get_input(
219*b6fb3261SAndroid Build Coastguard Worker          'Please input the desired Python library path to use.  '
220*b6fb3261SAndroid Build Coastguard Worker          'Default is [{}]\n'.format(python_lib_paths[0]))
221*b6fb3261SAndroid Build Coastguard Worker      if not python_lib_path:
222*b6fb3261SAndroid Build Coastguard Worker        python_lib_path = default_python_lib_path
223*b6fb3261SAndroid Build Coastguard Worker    environ_cp['PYTHON_LIB_PATH'] = python_lib_path
224*b6fb3261SAndroid Build Coastguard Worker
225*b6fb3261SAndroid Build Coastguard Worker  python_major_version = get_python_major_version(python_bin_path)
226*b6fb3261SAndroid Build Coastguard Worker  if python_major_version == '2':
227*b6fb3261SAndroid Build Coastguard Worker    write_to_bazelrc('build --host_force_python=PY2')
228*b6fb3261SAndroid Build Coastguard Worker
229*b6fb3261SAndroid Build Coastguard Worker  # Convert python path to Windows style before writing into bazel.rc
230*b6fb3261SAndroid Build Coastguard Worker  if is_windows() or is_cygwin():
231*b6fb3261SAndroid Build Coastguard Worker    python_lib_path = cygpath(python_lib_path)
232*b6fb3261SAndroid Build Coastguard Worker
233*b6fb3261SAndroid Build Coastguard Worker  # Set-up env variables used by python_configure.bzl
234*b6fb3261SAndroid Build Coastguard Worker  write_action_env_to_bazelrc('PYTHON_BIN_PATH', python_bin_path)
235*b6fb3261SAndroid Build Coastguard Worker  write_action_env_to_bazelrc('PYTHON_LIB_PATH', python_lib_path)
236*b6fb3261SAndroid Build Coastguard Worker  write_to_bazelrc('build --python_path=\"{}"'.format(python_bin_path))
237*b6fb3261SAndroid Build Coastguard Worker  environ_cp['PYTHON_BIN_PATH'] = python_bin_path
238*b6fb3261SAndroid Build Coastguard Worker
239*b6fb3261SAndroid Build Coastguard Worker  # If choosen python_lib_path is from a path specified in the PYTHONPATH
240*b6fb3261SAndroid Build Coastguard Worker  # variable, need to tell bazel to include PYTHONPATH
241*b6fb3261SAndroid Build Coastguard Worker  if environ_cp.get('PYTHONPATH'):
242*b6fb3261SAndroid Build Coastguard Worker    python_paths = environ_cp.get('PYTHONPATH').split(':')
243*b6fb3261SAndroid Build Coastguard Worker    if python_lib_path in python_paths:
244*b6fb3261SAndroid Build Coastguard Worker      write_action_env_to_bazelrc('PYTHONPATH', environ_cp.get('PYTHONPATH'))
245*b6fb3261SAndroid Build Coastguard Worker
246*b6fb3261SAndroid Build Coastguard Worker  # Write tools/python_bin_path.sh
247*b6fb3261SAndroid Build Coastguard Worker  with open(
248*b6fb3261SAndroid Build Coastguard Worker      os.path.join(_TF_WORKSPACE_ROOT, 'tools', 'python_bin_path.sh'),
249*b6fb3261SAndroid Build Coastguard Worker      'w') as f:
250*b6fb3261SAndroid Build Coastguard Worker    f.write('export PYTHON_BIN_PATH="{}"'.format(python_bin_path))
251*b6fb3261SAndroid Build Coastguard Worker
252*b6fb3261SAndroid Build Coastguard Worker
253*b6fb3261SAndroid Build Coastguard Workerdef reset_tf_configure_bazelrc():
254*b6fb3261SAndroid Build Coastguard Worker  """Reset file that contains customized config settings."""
255*b6fb3261SAndroid Build Coastguard Worker  open(_TF_BAZELRC, 'w').close()
256*b6fb3261SAndroid Build Coastguard Worker
257*b6fb3261SAndroid Build Coastguard Worker
258*b6fb3261SAndroid Build Coastguard Workerdef cleanup_makefile():
259*b6fb3261SAndroid Build Coastguard Worker  """Delete any leftover BUILD files from the Makefile build.
260*b6fb3261SAndroid Build Coastguard Worker
261*b6fb3261SAndroid Build Coastguard Worker  These files could interfere with Bazel parsing.
262*b6fb3261SAndroid Build Coastguard Worker  """
263*b6fb3261SAndroid Build Coastguard Worker  makefile_download_dir = os.path.join(_TF_WORKSPACE_ROOT, 'tensorflow',
264*b6fb3261SAndroid Build Coastguard Worker                                       'contrib', 'makefile', 'downloads')
265*b6fb3261SAndroid Build Coastguard Worker  if os.path.isdir(makefile_download_dir):
266*b6fb3261SAndroid Build Coastguard Worker    for root, _, filenames in os.walk(makefile_download_dir):
267*b6fb3261SAndroid Build Coastguard Worker      for f in filenames:
268*b6fb3261SAndroid Build Coastguard Worker        if f.endswith('BUILD'):
269*b6fb3261SAndroid Build Coastguard Worker          os.remove(os.path.join(root, f))
270*b6fb3261SAndroid Build Coastguard Worker
271*b6fb3261SAndroid Build Coastguard Worker
272*b6fb3261SAndroid Build Coastguard Workerdef get_var(environ_cp,
273*b6fb3261SAndroid Build Coastguard Worker            var_name,
274*b6fb3261SAndroid Build Coastguard Worker            query_item,
275*b6fb3261SAndroid Build Coastguard Worker            enabled_by_default,
276*b6fb3261SAndroid Build Coastguard Worker            question=None,
277*b6fb3261SAndroid Build Coastguard Worker            yes_reply=None,
278*b6fb3261SAndroid Build Coastguard Worker            no_reply=None):
279*b6fb3261SAndroid Build Coastguard Worker  """Get boolean input from user.
280*b6fb3261SAndroid Build Coastguard Worker
281*b6fb3261SAndroid Build Coastguard Worker  If var_name is not set in env, ask user to enable query_item or not. If the
282*b6fb3261SAndroid Build Coastguard Worker  response is empty, use the default.
283*b6fb3261SAndroid Build Coastguard Worker
284*b6fb3261SAndroid Build Coastguard Worker  Args:
285*b6fb3261SAndroid Build Coastguard Worker    environ_cp: copy of the os.environ.
286*b6fb3261SAndroid Build Coastguard Worker    var_name: string for name of environment variable, e.g. "TF_NEED_CUDA".
287*b6fb3261SAndroid Build Coastguard Worker    query_item: string for feature related to the variable, e.g. "CUDA for
288*b6fb3261SAndroid Build Coastguard Worker      Nvidia GPUs".
289*b6fb3261SAndroid Build Coastguard Worker    enabled_by_default: boolean for default behavior.
290*b6fb3261SAndroid Build Coastguard Worker    question: optional string for how to ask for user input.
291*b6fb3261SAndroid Build Coastguard Worker    yes_reply: optional string for reply when feature is enabled.
292*b6fb3261SAndroid Build Coastguard Worker    no_reply: optional string for reply when feature is disabled.
293*b6fb3261SAndroid Build Coastguard Worker
294*b6fb3261SAndroid Build Coastguard Worker  Returns:
295*b6fb3261SAndroid Build Coastguard Worker    boolean value of the variable.
296*b6fb3261SAndroid Build Coastguard Worker
297*b6fb3261SAndroid Build Coastguard Worker  Raises:
298*b6fb3261SAndroid Build Coastguard Worker    UserInputError: if an environment variable is set, but it cannot be
299*b6fb3261SAndroid Build Coastguard Worker      interpreted as a boolean indicator, assume that the user has made a
300*b6fb3261SAndroid Build Coastguard Worker      scripting error, and will continue to provide invalid input.
301*b6fb3261SAndroid Build Coastguard Worker      Raise the error to avoid infinitely looping.
302*b6fb3261SAndroid Build Coastguard Worker  """
303*b6fb3261SAndroid Build Coastguard Worker  if not question:
304*b6fb3261SAndroid Build Coastguard Worker    question = 'Do you wish to build TensorFlow with {} support?'.format(
305*b6fb3261SAndroid Build Coastguard Worker        query_item)
306*b6fb3261SAndroid Build Coastguard Worker  if not yes_reply:
307*b6fb3261SAndroid Build Coastguard Worker    yes_reply = '{} support will be enabled for TensorFlow.'.format(query_item)
308*b6fb3261SAndroid Build Coastguard Worker  if not no_reply:
309*b6fb3261SAndroid Build Coastguard Worker    no_reply = 'No {}'.format(yes_reply)
310*b6fb3261SAndroid Build Coastguard Worker
311*b6fb3261SAndroid Build Coastguard Worker  yes_reply += '\n'
312*b6fb3261SAndroid Build Coastguard Worker  no_reply += '\n'
313*b6fb3261SAndroid Build Coastguard Worker
314*b6fb3261SAndroid Build Coastguard Worker  if enabled_by_default:
315*b6fb3261SAndroid Build Coastguard Worker    question += ' [Y/n]: '
316*b6fb3261SAndroid Build Coastguard Worker  else:
317*b6fb3261SAndroid Build Coastguard Worker    question += ' [y/N]: '
318*b6fb3261SAndroid Build Coastguard Worker
319*b6fb3261SAndroid Build Coastguard Worker  var = environ_cp.get(var_name)
320*b6fb3261SAndroid Build Coastguard Worker  if var is not None:
321*b6fb3261SAndroid Build Coastguard Worker    var_content = var.strip().lower()
322*b6fb3261SAndroid Build Coastguard Worker    true_strings = ('1', 't', 'true', 'y', 'yes')
323*b6fb3261SAndroid Build Coastguard Worker    false_strings = ('0', 'f', 'false', 'n', 'no')
324*b6fb3261SAndroid Build Coastguard Worker    if var_content in true_strings:
325*b6fb3261SAndroid Build Coastguard Worker      var = True
326*b6fb3261SAndroid Build Coastguard Worker    elif var_content in false_strings:
327*b6fb3261SAndroid Build Coastguard Worker      var = False
328*b6fb3261SAndroid Build Coastguard Worker    else:
329*b6fb3261SAndroid Build Coastguard Worker      raise UserInputError(
330*b6fb3261SAndroid Build Coastguard Worker          'Environment variable %s must be set as a boolean indicator.\n'
331*b6fb3261SAndroid Build Coastguard Worker          'The following are accepted as TRUE : %s.\n'
332*b6fb3261SAndroid Build Coastguard Worker          'The following are accepted as FALSE: %s.\n'
333*b6fb3261SAndroid Build Coastguard Worker          'Current value is %s.' %
334*b6fb3261SAndroid Build Coastguard Worker          (var_name, ', '.join(true_strings), ', '.join(false_strings), var))
335*b6fb3261SAndroid Build Coastguard Worker
336*b6fb3261SAndroid Build Coastguard Worker  while var is None:
337*b6fb3261SAndroid Build Coastguard Worker    user_input_origin = get_input(question)
338*b6fb3261SAndroid Build Coastguard Worker    user_input = user_input_origin.strip().lower()
339*b6fb3261SAndroid Build Coastguard Worker    if user_input == 'y':
340*b6fb3261SAndroid Build Coastguard Worker      print(yes_reply)
341*b6fb3261SAndroid Build Coastguard Worker      var = True
342*b6fb3261SAndroid Build Coastguard Worker    elif user_input == 'n':
343*b6fb3261SAndroid Build Coastguard Worker      print(no_reply)
344*b6fb3261SAndroid Build Coastguard Worker      var = False
345*b6fb3261SAndroid Build Coastguard Worker    elif not user_input:
346*b6fb3261SAndroid Build Coastguard Worker      if enabled_by_default:
347*b6fb3261SAndroid Build Coastguard Worker        print(yes_reply)
348*b6fb3261SAndroid Build Coastguard Worker        var = True
349*b6fb3261SAndroid Build Coastguard Worker      else:
350*b6fb3261SAndroid Build Coastguard Worker        print(no_reply)
351*b6fb3261SAndroid Build Coastguard Worker        var = False
352*b6fb3261SAndroid Build Coastguard Worker    else:
353*b6fb3261SAndroid Build Coastguard Worker      print('Invalid selection: {}'.format(user_input_origin))
354*b6fb3261SAndroid Build Coastguard Worker  return var
355*b6fb3261SAndroid Build Coastguard Worker
356*b6fb3261SAndroid Build Coastguard Worker
357*b6fb3261SAndroid Build Coastguard Workerdef set_action_env_var(environ_cp,
358*b6fb3261SAndroid Build Coastguard Worker                       var_name,
359*b6fb3261SAndroid Build Coastguard Worker                       query_item,
360*b6fb3261SAndroid Build Coastguard Worker                       enabled_by_default,
361*b6fb3261SAndroid Build Coastguard Worker                       question=None,
362*b6fb3261SAndroid Build Coastguard Worker                       yes_reply=None,
363*b6fb3261SAndroid Build Coastguard Worker                       no_reply=None,
364*b6fb3261SAndroid Build Coastguard Worker                       bazel_config_name=None):
365*b6fb3261SAndroid Build Coastguard Worker  """Set boolean action_env variable.
366*b6fb3261SAndroid Build Coastguard Worker
367*b6fb3261SAndroid Build Coastguard Worker  Ask user if query_item will be enabled. Default is used if no input is given.
368*b6fb3261SAndroid Build Coastguard Worker  Set environment variable and write to .bazelrc.
369*b6fb3261SAndroid Build Coastguard Worker
370*b6fb3261SAndroid Build Coastguard Worker  Args:
371*b6fb3261SAndroid Build Coastguard Worker    environ_cp: copy of the os.environ.
372*b6fb3261SAndroid Build Coastguard Worker    var_name: string for name of environment variable, e.g. "TF_NEED_CUDA".
373*b6fb3261SAndroid Build Coastguard Worker    query_item: string for feature related to the variable, e.g. "CUDA for
374*b6fb3261SAndroid Build Coastguard Worker      Nvidia GPUs".
375*b6fb3261SAndroid Build Coastguard Worker    enabled_by_default: boolean for default behavior.
376*b6fb3261SAndroid Build Coastguard Worker    question: optional string for how to ask for user input.
377*b6fb3261SAndroid Build Coastguard Worker    yes_reply: optional string for reply when feature is enabled.
378*b6fb3261SAndroid Build Coastguard Worker    no_reply: optional string for reply when feature is disabled.
379*b6fb3261SAndroid Build Coastguard Worker    bazel_config_name: adding config to .bazelrc instead of action_env.
380*b6fb3261SAndroid Build Coastguard Worker  """
381*b6fb3261SAndroid Build Coastguard Worker  var = int(
382*b6fb3261SAndroid Build Coastguard Worker      get_var(environ_cp, var_name, query_item, enabled_by_default, question,
383*b6fb3261SAndroid Build Coastguard Worker              yes_reply, no_reply))
384*b6fb3261SAndroid Build Coastguard Worker
385*b6fb3261SAndroid Build Coastguard Worker  if not bazel_config_name:
386*b6fb3261SAndroid Build Coastguard Worker    write_action_env_to_bazelrc(var_name, var)
387*b6fb3261SAndroid Build Coastguard Worker  elif var:
388*b6fb3261SAndroid Build Coastguard Worker    write_to_bazelrc('build --config=%s' % bazel_config_name)
389*b6fb3261SAndroid Build Coastguard Worker  environ_cp[var_name] = str(var)
390*b6fb3261SAndroid Build Coastguard Worker
391*b6fb3261SAndroid Build Coastguard Worker
392*b6fb3261SAndroid Build Coastguard Workerdef convert_version_to_int(version):
393*b6fb3261SAndroid Build Coastguard Worker  """Convert a version number to a integer that can be used to compare.
394*b6fb3261SAndroid Build Coastguard Worker
395*b6fb3261SAndroid Build Coastguard Worker  Version strings of the form X.YZ and X.Y.Z-xxxxx are supported. The
396*b6fb3261SAndroid Build Coastguard Worker  'xxxxx' part, for instance 'homebrew' on OS/X, is ignored.
397*b6fb3261SAndroid Build Coastguard Worker
398*b6fb3261SAndroid Build Coastguard Worker  Args:
399*b6fb3261SAndroid Build Coastguard Worker    version: a version to be converted
400*b6fb3261SAndroid Build Coastguard Worker
401*b6fb3261SAndroid Build Coastguard Worker  Returns:
402*b6fb3261SAndroid Build Coastguard Worker    An integer if converted successfully, otherwise return None.
403*b6fb3261SAndroid Build Coastguard Worker  """
404*b6fb3261SAndroid Build Coastguard Worker  version = version.split('-')[0]
405*b6fb3261SAndroid Build Coastguard Worker  version_segments = version.split('.')
406*b6fb3261SAndroid Build Coastguard Worker  # Treat "0.24" as "0.24.0"
407*b6fb3261SAndroid Build Coastguard Worker  if len(version_segments) == 2:
408*b6fb3261SAndroid Build Coastguard Worker    version_segments.append('0')
409*b6fb3261SAndroid Build Coastguard Worker  for seg in version_segments:
410*b6fb3261SAndroid Build Coastguard Worker    if not seg.isdigit():
411*b6fb3261SAndroid Build Coastguard Worker      return None
412*b6fb3261SAndroid Build Coastguard Worker
413*b6fb3261SAndroid Build Coastguard Worker  version_str = ''.join(['%03d' % int(seg) for seg in version_segments])
414*b6fb3261SAndroid Build Coastguard Worker  return int(version_str)
415*b6fb3261SAndroid Build Coastguard Worker
416*b6fb3261SAndroid Build Coastguard Worker
417*b6fb3261SAndroid Build Coastguard Workerdef retrieve_bazel_version():
418*b6fb3261SAndroid Build Coastguard Worker  """Retrieve installed bazel version (or bazelisk).
419*b6fb3261SAndroid Build Coastguard Worker
420*b6fb3261SAndroid Build Coastguard Worker  Returns:
421*b6fb3261SAndroid Build Coastguard Worker    The bazel version detected.
422*b6fb3261SAndroid Build Coastguard Worker  """
423*b6fb3261SAndroid Build Coastguard Worker  bazel_executable = which('bazel')
424*b6fb3261SAndroid Build Coastguard Worker  if bazel_executable is None:
425*b6fb3261SAndroid Build Coastguard Worker    bazel_executable = which('bazelisk')
426*b6fb3261SAndroid Build Coastguard Worker    if bazel_executable is None:
427*b6fb3261SAndroid Build Coastguard Worker      print('Cannot find bazel. Please install bazel/bazelisk.')
428*b6fb3261SAndroid Build Coastguard Worker      sys.exit(1)
429*b6fb3261SAndroid Build Coastguard Worker
430*b6fb3261SAndroid Build Coastguard Worker  stderr = open(os.devnull, 'wb')
431*b6fb3261SAndroid Build Coastguard Worker  curr_version = run_shell([bazel_executable, '--version'],
432*b6fb3261SAndroid Build Coastguard Worker                           allow_non_zero=True,
433*b6fb3261SAndroid Build Coastguard Worker                           stderr=stderr)
434*b6fb3261SAndroid Build Coastguard Worker  if curr_version.startswith('bazel '):
435*b6fb3261SAndroid Build Coastguard Worker    curr_version = curr_version.split('bazel ')[1]
436*b6fb3261SAndroid Build Coastguard Worker
437*b6fb3261SAndroid Build Coastguard Worker  curr_version_int = convert_version_to_int(curr_version)
438*b6fb3261SAndroid Build Coastguard Worker
439*b6fb3261SAndroid Build Coastguard Worker  # Check if current bazel version can be detected properly.
440*b6fb3261SAndroid Build Coastguard Worker  if not curr_version_int:
441*b6fb3261SAndroid Build Coastguard Worker    print('WARNING: current bazel installation is not a release version.')
442*b6fb3261SAndroid Build Coastguard Worker    return curr_version
443*b6fb3261SAndroid Build Coastguard Worker
444*b6fb3261SAndroid Build Coastguard Worker  print('You have bazel %s installed.' % curr_version)
445*b6fb3261SAndroid Build Coastguard Worker  return curr_version
446*b6fb3261SAndroid Build Coastguard Worker
447*b6fb3261SAndroid Build Coastguard Worker
448*b6fb3261SAndroid Build Coastguard Workerdef set_cc_opt_flags(environ_cp):
449*b6fb3261SAndroid Build Coastguard Worker  """Set up architecture-dependent optimization flags.
450*b6fb3261SAndroid Build Coastguard Worker
451*b6fb3261SAndroid Build Coastguard Worker  Also append CC optimization flags to bazel.rc..
452*b6fb3261SAndroid Build Coastguard Worker
453*b6fb3261SAndroid Build Coastguard Worker  Args:
454*b6fb3261SAndroid Build Coastguard Worker    environ_cp: copy of the os.environ.
455*b6fb3261SAndroid Build Coastguard Worker  """
456*b6fb3261SAndroid Build Coastguard Worker  if is_ppc64le():
457*b6fb3261SAndroid Build Coastguard Worker    # gcc on ppc64le does not support -march, use mcpu instead
458*b6fb3261SAndroid Build Coastguard Worker    default_cc_opt_flags = '-mcpu=native'
459*b6fb3261SAndroid Build Coastguard Worker  elif is_windows():
460*b6fb3261SAndroid Build Coastguard Worker    default_cc_opt_flags = '/arch:AVX'
461*b6fb3261SAndroid Build Coastguard Worker  else:
462*b6fb3261SAndroid Build Coastguard Worker    # On all other platforms, no longer use `-march=native` as this can result
463*b6fb3261SAndroid Build Coastguard Worker    # in instructions that are too modern being generated. Users that want
464*b6fb3261SAndroid Build Coastguard Worker    # maximum performance should compile TF in their environment and can pass
465*b6fb3261SAndroid Build Coastguard Worker    # `-march=native` there.
466*b6fb3261SAndroid Build Coastguard Worker    # See https://github.com/tensorflow/tensorflow/issues/45744 and duplicates
467*b6fb3261SAndroid Build Coastguard Worker    default_cc_opt_flags = '-Wno-sign-compare'
468*b6fb3261SAndroid Build Coastguard Worker  question = ('Please specify optimization flags to use during compilation when'
469*b6fb3261SAndroid Build Coastguard Worker              ' bazel option "--config=opt" is specified [Default is %s]: '
470*b6fb3261SAndroid Build Coastguard Worker             ) % default_cc_opt_flags
471*b6fb3261SAndroid Build Coastguard Worker  cc_opt_flags = get_from_env_or_user_or_default(environ_cp, 'CC_OPT_FLAGS',
472*b6fb3261SAndroid Build Coastguard Worker                                                 question, default_cc_opt_flags)
473*b6fb3261SAndroid Build Coastguard Worker  for opt in cc_opt_flags.split():
474*b6fb3261SAndroid Build Coastguard Worker    write_to_bazelrc('build:opt --copt=%s' % opt)
475*b6fb3261SAndroid Build Coastguard Worker    write_to_bazelrc('build:opt --host_copt=%s' % opt)
476*b6fb3261SAndroid Build Coastguard Worker
477*b6fb3261SAndroid Build Coastguard Worker
478*b6fb3261SAndroid Build Coastguard Workerdef set_tf_cuda_clang(environ_cp):
479*b6fb3261SAndroid Build Coastguard Worker  """set TF_CUDA_CLANG action_env.
480*b6fb3261SAndroid Build Coastguard Worker
481*b6fb3261SAndroid Build Coastguard Worker  Args:
482*b6fb3261SAndroid Build Coastguard Worker    environ_cp: copy of the os.environ.
483*b6fb3261SAndroid Build Coastguard Worker  """
484*b6fb3261SAndroid Build Coastguard Worker  question = 'Do you want to use clang as CUDA compiler?'
485*b6fb3261SAndroid Build Coastguard Worker  yes_reply = 'Clang will be used as CUDA compiler.'
486*b6fb3261SAndroid Build Coastguard Worker  no_reply = 'nvcc will be used as CUDA compiler.'
487*b6fb3261SAndroid Build Coastguard Worker  set_action_env_var(
488*b6fb3261SAndroid Build Coastguard Worker      environ_cp,
489*b6fb3261SAndroid Build Coastguard Worker      'TF_CUDA_CLANG',
490*b6fb3261SAndroid Build Coastguard Worker      None,
491*b6fb3261SAndroid Build Coastguard Worker      False,
492*b6fb3261SAndroid Build Coastguard Worker      question=question,
493*b6fb3261SAndroid Build Coastguard Worker      yes_reply=yes_reply,
494*b6fb3261SAndroid Build Coastguard Worker      no_reply=no_reply,
495*b6fb3261SAndroid Build Coastguard Worker      bazel_config_name='cuda_clang')
496*b6fb3261SAndroid Build Coastguard Worker
497*b6fb3261SAndroid Build Coastguard Worker
498*b6fb3261SAndroid Build Coastguard Workerdef set_tf_download_clang(environ_cp):
499*b6fb3261SAndroid Build Coastguard Worker  """Set TF_DOWNLOAD_CLANG action_env."""
500*b6fb3261SAndroid Build Coastguard Worker  question = 'Do you wish to download a fresh release of clang? (Experimental)'
501*b6fb3261SAndroid Build Coastguard Worker  yes_reply = 'Clang will be downloaded and used to compile tensorflow.'
502*b6fb3261SAndroid Build Coastguard Worker  no_reply = 'Clang will not be downloaded.'
503*b6fb3261SAndroid Build Coastguard Worker  set_action_env_var(
504*b6fb3261SAndroid Build Coastguard Worker      environ_cp,
505*b6fb3261SAndroid Build Coastguard Worker      'TF_DOWNLOAD_CLANG',
506*b6fb3261SAndroid Build Coastguard Worker      None,
507*b6fb3261SAndroid Build Coastguard Worker      False,
508*b6fb3261SAndroid Build Coastguard Worker      question=question,
509*b6fb3261SAndroid Build Coastguard Worker      yes_reply=yes_reply,
510*b6fb3261SAndroid Build Coastguard Worker      no_reply=no_reply,
511*b6fb3261SAndroid Build Coastguard Worker      bazel_config_name='download_clang')
512*b6fb3261SAndroid Build Coastguard Worker
513*b6fb3261SAndroid Build Coastguard Worker
514*b6fb3261SAndroid Build Coastguard Workerdef get_from_env_or_user_or_default(environ_cp, var_name, ask_for_var,
515*b6fb3261SAndroid Build Coastguard Worker                                    var_default):
516*b6fb3261SAndroid Build Coastguard Worker  """Get var_name either from env, or user or default.
517*b6fb3261SAndroid Build Coastguard Worker
518*b6fb3261SAndroid Build Coastguard Worker  If var_name has been set as environment variable, use the preset value, else
519*b6fb3261SAndroid Build Coastguard Worker  ask for user input. If no input is provided, the default is used.
520*b6fb3261SAndroid Build Coastguard Worker
521*b6fb3261SAndroid Build Coastguard Worker  Args:
522*b6fb3261SAndroid Build Coastguard Worker    environ_cp: copy of the os.environ.
523*b6fb3261SAndroid Build Coastguard Worker    var_name: string for name of environment variable, e.g. "TF_NEED_CUDA".
524*b6fb3261SAndroid Build Coastguard Worker    ask_for_var: string for how to ask for user input.
525*b6fb3261SAndroid Build Coastguard Worker    var_default: default value string.
526*b6fb3261SAndroid Build Coastguard Worker
527*b6fb3261SAndroid Build Coastguard Worker  Returns:
528*b6fb3261SAndroid Build Coastguard Worker    string value for var_name
529*b6fb3261SAndroid Build Coastguard Worker  """
530*b6fb3261SAndroid Build Coastguard Worker  var = environ_cp.get(var_name)
531*b6fb3261SAndroid Build Coastguard Worker  if not var:
532*b6fb3261SAndroid Build Coastguard Worker    var = get_input(ask_for_var)
533*b6fb3261SAndroid Build Coastguard Worker    print('\n')
534*b6fb3261SAndroid Build Coastguard Worker  if not var:
535*b6fb3261SAndroid Build Coastguard Worker    var = var_default
536*b6fb3261SAndroid Build Coastguard Worker  return var
537*b6fb3261SAndroid Build Coastguard Worker
538*b6fb3261SAndroid Build Coastguard Worker
539*b6fb3261SAndroid Build Coastguard Workerdef set_clang_cuda_compiler_path(environ_cp):
540*b6fb3261SAndroid Build Coastguard Worker  """Set CLANG_CUDA_COMPILER_PATH."""
541*b6fb3261SAndroid Build Coastguard Worker  default_clang_path = which('clang') or ''
542*b6fb3261SAndroid Build Coastguard Worker  ask_clang_path = ('Please specify which clang should be used as device and '
543*b6fb3261SAndroid Build Coastguard Worker                    'host compiler. [Default is %s]: ') % default_clang_path
544*b6fb3261SAndroid Build Coastguard Worker
545*b6fb3261SAndroid Build Coastguard Worker  while True:
546*b6fb3261SAndroid Build Coastguard Worker    clang_cuda_compiler_path = get_from_env_or_user_or_default(
547*b6fb3261SAndroid Build Coastguard Worker        environ_cp, 'CLANG_CUDA_COMPILER_PATH', ask_clang_path,
548*b6fb3261SAndroid Build Coastguard Worker        default_clang_path)
549*b6fb3261SAndroid Build Coastguard Worker    if os.path.exists(clang_cuda_compiler_path):
550*b6fb3261SAndroid Build Coastguard Worker      break
551*b6fb3261SAndroid Build Coastguard Worker
552*b6fb3261SAndroid Build Coastguard Worker    # Reset and retry
553*b6fb3261SAndroid Build Coastguard Worker    print('Invalid clang path: %s cannot be found.' % clang_cuda_compiler_path)
554*b6fb3261SAndroid Build Coastguard Worker    environ_cp['CLANG_CUDA_COMPILER_PATH'] = ''
555*b6fb3261SAndroid Build Coastguard Worker
556*b6fb3261SAndroid Build Coastguard Worker  # Set CLANG_CUDA_COMPILER_PATH
557*b6fb3261SAndroid Build Coastguard Worker  environ_cp['CLANG_CUDA_COMPILER_PATH'] = clang_cuda_compiler_path
558*b6fb3261SAndroid Build Coastguard Worker  write_action_env_to_bazelrc('CLANG_CUDA_COMPILER_PATH',
559*b6fb3261SAndroid Build Coastguard Worker                              clang_cuda_compiler_path)
560*b6fb3261SAndroid Build Coastguard Worker
561*b6fb3261SAndroid Build Coastguard Worker
562*b6fb3261SAndroid Build Coastguard Workerdef prompt_loop_or_load_from_env(environ_cp,
563*b6fb3261SAndroid Build Coastguard Worker                                 var_name,
564*b6fb3261SAndroid Build Coastguard Worker                                 var_default,
565*b6fb3261SAndroid Build Coastguard Worker                                 ask_for_var,
566*b6fb3261SAndroid Build Coastguard Worker                                 check_success,
567*b6fb3261SAndroid Build Coastguard Worker                                 error_msg,
568*b6fb3261SAndroid Build Coastguard Worker                                 suppress_default_error=False,
569*b6fb3261SAndroid Build Coastguard Worker                                 resolve_symlinks=False,
570*b6fb3261SAndroid Build Coastguard Worker                                 n_ask_attempts=_DEFAULT_PROMPT_ASK_ATTEMPTS):
571*b6fb3261SAndroid Build Coastguard Worker  """Loop over user prompts for an ENV param until receiving a valid response.
572*b6fb3261SAndroid Build Coastguard Worker
573*b6fb3261SAndroid Build Coastguard Worker  For the env param var_name, read from the environment or verify user input
574*b6fb3261SAndroid Build Coastguard Worker  until receiving valid input. When done, set var_name in the environ_cp to its
575*b6fb3261SAndroid Build Coastguard Worker  new value.
576*b6fb3261SAndroid Build Coastguard Worker
577*b6fb3261SAndroid Build Coastguard Worker  Args:
578*b6fb3261SAndroid Build Coastguard Worker    environ_cp: (Dict) copy of the os.environ.
579*b6fb3261SAndroid Build Coastguard Worker    var_name: (String) string for name of environment variable, e.g. "TF_MYVAR".
580*b6fb3261SAndroid Build Coastguard Worker    var_default: (String) default value string.
581*b6fb3261SAndroid Build Coastguard Worker    ask_for_var: (String) string for how to ask for user input.
582*b6fb3261SAndroid Build Coastguard Worker    check_success: (Function) function that takes one argument and returns a
583*b6fb3261SAndroid Build Coastguard Worker      boolean. Should return True if the value provided is considered valid. May
584*b6fb3261SAndroid Build Coastguard Worker      contain a complex error message if error_msg does not provide enough
585*b6fb3261SAndroid Build Coastguard Worker      information. In that case, set suppress_default_error to True.
586*b6fb3261SAndroid Build Coastguard Worker    error_msg: (String) String with one and only one '%s'. Formatted with each
587*b6fb3261SAndroid Build Coastguard Worker      invalid response upon check_success(input) failure.
588*b6fb3261SAndroid Build Coastguard Worker    suppress_default_error: (Bool) Suppress the above error message in favor of
589*b6fb3261SAndroid Build Coastguard Worker      one from the check_success function.
590*b6fb3261SAndroid Build Coastguard Worker    resolve_symlinks: (Bool) Translate symbolic links into the real filepath.
591*b6fb3261SAndroid Build Coastguard Worker    n_ask_attempts: (Integer) Number of times to query for valid input before
592*b6fb3261SAndroid Build Coastguard Worker      raising an error and quitting.
593*b6fb3261SAndroid Build Coastguard Worker
594*b6fb3261SAndroid Build Coastguard Worker  Returns:
595*b6fb3261SAndroid Build Coastguard Worker    [String] The value of var_name after querying for input.
596*b6fb3261SAndroid Build Coastguard Worker
597*b6fb3261SAndroid Build Coastguard Worker  Raises:
598*b6fb3261SAndroid Build Coastguard Worker    UserInputError: if a query has been attempted n_ask_attempts times without
599*b6fb3261SAndroid Build Coastguard Worker      success, assume that the user has made a scripting error, and will
600*b6fb3261SAndroid Build Coastguard Worker      continue to provide invalid input. Raise the error to avoid infinitely
601*b6fb3261SAndroid Build Coastguard Worker      looping.
602*b6fb3261SAndroid Build Coastguard Worker  """
603*b6fb3261SAndroid Build Coastguard Worker  default = environ_cp.get(var_name) or var_default
604*b6fb3261SAndroid Build Coastguard Worker  full_query = '%s [Default is %s]: ' % (
605*b6fb3261SAndroid Build Coastguard Worker      ask_for_var,
606*b6fb3261SAndroid Build Coastguard Worker      default,
607*b6fb3261SAndroid Build Coastguard Worker  )
608*b6fb3261SAndroid Build Coastguard Worker
609*b6fb3261SAndroid Build Coastguard Worker  for _ in range(n_ask_attempts):
610*b6fb3261SAndroid Build Coastguard Worker    val = get_from_env_or_user_or_default(environ_cp, var_name, full_query,
611*b6fb3261SAndroid Build Coastguard Worker                                          default)
612*b6fb3261SAndroid Build Coastguard Worker    if check_success(val):
613*b6fb3261SAndroid Build Coastguard Worker      break
614*b6fb3261SAndroid Build Coastguard Worker    if not suppress_default_error:
615*b6fb3261SAndroid Build Coastguard Worker      print(error_msg % val)
616*b6fb3261SAndroid Build Coastguard Worker    environ_cp[var_name] = ''
617*b6fb3261SAndroid Build Coastguard Worker  else:
618*b6fb3261SAndroid Build Coastguard Worker    raise UserInputError('Invalid %s setting was provided %d times in a row. '
619*b6fb3261SAndroid Build Coastguard Worker                         'Assuming to be a scripting mistake.' %
620*b6fb3261SAndroid Build Coastguard Worker                         (var_name, n_ask_attempts))
621*b6fb3261SAndroid Build Coastguard Worker
622*b6fb3261SAndroid Build Coastguard Worker  if resolve_symlinks and os.path.islink(val):
623*b6fb3261SAndroid Build Coastguard Worker    val = os.path.realpath(val)
624*b6fb3261SAndroid Build Coastguard Worker  environ_cp[var_name] = val
625*b6fb3261SAndroid Build Coastguard Worker  return val
626*b6fb3261SAndroid Build Coastguard Worker
627*b6fb3261SAndroid Build Coastguard Worker
628*b6fb3261SAndroid Build Coastguard Workerdef create_android_ndk_rule(environ_cp):
629*b6fb3261SAndroid Build Coastguard Worker  """Set ANDROID_NDK_HOME and write Android NDK WORKSPACE rule."""
630*b6fb3261SAndroid Build Coastguard Worker  if is_windows() or is_cygwin():
631*b6fb3261SAndroid Build Coastguard Worker    default_ndk_path = cygpath('%s/Android/Sdk/ndk-bundle' %
632*b6fb3261SAndroid Build Coastguard Worker                               environ_cp['APPDATA'])
633*b6fb3261SAndroid Build Coastguard Worker  elif is_macos():
634*b6fb3261SAndroid Build Coastguard Worker    default_ndk_path = '%s/library/Android/Sdk/ndk-bundle' % environ_cp['HOME']
635*b6fb3261SAndroid Build Coastguard Worker  else:
636*b6fb3261SAndroid Build Coastguard Worker    default_ndk_path = '%s/Android/Sdk/ndk-bundle' % environ_cp['HOME']
637*b6fb3261SAndroid Build Coastguard Worker
638*b6fb3261SAndroid Build Coastguard Worker  def valid_ndk_path(path):
639*b6fb3261SAndroid Build Coastguard Worker    return (os.path.exists(path) and
640*b6fb3261SAndroid Build Coastguard Worker            os.path.exists(os.path.join(path, 'source.properties')))
641*b6fb3261SAndroid Build Coastguard Worker
642*b6fb3261SAndroid Build Coastguard Worker  android_ndk_home_path = prompt_loop_or_load_from_env(
643*b6fb3261SAndroid Build Coastguard Worker      environ_cp,
644*b6fb3261SAndroid Build Coastguard Worker      var_name='ANDROID_NDK_HOME',
645*b6fb3261SAndroid Build Coastguard Worker      var_default=default_ndk_path,
646*b6fb3261SAndroid Build Coastguard Worker      ask_for_var='Please specify the home path of the Android NDK to use.',
647*b6fb3261SAndroid Build Coastguard Worker      check_success=valid_ndk_path,
648*b6fb3261SAndroid Build Coastguard Worker      error_msg=('The path %s or its child file "source.properties" '
649*b6fb3261SAndroid Build Coastguard Worker                 'does not exist.'))
650*b6fb3261SAndroid Build Coastguard Worker  write_action_env_to_bazelrc('ANDROID_NDK_HOME', android_ndk_home_path)
651*b6fb3261SAndroid Build Coastguard Worker  write_action_env_to_bazelrc(
652*b6fb3261SAndroid Build Coastguard Worker      'ANDROID_NDK_API_LEVEL',
653*b6fb3261SAndroid Build Coastguard Worker      get_ndk_api_level(environ_cp, android_ndk_home_path))
654*b6fb3261SAndroid Build Coastguard Worker
655*b6fb3261SAndroid Build Coastguard Worker
656*b6fb3261SAndroid Build Coastguard Workerdef create_android_sdk_rule(environ_cp):
657*b6fb3261SAndroid Build Coastguard Worker  """Set Android variables and write Android SDK WORKSPACE rule."""
658*b6fb3261SAndroid Build Coastguard Worker  if is_windows() or is_cygwin():
659*b6fb3261SAndroid Build Coastguard Worker    default_sdk_path = cygpath('%s/Android/Sdk' % environ_cp['APPDATA'])
660*b6fb3261SAndroid Build Coastguard Worker  elif is_macos():
661*b6fb3261SAndroid Build Coastguard Worker    default_sdk_path = '%s/library/Android/Sdk' % environ_cp['HOME']
662*b6fb3261SAndroid Build Coastguard Worker  else:
663*b6fb3261SAndroid Build Coastguard Worker    default_sdk_path = '%s/Android/Sdk' % environ_cp['HOME']
664*b6fb3261SAndroid Build Coastguard Worker
665*b6fb3261SAndroid Build Coastguard Worker  def valid_sdk_path(path):
666*b6fb3261SAndroid Build Coastguard Worker    return (os.path.exists(path) and
667*b6fb3261SAndroid Build Coastguard Worker            os.path.exists(os.path.join(path, 'platforms')) and
668*b6fb3261SAndroid Build Coastguard Worker            os.path.exists(os.path.join(path, 'build-tools')))
669*b6fb3261SAndroid Build Coastguard Worker
670*b6fb3261SAndroid Build Coastguard Worker  android_sdk_home_path = prompt_loop_or_load_from_env(
671*b6fb3261SAndroid Build Coastguard Worker      environ_cp,
672*b6fb3261SAndroid Build Coastguard Worker      var_name='ANDROID_SDK_HOME',
673*b6fb3261SAndroid Build Coastguard Worker      var_default=default_sdk_path,
674*b6fb3261SAndroid Build Coastguard Worker      ask_for_var='Please specify the home path of the Android SDK to use.',
675*b6fb3261SAndroid Build Coastguard Worker      check_success=valid_sdk_path,
676*b6fb3261SAndroid Build Coastguard Worker      error_msg=('Either %s does not exist, or it does not contain the '
677*b6fb3261SAndroid Build Coastguard Worker                 'subdirectories "platforms" and "build-tools".'))
678*b6fb3261SAndroid Build Coastguard Worker
679*b6fb3261SAndroid Build Coastguard Worker  platforms = os.path.join(android_sdk_home_path, 'platforms')
680*b6fb3261SAndroid Build Coastguard Worker  api_levels = sorted(os.listdir(platforms))
681*b6fb3261SAndroid Build Coastguard Worker  api_levels = [x.replace('android-', '') for x in api_levels]
682*b6fb3261SAndroid Build Coastguard Worker
683*b6fb3261SAndroid Build Coastguard Worker  def valid_api_level(api_level):
684*b6fb3261SAndroid Build Coastguard Worker    return os.path.exists(
685*b6fb3261SAndroid Build Coastguard Worker        os.path.join(android_sdk_home_path, 'platforms',
686*b6fb3261SAndroid Build Coastguard Worker                     'android-' + api_level))
687*b6fb3261SAndroid Build Coastguard Worker
688*b6fb3261SAndroid Build Coastguard Worker  android_api_level = prompt_loop_or_load_from_env(
689*b6fb3261SAndroid Build Coastguard Worker      environ_cp,
690*b6fb3261SAndroid Build Coastguard Worker      var_name='ANDROID_API_LEVEL',
691*b6fb3261SAndroid Build Coastguard Worker      var_default=api_levels[-1],
692*b6fb3261SAndroid Build Coastguard Worker      ask_for_var=('Please specify the Android SDK API level to use. '
693*b6fb3261SAndroid Build Coastguard Worker                   '[Available levels: %s]') % api_levels,
694*b6fb3261SAndroid Build Coastguard Worker      check_success=valid_api_level,
695*b6fb3261SAndroid Build Coastguard Worker      error_msg='Android-%s is not present in the SDK path.')
696*b6fb3261SAndroid Build Coastguard Worker
697*b6fb3261SAndroid Build Coastguard Worker  build_tools = os.path.join(android_sdk_home_path, 'build-tools')
698*b6fb3261SAndroid Build Coastguard Worker  versions = sorted(os.listdir(build_tools))
699*b6fb3261SAndroid Build Coastguard Worker
700*b6fb3261SAndroid Build Coastguard Worker  def valid_build_tools(version):
701*b6fb3261SAndroid Build Coastguard Worker    return os.path.exists(
702*b6fb3261SAndroid Build Coastguard Worker        os.path.join(android_sdk_home_path, 'build-tools', version))
703*b6fb3261SAndroid Build Coastguard Worker
704*b6fb3261SAndroid Build Coastguard Worker  android_build_tools_version = prompt_loop_or_load_from_env(
705*b6fb3261SAndroid Build Coastguard Worker      environ_cp,
706*b6fb3261SAndroid Build Coastguard Worker      var_name='ANDROID_BUILD_TOOLS_VERSION',
707*b6fb3261SAndroid Build Coastguard Worker      var_default=versions[-1],
708*b6fb3261SAndroid Build Coastguard Worker      ask_for_var=('Please specify an Android build tools version to use. '
709*b6fb3261SAndroid Build Coastguard Worker                   '[Available versions: %s]') % versions,
710*b6fb3261SAndroid Build Coastguard Worker      check_success=valid_build_tools,
711*b6fb3261SAndroid Build Coastguard Worker      error_msg=('The selected SDK does not have build-tools version %s '
712*b6fb3261SAndroid Build Coastguard Worker                 'available.'))
713*b6fb3261SAndroid Build Coastguard Worker
714*b6fb3261SAndroid Build Coastguard Worker  write_action_env_to_bazelrc('ANDROID_BUILD_TOOLS_VERSION',
715*b6fb3261SAndroid Build Coastguard Worker                              android_build_tools_version)
716*b6fb3261SAndroid Build Coastguard Worker  write_action_env_to_bazelrc('ANDROID_SDK_API_LEVEL', android_api_level)
717*b6fb3261SAndroid Build Coastguard Worker  write_action_env_to_bazelrc('ANDROID_SDK_HOME', android_sdk_home_path)
718*b6fb3261SAndroid Build Coastguard Worker
719*b6fb3261SAndroid Build Coastguard Worker
720*b6fb3261SAndroid Build Coastguard Workerdef get_ndk_api_level(environ_cp, android_ndk_home_path):
721*b6fb3261SAndroid Build Coastguard Worker  """Gets the appropriate NDK API level to use for the provided Android NDK path."""
722*b6fb3261SAndroid Build Coastguard Worker
723*b6fb3261SAndroid Build Coastguard Worker  # First check to see if we're using a blessed version of the NDK.
724*b6fb3261SAndroid Build Coastguard Worker  properties_path = '%s/source.properties' % android_ndk_home_path
725*b6fb3261SAndroid Build Coastguard Worker  if is_windows() or is_cygwin():
726*b6fb3261SAndroid Build Coastguard Worker    properties_path = cygpath(properties_path)
727*b6fb3261SAndroid Build Coastguard Worker  with open(properties_path, 'r') as f:
728*b6fb3261SAndroid Build Coastguard Worker    filedata = f.read()
729*b6fb3261SAndroid Build Coastguard Worker
730*b6fb3261SAndroid Build Coastguard Worker  revision = re.search(r'Pkg.Revision = (\d+)', filedata)
731*b6fb3261SAndroid Build Coastguard Worker  if revision:
732*b6fb3261SAndroid Build Coastguard Worker    ndk_version = revision.group(1)
733*b6fb3261SAndroid Build Coastguard Worker  else:
734*b6fb3261SAndroid Build Coastguard Worker    raise Exception('Unable to parse NDK revision.')
735*b6fb3261SAndroid Build Coastguard Worker  if int(ndk_version) not in _SUPPORTED_ANDROID_NDK_VERSIONS:
736*b6fb3261SAndroid Build Coastguard Worker    print('WARNING: The NDK version in %s is %s, which is not '
737*b6fb3261SAndroid Build Coastguard Worker          'supported by Bazel (officially supported versions: %s). Please use '
738*b6fb3261SAndroid Build Coastguard Worker          'another version. Compiling Android targets may result in confusing '
739*b6fb3261SAndroid Build Coastguard Worker          'errors.\n' %
740*b6fb3261SAndroid Build Coastguard Worker          (android_ndk_home_path, ndk_version, _SUPPORTED_ANDROID_NDK_VERSIONS))
741*b6fb3261SAndroid Build Coastguard Worker
742*b6fb3261SAndroid Build Coastguard Worker  # Now grab the NDK API level to use. Note that this is different from the
743*b6fb3261SAndroid Build Coastguard Worker  # SDK API level, as the NDK API level is effectively the *min* target SDK
744*b6fb3261SAndroid Build Coastguard Worker  # version.
745*b6fb3261SAndroid Build Coastguard Worker  platforms = os.path.join(android_ndk_home_path, 'platforms')
746*b6fb3261SAndroid Build Coastguard Worker  api_levels = sorted(os.listdir(platforms))
747*b6fb3261SAndroid Build Coastguard Worker  api_levels = [
748*b6fb3261SAndroid Build Coastguard Worker      x.replace('android-', '') for x in api_levels if 'android-' in x
749*b6fb3261SAndroid Build Coastguard Worker  ]
750*b6fb3261SAndroid Build Coastguard Worker
751*b6fb3261SAndroid Build Coastguard Worker  def valid_api_level(api_level):
752*b6fb3261SAndroid Build Coastguard Worker    return os.path.exists(
753*b6fb3261SAndroid Build Coastguard Worker        os.path.join(android_ndk_home_path, 'platforms',
754*b6fb3261SAndroid Build Coastguard Worker                     'android-' + api_level))
755*b6fb3261SAndroid Build Coastguard Worker
756*b6fb3261SAndroid Build Coastguard Worker  android_ndk_api_level = prompt_loop_or_load_from_env(
757*b6fb3261SAndroid Build Coastguard Worker      environ_cp,
758*b6fb3261SAndroid Build Coastguard Worker      var_name='ANDROID_NDK_API_LEVEL',
759*b6fb3261SAndroid Build Coastguard Worker      var_default='21',  # 21 is required for ARM64 support.
760*b6fb3261SAndroid Build Coastguard Worker      ask_for_var=('Please specify the (min) Android NDK API level to use. '
761*b6fb3261SAndroid Build Coastguard Worker                   '[Available levels: %s]') % api_levels,
762*b6fb3261SAndroid Build Coastguard Worker      check_success=valid_api_level,
763*b6fb3261SAndroid Build Coastguard Worker      error_msg='Android-%s is not present in the NDK path.')
764*b6fb3261SAndroid Build Coastguard Worker
765*b6fb3261SAndroid Build Coastguard Worker  return android_ndk_api_level
766*b6fb3261SAndroid Build Coastguard Worker
767*b6fb3261SAndroid Build Coastguard Worker
768*b6fb3261SAndroid Build Coastguard Workerdef set_gcc_host_compiler_path(environ_cp):
769*b6fb3261SAndroid Build Coastguard Worker  """Set GCC_HOST_COMPILER_PATH."""
770*b6fb3261SAndroid Build Coastguard Worker  default_gcc_host_compiler_path = which('gcc') or ''
771*b6fb3261SAndroid Build Coastguard Worker  cuda_bin_symlink = '%s/bin/gcc' % environ_cp.get('CUDA_TOOLKIT_PATH')
772*b6fb3261SAndroid Build Coastguard Worker
773*b6fb3261SAndroid Build Coastguard Worker  if os.path.islink(cuda_bin_symlink):
774*b6fb3261SAndroid Build Coastguard Worker    # os.readlink is only available in linux
775*b6fb3261SAndroid Build Coastguard Worker    default_gcc_host_compiler_path = os.path.realpath(cuda_bin_symlink)
776*b6fb3261SAndroid Build Coastguard Worker
777*b6fb3261SAndroid Build Coastguard Worker  gcc_host_compiler_path = prompt_loop_or_load_from_env(
778*b6fb3261SAndroid Build Coastguard Worker      environ_cp,
779*b6fb3261SAndroid Build Coastguard Worker      var_name='GCC_HOST_COMPILER_PATH',
780*b6fb3261SAndroid Build Coastguard Worker      var_default=default_gcc_host_compiler_path,
781*b6fb3261SAndroid Build Coastguard Worker      ask_for_var='Please specify which gcc should be used by nvcc as the host '
782*b6fb3261SAndroid Build Coastguard Worker      'compiler.',
783*b6fb3261SAndroid Build Coastguard Worker      check_success=os.path.exists,
784*b6fb3261SAndroid Build Coastguard Worker      resolve_symlinks=True,
785*b6fb3261SAndroid Build Coastguard Worker      error_msg='Invalid gcc path. %s cannot be found.',
786*b6fb3261SAndroid Build Coastguard Worker  )
787*b6fb3261SAndroid Build Coastguard Worker
788*b6fb3261SAndroid Build Coastguard Worker  write_action_env_to_bazelrc('GCC_HOST_COMPILER_PATH', gcc_host_compiler_path)
789*b6fb3261SAndroid Build Coastguard Worker
790*b6fb3261SAndroid Build Coastguard Worker
791*b6fb3261SAndroid Build Coastguard Workerdef set_tf_cuda_paths(environ_cp):
792*b6fb3261SAndroid Build Coastguard Worker  """Set TF_CUDA_PATHS."""
793*b6fb3261SAndroid Build Coastguard Worker  ask_cuda_paths = (
794*b6fb3261SAndroid Build Coastguard Worker      'Please specify the comma-separated list of base paths to look for CUDA '
795*b6fb3261SAndroid Build Coastguard Worker      'libraries and headers. [Leave empty to use the default]: ')
796*b6fb3261SAndroid Build Coastguard Worker  tf_cuda_paths = get_from_env_or_user_or_default(environ_cp, 'TF_CUDA_PATHS',
797*b6fb3261SAndroid Build Coastguard Worker                                                  ask_cuda_paths, '')
798*b6fb3261SAndroid Build Coastguard Worker  if tf_cuda_paths:
799*b6fb3261SAndroid Build Coastguard Worker    environ_cp['TF_CUDA_PATHS'] = tf_cuda_paths
800*b6fb3261SAndroid Build Coastguard Worker
801*b6fb3261SAndroid Build Coastguard Worker
802*b6fb3261SAndroid Build Coastguard Workerdef set_tf_cuda_version(environ_cp):
803*b6fb3261SAndroid Build Coastguard Worker  """Set TF_CUDA_VERSION."""
804*b6fb3261SAndroid Build Coastguard Worker  ask_cuda_version = (
805*b6fb3261SAndroid Build Coastguard Worker      'Please specify the CUDA SDK version you want to use. '
806*b6fb3261SAndroid Build Coastguard Worker      '[Leave empty to default to CUDA %s]: ') % _DEFAULT_CUDA_VERSION
807*b6fb3261SAndroid Build Coastguard Worker  tf_cuda_version = get_from_env_or_user_or_default(environ_cp,
808*b6fb3261SAndroid Build Coastguard Worker                                                    'TF_CUDA_VERSION',
809*b6fb3261SAndroid Build Coastguard Worker                                                    ask_cuda_version,
810*b6fb3261SAndroid Build Coastguard Worker                                                    _DEFAULT_CUDA_VERSION)
811*b6fb3261SAndroid Build Coastguard Worker  environ_cp['TF_CUDA_VERSION'] = tf_cuda_version
812*b6fb3261SAndroid Build Coastguard Worker
813*b6fb3261SAndroid Build Coastguard Worker
814*b6fb3261SAndroid Build Coastguard Workerdef set_tf_cudnn_version(environ_cp):
815*b6fb3261SAndroid Build Coastguard Worker  """Set TF_CUDNN_VERSION."""
816*b6fb3261SAndroid Build Coastguard Worker  ask_cudnn_version = (
817*b6fb3261SAndroid Build Coastguard Worker      'Please specify the cuDNN version you want to use. '
818*b6fb3261SAndroid Build Coastguard Worker      '[Leave empty to default to cuDNN %s]: ') % _DEFAULT_CUDNN_VERSION
819*b6fb3261SAndroid Build Coastguard Worker  tf_cudnn_version = get_from_env_or_user_or_default(environ_cp,
820*b6fb3261SAndroid Build Coastguard Worker                                                     'TF_CUDNN_VERSION',
821*b6fb3261SAndroid Build Coastguard Worker                                                     ask_cudnn_version,
822*b6fb3261SAndroid Build Coastguard Worker                                                     _DEFAULT_CUDNN_VERSION)
823*b6fb3261SAndroid Build Coastguard Worker  environ_cp['TF_CUDNN_VERSION'] = tf_cudnn_version
824*b6fb3261SAndroid Build Coastguard Worker
825*b6fb3261SAndroid Build Coastguard Worker
826*b6fb3261SAndroid Build Coastguard Workerdef set_tf_tensorrt_version(environ_cp):
827*b6fb3261SAndroid Build Coastguard Worker  """Set TF_TENSORRT_VERSION."""
828*b6fb3261SAndroid Build Coastguard Worker  if not (is_linux() or is_windows()):
829*b6fb3261SAndroid Build Coastguard Worker    raise ValueError('Currently TensorRT is only supported on Linux platform.')
830*b6fb3261SAndroid Build Coastguard Worker
831*b6fb3261SAndroid Build Coastguard Worker  if not int(environ_cp.get('TF_NEED_TENSORRT', False)):
832*b6fb3261SAndroid Build Coastguard Worker    return
833*b6fb3261SAndroid Build Coastguard Worker
834*b6fb3261SAndroid Build Coastguard Worker  ask_tensorrt_version = (
835*b6fb3261SAndroid Build Coastguard Worker      'Please specify the TensorRT version you want to use. '
836*b6fb3261SAndroid Build Coastguard Worker      '[Leave empty to default to TensorRT %s]: ') % _DEFAULT_TENSORRT_VERSION
837*b6fb3261SAndroid Build Coastguard Worker  tf_tensorrt_version = get_from_env_or_user_or_default(
838*b6fb3261SAndroid Build Coastguard Worker      environ_cp, 'TF_TENSORRT_VERSION', ask_tensorrt_version,
839*b6fb3261SAndroid Build Coastguard Worker      _DEFAULT_TENSORRT_VERSION)
840*b6fb3261SAndroid Build Coastguard Worker  environ_cp['TF_TENSORRT_VERSION'] = tf_tensorrt_version
841*b6fb3261SAndroid Build Coastguard Worker
842*b6fb3261SAndroid Build Coastguard Worker
843*b6fb3261SAndroid Build Coastguard Workerdef set_tf_nccl_version(environ_cp):
844*b6fb3261SAndroid Build Coastguard Worker  """Set TF_NCCL_VERSION."""
845*b6fb3261SAndroid Build Coastguard Worker  if not is_linux():
846*b6fb3261SAndroid Build Coastguard Worker    raise ValueError('Currently NCCL is only supported on Linux platform.')
847*b6fb3261SAndroid Build Coastguard Worker
848*b6fb3261SAndroid Build Coastguard Worker  if 'TF_NCCL_VERSION' in environ_cp:
849*b6fb3261SAndroid Build Coastguard Worker    return
850*b6fb3261SAndroid Build Coastguard Worker
851*b6fb3261SAndroid Build Coastguard Worker  ask_nccl_version = (
852*b6fb3261SAndroid Build Coastguard Worker      'Please specify the locally installed NCCL version you want to use. '
853*b6fb3261SAndroid Build Coastguard Worker      '[Leave empty to use http://github.com/nvidia/nccl]: ')
854*b6fb3261SAndroid Build Coastguard Worker  tf_nccl_version = get_from_env_or_user_or_default(environ_cp,
855*b6fb3261SAndroid Build Coastguard Worker                                                    'TF_NCCL_VERSION',
856*b6fb3261SAndroid Build Coastguard Worker                                                    ask_nccl_version, '')
857*b6fb3261SAndroid Build Coastguard Worker  environ_cp['TF_NCCL_VERSION'] = tf_nccl_version
858*b6fb3261SAndroid Build Coastguard Worker
859*b6fb3261SAndroid Build Coastguard Worker
860*b6fb3261SAndroid Build Coastguard Workerdef get_native_cuda_compute_capabilities(environ_cp):
861*b6fb3261SAndroid Build Coastguard Worker  """Get native cuda compute capabilities.
862*b6fb3261SAndroid Build Coastguard Worker
863*b6fb3261SAndroid Build Coastguard Worker  Args:
864*b6fb3261SAndroid Build Coastguard Worker    environ_cp: copy of the os.environ.
865*b6fb3261SAndroid Build Coastguard Worker
866*b6fb3261SAndroid Build Coastguard Worker  Returns:
867*b6fb3261SAndroid Build Coastguard Worker    string of native cuda compute capabilities, separated by comma.
868*b6fb3261SAndroid Build Coastguard Worker  """
869*b6fb3261SAndroid Build Coastguard Worker  device_query_bin = os.path.join(
870*b6fb3261SAndroid Build Coastguard Worker      environ_cp.get('CUDA_TOOLKIT_PATH'), 'extras/demo_suite/deviceQuery')
871*b6fb3261SAndroid Build Coastguard Worker  if os.path.isfile(device_query_bin) and os.access(device_query_bin, os.X_OK):
872*b6fb3261SAndroid Build Coastguard Worker    try:
873*b6fb3261SAndroid Build Coastguard Worker      output = run_shell(device_query_bin).split('\n')
874*b6fb3261SAndroid Build Coastguard Worker      pattern = re.compile('[0-9]*\\.[0-9]*')
875*b6fb3261SAndroid Build Coastguard Worker      output = [pattern.search(x) for x in output if 'Capability' in x]
876*b6fb3261SAndroid Build Coastguard Worker      output = ','.join(x.group() for x in output if x is not None)
877*b6fb3261SAndroid Build Coastguard Worker    except subprocess.CalledProcessError:
878*b6fb3261SAndroid Build Coastguard Worker      output = ''
879*b6fb3261SAndroid Build Coastguard Worker  else:
880*b6fb3261SAndroid Build Coastguard Worker    output = ''
881*b6fb3261SAndroid Build Coastguard Worker  return output
882*b6fb3261SAndroid Build Coastguard Worker
883*b6fb3261SAndroid Build Coastguard Worker
884*b6fb3261SAndroid Build Coastguard Workerdef set_tf_cuda_compute_capabilities(environ_cp):
885*b6fb3261SAndroid Build Coastguard Worker  """Set TF_CUDA_COMPUTE_CAPABILITIES."""
886*b6fb3261SAndroid Build Coastguard Worker  while True:
887*b6fb3261SAndroid Build Coastguard Worker    native_cuda_compute_capabilities = get_native_cuda_compute_capabilities(
888*b6fb3261SAndroid Build Coastguard Worker        environ_cp)
889*b6fb3261SAndroid Build Coastguard Worker    if not native_cuda_compute_capabilities:
890*b6fb3261SAndroid Build Coastguard Worker      default_cuda_compute_capabilities = _DEFAULT_CUDA_COMPUTE_CAPABILITIES
891*b6fb3261SAndroid Build Coastguard Worker    else:
892*b6fb3261SAndroid Build Coastguard Worker      default_cuda_compute_capabilities = native_cuda_compute_capabilities
893*b6fb3261SAndroid Build Coastguard Worker
894*b6fb3261SAndroid Build Coastguard Worker    ask_cuda_compute_capabilities = (
895*b6fb3261SAndroid Build Coastguard Worker        'Please specify a list of comma-separated CUDA compute capabilities '
896*b6fb3261SAndroid Build Coastguard Worker        'you want to build with.\nYou can find the compute capability of your '
897*b6fb3261SAndroid Build Coastguard Worker        'device at: https://developer.nvidia.com/cuda-gpus. Each capability '
898*b6fb3261SAndroid Build Coastguard Worker        'can be specified as "x.y" or "compute_xy" to include both virtual and'
899*b6fb3261SAndroid Build Coastguard Worker        ' binary GPU code, or as "sm_xy" to only include the binary '
900*b6fb3261SAndroid Build Coastguard Worker        'code.\nPlease note that each additional compute capability '
901*b6fb3261SAndroid Build Coastguard Worker        'significantly increases your build time and binary size, and that '
902*b6fb3261SAndroid Build Coastguard Worker        'TensorFlow only supports compute capabilities >= 3.5 [Default is: '
903*b6fb3261SAndroid Build Coastguard Worker        '%s]: ' % default_cuda_compute_capabilities)
904*b6fb3261SAndroid Build Coastguard Worker    tf_cuda_compute_capabilities = get_from_env_or_user_or_default(
905*b6fb3261SAndroid Build Coastguard Worker        environ_cp, 'TF_CUDA_COMPUTE_CAPABILITIES',
906*b6fb3261SAndroid Build Coastguard Worker        ask_cuda_compute_capabilities, default_cuda_compute_capabilities)
907*b6fb3261SAndroid Build Coastguard Worker    # Check whether all capabilities from the input is valid
908*b6fb3261SAndroid Build Coastguard Worker    all_valid = True
909*b6fb3261SAndroid Build Coastguard Worker    # Remove all whitespace characters before splitting the string
910*b6fb3261SAndroid Build Coastguard Worker    # that users may insert by accident, as this will result in error
911*b6fb3261SAndroid Build Coastguard Worker    tf_cuda_compute_capabilities = ''.join(tf_cuda_compute_capabilities.split())
912*b6fb3261SAndroid Build Coastguard Worker    for compute_capability in tf_cuda_compute_capabilities.split(','):
913*b6fb3261SAndroid Build Coastguard Worker      m = re.match('[0-9]+.[0-9]+', compute_capability)
914*b6fb3261SAndroid Build Coastguard Worker      if not m:
915*b6fb3261SAndroid Build Coastguard Worker        # We now support sm_35,sm_50,sm_60,compute_70.
916*b6fb3261SAndroid Build Coastguard Worker        sm_compute_match = re.match('(sm|compute)_?([0-9]+[0-9]+)',
917*b6fb3261SAndroid Build Coastguard Worker                                    compute_capability)
918*b6fb3261SAndroid Build Coastguard Worker        if not sm_compute_match:
919*b6fb3261SAndroid Build Coastguard Worker          print('Invalid compute capability: %s' % compute_capability)
920*b6fb3261SAndroid Build Coastguard Worker          all_valid = False
921*b6fb3261SAndroid Build Coastguard Worker        else:
922*b6fb3261SAndroid Build Coastguard Worker          ver = int(sm_compute_match.group(2))
923*b6fb3261SAndroid Build Coastguard Worker          if ver < 30:
924*b6fb3261SAndroid Build Coastguard Worker            print(
925*b6fb3261SAndroid Build Coastguard Worker                'ERROR: TensorFlow only supports small CUDA compute'
926*b6fb3261SAndroid Build Coastguard Worker                ' capabilities of sm_30 and higher. Please re-specify the list'
927*b6fb3261SAndroid Build Coastguard Worker                ' of compute capabilities excluding version %s.' % ver)
928*b6fb3261SAndroid Build Coastguard Worker            all_valid = False
929*b6fb3261SAndroid Build Coastguard Worker          if ver < 35:
930*b6fb3261SAndroid Build Coastguard Worker            print('WARNING: XLA does not support CUDA compute capabilities '
931*b6fb3261SAndroid Build Coastguard Worker                  'lower than sm_35. Disable XLA when running on older GPUs.')
932*b6fb3261SAndroid Build Coastguard Worker      else:
933*b6fb3261SAndroid Build Coastguard Worker        ver = float(m.group(0))
934*b6fb3261SAndroid Build Coastguard Worker        if ver < 3.0:
935*b6fb3261SAndroid Build Coastguard Worker          print('ERROR: TensorFlow only supports CUDA compute capabilities 3.0 '
936*b6fb3261SAndroid Build Coastguard Worker                'and higher. Please re-specify the list of compute '
937*b6fb3261SAndroid Build Coastguard Worker                'capabilities excluding version %s.' % ver)
938*b6fb3261SAndroid Build Coastguard Worker          all_valid = False
939*b6fb3261SAndroid Build Coastguard Worker        if ver < 3.5:
940*b6fb3261SAndroid Build Coastguard Worker          print('WARNING: XLA does not support CUDA compute capabilities '
941*b6fb3261SAndroid Build Coastguard Worker                'lower than 3.5. Disable XLA when running on older GPUs.')
942*b6fb3261SAndroid Build Coastguard Worker
943*b6fb3261SAndroid Build Coastguard Worker    if all_valid:
944*b6fb3261SAndroid Build Coastguard Worker      break
945*b6fb3261SAndroid Build Coastguard Worker
946*b6fb3261SAndroid Build Coastguard Worker    # Reset and Retry
947*b6fb3261SAndroid Build Coastguard Worker    environ_cp['TF_CUDA_COMPUTE_CAPABILITIES'] = ''
948*b6fb3261SAndroid Build Coastguard Worker
949*b6fb3261SAndroid Build Coastguard Worker  # Set TF_CUDA_COMPUTE_CAPABILITIES
950*b6fb3261SAndroid Build Coastguard Worker  environ_cp['TF_CUDA_COMPUTE_CAPABILITIES'] = tf_cuda_compute_capabilities
951*b6fb3261SAndroid Build Coastguard Worker  write_action_env_to_bazelrc('TF_CUDA_COMPUTE_CAPABILITIES',
952*b6fb3261SAndroid Build Coastguard Worker                              tf_cuda_compute_capabilities)
953*b6fb3261SAndroid Build Coastguard Worker
954*b6fb3261SAndroid Build Coastguard Worker
955*b6fb3261SAndroid Build Coastguard Workerdef set_other_cuda_vars(environ_cp):
956*b6fb3261SAndroid Build Coastguard Worker  """Set other CUDA related variables."""
957*b6fb3261SAndroid Build Coastguard Worker  # If CUDA is enabled, always use GPU during build and test.
958*b6fb3261SAndroid Build Coastguard Worker  if environ_cp.get('TF_CUDA_CLANG') == '1':
959*b6fb3261SAndroid Build Coastguard Worker    write_to_bazelrc('build --config=cuda_clang')
960*b6fb3261SAndroid Build Coastguard Worker  else:
961*b6fb3261SAndroid Build Coastguard Worker    write_to_bazelrc('build --config=cuda')
962*b6fb3261SAndroid Build Coastguard Worker
963*b6fb3261SAndroid Build Coastguard Worker
964*b6fb3261SAndroid Build Coastguard Workerdef system_specific_test_config(environ_cp):
965*b6fb3261SAndroid Build Coastguard Worker  """Add default build and test flags required for TF tests to bazelrc."""
966*b6fb3261SAndroid Build Coastguard Worker  write_to_bazelrc('test --flaky_test_attempts=3')
967*b6fb3261SAndroid Build Coastguard Worker  write_to_bazelrc('test --test_size_filters=small,medium')
968*b6fb3261SAndroid Build Coastguard Worker
969*b6fb3261SAndroid Build Coastguard Worker  # Each instance of --test_tag_filters or --build_tag_filters overrides all
970*b6fb3261SAndroid Build Coastguard Worker  # previous instances, so we need to build up a complete list and write a
971*b6fb3261SAndroid Build Coastguard Worker  # single list of filters for the .bazelrc file.
972*b6fb3261SAndroid Build Coastguard Worker
973*b6fb3261SAndroid Build Coastguard Worker  # Filters to use with both --test_tag_filters and --build_tag_filters
974*b6fb3261SAndroid Build Coastguard Worker  test_and_build_filters = ['-benchmark-test', '-no_oss']
975*b6fb3261SAndroid Build Coastguard Worker  # Additional filters for --test_tag_filters beyond those in
976*b6fb3261SAndroid Build Coastguard Worker  # test_and_build_filters
977*b6fb3261SAndroid Build Coastguard Worker  test_only_filters = ['-oss_serial']
978*b6fb3261SAndroid Build Coastguard Worker  if is_windows():
979*b6fb3261SAndroid Build Coastguard Worker    test_and_build_filters.append('-no_windows')
980*b6fb3261SAndroid Build Coastguard Worker    if ((environ_cp.get('TF_NEED_CUDA', None) == '1') or
981*b6fb3261SAndroid Build Coastguard Worker        (environ_cp.get('TF_NEED_ROCM', None) == '1')):
982*b6fb3261SAndroid Build Coastguard Worker      test_and_build_filters += ['-no_windows_gpu', '-no_gpu']
983*b6fb3261SAndroid Build Coastguard Worker    else:
984*b6fb3261SAndroid Build Coastguard Worker      test_and_build_filters.append('-gpu')
985*b6fb3261SAndroid Build Coastguard Worker  elif is_macos():
986*b6fb3261SAndroid Build Coastguard Worker    test_and_build_filters += ['-gpu', '-nomac', '-no_mac']
987*b6fb3261SAndroid Build Coastguard Worker  elif is_linux():
988*b6fb3261SAndroid Build Coastguard Worker    if ((environ_cp.get('TF_NEED_CUDA', None) == '1') or
989*b6fb3261SAndroid Build Coastguard Worker        (environ_cp.get('TF_NEED_ROCM', None) == '1')):
990*b6fb3261SAndroid Build Coastguard Worker      test_and_build_filters.append('-no_gpu')
991*b6fb3261SAndroid Build Coastguard Worker      write_to_bazelrc('test --test_env=LD_LIBRARY_PATH')
992*b6fb3261SAndroid Build Coastguard Worker    else:
993*b6fb3261SAndroid Build Coastguard Worker      test_and_build_filters.append('-gpu')
994*b6fb3261SAndroid Build Coastguard Worker
995*b6fb3261SAndroid Build Coastguard Worker  # Disable tests with "v1only" tag in "v2" Bazel config, but not in "v1" config
996*b6fb3261SAndroid Build Coastguard Worker  write_to_bazelrc('test:v1 --test_tag_filters=%s' %
997*b6fb3261SAndroid Build Coastguard Worker                   ','.join(test_and_build_filters + test_only_filters))
998*b6fb3261SAndroid Build Coastguard Worker  write_to_bazelrc('test:v1 --build_tag_filters=%s' %
999*b6fb3261SAndroid Build Coastguard Worker                   ','.join(test_and_build_filters))
1000*b6fb3261SAndroid Build Coastguard Worker  write_to_bazelrc(
1001*b6fb3261SAndroid Build Coastguard Worker      'test:v2 --test_tag_filters=%s' %
1002*b6fb3261SAndroid Build Coastguard Worker      ','.join(test_and_build_filters + test_only_filters + ['-v1only']))
1003*b6fb3261SAndroid Build Coastguard Worker  write_to_bazelrc('test:v2 --build_tag_filters=%s' %
1004*b6fb3261SAndroid Build Coastguard Worker                   ','.join(test_and_build_filters + ['-v1only']))
1005*b6fb3261SAndroid Build Coastguard Worker
1006*b6fb3261SAndroid Build Coastguard Worker
1007*b6fb3261SAndroid Build Coastguard Workerdef set_system_libs_flag(environ_cp):
1008*b6fb3261SAndroid Build Coastguard Worker  syslibs = environ_cp.get('TF_SYSTEM_LIBS', '')
1009*b6fb3261SAndroid Build Coastguard Worker  if syslibs:
1010*b6fb3261SAndroid Build Coastguard Worker    if ',' in syslibs:
1011*b6fb3261SAndroid Build Coastguard Worker      syslibs = ','.join(sorted(syslibs.split(',')))
1012*b6fb3261SAndroid Build Coastguard Worker    else:
1013*b6fb3261SAndroid Build Coastguard Worker      syslibs = ','.join(sorted(syslibs.split()))
1014*b6fb3261SAndroid Build Coastguard Worker    write_action_env_to_bazelrc('TF_SYSTEM_LIBS', syslibs)
1015*b6fb3261SAndroid Build Coastguard Worker
1016*b6fb3261SAndroid Build Coastguard Worker  for varname in ('PREFIX', 'LIBDIR', 'INCLUDEDIR', 'PROTOBUF_INCLUDE_PATH'):
1017*b6fb3261SAndroid Build Coastguard Worker    if varname in environ_cp:
1018*b6fb3261SAndroid Build Coastguard Worker      write_to_bazelrc('build --define=%s=%s' % (varname, environ_cp[varname]))
1019*b6fb3261SAndroid Build Coastguard Worker
1020*b6fb3261SAndroid Build Coastguard Worker
1021*b6fb3261SAndroid Build Coastguard Workerdef set_windows_build_flags(environ_cp):
1022*b6fb3261SAndroid Build Coastguard Worker  """Set Windows specific build options."""
1023*b6fb3261SAndroid Build Coastguard Worker
1024*b6fb3261SAndroid Build Coastguard Worker  # First available in VS 16.4. Speeds up Windows compile times by a lot. See
1025*b6fb3261SAndroid Build Coastguard Worker  # https://groups.google.com/a/tensorflow.org/d/topic/build/SsW98Eo7l3o/discussion
1026*b6fb3261SAndroid Build Coastguard Worker  # pylint: disable=line-too-long
1027*b6fb3261SAndroid Build Coastguard Worker  write_to_bazelrc(
1028*b6fb3261SAndroid Build Coastguard Worker      'build --copt=/d2ReducedOptimizeHugeFunctions --host_copt=/d2ReducedOptimizeHugeFunctions'
1029*b6fb3261SAndroid Build Coastguard Worker  )
1030*b6fb3261SAndroid Build Coastguard Worker
1031*b6fb3261SAndroid Build Coastguard Worker  if get_var(
1032*b6fb3261SAndroid Build Coastguard Worker      environ_cp, 'TF_OVERRIDE_EIGEN_STRONG_INLINE', 'Eigen strong inline',
1033*b6fb3261SAndroid Build Coastguard Worker      True, ('Would you like to override eigen strong inline for some C++ '
1034*b6fb3261SAndroid Build Coastguard Worker             'compilation to reduce the compilation time?'),
1035*b6fb3261SAndroid Build Coastguard Worker      'Eigen strong inline overridden.', 'Not overriding eigen strong inline, '
1036*b6fb3261SAndroid Build Coastguard Worker      'some compilations could take more than 20 mins.'):
1037*b6fb3261SAndroid Build Coastguard Worker    # Due to a known MSVC compiler issue
1038*b6fb3261SAndroid Build Coastguard Worker    # https://github.com/tensorflow/tensorflow/issues/10521
1039*b6fb3261SAndroid Build Coastguard Worker    # Overriding eigen strong inline speeds up the compiling of
1040*b6fb3261SAndroid Build Coastguard Worker    # conv_grad_ops_3d.cc and conv_ops_3d.cc by 20 minutes,
1041*b6fb3261SAndroid Build Coastguard Worker    # but this also hurts the performance. Let users decide what they want.
1042*b6fb3261SAndroid Build Coastguard Worker    write_to_bazelrc('build --define=override_eigen_strong_inline=true')
1043*b6fb3261SAndroid Build Coastguard Worker
1044*b6fb3261SAndroid Build Coastguard Worker
1045*b6fb3261SAndroid Build Coastguard Workerdef config_info_line(name, help_text):
1046*b6fb3261SAndroid Build Coastguard Worker  """Helper function to print formatted help text for Bazel config options."""
1047*b6fb3261SAndroid Build Coastguard Worker  print('\t--config=%-12s\t# %s' % (name, help_text))
1048*b6fb3261SAndroid Build Coastguard Worker
1049*b6fb3261SAndroid Build Coastguard Worker
1050*b6fb3261SAndroid Build Coastguard Workerdef configure_ios(environ_cp):
1051*b6fb3261SAndroid Build Coastguard Worker  """Configures TensorFlow for iOS builds."""
1052*b6fb3261SAndroid Build Coastguard Worker  if not is_macos():
1053*b6fb3261SAndroid Build Coastguard Worker    return
1054*b6fb3261SAndroid Build Coastguard Worker  if not get_var(environ_cp, 'TF_CONFIGURE_IOS', 'iOS', False):
1055*b6fb3261SAndroid Build Coastguard Worker    return
1056*b6fb3261SAndroid Build Coastguard Worker  for filepath in APPLE_BAZEL_FILES:
1057*b6fb3261SAndroid Build Coastguard Worker    existing_filepath = os.path.join(_TF_WORKSPACE_ROOT, filepath + '.apple')
1058*b6fb3261SAndroid Build Coastguard Worker    renamed_filepath = os.path.join(_TF_WORKSPACE_ROOT, filepath)
1059*b6fb3261SAndroid Build Coastguard Worker    symlink_force(existing_filepath, renamed_filepath)
1060*b6fb3261SAndroid Build Coastguard Worker  for filepath in IOS_FILES:
1061*b6fb3261SAndroid Build Coastguard Worker    filename = os.path.basename(filepath)
1062*b6fb3261SAndroid Build Coastguard Worker    new_filepath = os.path.join(_TF_WORKSPACE_ROOT, filename)
1063*b6fb3261SAndroid Build Coastguard Worker    symlink_force(filepath, new_filepath)
1064*b6fb3261SAndroid Build Coastguard Worker
1065*b6fb3261SAndroid Build Coastguard Worker
1066*b6fb3261SAndroid Build Coastguard Workerdef validate_cuda_config(environ_cp):
1067*b6fb3261SAndroid Build Coastguard Worker  """Run find_cuda_config.py and return cuda_toolkit_path, or None."""
1068*b6fb3261SAndroid Build Coastguard Worker
1069*b6fb3261SAndroid Build Coastguard Worker  def maybe_encode_env(env):
1070*b6fb3261SAndroid Build Coastguard Worker    """Encodes unicode in env to str on Windows python 2.x."""
1071*b6fb3261SAndroid Build Coastguard Worker    if not is_windows() or sys.version_info[0] != 2:
1072*b6fb3261SAndroid Build Coastguard Worker      return env
1073*b6fb3261SAndroid Build Coastguard Worker    for k, v in env.items():
1074*b6fb3261SAndroid Build Coastguard Worker      if isinstance(k, unicode):
1075*b6fb3261SAndroid Build Coastguard Worker        k = k.encode('ascii')
1076*b6fb3261SAndroid Build Coastguard Worker      if isinstance(v, unicode):
1077*b6fb3261SAndroid Build Coastguard Worker        v = v.encode('ascii')
1078*b6fb3261SAndroid Build Coastguard Worker      env[k] = v
1079*b6fb3261SAndroid Build Coastguard Worker    return env
1080*b6fb3261SAndroid Build Coastguard Worker
1081*b6fb3261SAndroid Build Coastguard Worker  cuda_libraries = ['cuda', 'cudnn']
1082*b6fb3261SAndroid Build Coastguard Worker  if is_linux():
1083*b6fb3261SAndroid Build Coastguard Worker    if int(environ_cp.get('TF_NEED_TENSORRT', False)):
1084*b6fb3261SAndroid Build Coastguard Worker      cuda_libraries.append('tensorrt')
1085*b6fb3261SAndroid Build Coastguard Worker    if environ_cp.get('TF_NCCL_VERSION', None):
1086*b6fb3261SAndroid Build Coastguard Worker      cuda_libraries.append('nccl')
1087*b6fb3261SAndroid Build Coastguard Worker  if is_windows():
1088*b6fb3261SAndroid Build Coastguard Worker    if int(environ_cp.get('TF_NEED_TENSORRT', False)):
1089*b6fb3261SAndroid Build Coastguard Worker      cuda_libraries.append('tensorrt')
1090*b6fb3261SAndroid Build Coastguard Worker      print('WARNING: TensorRT support on Windows is experimental\n')
1091*b6fb3261SAndroid Build Coastguard Worker
1092*b6fb3261SAndroid Build Coastguard Worker  paths = glob.glob('**/third_party/gpus/find_cuda_config.py', recursive=True)
1093*b6fb3261SAndroid Build Coastguard Worker  if not paths:
1094*b6fb3261SAndroid Build Coastguard Worker    raise FileNotFoundError(
1095*b6fb3261SAndroid Build Coastguard Worker        "Can't find 'find_cuda_config.py' script inside working directory")
1096*b6fb3261SAndroid Build Coastguard Worker  proc = subprocess.Popen(
1097*b6fb3261SAndroid Build Coastguard Worker      [environ_cp['PYTHON_BIN_PATH'], paths[0]] + cuda_libraries,
1098*b6fb3261SAndroid Build Coastguard Worker      stdout=subprocess.PIPE,
1099*b6fb3261SAndroid Build Coastguard Worker      env=maybe_encode_env(environ_cp))
1100*b6fb3261SAndroid Build Coastguard Worker
1101*b6fb3261SAndroid Build Coastguard Worker  if proc.wait():
1102*b6fb3261SAndroid Build Coastguard Worker    # Errors from find_cuda_config.py were sent to stderr.
1103*b6fb3261SAndroid Build Coastguard Worker    print('Asking for detailed CUDA configuration...\n')
1104*b6fb3261SAndroid Build Coastguard Worker    return False
1105*b6fb3261SAndroid Build Coastguard Worker
1106*b6fb3261SAndroid Build Coastguard Worker  config = dict(
1107*b6fb3261SAndroid Build Coastguard Worker      tuple(line.decode('ascii').rstrip().split(': ')) for line in proc.stdout)
1108*b6fb3261SAndroid Build Coastguard Worker
1109*b6fb3261SAndroid Build Coastguard Worker  print('Found CUDA %s in:' % config['cuda_version'])
1110*b6fb3261SAndroid Build Coastguard Worker  print('    %s' % config['cuda_library_dir'])
1111*b6fb3261SAndroid Build Coastguard Worker  print('    %s' % config['cuda_include_dir'])
1112*b6fb3261SAndroid Build Coastguard Worker
1113*b6fb3261SAndroid Build Coastguard Worker  print('Found cuDNN %s in:' % config['cudnn_version'])
1114*b6fb3261SAndroid Build Coastguard Worker  print('    %s' % config['cudnn_library_dir'])
1115*b6fb3261SAndroid Build Coastguard Worker  print('    %s' % config['cudnn_include_dir'])
1116*b6fb3261SAndroid Build Coastguard Worker
1117*b6fb3261SAndroid Build Coastguard Worker  if 'tensorrt_version' in config:
1118*b6fb3261SAndroid Build Coastguard Worker    print('Found TensorRT %s in:' % config['tensorrt_version'])
1119*b6fb3261SAndroid Build Coastguard Worker    print('    %s' % config['tensorrt_library_dir'])
1120*b6fb3261SAndroid Build Coastguard Worker    print('    %s' % config['tensorrt_include_dir'])
1121*b6fb3261SAndroid Build Coastguard Worker
1122*b6fb3261SAndroid Build Coastguard Worker  if config.get('nccl_version', None):
1123*b6fb3261SAndroid Build Coastguard Worker    print('Found NCCL %s in:' % config['nccl_version'])
1124*b6fb3261SAndroid Build Coastguard Worker    print('    %s' % config['nccl_library_dir'])
1125*b6fb3261SAndroid Build Coastguard Worker    print('    %s' % config['nccl_include_dir'])
1126*b6fb3261SAndroid Build Coastguard Worker
1127*b6fb3261SAndroid Build Coastguard Worker  print('\n')
1128*b6fb3261SAndroid Build Coastguard Worker
1129*b6fb3261SAndroid Build Coastguard Worker  environ_cp['CUDA_TOOLKIT_PATH'] = config['cuda_toolkit_path']
1130*b6fb3261SAndroid Build Coastguard Worker  return True
1131*b6fb3261SAndroid Build Coastguard Worker
1132*b6fb3261SAndroid Build Coastguard Worker
1133*b6fb3261SAndroid Build Coastguard Workerdef get_gcc_compiler(environ_cp):
1134*b6fb3261SAndroid Build Coastguard Worker  gcc_env = environ_cp.get('CXX') or environ_cp.get('CC') or which('gcc')
1135*b6fb3261SAndroid Build Coastguard Worker  if gcc_env is not None:
1136*b6fb3261SAndroid Build Coastguard Worker    gcc_version = run_shell([gcc_env, '--version']).split()
1137*b6fb3261SAndroid Build Coastguard Worker    if gcc_version[0] in ('gcc', 'g++'):
1138*b6fb3261SAndroid Build Coastguard Worker      return gcc_env
1139*b6fb3261SAndroid Build Coastguard Worker  return None
1140*b6fb3261SAndroid Build Coastguard Worker
1141*b6fb3261SAndroid Build Coastguard Worker
1142*b6fb3261SAndroid Build Coastguard Workerdef main():
1143*b6fb3261SAndroid Build Coastguard Worker  global _TF_WORKSPACE_ROOT
1144*b6fb3261SAndroid Build Coastguard Worker  global _TF_BAZELRC
1145*b6fb3261SAndroid Build Coastguard Worker  global _TF_CURRENT_BAZEL_VERSION
1146*b6fb3261SAndroid Build Coastguard Worker
1147*b6fb3261SAndroid Build Coastguard Worker  parser = argparse.ArgumentParser()
1148*b6fb3261SAndroid Build Coastguard Worker  parser.add_argument(
1149*b6fb3261SAndroid Build Coastguard Worker      '--workspace',
1150*b6fb3261SAndroid Build Coastguard Worker      type=str,
1151*b6fb3261SAndroid Build Coastguard Worker      default=os.path.abspath(os.path.dirname(__file__)),
1152*b6fb3261SAndroid Build Coastguard Worker      help='The absolute path to your active Bazel workspace.')
1153*b6fb3261SAndroid Build Coastguard Worker  args = parser.parse_args()
1154*b6fb3261SAndroid Build Coastguard Worker
1155*b6fb3261SAndroid Build Coastguard Worker  _TF_WORKSPACE_ROOT = args.workspace
1156*b6fb3261SAndroid Build Coastguard Worker  _TF_BAZELRC = os.path.join(_TF_WORKSPACE_ROOT, _TF_BAZELRC_FILENAME)
1157*b6fb3261SAndroid Build Coastguard Worker
1158*b6fb3261SAndroid Build Coastguard Worker  # Make a copy of os.environ to be clear when functions and getting and setting
1159*b6fb3261SAndroid Build Coastguard Worker  # environment variables.
1160*b6fb3261SAndroid Build Coastguard Worker  environ_cp = dict(os.environ)
1161*b6fb3261SAndroid Build Coastguard Worker
1162*b6fb3261SAndroid Build Coastguard Worker  try:
1163*b6fb3261SAndroid Build Coastguard Worker    current_bazel_version = retrieve_bazel_version()
1164*b6fb3261SAndroid Build Coastguard Worker  except subprocess.CalledProcessError as e:
1165*b6fb3261SAndroid Build Coastguard Worker    print('Error retrieving bazel version: ', e.output.decode('UTF-8').strip())
1166*b6fb3261SAndroid Build Coastguard Worker    raise e
1167*b6fb3261SAndroid Build Coastguard Worker
1168*b6fb3261SAndroid Build Coastguard Worker  _TF_CURRENT_BAZEL_VERSION = convert_version_to_int(current_bazel_version)
1169*b6fb3261SAndroid Build Coastguard Worker
1170*b6fb3261SAndroid Build Coastguard Worker  reset_tf_configure_bazelrc()
1171*b6fb3261SAndroid Build Coastguard Worker
1172*b6fb3261SAndroid Build Coastguard Worker  cleanup_makefile()
1173*b6fb3261SAndroid Build Coastguard Worker  setup_python(environ_cp)
1174*b6fb3261SAndroid Build Coastguard Worker
1175*b6fb3261SAndroid Build Coastguard Worker  if is_windows():
1176*b6fb3261SAndroid Build Coastguard Worker    environ_cp['TF_NEED_OPENCL'] = '0'
1177*b6fb3261SAndroid Build Coastguard Worker    environ_cp['TF_CUDA_CLANG'] = '0'
1178*b6fb3261SAndroid Build Coastguard Worker    # TODO(ibiryukov): Investigate using clang as a cpu or cuda compiler on
1179*b6fb3261SAndroid Build Coastguard Worker    # Windows.
1180*b6fb3261SAndroid Build Coastguard Worker    environ_cp['TF_DOWNLOAD_CLANG'] = '0'
1181*b6fb3261SAndroid Build Coastguard Worker    environ_cp['TF_NEED_MPI'] = '0'
1182*b6fb3261SAndroid Build Coastguard Worker
1183*b6fb3261SAndroid Build Coastguard Worker  if is_macos():
1184*b6fb3261SAndroid Build Coastguard Worker    environ_cp['TF_NEED_TENSORRT'] = '0'
1185*b6fb3261SAndroid Build Coastguard Worker
1186*b6fb3261SAndroid Build Coastguard Worker  if is_ppc64le():
1187*b6fb3261SAndroid Build Coastguard Worker    # Enable MMA Dynamic Dispatch support if 'gcc' and if linker >= 2.35
1188*b6fb3261SAndroid Build Coastguard Worker    gcc_env = get_gcc_compiler(environ_cp)
1189*b6fb3261SAndroid Build Coastguard Worker    if gcc_env is not None:
1190*b6fb3261SAndroid Build Coastguard Worker
1191*b6fb3261SAndroid Build Coastguard Worker      # Get the linker version
1192*b6fb3261SAndroid Build Coastguard Worker      ld_version = run_shell([gcc_env, '-Wl,-version']).split()
1193*b6fb3261SAndroid Build Coastguard Worker
1194*b6fb3261SAndroid Build Coastguard Worker      ld_version_int = convert_version_to_int(ld_version[3])
1195*b6fb3261SAndroid Build Coastguard Worker      if ld_version_int is None:
1196*b6fb3261SAndroid Build Coastguard Worker        ld_version_int = convert_version_to_int(ld_version[4])
1197*b6fb3261SAndroid Build Coastguard Worker
1198*b6fb3261SAndroid Build Coastguard Worker      # Enable if 'ld' version >= 2.35
1199*b6fb3261SAndroid Build Coastguard Worker      if ld_version_int >= 2035000:
1200*b6fb3261SAndroid Build Coastguard Worker        write_to_bazelrc(
1201*b6fb3261SAndroid Build Coastguard Worker            'build --copt="-DEIGEN_ALTIVEC_ENABLE_MMA_DYNAMIC_DISPATCH=1"')
1202*b6fb3261SAndroid Build Coastguard Worker
1203*b6fb3261SAndroid Build Coastguard Worker  with_xla_support = environ_cp.get('TF_ENABLE_XLA', None)
1204*b6fb3261SAndroid Build Coastguard Worker  if with_xla_support is not None:
1205*b6fb3261SAndroid Build Coastguard Worker    write_to_bazelrc('build --define=with_xla_support=%s' %
1206*b6fb3261SAndroid Build Coastguard Worker                     ('true' if int(with_xla_support) else 'false'))
1207*b6fb3261SAndroid Build Coastguard Worker
1208*b6fb3261SAndroid Build Coastguard Worker  set_action_env_var(
1209*b6fb3261SAndroid Build Coastguard Worker      environ_cp, 'TF_NEED_ROCM', 'ROCm', False, bazel_config_name='rocm')
1210*b6fb3261SAndroid Build Coastguard Worker  if (environ_cp.get('TF_NEED_ROCM') == '1' and
1211*b6fb3261SAndroid Build Coastguard Worker      'LD_LIBRARY_PATH' in environ_cp and
1212*b6fb3261SAndroid Build Coastguard Worker      environ_cp.get('LD_LIBRARY_PATH') != '1'):
1213*b6fb3261SAndroid Build Coastguard Worker    write_action_env_to_bazelrc('LD_LIBRARY_PATH',
1214*b6fb3261SAndroid Build Coastguard Worker                                environ_cp.get('LD_LIBRARY_PATH'))
1215*b6fb3261SAndroid Build Coastguard Worker
1216*b6fb3261SAndroid Build Coastguard Worker  if (environ_cp.get('TF_NEED_ROCM') == '1' and environ_cp.get('ROCM_PATH')):
1217*b6fb3261SAndroid Build Coastguard Worker    write_action_env_to_bazelrc('ROCM_PATH', environ_cp.get('ROCM_PATH'))
1218*b6fb3261SAndroid Build Coastguard Worker    write_action_env_to_bazelrc('ROCBLAS_TENSILE_LIBPATH',
1219*b6fb3261SAndroid Build Coastguard Worker                                environ_cp.get('ROCM_PATH') + '/lib/library')
1220*b6fb3261SAndroid Build Coastguard Worker
1221*b6fb3261SAndroid Build Coastguard Worker  if (environ_cp.get('TF_NEED_ROCM') == '1' and environ_cp.get('HIP_PLATFORM')):
1222*b6fb3261SAndroid Build Coastguard Worker    write_action_env_to_bazelrc('HIP_PLATFORM', environ_cp.get('HIP_PLATFORM'))
1223*b6fb3261SAndroid Build Coastguard Worker
1224*b6fb3261SAndroid Build Coastguard Worker  environ_cp['TF_NEED_CUDA'] = str(
1225*b6fb3261SAndroid Build Coastguard Worker      int(get_var(environ_cp, 'TF_NEED_CUDA', 'CUDA', False)))
1226*b6fb3261SAndroid Build Coastguard Worker  if (environ_cp.get('TF_NEED_CUDA') == '1' and
1227*b6fb3261SAndroid Build Coastguard Worker      'TF_CUDA_CONFIG_REPO' not in environ_cp):
1228*b6fb3261SAndroid Build Coastguard Worker
1229*b6fb3261SAndroid Build Coastguard Worker    set_action_env_var(
1230*b6fb3261SAndroid Build Coastguard Worker        environ_cp,
1231*b6fb3261SAndroid Build Coastguard Worker        'TF_NEED_TENSORRT',
1232*b6fb3261SAndroid Build Coastguard Worker        'TensorRT',
1233*b6fb3261SAndroid Build Coastguard Worker        False,
1234*b6fb3261SAndroid Build Coastguard Worker        bazel_config_name='tensorrt')
1235*b6fb3261SAndroid Build Coastguard Worker
1236*b6fb3261SAndroid Build Coastguard Worker    environ_save = dict(environ_cp)
1237*b6fb3261SAndroid Build Coastguard Worker    for _ in range(_DEFAULT_PROMPT_ASK_ATTEMPTS):
1238*b6fb3261SAndroid Build Coastguard Worker
1239*b6fb3261SAndroid Build Coastguard Worker      if validate_cuda_config(environ_cp):
1240*b6fb3261SAndroid Build Coastguard Worker        cuda_env_names = [
1241*b6fb3261SAndroid Build Coastguard Worker            'TF_CUDA_VERSION',
1242*b6fb3261SAndroid Build Coastguard Worker            'TF_CUBLAS_VERSION',
1243*b6fb3261SAndroid Build Coastguard Worker            'TF_CUDNN_VERSION',
1244*b6fb3261SAndroid Build Coastguard Worker            'TF_TENSORRT_VERSION',
1245*b6fb3261SAndroid Build Coastguard Worker            'TF_NCCL_VERSION',
1246*b6fb3261SAndroid Build Coastguard Worker            'TF_CUDA_PATHS',
1247*b6fb3261SAndroid Build Coastguard Worker            # Items below are for backwards compatibility when not using
1248*b6fb3261SAndroid Build Coastguard Worker            # TF_CUDA_PATHS.
1249*b6fb3261SAndroid Build Coastguard Worker            'CUDA_TOOLKIT_PATH',
1250*b6fb3261SAndroid Build Coastguard Worker            'CUDNN_INSTALL_PATH',
1251*b6fb3261SAndroid Build Coastguard Worker            'NCCL_INSTALL_PATH',
1252*b6fb3261SAndroid Build Coastguard Worker            'NCCL_HDR_PATH',
1253*b6fb3261SAndroid Build Coastguard Worker            'TENSORRT_INSTALL_PATH'
1254*b6fb3261SAndroid Build Coastguard Worker        ]
1255*b6fb3261SAndroid Build Coastguard Worker        # Note: set_action_env_var above already writes to bazelrc.
1256*b6fb3261SAndroid Build Coastguard Worker        for name in cuda_env_names:
1257*b6fb3261SAndroid Build Coastguard Worker          if name in environ_cp:
1258*b6fb3261SAndroid Build Coastguard Worker            write_action_env_to_bazelrc(name, environ_cp[name])
1259*b6fb3261SAndroid Build Coastguard Worker        break
1260*b6fb3261SAndroid Build Coastguard Worker
1261*b6fb3261SAndroid Build Coastguard Worker      # Restore settings changed below if CUDA config could not be validated.
1262*b6fb3261SAndroid Build Coastguard Worker      environ_cp = dict(environ_save)
1263*b6fb3261SAndroid Build Coastguard Worker
1264*b6fb3261SAndroid Build Coastguard Worker      set_tf_cuda_version(environ_cp)
1265*b6fb3261SAndroid Build Coastguard Worker      set_tf_cudnn_version(environ_cp)
1266*b6fb3261SAndroid Build Coastguard Worker      if is_windows():
1267*b6fb3261SAndroid Build Coastguard Worker        set_tf_tensorrt_version(environ_cp)
1268*b6fb3261SAndroid Build Coastguard Worker      if is_linux():
1269*b6fb3261SAndroid Build Coastguard Worker        set_tf_tensorrt_version(environ_cp)
1270*b6fb3261SAndroid Build Coastguard Worker        set_tf_nccl_version(environ_cp)
1271*b6fb3261SAndroid Build Coastguard Worker
1272*b6fb3261SAndroid Build Coastguard Worker      set_tf_cuda_paths(environ_cp)
1273*b6fb3261SAndroid Build Coastguard Worker
1274*b6fb3261SAndroid Build Coastguard Worker    else:
1275*b6fb3261SAndroid Build Coastguard Worker      raise UserInputError(
1276*b6fb3261SAndroid Build Coastguard Worker          'Invalid CUDA setting were provided %d '
1277*b6fb3261SAndroid Build Coastguard Worker          'times in a row. Assuming to be a scripting mistake.' %
1278*b6fb3261SAndroid Build Coastguard Worker          _DEFAULT_PROMPT_ASK_ATTEMPTS)
1279*b6fb3261SAndroid Build Coastguard Worker
1280*b6fb3261SAndroid Build Coastguard Worker    set_tf_cuda_compute_capabilities(environ_cp)
1281*b6fb3261SAndroid Build Coastguard Worker    if 'LD_LIBRARY_PATH' in environ_cp and environ_cp.get(
1282*b6fb3261SAndroid Build Coastguard Worker        'LD_LIBRARY_PATH') != '1':
1283*b6fb3261SAndroid Build Coastguard Worker      write_action_env_to_bazelrc('LD_LIBRARY_PATH',
1284*b6fb3261SAndroid Build Coastguard Worker                                  environ_cp.get('LD_LIBRARY_PATH'))
1285*b6fb3261SAndroid Build Coastguard Worker
1286*b6fb3261SAndroid Build Coastguard Worker    set_tf_cuda_clang(environ_cp)
1287*b6fb3261SAndroid Build Coastguard Worker    if environ_cp.get('TF_CUDA_CLANG') == '1':
1288*b6fb3261SAndroid Build Coastguard Worker      # Ask whether we should download the clang toolchain.
1289*b6fb3261SAndroid Build Coastguard Worker      set_tf_download_clang(environ_cp)
1290*b6fb3261SAndroid Build Coastguard Worker      if environ_cp.get('TF_DOWNLOAD_CLANG') != '1':
1291*b6fb3261SAndroid Build Coastguard Worker        # Set up which clang we should use as the cuda / host compiler.
1292*b6fb3261SAndroid Build Coastguard Worker        set_clang_cuda_compiler_path(environ_cp)
1293*b6fb3261SAndroid Build Coastguard Worker      else:
1294*b6fb3261SAndroid Build Coastguard Worker        # Use downloaded LLD for linking.
1295*b6fb3261SAndroid Build Coastguard Worker        write_to_bazelrc('build:cuda_clang --config=download_clang_use_lld')
1296*b6fb3261SAndroid Build Coastguard Worker    else:
1297*b6fb3261SAndroid Build Coastguard Worker      # Set up which gcc nvcc should use as the host compiler
1298*b6fb3261SAndroid Build Coastguard Worker      # No need to set this on Windows
1299*b6fb3261SAndroid Build Coastguard Worker      if not is_windows():
1300*b6fb3261SAndroid Build Coastguard Worker        set_gcc_host_compiler_path(environ_cp)
1301*b6fb3261SAndroid Build Coastguard Worker    set_other_cuda_vars(environ_cp)
1302*b6fb3261SAndroid Build Coastguard Worker  else:
1303*b6fb3261SAndroid Build Coastguard Worker    # CUDA not required. Ask whether we should download the clang toolchain and
1304*b6fb3261SAndroid Build Coastguard Worker    # use it for the CPU build.
1305*b6fb3261SAndroid Build Coastguard Worker    set_tf_download_clang(environ_cp)
1306*b6fb3261SAndroid Build Coastguard Worker
1307*b6fb3261SAndroid Build Coastguard Worker  # ROCm / CUDA are mutually exclusive.
1308*b6fb3261SAndroid Build Coastguard Worker  # At most 1 GPU platform can be configured.
1309*b6fb3261SAndroid Build Coastguard Worker  gpu_platform_count = 0
1310*b6fb3261SAndroid Build Coastguard Worker  if environ_cp.get('TF_NEED_ROCM') == '1':
1311*b6fb3261SAndroid Build Coastguard Worker    gpu_platform_count += 1
1312*b6fb3261SAndroid Build Coastguard Worker  if environ_cp.get('TF_NEED_CUDA') == '1':
1313*b6fb3261SAndroid Build Coastguard Worker    gpu_platform_count += 1
1314*b6fb3261SAndroid Build Coastguard Worker  if gpu_platform_count >= 2:
1315*b6fb3261SAndroid Build Coastguard Worker    raise UserInputError('CUDA / ROCm are mututally exclusive. '
1316*b6fb3261SAndroid Build Coastguard Worker                         'At most 1 GPU platform can be configured.')
1317*b6fb3261SAndroid Build Coastguard Worker
1318*b6fb3261SAndroid Build Coastguard Worker  set_cc_opt_flags(environ_cp)
1319*b6fb3261SAndroid Build Coastguard Worker  set_system_libs_flag(environ_cp)
1320*b6fb3261SAndroid Build Coastguard Worker  if is_windows():
1321*b6fb3261SAndroid Build Coastguard Worker    set_windows_build_flags(environ_cp)
1322*b6fb3261SAndroid Build Coastguard Worker
1323*b6fb3261SAndroid Build Coastguard Worker  if get_var(environ_cp, 'TF_SET_ANDROID_WORKSPACE', 'android workspace', False,
1324*b6fb3261SAndroid Build Coastguard Worker             ('Would you like to interactively configure ./WORKSPACE for '
1325*b6fb3261SAndroid Build Coastguard Worker              'Android builds?'), 'Searching for NDK and SDK installations.',
1326*b6fb3261SAndroid Build Coastguard Worker             'Not configuring the WORKSPACE for Android builds.'):
1327*b6fb3261SAndroid Build Coastguard Worker    create_android_ndk_rule(environ_cp)
1328*b6fb3261SAndroid Build Coastguard Worker    create_android_sdk_rule(environ_cp)
1329*b6fb3261SAndroid Build Coastguard Worker
1330*b6fb3261SAndroid Build Coastguard Worker  system_specific_test_config(environ_cp)
1331*b6fb3261SAndroid Build Coastguard Worker
1332*b6fb3261SAndroid Build Coastguard Worker  configure_ios(environ_cp)
1333*b6fb3261SAndroid Build Coastguard Worker
1334*b6fb3261SAndroid Build Coastguard Worker  print('Preconfigured Bazel build configs. You can use any of the below by '
1335*b6fb3261SAndroid Build Coastguard Worker        'adding "--config=<>" to your build command. See .bazelrc for more '
1336*b6fb3261SAndroid Build Coastguard Worker        'details.')
1337*b6fb3261SAndroid Build Coastguard Worker  config_info_line('mkl', 'Build with MKL support.')
1338*b6fb3261SAndroid Build Coastguard Worker  config_info_line(
1339*b6fb3261SAndroid Build Coastguard Worker      'mkl_aarch64',
1340*b6fb3261SAndroid Build Coastguard Worker      'Build with oneDNN and Compute Library for the Arm Architecture (ACL).')
1341*b6fb3261SAndroid Build Coastguard Worker  config_info_line('monolithic', 'Config for mostly static monolithic build.')
1342*b6fb3261SAndroid Build Coastguard Worker  config_info_line('numa', 'Build with NUMA support.')
1343*b6fb3261SAndroid Build Coastguard Worker  config_info_line(
1344*b6fb3261SAndroid Build Coastguard Worker      'dynamic_kernels',
1345*b6fb3261SAndroid Build Coastguard Worker      '(Experimental) Build kernels into separate shared objects.')
1346*b6fb3261SAndroid Build Coastguard Worker  config_info_line('v1', 'Build with TensorFlow 1 API instead of TF 2 API.')
1347*b6fb3261SAndroid Build Coastguard Worker
1348*b6fb3261SAndroid Build Coastguard Worker  print('Preconfigured Bazel build configs to DISABLE default on features:')
1349*b6fb3261SAndroid Build Coastguard Worker  config_info_line('nogcp', 'Disable GCP support.')
1350*b6fb3261SAndroid Build Coastguard Worker  config_info_line('nonccl', 'Disable NVIDIA NCCL support.')
1351*b6fb3261SAndroid Build Coastguard Worker
1352*b6fb3261SAndroid Build Coastguard Worker
1353*b6fb3261SAndroid Build Coastguard Workerif __name__ == '__main__':
1354*b6fb3261SAndroid Build Coastguard Worker  main()
1355