PyTorch torch.utils.cpp_extension

原文： PyTorch 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']),
        ],
        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(而不是通常的列表）的字典，该字典从语言(cxx或nvcc）映射到其他编译器标志的列表提供给编译器。这样就可以在混合编译期间向 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_cflags或extra_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_module为True，则将加载的 PyTorch 扩展名作为 Python 模块返回。如果is_python_module为False，则什么都不返回(作为副作用，共享库已加载到进程中）。

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 <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() 。

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。

PyTorch torch.utils.cpp_extension

推荐文章

推荐教程

推荐课程