Python CUDA 編程 - 5 - 多流

之前討論的並行，都是線程級別的，即CUDA開啟多個線程，並行執行核函數內的代碼。GPU最多就上千個核心，同一時間只能並行執行上千個任務。當我們處理千萬級別的數據，整個大任務無法被GPU一次執行，所有的計算任務需要放在一個隊列中，排隊順序執行。CUDA將放入隊列順序執行的一系列操作稱為流（Stream）。

來源

由於異構計算的硬件特性，CUDA中以下操作是相互獨立的，通過編程，是可以操作他們並發地執行的：

• 主機端上的計算
• 設備端的計算（核函數）
• 數據從主機和設備間相互拷貝
• 數據從設備內拷貝或轉移
• 數據從多個GPU設備間拷貝或轉移

針對這種互相獨立的硬件架構，CUDA使用多流作為一種高並發的方案：

• 把一個大任務中的上述幾部分拆分開，放到多個流中，每次只對一部分數據進行拷貝、計算和回寫，並把這個流程做成流水線。
• 因為數據拷貝不占用計算資源，計算不占用數據拷貝的總線（Bus）資源，因此計算和數據拷貝完全可以並發執行。
• 如圖所示，將數據拷貝和函數計算重疊起來的，形成流水線，能獲得非常大的性能提升。

實際上，流水線作業的思想被廣泛應用於CPU和GPU等計算機芯片設計上，以加速程序。

多流

以向量加法為例，上圖中第一行的Stream 0部分是我們之前的邏輯，沒有使用多流技術，程序的三大步驟是順序執行的：

1. 先從主機拷貝初始化數據到設備（Host To Device）；
2. 在設備上執行核函數（Kernel）；
3. 將計算結果從設備拷貝回主機（Device To Host）。

當數據量很大時，每個步驟的耗時很長，後面的步驟必須等前面執行完畢才能繼續，整體的耗時相當長。

以2000萬維的向量加法為例，向量大約有幾十M大小，將整個向量在主機和設備間拷貝將占用占用上百毫秒的時間，有可能遠比核函數計算的時間多得多。將程序改為多流後，每次只計算一小部分，流水線並發執行，會得到非常大的性能提升。

規則

默認情況下，CUDA使用0號流，又稱默認流。不使用多流時，所有任務都在默認流中順序執行，效率較低。在使用多流之前，必須先了解多流的一些規則：

• 給定流內的所有操作會按序執行。
• 非默認流之間的不同操作，無法保證其執行順序。
• 所有非默認流執行完後，才能執行默認流；默認流執行完後，才能執行其他非默認流。

參照上圖，可將這三個規則解釋為：

• 非默認流1中，根據進流的先後順序，核函數1和2是順序執行的。
• 無法保證核函數2與核函數4的執行先後順序，因為他們在不同的流中。他們執行的開始時間依賴於該流中前一個操作結束時間，例如核函數2的開始依賴於核函數1的結束，與核函數3、4完全不相關。
• 默認流有阻塞的作用。如圖中紅線所示，如果調用默認流，那麼默認流會等非默認流都執行完才能執行；同樣，默認流執行完，才能再次執行其他非默認流。

某個流內的操作是順序的，非默認流之間是異步的，默認流有阻塞作用。

使用

定義

如果想使用多流時，必須先定義流：

stream = numba.cuda.stream()

CUDA的數據拷貝以及核函數都有專門的stream參數來接收流，以告知該操作放入哪個流中執行：

• numba.cuda.to_device(obj, stream=0, copy=True, to=None)
• numba.cuda.copy_to_host(self, ary=None, stream=0)

核函數調用的地方除了要寫清執行配置，還要加一項stream參數：

• kernel[blocks_per_grid, threads_per_block, stream=0]

根據這些函數定義也可以知道，不指定stream參數時，這些函數都使用默認的0號流。

對於程序員來說，需要將數據和計算做拆分，分別放入不同的流裡，構成一個流水線操作。

將之前的向量加法的例子改為多流處理，完整的代碼為：

from numba import cuda
import numpy as np
import math
from time import time
@cuda.jit
def vector_add(a, b, result, n):
idx = cuda.threadIdx.x + cuda.blockDim.x * cuda.blockIdx.x
if idx < n :
result[idx] = a[idx] + b[idx]
def main():
n = 20000000
x = np.random.uniform(10,20,n)
y = np.random.uniform(10,20,n)
# x = np.arange(n).astype(np.int32)
# y = 2 * x
start = time()
# 使用默認流
# Host To Device
x_device = cuda.to_device(x)
y_device = cuda.to_device(y)
z_device = cuda.device_array(n)
z_streams_device = cuda.device_array(n)
threads_per_block = 1024
blocks_per_grid = math.ceil(n / threads_per_block)
# Kernel
vector_add[blocks_per_grid, threads_per_block](x_device, y_device, z_device, n)
# Device To Host
default_stream_result = z_device.copy_to_host()
cuda.synchronize()
print("gpu vector add time " + str(time() - start))
start = time()
# 使用5個流
number_of_streams = 5
# 每個流處理的數據量為原來的 1/5
# 符號//得到一個整數結果
segment_size = n // number_of_streams
# 創建5個cuda stream
stream_list = list()
for i in range (0, number_of_streams):
stream = cuda.stream()
stream_list.append(stream)
threads_per_block = 1024
# 每個stream的處理的數據變為原來的1/5
blocks_per_grid = math.ceil(segment_size / threads_per_block)
streams_result = np.empty(n)
# 啟動多個stream
for i in range(0, number_of_streams):
# 傳入不同的參數，讓函數在不同的流執行
# Host To Device
x_i_device = cuda.to_device(x[i * segment_size : (i + 1) * segment_size], stream=stream_list[i])
y_i_device = cuda.to_device(y[i * segment_size : (i + 1) * segment_size], stream=stream_list[i])
# Kernel
vector_add[blocks_per_grid, threads_per_block, stream_list[i]](
x_i_device,
y_i_device,
z_streams_device[i * segment_size : (i + 1) * segment_size],
segment_size)
# Device To Host
streams_result[i * segment_size : (i + 1) * segment_size] = z_streams_device[i * segment_size : (i + 1) * segment_size].copy_to_host(stream=stream_list[i])
cuda.synchronize()
print("gpu streams vector add time " + str(time() - start))
if (np.array_equal(default_stream_result, streams_result)):
print("result correct")
if __name__ == "__main__":
main()

運行結果：

gpu vector add time 7.996402740478516
gpu streams vector add time 0.3867764472961426

多流不僅需要程序員掌握流水線思想，還需要用戶對數據和計算進行拆分，並編寫更多的代碼，但是收益非常明顯。對於計算密集型的程序，這種技術非常值得認真研究。

參考資料

• https://lulaoshi.info/gpu/python-cuda/streams.html