1#!/usr/bin/env python
2"""Crosstool wrapper for compiling ROCm programs.
3
4SYNOPSIS:
5  crosstool_wrapper_driver_rocm [options passed in by cc_library()
6                                or cc_binary() rule]
7
8DESCRIPTION:
9  This script is expected to be called by the cc_library() or cc_binary() bazel
10  rules. When the option "-x rocm" is present in the list of arguments passed
11  to this script, it invokes the hipcc compiler. Most arguments are passed
12  as is as a string to --compiler-options of hipcc. When "-x rocm" is not
13  present, this wrapper invokes gcc with the input arguments as is.
14"""
15
16__author__ = '[email protected] (Wen-Heng (Jack) Chung)'
17
18from argparse import ArgumentParser
19import os
20import subprocess
21import re
22import sys
23import pipes
24
25# Template values set by rocm_configure.bzl.
26CPU_COMPILER = ('%{cpu_compiler}')
27
28HIPCC_PATH = '%{hipcc_path}'
29HIPCC_ENV = '%{hipcc_env}'
30HIP_RUNTIME_PATH = '%{hip_runtime_path}'
31HIP_RUNTIME_LIBRARY = '%{hip_runtime_library}'
32ROCR_RUNTIME_PATH = '%{rocr_runtime_path}'
33ROCR_RUNTIME_LIBRARY = '%{rocr_runtime_library}'
34VERBOSE = '%{crosstool_verbose}'=='1'
35
36def Log(s):
37  print('gpus/crosstool: {0}'.format(s))
38
39
40def GetOptionValue(argv, option):
41  """Extract the list of values for option from the argv list.
42
43  Args:
44    argv: A list of strings, possibly the argv passed to main().
45    option: The option whose value to extract, without the leading '-'.
46
47  Returns:
48    A list of values, either directly following the option,
49    (eg., -opt val1 val2) or values collected from multiple occurrences of
50    the option (eg., -opt val1 -opt val2).
51  """
52
53  parser = ArgumentParser()
54  parser.add_argument('-' + option, nargs='*', action='append')
55  args, _ = parser.parse_known_args(argv)
56  if not args or not vars(args)[option]:
57    return []
58  else:
59    return sum(vars(args)[option], [])
60
61
62def GetHostCompilerOptions(argv):
63  """Collect the -isystem, -iquote, and --sysroot option values from argv.
64
65  Args:
66    argv: A list of strings, possibly the argv passed to main().
67
68  Returns:
69    The string that can be used as the --compiler-options to hipcc.
70  """
71
72  parser = ArgumentParser()
73  parser.add_argument('-isystem', nargs='*', action='append')
74  parser.add_argument('-iquote', nargs='*', action='append')
75  parser.add_argument('--sysroot', nargs=1)
76  parser.add_argument('-g', nargs='*', action='append')
77  parser.add_argument('-fno-canonical-system-headers', action='store_true')
78
79  args, _ = parser.parse_known_args(argv)
80
81  opts = ''
82
83  if args.isystem:
84    opts += ' -isystem ' + ' -isystem '.join(sum(args.isystem, []))
85  if args.iquote:
86    opts += ' -iquote ' + ' -iquote '.join(sum(args.iquote, []))
87  if args.g:
88    opts += ' -g' + ' -g'.join(sum(args.g, []))
89  #if args.fno_canonical_system_headers:
90  #  opts += ' -fno-canonical-system-headers'
91  if args.sysroot:
92    opts += ' --sysroot ' + args.sysroot[0]
93
94  return opts
95
96def system(cmd):
97  """Invokes cmd with os.system().
98
99  Args:
100    cmd: The command.
101
102  Returns:
103    The exit code if the process exited with exit() or -signal
104    if the process was terminated by a signal.
105  """
106  retv = os.system(cmd)
107  if os.WIFEXITED(retv):
108    return os.WEXITSTATUS(retv)
109  else:
110    return -os.WTERMSIG(retv)
111
112
113def InvokeHipcc(argv, log=False):
114  """Call hipcc with arguments assembled from argv.
115
116  Args:
117    argv: A list of strings, possibly the argv passed to main().
118    log: True if logging is requested.
119
120  Returns:
121    The return value of calling os.system('hipcc ' + args)
122  """
123
124  host_compiler_options = GetHostCompilerOptions(argv)
125  opt_option = GetOptionValue(argv, 'O')
126  m_options = GetOptionValue(argv, 'm')
127  m_options = ''.join([' -m' + m for m in m_options if m in ['32', '64']])
128  include_options = GetOptionValue(argv, 'I')
129  out_file = GetOptionValue(argv, 'o')
130  depfiles = GetOptionValue(argv, 'MF')
131  defines = GetOptionValue(argv, 'D')
132  defines = ''.join([' -D' + define for define in defines])
133  undefines = GetOptionValue(argv, 'U')
134  undefines = ''.join([' -U' + define for define in undefines])
135  std_options = GetOptionValue(argv, 'std')
136  hipcc_allowed_std_options = ["c++11", "c++14", "c++17"]
137  std_options = ''.join([' -std=' + define
138      for define in std_options if define in hipcc_allowed_std_options])
139
140  # The list of source files get passed after the -c option. I don't know of
141  # any other reliable way to just get the list of source files to be compiled.
142  src_files = GetOptionValue(argv, 'c')
143
144  if len(src_files) == 0:
145    return 1
146  if len(out_file) != 1:
147    return 1
148
149  opt = (' -O2' if (len(opt_option) > 0 and int(opt_option[0]) > 0)
150         else ' -g')
151
152  includes = (' -I ' + ' -I '.join(include_options)
153              if len(include_options) > 0
154              else '')
155
156  # Unfortunately, there are other options that have -c prefix too.
157  # So allowing only those look like C/C++ files.
158  src_files = [f for f in src_files if
159               re.search('\.cpp$|\.cc$|\.c$|\.cxx$|\.C$', f)]
160  srcs = ' '.join(src_files)
161  out = ' -o ' + out_file[0]
162
163  hipccopts = ' '
164  # In hip-clang environment, we need to make sure that hip header is included
165  # before some standard math header like <complex> is included in any source.
166  # Otherwise, we get build error.
167  # Also we need to retain warning about uninitialised shared variable as
168  # warning only, even when -Werror option is specified.
169  hipccopts += ' --include=hip/hip_runtime.h '
170  # Force C++17 dialect (note, everything in just one string!)
171  hipccopts += ' --std=c++17 '
172  # Use -fno-gpu-rdc by default for early GPU kernel finalization
173  # This flag would trigger GPU kernels be generated at compile time, instead
174  # of link time. This allows the default host compiler (gcc) be used as the
175  # linker for TensorFlow on ROCm platform.
176  hipccopts += ' -fno-gpu-rdc '
177  hipccopts += ' -fcuda-flush-denormals-to-zero '
178  hipccopts += undefines
179  hipccopts += defines
180  hipccopts += std_options
181  hipccopts += m_options
182
183  if depfiles:
184    # Generate the dependency file
185    depfile = depfiles[0]
186    cmd = (HIPCC_PATH + ' ' + hipccopts +
187           host_compiler_options +
188           ' -I .' + includes + ' ' + srcs + ' -M -o ' + depfile)
189    cmd = HIPCC_ENV.replace(';', ' ') + ' ' + cmd
190    if log: Log(cmd)
191    if VERBOSE: print(cmd)
192    exit_status = os.system(cmd)
193    if exit_status != 0:
194      return exit_status
195
196  cmd = (HIPCC_PATH + ' ' + hipccopts +
197         host_compiler_options + ' -fPIC' +
198         ' -I .' + opt + includes + ' -c ' + srcs + out)
199
200  cmd = HIPCC_ENV.replace(';', ' ') + ' '\
201        + cmd
202  if log: Log(cmd)
203  if VERBOSE: print(cmd)
204  return system(cmd)
205
206
207def main():
208  # ignore PWD env var
209  os.environ['PWD']=''
210
211  parser = ArgumentParser(fromfile_prefix_chars='@')
212  parser.add_argument('-x', nargs=1)
213  parser.add_argument('--rocm_log', action='store_true')
214  parser.add_argument('-pass-exit-codes', action='store_true')
215  args, leftover = parser.parse_known_args(sys.argv[1:])
216
217  if VERBOSE: print('PWD=' + os.getcwd())
218  if VERBOSE: print('HIPCC_ENV=' + HIPCC_ENV)
219
220  if args.x and args.x[0] == 'rocm':
221    # compilation for GPU objects
222    if args.rocm_log: Log('-x rocm')
223    leftover = [pipes.quote(s) for s in leftover]
224    if args.rocm_log: Log('using hipcc')
225    return InvokeHipcc(leftover, log=args.rocm_log)
226
227  elif args.pass_exit_codes:
228    # link
229    # with hipcc compiler invoked with -fno-gpu-rdc by default now, it's ok to
230    # use host compiler as linker, but we have to link with HCC/HIP runtime.
231    # Such restriction would be revised further as the bazel script get
232    # improved to fine tune dependencies to ROCm libraries.
233    gpu_linker_flags = [flag for flag in sys.argv[1:]
234                               if not flag.startswith(('--rocm_log'))]
235
236    gpu_linker_flags.append('-L' + ROCR_RUNTIME_PATH)
237    gpu_linker_flags.append('-Wl,-rpath=' + ROCR_RUNTIME_PATH)
238    gpu_linker_flags.append('-l' + ROCR_RUNTIME_LIBRARY)
239    gpu_linker_flags.append('-L' + HIP_RUNTIME_PATH)
240    gpu_linker_flags.append('-Wl,-rpath=' + HIP_RUNTIME_PATH)
241    gpu_linker_flags.append('-l' + HIP_RUNTIME_LIBRARY)
242    gpu_linker_flags.append("-lrt")
243    gpu_linker_flags.append("-lstdc++")
244
245    if VERBOSE: print(' '.join([CPU_COMPILER] + gpu_linker_flags))
246    return subprocess.call([CPU_COMPILER] + gpu_linker_flags)
247
248  else:
249    # compilation for host objects
250
251    # Strip our flags before passing through to the CPU compiler for files which
252    # are not -x rocm. We can't just pass 'leftover' because it also strips -x.
253    # We not only want to pass -x to the CPU compiler, but also keep it in its
254    # relative location in the argv list (the compiler is actually sensitive to
255    # this).
256    cpu_compiler_flags = [flag for flag in sys.argv[1:]
257                               if not flag.startswith(('--rocm_log'))]
258
259    # XXX: SE codes need to be built with gcc, but need this macro defined
260    cpu_compiler_flags.append("-D__HIP_PLATFORM_HCC__")
261    if VERBOSE: print(' '.join([CPU_COMPILER] + cpu_compiler_flags))
262    return subprocess.call([CPU_COMPILER] + cpu_compiler_flags)
263
264if __name__ == '__main__':
265  sys.exit(main())
266