【TVM 教程】如何在 GPU 上最佳化卷積

超神经HyperAI發表於2024-12-10
GPU卷積

Apache TVM 是一個端到端的深度學習編譯框架，適用於 CPU、GPU 和各種機器學習加速晶片。更多 TVM 中文文件可訪問 →Apache TVM 是一個端到端的深度學習編譯框架，適用於 CPU、GPU 和各種機器學習加速晶片。 | Apache TVM 中文站

本教程演示瞭如何在 TVM 中編寫高效能卷積實現。以正方形大小的輸入張量和濾波器為例，假設卷積輸入的 batch 較大。在此示例中，使用不同的佈局來儲存資料，以實現更好的資料區域性性。緩衝區佈局是 HWCN，分別代表高度、寬度、通道、batch。

準備和演算法​

對具有 256 個通道和 14 x 14 維度的輸入張量使用固定大小。batch size 為 256，卷積過濾器包含 512 個大小為 3 x 3 的過濾器，用步長為 1 和 padding size 為 1 進行卷積。以下程式碼定義了 TVM 中的卷積演算法。

import numpy as np
import tvm
from tvm import te

# 輸入和過濾器的大小
batch = 256
in_channel = 256
out_channel = 512
in_size = 14
kernel = 3
pad = 1
stride = 1

# 演算法
A = te.placeholder((in_size, in_size, in_channel, batch), name="A")
W = te.placeholder((kernel, kernel, in_channel, out_channel), name="W")
out_size = (in_size - kernel + 2 * pad) // stride + 1
# Pad 輸入
Apad = te.compute(
 (in_size + 2 * pad, in_size + 2 * pad, in_channel, batch),
 lambda yy, xx, cc, nn: tvm.tir.if_then_else(
        tvm.tir.all(yy >= pad, yy - pad < in_size, xx >= pad, xx - pad < in_size),
        A[yy - pad, xx - pad, cc, nn],
        tvm.tir.const(0.0, "float32"),
 ),
    name="Apad",
)
# 建立歸約變數
rc = te.reduce_axis((0, in_channel), name="rc")
ry = te.reduce_axis((0, kernel), name="ry")
rx = te.reduce_axis((0, kernel), name="rx")
# 計算卷積
B = te.compute(
 (out_size, out_size, out_channel, batch),
 lambda yy, xx, ff, nn: te.sum(
        Apad[yy * stride + ry, xx * stride + rx, rc, nn] * W[ry, rx, rc, ff], axis=[ry, rx, rc]
 ),
    name="B",
)

記憶體層次結構​

首先指定緩衝區的記憶體層次結構。下圖顯示了 GPU 記憶體層次結構，與 CPU 記憶體層次結構的重要區別是 GPU 提供了一個稱為共享記憶體的快取緩衝區，由程式設計師管理。因此，如何最大化共享記憶體中的資料重用對於在 GPU 核心中實現高效能至關重要。

在本例中，將 Apad 和 W 載入到緩衝區 AA 和 WW 中（儲存在共享記憶體中）。這些緩衝區稍後將由同一執行緒塊中的所有執行緒共享以計算卷積，然後每個執行緒將自己的部分從共享緩衝區載入到它們的本地暫存器 AL 和 WL 中。BL 是輸出 B 的本地快取，也儲存線上程本地暫存器中。

# 指定記憶體層次結構
s = te.create_schedule(B.op)
s[Apad].compute_inline() # compute Apad inline
AA = s.cache_read(Apad, "shared", [B])
WW = s.cache_read(W, "shared", [B])
AL = s.cache_read(AA, "local", [B])
WL = s.cache_read(WW, "local", [B])
BL = s.cache_write(B, "local")

分塊​

以下程式碼將工作負載拆分為執行緒塊和單獨的執行緒，遵循矩陣乘法中的分塊方案。如下圖所示，給定一個畫素座標（y、x），一個執行緒塊負責計算一個 block_factor x block_factor (64 x 64) 的區域，用於輸出通道和 batch。由於共享記憶體空間的限制，每次只從 Apad 和 B 載入 step x block_factor (8 x 64) 資料到共享記憶體中的緩衝區。

# 平鋪常量
tile = 8
num_thread = 8
block_factor = tile * num_thread
step = 8
vthread = 2

