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

警告

PyTorch python API（由 libtorch_python 提供）不能使用 py_limited_api=True 標誌構建。傳遞此標誌時，使用者有責任在其庫中不使用來自 libtorch_python（特別是 pytorch/python 繫結）的 API，而僅使用來自 libtorch（aten 物件、運算子和排程器）的 API。例如，要從 python 訪問自定義運算子，庫應該透過排程器註冊這些運算子。

與 CPython setuptools 不同，後者在 setup 中的 “bdist_wheel” 命令指定 py_limited_api 選項時不會將 -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 include 路徑、library 路徑和 runtime library。

所有引數都會轉發給 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 不同，後者在 setup 中的 “bdist_wheel” 命令指定 py_limited_api 選項時不會將 -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, 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
...     })

計算能力 (Compute capabilities)

預設情況下，擴充套件將被編譯以在構建過程中可見的卡的所有架構（archs）上執行，外加 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 bug 的開放問題：https://github.com/pytorch/pytorch/issues/69460 完整的 workaround 程式碼示例：https://github.com/facebookresearch/pytorch3d/commit/cb170ac024a949f1f9614ffe6af1c38d972f7d48

可重定位裝置程式碼連結 (Relocatable device code linking)

如果您想跨編譯單元（跨目標檔案）引用裝置符號，則需要使用 relocatable device code （-rdc=true 或 -dc）構建目標檔案。此規則的一個例外是“dynamic parallelism”（巢狀核心啟動），目前已不太常用。Relocatable device code 最佳化程度較低，因此僅應在需要它的目標檔案上使用。在裝置程式碼編譯步驟和 dlink 步驟中使用 -dlto (Device Link Time Optimization) 有助於減少 -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 建構函式。

警告

PyTorch python API（由 libtorch_python 提供）不能使用 py_limited_api=True 標誌構建。傳遞此標誌時，使用者有責任在其庫中不使用來自 libtorch_python（特別是 pytorch/python 繫結）的 API，而僅使用來自 libtorch（aten 物件、運算子和排程器）的 API。例如，要從 python 訪問自定義運算子，庫應該透過排程器註冊這些運算子。

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

示例

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

預設情況下，擴充套件將被編譯以在構建過程中可見的卡的所有架構（archs）上執行。如果之後安裝了新卡，擴充套件可能需要重新編譯。您可以使用 TORCH_XPU_ARCH_LIST 明確指定您希望擴充套件支援哪些裝置架構（device architectures），從而覆蓋預設行為

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 個 worker 來構建擴充套件。這可能會在某些系統上佔用過多資源。可以透過將 MAX_JOBS 環境變數設定為非負數來控制 worker 的數量。

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

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

提供對 SYCL 的混合編譯支援。只需傳遞 SYCL 原始檔（.sycl）以及其他原始檔。這些檔案將被檢測到並由 SYCL 編譯器（例如 Intel DPC++ Compiler）而非 C++ 編譯器編譯。您可以透過 extra_sycl_cflags 向 SYCL 編譯器傳遞附加標誌，就像使用 extra_cflags 傳遞給 C++ 一樣。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 – 可選的 include 目錄列表，用於轉發到構建過程。
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，則構建獨立可執行檔案。

返回值

返回載入的 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)[source][source]¶

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

此函式與 load() 函式的行為完全一致，但它接收的是字串形式的源而不是檔名。這些字串被儲存到構建目錄中的檔案裡，此後 load_inline() 的行為與 load() 函式相同。

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

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

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

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。

示例

>>> 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')[source][source]¶

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

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

torch.utils.cpp_extension.get_compiler_abi_compatibility_and_version(compiler)[source][source]¶

確定給定的編譯器是否與 PyTorch ABI 相容，並獲取其版本。

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

torch.utils.cpp_extension.verify_ninja_availability()[source][source]¶: 如果系統上沒有 ninja 構建系統，則丟擲 RuntimeError，否則不執行任何操作。

torch.utils.cpp_extension.is_ninja_available()[source][source]¶: 如果系統上可以使用 ninja 構建系統，則返回 True，否則返回 False。

torch.utils.cpp_extension¶

文件

教程

資源