實(shí)操教程｜PyTorch自定義CUDA算子教程與運(yùn)行時(shí)間分析

最近因?yàn)楣ぷ餍枰瑢W(xué)習(xí)了一波CUDA。這里簡(jiǎn)單記錄一下PyTorch自定義CUDA算子的方法，寫(xiě)了一個(gè)非常簡(jiǎn)單的example，再介紹一下正確的PyTorch中CUDA運(yùn)行時(shí)間分析方法。

所有的代碼都放在了github上，地址是：

https://github.com/godweiyang/torch-cuda-examplegithub.com

完整流程

下面我們就來(lái)詳細(xì)了解一下PyTorch是如何調(diào)用自定義的CUDA算子的。

首先我們可以看到有四個(gè)代碼文件：

• main.py，這是python入口，也就是你平時(shí)寫(xiě)模型的地方。
• add2.cpp，這是torch和CUDA連接的地方，將CUDA程序封裝成了python可以調(diào)用的庫(kù)。
• add2.h，CUDA函數(shù)聲明。
• add2.cu，CUDA函數(shù)實(shí)現(xiàn)。

然后逐個(gè)文件看一下是怎么調(diào)用的。

CUDA算子實(shí)現(xiàn)

首先最簡(jiǎn)單的當(dāng)屬add2.h和add2.cu，這就是普通的CUDA實(shí)現(xiàn)。

void launch_add2(float *c,                 const float *a,                 const float *b,                 int n);
__global__ void add2_kernel(float* c,                            const float* a,                            const float* b,                            int n) {    for (int i = blockIdx.x * blockDim.x + threadIdx.x; \            i < n; i += gridDim.x * blockDim.x) {        c[i] = a[i] + b[i];    }}
void launch_add2(float* c,                 const float* a,                 const float* b,                 int n) {    dim3 grid((n + 1023) / 1024);    dim3 block(1024);    add2_kernel<<<grid, block>>>(c, a, b, n);}

這里實(shí)現(xiàn)的功能是兩個(gè)長(zhǎng)度為的tensor相加，每個(gè)block有1024個(gè)線程，一共有個(gè)block。具體CUDA細(xì)節(jié)就不講了，本文重點(diǎn)不在于這個(gè)。

add2_kernel是kernel函數(shù)，運(yùn)行在GPU端的。而launch_add2是CPU端的執(zhí)行函數(shù)，調(diào)用kernel。注意它是異步的，調(diào)用完之后控制權(quán)立刻返回給CPU，所以之后計(jì)算時(shí)間的時(shí)候要格外小心，很容易只統(tǒng)計(jì)到調(diào)用的時(shí)間。

Torch C++封裝

這里涉及到的是add2.cpp，這個(gè)文件主要功能是提供一個(gè)PyTorch可以調(diào)用的接口。

#include <torch/extension.h>#include "add2.h"
void torch_launch_add2(torch::Tensor &c,                       const torch::Tensor &a,                       const torch::Tensor &b,                       int n) {    launch_add2((float *)c.data_ptr(),                (const float *)a.data_ptr(),                (const float *)b.data_ptr(),                n);}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {    m.def("torch_launch_add2",          &torch_launch_add2,          "add2 kernel warpper");}

torch_launch_add2函數(shù)傳入的是C++版本的torch tensor，然后轉(zhuǎn)換成C++指針數(shù)組，調(diào)用CUDA函數(shù)launch_add2來(lái)執(zhí)行核函數(shù)。

這里用pybind11來(lái)對(duì)torch_launch_add2函數(shù)進(jìn)行封裝，然后用cmake編譯就可以產(chǎn)生python可以調(diào)用的.so庫(kù)。但是我們這里不直接手動(dòng)cmake編譯，具體方法看下面的章節(jié)。

Python調(diào)用

最后就是python層面，也就是我們用戶編寫(xiě)代碼去調(diào)用上面生成的庫(kù)了。

