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