CUDA 代码 Python 暴露实战笔记 + C++示例代码（第九章重点）

#ifndef VECTOR_ADD_H
#define VECTOR_ADD_H

// extern "C"：避免C++名称修饰，确保Ctypes/Python C API能找到函数
extern "C" {
    /**
     * @brief CUDA加速向量加法（主机端接口）
     * @param a 输入向量a（主机内存，Python分配）
     * @param b 输入向量b（主机内存，Python分配）
     * @param c 输出向量c（主机内存，Python分配）
     * @param N 向量长度
     */
    void vectorAdd(int* a, int* b, int* c, int N);
}

#endif

#include "../include/vector_add.h"
#include <cuda_runtime.h>

// 错误检查宏（简化CUDA API错误处理）
#define CHECK_CUDA_ERR(err) \
    if (err != cudaSuccess) { \
        fprintf(stderr, "CUDA错误：%s（行号：%d）\n", cudaGetErrorString(err), __LINE__); \
        exit(1); \
    }

// CUDA核函数：向量加法（设备端执行）
__global__ void vectorAddKernel(int* d_a, int* d_b, int* d_c, int N) {
    int idx = threadIdx.x + blockIdx.x * blockDim.x;
    if (idx < N) {
        d_c[idx] = a[idx] + b[idx];  // 修复第九章示例中的故意bug
    }
}

// 主机端接口实现
void vectorAdd(int* a, int* b, int* c, int N) {
    int *d_a, *d_b, *d_c;
    size_t data_size = N * sizeof(int);

    // 1. 分配GPU内存
    CHECK_CUDA_ERR(cudaMalloc((void**)&d_a, data_size));
    CHECK_CUDA_ERR(cudaMalloc((void**)&d_b, data_size));
    CHECK_CUDA_ERR(cudaMalloc((void**)&d_c, data_size));

    // 2. 主机→设备数据拷贝（同步传输）
    CHECK_CUDA_ERR(cudaMemcpy(d_a, a, data_size, cudaMemcpyHostToDevice));
    CHECK_CUDA_ERR(cudaMemcpy(d_b, b, data_size, cudaMemcpyHostToDevice));

    // 3. 启动核函数（256线程/块，动态计算网格大小）
    int threadsPerBlock = 256;
    int blocksPerGrid = (N + threadsPerBlock - 1) / threadsPerBlock;
    vectorAddKernel<<<blocksPerGrid, threadsPerBlock>>>(d_a, d_b, d_c, N);
    CHECK_CUDA_ERR(cudaGetLastError());  // 检查核函数启动错误

    // 4. 设备→主机数据拷贝（同步传输）
    CHECK_CUDA_ERR(cudaMemcpy(c, d_c, data_size, cudaMemcpyDeviceToHost));

    // 5. 释放GPU内存
    CHECK_CUDA_ERR(cudaFree(d_a));
    CHECK_CUDA_ERR(cudaFree(d_b));
    CHECK_CUDA_ERR(cudaFree(d_c));
}

cmake_minimum_required(VERSION 3.18)
project(cuda_vector_add LANGUAGES CXX CUDA)

# 设置C++和CUDA标准
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_CUDA_STANDARD 11)

# 查找CUDA Toolkit
find_package(CUDAToolkit REQUIRED)

# 创建共享库（SHARED表示动态库）
add_library(vector_add SHARED src/vector_add.cu)

# 启用CUDA分离编译（必须，否则共享库无法正确链接CUDA核函数）
set_target_properties(vector_add PROPERTIES CUDA_SEPARABLE_COMPILATION ON)

# 链接CUDA库
target_link_libraries(vector_add PRIVATE CUDA::cudart)

# 指定头文件目录（让外部能找到vector_add.h）
target_include_directories(vector_add PUBLIC include)

# 设置输出目录（共享库生成到build/lib）
set_target_properties(vector_add PROPERTIES
    LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/lib
)

# 创建构建目录
mkdir -p build && cd build

