您当前的位置:首页 > 计算机 > 编程开发 > Python

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

时间:09-09来源:作者:点击数:

借助于扩展库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倍左右。

方便获取更多学习、工作、生活信息请关注本站微信公众号城东书院 微信服务号城东书院 微信订阅号
推荐内容
相关内容
栏目更新
栏目热门