b6ee53f9af602a963f352ed37c717dead1f8cef3

1 files changed, 9 insertions, 0 deletions

diff --git a/build/scripts/link_dyn_lib.py b/build/scripts/link_dyn_lib.py
index 3b7c0d3e51..72937544b6 100644
--- a/build/scripts/link_dyn_lib.py
+++ b/build/scripts/link_dyn_lib.py
@@ -7,6 +7,7 @@ import optparse
 import pipes
 
 import thinlto_cache
+import link_exe
 
 from process_whole_archive_option import ProcessWholeArchiveOption
 from fix_py2_protobuf import fix_py2
@@ -215,6 +216,10 @@ def parse_args():
     parser.add_option('--linker-output')
     parser.add_option('--musl', action='store_true')
     parser.add_option('--dynamic-cuda', action='store_true')
+    parser.add_option('--cuda-architectures',
+                      help='List of supported CUDA architectures, separated by ":" (e.g. "sm_52:compute_70:lto_90a"')
+    parser.add_option('--nvprune-exe')
+    parser.add_option('--objcopy-exe')
     parser.add_option('--whole-archive-peers', action='append')
     parser.add_option('--whole-archive-libs', action='append')
     parser.add_option('--custom-step')
@@ -237,6 +242,10 @@ if __name__ == '__main__':
         cmd = fix_cmd_for_musl(cmd)
     if opts.dynamic_cuda:
         cmd = fix_cmd_for_dynamic_cuda(cmd)
+    else:
+        cuda_manager = link_exe.CUDAManager(opts.cuda_architectures, opts.nvprune_exe)
+        cmd = link_exe.process_cuda_libraries_by_nvprune(cmd, cuda_manager, opts.build_root)
+        cmd = link_exe.process_cuda_libraries_by_objcopy(cmd, opts.build_root, opts.objcopy_exe)
 
     cmd = ProcessWholeArchiveOption(opts.arch, opts.whole_archive_peers, opts.whole_archive_libs).construct_cmd(cmd)
     thinlto_cache.preprocess(opts, cmd)