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

完整流程

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

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

main.py，這是python入口，也就是你平時寫模型的地方。

add2.cpp，這是torch和CUDA連接的地方，將CUDA程序封裝成了python可以調(diào)用的庫。

add2.h，CUDA函數(shù)聲明。

add2.cu，CUDA函數(shù)實(shí)現(xiàn)。

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

CUDA算子實(shí)現(xià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)的功能是兩個長度為的tensor相加，每個block有1024個線程，一共有個block。具體CUDA細(xì)節(jié)就不講了，本文重點(diǎn)不在于這個。

add2_kernel是kernel函數(shù)，運(yùn)行在GPU端的。而launch_add2是CPU端的執(zhí)行函數(shù)，調(diào)用kernel。注意它是異步的，調(diào)用完之后控制權(quán)立刻返回給CPU，所以之后計(jì)算時間的時候要格外小心，很容易只統(tǒng)計(jì)到調(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ù)進(jìn)行封裝，然后用cmake編譯就可以產(chǎn)生python可以調(diào)用的.so庫。但是我們這里不直接手動cmake編譯，具體方法看下面的章節(jié)。

Python調(diào)用

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

import time

import numpy as np

import torch

from 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 * 1024

a = 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.。?！保?/p>

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，也就是我們封裝好的接口來進(jìn)行調(diào)用。

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

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

先編寫CUDA算子和對應(yīng)的調(diào)用函數(shù)。

然后編寫torch cpp函數(shù)建立PyTorch和CUDA之間的聯(lián)系，用pybind11封裝。

最后用PyTorch的cpp擴(kuò)展庫進(jìn)行編譯和調(diào)用。

運(yùn)行時間分析

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

所以我們要加上線程同步函數(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端運(yùn)行的指令，第二次同步就是為了等fun（）所有線程執(zhí)行完畢后再統(tǒng)計(jì)時間。

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

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

兩次同步

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

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

兩次不同步

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

安裝命令為：

sudo apt install nsight-systems

然后在運(yùn)行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)計(jì)時間的方法，我們打開Nsight Systems，放大kernel運(yùn)行那一段可以看到下圖：

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

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

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

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

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

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

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

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

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

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

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

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

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

兩次不同步

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

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

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

小結(jié)

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

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