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 中的 API (特別是 pytorch/python 繫結)，而僅使用 libtorch 中的 API (aten 物件、運算子和分派器)。 例如，若要從 python 授權存取自定義運算，函式庫應通過分派器註冊運算。

範例

>>> 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 路徑、函式庫路徑和執行階段函式庫。

所有引數都會轉發到 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 中的 API (特別是 pytorch/python 繫結)，而僅使用 libtorch 中的 API (aten 物件、運算子和分派器)。 例如，若要從 python 授權存取自定義運算，函式庫應通過分派器註冊運算。

範例

>>> 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 錯誤的公開問題：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 編譯物件的靜態 lib 時，例如 [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 (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_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，就像使用 C++ 的 extra_cflags 一樣。使用各種尋找 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 包含作為前置詞。 .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'])

注意

由於 load_inline 會即時編譯原始碼，請確保您在執行階段安裝了正確的工具鏈。 例如，在載入 C++ 時，請確保 C++ 編譯器可用。 如果您要載入 CUDA 擴充，您還需要額外安裝相應的 CUDA 工具包（nvcc 以及您的程式碼所具有的任何其他依賴項）。 當您安裝 torch 時，不會包含編譯工具鏈，必須額外安裝。

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

torch.utils.cpp_extension.include_paths(device_type='cpu')

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

參數

device_type (str) – 預設為 "cpu"。

傳回

include 路徑字串的列表。

回傳型別

List[str]

torch.utils.cpp_extension.get_compiler_abi_compatibility_and_version(compiler)

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

參數

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

傳回

一個 tuple，其中包含一個布林值，用於定義編譯器是否（可能）與 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。