setuptools编译CUDA代码并调用

setuptools简介

setuptools是Python的一个库，它被设计用来简化Python包的创建、构建、分发和安装过程。它是distutils的增强版，后者是Python标准库的一部分，用于同样的目的。相比于distutils，setuptools提供了更多的功能，例如自动依赖性下载、入口点（entry points）和“开发模式”安装。

主要特性：

• 易于安装依赖：setuptools 允许在 setup.py 文件中声明项目的依赖，它会自动处理这些依赖的下载和安装。
• 自动发现包目录：使用 find_packages() 函数可以自动查找项目中需要包含的子包和模块。
• 条目点（Entry Points）：允许创建可插拔的应用程序，可以在安装时提供可执行脚本或者为其他包提供插件。
• 扩展包的构建和支持：简化了 C/C++ 扩展的编译过程，支持将这些扩展集成进 Python 包中。
• 自定义构建命令：可以自定义构建、安装和打包流程的各个环节。
• 生成二进制分发包：可以生成轮（wheel）文件，这是一种现代的二进制分发格式，它提供了快速的安装体验。

setuptools的安装方法如下：

pip install setuptools

setuptools通常通过一个setup.py脚本来使用，该脚本位于Python项目的根目录下。在setup.py文件中，开发者可以定义他们的包信息和配置，如包名、版本、描述、依赖关系、包含的文件以及其他元数据。与CMake相类似，在使用setuptools安装算子包时，只需手动编译一次，后续可重复使用，但使用setuptools要比CMake更加简单一些，在C++项目中，则需要编写CMakeLists.txt文件，使用CMake编译。此外还有JIT编译，即just-in-time，也就是即时编译或动态编译，就是说在python代码运行的时候再去编译cpp和cuda文件。

以下是一个基本的setup.py示例：

from setuptools import setup, find_packages

setup(
    name='example_package',  # 包名称
    version='0.1',  # 版本号
    packages=find_packages(),  # 自动发现并包含项目中的所有包
    install_requires=[  # 依赖列表
        'requests>=2.25.1',  
        'numpy>=1.19.5',
    ],
    entry_points={  # 可执行条目点定义
        'console_scripts': [
            'example-script=example_package:main',
        ],
    },
    # 更多项目元数据
    author='Author Name',
    author_email='author@example.com',
    description='A simple example package',
    long_description=open('README.md').read(),
    long_description_content_type='text/markdown',
    url='https://github.com/username/example_package',
    classifiers=[
        'Programming Language :: Python :: 3',
        'License :: OSI Approved :: MIT License',
        'Operating System :: OS Independent',
    ],
    python_requires='>=3.6',  # Python 版本要求
)

使用setuptools的一些主要命令和功能包括：

• python setup.py build: 构建项目。
• python setup.py install: 安装项目到Python环境中。
• python setup.py sdist: 创建源代码分发包（通常是.tar.gz）。
• python setup.py bdist_wheel: 创建一个wheel包（.whl），它是一种更现代的Python分发格式。
• python setup.py develop: 以“开发模式”安装包，这样对代码的更改无需重新安装即可反映。
• python setup.py test: 运行项目的测试套件。

编译代码与调用

include文件夹用来放cuda算子的头文件（.h文件），里面是cuda算子的定义。kernel文件夹放cuda算子的具体实现（.cu文件）和cpp torch的接口封装（.cpp文件）。

目录结构：

├── include
│   └── add2.h # cuda算子的头文件
├── kernel
│   ├── add2_kernel.cu # cuda算子的具体实现
│   └── add2.cpp # cuda算子的cpp torch封装
├── CMakeLists.txt
├── LICENSE
├── README.md
├── setup.py

各文件代码如下：
add2.h

void launch_add2(float *c,
                 const float *a,
                 const float *b,
                 int n);

add2.cu

__global__ void add2_kernel(float* c,
                            const float* a,
                            const float* b,
                            int n) {
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; \
            i < n; i += gridDim.x * blockDim.x) {
        c[i] = a[i] + b[i];
    }
}

void launch_add2(float* c,
                 const float* a,
                 const float* b,
                 int n) {
    dim3 grid((n + 1023) / 1024);
    dim3 block(1024);
    add2_kernel<<<grid, block>>>(c, a, b, n);
}

add2.cpp

#include <torch/extension.h>
#include "add2.h"

void torch_launch_add2(torch::Tensor &c,
                       const torch::Tensor &a,
                       const torch::Tensor &b,
                       int n) {
    launch_add2((float *)c.data_ptr(),
                (const float *)a.data_ptr(),
                (const float *)b.data_ptr(),
                n);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("torch_launch_add2",
          &torch_launch_add2,
          "add2 kernel warpper");
}

代码编译

setup.py

from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension

setup(
    name="add2",
    include_dirs=["include"],
    ext_modules=[
        CUDAExtension(
            "add2",
            ["kernel/add2.cpp", "kernel/add2_kernel.cu"],
        )
    ],
    cmdclass={
        "build_ext": BuildExtension
    }
)

运行以下命令，进行代码编译：

python3 setup.py install

如果执行正常的话，可以看到两条编译命令的：

[1/2] nvcc -c add2_kernel.cu -o add2_kernel.o
[2/2] c++ -c add2.cpp -o add2.o

然后会执行第三条：

x86_64-linux-gnu-g++ -shared add2.o add2_kernel.o -o add2.cpython-37m-x86_64-linux-gnu.so

模块调用

编译完成后会生成了一个动态链接库，python端我们不需要加载这个动态链接库，因为setuptools已经帮我们把cuda算子调用的接口注册到python模块里了，直接import即可：

import torch
import add2
add2.torch_launch_add2(c, a, b, n)

注意，.cpp和.cu文件名不要相同，也最好不要取容易与python自带库重复的名字。此外要先import torch，然后再import add2，不然也会报错。

参考文献

1. https://zhuanlan.zhihu.com/p/545221832
2. https://zhuanlan.zhihu.com/p/358778742