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 的程式。

import pycuda.driver as cudaimport pycuda.autoinitfrom 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))""")

完整程式码如下：

import pycuda.driver as cudaimport pycuda.autoinitfrom pycuda.compiler import SourceModulemod = 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 numpya = 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 cudaimport pycuda.autoinitfrom pycuda.compiler import SourceModuleimport 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 目录。