如何在GPU上優化卷積
本文將演示如何在TVM中編寫高性能的卷積實現。以平方大小的輸入張量和濾波器為例，并假設卷積的輸入量很大。使用不同的布局來存儲數據，以實現更好的數據局部性。緩沖區布局為HWCN，代表高度，寬度，通道，批次。
準備和算法
將固定大小用于256通道和14 x 14尺寸的輸入張量。批處理大小為256。卷積過濾器包含512個大小為3 x 3的過濾器。對于卷積，使用步幅大小1和填充大小1。以下代碼定義了TVM中的卷積算法。
import numpy as np
import tvm
from tvm import te

# The sizes of inputs and filters
batch = 256
in_channel = 256
out_channel = 512
in_size = 14
kernel = 3
pad = 1
stride = 1

# Algorithm
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 input
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”,
)
# Create reduction variables
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”)
# Compute the convolution
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的本地緩存，它也存儲在線程本地寄存器中。
# Designate the memory hierarchy
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”)
阻塞Blocking
以下代碼將工作負載分為線程塊和單個線程。在矩陣乘法中遵循阻塞方案。如下圖所示，給定像素坐標（y，x），線程塊負責為輸出通道和批處理計算block_factor x block_factor（64 x 64）的區域。由于共享內存空間的限制，每次僅將Apad和B中的step x block_factor（8 x 64）數據加載到共享內存中的緩沖區中。

# tile consts
tile = 8
num_thread = 8
block_factor = tile * num_thread
step = 8
vthread = 2

# Get the GPU thread indices
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 the workloads
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)

# Bind the iteration variables to GPU thread indices
s[B].bind(bz, block_z)
s[B].bind(by, block_y)
s[B].bind(bx, block_x)
虛擬線程分割Virtual Thread Split
將工作負載從線程塊劃分到各個線程。為避免內存庫沖突，使用虛擬線程將區域劃分為4個部分，然后平鋪為8x8網格。如下圖所示，每個線程計算4個網格，每個網格的大小為4 x 4。

tyz, fi = s[B].split(fi, nparts=vthread) # virtual thread split
txz, ni = s[B].split(ni, nparts=vthread) # virtual thread 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
每個時間步驟都需要將步驟x block_factor數據從GPU全局內存傳輸到共享內存。為了減少每個線程的內存傳輸，以下代碼使同一線程塊中的線程，可以協作地從全局內存中獲取相關數據。
# Schedule BL local write
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)

# Attach computation to iteration variables
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)

# Schedule for A’s shared memory load
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) # vectorize memory load

# Schedule for W’s shared memory load
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) # vectorize memory load
生成CUDA內核
最后，使用TVM生成和編譯CUDA內核，并評估卷積的延遲。
func = tvm.build(s, [A, W, B], “cuda”)
ctx = tvm.gpu(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, ctx)
w = tvm.nd.array(w_np, ctx)
b = tvm.nd.array(np.zeros((out_size, out_size, out_channel, batch), dtype=B.dtype), ctx)
func(a, w, b)
evaluator = func.time_evaluator(func.entry_name, ctx, number=1)
print(“Convolution: %f ms” % (evaluator(a, w, b).mean * 1e3))
出：
Convolution: 53.197723 ms

https://tvm.apache.org/docs/tutorials/optimize/opt_conv_cuda.html#sphx-glr-tutorials-optimize-opt-conv-cuda-py

總結

以上是生活随笔為你收集整理的如何在GPU上优化卷积的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。