PyTorch torch.utils.cpp_extension

原文： PyTorch torch.utils.cpp_extension

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

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

一種便捷方法，它使用最少的(但通常是足夠的）參數(shù)創(chuàng)建setuptools.Extension來(lái)構(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)建一個(gè)setuptools.Extension。

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

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ò)展。

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

使用 BuildExtension 時(shí)，可以提供extra_compile_args(而不是通常的列表）的字典，該字典從語(yǔ)言(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)?

即時(shí)加載 PyTorch C ++擴(kuò)展(JIT）。

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

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

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

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

參數(shù)

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

退貨

如果is_python_module為True，則將加載的 PyTorch 擴(kuò)展名作為 Python 模塊返回。 如果is_python_module為False，則什么都不返回(作為副作用，共享庫(kù)已加載到進(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í)時(shí)加載 PyTorch C ++擴(kuò)展(JIT）。

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

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

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

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

有關(guān)以下省略的自變量的說(shuō)明，請(qǐng)參見(jiàn) load() 。

Parameters

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

Returns

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

torch.utils.cpp_extension.verify_ninja_availability()?

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