2021-06-21

python 和 cuda 混合编程

主要想法是，结合gcc 和nvcc 把，cuda程序编译成.so 库，然后使用python去调用cuda程序。以下是主要流程：

注意点

• cuda的环境要配好，版本对应好。
• extern “C” void test() 才能导出 test函数。
• -L/usr/local/cuda/lib64 -lcudart 链接好库，才能调用成功。

代码

libtest.cu

#include <stdio.h>
#include "cuda_runtime.h"
#include "cuda.h"

__global__ void hello() {
    int g = blockIdx.x;
	int t = threadIdx.x;
    printf("Hello world from device(%d,%d)\n",g,t);
}

extern "C" void test() {
    hello<<<5, 2>>>();
    printf("Hello world from host\n");
    cudaDeviceSynchronize();
}

编译

nvcc -c libtest.cu -Xcompiler -fPIC  -shared  -o libtest.o
g++ libtest.o -L/usr/local/cuda/lib64 -lcudart -fPIC -shared -o libtest.so

调用

import time
from ctypes import cdll
matcur = cdll.LoadLibrary('libtest.so')
matcur.test()

结果