通过Numba调用CUDA用GPU为Python加速：进阶理解网格跨步、多流、共享内存

• 前导知识

  理解本文需要先了解：

  1. 计算机底层基础知识，CPU、机器码、编译等《编译型语言与解释型语言如何在计算机底层运行》《计算机底层运转机制：多核、缓存、CPU、CU、ALU、Cache》
  2. Python代码与GPU加速的关系《Python程序如何用GPU加速：Tesla、CUDA、Numba》
  3. 在CPU入门numba《Python代码在CPU下加速：Numba入门》
  4. 在GPU入门numba《Python通过Numba实现GPU加速》

• CUDA优化方向

  CPU + GPU 是一种异构计算的组合，各有独立的内存，GPU的优势是更多的计算核心。该架构在并行计算上有很大优势，但是数据需要从主机和设备间相互拷贝，会造成一定的延迟。因此，要从下面两个方面来优化GPU程序：

  • 充分利用GPU的多核心，最大化并行执行度
  • 优化内存使用，最大化数据吞吐量，减少不必要的数据拷贝

  哪个方向有更大收益，最终还是要看具体的计算场景。

  英伟达提供了非常强大的性能分析器nvprof和可视化版nvvp，使用性能分析器能监控到当前程序的瓶颈。分析器只支持C/C++编译后的可执行文件，Python Numba目前应该不支持。

• 并行计算优化

  CUDA的执行配置：[gridDim, blockDim]中的blockDim最大只能是1024，gridDim最大为一个32位整数的最大值（2,147,483,648）,即20多亿。

• 网格跨度

  如果对并行计算的维度有更高需求，就需要网格跨度。

  

  以[2,4]的执行配置为例，该执行配置中整个grid只能并行执行8个线程，如上图，如果我们需要并行计算的数据是32，会发现8-32号共计24个数据无法被计算。

  

  如上图（网格跨度），可以在0号线程中，处理第0、8、16、24号数据，就能解决数据远大于执行配置中线程总数的问题，用程序表示，就是在核函数里再写个for循环。

  from numba import cuda
  
  @cuda.jit
  def gpu_print(N):
      idxWithinGrid = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x 
      gridStride = cuda.gridDim.x * cuda.blockDim.x
      # 从 idxWithinGrid 开始
      # 每次以整个网格线程总数为跨步数
      for i in range(idxWithinGrid, N, gridStride):
          print(i)
  
  def main():
      gpu_print[2, 4](32)
      cuda.synchronize()
  
  if __name__ == "__main__":
      main()
  

  跨步大小：网格(grid)中线程(thread)总数=gridDim.x * blockDim.x，即for循环中的step，也是网格跨步名称的由来。

  举例，如果网格总线程数(跨步大小)为1024，那么0号线程将