intermediate changes

ref:f1a76bb520860c96f863dde2f5dfa5e45b9ea67b

1 files changed, 46 insertions, 6 deletions

diff --git a/build/scripts/link_exe.py b/build/scripts/link_exe.py
index f469e3b442..bec23e7c30 100644
--- a/build/scripts/link_exe.py
+++ b/build/scripts/link_exe.py
@@ -15,11 +15,45 @@ def get_leaks_suppressions(cmd):
     return supp, newcmd
 
 
-musl_libs = '-lc', '-lcrypt', '-ldl', '-lm', '-lpthread', '-lrt', '-lutil'
-
-
-def fix_cmd(musl, c):
-    return [i for i in c if (not musl or i not in musl_libs) and not i.endswith('.ios.interface') and not i.endswith('.pkg.fake')]
+MUSL_LIBS = '-lc', '-lcrypt', '-ldl', '-lm', '-lpthread', '-lrt', '-lutil'
+
+
+CUDA_LIBRARIES = {
+    '-llapack_static': '-llapack',
+    '-lcublas_static': '-lcublas',
+    '-lcublasLt_static': '-lcublasLt',
+    '-lcudart_static': '-lcudart',
+    '-lcufft_static_nocallback': '-lcufft',
+    '-lcurand_static': '-lcurand',
+    '-lcusolver_static': '-lcusolver',
+    '-lcusparse_static': '-lcusparse',
+}
+
+
+def remove_excessive_flags(cmd):
+    flags = []
+    for flag in cmd:
+        if not flag.endswith('.ios.interface') and not flag.endswith('.pkg.fake'):
+            flags.append(flag)
+    return flags
+
+
+def fix_cmd_for_musl(cmd):
+    flags = []
+    for flag in cmd:
+        if flag not in MUSL_LIBS:
+            flags.append(flag)
+    return flags
+
+
+def fix_cmd_for_dynamic_cuda(cmd):
+    flags = []
+    for flag in cmd:
+        if flag in CUDA_LIBRARIES:
+            flags.append(CUDA_LIBRARIES[flag])
+        else:
+            flags.append(flag)
+    return flags
 
 
 def gen_default_suppressions(inputs, output, source_root):
@@ -50,6 +84,7 @@ def parse_args():
     parser.add_option('--custom-step')
     parser.add_option('--python')
     parser.add_option('--source-root')
+    parser.add_option('--dynamic-cuda', action='store_true')
     parser.add_option('--arch')
     parser.add_option('--linker-output')
     parser.add_option('--whole-archive-peers', action='append')
@@ -60,7 +95,12 @@ def parse_args():
 if __name__ == '__main__':
     opts, args = parse_args()
 
-    cmd = fix_cmd(opts.musl, args)
+    cmd = remove_excessive_flags(args)
+    if opts.musl:
+        cmd = fix_cmd_for_musl(cmd)
+    
+    if opts.dynamic_cuda:
+        cmd = fix_cmd_for_dynamic_cuda(cmd)
     cmd = ProcessWholeArchiveOption(opts.arch, opts.whole_archive_peers, opts.whole_archive_libs).construct_cmd(cmd)
 
     if opts.custom_step: