torch.utils.cpp_extension

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'],
...             extra_link_args=['-Wl,--no-as-needed', '-lm'])
...     ],
...     cmdclass={
...         'build_ext': BuildExtension
...     })

torch.utils.cpp_extension.CUDAExtension(name, sources, *args, **kwargs)[源代码]

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

这是一种方便的方法，用于创建一个setuptools.Extension对象，其中包含了构建CUDA/C++扩展所需的基本参数（这些参数通常是足够的）。这包括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']},
...                 extra_link_args=['-Wl,--no-as-needed', '-lcuda'])
...     ],
...     cmdclass={
...         'build_ext': BuildExtension
...     })

计算能力：

默认情况下，该扩展程序将在构建过程中可见的所有卡架构上编译运行，并且还包括 PTX。如果将来安装了新卡，则可能需要重新编译扩展程序。如果可见的卡具有比你的 nvcc 可以构建完整编译二进制文件更高的计算能力（CC），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选项会使扩展内核二进制文件包含指定计算能力（Compute Capability，简称CC）的PTX指令。PTX是一种中间表示形式，允许内核在运行时编译为任何大于或等于指定CC的设备（例如，8.6+PTX生成可以在所有具有CC >= 8.6 的GPU上运行时编译的PTX）。这提高了二进制文件的向前兼容性。然而，依赖于旧版PTX在新版本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++文件中。

示例用法：: #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（Device Link Time Optimization），可以减少-rdc带来的潜在性能下降。需要注意的是，在这两个步骤中都需要使用-dlto才能有效。

如果你有rdc对象，需要在CPU符号链接步骤之前执行一个额外的-dlink（设备链接）步骤。还有一种情况是在没有-rdc的情况下使用-dlink：当扩展程序链接到包含类似[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.BuildExtension(*args, **kwargs)[源代码]

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

这个 setuptools.build_ext 子类负责传递所需的最小编译标志（例如 -std=c++17），并处理混合 C++/CUDA 编译，同时提供对 CUDA 文件的一般支持。

当使用BuildExtension时，可以为extra_compile_args提供一个字典（而不是通常的列表），该字典将语言(cxx或nvcc)映射到一系列额外的编译标志。这样可以在混合编译期间向C++和CUDA编译器传递不同的编译选项。

use_ninja (bool): 如果 use_ninja 为 True（默认值），则尝试使用 Ninja 后端进行构建。Ninja 能显著加快编译速度，相比标准的 setuptools.build_ext 更快。如果没有安装 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)[源代码]

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

要加载一个扩展，会生成一个Ninja构建文件，用于将给定的源代码编译为动态库。然后，这个库会被加载到当前的Python进程作为一个模块，并从该函数中返回，准备好被使用。

默认情况下，构建文件输出的目录和生成的库所在的目录是 <tmp>/torch_extensions/<name>，其中 <tmp> 是当前平台上的临时文件夹，<name> 是扩展名。此位置可以通过两种方式更改。首先，如果设置了环境变量 TORCH_EXTENSIONS_DIR，则它会替换默认路径中的 <tmp>/torch_extensions，所有扩展将编译到该目录的子文件夹中。其次，如果向此函数提供了参数 build_directory，则它将覆盖整个路径，库将直接编译到指定的文件夹。

要编译源代码，默认情况下会使用系统编译器（c++），但可以通过设置环境变量 CXX 来覆盖默认编译器。为了向编译过程添加额外参数，可以提供 extra_cflags 或 extra_ldflags。例如，要使用优化选项（如 -O3）来编译扩展模块，可以在 extra_cflags 中传入相应的值。此外，你还可以通过 extra_cflags 传递额外的包含目录。

提供了混合编译的CUDA支持。只需将CUDA源文件（如.cu或.cuh）与其他源文件一起传递即可。这些文件会被检测并使用nvcc进行编译，而不是C++编译器。此外，会将CUDA lib64目录作为库目录传递，并链接cudart。你可以通过extra_cuda_cflags向nvcc传递额外的标志，就像使用extra_cflags为C++一样。系统会尝试各种方法来查找CUDA安装目录，通常效果良好。如果找不到，则设置CUDA_HOME环境变量是最安全的选择。

参数

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

返回值

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

如果 is_python_module 和 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, 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)[源代码]

从字符串源即时（JIT）加载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 中包含其名称）。

参见load()以了解以下省略的参数描述。

参数

cpp_sources - 包含 C++ 源代码的字符串或字符串列表。
cuda_sources - 包含 CUDA 源代码的字符串或字符串列表。
functions - 一个包含函数名称的列表，用于生成函数绑定。如果提供的参数是字典，则该字典应该将函数名称映射到对应的文档字符串；如果没有提供文档字符串，默认使用函数名称作为文档字符串。
with_cuda – 确定是否将 CUDA 头文件和库添加到构建中。默认情况下，如果提供了 cuda_sources，则会自动确定此值。将其设置为True以强制包含 CUDA 头文件和库。
with_pytorch_error_handling – 决定是否由 PyTorch 而不是 pybind 处理 PyTorch 的错误和警告宏。为此，每个函数 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'])

注意

torch.utils.cpp_extension.include_paths(cuda=False)[源代码]

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

参数: cuda (bool) – 如果为True，则包含与CUDA相关的特定路径。
返回值: 包含路径的字符串列表。
返回类型: List[\str]

torch.utils.cpp_extension.get_compiler_abi_compatibility_and_version(compiler)[源代码]

判断给定的编译器是否与 PyTorch 二进制兼容（ABI 兼容），并确认其版本。

参数: compiler (str) – 要检查的编译器可执行文件名（例如 g++）。该文件必须可以在 shell 进程中运行。
返回值: 一个元组，第一个元素是一个布尔值，表示编译器是否（可能）与 PyTorch 的 ABI 不兼容；第二个元素是 TorchVersion 字符串，通过点分隔符包含编译器的版本号。
返回类型: 元组[bool, TorchVersion]

torch.utils.cpp_extension.verify_ninja_availability()[源代码]: 如果没有安装ninja构建系统，则抛出RuntimeError异常，否则不执行任何操作。

torch.utils.cpp_extension.is_ninja_available()[源代码]: 如果系统中安装了ninja构建系统，则返回True，否则返回False。