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)