• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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
37from __future__ import print_function
38
39__author__ = 'keveman@google.com (Manjunath Kudlur)'
40
41from argparse import ArgumentParser
42import os
43import subprocess
44import re
45import sys
46import pipes
47
48# Template values set by cuda_autoconf.
49CPU_COMPILER = ('%{cpu_compiler}')
50GCC_HOST_COMPILER_PATH = ('%{gcc_host_compiler_path}')
51
52NVCC_PATH = '%{nvcc_path}'
53PREFIX_DIR = os.path.dirname(GCC_HOST_COMPILER_PATH)
54NVCC_VERSION = '%{cuda_version}'
55
56def Log(s):
57  print('gpus/crosstool: {0}'.format(s))
58
59
60def GetOptionValue(argv, option):
61  """Extract the list of values for option from the argv list.
62
63  Args:
64    argv: A list of strings, possibly the argv passed to main().
65    option: The option whose value to extract, with the leading '-'.
66
67  Returns:
68    A list of values, either directly following the option,
69    (eg., -opt val1 val2) or values collected from multiple occurrences of
70    the option (eg., -opt val1 -opt val2).
71  """
72
73  parser = ArgumentParser()
74  parser.add_argument(option, nargs='*', action='append')
75  option = option.lstrip('-').replace('-', '_')
76  args, _ = parser.parse_known_args(argv)
77  if not args or not vars(args)[option]:
78    return []
79  else:
80    return sum(vars(args)[option], [])
81
82
83def GetHostCompilerOptions(argv):
84  """Collect the -isystem, -iquote, and --sysroot option values from argv.
85
86  Args:
87    argv: A list of strings, possibly the argv passed to main().
88
89  Returns:
90    The string that can be used as the --compiler-options to nvcc.
91  """
92
93  parser = ArgumentParser()
94  parser.add_argument('-isystem', nargs='*', action='append')
95  parser.add_argument('-iquote', nargs='*', action='append')
96  parser.add_argument('--sysroot', nargs=1)
97  parser.add_argument('-g', nargs='*', action='append')
98  parser.add_argument('-fno-canonical-system-headers', action='store_true')
99  parser.add_argument('-no-canonical-prefixes', action='store_true')
100
101  args, _ = parser.parse_known_args(argv)
102
103  opts = ''
104
105  if args.isystem:
106    opts += ' -isystem ' + ' -isystem '.join(sum(args.isystem, []))
107  if args.iquote:
108    opts += ' -iquote ' + ' -iquote '.join(sum(args.iquote, []))
109  if args.g:
110    opts += ' -g' + ' -g'.join(sum(args.g, []))
111  if args.fno_canonical_system_headers:
112    opts += ' -fno-canonical-system-headers'
113  if args.no_canonical_prefixes:
114    opts += ' -no-canonical-prefixes'
115  if args.sysroot:
116    opts += ' --sysroot ' + args.sysroot[0]
117
118  return opts
119
120def _update_options(nvcc_options):
121  if NVCC_VERSION in ("7.0",):
122    return nvcc_options
123
124  update_options = { "relaxed-constexpr" : "expt-relaxed-constexpr" }
125  return [ update_options[opt] if opt in update_options else opt
126                    for opt in nvcc_options ]
127
128def GetNvccOptions(argv):
129  """Collect the -nvcc_options values from argv.
130
131  Args:
132    argv: A list of strings, possibly the argv passed to main().
133
134  Returns:
135    The string that can be passed directly to nvcc.
136  """
137
138  parser = ArgumentParser()
139  parser.add_argument('-nvcc_options', nargs='*', action='append')
140
141  args, _ = parser.parse_known_args(argv)
142
143  if args.nvcc_options:
144    options = _update_options(sum(args.nvcc_options, []))
145    return ' '.join(['--'+a for a in options])
146  return ''
147
148def system(cmd):
149  """Invokes cmd with os.system().
150
151  Args:
152    cmd: The command.
153
154  Returns:
155    The exit code if the process exited with exit() or -signal
156    if the process was terminated by a signal.
157  """
158  retv = os.system(cmd)
159  if os.WIFEXITED(retv):
160    return os.WEXITSTATUS(retv)
161  else:
162    return -os.WTERMSIG(retv)
163
164def InvokeNvcc(argv, log=False):
165  """Call nvcc with arguments assembled from argv.
166
167  Args:
168    argv: A list of strings, possibly the argv passed to main().
169    log: True if logging is requested.
170
171  Returns:
172    The return value of calling system('nvcc ' + args)
173  """
174
175  host_compiler_options = GetHostCompilerOptions(argv)
176  nvcc_compiler_options = GetNvccOptions(argv)
177  opt_option = GetOptionValue(argv, '-O')
178  m_options = GetOptionValue(argv, '-m')
179  m_options = ''.join([' -m' + m for m in m_options if m in ['32', '64']])
180  include_options = GetOptionValue(argv, '-I')
181  out_file = GetOptionValue(argv, '-o')
182  depfiles = GetOptionValue(argv, '-MF')
183  defines = GetOptionValue(argv, '-D')
184  defines = ''.join([' -D' + define for define in defines])
185  undefines = GetOptionValue(argv, '-U')
186  undefines = ''.join([' -U' + define for define in undefines])
187  std_options = GetOptionValue(argv, '-std')
188  # Supported -std flags as of CUDA 9.0. Only keep last to mimic gcc/clang.
189  nvcc_allowed_std_options = ["c++03", "c++11", "c++14"]
190  std_options = ''.join([' -std=' + define
191      for define in std_options if define in nvcc_allowed_std_options][-1:])
192  fatbin_options = ''.join([' --fatbin-options=' + option
193      for option in GetOptionValue(argv, '-Xcuda-fatbinary')])
194
195  # The list of source files get passed after the -c option. I don't know of
196  # any other reliable way to just get the list of source files to be compiled.
197  src_files = GetOptionValue(argv, '-c')
198
199  # Pass -w through from host to nvcc, but don't do anything fancier with
200  # warnings-related flags, since they're not necessarily the same across
201  # compilers.
202  warning_options = ' -w' if '-w' in argv else ''
203
204  if len(src_files) == 0:
205    return 1
206  if len(out_file) != 1:
207    return 1
208
209  opt = (' -O2' if (len(opt_option) > 0 and int(opt_option[0]) > 0)
210         else ' -g')
211
212  includes = (' -I ' + ' -I '.join(include_options)
213              if len(include_options) > 0
214              else '')
215
216  # Unfortunately, there are other options that have -c prefix too.
217  # So allowing only those look like C/C++ files.
218  src_files = [f for f in src_files if
219               re.search('\.cpp$|\.cc$|\.c$|\.cxx$|\.C$', f)]
220  srcs = ' '.join(src_files)
221  out = ' -o ' + out_file[0]
222
223  nvccopts = '-D_FORCE_INLINES '
224  for capability in GetOptionValue(argv, "--cuda-gpu-arch"):
225    capability = capability[len('sm_'):]
226    nvccopts += r'-gencode=arch=compute_%s,\"code=sm_%s\" ' % (capability,
227                                                               capability)
228  for capability in GetOptionValue(argv, '--cuda-include-ptx'):
229    capability = capability[len('sm_'):]
230    nvccopts += r'-gencode=arch=compute_%s,\"code=compute_%s\" ' % (capability,
231                                                                    capability)
232  nvccopts += nvcc_compiler_options
233  nvccopts += undefines
234  nvccopts += defines
235  nvccopts += std_options
236  nvccopts += m_options
237  nvccopts += warning_options
238  nvccopts += fatbin_options
239
240  if depfiles:
241    # Generate the dependency file
242    depfile = depfiles[0]
243    cmd = (NVCC_PATH + ' ' + nvccopts +
244           ' --compiler-options "' + host_compiler_options + '"' +
245           ' --compiler-bindir=' + GCC_HOST_COMPILER_PATH +
246           ' -I .' +
247           ' -x cu ' + opt + includes + ' ' + srcs + ' -M -o ' + depfile)
248    if log: Log(cmd)
249    exit_status = system(cmd)
250    if exit_status != 0:
251      return exit_status
252
253  cmd = (NVCC_PATH + ' ' + nvccopts +
254         ' --compiler-options "' + host_compiler_options + ' -fPIC"' +
255         ' --compiler-bindir=' + GCC_HOST_COMPILER_PATH +
256         ' -I .' +
257         ' -x cu ' + opt + includes + ' -c ' + srcs + out)
258
259  # TODO(zhengxq): for some reason, 'gcc' needs this help to find 'as'.
260  # Need to investigate and fix.
261  cmd = 'PATH=' + PREFIX_DIR + ':$PATH ' + cmd
262  if log: Log(cmd)
263  return system(cmd)
264
265
266def main():
267  parser = ArgumentParser()
268  parser.add_argument('-x', nargs=1)
269  parser.add_argument('--cuda_log', action='store_true')
270  args, leftover = parser.parse_known_args(sys.argv[1:])
271
272  if args.x and args.x[0] == 'cuda':
273    if args.cuda_log: Log('-x cuda')
274    leftover = [pipes.quote(s) for s in leftover]
275    if args.cuda_log: Log('using nvcc')
276    return InvokeNvcc(leftover, log=args.cuda_log)
277
278  # Strip our flags before passing through to the CPU compiler for files which
279  # are not -x cuda. We can't just pass 'leftover' because it also strips -x.
280  # We not only want to pass -x to the CPU compiler, but also keep it in its
281  # relative location in the argv list (the compiler is actually sensitive to
282  # this).
283  cpu_compiler_flags = [flag for flag in sys.argv[1:]
284                             if not flag.startswith(('--cuda_log'))]
285
286  return subprocess.call([CPU_COMPILER] + cpu_compiler_flags)
287
288if __name__ == '__main__':
289  sys.exit(main())
290