Python使用pyopencl在GPU上并行处理批量判断素数

扩展库pyopencl使得可以在Python中调用OpenCL的并行计算API。OpenCL（Open Computing Language）是跨平台的并行编程标准，可以运行在个人电脑、服务器、移动终端以及嵌入式系统等多种平台，既可以运行在CPU上又可以运行于GPU上，大幅度提高了各类应用中的数据处理速度，包括游戏、娱乐、医学软件以及科学计算等等。

import numpy as np

import pyopencl as cl

import pyopencl.array

from pyopencl.elementwise import ElementwiseKernel

#判断素数的C语言版GPU代码

isPrime = ElementwiseKernel(ctx,

'long *a_g, long *b_g, long *res_g',

'''

int j;

for(j=2; j<b_g[i]; j++)

{

if(a_g[i]%j == 0)

{

break;

}

}

if(j >= b_g[i])

{

res_g[i] = a_g[i];

}''',

'isPrime'

)

#定义待测数值范围，和每次处理的数字数量

end = 100000000

start_end = range(2, end)

size = 1000

result = 0

ctx = cl.create_some_context()

queue = cl.CommandQueue(ctx)

#对指定范围内的数字进行分批处理

for i in range(end//size + 1):

startN = i * size

#本次要处理的数字范围

a_np = np.array(start_end[startN: startN+size]).astype(np.int64)

#b_np里的数字是a_np中数字的平方根取整后加1

b_np = np.array(list(map(lambda x: int(x**0.5)+1, a_np))).astype(np.int64)

#把数据写入GPU

a_g = cl.array.to_device(queue, a_np)

b_g = cl.array.to_device(queue, b_np)

res_g = cl.array.zeros_like(a_g)

#批量判断

isPrime(a_g, b_g, res_g)

t = set(filter(None, res_g.get()))

#记录本批数字中素数的个数

result += len(t)

print(result)