注意這個和前面的《Python與C語言混合編程：通過distutils或setuptools實現的一個簡單的C擴展》不同，這個是pytorch的擴展，不是python的擴展。
在pytorch的utils中，集成了setuptools模塊。
官方文檔在這裏：https://pytorch.org/docs/master/cpp_extension.html
中文說明在這裏：https://ptorch.com/news/188.html
torch.utils.cpp_extension.CppExtension(name, sources, *args, **kwargs)
創建一個C++的setuptools.Extension。
便捷地創建一個setuptools.Extension具有最小（但通常是足夠）的參數來構建C++擴展的方法。
所有參數都被轉發給setuptools.Extension構造函數。
例
>>> 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++創建一個setuptools.Extension。
創建一個setuptools.Extension用於構建CUDA/C ++擴展的最少參數（但通常是足夠的）的便捷方法。這裏包括CUDA路徑，庫路徑和運行庫。 所有參數都被轉發給setuptools.Extension構造函數。
例
>>> from setuptools import setup
>>> from torch.utils.cpp_extension import BuildExtension, CppExtension
>>> 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（dist，** kw ）
自定義setuptools構建擴展。
setuptools.build_ext子類負責傳遞所需的最小編譯器參數（例如-std=c++11）以及混合的C ++/CUDA編譯（以及一般對CUDA文件的支持）。
當使用BuildExtension時，它將提供一個用於extra_compile_args（不是普通列表）的詞典，通過語言（cxx或cuda）映射到參數列表提供給編譯器。這樣可以在混合編譯期間爲C ++和CUDA編譯器提供不同的參數。
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)
即時加載(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鏈接。您可以將其他參數傳遞給nvcc extra_cuda_cflags，就像使用C ++的extra_cflags一樣。使用了各種原始方法來查找CUDA安裝目錄，通常情況下可以正常運行。如果不可以，最好設置CUDA_HOME環境變量。
- 參數：
- name - 要構建的擴展名。這個必須和
pybind11模塊的名字一樣！ - sources -
C++源文件的相對或絕對路徑列表。 - extra_cflags - 編譯器參數的可選列表，用於轉發到構建。
- extra_cuda_cflags - 編譯器標記的可選列表，在構建
CUDA源時轉發給nvcc。 - extra_ldflags - 鏈接器參數的可選列表，用於轉發到構建。
- extra_include_paths - 轉發到構建的包含目錄的可選列表。
- build_directory - 可選路徑作爲構建區域。
- verbose - 如果爲
True，打開加載步驟的詳細記錄。
- name - 要構建的擴展名。這個必須和
- 返回：
- 加載
PyTorch擴展作爲Python模塊。
- 加載
例
>>> from torch.utils.cpp_extension import load
>>> module = load(
name='extension',
sources=['extension.cpp', 'extension_kernel.cu'],
extra_cflags=['-O2'],
verbose=True)
後面還新添加了一個load_inline,不過沒有中文翻譯
Loads a PyTorch C++ extension just-in-time (JIT) from string sources.
This function behaves exactly like load(), but takes its sources as strings rather than filenames. These strings are stored to files in the build directory, after which the behavior of load_inline() is identical to load().
See the tests for good examples of using this function.
Sources may omit two required parts of a typical non-inline C++ extension: the necessary header includes, as well as the (pybind11) binding code. More precisely, strings passed to cpp_sources are first concatenated into a single .cpp file. This file is then prepended with #include <torch/extension.h>.
Furthermore, if the functions argument is supplied, bindings will be automatically generated for each function specified. functions can either be a list of function names, or a dictionary mapping from function names to docstrings. If a list is given, the name of each function is used as its docstring.
The sources in cuda_sources are concatenated into a separate .cu file and prepended with torch/types.h, cuda.hand cuda_runtime.h includes. The .cpp and .cu files are compiled separately, but ultimately linked into a single library. Note that no bindings are generated for functions in cuda_sources per se. To bind to a CUDA kernel, you must create a C++ function that calls it, and either declare or define this C++ function in one of the cpp_sources (and include its name in functions).
See load() for a description of arguments omitted below.
| Parameters: |
|
|---|
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)
獲取構建C++或CUDA擴展所需的路徑。
- 參數：
cuda- 如果爲True，則包含CUDA特定的包含路徑。 - 返回： 包含路徑字符串的列表。
例如：
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CppExtension
torch.utils.cpp_extension.include_paths(cuda=False)
#
['/usr/local/lib/python3.6/site-packages/torch/lib/include',
'/usr/local/lib/python3.6/site-packages/torch/lib/include/TH',
'/usr/local/lib/python3.6/site-packages/torch/lib/include/THC']
torch.utils.cpp_extension.check_compiler_abi_compatibility(compiler)
驗證給定的編譯器是否與PyTorch ABI兼容。
- 參數：compiler(str) - 要檢查可執行的編譯器文件名(例如g++),必須在
shell進程中可執行。 - 返回：如果編譯器（可能）與
PyTorchABI不兼容，則爲False，否則返回True。
torch.utils.cpp_extension.verify_ninja_availability()
如果可以在ninja上運行則返回True。