# torch.utils.cpp_extension

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

* * *

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

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

一种便捷方法，它使用最少的(但通常是足够的）参数创建`setuptools.Extension`来构建 C ++扩展。

所有参数都转发到`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`。

一种便捷方法，它使用最少的(但通常是足够的）参数创建`setuptools.Extension`，以构建 CUDA / C ++扩展。 这包括 CUDA 包含路径，库路径和运行时库。

All arguments are forwarded to the `setuptools.Extension` constructor.

Example

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

```

* * *

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

自定义`setuptools`构建扩展。

这个`setuptools.build_ext`子类负责传递所需的最低编译器标志(例如`-std=c++11`）以及混合的 C ++ / CUDA 编译(并通常支持 CUDA 文件）。

使用 [`BuildExtension`](#torch.utils.cpp_extension.BuildExtension "torch.utils.cpp_extension.BuildExtension") 时，可以提供`extra_compile_args`(而不是通常的列表）的字典，该字典从语言(`cxx`或`nvcc`）映射到其他编译器标志的列表 提供给编译器。 这样就可以在混合编译期间向 C ++和 CUDA 编译器提供不同的标志。

* * *

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

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

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

默认情况下，生成文件的发布目录和编译到的结果库为`&lt;tmp&gt;/torch_extensions/&lt;name&gt;`，其中`&lt;tmp&gt;`是当前平台上的临时文件夹，`&lt;name&gt;`是扩展名。 可以通过两种方式覆盖此位置。 首先，如果设置了`TORCH_EXTENSIONS_DIR`环境变量，它将替换`&lt;tmp&gt;/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`环境变量是最安全的选择。

参数

*   **名称** –要构建的扩展名。 该名称必须与 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`自动确定该值。 将其设置为 &lt;cite&gt;True`&lt;/cite&gt; 以强制包含 CUDA 标头和库。

*   **is_python_module** –如果为`True`(默认），则将生成的共享库作为 Python 模块导入。 如果为`False`，则将其作为纯动态库加载到进程中。

退货

如果`is_python_module`为`True`，则将加载的 PyTorch 扩展名作为 Python 模块返回。 如果`is_python_module`为`False`，则什么都不返回(作为副作用，共享库已加载到进程中）。

Example

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

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

此函数的行为与 [`load()`](#torch.utils.cpp_extension.load "torch.utils.cpp_extension.load") 完全相同，但是将其源作为字符串而不是文件名使用。 这些字符串存储到构建目录中的文件中，之后 [`load_inline()`](#torch.utils.cpp_extension.load_inline "torch.utils.cpp_extension.load_inline") 的行为与 [`load()`](#torch.utils.cpp_extension.load "torch.utils.cpp_extension.load") 相同。

有关使用此功能的良好示例，请参见[测试](https://github.com/pytorch/pytorch/blob/master/test/test_cpp_extensions.py)。

源可能会省略典型的非内联 C ++扩展的两个必需部分：必需的头文件以及(pybind11）绑定代码。 更准确地说，首先将传递给`cpp_sources`的字符串连接到单个`.cpp`文件中。 该文件然后以`#include &lt;torch/extension.h&gt;`开头。

此外，如果提供`functions`参数，则将为指定的每个函数自动生成绑定。 `functions`可以是函数名称列表，也可以是从函数名称到文档字符串的字典映射。 如果给出了列表，则将每个函数的名称用作其文档字符串。

`cuda_sources`中的源被连接到单独的`.cu`文件中，并以`torch/types.h`，`cuda.h`和`cuda_runtime.h`包括在内。 `.cpp`和`.cu`文件是分别编译的，但最终链接到一个库中。 注意，`cuda_sources`本身不为函数生成任何绑定。 要绑定到 CUDA 内核，您必须创建一个调用它的 C ++函数，并在`cpp_sources`之一中声明或定义此 C ++函数(并在`functions`中包括其名称）。

有关以下省略的自变量的说明，请参见 [`load()`](#torch.utils.cpp_extension.load "torch.utils.cpp_extension.load") 。

Parameters

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

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

*   **函数** –为其生成函数绑定的函数名称列表。 如果提供了字典，则应将函数名称映射到文档字符串(否则仅是函数名称）。

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

*   **with_pytorch_error_handling** –确定 pytorch 而不是 pybind 处理 pytorch 错误和警告宏。 为此，每个功能`foo`都通过中间`_safe_foo`功能调用。 这种重定向在 cpp 晦涩的情况下可能会引起问题。 当此重定向导致问题时，应将此标志设置为`False`。

Example

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

```

* * *

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

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

Parameters

**cuda** -如果&lt;cite&gt;为真&lt;/cite&gt;，则包含特定于 CUDA 的包含路径。

Returns

包含路径字符串的列表。

* * *

```
torch.utils.cpp_extension.check_compiler_abi_compatibility(compiler)¶
```

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

Parameters

**编译器** (_str_ )–要检查的编译器可执行文件名称(例如`g++`）。 必须在 Shell 进程中可执行。

Returns

如果编译器(可能）与 PyTorch 不兼容，则为 False，否则为 True。

* * *

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

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