torch.utils.cpp_extension#

創建於：2018年3月7日 | 最後更新於：2025年2月16日

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 找到。

警告

PyTorch 的 Python API（如 libtorch_python 中提供）不能使用 py_limited_api=True 標誌進行構建。當傳遞此標誌時，使用者有責任在其庫中不使用 libtorch_python (尤其是 pytorch/python 繫結) 中的 API，而只使用 libtorch (aten 物件、運算子和分發器) 中的 API。例如，要從 Python 訪問自定義操作，庫應透過分發器註冊這些操作。

與 CPython setuptools 不同，CPython setuptools 在 setup 中將 py_limited_api 指定為 "bdist_wheel" 命令的選項時，不會將 -DPy_LIMITED_API 定義為編譯標誌，而 PyTorch 會！我們將指定 -DPy_LIMITED_API=min_supported_cpython 以最好地強制執行一致性、安全性和合理性，從而鼓勵最佳實踐。要針對不同的版本，請將 min_supported_cpython 設定為所選 CPython 版本的十六進位制程式碼。

示例

>>> 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_args=['-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 的包含路徑、庫路徑和執行時庫。

所有引數都會轉發到 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_args=['-Wl,--no-as-needed', '-lcuda'])
...     ],
...     cmdclass={
...         'build_ext': BuildExtension
...     })

計算能力

預設情況下，擴充套件將編譯為在構建過程中可見的所有顯示卡上執行，以及 PTX。如果將來安裝了新顯示卡，則可能需要重新編譯擴充套件。如果可見顯示卡的計算能力 (CC) 比 nvcc 可構建完全編譯的二進位制檔案的最新版本新，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 透過執行時編譯為較新的 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 錯誤開放問題：pytorch/pytorch#69460 完整的解決方法程式碼示例：facebookresearch/pytorch3d

可重定位裝置程式碼連結

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

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

注意：使用 RDC 連結構建 CUDA 擴充套件需要 Ninja。

示例

>>> 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.SyclExtension(name, sources, *args, **kwargs)[原始碼]#

建立一個用於 SYCL/C++ 的 setuptools.Extension。

方便的方法，用於建立一個具有最少（但通常足夠）引數的 setuptools.Extension，以構建 SYCL/C++ 擴充套件。

所有引數都會轉發到 setuptools.Extension 建構函式。

警告

示例

>>> from torch.utils.cpp_extension import BuildExtension, SyclExtension
>>> setup(
...     name='xpu_extension',
...     ext_modules=[
...     SyclExtension(
...                 name='xpu_extension',
...                 sources=['extension.cpp', 'extension_kernel.cpp'],
...                 extra_compile_args={'cxx': ['-g', '-std=c++20', '-fPIC']})
...     ],
...     cmdclass={
...         'build_ext': BuildExtension
...     })

預設情況下，擴充套件將編譯為在構建過程中可見的所有裝置的架構上執行。如果將來安裝了新裝置，則可能需要重新編譯擴充套件。您可以使用 TORCH_XPU_ARCH_LIST 覆蓋預設行為，以明確指定您希望擴充套件支援的裝置架構。

TORCH_XPU_ARCH_LIST="pvc,xe-lpg" python build_my_extension.py

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

注意：構建 SyclExtension 需要 Ninja。

torch.utils.cpp_extension.BuildExtension(*args, **kwargs)[原始碼]#

一個自定義的 setuptools 構建擴充套件。

這個 setuptools.build_ext 子類負責傳遞最低要求的編譯器標誌（例如 -std=c++17）以及混合 C++/CUDA/SYCL 編譯（以及對 CUDA/SYCL 檔案的一般支援）。

當使用 BuildExtension 時，允許為 extra_compile_args 提供一個字典（而不是通常的列表），該字典將語言/編譯器（唯一預期的值是 cxx、nvcc 或 sycl）對映到要傳遞給編譯器的附加編譯器標誌列表。這使得在混合編譯期間可以為 C++、CUDA 和 SYCL 編譯器提供不同的標誌。

use_ninja (bool)：如果 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_sycl_cflags=None, extra_ldflags=None, extra_include_paths=None, build_directory=None, verbose=False, with_cuda=None, with_sycl=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 來傳遞其他包含目錄。

