torch.utils.cpp_extension¶

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

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

创建一个 setuptools.Extension 的便捷方法，该方法使用构建 C++ 扩展所需的最少（但通常足够）的参数。

所有参数都被转发到 setuptools.Extension 构造函数。完整的参数列表可以在 https://setuptools.pypa.io/en/latest/userguide/ext_modules.html#extension-api-reference 找到

示例

>>> 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_flags=['-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 构造函数。完整的参数列表可以在 https://setuptools.pypa.io/en/latest/userguide/ext_modules.html#extension-api-reference 找到

示例

>>> 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_flags=['-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 选项使扩展内核二进制文件包含指定 CC 的 PTX 指令。PTX 是一种中间表示形式，允许内核在运行时为任何 CC >= 指定 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 at::Tensor SigmoidAlphaBlendForwardCuda(….)
Instead of:: #include torch::Tensor SigmoidAlphaBlendForwardCuda(…)

目前关于nvcc错误的开放问题: https://github.com/pytorch/pytorch/issues/69460 完整的解决代码示例: https://github.com/facebookresearch/pytorch3d/commit/cb170ac024a949f1f9614ffe6af1c38d972f7d48

可重定位设备代码链接：

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

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

注意：使用Ninja构建带有RDC链接的CUDA扩展是必需的。

示例

>>> 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 后端进行构建。与标准的 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)[源代码]¶

加载一个PyTorch C++扩展，即时（JIT）。

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

默认情况下，构建文件生成的目录和编译后的库所在的目录是 /torch_extensions/，其中是当前平台上的临时文件夹，是扩展的名称。此位置可以通过两种方式覆盖。首先，如果设置了 TORCH_EXTENSIONS_DIR 环境变量，它将替换 /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传递额外的标志，就像通过extra_cflags为C++传递标志一样。使用了各种启发式方法来查找CUDA安装目录，通常效果良好。如果没有，设置CUDA_HOME环境变量是最安全的选择。

Parameters

名称 – 要构建的扩展的名称。这必须与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 (可选[布尔值]) – 确定是否将CUDA头文件和库添加到构建中。如果设置为None（默认），此值将根据sources中是否存在.cu或.cuh自动确定。设置为True`以强制包含CUDA头文件和库。
is_python_module – 如果 True（默认），将生成的共享库作为Python模块导入。如果 False，行为取决于 is_standalone。
is_standalone – 如果 False（默认），将构建的扩展加载到进程中作为普通的动态库。如果 True，则构建一个独立的可执行文件。

Returns

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

如果 is_python_module 是 False 并且 is_standalone 是 False：: 不返回任何内容。（共享库作为副作用加载到进程中。）
如果 is_standalone 是 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, use_pch=False)[源代码]¶

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

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

参见测试以获取使用此函数的良好示例。

源文件可能会省略典型非内联 C++ 扩展的两个必需部分：必要的头文件包含以及（pybind11）绑定代码。更准确地说，传递给 cpp_sources 的字符串首先被连接成一个单独的 .cpp 文件。然后，这个文件会被预先加上 #include。

此外，如果提供了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源代码。
函数 – 要为其生成函数绑定的函数名称列表。如果给定的是字典，则应将函数名称映射到文档字符串（否则文档字符串仅为函数名称）。
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)[源代码]¶

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

Parameters: cuda (布尔值) – 如果为True，则包含CUDA特定的包含路径。
Returns: 包含路径字符串的列表。
Return type: 列表[字符串]

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

确定给定的编译器是否与其版本一起与PyTorch ABI兼容。

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

torch.utils.cpp_extension.verify_ninja_availability()[源代码]¶: 如果系统上没有可用的ninja构建系统，则引发RuntimeError，否则不执行任何操作。

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