torch.utils.cpp_extension

创建于: Mar 07, 2018 | 最后更新于: Feb 16, 2025

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

创建用于 C++ 的 setuptools.Extension。

方便的方法，它创建了一个 setuptools.Extension，并包含构建 C++ 扩展所需的最低限度（但通常足够）的参数。

所有参数都会被转发给 setuptools.Extension 构造函数。完整的参数列表可以在 https://setuptools.pypa.io/en/latest/userguide/ext_modules.html#extension-api-reference 找到。

警告

PyTorch 的 Python API（以 libtorch_python 的形式提供）不能使用 py_limited_api=True 标志进行构建。当传递此标志时，用户有责任在其库中不使用 libtorch_python 中的 API（特别是 pytorch/python 绑定），而只使用 libtorch 中的 API（aten 对象、运算符和调度器）。例如，要从 Python 中访问自定义运算符，库应该通过调度器注册这些运算符。

与 CPython 的 setuptools 不同，后者在 `setup` 的 `bdist_wheel` 命令中将 `py_limited_api` 指定为选项时，不会将 `-DPy_LIMITED_API` 定义为编译标志，而 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)

创建用于 CUDA/C++ 的 setuptools.Extension。

方便的方法，它创建了一个 setuptools.Extension，并包含构建 CUDA/C++ 扩展所需的最低限度（但通常足够）的参数。这包括 CUDA 包含路径、库路径和运行时库。

所有参数都会被转发给 setuptools.Extension 构造函数。完整的参数列表可以在 https://setuptools.pypa.io/en/latest/userguide/ext_modules.html#extension-api-reference 找到。

警告

PyTorch 的 Python API（以 libtorch_python 的形式提供）不能使用 py_limited_api=True 标志进行构建。当传递此标志时，用户有责任在其库中不使用 libtorch_python 中的 API（特别是 pytorch/python 绑定），而只使用 libtorch 中的 API（aten 对象、运算符和调度器）。例如，要从 Python 中访问自定义运算符，库应该通过调度器注册这些运算符。

与 CPython 的 setuptools 不同，后者在 `setup` 的 `bdist_wheel` 命令中将 `py_limited_api` 指定为选项时，不会将 `-DPy_LIMITED_API` 定义为编译标志，而 PyTorch 则会！我们将指定 `-DPy_LIMITED_API=min_supported_cpython` 以最大程度地强制执行一致性、安全性和健全性，以鼓励最佳实践。要针对不同的版本，请将 `min_supported_cpython` 设置为所需 CPython 版本的十六进制代码。

示例

>>> 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 回退到使用 nvcc 支持的最新 PTX 版本进行内核构建（有关 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 选项会导致扩展内核二进制文件包含指定 CC 的 PTX 指令。PTX 是一种中间表示，允许内核在运行时编译任何 CC >= 指定 CC（例如，8.6+PTX 生成的 PTX 可以运行时编译任何 CC >= 8.6 的 GPU）。这提高了二进制文件的前向兼容性。但是，依赖较旧的 PTX 通过运行时编译较新 CC 来提供前向兼容性，可能会略微降低这些较新 CC 的性能。如果您知道要针对的确切 CC 显卡，最好单独指定它们。例如，如果您希望扩展在 8.0 和 8.6 上运行，“8.0+PTX”在功能上可行，因为它包含了可以运行时编译 8.6 的 PTX，但“8.0 8.6”会更好。

请注意，虽然包含所有支持的架构是可能的，但包含的架构越多，构建过程就越慢，因为它将为每个架构构建一个单独的内核映像。

请注意，CUDA-11.5 nvcc 在 Windows 上解析 torch/extension.h 时会遇到内部编译器错误。为了解决这个问题，请将 Python 绑定逻辑移至纯 C++ 文件。

示例用法

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

而不是

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

nvcc 错误目前的开放问题：pytorch/pytorch#69460 完整的解决方法代码示例：facebookresearch/pytorch3d

可重定位设备代码链接

如果您想在编译单元（对象文件）之间引用设备符号，则需要使用 可重定位设备代码（-rdc=true 或 -dc）构建对象文件。此规则的一个例外是“动态并行”（嵌套内核启动），这种方式已不再常用。 可重定位设备代码 优化程度较低，因此只能用于需要它的对象文件。在设备代码编译步骤和 dlink 步骤中使用 -dlto（设备链接时优化）有助于减少 -rdc 的潜在性能下降。请注意，必须同时在两个步骤中使用它才有意义。

如果您有 rdc 对象，则需要在 CPU 符号链接步骤之前进行额外的 -dlink（设备链接）步骤。在没有 -rdc 的情况下使用 -dlink 的情况也存在：当一个扩展链接到一个包含 rdc 编译对象的静态库时，例如 [NVSHMEM 库](https://developer.nvidia.com/nvshmem)。

注意：构建具有 RDC 链接的 CUDA 扩展需要 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)

创建用于 SYCL/C++ 的 setuptools.Extension。

方便的方法，它创建了一个 setuptools.Extension，并包含构建 SYCL/C++ 扩展所需的最低限度（但通常足够）的参数。

所有参数都会被转发给 setuptools.Extension 构造函数。

警告

PyTorch 的 Python API（以 libtorch_python 的形式提供）不能使用 py_limited_api=True 标志进行构建。当传递此标志时，用户有责任在其库中不使用 libtorch_python 中的 API（特别是 pytorch/python 绑定），而只使用 libtorch 中的 API（aten 对象、运算符和调度器）。例如，要从 Python 中访问自定义运算符，库应该通过调度器注册这些运算符。