提供了 CUDA 支援和混合編譯。只需將 CUDA 原始檔（.cu 或 .cuh）與其它原始檔一起傳遞。這些檔案將被檢測到並使用 nvcc 而不是 C++ 編譯器進行編譯。這包括將 CUDA lib64 目錄作為庫目錄傳遞，並連結 cudart。您可以透過 extra_cuda_cflags 向 nvcc 傳遞其他標誌，就像對 C++ 使用 extra_cflags 一樣。使用了各種查詢 CUDA 安裝目錄的啟發式方法，通常都能正常工作。如果不行，設定 CUDA_HOME 環境變數是最安全的選擇。

提供了 SYCL 支援和混合編譯。只需將 SYCL 原始檔（.sycl）與其它原始檔一起傳遞。這些檔案將被檢測到並使用 SYCL 編譯器（例如 Intel DPC++ Compiler）而不是 C++ 編譯器進行編譯。您可以將其他標誌透過 extra_sycl_cflags 傳遞給 SYCL 編譯器，就像對 C++ 使用 extra_cflags 一樣。SYCL 編譯器預計可以透過系統 PATH 環境變數找到。

引數

name – 要構建的擴充套件的名稱。這必須與 pybind11 模組的名稱相同！
sources（Union[str, list[str]]）– C++ 原始檔的相對或絕對路徑列表。
extra_cflags – 要轉發給構建的可選編譯器標誌列表。
extra_cuda_cflags – 在構建 CUDA 原始檔時要轉發給 nvcc 的可選編譯器標誌列表。
extra_sycl_cflags – 在構建 SYCL 原始檔時要轉發給 SYCL 編譯器的可選編譯器標誌列表。
extra_ldflags – 要轉發給構建的可選連結器標誌列表。
extra_include_paths – 要轉發給構建的可選包含目錄列表。
build_directory – 用於作為構建工作區的可選路徑。
verbose – 如果為 True，則開啟載入步驟的詳細日誌記錄。
with_cuda（Optional[bool]）– 確定是否將 CUDA 標頭檔案和庫新增到構建中。如果設定為 None（預設），此值將根據 sources 中是否存在 .cu 或 .cuh 自動確定。將其設定為 True` 以強制包含 CUDA 標頭檔案和庫。
with_sycl（Optional[bool]）– 確定是否將 SYCL 標頭檔案和庫新增到構建中。如果設定為 None（預設），此值將根據 sources 中是否存在 .sycl 自動確定。將其設定為 True` 以強制包含 SYCL 標頭檔案和庫。
is_python_module – 如果為 True（預設），則將生成的共享庫匯入為 Python 模組。如果為 False，則行為取決於 is_standalone。
is_standalone – 如果為 False（預設），則將構造的擴充套件載入到程序中作為普通的動態庫。如果為 True，則構建一個獨立的 C++ 可執行檔案。

返回

返回載入的 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, sycl_sources=None, functions=None, extra_cflags=None, extra_cuda_cflags=None, extra_sycl_cflags=None, extra_ldflags=None, extra_include_paths=None, build_directory=None, verbose=False, with_cuda=None, with_sycl=None, is_python_module=True, with_pytorch_error_handling=True, keep_intermediates=True, use_pch=False, no_implicit_headers=False)[原始碼]#

從字串源即時（JIT）載入 PyTorch C++ 擴充套件。

此函式與 load() 的行為完全相同，但它將源作為字串而不是檔名接收。這些字串將被寫入構建目錄中的檔案，之後 load_inline() 的行為與 load() 相同。

請參閱 tests 以獲取使用此函式的良好示例。

原始檔可以省略典型非內聯 C++ 擴充套件的兩個必需部分：必要的標頭檔案包含，以及（pybind11）繫結程式碼。更具體地說，傳遞給 cpp_sources 的字串首先被連線成一個單獨的 .cpp 檔案。然後，在此檔案前新增 #include <torch/extension.h>。

此外，如果提供了 functions 引數，則會自動為每個指定的函式生成繫結。 functions 可以是函式名列表，也可以是從函式名對映到文件字串的字典。如果給出了列表，則每個函式的名稱將用作其文件字串。

