Use `CUDA_ARCHITECTURES` flag to prune architectures

* Добавляю переменную `CUDA_ARCHITECTURES`, в которых указываю, для каких архитектур оставлять device-код в гпушных либах CUDA / cuDNN
* Выставляю для неё разумные дефолты для тех архитектур, которые потенциально страдают от избыточного размера кода (сейчас это 11.4, потом можно будет расширить для других архитектур)
* Пробрасываю эту переменную и путь до бинарника `nvprune` в скрипт линковки
* Ищу и пруню гпушные либы

Как проверить, что работает:
* например, успешно собирается `ya make --build=release -DTENSORFLOW_WITH_CUDA -DCUDA_VERSION=11.4 -DCUDNN_VERSION=8.0.5 -DCUDA_ARCHITECTURES=sm_70 yweb/webdaemons/ocrdaemon` (с дефолтным `-DCUDA_ARCHITECTURES` падает по relocation overflow)
* елси хочется посмотреть на изменение размеров, то можно собрать `ya make ml/zeliboba/libs/ynmt_lm/score/bin/ --build=relwithdebinfo -DCUDA_VERSION=11.4 -DCUDNN_VERSION=8.0.5` один раз с дефолтным значением, другой раз с `-DCUDA_ARCHITECTURES=sm_70,sm_80,compute_80` - размер бинаря уменьшится

Особенности:
* ~~Сейчас дефолт для каждой новой версии куды содержит в себе тупо все поддерживаемые архитектуры. Имеет смысл его порезать до чего-то более разумного~~ - порезал
* Обход аргументов в обратном порядке сделан для того, чтобы эмулировать поведение линкера (для линкера важен порядок пробрасывания между либами для линковки и директориями, в которых их нужно искать)
* ~~Сейчас артефакты прунинга кладутся рядом с оригинальными либами; возможно, стоит их складывать в отдельное место в build-директории, но я пока не разобрался как это делать; буду рад, если кто-то подскажет~~ уже неактуально - кладу в build_root
* nvprune имеет свойство мусорить варнингом в stderr, когда прунит `libcudart_static.a` (это норма - там нет кода, который можно было бы попрунить); как вариант можно закостылить и не прунить `libcudart_static.a`, но я открыт к предложениям

1 files changed, 56 insertions, 0 deletions

diff --git a/build/scripts/link_exe.py b/build/scripts/link_exe.py
index e0b7c94809..a5744e5f25 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,9 @@ 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')
+    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 +230,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: