Python使用pycuda在GPU上并行处理批量判断素数
借助于扩展库pycuda，可以在Python中访问NVIDIA显卡提供的CUDA并行计算API，使用非常方便。安装pycuda时要求已正确安装合适版本的CUDA和Visual Studio（注意，并不是版本越新越合适，目前2015暂时还不行，最好使用VS2013），然后再使用pip安装pycuda。
下面的代码用来统计100000000之内的所有素数个数。
import time
import pycuda.autoinit
import pycuda.driver as drv
import numpy as np
from pycuda.compiler import SourceModule
#编译C代码进入显卡，并行判断素数
mod = SourceModule('''
__global__ void isPrime(int *dest, int *a, int *b)
{
const int i = threadIdx.x+blockDim.x*blockIdx.x;
int j;
for(j=2;j<b[i];j++)
{
if(a[i]%j == 0)
{
break;
}
}
if(j >= b[i])
{
dest[i] = a[i];
}
}
''')
#定义待测数值范围，和每次处理的数字数量
end = 100000000
size = 1000
#获取函数
isPrime = mod.get_function("isPrime")
result = 0
start = time.time()
#分段处理，每次处理1000个数字
for i in range(end//size):
startN = i * size
a = np.array(range(startN, startN+size)).astype(np.int64)
b = np.array(list(map(lambda x: int(x**0.5)+1, a))).astype(np.int64)
dest = np.zeros_like(a)
isPrime(drv.Out(dest), drv.In(a), drv.In(b),
block=(size,1,1), grid=(2,1))
result += len(set(filter(None, dest)))
print(time.time()-start)
#上面的代码中把1也算上了，这里减去
print(result-1)
测试结果：在4核CPU、640核GPU的笔记本上运行，本文代码为在CPU上运行的类似代码运行速度的8倍左右。