前言

前面我們介紹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 目錄。