GPU程序设计(5) -- Python

前言

前面我们介绍C++的使用，有些读者可能会希望使用Python撰写(包括我)，因此，我们就来看看 PyCuda 这个套件的用法。

PyCuda 安装及文件

PyCuda 可以将C/C++程序包在Python字串中，执行时会先使用NVCC编译，所以，读者还是要安装CUDA toolkit及 VC Studio，PyCuda安装很简单，以下列指令执行：

pip install pycuda

官网的文件：https://documen.tician.de/pycuda/tutorial.html
官网的范例：https://wiki.tiker.net/PyCuda/Examples/

程序说明

先写一支 Hello 的程序。

1. 引进相关套件/模组。

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

2. 撰写 GPU 函数，使用 C/C++ 语法。

mod = SourceModule("""
    #include <stdio.h>

     __global__ void GPU_function()
       {
        printf("Hello PyCUDA!!!");
      }
""")

3. 以 Python 呼叫上述程序，指定多执行绪参数(1,1,1)，表单一区块、单一执行绪。

function = mod.get_function("GPU_function")
function(block=(1,1,1))
""")

完整程序码如下：

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

mod = SourceModule("""
    #include <stdio.h>

     __global__ void GPU_function()
       {
        printf("Hello PyCUDA!!!");
      }
""")
 
function = mod.get_function("GPU_function")
function(block=(1,1,1))
""")

执行上述程序，输出如下：

Hello PyCUDA!!!

第一次执行会很慢，我猜是Python桥接C的关系，因此，若不是很复杂的程序，这种混合语言的写法并不会得到好处。

进阶

你可以把C程序放在一个档案中，例如 c_code.cu，然後以 python 读入执行，这样就类似函数库(Library)的概念，可以尽情扩充 c_code.cu。

mod = SourceModule(open('./c_code.cu', encoding='utf8').read())

变数传递

GPU只支援单精度(Single)浮点数，要将 Python 变数复制到 GPU 上，双精度的变数须转型。

import numpy
a = numpy.random.randn(4,4)
# 双精度的变数须转型为单精度(Single)浮点数
a = a.astype(numpy.float32)

# 配置GPU记忆体
d_a = cuda.mem_alloc(a.nbytes)

# 复制到 GPU 上
cuda.memcpy_htod(d_a, a)

mod = SourceModule("""
  __global__ void square(float *a)
  {
    int idx = threadIdx.x + threadIdx.y*4;
    a[idx] *= 2;
  }
  """)

# Python 呼叫 C 程序  
func = mod.get_function("square")
func(d_a, block=(4,4,1))  

# 复制到 CPU 上
h_a = numpy.empty_like(a)
cuda.memcpy_dtoh(h_a, d_a)
print("\n平方：")
print(h_a)

Github档案为 03_pass_variable.py。

精简写法

之前呼叫GPU函数前，变数都要先复制到GPU，PyCuda 提供 cuda.InOut() 函数，自动完成这些转换，缩减的程序如下：

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

# 双精度的变数须转型为单精度(Single)浮点数
a = numpy.random.randn(4,4)

mod = SourceModule("""
  __global__ void square(float *a)
  {
    int idx = threadIdx.x + threadIdx.y*4;
    a[idx] *= 2;
  }
  """)
  
# Python 呼叫 C 程序  
func = mod.get_function("square")
func(cuda.InOut(a), block=(4,4,1))  

print("\n平方：")
print(a)

Github档案为 04_inout.py。

减少函数调用的成本

如果会重复呼叫GPU函数多次，可以像资料库的预存程序(Stored Procedure)一样，将编译的程序码储存起来，之後就直接呼叫编译的程序码即可。

# Python 呼叫 C 程序  
func = mod.get_function("square")

# 编译程序码
func.prepare("P")
grid = (1, 1)
block = (4, 4, 1)
func.prepared_call(grid, block, d_a)

完整程序码请参照 05_prepare.py。

结语

本系列的文章到此告一段落，还有许多宝藏待挖掘，有待後续再慢慢咀嚼了。

相关程序可至『GitHub』下载，本篇程序在 python 目录。