實操教程｜PyTorch自定義CUDA算子教程與運行時間分析

最近因為工作需要，學習了一波CUDA。這里簡單記錄一下PyTorch自定義CUDA算子的方法，寫了一個非常簡單的example，再介紹一下正確的PyTorch中CUDA運行時間分析方法。

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

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

完整流程

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

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

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

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

CUDA算子實現(xiàn)

首先最簡單的當屬add2.h和add2.cu，這就是普通的CUDA實現(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);}

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

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

Torch C++封裝

這里涉及到的是add2.cpp，這個文件主要功能是提供一個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來執(zhí)行核函數(shù)。

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

Python調(diào)用

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

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ù)就是用來自動編譯上面的幾個cpp和cu文件的。最主要的就是sources參數(shù)，指定了需要編譯的文件列表。然后就可以通過cuda_module.torch_launch_add2，也就是我們封裝好的接口來進行調(diào)用。

接下來的代碼就隨心所欲了，這里簡單寫了一個測量運行時間，對比和torch速度的代碼，這部分留著下一章節(jié)講解。

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

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

運行時間分析

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

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

具體來說就是形如下面代碼：

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

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

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

我們分別測試四種情況，分別是：

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

這里我們采用英偉達的Nsight Systems來可視化運行的每個時刻指令執(zhí)行的情況。

安裝命令為：

sudo apt install nsight-systems

然后在運行python代碼時，在命令前面加上nsys profile就行了：

nsys profile python3 main.py

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

QdstrmImporter -i report1.qdstrm

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

兩次同步

這是正確的統(tǒng)計時間的方法，我們打開Nsight Systems，放大kernel運行那一段可以看到下圖：

其中第1和第3個框分別是cuda和torch的GPU warm up過程，這部分沒有進行線程同步（上面的黃色塊）。

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

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

所以最后算出來的時間就是這三個步驟的耗時，也就是下圖選中的范圍：

時間大概在29us左右，和我們實際代碼測出來的也是比較接近的：

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

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

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

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

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

我們先來看一下實際統(tǒng)計的時間：

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

可以看出，因為第一次開始計時前沒有同步線程，所以在GPU warm up調(diào)用api完畢后，第一次cuda kernel調(diào)用就開始了。然后一直等到warm up執(zhí)行完畢，才開始執(zhí)行第一次cuda kernel，然后是線程同步，結(jié)束后才結(jié)束計時。這個過程非常長，差不多有130us左右。然后第二次開始執(zhí)行就很正常了，因為kernel結(jié)束的同步相當于是下一次執(zhí)行之前的同步。

兩次不同步

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

可以看出因為沒有任何同步，所有GPU warm up和cuda kernel的api調(diào)用全接在一起了，執(zhí)行也是。所以計時只計算到了每個api調(diào)用的時間，差不多在7us左右。

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

小結(jié)

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

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