PyCUDA矩阵按元素平方运算

import pycuda.driver as drv
import pycuda.autoinit  
from pycuda.compiler import SourceModule
import numpy
mod = SourceModule("""
    __global__ void square(float *d_a)
    {
      int idx = threadIdx.x + threadIdx.y*5;
      d_a[idx] = d_a[idx]*d_a[idx];
    }
    """)
start = drv.Event()
end=drv.Event()
h_a = numpy.random.randint(1,5,(5, 5))
h_a = h_a.astype(numpy.float32)
h_b=h_a.copy()
start.record()
d_a = drv.mem_alloc(h_a.size * h_a.dtype.itemsize)
drv.memcpy_htod(d_a, h_a)
square = mod.get_function("square")
square(d_a, block=(5, 5, 1), grid=(1, 1), shared=0)
h_result = numpy.empty_like(h_a)
drv.memcpy_dtoh(h_result, d_a)
end.record()
end.synchronize()
secs = start.time_till(end)*1e-3
print("Time of Squaring on GPU without inout")
print("%fs" % (secs))
print("original array:")
print(h_a)
print("Square with kernel:")
print(h_result)

#----Using inout functionality of driver class -------------------------------------------------
start.record()
start.synchronize()
square(drv.InOut(h_a), block=(5, 5, 1))
end.record()
end.synchronize()

print("Square with InOut:")
print(h_a)
secs = start.time_till(end)*1e-3
print("Time of Squaring on GPU with inout")
print("%fs" % (secs))

# ---------------Using gpuarray class----------------------------------------------------#
import pycuda.gpuarray as gpuarray
start.record()
start.synchronize()
h_b = numpy.random.randint(1,5,(5, 5))
d_b = gpuarray.to_gpu(h_b.astype(numpy.float32))
h_result = (d_b**2).get()
end.record()
end.synchronize()
print("original array:")
print(h_b)
print("Squared with gpuarray:")
print(h_result)
secs = start.time_till(end)*1e-3
print("Time of Squaring on GPU with gpuarray")
print("%fs" % (secs))

首先创建两个事件来测量内核函数的计时。接着5×5的矩阵在主机上用随机数初始化，通过numpy.random模块的randint函数完成，这里需要3个参数:前两个参数定义用于生成随机数的数字范围，第一个参数用于生成数字的最小值，第二个参数是用于生成数字的最大值;第三个参数是数组大小，这里指定为元组(5，5)。将生成的矩阵再次转换为单精度数字，以加快处理速度。该矩阵的内存被分配到设备上，生成的随机数矩阵被复制到设备上。
创建指向内核函数的指针引用，并通过将设备显存指针作为参数传递来调用内核。内核使用多维线程调用，x和y方向的值为5，因此启动的线程总数是25，每个线程计算矩阵中单个元素的平方。内核计算的结果被复制回主机并显示在控制台上，计算所需的时间与输入和输出矩阵一起显示在控制台上。
同一个数组同时作为输入和输出，PyCUDA中的driver模块为这种情况提供一个inout 指令，可以消除为数组单独分配内存的需要，直接将其上载到设备并将结果下载回主机。所有操作都在内核调用期间同时执行，这使得代码更简单，更易于阅读。
Python为数值计算提供一个numpy库，PyCUDA则提供一个类似于numpy的 gpuarray类，用来存储数据并在GPU设备上执行计算。数组形状和数据类型与numpy完全相同,gpuarray类为计算提供许多算术方法，消除C或C++内核代码指定SourceModule的要求,因此PyCUDA代码将只包含Python代码。
导入pycuda.gpuarray模块中可用的 gpuarray类以便在代码中使用，矩阵初始化为从1到5的随机整数进行计算。此矩阵通过使用gpuarray类的to_gpu()方法上载到设备显存,将上载的矩阵作为此方法的参数，然后将矩阵转换成单精度数字，上载到矩阵上的所有操作都在设备上执行。平方操作的执行方式与我们在Python代码中的执行方式类似，但是由于变量使用gpuarray类存储在设备上，因此该操作也将在设备上执行。使用get方法将结果下载回主机。