# 配置CMake（指定CUDA Toolkit路径，可选）
cmake .. -DCMAKE_BUILD_TYPE=Release

# 编译（-j4表示4线程并行编译）
make -j4

import ctypes
import numpy as np

def main():
    # 1. 加载CUDA共享库
    # 注意：路径需指向编译生成的libvector_add.so
    lib = ctypes.CDLL("../build/lib/libvector_add.so")

    # 2. 声明函数签名（必须匹配C接口，否则会内存错误）
    # void vectorAdd(int* a, int* b, int* c, int N)
    lib.vectorAdd.argtypes = [
        ctypes.POINTER(ctypes.c_int),  # a
        ctypes.POINTER(ctypes.c_int),  # b
        ctypes.POINTER(ctypes.c_int),  # c
        ctypes.c_int                  # N
    ]
    lib.vectorAdd.restype = None  # 无返回值

    # 3. 生成测试数据（NumPy数组，int32类型，与C接口匹配）
    N = 10_000_000  # 1000万元素（大数组测试零拷贝性能）
    a = np.random.randint(0, N, size=N, dtype=np.int32)
    b = np.random.randint(N, 2*N, size=N, dtype=np.int32)
    c = np.zeros(N, dtype=np.int32)  # 输出数组（Python分配内存）

    # 4. 调用CUDA函数（传递NumPy数组内部指针，零拷贝）
    lib.vectorAdd(
        a.ctypes.data_as(ctypes.POINTER(ctypes.c_int)),
        b.ctypes.data_as(ctypes.POINTER(ctypes.c_int)),
        c.ctypes.data_as(ctypes.POINTER(ctypes.c_int)),
        ctypes.c_int(N)
    )

    # 5. 验证结果（前10个元素）
    valid = np.allclose(a + b, c)
    print(f"结果验证：{'正确' if valid else '错误'}")
    print(f"前10个元素对比：")
    print(f"a + b: {a[:10] + b[:10]}")
    print(f"c    : {c[:10]}")

if __name__ = "__main__":
    main()

# 确保已编译共享库，然后运行Python脚本
python3 ctypes_demo.py

#include <Python.h>
#include <dlfcn.h>
#include "../include/vector_add.h"

// 函数指针：指向CUDA共享库中的vectorAdd
typedef void (*VectorAddFunc)(int* a, int* b, int* c, int N);
static VectorAddFunc vectorAdd = NULL;

/**
 * @brief Python包装函数：vectorAdd（处理Python列表）
 * @param self 模块实例（未使用）
 * @param args 传入的Python参数（(list_a, list_b, N)）
 * @return Python列表（结果）
 */
static PyObject* py_vector_add(PyObject* self, PyObject* args) {
    PyObject *py_a, *py_b;
    int N;

    // 1. 解析Python参数："OOi"表示两个对象（列表）+ 一个整数
    if (!PyArg_ParseTuple(args, "OOi", &py_a, &py_b, &N)) {
        return NULL;  // 解析失败返回NULL，触发Python异常
    }

    // 2. 验证参数类型（必须是列表）
    if (!PyList_Check(py_a) || !PyList_Check(py_b)) {
        PyErr_SetString(PyExc_TypeError, "参数必须是列表");
        return NULL;
    }

    // 3. 分配主机内存（C侧）
    int* a = (int*)malloc(N * sizeof(int));
    int* b = (int*)malloc(N * sizeof(int));
    int* c = (int*)malloc(N * sizeof(int));
    if (!a || !b || !c) {
        PyErr_SetString(PyExc_MemoryError, "内存分配失败");
        free(a); free(b); free(c);
        return NULL;
    }

    // 4. Python列表 → C数组（数据拷贝，开销较大）
    for (int i = 0; i < N; ++i) {
        a[i] = (int)PyLong_AsLong(PyList_GetItem(py_a, i));
        b[i] = (int)PyLong_AsLong(PyList_GetItem(py_b, i));
    }

    // 5. 调用CUDA共享库函数
    vectorAdd(a, b, c, N);

    // 6. C数组 → Python列表（数据拷贝）
    PyObject* py_result = PyList_New(N);
    for (int i = 0; i < N; ++i) {
        PyList_SetItem(py_result, i, PyLong_FromLong(c[i]));
    }

    // 7. 释放C侧内存
    free(a);
    free(b);
    free(c);

    return py_result;
}

