torch.utils.cpp_extension

Создано: 07 мар. 2018 | Последнее обновление: 12 мая 2026

torch.utils.cpp_extension.CppExtension(name, sources, *args, **kwargs) [исходный код]

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

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

Все аргументы передаются конструктору setuptools.Extension. Полный список аргументов см. по адресу https://setuptools.pypa.io/en/latest/userguide/ext_modules.html#extension-api-reference

Предупреждение

API Python PyTorch (предоставляемый в libtorch_python) нельзя собирать с флагом py_limited_api=True. При передаче этого флага пользователь несет ответственность за то, чтобы его библиотека не использовала API из libtorch_python (в частности, привязки pytorch/python), а использовала только API из libtorch (объекты aten, операторы и диспетчер). Например, чтобы предоставить доступ к пользовательским операциям из Python, библиотека должна зарегистрировать их через диспетчер.

В отличие от setuptools CPython, который не задает -DPy_LIMITED_API как флаг компиляции, когда py_limited_api указан в качестве параметра команды “bdist_wheel” в setup, PyTorch это делает! Мы укажем -DPy_LIMITED_API=min_supported_cpython, чтобы обеспечить согласованность, безопасность и надежность, а также поощрять использование передовых практик. Чтобы выбрать другую версию, задайте min_supported_cpython в виде шестнадцатеричного кода нужной версии CPython.

Пример

>>> 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'],
...             extra_link_args=['-Wl,--no-as-needed', '-lm'])
...     ],
...     cmdclass={
...         'build_ext': BuildExtension
...     })

torch.utils.cpp_extension.CUDAExtension(name, sources, *args, **kwargs) [исходный код]

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

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

Предупреждение

Пример

>>> 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']},
...                 extra_link_args=['-Wl,--no-as-needed', '-lcuda'])
...     ],
...     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 >= указанной (например, 8.6+PTX генерирует PTX, который можно скомпилировать во время выполнения для любой GPU с CC >= 8.6). Это повышает совместимость двоичного файла с будущими устройствами. Однако использование более старого PTX для обеспечения совместимости с будущими устройствами посредством компиляции во время выполнения для новых версий CC может несколько снизить производительность на таких устройствах. Если вам известны точные версии CC целевых GPU, лучше всегда указывать их по отдельности. Например, если расширение должно работать на версиях 8.0 и 8.6, “8.0+PTX” функционально подойдет, поскольку включает PTX, который можно скомпилировать во время выполнения для 8.6, но вариант “8.0 8.6” будет лучше.

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

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

Пример использования:: #include <ATen/ATen.h> at::Tensor SigmoidAlphaBlendForwardCuda(….)
Вместо:: #include <torch/extension.h> torch::Tensor SigmoidAlphaBlendForwardCuda(…)

Текущая открытая задача об ошибке nvcc: pytorch/pytorch#69460 Полный пример обходного решения: facebookresearch/pytorch3d

Компоновка перемещаемого кода устройства:

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

Если у вас есть объекты rdc, перед этапом компоновки символов CPU необходимо выполнить дополнительный этап -dlink (компоновка кода устройства). Также -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.SyclExtension(name, sources, *args, **kwargs) [исходный код]

Создает setuptools.Extension для SYCL/C++.

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

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

Предупреждение

Пример

>>> from torch.utils.cpp_extension import BuildExtension, SyclExtension
>>> setup(
...     name='xpu_extension',
...     ext_modules=[
...     SyclExtension(
...                 name='xpu_extension',
...                 sources=['extension.cpp', 'extension_kernel.cpp'],
...                 extra_compile_args={'cxx': ['-g', '-std=c++20', '-fPIC']})
...     ],
...     cmdclass={
...         'build_ext': BuildExtension
...     })

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

TORCH_XPU_ARCH_LIST="pvc,xe-lpg" python build_my_extension.py

Примечание: для сборки SyclExtension требуется Ninja.

torch.utils.cpp_extension.BuildExtension(*args, **kwargs) [исходный код]

Пользовательское расширение сборки setuptools.

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

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

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_sycl_cflags=None, extra_ldflags=None, extra_include_paths=None, build_directory=None, verbose=False, with_cuda=None, with_sycl=None, is_python_module=True, is_standalone=False, keep_intermediates=True) [исходный код]