cuda_sources 中的原始檔被連線成一個單獨的 .cu 檔案，並在前面新增 torch/types.h、cuda.h 和 cuda_runtime.h 包含。 .cpp 和 .cu 檔案分別編譯，但最終連結成一個庫。請注意，cuda_sources 中的函式本身不會生成繫結。要繫結到 CUDA 核心，您必須建立一個呼叫它的 C++ 函式，並在 cpp_sources 中的一個檔案中宣告或定義此 C++ 函式（並將其名稱包含在 functions 中）。

sycl_sources 中的原始檔被連線成一個單獨的 .sycl 檔案，並在前面新增 torch/types.h、sycl/sycl.hpp 包含。 .cpp 和 .sycl 檔案分別編譯，但最終連結成一個庫。請注意，sycl_sources 中的函式本身不會生成繫結。要繫結到 SYCL 核心，您必須建立一個呼叫它的 C++ 函式，並在 cpp_sources 中的一個檔案中宣告或定義此 C++ 函式（並將其名稱包含在 functions 中）。

有關下面省略的引數的說明，請參閱 load()。

引數

cpp_sources – 包含 C++ 原始碼的字串或字串列表。
cuda_sources – 包含 CUDA 原始碼的字串或字串列表。
sycl_sources – 包含 SYCL 原始碼的字串或字串列表。
functions – 要為其生成函式繫結的函式名稱列表。如果提供了字典，則應將函式名稱對映到文件字串（否則為函式名稱本身）。
with_cuda – 確定是否將 CUDA 標頭檔案和庫新增到構建中。如果設定為 None（預設），此值將根據是否提供了 cuda_sources 自動確定。將其設定為 True 以強制包含 CUDA 標頭檔案和庫。
with_sycl – 確定是否將 SYCL 標頭檔案和庫新增到構建中。如果設定為 None（預設），此值將根據是否提供了 sycl_sources 自動確定。將其設定為 True 以強制包含 SYCL 標頭檔案和庫。
with_pytorch_error_handling – 確定是讓 pytorch 宏（錯誤和警告）由 pytorch 而不是 pybind 進行處理。為此，每個函式 foo 都透過一箇中間函式 _safe_foo 呼叫。這種重定向在某些晦澀的 C++ 情況中可能會引起問題。當此重定向引起問題時，應將此標誌設定為 False。
no_implicit_headers – 如果為 True，則跳過自動新增標頭檔案，特別是 #include <torch/extension.h> 和 #include <torch/types.h> 行。當您已經在原始碼中包含了必要的標頭檔案時，使用此選項可以提高冷啟動時間。預設值：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'])

注意

由於 load_inline 將即時編譯原始碼，請確保您在執行時環境中安裝了正確的工具鏈。例如，在載入 C++ 時，請確保 C++ 編譯器可用。如果您正在載入 CUDA 擴充套件，您還需要另外安裝相應的 CUDA 工具包（nvcc 和您的程式碼可能具有的任何其他依賴項）。編譯工具鏈在安裝 torch 時不包含，並且必須另外安裝。

在編譯過程中，預設情況下，Ninja 後端使用 #CPUS + 2 個工作程序來構建擴充套件。這可能會在某些系統上佔用過多資源。可以透過將 MAX_JOBS 環境變數設定為非負數來控制工作程序的數量。

torch.utils.cpp_extension.include_paths(device_type='cpu')[原始碼]#

獲取構建 C++ 或 CUDA 或 SYCL 擴充套件所需的包含路徑。

引數: device_type（str）– 預設為“cpu”。
返回: 包含路徑字串列表。
返回型別: list[str]

torch.utils.cpp_extension.get_compiler_abi_compatibility_and_version(compiler)[原始碼]#

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

引數: compiler（str）– 要檢查的編譯器可執行檔名（例如 g++）。必須在 shell 程序中可執行。
返回: 一個元組，其中包含一個布林值，該布林值定義了編譯器是否（可能）與 PyTorch ABI 不相容，後面跟著一個 TorchVersion 字串，其中包含用點分隔的編譯器版本。
返回型別: tuple[bool, torch.torch_version.TorchVersion]

torch.utils.cpp_extension.verify_ninja_availability()[原始碼]#: 如果系統中 ninja 構建系統不可用，則引發 RuntimeError，否則不做任何事情。

torch.utils.cpp_extension.is_ninja_available()[原始碼]#: 如果系統中 ninja 構建系統可用，則返回 True，否則返回 False。

torch.utils.cpp_extension#

文件

教程

資源