// 2. 定义模块函数列表
static PyMethodDef VectorAddMethods[] = {
    {
        "vector_add",          // Python中调用的函数名
        py_vector_add,         // 对应的C包装函数
        METH_VARARGS,          // 参数传递方式（元组）
        "CUDA加速向量加法（输入Python列表）"  // 函数描述
    },
    {NULL, NULL, 0, NULL}  // 结束标记
};

// 3. 定义模块结构
static struct PyModuleDef vector_add_module = {
    PyModuleDef_HEAD_INIT,
    "vector_add_wrapper",    // 模块名（Python中import的名称）
    "CUDA向量加法Python扩展",  // 模块描述
    -1,                     // 不维护模块状态（-1表示全局）
    VectorAddMethods        // 模块函数列表
};

// 4. 模块初始化函数（Python导入时调用）
PyMODINIT_FUNC PyInit_vector_add_wrapper(void) {
    // 加载CUDA共享库
    void* handle = dlopen("../build/lib/libvector_add.so", RTLD_LAZY);
    if (!handle) {
        PyErr_SetString(PyExc_ImportError, dlerror());
        return NULL;
    }

    // 获取vectorAdd函数指针
    vectorAdd = (VectorAddFunc)dlsym(handle, "vectorAdd");
    if (!vectorAdd) {
        PyErr_SetString(PyExc_ImportError, dlerror());
        dlclose(handle);
        return NULL;
    }

    // 创建并返回模块
    return PyModule_Create(&vector_add_module);
}

#include <Python.h>
#include <numpy/arrayobject.h>  // NumPy C API
#include <dlfcn.h>
#include "../include/vector_add.h"

// 函数指针：指向CUDA共享库中的vectorAdd
typedef void (*VectorAddFunc)(int* a, int* b, int* c, int N);
static VectorAddFunc vectorAdd = NULL;

/**
 * @brief Python包装函数：vector_add（处理NumPy数组，零拷贝）
 * @param self 模块实例（未使用）
 * @param args 传入的Python参数（(np_a, np_b, np_c, N)）
 * @return None
 */
static PyObject* py_vector_add_np(PyObject* self, PyObject* args) {
    PyArrayObject *np_a, *np_b, *np_c;
    int N;

    // 1. 解析Python参数："OOOi"表示三个NumPy数组 + 一个整数
    if (!PyArg_ParseTuple(args, "OOOi", &np_a, &np_b, &np_c, &N)) {
        return NULL;
    }

    // 2. 验证NumPy数组类型（必须是int32，连续内存）
    if (PyArray_TYPE(np_a) != NPY_INT32 || !PyArray_IS_C_CONTIGUOUS(np_a) ||
        PyArray_TYPE(np_b) != NPY_INT32 || !PyArray_IS_C_CONTIGUOUS(np_b) ||
        PyArray_TYPE(np_c) != NPY_INT32 || !PyArray_IS_C_CONTIGUOUS(np_c)) {
        PyErr_SetString(PyExc_TypeError, "NumPy数组必须是int32类型且连续内存");
        return NULL;
    }

    // 3. 获取NumPy数组内部指针（零拷贝，直接操作Python分配的内存）
    int* a = (int*)PyArray_DATA(np_a);
    int* b = (int*)PyArray_DATA(np_b);
    int* c = (int*)PyArray_DATA(np_c);

    // 4. 调用CUDA共享库函数
    vectorAdd(a, b, c, N);

    // 5. 返回None（结果已写入np_c）
    Py_RETURN_NONE;
}