Загружает расширение C++ PyTorch «на лету» (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++. При этом каталог lib64 CUDA будет добавлен в качестве каталога библиотек и выполнится компоновка с cudart. Дополнительные флаги nvcc можно передать через extra_cuda_cflags, так же как extra_cflags для C++. Используются различные эвристики для поиска каталога установки CUDA, которые обычно работают корректно. Если это не так, надежнее всего задать переменную окружения CUDA_HOME.

Поддерживается смешанная компиляция с SYCL. Достаточно передать исходные файлы SYCL (.sycl) вместе с остальными исходными файлами. Такие файлы будут обнаружены и скомпилированы компилятором SYCL (например, Intel DPC++ Compiler), а не компилятором C++. Дополнительные флаги компилятора SYCL можно передать через extra_sycl_cflags, так же как extra_cflags для C++. Предполагается, что компилятор SYCL доступен через системную переменную окружения PATH.

Параметры:

name – Имя собираемого расширения. Оно ДОЛЖНО совпадать с именем модуля pybind11!
sources (str | list[str]) – Список относительных или абсолютных путей к исходным файлам C++.
extra_cflags – Необязательный список флагов компилятора, передаваемых сборке.
extra_cuda_cflags – Необязательный список флагов компилятора, передаваемых nvcc при сборке исходных файлов CUDA.
extra_sycl_cflags – Необязательный список флагов компилятора, передаваемых компилятору SYCL при сборке исходных файлов SYCL.
extra_ldflags – Необязательный список флагов компоновщика, передаваемых сборке.
extra_include_paths – Необязательный список каталогов включаемых файлов, передаваемых сборке.
build_directory – Необязательный путь, используемый как рабочий каталог сборки.
verbose – Если True, включает подробное журналирование этапов загрузки.
with_cuda (bool | None) – Определяет, будут ли в сборку добавлены заголовочные файлы и библиотеки CUDA. Если задано значение None (по умолчанию), оно определяется автоматически по наличию .cu или .cuh в sources. Установите значение True`, чтобы принудительно включить заголовочные файлы и библиотеки CUDA.
with_sycl (bool | None) – Определяет, будут ли в сборку добавлены заголовочные файлы и библиотеки SYCL. Если задано значение None (по умолчанию), оно определяется автоматически по наличию .sycl в sources. Установите значение True`, чтобы принудительно включить заголовочные файлы и библиотеки SYCL.
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, sycl_sources=None, functions=None, extra_cflags=None, extra_cuda_cflags=None, extra_sycl_cflags=None, extra_ldflags=None, extra_include_paths=None, build_directory=None, verbose=False, with_cuda=None, with_sycl=None, is_python_module=True, with_pytorch_error_handling=True, keep_intermediates=True, use_pch=False, no_implicit_headers=False) [исходный код]

Загружает расширение C++ PyTorch «на лету» (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).

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

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

Параметры:

cpp_sources – Строка или список строк с исходным кодом C++.
cuda_sources – Строка или список строк с исходным кодом CUDA.
sycl_sources – Строка или список строк с исходным кодом SYCL.
functions – Список имен функций, для которых нужно создать привязки. Если передан словарь, он должен сопоставлять имена функций со строками документации (в противном случае используются только имена функций).
with_cuda – Определяет, будут ли в сборку добавлены заголовочные файлы и библиотеки CUDA. Если задано значение None (по умолчанию), оно определяется автоматически в зависимости от того, передан ли аргумент cuda_sources. Установите значение True, чтобы принудительно включить заголовочные файлы и библиотеки CUDA.
with_sycl – Определяет, будут ли в сборку добавлены заголовочные файлы и библиотеки SYCL. Если задано значение None (по умолчанию), оно определяется автоматически в зависимости от того, передан ли аргумент sycl_sources. Установите значение True, чтобы принудительно включить заголовочные файлы и библиотеки SYCL.
with_pytorch_error_handling – Определяет, будут ли обработка ошибок и макросы предупреждений PyTorch выполняться самим PyTorch, а не pybind. Для этого каждая функция foo вызывается через промежуточную функцию _safe_foo. В редких случаях C++ такое перенаправление может вызывать проблемы. Если перенаправление приводит к проблемам, установите для этого флага значение False.
no_implicit_headers – Если True, пропускает автоматическое добавление заголовочных файлов, прежде всего строк #include <torch/extension.h> и #include <torch/types.h>. Используйте этот параметр, чтобы сократить время холодного запуска, если необходимые заголовочные файлы уже включены в исходный код. По умолчанию: 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'])

Примечание

Поскольку load_inline выполняет компиляцию исходного кода «на лету», убедитесь, что в среде выполнения установлены нужные цепочки инструментов. Например, для загрузки кода C++ убедитесь, что доступен компилятор C++. Для загрузки расширения CUDA дополнительно установите соответствующий инструментарий CUDA (nvcc и другие необходимые вашему коду зависимости). Цепочки инструментов для компиляции не входят в пакет torch и должны устанавливаться отдельно.

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

torch.utils.cpp_extension.include_paths(device_type='cpu', torch_include_dirs=True) [исходный код]

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

Параметры:: device_type (str) – По умолчанию «cpu».
Возвращает:: Список строк с путями к включаемым файлам.
Тип возвращаемого значения:: list[str]

torch.utils.cpp_extension.get_compiler_abi_compatibility_and_version(compiler) [исходный код]

Определяет, совместим ли указанный компилятор с PyTorch по ABI, и возвращает его версию.

Параметры:: compiler (str) – Имя исполняемого файла компилятора для проверки (например, g++). Должен запускаться в оболочке.
Возвращает:: Кортеж, содержащий логическое значение, указывающее, несовместим ли компилятор (предположительно) с PyTorch по ABI, и строку TorchVersion с версией компилятора, компоненты которой разделены точками.
Тип возвращаемого значения:: tuple[bool, TorchVersion]

torch.utils.cpp_extension.verify_ninja_availability() [исходный код]: Вызывает RuntimeError, если система сборки ninja недоступна в системе; в противном случае ничего не делает.

torch.utils.cpp_extension.is_ninja_available() [исходный код]

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

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

© 2026, PyTorch Contributors
PyTorch has a BSD-style license, as found in the LICENSE file.
https://docs.pytorch.org/docs/2.14/cpp_extension.html