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`, но я открыт к предложениям

3 files changed, 99 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:
diff --git a/build/ymake.core.conf b/build/ymake.core.conf
index fd7c7877d7..d00c3c7a86 100644
--- a/build/ymake.core.conf
+++ b/build/ymake.core.conf
@@ -1042,6 +1042,12 @@ module _LINK_UNIT: _BASE_UNIT {
         LINK_SCRIPT_EXE_FLAGS += --dynamic-cuda
     }
 
+    when ($CUDA_ARCHITECTURES && $USE_DYNAMIC_CUDA != "yes") {
+        LINK_SCRIPT_EXE_FLAGS+=--cuda-architectures $CUDA_ARCHITECTURES
+        LINK_SCRIPT_EXE_FLAGS+=--nvprune-exe $CUDA_ROOT/bin/nvprune
+        LINK_SCRIPT_EXE_FLAGS+=--build-root $(BUILD_ROOT)
+    }
+
     when ($OPENSOURCE == "yes" && $AUTOCHECK == "yes") {
         # FIXME: Replace AUTOCHECK == yes with _not a host platform_ check after YMAKE-218
         MODULE_LICENSES_RESTRICTION_TYPES = ALLOW_ONLY
diff --git a/build/ymake_conf.py b/build/ymake_conf.py
index 431c8a4f04..acf6dfd9a0 100755
--- a/build/ymake_conf.py
+++ b/build/ymake_conf.py
@@ -2260,6 +2260,7 @@ class Cuda(object):
 
         self.cuda_root = Setting('CUDA_ROOT')
         self.cuda_version = Setting('CUDA_VERSION', auto=self.auto_cuda_version, convert=self.convert_major_version, rewrite=True)
+        self.cuda_architectures = Setting('CUDA_ARCHITECTURES', auto=self.auto_cuda_architectures, rewrite=True)
         self.use_arcadia_cuda = Setting('USE_ARCADIA_CUDA', auto=self.auto_use_arcadia_cuda, convert=to_bool)
         self.use_arcadia_cuda_host_compiler = Setting('USE_ARCADIA_CUDA_HOST_COMPILER', auto=self.auto_use_arcadia_cuda_host_compiler, convert=to_bool)
         self.cuda_use_clang = Setting('CUDA_USE_CLANG', auto=False, convert=to_bool)
@@ -2307,6 +2308,7 @@ class Cuda(object):
 
         self.cuda_root.emit()
         self.cuda_version.emit()
+        self.cuda_architectures.emit()
         self.use_arcadia_cuda.emit()
         self.use_arcadia_cuda_host_compiler.emit()
         self.cuda_use_clang.emit()
@@ -2387,6 +2389,41 @@ class Cuda(object):
         else:
             return value
 
+    def auto_cuda_architectures(self):
+        # empty list does not mean "no architectures"
+        # it means "no restriction -- any available architecture"
+
+        host, target = self.build.host_target
+        if not target.is_linux_x86_64:
+            # do not impose any restrictions, when build not for "linux 64-bit"
+            return []
+
+        # do not include 'lto' type,
+        # because we already perform static linking
+        supported_types = ['compute', 'sm']
+
+        # Equality to CUDA 11.4 is rather strict comparison
+        # TODO: find out how we can relax check (e.g. to include more version of CUDA toolkit)
+        if self.cuda_version.value == '11.4':
+            # * use output of CUDA 11.4 `nvcc --help`
+            # * drop support for '53', '62', '72' and '87'
+            #   (these devices run only on arm64)
+            # * drop support for '37'
+            #   the single place it's used in Arcadia is https://a.yandex-team.ru/arcadia/sdg/sdc/third_party/cub/common.mk?rev=r13268523#L69
+            supported_vers = ['35',
+                              '50', '52',
+                              '60', '61',
+                              '70', '75',
+                              '80', '86']
+        else:
+            supported_vers = []
+
+        cuda_architectures = ['{typ}_{ver}'.format(typ=typ, ver=ver)
+                              for typ in supported_types
+                              for ver in supported_vers]
+
+        return ','.join(cuda_architectures)
+
     def auto_use_arcadia_cuda(self):
         return not self.cuda_root.from_user