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

完整流程

下面我們就來(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 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.。。”）

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)行的指令，第二次同步就是為了等f(wàn)un（）所有線程執(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算子等等。
編輯：lyn