torch.utils.cpp_extension ¬

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

为 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 中将 py_limited_api 选项指定为“bdist_wheel”命令的选项时，setuptools 不会将 -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)[source][source]¶

为 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 中将 py_limited_api 选项指定为“bdist_wheel”命令的选项时，setuptools 不会将 -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 的性能。如果您知道要针对的 GPU 的确切 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 错误问题：https://github.com/pytorch/pytorch/issues/69460 完整的解决方案代码示例：https://github.com/facebookresearch/pytorch3d/commit/cb170ac024a949f1f9614ffe6af1c38d972f7d48

可移动设备代码链接：

如果您想在编译单元（对象文件）之间引用设备符号（跨对象文件），则需要使用可移动设备代码（-rdc=true 或 -dc）。此规则的例外是“动态并行性”（嵌套内核启动），现在使用得越来越少。可移动设备代码的优化程度较低，因此仅需要在需要它的对象文件中使用。在设备代码编译步骤和 dlink 步骤中使用-dlto（设备链接时间优化）有助于减少-rdc 可能导致的性能下降。请注意，它需要在两个步骤中使用才有用。

如果您有 rdc 对象，您需要在 CPU 符号链接步骤之前进行额外的-dlink（设备链接）步骤。还有一种情况是使用-dlink 而不使用-rdc：当扩展库链接到一个包含 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)[source][source]¶

创建一个用于 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 中将 py_limited_api 选项指定为“bdist_wheel”命令的选项时，setuptools 不会将 -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)[source][source]

自定义 setuptools 构建扩展。

此 setuptools.build_ext 子类负责传递所需的最低编译器标志（例如 -std=c++17 ）以及混合 C++/CUDA/SYCL 编译（以及一般对 CUDA/SYCL 文件的支持）。

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

use_ninja （布尔值）：如果 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)[source][source]¶

即时加载 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++编译器）编译，而不是 C++编译器。您可以通过 extra_sycl_cflags 向 SYCL 编译器传递额外的标志，就像向 extra_cflags 传递 C++标志一样。期望通过系统 PATH 环境变量找到 SYCL 编译器。

参数:

名称 – 要构建的扩展的名称。这必须与 pybind11 模块的名称相同！
源文件（Union[str, list[str]]） – 相对或绝对路径的 C++源文件列表。
extra_cflags – 向构建过程传递的编译器标志的可选列表。
extra_cuda_cflags – 构建 CUDA 源代码时传递给 nvcc 的编译器标志的可选列表。
extra_sycl_cflags – 构建 SYCL 源代码时传递给 SYCL 编译器的编译器标志的可选列表。
extra_ldflags – 向构建过程传递的链接器标志的可选列表。
额外包含路径 – 可选的传递给构建的包含目录列表。
构建目录 – 可选的用作构建工作区的路径。
详细模式 – 如果 True ，则开启加载步骤的详细日志记录。
with_cuda（可选[bool]）– 决定是否将 CUDA 头文件和库添加到构建中。如果设置为 None （默认），此值将根据 .cu 或 .cuh 在 sources 中的存在自动确定。将其设置为 True 以强制包含 CUDA 头文件和库。
with_sycl (可选[布尔值]) – 决定是否将 SYCL 头文件和库添加到构建中。如果设置为 None （默认），此值将根据 .sycl 在 sources 中的存在自动确定。将其设置为 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)[source][source]¶

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

此函数的行为与 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++函数，并在其中一个 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 错误和警告宏是否由 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'])

注意

由于 load_inline 会即时编译源代码，请确保在运行时已安装正确的工具链。例如，当加载 C++时，请确保有 C++编译器可用。如果您正在加载 CUDA 扩展，则需要额外安装相应的 CUDA 工具包（nvcc 和您的代码所需的任何其他依赖项）。编译工具链在安装 torch 时不会包括在内，必须额外安装。

在编译过程中，默认情况下，Ninja 后端使用#CPUS + 2 个工作线程来构建扩展。这可能在某些系统上消耗过多资源。可以通过将 MAX_JOBS 环境变量设置为非负数来控制工作线程的数量。

torch.utils.cpp_extension.include_paths(device_type='cpu')[source][source]

获取构建 C++ 或 CUDA 或 SYCL 扩展所需的包含路径。

参数:: device_type（字符串）- 默认为“cpu”。
返回:: 包含路径字符串列表。
返回类型:: list[str]

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

判断给定的编译器是否与 PyTorch 兼容其版本。

参数:: compiler (str) – 要检查的编译器可执行文件名（例如 g++ ）。必须在 shell 进程中可执行。
返回:: 一个元组，包含一个布尔值，表示编译器（可能）与 PyTorch 不兼容，后跟一个包含编译器版本（用点分隔）的 TorchVersion 字符串。
返回类型:: tuple[bool, torch.torch_version.TorchVersion]

torch.utils.cpp_extension.verify_ninja_availability()[source][source]¶: 如果系统上不可用 ninja 构建系统，则抛出 RuntimeError ，否则不执行任何操作。

torch.utils.cpp_extension.is_ninja_available()[source][source]¶: 返回 True 如果系统上安装了忍者构建系统， False 否则。

torch.utils.cpp_extension ¬

文档

教程

资源