# torch.utils.cpp_extension

> 原文：[`pytorch.org/docs/stable/cpp_extension.html`](https://pytorch.org/docs/stable/cpp_extension.html)

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

为 C++创建一个`setuptools.Extension`。

使用最少的参数创建一个`setuptools.Extension`的便捷方法，用于构建 C++扩展。

所有参数都会传递给`setuptools.Extension`构造函数。

示例

```py
>>> 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
...     }) 
```

```py
torch.utils.cpp_extension.CUDAExtension(name, sources, *args, **kwargs)
```

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

使用最少的参数创建一个`setuptools.Extension`的便捷方法，用于构建 CUDA/C++扩展。这包括 CUDA 包含路径、库路径和运行时库。

所有参数都会传递给`setuptools.Extension`构造函数。

示例

```py
>>> 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
...     }) 
```

计算能力：

默认情况下，扩展将被编译以在构建过程中可见的所有卡的所有 archs 上运行，再加上 PTX。如果将来安装了新的卡，可能需要重新编译扩展。如果可见卡的计算能力（CC）比您的 nvcc 可以完全编译的最新版本要新，Pytorch 将使 nvcc 退回到使用您的 nvcc 支持的最新版本的 PTX 构建内核（有关 PTX 的详细信息，请参见下文）。

您可以使用 TORCH_CUDA_ARCH_LIST 来覆盖默认行为，明确指定扩展要支持的 CCs：

`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(s)，最好是分别指定它们。例如，如果您希望您的扩展在 8.0 和 8.6 上运行，“8.0+PTX”在功能上可以工作，因为它包含了可以为 8.6 运行时编译的 PTX，但“8.0 8.6”会更好。

请注意，虽然可以包含所有支持的 archs，但是包含的 archs 越多，构建过程就会变得越慢，因为它将为每个 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 bug 的未解决问题：[`github.com/pytorch/pytorch/issues/69460`](https://github.com/pytorch/pytorch/issues/69460) 完整的解决方案代码示例：[`github.com/facebookresearch/pytorch3d/commit/cb170ac024a949f1f9614ffe6af1c38d972f7d48`](https://github.com/facebookresearch/pytorch3d/commit/cb170ac024a949f1f9614ffe6af1c38d972f7d48)

可重定位设备代码链接：

如果要在编译单元之间引用设备符号（跨目标文件），则需要使用可重定位设备代码（-rdc=true 或-dc）构建目标文件。这个规则的一个例外是“动态并行性”（嵌套内核启动），这在现在已经不太常用了。可重定位设备代码不太优化，因此只有在需要的目标文件上使用它时才需要使用。在设备代码编译步骤和 dlink 步骤中使用-dlto（设备链接时间优化）有助于减少-rdc 的潜在性能降低。请注意，它需要在两个步骤中使用才能发挥作用。

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

注意：构建 CUDA 扩展需要使用 Ninja 进行 RDC 链接。

示例

```py
>>> 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']}) 
```

```py
torch.utils.cpp_extension.BuildExtension(*args, **kwargs)
```

一个自定义的`setuptools`构建扩展。

这个`setuptools.build_ext`子类负责传递最低要求的编译器标志（例如`-std=c++17`）以及混合 C++/CUDA 编译（以及对 CUDA 文件的支持）。

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

`use_ninja`（布尔值）：如果`use_ninja`为`True`（默认值），则我们尝试使用 Ninja 后端进行构建。与标准的`setuptools.build_ext`相比，Ninja 大大加快了编译速度。如果 Ninja 不可用，则回退到标准的 distutils 后端。

注意

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

```py
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)
```

即时加载 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`环境变量是最安全的选择。

参数

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