// 2. 定义模块函数列表
static PyMethodDef VectorAddNpMethods[] = {
    {
        "vector_add",          // Python中调用的函数名
        py_vector_add_np,      // 对应的C包装函数
        METH_VARARGS,          // 参数传递方式（元组）
        "CUDA加速向量加法（输入NumPy数组，零拷贝）"
    },
    {NULL, NULL, 0, NULL}  // 结束标记
};

// 3. 定义模块结构
static struct PyModuleDef vector_add_np_module = {
    PyModuleDef_HEAD_INIT,
    "vector_add_np_wrapper",  // 模块名
    "CUDA向量加法Python扩展（NumPy优化）",
    -1,
    VectorAddNpMethods
};

// 4. 模块初始化函数
PyMODINIT_FUNC PyInit_vector_add_np_wrapper(void) {
    // 加载CUDA共享库
    void* handle = dlopen("../build/lib/libvector_add.so", RTLD_LAZY);
    if (!handle) {
        PyErr_SetString(PyExc_ImportError, dlerror());
        return NULL;
    }

    // 获取vectorAdd函数指针
    vectorAdd = (VectorAddFunc)dlsym(handle, "vectorAdd");
    if (!vectorAdd) {
        PyErr_SetString(PyExc_ImportError, dlerror());
        dlclose(handle);
        return NULL;
    }

    // 初始化NumPy C API（必须调用，否则PyArray_DATA会崩溃）
    import_array();

    // 创建并返回模块
    return PyModule_Create(&vector_add_np_module);
}

from setuptools import setup, Extension
import numpy as np

# 定义普通版本扩展
extension_normal = Extension(
    name="vector_add_wrapper",  # 模块名（import时用）
    sources=["vector_add_wrapper.c"],  # C包装代码
    extra_link_args=["-ldl"],  # 链接dl库（用于dlopen加载CUDA共享库）
    include_dirs=["include"]  # 头文件目录
)

# 定义NumPy优化版本扩展
extension_np = Extension(
    name="vector_add_np_wrapper",  # 模块名
    sources=["vector_add_np_wrapper.c"],
    extra_link_args=["-ldl"],
    include_dirs=["include", np.get_include()]  # 加入NumPy头文件目录
)

# 构建扩展
setup(
    name="cuda_vector_add_extensions",
    version="1.0",
    description="CUDA向量加法Python扩展",
    ext_modules=[extension_normal, extension_np]
)

# 1. 构建扩展（生成.so文件，与Python脚本同目录）
python3 setup.py build_ext --inplace

# 2. Python调用脚本（extension_demo.py）

import numpy as np
import vector_add_wrapper as vaw  # 普通版本
import vector_add_np_wrapper as vaw_np  # NumPy优化版本

def test_normal():
    """测试普通版本（Python列表，数据拷贝）"""
    N = 10_000  # 小数组（大数据组拷贝开销过大）
    a = [np.random.randint(0, N) for _ in range(N)]
    b = [np.random.randint(N, 2*N) for _ in range(N)]
    result = vaw.vector_add(a, b, N)
    valid = all(result[i] = a[i] + b[i] for i in range(N))
    print(f"普通版本（列表）结果验证：{'正确' if valid else '错误'}")

def test_np():
    """测试NumPy版本（零拷贝）"""
    N = 10_000_000  # 1000万元素
    a = np.random.randint(0, N, size=N, dtype=np.int32)
    b = np.random.randint(N, 2*N, size=N, dtype=np.int32)
    c = np.zeros(N, dtype=np.int32)
    vaw_np.vector_add(a, b, c, N)
    valid = np.allclose(a + b, c)
    print(f"NumPy版本（零拷贝）结果验证：{'正确' if valid else '错误'}")
    print(f"前10个元素对比：")
    print(f"a + b: {a[:10] + b[:10]}")
    print(f"c    : {c[:10]}")

if __name__ = "__main__":
    test_normal()
    test_np()