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 with nvcc on Windows. 18 19DESCRIPTION: 20 This script is the Windows version of //third_party/gpus/crosstool/crosstool_wrapper_is_not_gcc 21""" 22 23from argparse import ArgumentParser 24import os 25import subprocess 26import re 27import sys 28import tempfile 29 30# Template values set by cuda_autoconf. 31CPU_COMPILER = ('%{cpu_compiler}') 32GCC_HOST_COMPILER_PATH = ('%{gcc_host_compiler_path}') 33 34NVCC_PATH = '%{nvcc_path}' 35NVCC_VERSION = '%{cuda_version}' 36NVCC_TEMP_DIR = "%{nvcc_tmp_dir}" 37 38def Log(s): 39 print('gpus/crosstool: {0}'.format(s)) 40 41 42def GetOptionValue(argv, option): 43 """Extract the list of values for option from options. 44 45 Args: 46 option: The option whose value to extract. 47 48 Returns: 49 1. A list of values, either directly following the option, 50 (eg., /opt val1 val2) or values collected from multiple occurrences of 51 the option (eg., /opt val1 /opt val2). 52 2. The leftover options. 53 """ 54 55 parser = ArgumentParser(prefix_chars='-/') 56 parser.add_argument(option, nargs='*', action='append') 57 option = option.lstrip('-/').replace('-', '_') 58 args, leftover = parser.parse_known_args(argv) 59 if args and vars(args)[option]: 60 return (sum(vars(args)[option], []), leftover) 61 return ([], leftover) 62 63def _update_options(nvcc_options): 64 if NVCC_VERSION in ("7.0",): 65 return nvcc_options 66 67 update_options = { "relaxed-constexpr" : "expt-relaxed-constexpr" } 68 return [ update_options[opt] if opt in update_options else opt 69 for opt in nvcc_options ] 70 71def GetNvccOptions(argv): 72 """Collect the -nvcc_options values from argv. 73 74 Args: 75 argv: A list of strings, possibly the argv passed to main(). 76 77 Returns: 78 1. The string that can be passed directly to nvcc. 79 2. The leftover options. 80 """ 81 82 parser = ArgumentParser() 83 parser.add_argument('-nvcc_options', nargs='*', action='append') 84 85 args, leftover = parser.parse_known_args(argv) 86 87 if args.nvcc_options: 88 options = _update_options(sum(args.nvcc_options, [])) 89 return (['--' + a for a in options], leftover) 90 return ([], leftover) 91 92 93def InvokeNvcc(argv, log=False): 94 """Call nvcc with arguments assembled from argv. 95 96 Args: 97 argv: A list of strings, possibly the argv passed to main(). 98 log: True if logging is requested. 99 100 Returns: 101 The return value of calling os.system('nvcc ' + args) 102 """ 103 104 src_files = [f for f in argv if 105 re.search('\.cpp$|\.cc$|\.c$|\.cxx$|\.C$', f)] 106 if len(src_files) == 0: 107 raise Error('No source files found for cuda compilation.') 108 109 out_file = [ f for f in argv if f.startswith('/Fo') ] 110 if len(out_file) != 1: 111 raise Error('Please specify exactly one output file for cuda compilation.') 112 out = ['-o', out_file[0][len('/Fo'):]] 113 114 nvcc_compiler_options, argv = GetNvccOptions(argv) 115 116 opt_option, argv = GetOptionValue(argv, '/O') 117 opt = ['-g'] 118 if (len(opt_option) > 0 and opt_option[0] != 'd'): 119 opt = ['-O2'] 120 121 include_options, argv = GetOptionValue(argv, '/I') 122 includes = ["-I " + include for include in include_options] 123 124 defines, argv = GetOptionValue(argv, '/D') 125 defines = ['-D' + define for define in defines] 126 127 undefines, argv = GetOptionValue(argv, '/U') 128 undefines = ['-U' + define for define in undefines] 129 130 fatbin_options, argv = GetOptionValue(argv, '-Xcuda-fatbinary') 131 fatbin_options = ['--fatbin-options=' + option for option in fatbin_options] 132 133 # The rest of the unrecognized options should be passed to host compiler 134 host_compiler_options = [option for option in argv if option not in (src_files + out_file)] 135 136 m_options = ["-m64"] 137 138 nvccopts = ['-D_FORCE_INLINES'] 139 compute_capabilities, argv = GetOptionValue(argv, "--cuda-gpu-arch") 140 for capability in compute_capabilities: 141 capability = capability[len('sm_'):] 142 nvccopts += [ 143 r'-gencode=arch=compute_%s,"code=sm_%s"' % (capability, capability) 144 ] 145 compute_capabilities, argv = GetOptionValue(argv, '--cuda-include-ptx') 146 for capability in compute_capabilities: 147 capability = capability[len('sm_'):] 148 nvccopts += [ 149 r'-gencode=arch=compute_%s,"code=compute_%s"' % (capability, capability) 150 ] 151 _, argv = GetOptionValue(argv, '--no-cuda-include-ptx') 152 153 # nvcc doesn't respect the INCLUDE and LIB env vars from MSVC, 154 # so we explicity specify the system include paths and library search paths. 155 if 'INCLUDE' in os.environ: 156 nvccopts += [('--system-include="%s"' % p) for p in os.environ['INCLUDE'].split(";")] 157 if 'LIB' in os.environ: 158 nvccopts += [('--library-path="%s"' % p) for p in os.environ['LIB'].split(";")] 159 160 nvccopts += nvcc_compiler_options 161 nvccopts += undefines 162 nvccopts += defines 163 nvccopts += m_options 164 nvccopts += fatbin_options 165 nvccopts += ['--compiler-options=' + ",".join(host_compiler_options)] 166 nvccopts += ['-x', 'cu'] + opt + includes + out + ['-c'] + src_files 167 # Specify a unique temp directory for nvcc to generate intermediate files, 168 # then Bazel can ignore files under NVCC_TEMP_DIR during dependency check 169 # http://docs.nvidia.com/cuda/cuda-compiler-driver-nvcc/index.html#options-for-guiding-compiler-driver 170 # Different actions are sharing NVCC_TEMP_DIR, so we cannot remove it if the directory already exists. 171 if os.path.isfile(NVCC_TEMP_DIR): 172 os.remove(NVCC_TEMP_DIR) 173 if not os.path.exists(NVCC_TEMP_DIR): 174 os.makedirs(NVCC_TEMP_DIR) 175 # Provide a unique dir for each compiling action to avoid conflicts. 176 tempdir = tempfile.mkdtemp(dir = NVCC_TEMP_DIR) 177 nvccopts += ['--keep', '--keep-dir', tempdir] 178 # Force C++17 dialect (note, everything in just one string!) 179 nvccopts += ['--std c++17'] 180 if log: 181 Log([NVCC_PATH] + nvccopts) 182 183 # Store command line options in a file to avoid hitting the character limit. 184 optsfile = tempfile.NamedTemporaryFile(mode='w', dir=tempdir, delete=False) 185 optsfile.write("\n".join(nvccopts)) 186 optsfile.close() 187 188 proc = subprocess.Popen([NVCC_PATH, "--options-file", optsfile.name], 189 stdout=sys.stdout, 190 stderr=sys.stderr, 191 env=os.environ.copy(), 192 shell=True) 193 proc.wait() 194 return proc.returncode 195 196def main(): 197 parser = ArgumentParser() 198 parser.add_argument('-x', nargs=1) 199 parser.add_argument('--cuda_log', action='store_true') 200 args, leftover = parser.parse_known_args(sys.argv[1:]) 201 202 if args.x and args.x[0] == 'cuda': 203 if args.cuda_log: Log('-x cuda') 204 if args.cuda_log: Log('using nvcc') 205 return InvokeNvcc(leftover, log=args.cuda_log) 206 207 # Strip our flags before passing through to the CPU compiler for files which 208 # are not -x cuda. We can't just pass 'leftover' because it also strips -x. 209 # We not only want to pass -x to the CPU compiler, but also keep it in its 210 # relative location in the argv list (the compiler is actually sensitive to 211 # this). 212 cpu_compiler_flags = [flag for flag in sys.argv[1:] 213 if not flag.startswith(('--cuda_log')) 214 and not flag.startswith(('-nvcc_options'))] 215 216 return subprocess.call([CPU_COMPILER] + cpu_compiler_flags) 217 218if __name__ == '__main__': 219 sys.exit(main()) 220