PyCuda学习二之矩阵的加法

矩阵的加法

1. 导入必要的库

   import numpy as np
   import pycuda.autoinit
   import pycuda.driver as cuda
   from pycuda.compiler import SourceModule
   

2. 编写核函数 注意: 三引号里面的语法是c/c++的语法

   mod = SourceModule("""
   	__global__ void gpu_add(float* x, float* y, size_t N)
   	{
   		int index = threadIdx.x + blockDim.x * blockIdx.x;
   		int stride = blockDim.x * gridDim.x;
   		for (int i = index; i < N; i += stride) {
   			y[i] = x[i] + y[i]
   		}
   	}
   """)
   
   

3. 编写测试函数

   
   def test(N):
     	a = np.empty(N, dtype=np.float32)
     	a[:] = 10.0
       b = np.empty(N, dtype=np.float32)
       b[:] = 20.0
       # N一定要转成int不然会出错, 转成32位和64位的都行
       N = np.int32(N)
       # 设定线程块为256 可以改变, 但是256对于当前这个程序来说是最优的, 至于为什么最优目前还不知道
       thread_size = 256
       # 计算block的大小
       block_size = int((N + thread_size - 1) / thread_size)
       print(block_size)
       # 得到核函数的引用
       add = mod.get_function("add")
       # 调用
       add(
           cuda.In(a), cuda.InOut(b), N,
           block=(thread_size, 1, 1), grid=(block_size, 1)
       )
       print(b[:10])
   
   if __name__ == '__main__':
     	N = 1024 * 1024
     	test(N)
   
   

4. 全部的代码

   import numpy as np
   import pycuda.autoinit
   import pycuda.driver as cuda
   from pycuda.compiler import SourceModule
   
   
   mod = SourceModule("""
   
       __global__ void add(float* x, float* y, size_t N)
       {
           // 获取全局索引
           int index = threadIdx.x + blockIdx.x * blockDim.x;
           // 步长
           int stride = blockDim.x * gridDim.x;
           for (int i = index; i < N; i += stride)
           {
               y[i] = x[i] + y[i];
           }
       }
   
   """)
   
   
   def test(N):
   
       a = np.empty(N, dtype=np.float32)
       a[:] = 10.0
       b = np.empty(N, dtype=np.float32)
       b[:] = 20.0
       # N一定要转成int不然会出错, 转成32位和64位的都行
       N = np.int32(N)
       # 设定线程块为256 可以改变, 但是256对于当前这个程序来说是最优的, 至于为什么最优目前还不知道
       thread_size = 256
       # 计算block的大小
       block_size = int((N + thread_size - 1) / thread_size)
       print(block_size)
       # 得到核函数的引用
       add = mod.get_function("add")
       # 调用
       add(
           cuda.In(a), cuda.InOut(b), N,
           block=(thread_size, 1, 1), grid=(block_size, 1)
       )
       print(b[:10])
   
   
   if __name__ == '__main__':
       N = 1024 * 1024
       test(N)