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 (связывание устройств) шаг перед шагом связывания символов CPU. Также существует случай, когда -dlink используется без -rdc: когда расширение ссылается на статическую библиотеку, содержащую объекты, скомпилированные с помощью rdc, такие как библиотека [NVSHMEM](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++17) а также смешанной компиляции C++/CUDA (и поддержке файлов CUDA в целом).

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

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

Примечание

По умолчанию бэкенд 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 (Just-in-time).

Для загрузки расширения генерируется файл сборки 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, use_pch=False) [source]

Загружает расширение PyTorch на C++ в режиме JIT (Just-in-time) из исходных кодов в виде строк.

Эта функция ведет себя точно так же, как 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 использует #CPU + 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 исключение, если система не поддерживает RuntimeError build system ninja, в противном случае ничего не делает.

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

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