torch.utils.cpp_extension¶

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

创建一个 setuptools.Extension 用于 C++。

一种便捷的方法，使用构建 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)[source]¶

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

一种便捷的方法，使用构建 CUDA/C++ 扩展所需的最基本（但通常足够）的参数来创建一个 setuptools.Extension。这包括 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回退到使用你的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 的版本（例如，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”会更好。

请注意，虽然可以包含所有受支持的架构，但包含的架构越多，构建过程就会越慢，因为它将为每个架构分别构建一个内核镜像。

请注意，在 Windows 上使用 CUDA-11.5 的 nvcc 解析 torch/extension.h 时会发生内部编译器错误。要解决此问题，请将 Python 绑定逻辑移至纯 C++ 文件中。

Example use:

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

Instead of:

>>> #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)[source]¶

一个自定义的 setuptools 构建扩展 .

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

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

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_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）编译方式。

要加载一个扩展，将生成一个 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传递额外标志，就像为C++使用 extra_cflags 一样。使用了多种用于查找CUDA安装目录的启发式方法，通常都能正常工作。如果不成功，设置 CUDA_HOME 环境变量是最安全的选择。

Parameters

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

Returns

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

If is_python_module is False and is_standalone is False:: 不返回任何内容。（共享库会作为副作用加载到进程中。）
If is_standalone is True.: 返回可执行文件的路径。（在 Windows 系统上，TORCH_LIB_PATH 会作为副作用被添加到 PATH 环境变量中。）

Return type

如果 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)[source]¶

从字符串源加载一个 PyTorch C++ 扩展的即时编译（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++函数，并在其中一个cpp_sources中声明或定义此C++函数（并将其名称包含在functions中）。

请参阅 load() 以获取对下文省略参数的描述。

Parameters

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 后端使用 #CPUS + 2 个工作人员来构建扩展。这在某些系统上可能会占用太多资源。可以通过将 MAX_JOBS 环境变量设置为非负数来控制工作人员的数量。

torch.utils.cpp_extension.include_paths(cuda=False)[source]¶

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

Parameters: cuda – 如果 True，包含 CUDA 特定的包含路径。
Returns: 一个包含路径字符串的列表。

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

确定给定的编译器是否与 PyTorch 及其版本具有 ABI 兼容性。

Parameters: 编译器 (str) – 要检查的编译器可执行文件名称（例如 g++）。必须可以在 shell 进程中执行。
Returns: 一个元组，其中包含一个布尔值，用于定义编译器是否（可能）与 PyTorch 的 ABI 不兼容，后跟一个 TorchVersion 字符串，该字符串以点分隔的形式包含编译器版本号。

torch.utils.cpp_extension.verify_ninja_availability()[source]¶: 如果系统上未安装 RuntimeError ninja 构建系统，则引发否则不执行任何操作。

torch.utils.cpp_extension.is_ninja_available()[source]¶: 如果系统上可用 ninja 构建系统，则返回 True，否则返回 False。

torch.utils.cpp_extension¶

文档

教程

资源