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 找到警告
PyTorch python API(由 libtorch_python 提供)不能使用
py_limited_api=True
标志构建。传递此标志时,用户有责任在其库中不使用来自 libtorch_python(特别是 pytorch/python 绑定)的 API,而仅使用来自 libtorch(aten 对象、操作符和调度器)的 API。例如,要从 python 访问自定义操作符,库应该通过调度器注册这些操作符。与 CPython setuptools 不同,后者在
setup
中的 “bdist_wheel” 命令指定 py_limited_api 选项时不会将 -DPy_LIMITED_API 定义为编译标志,而 PyTorch 会这样做!我们将指定 -DPy_LIMITED_API=min_supported_cpython 以最大限度地强制一致性、安全性和合理性,从而鼓励最佳实践。要指定不同的版本,请将 min_supported_cpython 设置为所需 CPython 版本的十六进制代码。示例
>>> 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 include 路径、library 路径和 runtime library。所有参数都会转发给
setuptools.Extension
构造函数。参数完整列表可在 https://setuptools.pypa.io/en/latest/userguide/ext_modules.html#extension-api-reference 找到警告
PyTorch python API(由 libtorch_python 提供)不能使用
py_limited_api=True
标志构建。传递此标志时,用户有责任在其库中不使用来自 libtorch_python(特别是 pytorch/python 绑定)的 API,而仅使用来自 libtorch(aten 对象、操作符和调度器)的 API。例如,要从 python 访问自定义操作符,库应该通过调度器注册这些操作符。与 CPython setuptools 不同,后者在
setup
中的 “bdist_wheel” 命令指定 py_limited_api 选项时不会将 -DPy_LIMITED_API 定义为编译标志,而 PyTorch 会这样做!我们将指定 -DPy_LIMITED_API=min_supported_cpython 以最大限度地强制一致性、安全性和合理性,从而鼓励最佳实践。要指定不同的版本,请将 min_supported_cpython 设置为所需 CPython 版本的十六进制代码。示例
>>> 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 ... })
计算能力 (Compute capabilities)
默认情况下,扩展将被编译以在构建过程中可见的卡的所有架构(archs)上运行,外加 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 生成的 PTX 可以为任何 CC >= 8.6 的 GPU 进行运行时编译)。这提高了二进制文件的向前兼容性。然而,依赖较旧的 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 bug 的开放问题:https://github.com/pytorch/pytorch/issues/69460 完整的 workaround 代码示例:https://github.com/facebookresearch/pytorch3d/commit/cb170ac024a949f1f9614ffe6af1c38d972f7d48
可重定位设备代码链接 (Relocatable device code linking)
如果您想跨编译单元(跨目标文件)引用设备符号,则需要使用 relocatable device code (-rdc=true 或 -dc)构建目标文件。此规则的一个例外是“dynamic parallelism”(嵌套内核启动),目前已不太常用。Relocatable device code 优化程度较低,因此仅应在需要它的目标文件上使用。在设备代码编译步骤和 dlink 步骤中使用 -dlto (Device Link Time Optimization) 有助于减少 -rdc 可能导致的性能下降。请注意,必须在两个步骤中都使用它才能有效。
如果您有 rdc 对象,则需要在 CPU 符号链接步骤之前额外进行 -dlink (设备链接)步骤。还有一种情况是在没有 -rdc 的情况下使用 -dlink:当扩展链接到包含由 rdc 编译的对象的静态库时,例如 [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.SyclExtension(name, sources, *args, **kwargs)[源代码][源代码]¶
为 SYCL/C++ 创建
setuptools.Extension
。便利方法,用于创建具有最基本(但通常已足够)参数的
setuptools.Extension
来构建 SYCL/C++ 扩展。所有参数都会转发给
setuptools.Extension
构造函数。警告
PyTorch python API(由 libtorch_python 提供)不能使用
py_limited_api=True
标志构建。传递此标志时,用户有责任在其库中不使用来自 libtorch_python(特别是 pytorch/python 绑定)的 API,而仅使用来自 libtorch(aten 对象、操作符和调度器)的 API。例如,要从 python 访问自定义操作符,库应该通过调度器注册这些操作符。与 CPython setuptools 不同,后者在
setup
中的 “bdist_wheel” 命令指定 py_limited_api 选项时不会将 -DPy_LIMITED_API 定义为编译标志,而 PyTorch 会这样做!我们将指定 -DPy_LIMITED_API=min_supported_cpython 以最大限度地强制一致性、安全性和合理性,从而鼓励最佳实践。要指定不同的版本,请将 min_supported_cpython 设置为所需 CPython 版本的十六进制代码。示例
>>> from torch.utils.cpp_extension import BuildExtension, SyclExtension >>> setup( ... name='xpu_extension', ... ext_modules=[ ... SyclExtension( ... name='xpu_extension', ... sources=['extension.cpp', 'extension_kernel.cpp'], ... extra_compile_args={'cxx': ['-g', '-std=c++20', '-fPIC']}) ... ], ... cmdclass={ ... 'build_ext': BuildExtension ... })
默认情况下,扩展将被编译以在构建过程中可见的卡的所有架构(archs)上运行。如果之后安装了新卡,扩展可能需要重新编译。您可以使用 TORCH_XPU_ARCH_LIST 明确指定您希望扩展支持哪些设备架构(device architectures),从而覆盖默认行为
TORCH_XPU_ARCH_LIST="pvc,xe-lpg" python build_my_extension.py
请注意,虽然可以包含所有受支持的架构,但包含的架构越多,构建过程就会越慢,因为它会为每个架构构建一个单独的内核映像。
注意:构建 SyclExtension 需要 Ninja。
- torch.utils.cpp_extension.BuildExtension(*args, **kwargs)[源代码][源代码]¶
一个自定义的
setuptools
构建扩展。这个
setuptools.build_ext
子类负责传递所需的最低编译器标志(例如-std=c++17
)以及混合 C++/CUDA/SYCL 编译(以及对一般 CUDA/SYCL 文件的支持)。使用
BuildExtension
时,允许为extra_compile_args
提供一个字典(而不是通常的列表),该字典将语言/编译器(唯一预期值为cxx
、nvcc
或sycl
)映射到要提供给编译器的附加编译器标志列表。这使得在混合编译期间可以向 C++、CUDA 和 SYCL 编译器提供不同的标志。use_ninja
(bool):如果use_ninja
为True
(默认),则尝试使用 Ninja 后端进行构建。与标准的setuptools.build_ext
相比,Ninja 大大加快了编译速度。如果 Ninja 不可用,则回退到标准的 distutils 后端。注意
默认情况下,Ninja 后端使用 #CPUS + 2 个 worker 来构建扩展。这可能会在某些系统上占用过多资源。可以通过将 MAX_JOBS 环境变量设置为非负数来控制 worker 的数量。
- torch.utils.cpp_extension.load(name, sources, extra_cflags=None, extra_cuda_cflags=None, extra_sycl_cflags=None, extra_ldflags=None, extra_include_paths=None, build_directory=None, verbose=False, with_cuda=None, with_sycl=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
。例如,要使用优化编译扩展,请传递extra_cflags=['-O3']
。您也可以使用extra_cflags
传递其他 include 目录。提供对 CUDA 的混合编译支持。只需传递 CUDA 源文件(
.cu
或.cuh
)以及其他源文件。这些文件将被检测到并由 nvcc 而非 C++ 编译器编译。这包括将 CUDA lib64 目录作为库目录传递,并链接cudart
。您可以通过extra_cuda_cflags
向 nvcc 传递附加标志,就像使用extra_cflags
传递给 C++ 一样。使用各种启发式方法查找 CUDA 安装目录,通常效果不错。如果不行,设置CUDA_HOME
环境变量是最安全的选项。提供对 SYCL 的混合编译支持。只需传递 SYCL 源文件(
.sycl
)以及其他源文件。这些文件将被检测到并由 SYCL 编译器(例如 Intel DPC++ Compiler)而非 C++ 编译器编译。您可以通过extra_sycl_cflags
向 SYCL 编译器传递附加标志,就像使用extra_cflags
传递给 C++ 一样。SYCL 编译器预计通过系统 PATH 环境变量找到。- 参数
name – 要构建的扩展的名称。这必须与 pybind11 模块的名称相同!
extra_cflags – 可选的编译器标志列表,用于转发到构建过程。
extra_cuda_cflags – 构建 CUDA 源文件时,可选的编译器标志列表,用于转发到 nvcc。
extra_sycl_cflags – 构建 SYCL 源文件时,可选的编译器标志列表,用于转发到 SYCL 编译器。
extra_ldflags – 可选的链接器标志列表,用于转发到构建过程。
extra_include_paths – 可选的 include 目录列表,用于转发到构建过程。
build_directory – 可选的路径,用作构建工作空间。
verbose – 如果为
True
,则开启加载步骤的详细日志记录。with_cuda (Optional[bool]) – 确定是否将 CUDA 头文件和库添加到构建中。如果设置为
None
(默认),则根据sources
中是否存在.cu
或.cuh
自动确定此值。设置为 True 强制包含 CUDA 头文件和库。with_sycl (Optional[bool]) – 确定是否将 SYCL 头文件和库添加到构建中。如果设置为
None
(默认),则根据sources
中是否存在.sycl
自动确定此值。设置为 True 强制包含 SYCL 头文件和库。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, sycl_sources=None, functions=None, extra_cflags=None, extra_cuda_cflags=None, extra_sycl_cflags=None, extra_ldflags=None, extra_include_paths=None, build_directory=None, verbose=False, with_cuda=None, with_sycl=None, is_python_module=True, with_pytorch_error_handling=True, keep_intermediates=True, use_pch=False)[source][source]¶
从字符串源即时(JIT)加载 PyTorch C++ 扩展。
此函数与
load()
函数的行为完全一致,但它接收的是字符串形式的源而不是文件名。这些字符串被存储到构建目录中的文件里,此后load_inline()
的行为与load()
函数相同。请参阅测试用例,以获取使用此函数的良好示例。
源文件可以省略典型非内联 C++ 扩展的两个必需部分:必要的头文件包含以及 (pybind11) 绑定代码。更准确地说,传递给
cpp_sources
的字符串首先被连接成一个单独的.cpp
文件。然后此文件会在开头被加上#include <torch/extension.h>
。此外,如果提供了
functions
参数,则会为指定的每个函数自动生成绑定。functions
可以是函数名称列表,也可以是函数名称到文档字符串(docstrings)的字典。如果提供的是列表,则每个函数的名称将用作其文档字符串。cuda_sources
中的源被连接成一个单独的.cu
文件,并在开头被加上torch/types.h
、cuda.h
和cuda_runtime.h
等包含。.cpp
和.cu
文件会分开编译,但最终会链接到一个单独的库中。请注意,cuda_sources
中的函数本身不会生成绑定。要绑定到 CUDA 内核,您必须创建一个调用该内核的 C++ 函数,并在其中一个cpp_sources
中声明或定义此 C++ 函数(并将其名称包含在functions
中)。sycl_sources
中的源被连接成一个单独的.sycl
文件,并在开头被加上torch/types.h
、sycl/sycl.hpp
等包含。.cpp
和.sycl
文件会分开编译,但最终会链接到一个单独的库中。请注意,sycl_sources
中的函数本身不会生成绑定。要绑定到 SYCL 内核,您必须创建一个调用该内核的 C++ 函数,并在其中一个cpp_sources
中声明或定义此 C++ 函数(并将其名称包含在functions
中)。有关以下省略的参数的描述,请参阅
load()
。- 参数
cpp_sources – 包含 C++ 源代码的字符串或字符串列表。
cuda_sources – 包含 CUDA 源代码的字符串或字符串列表。
sycl_sources – 包含 SYCL 源代码的字符串或字符串列表。
functions – 要为其生成函数绑定的函数名称列表。如果提供字典,则应将函数名称映射到文档字符串(否则文档字符串就是函数名称本身)。
with_cuda – 确定是否将 CUDA 头文件和库添加到构建中。如果设置为
None
(默认值),则此值根据是否提供了cuda_sources
自动确定。将其设置为True
以强制包含 CUDA 头文件和库。with_sycl – 确定是否将 SYCL 头文件和库添加到构建中。如果设置为
None
(默认值),则此值根据是否提供了sycl_sources
自动确定。将其设置为True
以强制包含 SYCL 头文件和库。with_pytorch_error_handling – 确定 PyTorch 的错误和警告宏是由 PyTorch 而非 pybind 处理。为此,每个函数
foo
都通过一个中间函数_safe_foo
来调用。这种重定向在一些晦涩的 C++ 用例中可能会导致问题。当此重定向导致问题时,应将此标志设置为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'])
注意
由于 load_inline 会即时编译源代码,请确保您的运行时环境中安装了正确的工具链。例如,加载 C++ 时,请确保有可用的 C++ 编译器。如果加载 CUDA 扩展,则需要额外安装相应的 CUDA 工具包(nvcc 以及代码的任何其他依赖项)。安装 torch 时不包含编译工具链,必须额外安装。
编译期间,默认情况下,Ninja 后端使用 #CPUS + 2 个工作进程来构建扩展。这在某些系统上可能会占用过多资源。可以通过将 MAX_JOBS 环境变量设置为非负数来控制工作进程的数量。
- torch.utils.cpp_extension.include_paths(device_type='cpu')[source][source]¶
获取构建 C++、CUDA 或 SYCL 扩展所需的包含路径。
- torch.utils.cpp_extension.get_compiler_abi_compatibility_and_version(compiler)[source][source]¶
确定给定的编译器是否与 PyTorch ABI 兼容,并获取其版本。