torch.utils.cpp_extension

torch.utils.cpp_extension.CppExtension(name, sources, *args, **kwargs) [source]

Создаёт setuptools.Extension для C++.

Удобный метод, создающий setuptools.Extension с минимально необходимыми (но часто достаточными) аргументами для построения расширения C++.

Все аргументы передаются конструктору setuptools.Extension.

Пример

>>> from setuptools import setup
>>> from torch.utils.cpp_extension import BuildExtension, CppExtension
>>> setup(
...     name='extension',
...     ext_modules=[
...         CppExtension(
...             name='extension',
...             sources=['extension.cpp'],
...             extra_compile_args=['-g']),
...     ],
...     cmdclass={
...         'build_ext': BuildExtension
...     })

torch.utils.cpp_extension.CUDAExtension(name, sources, *args, **kwargs) [source]

Создаёт setuptools.Extension для CUDA/C++.

Удобный метод, создающий setuptools.Extension с минимально необходимыми (но часто достаточными) аргументами для построения расширения CUDA/C++. Включает путь к заголовочным файлам CUDA, путь к библиотекам и библиотекам времени выполнения.

Все аргументы передаются конструктору setuptools.Extension.

Пример

>>> from setuptools import setup
>>> from torch.utils.cpp_extension import BuildExtension, CUDAExtension
>>> setup(
...     name='cuda_extension',
...     ext_modules=[
...         CUDAExtension(
...                 name='cuda_extension',
...                 sources=['extension.cpp', 'extension_kernel.cu'],
...                 extra_compile_args={'cxx': ['-g'],
...                                     'nvcc': ['-O2']})
...     ],
...     cmdclass={
...         'build_ext': BuildExtension
...     })

Возможности вычислений:

По умолчанию расширение будет скомпилировано для работы на всех архитектурах карт, видимых во время процесса построения расширения, плюс PTX. Если в будущем будет установлена новая карта, расширение может потребовать перекомпиляции. Если у видимой карты есть возможность вычисления (CC), которая новее, чем самая последняя версия, для которой ваш nvcc может построить полностью скомпилированные двоичные файлы, Pytorch заставит nvcc вернуться к построению ядер с самой последней версией PTX, которую поддерживает ваш nvcc (см. подробности ниже о PTX).

Вы можете переопределить стандартное поведение, используя TORCH_CUDA_ARCH_LIST для явного указания поддерживаемых CC:

TORCH_CUDA_ARCH_LIST=”6.1 8.6” python build_my_extension.py TORCH_CUDA_ARCH_LIST=”5.2 6.0 6.1 7.0 7.5 8.0 8.6+PTX” python build_my_extension.py

Опция +PTX заставляет двоичные файлы ядер расширения включать инструкции PTX для указанного CC. PTX — это промежуточное представление, которое позволяет ядрам динамически компилироваться для любого CC ≥ указанного CC (например, 8.6+PTX генерирует PTX, который может динамически компилироваться для любой видеокарты с CC ≥ 8.6). Это улучшает совместимость с будущими версиями.

Однако, полагаться на более старую PTX для обеспечения совместимости путём динамической компиляции для более новых CC может незначительно снизить производительность на этих новых CC. Если вы знаете точные CC(ы) целевых видеокарт, вам всегда лучше указать их индивидуально. Например, если вы хотите, чтобы ваше расширение работало на 8.0 и 8.6, «8.0+PTX» будет функционально работать, потому что оно включает PTX, который может динамически компилироваться для 8.6, но «8.0 8.6» будет лучше.

Обратите внимание, что, хотя включение всех поддерживаемых архитектур возможно, чем больше архитектур включено, тем медленнее будет процесс построения, так как он будет создавать отдельный образ ядра для каждой архитектуры.

Обратите внимание, что CUDA-11.5 nvcc столкнётся с внутренней ошибкой компилятора при обработке torch/extension.h в Windows. Чтобы обойти проблему, перенесите логику привязки Python в чистый файл C++.

Пример использования:

>>> #include <ATen/ATen.h>
>>> at::Tensor SigmoidAlphaBlendForwardCuda(....)

