Merge pull request #1260 from ydb-platform/mergelibs10

mergelibs10

1 files changed, 57 insertions, 0 deletions

diff --git a/build/scripts/link_exe.py b/build/scripts/link_exe.py
index e0b7c94809..980ac05d7e 100644
--- a/build/scripts/link_exe.py
+++ b/build/scripts/link_exe.py
@@ -1,3 +1,6 @@
+import itertools
+import os
+import os.path
 import sys
 import subprocess
 import optparse
@@ -40,6 +43,54 @@ CUDA_LIBRARIES = {
 }
 
 
+def prune_cuda_libraries(cmd, prune_arches, nvprune_exe, build_root):
+    def name_generator(prefix):
+        for idx in itertools.count():
+            yield prefix + '_' + str(idx)
+
+    def compute_arch(arch):
+        _, ver = arch.split('_', 1)
+        return 'compute_{}'.format(ver)
+
+    tmp_names_gen = name_generator('cuda_pruned_libs')
+
+    arch_args = []
+    for arch in prune_arches.split(':'):
+        arch_args.append('-gencode')
+        arch_args.append('arch={},code={}'.format(compute_arch(arch), arch))
+
+    flags = []
+    cuda_deps = set()
+    for flag in reversed(cmd):
+        if flag in CUDA_LIBRARIES:
+            cuda_deps.add('lib' + flag[2:] + '.a')
+            flag += '_pruned'
+        elif flag.startswith('-L') and any(f in cuda_deps for f in os.listdir(flag[2:])):
+            from_dirpath = flag[2:]
+            from_deps = list(cuda_deps & set(os.listdir(from_dirpath)))
+
+            if from_deps:
+                to_dirpath = os.path.abspath(os.path.join(build_root, next(tmp_names_gen)))
+                os.makedirs(to_dirpath)
+
+                for f in from_deps:
+                    # prune lib
+                    from_path = os.path.join(from_dirpath, f)
+                    to_path = os.path.join(to_dirpath, f[:-2] + '_pruned.a')
+                    subprocess.check_call([nvprune_exe] + arch_args + ['--output-file', to_path, from_path])
+                    cuda_deps.remove(f)
+
+                # do not remove current directory
+                # because it can contain other libraries we want link to
+                # instead we just add new directory with pruned libs
+                flags.append('-L' + to_dirpath)
+
+        flags.append(flag)
+
+    assert not cuda_deps, ('Unresolved CUDA deps: ' + ','.join(cuda_deps))
+    return reversed(flags)
+
+
 def remove_excessive_flags(cmd):
     flags = []
     for flag in cmd:
@@ -149,6 +200,10 @@ def parse_args():
     parser.add_option('--source-root')
     parser.add_option('--clang-ver')
     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('--build-root')
     parser.add_option('--arch')
     parser.add_option('--linker-output')
     parser.add_option('--whole-archive-peers', action='append')
@@ -176,6 +231,8 @@ if __name__ == '__main__':
 
     if opts.dynamic_cuda:
         cmd = fix_cmd_for_dynamic_cuda(cmd)
+    elif opts.cuda_architectures:
+        cmd = prune_cuda_libraries(cmd, opts.cuda_architectures, opts.nvprune_exe, opts.build_root)
     cmd = ProcessWholeArchiveOption(opts.arch, opts.whole_archive_peers, opts.whole_archive_libs).construct_cmd(cmd)
 
     if opts.custom_step: