setuptools编译CUDA代码并调用

setuptools编译CUDA代码并调用

  • setuptools简介
  • 编译代码与调用
    • 代码编译
    • 模块调用
  • 参考文献


在这里插入图片描述

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='[email protected]',
    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