Вместо:

>>> #include <torch/extension.h>
>>> torch::Tensor SigmoidAlphaBlendForwardCuda(...)

В настоящее время открыта проблема с ошибкой nvcc: https://github.com/pytorch/pytorch/issues/69460 Полный пример решения: https://github.com/facebookresearch/pytorch3d/commit/cb170ac024a949f1f9614ffe6af1c38d972f7d48

Связывание переносимого кода устройства:

Если вы хотите ссылаться на символы устройства в разных модулях компиляции (в разных объектных файлах), объектные файлы должны быть построены с relocatable device code (-rdc=true или -dc). Исключением из этого правила является «динамичная параллельность» (вложенные запуски ядер), которая используется уже не так часто. Relocatable device code менее оптимизирован, поэтому его следует использовать только в объектных файлах, которые в нём нуждаются. Использование -dlto (оптимизация времени связывания устройств) на этапе компиляции кода устройства и dlink поможет уменьшить потенциальное снижение производительности -rdc. Обратите внимание, что для достижения эффекта его нужно использовать на обоих этапах.

Если у вас есть rdc объекты, вам нужно выполнить дополнительный -dlink (связывание устройств) шаг перед шагом связывания символов процессора. Также существует случай, когда -dlink используется без -rdc: когда расширение связано со статической библиотекой, содержащей объекты, скомпилированные с rdc, такие как [NVSHMEM library](https://developer.nvidia.com/nvshmem).

Примечание: Для построения расширения CUDA с связыванием RDC требуется Ninja.

Пример

>>> CUDAExtension(
...        name='cuda_extension',
...        sources=['extension.cpp', 'extension_kernel.cu'],
...        dlink=True,
...        dlink_libraries=["dlink_lib"],
...        extra_compile_args={'cxx': ['-g'],
...                            'nvcc': ['-O2', '-rdc=true']})

torch.utils.cpp_extension.BuildExtension(*args, **kwargs) [source]

Настраиваемое setuptools построение расширения.

Этот setuptools.build_ext подкласс заботится о передаче минимально необходимых флагов компилятора (например, -std=c++14) а также смешанной C++/CUDA компиляции (и поддержке файлов CUDA в целом).

При использовании BuildExtension, разрешается передать словарь для extra_compile_args (вместо обычного списка), который сопоставляет языки (cxx или nvcc) со списком дополнительных флагов компилятора для передачи компилятору. Это позволяет передавать разные флаги C++ и CUDA компилятору во время смешанной компиляции.

use_ninja (bool): Если use_ninja равняется True (по умолчанию), мы пытаемся построить используя бэкэнд Ninja. Ninja значительно ускоряет компиляцию по сравнению со стандартным setuptools.build_ext. Возвращается к стандартному бэкэнду distutils, если Ninja недоступен.

Примечание

По умолчанию бэкэнд Ninja использует #CPUS + 2 работника для построения расширения. Это может использовать слишком много ресурсов на некоторых системах. Количество работников можно контролировать, задав переменную окружения MAX_JOBS на любое неотрицательное число.

torch.utils.cpp_extension.load(name, sources, extra_cflags=None, extra_cuda_cflags=None, extra_ldflags=None, extra_include_paths=None, build_directory=None, verbose=False, with_cuda=None, is_python_module=True, is_standalone=False, keep_intermediates=True) [source]

Загрузка расширения PyTorch на C++ в режиме реального времени (JIT).

Для загрузки расширения генерируется файл сборки Ninja, который используется для компиляции указанных исходных файлов в динамическую библиотеку. Эта библиотека затем загружается в текущий процесс Python в качестве модуля и возвращается этой функцией, готовой к использованию.

По умолчанию, каталог, в который генерируется файл сборки и в который компилируется результирующая библиотека, это <tmp>/torch_extensions/<name>, где <tmp> — временная папка текущей платформы, а <name> — имя расширения. Это расположение можно переопределить двумя способами. Во-первых, если переменная среды TORCH_EXTENSIONS_DIR установлена, она заменяет <tmp>/torch_extensions и все расширения будут компилироваться в подкаталоги этого каталога. Во-вторых, если аргумент build_directory в этой функции указан, он переопределяет весь путь, то есть библиотека будет скомпилирована непосредственно в этот каталог.

Для компиляции исходных кодов используется компилятор системы по умолчанию (c++), который можно переопределить, установив переменную среды CXX. Чтобы передать дополнительные аргументы в процесс компиляции, можно использовать extra_cflags или extra_ldflags. Например, чтобы скомпилировать расширение с оптимизациями, передайте extra_cflags=['-O3']. Также можно использовать extra_cflags для передачи дополнительных каталогов включаемых файлов.

Поддержка CUDA с смешанной компиляцией предоставляется. Просто передайте файлы исходных кодов CUDA (.cu или .cuh) вместе с другими исходными файлами. Такие файлы будут обнаружены и скомпилированы с помощью nvcc вместо компилятора C++. Это включает передачу каталога CUDA lib64 в качестве каталога библиотек и подключение cudart. Вы можете передать дополнительные флаги в nvcc через extra_cuda_cflags, так же, как и extra_cflags для C++.

Используются различные эвристики для поиска каталога установки CUDA, которые обычно работают хорошо. Если нет, установка переменной среды CUDA_HOME является самым безопасным вариантом.

Параметры:

• name – Имя расширения для построения. ОНО ДОЛЖНО быть таким же, как имя модуля pybind11!
• sources (Union[str, List[str]]) – Список относительных или абсолютных путей к файлам исходных кодов C++.
• extra_cflags – необязательный список флагов компилятора, передаваемых для сборки.
• extra_cuda_cflags – необязательный список флагов компилятора, передаваемых для nvcc при построении исходных кодов CUDA.
• extra_ldflags – необязательный список флагов компоновщика, передаваемых для сборки.
• extra_include_paths – необязательный список каталогов включаемых файлов, передаваемых для сборки.
• build_directory – необязательный путь для использования в качестве рабочей области сборки.
• verbose – Если True, включает подробную запись логов шагов загрузки.
• with_cuda (Optional[bool]) – Определяет, добавляются ли заголовки и библиотеки CUDA к сборке. Если установлено None (по умолчанию), это значение автоматически определяется на основе наличия .cu или .cuh в sources. Установите его в True` для принудительного включения заголовков и библиотек CUDA.
• is_python_module – Если True (по умолчанию), импортирует созданную общую библиотеку в качестве модуля Python. Если False, поведение зависит от is_standalone.
• is_standalone – Если False (по умолчанию) загружает созданное расширение в процесс как обычную динамическую библиотеку. Если True, создает автономный исполняемый файл.

Возвращаемое значение:

Возвращает загруженное расширение PyTorch в качестве модуля Python.

If is_python_module is False and is_standalone is False:

Ничего не возвращает. (Общая библиотека загружается в процесс как побочный эффект.)

If is_standalone is True.

Возвращает путь к исполняемому файлу. (В Windows, TORCH_LIB_PATH добавляется в переменную среды PATH в качестве побочного эффекта.)

Тип возвращаемого значения:

Если is_python_module это True

Пример

>>> from torch.utils.cpp_extension import load
>>> module = load(
...     name='extension',
...     sources=['extension.cpp', 'extension_kernel.cu'],
...     extra_cflags=['-O2'],
...     verbose=True)

torch.utils.cpp_extension.load_inline(name, cpp_sources, cuda_sources=None, functions=None, extra_cflags=None, extra_cuda_cflags=None, extra_ldflags=None, extra_include_paths=None, build_directory=None, verbose=False, with_cuda=None, is_python_module=True, with_pytorch_error_handling=True, keep_intermediates=True) [source]

Загрузка расширения PyTorch на C++ в режиме реального времени (JIT) из исходных кодов в виде строк.

Эта функция ведет себя точно так же, как load(), но принимает исходные коды в виде строк вместо имен файлов. Эти строки записываются в файлы в каталоге сборки, после чего поведение load_inline() идентично load().

Примеры использования этой функции см. в тестах.

Исходные коды могут опустить две обязательные части обычного расширения C++ без встроенного кода: необходимые включаемые заголовки, а также код связывания (pybind11). Точнее, строки, переданные в cpp_sources, сначала конкатенируются в один файл .cpp. Затем к этому файлу добавляется #include <torch/extension.h>.

Кроме того, если аргумент functions указан, для каждой указанной функции будут автоматически сгенерированы связи. functions может быть списком имен функций или словарем, сопоставляющим имена функций с строками документации. Если задан список, имя каждой функции используется как её строка документации.

Исходные коды в cuda_sources конкатенируются в отдельный файл .cu и добавляются torch/types.h, cuda.h и cuda_runtime.h включаемые файлы. Файлы .cpp и .cu компилируются отдельно, но в конечном итоге объединяются в одну библиотеку. Обратите внимание, что для функций в cuda_sources связи не генерируются сами по себе. Чтобы связаться с ядром CUDA, необходимо создать функцию C++, которая ее вызывает, и либо объявить, либо определить эту функцию C++ в одном из cpp_sources (и включить ее имя в functions).

Описание аргументов, опущенных ниже, см. в load().

Параметры:

• cpp_sources – Строка или список строк, содержащих исходный код C++.
• cuda_sources – Строка или список строк, содержащих исходный код CUDA.
• functions – Список имен функций, для которых нужно сгенерировать связи функций. Если задан словарь, он должен сопоставлять имена функций со строками документации (в противном случае используются просто имена функций).
• with_cuda – Определяет, добавляются ли заголовки и библиотеки CUDA к сборке. Если установлено None (по умолчанию), это значение автоматически определяется на основе наличия cuda_sources. Установите его в True для принудительного включения заголовков и библиотек CUDA.
• with_pytorch_error_handling – Определяет, обрабатываются ли макросы ошибок и предупреждений pytorch pytorch, а не pybind. Для этого каждая функция foo вызывается через промежуточную функцию _safe_foo. Это перенаправление может вызвать проблемы в особых случаях cpp. Этот флаг следует установить в False при возникновении проблем с этим перенаправлением.

Пример

>>> from torch.utils.cpp_extension import load_inline
>>> source = """
at::Tensor sin_add(at::Tensor x, at::Tensor y) {
  return x.sin() + y.sin();
}
"""
>>> module = load_inline(name='inline_extension',
...                      cpp_sources=[source],
...                      functions=['sin_add'])

Примечание

По умолчанию бэкенд Ninja использует #CPUS + 2 потока для сборки расширения. Это может использовать слишком много ресурсов на некоторых системах. Количество потоков можно контролировать, установив переменную среды MAX_JOBS на неотрицательное число.

torch.utils.cpp_extension.include_paths(cuda=False) [source]

Получение путей к включаемым файлам, необходимых для сборки расширения C++ или CUDA.

Параметры:

cuda (bool) – Если True, включает пути к включаемым файлам CUDA.

Возвращаемое значение:

Список строк путей к включаемым файлам.

Тип возвращаемого значения:

List[str]

torch.utils.cpp_extension.get_compiler_abi_compatibility_and_version(compiler) [source]

Определение совместимости компилятора с PyTorch по ABI и его версии.

Параметры:

compiler (str) – Имя исполняемого файла компилятора для проверки (например, g++). Должен быть исполняемым в оболочке.

Возвращаемое значение:

Кортеж, содержащий булево значение, определяющее (вероятную) несовместимость компилятора с PyTorch по ABI, а также строку TorchVersion, содержащую версию компилятора, разделенную точками.

Тип возвращаемого значения:

Tuple[bool, TorchVersion]

torch.utils.cpp_extension.verify_ninja_availability() [source]

Вызывает исключение RuntimeError, если система не поддерживает систему сборки ninja, в противном случае ничего не делает.

torch.utils.cpp_extension.is_ninja_available() [source]

Возвращает True, если система поддерживает систему сборки ninja, False в противном случае.