torch.utils.cpp_extension

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

為 C++ 建立 setuptools.Extension。

便捷方法，使用最少（但通常足夠）的參數建立 setuptools.Extension 以建置 C++ 擴充功能。

所有參數都會轉發至 setuptools.Extension 建構函式。完整的參數清單可以在 https://setuptools.pypa.io/en/latest/userguide/ext_modules.html#extension-api-reference 找到

範例

>>> 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'],
...             extra_link_flags=['-Wl,--no-as-needed', '-lm'])
...     ],
...     cmdclass={
...         'build_ext': BuildExtension
...     })

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

為 CUDA/C++ 建立 setuptools.Extension。

便捷方法，使用最少（但通常足夠）的參數建立 setuptools.Extension 以建置 CUDA/C++ 擴充功能。這包括 CUDA include 路徑、函式庫路徑和執行階段函式庫。

所有參數都會轉發至 setuptools.Extension 建構函式。完整的參數清單可以在 https://setuptools.pypa.io/en/latest/userguide/ext_modules.html#extension-api-reference 找到

範例

>>> 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']},
...                 extra_link_flags=['-Wl,--no-as-needed', '-lcuda'])
...     ],
...     cmdclass={
...         'build_ext': BuildExtension
...     })

計算能力

根據預設，擴充功能將被編譯為在擴充功能建置過程中可見的所有架構卡上執行，以及 PTX。如果稍後安裝了新的顯示卡，則可能需要重新編譯擴充功能。如果可見的顯示卡具有比 nvcc 可以建置完整編譯的二進制檔案的最新版本更新的計算能力 (CC)，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 會產生可以針對任何 CC >= 8.6 的 GPU 進行執行階段編譯的 PTX）。這提高了二進制檔案的向前相容性。但是，依靠較舊的 PTX 通過針對較新的 CC 進行執行階段編譯來提供向前相容性，可能會略微降低這些較新的 CC 上的效能。如果您知道要鎖定的 GPU 的確切 CC，最好單獨指定它們。例如，如果您希望擴充功能在 8.0 和 8.6 上執行，“8.0+PTX”在功能上可以運作，因為它包含可以針對 8.6 進行執行階段編譯的 PTX，但“8.0 8.6”會更好。

請注意，雖然可以包含所有支援的架構，但包含的架構越多，建置過程就會越慢，因為它會為每個架構建置一個單獨的核心映像檔。

請注意，CUDA-11.5 nvcc 在解析 Windows 上的 torch/extension.h 時會遇到內部編譯器錯誤。若要解決此問題，請將 python 繫結邏輯移至純 C++ 檔案。

範例用法

#include <ATen/ATen.h> at::Tensor SigmoidAlphaBlendForwardCuda(….)

而不是

#include <torch/extension.h> torch::Tensor SigmoidAlphaBlendForwardCuda(…)

目前針對 nvcc 錯誤開啟的問題：https://github.com/pytorch/pytorch/issues/69460 完整的解決方案程式碼範例：https://github.com/facebookresearch/pytorch3d/commit/cb170ac024a949f1f9614ffe6af1c38d972f7d48

可重定位裝置程式碼連結

如果您想跨編譯單元（跨物件檔案）參考裝置符號，則需要使用 可重定位裝置程式碼 (-rdc=true 或 -dc) 建置物件檔案。此規則的一個例外是“動態平行”（巢狀核心啟動），它現在已經不常使用了。可重定位裝置程式碼 的最佳化程度較低，因此只需要在需要它的物件檔案上使用它。在裝置程式碼編譯步驟和 dlink 步驟中使用 -dlto（裝置連結時間最佳化）有助於降低 -rdc 的潛在效能下降。請注意，它需要在兩個步驟中都使用才能發揮作用。

如果您有 rdc 物件，則在 CPU 符號連結步驟之前，您需要執行額外的 -dlink（裝置連結）步驟。還有一種情況是在沒有 -rdc 的情況下使用 -dlink：當擴充功能連結到包含 rdc 編譯物件的靜態庫時，例如 [NVSHMEM 庫](https://developer.nvidia.com/nvshmem)。

備註：需要使用 Ninja 來建置具有 RDC 連結的 CUDA 擴充功能。

範例

>>> CUDAExtension(
...        name='cuda_extension',
...        sources=['extension.cpp', 'extension_kernel.cu'],
...        dlink=True,
...        dlink_libraries=["dlink_lib"],
...        extra_compile_args={'cxx': ['-g'],
...                            'nvcc': ['-O2', '-rdc=true']})

torch.utils.cpp_extension.BuildExtension(*args, **kwargs)

一個自訂的 setuptools 建置擴充功能。

這個 setuptools.build_ext 子類別負責傳遞所需的最小編譯器旗標（例如 -std=c++17）以及混合的 C++/CUDA 編譯（以及對 CUDA 檔案的支援）。

使用 BuildExtension 時，允許為 extra_compile_args 提供一個字典（而不是通常的清單），該字典從語言（cxx 或 nvcc）映射到要提供給編譯器的其他編譯器旗標清單。這使得在混合編譯期間可以將不同的旗標提供給 C++ 和 CUDA 編譯器。

use_ninja (布林值)：如果 use_ninja 為 True（預設值），則我們嘗試使用 Ninja 後端進行建置。與標準的 setuptools.build_ext 相比，Ninja 大大加快了編譯速度。如果 Ninja 不可用，則回退到標準的 distutils 後端。

備註

根據預設，Ninja 後端使用 #CPUS + 2 個工作線程來建置擴充功能。這在某些系統上可能會耗盡太多資源。您可以透過將 MAX_JOBS 環境變數設定為非負數來控制工作線程的數量。

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, is_standalone=False, keep_intermediates=True)

