[SE] Let users specify CUDA path
[lldb.git] / parallel-libs / streamexecutor / tools / streamexecutor-config / streamexecutor-config.in
1 #!@PYTHON_EXECUTABLE@
2 #
3 #===- streamexecutor-config - Build config script for SE -----*- python -*--===#
4 #
5 #                     The LLVM Compiler Infrastructure
6 #
7 # This file is distributed under the University of Illinois Open Source
8 # License. See LICENSE.TXT for details.
9 #
10 #===------------------------------------------------------------------------===#
11
12 r"""
13 Get configuration info needed to compile programs which use StreamExecutor.
14
15 Runs llvm-config and adds StreamExecutor-specific flags to the output. Supports
16 only the subset of llvm-config flags that are relevant for applications
17 compiling against StreamExecutor.
18
19 This utility will typically be used to construct a compile command line for an
20 application which depends on the StreamExecutor library.
21
22 For example:
23   c++ example.cpp -o example \
24       $(streamexecutor-config \
25           --cppflags --cxxflags --ldflags --libs --system-libs)
26 """
27
28 import argparse
29 import errno
30 import os
31 import shlex
32 import subprocess
33 import sys
34
35 # The following functions are configured by cmake. They use raw triple-quoted
36 # strings to surround values that are substituted by cmake at configure time.
37 # This kind of quoting should allow for paths that contain spaces.
38
39 def get_llvm_config_dir():
40   """Gets the path to the llvm-config executable."""
41   return r"""@LLVM_BINARY_DIR@/bin"""
42
43 def get_cmake_install_prefix():
44   """Gets the value of the cmake variable CMAKE_INSTALL_PREFIX."""
45   return r"""@CMAKE_INSTALL_PREFIX@"""
46
47 def convert_library_name(library_name):
48   """Converts a library name ending in '.framework' into a '-framework' flag.
49
50   This is used to support OS X.
51
52   >>> convert_library_name('')
53   ''
54
55   >>> convert_library_name('/usr/local/lib64/libcuda.so')
56   '/usr/local/lib64/libcuda.so'
57
58   >>> convert_library_name('/Library/Frameworks/cuda.framework')
59   '-framework cuda'
60   """
61   framework_suffix = '.framework'
62   if library_name.endswith(framework_suffix):
63     framework_name = os.path.basename(library_name)[:-len(framework_suffix)]
64     library_name = '-framework ' + framework_name
65   return library_name
66
67 def get_cuda_driver_library():
68   """Gets the value of the cmake variable CUDA_DRIVER_LIBRARY."""
69   return convert_library_name(r"""@CUDA_DRIVER_LIBRARY@""")
70
71 def cuddle_flag(flag, tokens):
72   """If flag appears by itself in tokens, combines it with the next token.
73
74   >>> tokens = ['-I', '/usr/include']
75   >>> cuddle_flag('-I', tokens)
76   >>> tokens
77   ['-I/usr/include']
78
79   >>> tokens = ['-L', '/usr/lib']
80   >>> cuddle_flag('-L', tokens)
81   >>> tokens
82   ['-L/usr/lib']
83
84   >>> tokens = ['-I']
85   >>> cuddle_flag('-I', tokens)
86   >>> tokens
87   ['-I']
88
89   >>> tokens = ['-I', '/usr/include', '-I', '/usr/local/include']
90   >>> cuddle_flag('-I', tokens)
91   >>> tokens
92   ['-I/usr/include', '-I/usr/local/include']
93   """
94   start = 0
95   while True:
96     try:
97       index = tokens.index(flag, start)
98     except ValueError:
99       return
100     if index + 1 < len(tokens):
101       follower = tokens.pop(index + 1)
102       tokens[index] = flag + follower
103     start = index + 1
104
105 def get_llvm_config_output_for_dir(llvm_config_dir, flags_string):
106   """Calls llvm-config at the given path and returns the output with -I and -L
107   flags cuddled."""
108   output = subprocess.check_output(
109       ['%s/llvm-config' % llvm_config_dir] + flags_string.split()).strip()
110   tokens = shlex.split(output)
111   cuddle_flag('-I', tokens)
112   cuddle_flag('-L', tokens)
113   return ' '.join(tokens)
114
115 def has_token(token, string):
116   """Checks if the given token appears in the string.
117
118   The token argument must be a single shell token.
119
120   >>> string = '-I/usr/include -L"/usr/lib"'
121   >>> has_token('-I/usr/include', string)
122   True
123   >>> has_token('-I/usr/local/include', string)
124   False
125   >>> has_token('-I"/usr/include"', string)
126   True
127   >>> has_token('-L"/usr/lib"', string)
128   True
129   >>> has_token('-L/usr/lib', string)
130   True
131   """
132   split_token = shlex.split(token)
133   if len(split_token) > 1:
134     raise ValueError('has_token called with a multi-token token: ' + token)
135   escaped_token = split_token[0]
136   return escaped_token in shlex.split(string)
137
138 def main():
139   parser = argparse.ArgumentParser(
140       prog='streamexecutor-config',
141       formatter_class=argparse.RawDescriptionHelpFormatter,
142       description=__doc__)
143
144   parser.add_argument(
145       '--cppflags',
146       action='store_true',
147       help=
148         'C preprocessor flags for files that include StreamExecutor headers.')
149
150   parser.add_argument(
151       '--cxxflags',
152       action='store_true',
153       help='C++ compiler flags for files that include StreamExecutor headers.')
154
155   parser.add_argument(
156       '--ldflags',
157       action='store_true',
158       help='Print linker flags.')
159
160   parser.add_argument(
161       '--libs',
162       action='store_true',
163       help='Libraries needed to link against StreamExecutor.')
164
165   parser.add_argument(
166       '--system-libs',
167       action='store_true',
168       help='System libraries needed to link against StreamExecutor.')
169
170   parser.add_argument(
171       '--llvm-config-dir',
172       default=get_llvm_config_dir(),
173       help='Directory containing the llvm-config executable. '\
174           'If not specified, defaults to the cmake-configured location')
175
176   args = parser.parse_args()
177
178   # Print the help message if the user did not pass any flag arguments.
179   if not any(
180       getattr(args, flag)
181         for flag in ('cppflags', 'cxxflags', 'ldflags', 'libs', 'system_libs')):
182     parser.print_help()
183     sys.exit(1)
184
185   # Check for the presence of the llvm-config executable.
186   if not os.path.isfile('%s/llvm-config' % args.llvm_config_dir):
187     sys.exit('llvm-config not found in: ' + args.llvm_config_dir)
188   if not os.access('%s/llvm-config' % args.llvm_config_dir, os.X_OK):
189     sys.exit('llvm-config not executable in: ' + args.llvm_config_dir)
190
191   # We will always use args.llvm_config_dir as the second argument to
192   # get_llvm_config_output_for_path.
193   get_llvm_config_output = lambda flags : get_llvm_config_output_for_dir(
194       args.llvm_config_dir, flags)
195
196   all_flags = []
197
198   if args.cppflags:
199     llvm_flags = get_llvm_config_output('--cppflags')
200     all_flags.append(llvm_flags)
201     se_flag = "-I%s/include" % get_cmake_install_prefix()
202     if not has_token(token=se_flag, string=llvm_flags):
203       all_flags.append(se_flag)
204
205   if args.cxxflags:
206     all_flags.append(get_llvm_config_output('--cxxflags'))
207
208   if args.ldflags:
209     llvm_flags = get_llvm_config_output('--ldflags')
210     all_flags.append(llvm_flags)
211     se_flag = "-L%s/lib" % get_cmake_install_prefix()
212     if not has_token(token=se_flag, string=llvm_flags):
213       all_flags.append(se_flag)
214
215   if args.libs:
216     llvm_flags = get_llvm_config_output('--libs support symbolize')
217     se_flag = '-lstreamexecutor'
218     if not has_token(token=se_flag, string=llvm_flags):
219       all_flags.append(se_flag)
220     cuda_driver_library = get_cuda_driver_library()
221     if cuda_driver_library:
222       all_flags.append(cuda_driver_library)
223     all_flags.append(llvm_flags)
224
225   if args.system_libs:
226     all_flags.append(get_llvm_config_output('--system-libs'))
227
228   print(' '.join(all_flags))
229
230 if __name__ == '__main__':
231   main()