import timeimport numpy as npimport torchfrom torch.utils.cpp_extension import load
cuda_module = load(name="add2",                   sources=["add2.cpp", "add2.cu"],                   verbose=True)
# c = a + b (shape: [n])n = 1024 * 1024a = torch.rand(n, device="cuda:0")b = torch.rand(n, device="cuda:0")cuda_c = torch.rand(n, device="cuda:0")
ntest = 10
def show_time(func):    times = list()    res = list()    # GPU warm up    for _ in range(10):        func()    for _ in range(ntest):        # sync the threads to get accurate cuda running time        torch.cuda.synchronize(device="cuda:0")        start_time = time.time()        r = func()        torch.cuda.synchronize(device="cuda:0")        end_time = time.time()
        times.append((end_time-start_time)*1e6)        res.append(r)    return times, res
def run_cuda():    cuda_module.torch_launch_add2(cuda_c, a, b, n)    return cuda_c
def run_torch():    # return None to avoid intermediate GPU memory application    # for accurate time statistics    a + b    return None
print("Running cuda...")cuda_time, _ = show_time(run_cuda)print("Cuda time:  {:.3f}us".format(np.mean(cuda_time)))
print("Running torch...")torch_time, _ = show_time(run_torch)print("Torch time:  {:.3f}us".format(np.mean(torch_time)))

這里6-8行的torch.utils.cpp_extension.load函數(shù)就是用來(lái)自動(dòng)編譯上面的幾個(gè)cpp和cu文件的。最主要的就是sources參數(shù)，指定了需要編譯的文件列表。然后就可以通過(guò)cuda_module.torch_launch_add2，也就是我們封裝好的接口來(lái)進(jìn)行調(diào)用。

接下來(lái)的代碼就隨心所欲了，這里簡(jiǎn)單寫(xiě)了一個(gè)測(cè)量運(yùn)行時(shí)間，對(duì)比和torch速度的代碼，這部分留著下一章節(jié)講解。

總結(jié)一下，主要分為三個(gè)模塊：

• 先編寫(xiě)CUDA算子和對(duì)應(yīng)的調(diào)用函數(shù)。
• 然后編寫(xiě)torch cpp函數(shù)建立PyTorch和CUDA之間的聯(lián)系，用pybind11封裝。
• 最后用PyTorch的cpp擴(kuò)展庫(kù)進(jìn)行編譯和調(diào)用。

運(yùn)行時(shí)間分析

我們知道，CUDA kernel函數(shù)是異步的，所以不能直接在CUDA函數(shù)兩端加上time.time()測(cè)試時(shí)間，這樣測(cè)出來(lái)的只是調(diào)用CUDA api的時(shí)間，不包括GPU端運(yùn)行的時(shí)間。

所以我們要加上線程同步函數(shù)，等待kernel中所有線程全部執(zhí)行完畢再執(zhí)行CPU端后續(xù)指令。這里我們將同步指令加在了python端，用的是torch.cuda.synchronize函數(shù)。

具體來(lái)說(shuō)就是形如下面代碼：

torch.cuda.synchronize()start_time = time.time()func()torch.cuda.synchronize()end_time = time.time()

其中第一次同步是為了防止前面的代碼中有未同步還在GPU端運(yùn)行的指令，第二次同步就是為了等fun()所有線程執(zhí)行完畢后再統(tǒng)計(jì)時(shí)間。

這里我們torch和cuda分別執(zhí)行10次看看平均時(shí)間，此外執(zhí)行前需要先執(zhí)行10次做一下warm up，讓GPU達(dá)到正常狀態(tài)。

我們分別測(cè)試四種情況，分別是：

• 兩次同步
• 第一次同步，第二次不同步
• 第一次不同步，第二次同步
• 兩次不同步

這里我們采用英偉達(dá)的Nsight Systems來(lái)可視化運(yùn)行的每個(gè)時(shí)刻指令執(zhí)行的情況。

安裝命令為：

sudo apt install nsight-systems

