torch.utils.cpp_extension¶

torch.utils.cpp_extension.CppExtension(name， sources， *args， **kwargs）[来源]¶

为 C++.setuptools.Extension

Convenience 方法，该方法使用构建 C++ 扩展所需的最小（但通常足够）参数。setuptools.Extension

所有参数都转发给构造函数。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）[来源]¶

为 CUDA/C++ 创建一个。setuptools.Extension

Convenience 方法，该方法使用构建 CUDA/C++ 所需的最小（但通常足够）参数外延。这包括 CUDA 包含路径、库路径和运行时库。setuptools.Extension

所有参数都转发给构造函数。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 回退到使用 NVCC 所做的最新版本的 PTX 构建内核支持（有关 PTX 的详细信息，请参阅下文）。

您可以使用 TORCH_CUDA_ARCH_LIST 来显式指定哪个您希望扩展支持的抄送：

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 >= 指定的 CC（例如，8.6+PTX 生成的 PTX 可以为任何具有 CC >= 8.6）的这可以提高二进制文件的向前兼容性。但是，依靠较旧的 PTX 来通过运行时编译为较新的 CC 提供前向兼容可能会适度降低那些较新的 CC。如果您知道要定位的 GPU 的确切 CC，那么您总是更好 off 单独指定它们。例如，如果您希望扩展在 8.0 和 8.6 上运行， “8.0+PTX” 在功能上可以工作，因为它包含可以针对 8.6 进行运行时编译的 PTX，但是 “8.0 8.6” 会更好。

请注意，虽然可以包含所有支持的 arch，但包含的 arch 越多，构建过程会变慢，因为它会为每个 Arch 构建一个单独的内核镜像。

请注意，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

torch.utils.cpp_extension.BuildExtension(*args， **kwargs）[来源]¶

自定义构建扩展。setuptools

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

使用时，允许提供字典 for（而不是通常的列表），映射自 languages （或）添加到其他编译器标志列表中，以 supply 提供给编译器。这使得向混合编译期间的 C++ 和 CUDA 编译器。extra_compile_argscxxnvcc

use_ninja（bool）：如果是（default），那么我们尝试使用 Ninja 后端进行构建。Ninja 大大加快了速度编译与标准 . 如果 Ninja 不可用，则回退到标准 distutils 后端。use_ninjaTruesetuptools.build_ext

注意

默认情况下，Ninja 后端使用 #CPUS + 2 个 worker 来构建外延。这可能会占用某些系统上的太多资源。一可以通过设置 MAX_JOBS 环境来控制 worker 的数量 variable 设置为非负数。

torch.utils.cpp_extension.load(name， sources， extra_cflags=无， extra_cuda_cflags=无， extra_ldflags=无， extra_include_paths=无、build_directory=无、详细=False、with_cuda=无、is_python_module=True、is_standalone=False， keep_intermediates=True）[来源]¶

实时加载 PyTorch C++ 扩展 JIT。

要加载扩展，会发出一个 Ninja 构建文件，该文件用于将给定的源编译到动态库中。这个库是随后作为模块加载到当前的 Python 进程中，并且从此函数返回，可供使用。

默认情况下，构建文件将发送到的目录和编译为的结果库为，其中是当前平台上的临时文件夹和扩展的名称。可以通过两种方式覆盖此位置。首先，如果设置了环境变量，则替换，并且所有扩展都将被编译添加到此目录的子文件夹中。其次，如果提供了此函数的参数，它将覆盖整个路径，即该库将直接编译到该文件夹中。<tmp>/torch_extensions/<name><tmp><name>TORCH_EXTENSIONS_DIR<tmp>/torch_extensionsbuild_directory

为了编译源代码，使用默认的系统编译器（），这可以通过设置 environment 变量来覆盖。要传递其他参数添加到编译过程，或者可以提供。例如，要编译扩展通过优化，将 .您还可以使用 to 传递进一步的 include 目录。c++CXXextra_cflagsextra_ldflagsextra_cflags=['-O3']extra_cflags

提供混合编译的 CUDA 支持。只需传递 CUDA 源文件（或）以及其他源。此类文件将是使用 NVCC 而不是 C++ 编译器进行检测和编译。这包括将 CUDA lib64 目录作为库目录传递，并链接 .你可以通过将额外的标志传递给 nvcc，就像 C++ 一样。各种使用启发式方法来查找 CUDA 安装目录，通常工作正常。如果不是，则设置环境变量为最安全的选择。.cu.cuhcudartextra_cuda_cflagsextra_cflagsCUDA_HOME

参数

name – 要构建的扩展的名称。这必须与 Pybind11 模块的名称！
sources （源） – C++ 源文件的相对或绝对路径列表。
extra_cflags – 要转发到构建的编译器标志的可选列表。
extra_cuda_cflags – 要转发到 NVCC 的编译器标志的可选列表在构建 CUDA 源时。
extra_ldflags – 要转发到构建的链接器标志的可选列表。
extra_include_paths – 要转发的包含目录的可选列表添加到构建中。
build_directory – 用作构建工作区的可选路径。
verbose – 如果，则打开加载步骤的详细日志记录。True
with_cuda – 确定是否将 CUDA 头文件和库添加到构建。如果设置为（default），则此值为根据或的存在自动确定。将其设置为 True' 以强制使用 CUDA 标头和库。None.cu.cuhsources
is_python_module – 如果（默认），则导入生成的共享 library 作为 Python 模块。如果，则行为取决于。TrueFalseis_standalone
is_standalone – 如果（默认）加载构建的扩展作为普通动态库进入流程中。如果，则构建一个 standalone 可执行文件。FalseTrue

返回

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

如果 is 和 is ：is_python_moduleFalseis_standaloneFalse: 不返回任何内容。（共享库作为副作用。
如果是 .is_standaloneTrue: 返回可执行文件的路径。（在 Windows 上，TORCH_LIB_PATH 添加到 PATH 环境变量中作为副作用。

返回类型

如果是is_python_moduleTrue

例

>>> 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=无， functions=无， extra_cflags=无， extra_cuda_cflags=无、extra_ldflags=无、extra_include_paths=无、build_directory=无、verbose=False、with_cuda=无，is_python_module=True，with_pytorch_error_handling=True，keep_intermediates=True）[来源]¶

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

此函数的行为与，但将其源视为 strings 而不是 filenames。这些字符串存储在 build 目录下，之后的行为为等同于 .

请参阅测试使用此函数的良好示例。

源可以省略典型非内联 C++ 扩展的两个必需部分：必要的标头包括（PyBind11）绑定代码。更多准确地说，传递给的字符串首先被连接成一个单个文件。然后，此文件前面加上。cpp_sources.cpp#include <torch/extension.h>

此外，如果提供了参数，则 bindings 将为为每个指定的函数自动生成。能可以是函数名称列表，也可以是 function 的字典映射 names 设置为 docstrings。如果给出了列表，则使用每个函数的名称作为其文档字符串。functionsfunctions

中的源被连接到一个单独的文件中，并在前面加上、和 includes。编译 and 文件单独，但最终链接到单个库中。请注意，没有本身为函数生成绑定。绑定添加到 CUDA 内核中，您必须创建一个调用它的 C++ 函数，并且 C++在（和将其名称包含在中）。cuda_sources.cutorch/types.hcuda.hcuda_runtime.h.cpp.cucuda_sourcescpp_sourcesfunctions

有关下面省略的参数的描述，请参见。

参数

cpp_sources – 包含源代码C++字符串或字符串列表。
cuda_sources – 包含 CUDA 源代码的字符串或字符串列表。
functions – 要为其生成函数的函数名称列表绑定。如果给出了字典，它应该将函数名称映射到文档字符串（否则只是函数名称）。
with_cuda – 确定是否将 CUDA 头文件和库添加到构建。如果设置为（default），则此值为根据 IS 是否自动确定提供。将其设置为以强制使用 CUDA 标头和库。Nonecuda_sourcesTrue
with_pytorch_error_handling – 确定 pytorch 错误和 warning 宏由 PyTorch 而不是 Pybind 处理。待办事项这样，每个函数都是通过一个中间函数调用的。此重定向可能会导致问题 CPP 的。此标志应设置为此重定向导致问题。foo_safe_fooFalse

例

>>> 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 个 worker 来构建外延。这可能会占用某些系统上的太多资源。一可以通过设置 MAX_JOBS 环境来控制 worker 的数量 variable 设置为非负数。

torch.utils.cpp_extension.include_paths(cuda=False）[来源]¶

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

参数: cuda — 如果为 True，则包含特定于 CUDA 的包含路径。
返回: 包含路径字符串的列表。

torch.utils.cpp_extension.check_compiler_abi_compatibility(编译器）[来源]¶

验证给定编译器是否与 PyTorch ABI 兼容。

参数: compiler （str） - 要检查的编译器可执行文件名称（例如）。必须在 shell 进程中可执行。g++
返回: 如果编译器（可能）与 PyTorch 不兼容，则为 False，否则为 True。

torch.utils.cpp_extension.verify_ninja_availability()[来源]¶: 如果 ninja 构建系统不是，则引发 available 的 S 系统，则不执行其他任何操作。RuntimeError

torch.utils.cpp_extension.is_ninja_available()[来源]¶: 如果 ninja 构建系统为 available 的 URL，否则。TrueFalse

torch.utils.cpp_extension¶

文档

教程

资源