即時 (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 來傳遞其他 include 目錄。

提供具有混合編譯的 CUDA 支援。只需將 CUDA 原始碼檔案（.cu 或 .cuh）與其他原始碼一起傳遞即可。將會偵測到此類檔案，並使用 nvcc 而不是 C++ 編譯器進行編譯。這包括將 CUDA lib64 目錄作為程式庫目錄傳遞，以及連結 cudart。您可以透過 extra_cuda_cflags 將其他旗標傳遞給 nvcc，就像使用 extra_cflags 來傳遞給 C++ 一樣。我們使用各種啟發式方法來尋找 CUDA 安裝目錄，這些方法通常都很好用。如果沒有，設定 CUDA_HOME 環境變數是最安全的選擇。

參數

• name – 要建置的擴充功能的名稱。這必須與 pybind11 模組的名稱相同！

• sources (Union[str, List[str]]) – C++ 原始碼檔案的相對或絕對路徑清單。

• extra_cflags – 要轉發給建置的編譯器旗標的選用清單。

• extra_cuda_cflags – 要在建置 CUDA 原始碼時轉發給 nvcc 的編譯器旗標的選用清單。

• extra_ldflags – 要轉發給建置的連結器旗標的選用清單。

• extra_include_paths – 要轉發給建置的 include 目錄的選用清單。

• build_directory – 要用作建置工作區的路徑（選用）。

• verbose – 如果為 True，則開啟載入步驟的詳細記錄。

• with_cuda (Optional[bool]) – 決定是否將 CUDA 標頭和程式庫新增到建置中。如果設定為 None（預設值），則會根據 sources 中是否存在 .cu 或 .cuh 自動決定此值。將其設定為 True` 可以強制包含 CUDA 標頭和程式庫。

• is_python_module – 如果為 True（預設值），則將產生的共用程式庫作為 Python 模組匯入。如果為 False，則行為取決於 is_standalone。

• is_standalone – 如果為 False（預設值），則將建構的擴充功能作為純動態程式庫載入到程序中。如果為 True，則建置獨立的可執行檔。

傳回值

將載入的 PyTorch 擴充功能作為 Python 模組傳回。

如果 is_python_module 為 False 且 is_standalone 為 False

不傳回任何值。（共用程式庫會作為副作用載入到程序中。）

如果 is_standalone 為 True。

傳回可執行檔的路徑。（在 Windows 上，TORCH_LIB_PATH 會作為副作用新增到 PATH 環境變數中。）

傳回類型

如果 is_python_module 為 True

範例

>>> 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, keep_intermediates=True, use_pch=False)

從字串來源即時 (JIT) 載入 PyTorch C++ 擴充功能。

這個函式的行為與 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 include 為開頭。 .cpp 和 .cu 檔案會單獨編譯，但最終會連結到單個程式庫中。請注意，不會為 cuda_sources 中的函式本身產生任何繫結。若要繫結到 CUDA 核心，您必須建立一個呼叫它的 C++ 函式，並在其中一個 cpp_sources 中宣告或定義這個 C++ 函式（並將其名稱包含在 functions 中）。

請參閱 load() 以取得以下省略的引數的描述。

參數

• cpp_sources - 包含 C++ 原始碼的字串或字串清單。

• cuda_sources - 包含 CUDA 原始碼的字串或字串清單。

• functions - 要為其產生函式繫結的函式名稱清單。如果給定字典，則它應該將函式名稱映射到文件字串（否則只是函式名稱）。

• with_cuda - 決定是否將 CUDA 標頭和程式庫新增到建置中。如果設定為 None（預設值），則會根據是否提供 cuda_sources 自動確定此值。將其設定為 True 可強制包含 CUDA 標頭和程式庫。

• with_pytorch_error_handling - 決定 pytorch 錯誤和警告宏是由 pytorch 還是 pybind 處理。為此，每個函式 foo 都會透過中介函式 _safe_foo 呼叫。這種重新導向可能會在 cpp 的少數情況下造成問題。當此重新導向造成問題時，應將此旗標設定為 False。

範例

>>> 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'])

備註

根據預設，Ninja 後端使用 #CPUS + 2 個工作線程來建置擴充功能。這在某些系統上可能會耗盡太多資源。您可以透過將 MAX_JOBS 環境變數設定為非負數來控制工作線程的數量。

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

取得建置 C++ 或 CUDA 擴充功能所需的 include 路徑。

參數

cuda (bool) - 如果為 True，則包含 CUDA 特定的 include 路徑。

傳回值

include 路徑字串的清單。

傳回類型

List[str]

torch.utils.cpp_extension.get_compiler_abi_compatibility_and_version(compiler)

判斷給定的編譯器是否與 PyTorch ABI 相容，以及其版本。

參數

compiler (str) - 要檢查的編譯器可執行檔名稱（例如 g++）。必須可以在 Shell 程序中執行。

傳回值

一個元組，其中包含一個布林值，用於定義編譯器是否（可能）與 PyTorch ABI 不相容，以及一個包含以點分隔的編譯器版本的 TorchVersion 字串。

傳回類型

Tuple[bool, TorchVersion]

torch.utils.cpp_extension.verify_ninja_availability()

如果系統上沒有 ninja 建置系統，則引發 RuntimeError，否則不執行任何操作。

torch.utils.cpp_extension.is_ninja_available()

如果系統上有 ninja 建置系統，則傳回 True，否則傳回 False。