1#!/usr/bin/env python
2# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15# ==============================================================================
16
17"""Crosstool wrapper for compiling CUDA programs.
18
19SYNOPSIS:
20  crosstool_wrapper_is_not_gcc [options passed in by cc_library()
21                                or cc_binary() rule]
22
23DESCRIPTION:
24  This script is expected to be called by the cc_library() or cc_binary() bazel
25  rules. When the option "-x cuda" is present in the list of arguments passed
26  to this script, it invokes the nvcc CUDA compiler. Most arguments are passed
27  as is as a string to --compiler-options of nvcc. When "-x cuda" is not
28  present, this wrapper invokes hybrid_driver_is_not_gcc with the input
29  arguments as is.
30
31NOTES:
32  Changes to the contents of this file must be propagated from
33  //third_party/gpus/crosstool/crosstool_wrapper_is_not_gcc to
34  //third_party/gpus/crosstool/v*/*/clang/bin/crosstool_wrapper_is_not_gcc
35"""
36
37__author__ = '[email protected] (Manjunath Kudlur)'
38
39from argparse import ArgumentParser
40import os
41import subprocess
42import re
43import sys
44import pipes
45
46# Template values set by cuda_autoconf.
47CPU_COMPILER = ('%{cpu_compiler}')
48GCC_HOST_COMPILER_PATH = ('%{gcc_host_compiler_path}')
49
50NVCC_PATH = '%{nvcc_path}'
51PREFIX_DIR = os.path.dirname(GCC_HOST_COMPILER_PATH)
52NVCC_VERSION = '%{cuda_version}'
53
54def Log(s):
55  print('gpus/crosstool: {0}'.format(s))
56
57
58def GetOptionValue(argv, option):
59  """Extract the list of values for option from the argv list.
60
61  Args:
62    argv: A list of strings, possibly the argv passed to main().
63    option: The option whose value to extract, with the leading '-'.
64
65  Returns:
66    A list of values, either directly following the option,
67    (eg., -opt val1 val2) or values collected from multiple occurrences of
68    the option (eg., -opt val1 -opt val2).
69  """
70
71  parser = ArgumentParser()
72  parser.add_argument(option, nargs='*', action='append')
73  option = option.lstrip('-').replace('-', '_')
74  args, _ = parser.parse_known_args(argv)
75  if not args or not vars(args)[option]:
76    return []
77  else:
78    return sum(vars(args)[option], [])
79
80
81def GetHostCompilerOptions(argv):
82  """Collect the -isystem, -iquote, and --sysroot option values from argv.
83
84  Args:
85    argv: A list of strings, possibly the argv passed to main().
86
87  Returns:
88    The string that can be used as the --compiler-options to nvcc.
89  """
90
91  parser = ArgumentParser()
92  parser.add_argument('-isystem', nargs='*', action='append')
93  parser.add_argument('-iquote', nargs='*', action='append')
94  parser.add_argument('--sysroot', nargs=1)
95  parser.add_argument('-g', nargs='*', action='append')
96  parser.add_argument('-fno-canonical-system-headers', action='store_true')
97  parser.add_argument('-no-canonical-prefixes', action='store_true')
98
99  args, _ = parser.parse_known_args(argv)
100
101  opts = ''
102
103  if args.isystem:
104    opts += ' -isystem ' + ' -isystem '.join(sum(args.isystem, []))
105  if args.iquote:
106    opts += ' -iquote ' + ' -iquote '.join(sum(args.iquote, []))
107  if args.g:
108    opts += ' -g' + ' -g'.join(sum(args.g, []))
109  if args.fno_canonical_system_headers:
110    opts += ' -fno-canonical-system-headers'
111  if args.no_canonical_prefixes:
112    opts += ' -no-canonical-prefixes'
113  if args.sysroot:
114    opts += ' --sysroot ' + args.sysroot[0]
115
116  return opts
117
118def _update_options(nvcc_options):
119  if NVCC_VERSION in ("7.0",):
120    return nvcc_options
121
122  update_options = { "relaxed-constexpr" : "expt-relaxed-constexpr" }
123  return [ update_options[opt] if opt in update_options else opt
124                    for opt in nvcc_options ]
125
126def GetNvccOptions(argv):
127  """Collect the -nvcc_options values from argv.
128
129  Args:
130    argv: A list of strings, possibly the argv passed to main().
131
132  Returns:
133    The string that can be passed directly to nvcc.
134  """
135
136  parser = ArgumentParser()
137  parser.add_argument('-nvcc_options', nargs='*', action='append')
138
139  args, _ = parser.parse_known_args(argv)
140
141  if args.nvcc_options:
142    options = _update_options(sum(args.nvcc_options, []))
143    return ' '.join(['--'+a for a in options])
144  return ''
145
146def system(cmd):
147  """Invokes cmd with os.system().
148
149  Args:
150    cmd: The command.
151
152  Returns:
153    The exit code if the process exited with exit() or -signal
154    if the process was terminated by a signal.
155  """
156  retv = os.system(cmd)
157  if os.WIFEXITED(retv):
158    return os.WEXITSTATUS(retv)
159  else:
160    return -os.WTERMSIG(retv)
161
162def InvokeNvcc(argv, log=False):
163  """Call nvcc with arguments assembled from argv.
164
165  Args:
166    argv: A list of strings, possibly the argv passed to main().
167    log: True if logging is requested.
168
169  Returns:
170    The return value of calling system('nvcc ' + args)
171  """
172
173  host_compiler_options = GetHostCompilerOptions(argv)
174  nvcc_compiler_options = GetNvccOptions(argv)
175  opt_option = GetOptionValue(argv, '-O')
176  m_options = GetOptionValue(argv, '-m')
177  m_options = ''.join([' -m' + m for m in m_options if m in ['32', '64']])
178  include_options = GetOptionValue(argv, '-I')
179  out_file = GetOptionValue(argv, '-o')
180  depfiles = GetOptionValue(argv, '-MF')
181  defines = GetOptionValue(argv, '-D')
182  defines = ''.join([' -D' + define for define in defines])
183  undefines = GetOptionValue(argv, '-U')
184  undefines = ''.join([' -U' + define for define in undefines])
185  std_options = GetOptionValue(argv, '-std')
186  # Supported -std flags as of CUDA 9.0. Only keep last to mimic gcc/clang.
187  nvcc_allowed_std_options = ["c++03", "c++11", "c++14"]
188  nvcc_std_map = {}
189  if int(NVCC_VERSION.split('.')[0]) >= 11:
190      nvcc_std_map["c++1z"] = "c++17"
191      nvcc_allowed_std_options += ["c++17", "c++1z"]
192  std_options = ''.join([' -std=' +
193      (nvcc_std_map[define] if define in nvcc_std_map else define)
194      for define in std_options if define in nvcc_allowed_std_options][-1:])
195  fatbin_options = ''.join([' --fatbin-options=' + option
196      for option in GetOptionValue(argv, '-Xcuda-fatbinary')])
197
198  # The list of source files get passed after the -c option. I don't know of
199  # any other reliable way to just get the list of source files to be compiled.
200  src_files = GetOptionValue(argv, '-c')
201
202  # Pass -w through from host to nvcc, but don't do anything fancier with
203  # warnings-related flags, since they're not necessarily the same across
204  # compilers.
205  warning_options = ' -w' if '-w' in argv else ''
206
207  if len(src_files) == 0:
208    return 1
209  if len(out_file) != 1:
210    return 1
211
212  opt = (' -O2' if (len(opt_option) > 0 and int(opt_option[0]) > 0)
213         else ' -g')
214
215  includes = (' -I ' + ' -I '.join(include_options)
216              if len(include_options) > 0
217              else '')
218
219  # Unfortunately, there are other options that have -c prefix too.
220  # So allowing only those look like C/C++ files.
221  src_files = [f for f in src_files if
222               re.search('\.cpp$|\.cc$|\.c$|\.cxx$|\.C$', f)]
223  srcs = ' '.join(src_files)
224  out = ' -o ' + out_file[0]
225
226  nvccopts = '-D_FORCE_INLINES '
227  capabilities_sm = set(GetOptionValue(argv, "--cuda-gpu-arch"))
228  capabilities_compute = set(GetOptionValue(argv, '--cuda-include-ptx'))
229  # When both "code=sm_xy" and "code=compute_xy" are requested for a single
230  # arch, they can be combined using "code=xy,compute_xy" which avoids a
231  # redundant PTX generation during compilation.
232  capabilities_both = capabilities_sm.intersection(capabilities_compute)
233  for capability in capabilities_both:
234    capability = capability[len('sm_'):]
235    nvccopts += r'-gencode=arch=compute_%s,code=\"sm_%s,compute_%s\" ' % (
236        capability, capability, capability)
237  for capability in capabilities_sm - capabilities_both:
238    capability = capability[len('sm_'):]
239    nvccopts += r'-gencode=arch=compute_%s,\"code=sm_%s\" ' % (capability,
240                                                               capability)
241  for capability in capabilities_compute - capabilities_both:
242    capability = capability[len('sm_'):]
243    nvccopts += r'-gencode=arch=compute_%s,\"code=compute_%s\" ' % (capability,
244                                                                    capability)
245  nvccopts += nvcc_compiler_options
246  nvccopts += undefines
247  nvccopts += defines
248  nvccopts += std_options
249  nvccopts += m_options
250  nvccopts += warning_options
251  # Force C++17 dialect (note, everything in just one string!)
252  nvccopts += ' --std c++17 '
253  nvccopts += fatbin_options
254
255  if depfiles:
256    # Generate the dependency file
257    depfile = depfiles[0]
258    cmd = (NVCC_PATH + ' ' + nvccopts +
259           ' --compiler-options "' + host_compiler_options + '"' +
260           ' --compiler-bindir=' + GCC_HOST_COMPILER_PATH +
261           ' -I .' +
262           ' -x cu ' + opt + includes + ' ' + srcs + ' -M -o ' + depfile)
263    if log: Log(cmd)
264    exit_status = system(cmd)
265    if exit_status != 0:
266      return exit_status
267
268  cmd = (NVCC_PATH + ' ' + nvccopts +
269         ' --compiler-options "' + host_compiler_options + ' -fPIC"' +
270         ' --compiler-bindir=' + GCC_HOST_COMPILER_PATH +
271         ' -I .' +
272         ' -x cu ' + opt + includes + ' -c ' + srcs + out)
273
274  # TODO(zhengxq): for some reason, 'gcc' needs this help to find 'as'.
275  # Need to investigate and fix.
276  cmd = 'PATH=' + PREFIX_DIR + ':$PATH ' + cmd
277  if log: Log(cmd)
278  return system(cmd)
279
280
281def main():
282  parser = ArgumentParser()
283  parser.add_argument('-x', nargs=1)
284  parser.add_argument('--cuda_log', action='store_true')
285  args, leftover = parser.parse_known_args(sys.argv[1:])
286
287  if args.x and args.x[0] == 'cuda':
288    if args.cuda_log: Log('-x cuda')
289    leftover = [pipes.quote(s) for s in leftover]
290    if args.cuda_log: Log('using nvcc')
291    return InvokeNvcc(leftover, log=args.cuda_log)
292
293  # Strip our flags before passing through to the CPU compiler for files which
294  # are not -x cuda. We can't just pass 'leftover' because it also strips -x.
295  # We not only want to pass -x to the CPU compiler, but also keep it in its
296  # relative location in the argv list (the compiler is actually sensitive to
297  # this).
298  cpu_compiler_flags = [flag for flag in sys.argv[1:]
299                             if not flag.startswith(('--cuda_log'))]
300
301  return subprocess.call([CPU_COMPILER] + cpu_compiler_flags)
302
303if __name__ == '__main__':
304  sys.exit(main())
305