Python PyTorch CUDAExtension用法及代碼示例

本文簡要介紹python語言中 torch.utils.cpp_extension.CUDAExtension 的用法。

用法:

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, 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
        })

計算能力：

默認情況下，擴展程序將被編譯為在擴展程序構建過程中可見的所有卡牌以及 PTX 上運行。如果以後安裝了新卡，則可能需要重新編譯擴展。如果可見卡的計算能力 (CC) 比您的 nvcc 可以為其構建 fully-compiled 二進製文件的最新版本更新，Pytorch 將使 nvcc 回退到使用您的 nvcc 支持的最新版本 PTX 構建內核(見下文有關 PTX 的詳細信息)。

您可以使用 TORCH_CUDA_ARCH_LIST 覆蓋默認行為，以明確指定您希望擴展支持哪些 CC：

TORCH_CUDA_ARCH_LIST=”6.1 8.6” python build_my_extension.py TORCH_CUDA_ARCH_LIST=”5.2 6.0 6.1 7.0 7.5 8.0 8.6+PTX” python build_my_extension.py

+PTX 選項使擴展內核二進製文件包含指定 CC 的 PTX 指令。 PTX 是一種中間表示，它允許內核為任何 CC >= 指定的 CC 運行時編譯(例如，8.6+PTX 生成的 PTX 可以為任何 CC >= 8.6 的 GPU 運行時編譯)。這提高了二進製文件的前向兼容性。但是，依靠較舊的 PTX 通過runtime-compiling 為較新的 CC 提供前向兼容可能會適度降低這些較新 CC 的性能。如果您知道要定位的 GPU 的確切 CC，最好單獨指定它們。例如，如果您希望您的擴展在 8.0 和 8.6 上運行，“8.0+PTX” 可以正常工作，因為它包含可以為 8.6 運行時編譯的 PTX，但“8.0 8.6”會更好。

請注意，雖然可以包含所有受支持的拱門，但包含的拱門越多，構建過程就越慢，因為它將為每個拱門構建單獨的內核映像。