与 CPython 的 setuptools 不同，后者在 `setup` 的 `bdist_wheel` 命令中将 `py_limited_api` 指定为选项时，不会将 `-DPy_LIMITED_API` 定义为编译标志，而 PyTorch 则会！我们将指定 `-DPy_LIMITED_API=min_supported_cpython` 以最大程度地强制执行一致性、安全性和健全性，以鼓励最佳实践。要针对不同的版本，请将 `min_supported_cpython` 设置为所需 CPython 版本的十六进制代码。

示例

>>> 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++17）以及 C++/CUDA/SYCL 混合编译（以及对 CUDA/SYCL 文件的一般支持）。

在使用 BuildExtension 时，允许为 extra_compile_args（而不是通常的列表）提供一个字典，该字典将语言/编译器（预期的唯一值是 cxx、nvcc 或 sycl）映射到要传递给编译器的其他编译器标志列表。这使得在混合编译期间可以向 C++、CUDA 和 SYCL 编译器提供不同的标志。

use_ninja (bool)：如果 use_ninja 为 True（默认），则我们尝试使用 Ninja 后端进行构建。与标准的 setuptools.build_ext 相比，Ninja 大大加快了编译速度。如果 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)

即时 (JIT) 加载 PyTorch C++ 扩展。

要加载一个扩展，会生成一个 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。您可以通过 extra_cuda_cflags 向 nvcc 传递其他标志，就像使用 extra_cflags 处理 C++ 一样。使用各种启发式方法查找 CUDA 安装目录，这些方法通常都能正常工作。如果不行，设置 CUDA_HOME 环境变量是最稳妥的选择。

提供了对混合编译的 SYCL 支持。只需将 SYCL 源文件（.sycl）与其他源文件一起传递。这些文件将被检测并用 SYCL 编译器（如 Intel DPC++ Compiler）而不是 C++ 编译器进行编译。您可以通过 extra_sycl_cflags 向 SYCL 编译器传递其他标志，就像使用 extra_cflags 处理 C++ 一样。SYCL 编译器预计可以通过系统 PATH 环境变量找到。

参数:

• name – 要构建的扩展的名称。这必须与 pybind11 模块的名称相同！

• sources (str | list[str]) – C++ 源文件的相对或绝对路径列表。

• extra_cflags – 可选的编译器标志列表，将转发给构建过程。

• extra_cuda_cflags – 可选的编译器标志列表，在构建 CUDA 源文件时将转发给 nvcc。

• extra_sycl_cflags – 可选的编译器标志列表，在构建 SYCL 源文件时将转发给 SYCL 编译器。

• extra_ldflags – 可选的链接器标志列表，将转发给构建过程。

• extra_include_paths – 可选的包含目录列表，将转发给构建过程。

• build_directory – 可选的用作构建工作区的路径。

• verbose – 如果为 True，则启用加载步骤的详细日志记录。

• with_cuda (bool | None) – 确定是否将 CUDA 头文件和库添加到构建中。如果设置为 None（默认），则根据 sources 中是否存在 .cu 或 .cuh 自动确定此值。将其设置为 True` 以强制包含 CUDA 头文件和库。

• with_sycl (bool | None) – 确定是否将 SYCL 头文件和库添加到构建中。如果设置为 None（默认），则根据 sources 中是否存在 .sycl 自动确定此值。将其设置为 True` 以强制包含 SYCL 头文件和库。

• is_python_module – 如果为 True（默认），则将生成的共享库导入为 Python 模块。如果为 False，行为取决于 is_standalone。

• is_standalone – 如果为 False（默认），则将构建的扩展作为普通动态库加载到进程中。如果为 True，则构建一个独立的可执行文件。

返回:

返回加载的 PyTorch 扩展作为一个 Python 模块。

如果 is_python_module 为 False 且 is_standalone 为 False

不返回任何内容。（共享库作为副作用加载到进程中。）

如果 is_standalone 为 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)

从字符串源即时 (JIT) 加载 PyTorch C++ 扩展。

此函数与 load() 的行为完全相同，但它接受字符串而不是文件名作为其源。这些字符串会被存储到构建目录中的文件中，之后 load_inline() 的行为与 load() 相同。

有关使用此函数的优秀示例，请参见 tests。

源文件可以省略典型的非内联 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++ 函数，并在 cpp_sources 之一中声明或定义此 C++ 函数（并将其名称包含在 functions 中）。

sycl_sources 中的源文件将被连接成一个单独的 .sycl 文件，并在前面加上 torch/types.h、sycl/sycl.hpp 包含。 .cpp 和 .sycl 文件分开编译，但最终链接成一个库。请注意，本身不会为 sycl_sources 中的函数生成绑定。要绑定到 SYCL 内核，您必须创建一个调用它的 C++ 函数，并在 cpp_sources 之一中声明或定义此 C++ 函数（并将其名称包含在 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 还是 pybind 来处理 PyTorch 的错误和警告宏。为此，每个函数 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++）。必须可在 shell 进程中执行。

返回:

一个元组，其中包含一个布尔值，指示编译器是否（可能）与 PyTorch ABI 不兼容，后跟一个 TorchVersion 字符串，其中包含用点分隔的编译器版本。

返回类型:

tuple[bool, TorchVersion]

torch.utils.cpp_extension.verify_ninja_availability()

如果系统上 ninja 构建系统不可用，则引发 RuntimeError，否则不执行任何操作。

torch.utils.cpp_extension.is_ninja_available()

如果系统上 ninja 构建系统可用，则返回 True，否则返回 False。

返回类型:

布尔值