PyTorch torch.utils.cpp_extension

原文： PyTorch torch.utils.cpp_extension

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

為 C ++創(chuàng)建一個setuptools.Extension。

一種便捷方法，它使用最少的(但通常是足夠的）參數(shù)創(chuàng)建setuptools.Extension來構(gòu)建 C ++擴(kuò)展。

所有參數(shù)都轉(zhuǎn)發(fā)到setuptools.Extension構(gòu)造函數(shù)。

例

>>> 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 ++創(chuàng)建一個setuptools.Extension。

一種便捷方法，它使用最少的(但通常是足夠的）參數(shù)創(chuàng)建setuptools.Extension，以構(gòu)建 CUDA / C ++擴(kuò)展。 這包括 CUDA 包含路徑，庫路徑和運(yùn)行時庫。

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構(gòu)建擴(kuò)展。

這個setuptools.build_ext子類負(fù)責(zé)傳遞所需的最低編譯器標(biāo)志(例如-std=c++11）以及混合的 C ++ / CUDA 編譯(并通常支持 CUDA 文件）。

使用 BuildExtension 時，可以提供extra_compile_args(而不是通常的列表）的字典，該字典從語言(cxx或nvcc）映射到其他編譯器標(biāo)志的列表 提供給編譯器。 這樣就可以在混合編譯期間向 C ++和 CUDA 編譯器提供不同的標(biāo)志。

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 ++擴(kuò)展(JIT）。

要加載擴(kuò)展，將發(fā)出 Ninja 構(gòu)建文件，該文件用于將給定的源編譯到動態(tài)庫中。 隨后將該庫作為模塊加載到當(dāng)前的 Python 進(jìn)程中，并從此函數(shù)返回，以供使用。

默認(rèn)情況下，生成文件的發(fā)布目錄和編譯到的結(jié)果庫為<tmp>/torch_extensions/<name>，其中<tmp>是當(dāng)前平臺上的臨時文件夾，<name>是擴(kuò)展名。 可以通過兩種方式覆蓋此位置。 首先，如果設(shè)置了TORCH_EXTENSIONS_DIR環(huán)境變量，它將替換<tmp>/torch_extensions，所有擴(kuò)展名都將編譯到該目錄的子文件夾中。 第二，如果提供了此函數(shù)的build_directory參數(shù)，它將覆蓋整個路徑，即庫將直接編譯到該文件夾中。

要編譯源，使用默認(rèn)的系統(tǒng)編譯器(c++），可以通過設(shè)置CXX環(huán)境變量來覆蓋它。 要將其他參數(shù)傳遞給編譯過程，可以提供extra_cflags或extra_ldflags。 例如，要使用優(yōu)化來編譯擴(kuò)展，請傳遞extra_cflags=['-O3']。 您也可以使用extra_cflags傳遞更多的包含目錄。

提供帶有混合編譯的 CUDA 支持。 只需將 CUDA 源文件(.cu或.cuh）與其他源一起傳遞即可。 將使用 nvcc 而不是 C ++編譯器檢測并編譯此類文件。 這包括將 CUDA lib64 目錄作為庫目錄傳遞，并鏈接cudart。 您可以通過extra_cuda_cflags將其他標(biāo)志傳遞給 nvcc，就像 C ++的extra_cflags一樣。 使用各種啟發(fā)式方法來查找 CUDA 安裝目錄，通常可以正常工作。 否則，設(shè)置CUDA_HOME環(huán)境變量是最安全的選擇。

參數(shù)

• 名稱 –要構(gòu)建的擴(kuò)展名。 該名稱必須與 pybind11 模塊的名稱相同！
• 源 – C ++源文件的相對或絕對路徑的列表。
• extra_cflags –編譯器標(biāo)志的可選列表，以轉(zhuǎn)發(fā)到構(gòu)建。
• extra_cuda_cflags –生成 CUDA 源時轉(zhuǎn)發(fā)到 nvcc 的編譯器標(biāo)志的可選列表。
• extra_ldflags –鏈接標(biāo)志的可選列表，以轉(zhuǎn)發(fā)到構(gòu)建。
• extra_include_paths –包含目錄的可選列表，以轉(zhuǎn)發(fā)到構(gòu)建。
• build_directory –用作構(gòu)建工作區(qū)的可選路徑。
• verbose –如果True，則打開加載步驟的詳細(xì)日志記錄。
• with_cuda –確定是否將 CUDA 標(biāo)頭和庫添加到構(gòu)建中。 如果設(shè)置為None(默認(rèn)值），則根據(jù)sources中是否存在.cu或.cuh自動確定該值。 將其設(shè)置為 <cite>True`</cite> 以強(qiáng)制包含 CUDA 標(biāo)頭和庫。
• is_python_module –如果為True(默認(rèn)），則將生成的共享庫作為 Python 模塊導(dǎo)入。 如果為False，則將其作為純動態(tài)庫加載到進(jìn)程中。

退貨

如果is_python_module為True，則將加載的 PyTorch 擴(kuò)展名作為 Python 模塊返回。 如果is_python_module為False，則什么都不返回(作為副作用，共享庫已加載到進(jìn)程中）。

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)?

從字符串源實(shí)時加載 PyTorch C ++擴(kuò)展(JIT）。

此函數(shù)的行為與 load() 完全相同，但是將其源作為字符串而不是文件名使用。 這些字符串存儲到構(gòu)建目錄中的文件中，之后 load_inline() 的行為與 load() 相同。

源可能會省略典型的非內(nèi)聯(lián) C ++擴(kuò)展的兩個必需部分：必需的頭文件以及(pybind11）綁定代碼。 更準(zhǔn)確地說，首先將傳遞給cpp_sources的字符串連接到單個.cpp文件中。 該文件然后以#include <torch/extension.h>開頭。

此外，如果提供functions參數(shù)，則將為指定的每個函數(shù)自動生成綁定。 functions可以是函數(shù)名稱列表，也可以是從函數(shù)名稱到文檔字符串的字典映射。 如果給出了列表，則將每個函數(shù)的名稱用作其文檔字符串。

cuda_sources中的源被連接到單獨(dú)的.cu文件中，并以torch/types.h，cuda.h和cuda_runtime.h包括在內(nèi)。 .cpp和.cu文件是分別編譯的，但最終鏈接到一個庫中。 注意，cuda_sources本身不為函數(shù)生成任何綁定。 要綁定到 CUDA 內(nèi)核，您必須創(chuàng)建一個調(diào)用它的 C ++函數(shù)，并在cpp_sources之一中聲明或定義此 C ++函數(shù)(并在functions中包括其名稱）。

有關(guān)以下省略的自變量的說明，請參見 load() 。

Parameters

• cpp_sources –包含 C ++源代碼的字符串或字符串列表。
• cuda_sources –包含 CUDA 源代碼的字符串或字符串列表。
• 函數(shù) –為其生成函數(shù)綁定的函數(shù)名稱列表。 如果提供了字典，則應(yīng)將函數(shù)名稱映射到文檔字符串(否則僅是函數(shù)名稱）。
• with_cuda –確定是否將 CUDA 標(biāo)頭和庫添加到構(gòu)建中。 如果設(shè)置為None(默認(rèn)），則根據(jù)是否提供cuda_sources自動確定該值。 將其設(shè)置為True以強(qiáng)制包含 CUDA 標(biāo)頭和庫。
• with_pytorch_error_handling –確定 pytorch 而不是 pybind 處理 pytorch 錯誤和警告宏。 為此，每個功能foo都通過中間_safe_foo功能調(diào)用。 這種重定向在 cpp 晦澀的情況下可能會引起問題。 當(dāng)此重定向?qū)е聠栴}時，應(yīng)將此標(biāo)志設(shè)置為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)?

獲取構(gòu)建 C ++或 CUDA 擴(kuò)展所需的包含路徑。

Parameters

cuda -如果<cite>為真</cite>，則包含特定于 CUDA 的包含路徑。

Returns

包含路徑字符串的列表。

torch.utils.cpp_extension.check_compiler_abi_compatibility(compiler)?

驗(yàn)證給定的編譯器是否與 PyTorch 兼容。

Parameters

編譯器 (str )–要檢查的編譯器可執(zhí)行文件名稱(例如g++）。 必須在 Shell 進(jìn)程中可執(zhí)行。

Returns

如果編譯器(可能）與 PyTorch 不兼容，則為 False，否則為 True。

torch.utils.cpp_extension.verify_ninja_availability()?

如果系統(tǒng)上有 ninja 構(gòu)建系統(tǒng)，則返回True。