# 獲取 GPU 執行緒索引
block_x = te.thread_axis("blockIdx.x")
block_y = te.thread_axis("blockIdx.y")
block_z = te.thread_axis("blockIdx.z")
thread_x = te.thread_axis((0, num_thread), "threadIdx.x")
thread_y = te.thread_axis((0, num_thread), "threadIdx.y")
thread_xz = te.thread_axis((0, vthread), "vthread", name="vx")
thread_yz = te.thread_axis((0, vthread), "vthread", name="vy")

# split 工作負載
hi, wi, fi, ni = s[B].op.axis
bz = s[B].fuse(hi, wi)
by, fi = s[B].split(fi, factor=block_factor)
bx, ni = s[B].split(ni, factor=block_factor)

# 將迭代變數繫結到 GPU 執行緒索引
s[B].bind(bz, block_z)
s[B].bind(by, block_y)
s[B].bind(bx, block_x)

虛擬執行緒分割​

進一步將工作負載從執行緒塊拆分為單個執行緒。為了避免 memory bank conflict，使用虛擬執行緒將區域分成 4 個部分，然後平鋪成 8x8 的網格。因此，如下圖所示，每個執行緒計算 4 個跨步網格，每個網格的大小為 4 x 4。

tyz, fi = s[B].split(fi, nparts=vthread) # 虛擬執行緒 split
txz, ni = s[B].split(ni, nparts=vthread) # 虛擬執行緒 split
ty, fi = s[B].split(fi, nparts=num_thread)
tx, ni = s[B].split(ni, nparts=num_thread)
s[B].reorder(bz, by, bx, tyz, txz, ty, tx, fi, ni)

s[B].bind(tyz, thread_yz)
s[B].bind(txz, thread_xz)
s[B].bind(ty, thread_y)
s[B].bind(tx, thread_x)

協同獲取（Cooperative Fetching）​

如前所述，每個時間步長都要將 step x block_factor 資料從 GPU 全域性記憶體傳輸到共享記憶體。為了減少每個執行緒的記憶體傳輸，以下程式碼讓同一執行緒塊中的執行緒協同從全域性記憶體中獲取相關資料。

# Schedule BL 本地寫入
s[BL].compute_at(s[B], tx)
yi, xi, fi, ni = s[BL].op.axis
ry, rx, rc = s[BL].op.reduce_axis
rco, rci = s[BL].split(rc, factor=step)
s[BL].reorder(rco, ry, rx, rci, fi, ni)

# 將計算附加到迭代變數
s[AA].compute_at(s[BL], rx)
s[WW].compute_at(s[BL], rx)
s[AL].compute_at(s[BL], rci)
s[WL].compute_at(s[BL], rci)

# A 的共享記憶體負載排程
yi, xi, ci, ni = s[AA].op.axis
ty, ci = s[AA].split(ci, nparts=num_thread)
tx, ni = s[AA].split(ni, nparts=num_thread)
_, ni = s[AA].split(ni, factor=4)
s[AA].reorder(ty, tx, yi, xi, ci, ni)
s[AA].bind(ty, thread_y)
s[AA].bind(tx, thread_x)
s[AA].vectorize(ni) # 向量化記憶體載入

# W 的共享記憶體負載排程
yi, xi, ci, fi = s[WW].op.axis
ty, ci = s[WW].split(ci, nparts=num_thread)
tx, fi = s[WW].split(fi, nparts=num_thread)
_, fi = s[WW].split(fi, factor=4)
s[WW].reorder(ty, tx, yi, xi, ci, fi)
s[WW].bind(ty, thread_y)
s[WW].bind(tx, thread_x)
s[WW].vectorize(fi) # 向量化記憶體載入

生成 CUDA 核心​

最後用 TVM 生成和編譯 CUDA 核心，並評估卷積的延遲。

func = tvm.build(s, [A, W, B], "cuda")
dev = tvm.cuda(0)
a_np = np.random.uniform(size=(in_size, in_size, in_channel, batch)).astype(A.dtype)
w_np = np.random.uniform(size=(kernel, kernel, in_channel, out_channel)).astype(W.dtype)
a = tvm.nd.array(a_np, dev)
w = tvm.nd.array(w_np, dev)
b = tvm.nd.array(np.zeros((out_size, out_size, out_channel, batch), dtype=B.dtype), dev)
func(a, w, b)
evaluator = func.time_evaluator(func.entry_name, dev, number=1)
print("Convolution: %f ms" % (evaluator(a, w, b).mean * 1e3))

輸出結果：

Convolution: 54.146944 ms

下載 Python 原始碼：opt_conv_cuda.py

下載 Jupyter notebook：opt_conv_cuda.ipynb

相關文章