torch.utils.cpp_extension

torch.utils.cpp_extension.CppExtension(name, sources, args, *kwargs)

创建一个C++setuptools.Extension

便捷地创建一个setuptools.Extension具有最小(但通常是足够)的参数来构建C++扩展的方法。

所有参数都被转发给setuptools.Extension构造函数。

  1. >>> from setuptools import setup
  2. >>> from torch.utils.cpp_extension import BuildExtension, CppExtension
  3. >>> setup(
  4. name='extension',
  5. ext_modules=[
  6. CppExtension(
  7. name='extension',
  8. sources=['extension.cpp'],
  9. extra_compile_args=['-g'])),
  10. ],
  11. cmdclass={
  12. 'build_ext': BuildExtension
  13. })

torch.utils.cpp_extension.CUDAExtension(name, sources, args, *kwargs)

CUDA/C++创建一个setuptools.Extension。 创建一个setuptools.Extension用于构建CUDA/C ++扩展的最少参数(但通常是足够的)的便捷方法。这里包括CUDA路径,库路径和运行库。 所有参数都被转发给setuptools.Extension构造函数。

  1. >>> from setuptools import setup
  2. >>> from torch.utils.cpp_extension import BuildExtension, CppExtension
  3. >>> setup(
  4. name='cuda_extension',
  5. ext_modules=[
  6. CUDAExtension(
  7. name='cuda_extension',
  8. sources=['extension.cpp', 'extension_kernel.cu'],
  9. extra_compile_args={'cxx': ['-g'],
  10. 'nvcc': ['-O2']})
  11. ],
  12. cmdclass={
  13. 'build_ext': BuildExtension
  14. })

torch.utils.cpp_extension.BuildExtension(dist,** kw )[source]

自定义setuptools构建扩展。

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

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

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

torch.utils.cpp_extension.include_paths(cuda=False)

获取构建C++CUDA扩展所需的路径。

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

例如:

  1. from setuptools import setup
  2. from torch.utils.cpp_extension import BuildExtension, CppExtension
  3. torch.utils.cpp_extension.include_paths(cuda=False)
  4. # ['/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

](https://ptorch.com/docs/8/torch-utils-cpp_extension)

译者署名

用户名 头像 职能 签名
Song torch.utils.cpp_extension - 图1 翻译 人生总要追求点什么