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

完整流程

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

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

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

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

add2.h，CUDA函數聲明。

add2.cu，CUDA函數實現。

然后逐個文件看一下是怎么調用的。

CUDA算子實現

首先最簡單的當屬add2.h和add2.cu，這就是普通的CUDA實現。

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）;

}

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

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

Torch C++封裝

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

#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函數傳入的是C++版本的torch tensor，然后轉換成C++指針數組，調用CUDA函數launch_add2來執行核函數。

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

Python調用

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

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

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

總結一下，主要分為三個模塊：

先編寫CUDA算子和對應的調用函數。

然后編寫torch cpp函數建立PyTorch和CUDA之間的聯系，用pybind11封裝。

最后用PyTorch的cpp擴展庫進行編譯和調用。

運行時間分析

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

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

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

torch.cuda.synchronize（）

start_time = time.time（）

func（）

torch.cuda.synchronize（）

end_time = time.time（）

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

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

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

兩次同步

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

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

兩次不同步

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

安裝命令為：

sudo apt install nsight-systems

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

nsys profile python3 main.py

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

QdstrmImporter -i report1.qdstrm

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

兩次同步

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

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

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

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

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

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

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

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

放大每次執行的過程：

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

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

我們先來看一下實際統計的時間：

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

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

兩次不同步

先來看看執行情況：

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

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

小結

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

當然還有一些內容留作今后講解，比如如何實現PyTorch神經網絡的自定義前向和反向傳播CUDA算子、如何用TensorFlow調用CUDA算子等等。
編輯：lyn