Skip to content

torch.utils.cpp_extension

原文: 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 时,可以提供extra_compile_args(而不是通常的列表)的字典,该字典从语言(cxxnvcc)映射到其他编译器标志的列表 提供给编译器。 这样就可以在混合编译期间向 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 进程中,并从此函数返回,以供使用。

默认情况下,生成文件的发布目录和编译到的结果库为<tmp>/torch_extensions/<name>,其中<tmp>是当前平台上的临时文件夹,<name>是扩展名。 可以通过两种方式覆盖此位置。 首先,如果设置了TORCH_EXTENSIONS_DIR环境变量,它将替换<tmp>/torch_extensions,所有扩展名都将编译到该目录的子文件夹中。 第二,如果提供了此函数的build_directory参数,它将覆盖整个路径,即库将直接编译到该文件夹​​中。

要编译源,使用默认的系统编译器(c++),可以通过设置CXX环境变量来覆盖它。 要将其他参数传递给编译过程,可以提供extra_cflagsextra_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自动确定该值。 将其设置为 <cite>True`</cite> 以强制包含 CUDA 标头和库。

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

退货

如果is_python_moduleTrue,则将加载的 PyTorch 扩展名作为 Python 模块返回。 如果is_python_moduleFalse,则什么都不返回(作为副作用,共享库已加载到进程中)。

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() 完全相同,但是将其源作为字符串而不是文件名使用。 这些字符串存储到构建目录中的文件中,之后 load_inline() 的行为与 load() 相同。

有关使用此功能的良好示例,请参见测试

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

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

cuda_sources中的源被连接到单独的.cu文件中,并以torch/types.hcuda.hcuda_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 而不是 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 -如果<cite>为真</cite>,则包含特定于 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 构建系统,则返回True



回到顶部