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
。