torch.utils.cpp_extension¶
- torch.utils.cpp_extension.CppExtension(name, sources, *args, **kwargs)[source]¶
为 C++ 创建
setuptools.Extension
。便利方法,使用构建 C++ 扩展的必要最小参数(但通常足够)来创建
setuptools.Extension
。所有参数都将转发到
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_args=['-Wl,--no-as-needed', '-lm']) ... ], ... cmdclass={ ... 'build_ext': BuildExtension ... })
- torch.utils.cpp_extension.CUDAExtension(name, sources, *args, **kwargs)[source]¶
为 CUDA/C++ 创建
setuptools.Extension
。便利方法,使用构建 CUDA/C++ 扩展的必要最小参数(但通常足够)来创建
setuptools.Extension
。这包括 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_args=['-Wl,--no-as-needed', '-lcuda']) ... ], ... 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 >= 指定的 CC 运行时编译(例如,8.6+PTX 生成可以为任何 CC >= 8.6 的 GPU 运行时编译的 PTX)。这提高了二进制文件的向前兼容性。但是,依赖旧的 PTX 通过为较新的 CC 运行时编译来提供向前兼容性可能会略微降低这些较新 CC 上的性能。如果您知道要定位的 GPU 的确切 CC,那么最好单独指定它们。例如,如果您希望扩展在 8.0 和 8.6 上运行,“8.0+PTX”在功能上可以工作,因为它包含可以为 8.6 运行时编译的 PTX,但“8.0 8.6”会更好。
请注意,虽然可以包含所有支持的架构,但包含的架构越多,构建过程就会越慢,因为它将为每个架构构建一个单独的内核映像。
请注意,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 错误的开放问题: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)[source]¶
自定义的
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)[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,就像使用extra_cflags
用于 C++ 一样。 使用各种启发式方法来查找 CUDA 安装目录,这些方法通常有效。 如果无效,设置CUDA_HOME
环境变量是最安全的选项。- 参数
name – 要构建的扩展的名称。 这必须与 pybind11 模块的名称相同!
extra_cflags – 要转发到构建的可选编译器标志列表。
extra_cuda_cflags – 在构建 CUDA 源文件时要转发给 nvcc 的可选编译器标志列表。
extra_ldflags – 要转发到构建的可选链接器标志列表。
extra_include_paths – 要转发到构建的可选包含目录列表。
build_directory – 用作构建工作区的可选路径。
verbose – 如果为
True
,则打开加载步骤的详细日志记录。with_cuda (Optional[bool]) – 确定是否将 CUDA 头文件和库添加到构建中。 如果设置为
None
(默认值),则此值将根据.cu
或.cuh
在sources
中的存在自动确定。 将其设置为 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
示例
>>> 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)[source]¶
从字符串源代码实时(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 头文件和库添加到构建中。如果设置为
None
(默认值),则此值将根据是否提供了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'])
注意
默认情况下,Ninja 后端使用 #CPUS + 2 个工作进程来构建扩展。 这可能会在某些系统上占用太多资源。 可以通过将 MAX_JOBS 环境变量设置为一个非负数来控制工作进程的数量。
- torch.utils.cpp_extension.get_compiler_abi_compatibility_and_version(compiler)[source]¶
确定给定的编译器是否与 PyTorch 的 ABI 兼容,以及它的版本。