CUDAExtension

在使用pytorch搭建深度学习模型的时候，用到了CUDAExtension；做一下笔记，防止遗忘

• 样例如下

# setup.py
# 编译支持 CPU/GPU 的算子
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension

setup(
    name='roipool3d',
    ext_modules=[
        CUDAExtension(
        			name='roipool3d_cuda', 
        			source=['src/roipool3d.cpp','src/roipool3d_kernel.cu'],
        			extra_compile_args={'cxx': ['-g'],'nvcc': ['-O2']})],
    cmdclass={'build_ext': BuildExtension})

此接口用于配置自定义 OP 的源文件信息，编译生成同时支持 CPU 和 GPU 设备上执行的算子。若要编译仅支持 CPU 设备的算子，请使用 CppExtension 。

此接口是对 Python 内建库 setuptools.Extension 的进一步封装

搭配 setup 接口使用，编译生成的动态库名称与 setup 接口中的 name 一致。

简单来说就是把cpp文件的功能，搞成一个动态库。动态库的名字在setup的接口中自行设置；然后让python可以调用；