+   **sources**（[*Union*](https://docs.python.org/3/library/typing.html#typing.Union "(在 Python v3.12 中)")*[*[*str*](https://docs.python.org/3/library/stdtypes.html#str "(在 Python v3.12 中)")*,* [*List*](https://docs.python.org/3/library/typing.html#typing.List "(在 Python v3.12 中)")*[*[*str*](https://docs.python.org/3/library/stdtypes.html#str "(在 Python v3.12 中)")*]**]*) – 一个包含 C++源文件的相对或绝对路径列表。

+   **extra_cflags** – 转发到构建的编译器标志的可选列表。

+   **extra_cuda_cflags** – 转发到构建 CUDA 源代码时传递给 nvcc 的编译器标志的可选列表。

+   **extra_ldflags** – 转发到构建的链接器标志的可选列表。

+   **extra_include_paths** – 转发到构建的包含目录的可选列表。

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

+   **verbose** – 如果为`True`，则打开加载步骤的详细日志记录。

+   **with_cuda**（[*可选*](https://docs.python.org/3/library/typing.html#typing.Optional "(在 Python v3.12 中)")*[*[*bool*](https://docs.python.org/3/library/functions.html#bool "(在 Python v3.12 中)")*]*) – 确定是否将 CUDA 头文件和库添加到构建中。如果设置为`None`（默认值），则根据`sources`中是否存在`.cu`或`.cuh`自动确定此值。将其设置为 True`以强制包含 CUDA 头文件和库。

+   **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`

示例

```py
>>> from torch.utils.cpp_extension import load
>>> module = load(
...     name='extension',
...     sources=['extension.cpp', 'extension_kernel.cu'],
...     extra_cflags=['-O2'],
...     verbose=True) 
```

```py
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, use_pch=False)
```

从字符串源加载 PyTorch C++扩展程序。

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

查看[测试](https://github.com/pytorch/pytorch/blob/master/test/test_cpp_extensions_jit.py)以获取使用此函数的良好示例。

源代码可能省略典型非内联 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()`。

参数

+   **cpp_sources** - 包含 C++源代码的字符串或字符串列表。

+   **cuda_sources** - 包含 CUDA 源代码的字符串或字符串列表。

+   **函数** - 一个要生成函数绑定的函数名称列表。如果给定一个字典，它应该将函数名称映射到文档字符串（否则只是函数名称）。

+   **with_cuda** - 确定是否将 CUDA 头文件和库添加到构建中。如果设置为`None`（默认值），则此值将根据是否提供了`cuda_sources`自动确定。将其设置为`True`以强制包含 CUDA 头文件和库。

+   **with_pytorch_error_handling** - 确定是否由 pytorch 处理 pytorch 错误和警告宏，而不是由 pybind 处理。为此，每个函数`foo`都通过一个中间函数`_safe_foo`调用。这种重定向可能会在 cpp 的一些复杂情况下引起问题。当此重定向引起问题时，应将此标志设置为`False`。

示例

```py
>>> 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 环境变量设置为非负数来控制工作进程的数量。

```py
torch.utils.cpp_extension.include_paths(cuda=False)
```

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

参数

**cuda**（[*bool*](https://docs.python.org/3/library/functions.html#bool "(在 Python v3.12 中)")） - 如果为 True，则包含特定于 CUDA 的包含路径。

返回

包含包含路径字符串的列表。

返回类型

[*列表*](https://docs.python.org/3/library/typing.html#typing.List "(在 Python v3.12 中)")[[str](https://docs.python.org/3/library/stdtypes.html#str "(在 Python v3.12 中)")]

```py
torch.utils.cpp_extension.get_compiler_abi_compatibility_and_version(compiler)
```

确定给定编译器是否与 PyTorch ABI 兼容以及其版本。

参数

**编译器**（[*str*](https://docs.python.org/3/library/stdtypes.html#str "(在 Python v3.12 中)")） - 要检查的编译器可执行文件名称（例如`g++`）。必须在 shell 进程中可执行。

返回

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

返回类型

[*元组*](https://docs.python.org/3/library/typing.html#typing.Tuple "(在 Python v3.12 中)")[[bool](https://docs.python.org/3/library/functions.html#bool "(在 Python v3.12 中)"), *TorchVersion*]

```py
torch.utils.cpp_extension.verify_ninja_availability()
```

如果系统上没有[ninja](https://ninja-build.org/)构建系统，则引发`RuntimeError`，否则不执行任何操作。

```py
torch.utils.cpp_extension.is_ninja_available()
```

如果系统上可用[ninja](https://ninja-build.org/)构建系统，则返回`True`，否则返回`False`。