然后在運(yùn)行python代碼時(shí)，在命令前面加上nsys profile就行了：

nsys profile python3 main.py

然后就會(huì)生成report1.qdstrm和report1.sqlite兩個(gè)文件，將report1.qdstrm轉(zhuǎn)換為report1.qdrep文件：

QdstrmImporter -i report1.qdstrm

最后將生成的report1.qdrep文件用Nsight Systems軟件打開(kāi)，我這里是mac系統(tǒng)。

兩次同步

這是正確的統(tǒng)計(jì)時(shí)間的方法，我們打開(kāi)Nsight Systems，放大kernel運(yùn)行那一段可以看到下圖：

其中第1和第3個(gè)框分別是cuda和torch的GPU warm up過(guò)程，這部分沒(méi)有進(jìn)行線程同步（上面的黃色塊）。

而第2和第4個(gè)框就分別是cuda和torch的加法執(zhí)行過(guò)程了，我們可以放大來(lái)看看。

可以看出，每執(zhí)行一次（一個(gè)框）都經(jīng)過(guò)了三個(gè)步驟：先是調(diào)用api（左上角藍(lán)色框），然后執(zhí)行kernel（下方藍(lán)色框），最后線程同步（右上角黃色框）。

所以最后算出來(lái)的時(shí)間就是這三個(gè)步驟的耗時(shí)，也就是下圖選中的范圍：

時(shí)間大概在29us左右，和我們實(shí)際代碼測(cè)出來(lái)的也是比較接近的：

其實(shí)我們實(shí)際想要知道的耗時(shí)并不包括api調(diào)用和線程同步的時(shí)間，但是這部分時(shí)間在python端不好去掉，所以就加上了。

第一次同步，第二次不同步

放大每次執(zhí)行的過(guò)程：

可以看出，雖然長(zhǎng)的和上一種情況幾乎一模一樣，但是在api調(diào)用完之后，立刻就進(jìn)行計(jì)時(shí)了，所以耗時(shí)只有8us左右，實(shí)際測(cè)出來(lái)情況也是這樣的：

第一次不同步，第二次同步

我們先來(lái)看一下實(shí)際統(tǒng)計(jì)的時(shí)間：

很奇怪是不是，第一次運(yùn)行耗時(shí)非常久，那我們可視化看看到底怎么回事：

可以看出，因?yàn)榈谝淮伍_(kāi)始計(jì)時(shí)前沒(méi)有同步線程，所以在GPU warm up調(diào)用api完畢后，第一次cuda kernel調(diào)用就開(kāi)始了。然后一直等到warm up執(zhí)行完畢，才開(kāi)始執(zhí)行第一次cuda kernel，然后是線程同步，結(jié)束后才結(jié)束計(jì)時(shí)。這個(gè)過(guò)程非常長(zhǎng)，差不多有130us左右。然后第二次開(kāi)始執(zhí)行就很正常了，因?yàn)閗ernel結(jié)束的同步相當(dāng)于是下一次執(zhí)行之前的同步。

兩次不同步

先來(lái)看看執(zhí)行情況：

可以看出因?yàn)闆](méi)有任何同步，所有GPU warm up和cuda kernel的api調(diào)用全接在一起了，執(zhí)行也是。所以計(jì)時(shí)只計(jì)算到了每個(gè)api調(diào)用的時(shí)間，差不多在7us左右。

上面四種情況，torch指令情形幾乎一樣，因此不再贅述。

小結(jié)

通過(guò)這篇文章，應(yīng)該可以大致了解PyTorch實(shí)現(xiàn)自定義CUDA算子并調(diào)用的方法，也能知道怎么正確的測(cè)量CUDA程序的耗時(shí)。

當(dāng)然還有一些內(nèi)容留作今后講解，比如如何實(shí)現(xiàn)PyTorch神經(jīng)網(wǎng)絡(luò)的自定義前向和反向傳播CUDA算子、如何用TensorFlow調(diào)用CUDA算子等等。