pytorch通过torch.utils.cpp_extension构建CUDA/C++拓展

注意这个和前面的《Python与C语言混合编程：通过distutils或setuptools实现的一个简单的C扩展》不同，这个是pytorch的扩展，不是python的扩展。

在pytorch的utils中，集成了setuptools模块。

官方文档在这里：https://pytorch.org/docs/master/cpp_extension.html

中文说明在这里：https://ptorch.com/news/188.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路径，库路径和运行库。 所有参数都被转发给setuptools.Extension构造函数。

例

>>> from setuptools import setup
>>> from torch.utils.cpp_extension import BuildExtension, CppExtension
>>> 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（dist，** kw ）

自定义setuptools构建扩展。

setuptools.build_ext子类负责传递所需的最小编译器参数（例如-std=c++11）以及混合的C ++/CUDA编译（以及一般对CUDA文件的支持）。

当使用BuildExtension时，它将提供一个用于extra_compile_args（不是普通列表）的词典，通过语言（cxx或cuda）映射到参数列表提供给编译器。这样可以在混合编译期间为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)

即时加载(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传递进一步包含目录。

提供了混合编译的CUDA支持。只需将CUDA源文件（.cu或.cuh）与其他源一起传递即可。这些文件将被检测，并且使用nvcc而不是C ++编译器进行编译。包括将CUDA lib64目录作为库目录传递并进行cudart链接。您可以将其他参数传递给nvcc extra_cuda_cflags，就像使用C ++的extra_cflags一样。使用了各种原始方法来查找CUDA安装目录，通常情况下可以正常运行。如果不可以，最好设置CUDA_HOME环境变量。

• 参数：
  • name - 要构建的扩展名。这个必须和pybind11模块的名字一样！
  • sources - C++源文件的相对或绝对路径列表。
  • extra_cflags - 编译器参数的可选列表，用于转发到构建。
  • extra_cuda_cflags - 编译器标记的可选列表，在构建CUDA源时转发给nvcc。
  • extra_ldflags - 链接器参数的可选列表，用于转发到构建。
  • extra_include_paths - 转发到构建的包含目录的可选列表。
  • build_directory - 可选路径作为构建区域。
  • verbose - 如果为True，打开加载步骤的详细记录。
• 返回：
  • 加载PyTorch扩展作为Python模块。

例

>>> from torch.utils.cpp_extension import load
>>> module = load(
        name='extension',
        sources=['extension.cpp', 'extension_kernel.cu'],
        extra_cflags=['-O2'],
        verbose=True)

后面还新添加了一个load_inline,不过没有中文翻译

Loads a PyTorch C++ extension just-in-time (JIT) from string sources.

This function behaves exactly like load(), but takes its sources as strings rather than filenames. These strings are stored to files in the build directory, after which the behavior of load_inline() is identical to load().

See the tests for good examples of using this function.

Sources may omit two required parts of a typical non-inline C++ extension: the necessary header includes, as well as the (pybind11) binding code. More precisely, strings passed to cpp_sources are first concatenated into a single .cpp file. This file is then prepended with #include <torch/extension.h>.

Furthermore, if the functions argument is supplied, bindings will be automatically generated for each function specified. functions can either be a list of function names, or a dictionary mapping from function names to docstrings. If a list is given, the name of each function is used as its docstring.

The sources in cuda_sources are concatenated into a separate .cu file and prepended with torch/types.h, cuda.hand cuda_runtime.h includes. The .cpp and .cu files are compiled separately, but ultimately linked into a single library. Note that no bindings are generated for functions in cuda_sources per se. To bind to a CUDA kernel, you must create a C++ function that calls it, and either declare or define this C++ function in one of the cpp_sources (and include its name in functions).

See load() for a description of arguments omitted below.

Parameters:

cpp_sources – A string, or list of strings, containing C++ source code. cuda_sources – A string, or list of strings, containing CUDA source code. functions – A list of function names for which to generate function bindings. If a dictionary is given, it should map function names to docstrings (which are otherwise just the function names). with_cuda – Determines whether CUDA headers and libraries are added to the build. If set to None (default), this value is automatically determined based on whether cuda_sources is provided. Set it to True` to force CUDA headers and libraries to be included.

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扩展所需的路径。

• 参数： cuda - 如果为True，则包含CUDA特定的包含路径。
• 返回： 包含路径字符串的列表。

例如：

from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CppExtension

torch.utils.cpp_extension.include_paths(cuda=False)
# 
['/usr/local/lib/python3.6/site-packages/torch/lib/include', 
'/usr/local/lib/python3.6/site-packages/torch/lib/include/TH', 
'/usr/local/lib/python3.6/site-packages/torch/lib/include/THC']

torch.utils.cpp_extension.check_compiler_abi_compatibility(compiler)

验证给定的编译器是否与PyTorch ABI兼容。

• 参数：compiler(str) - 要检查可执行的编译器文件名(例如g++),必须在shell进程中可执行。
• 返回：如果编译器（可能）与PyTorchABI不兼容，则为False，否则返回True。

torch.utils.cpp_extension.verify_ninja_availability()

如果可以在ninja上运行则返回True。

文档地址：torch.utils.cpp_extension