準備工作:

要對Pytorch的c++代碼進行比較方便的debug, 需要先做完以下幾件事兒:

1. 購買NVIDIA顯卡裝在自己的電腦上. (需要會拆裝顯卡, 連接電源線之類的操作)

2. 安裝nvidia顯卡驅(qū)動, 安裝cuda, 安裝cudnn. 這里會遇到顯卡驅(qū)動版本/顯卡本身/cuda/cudnn版本是否匹配的問題. 安裝過程中會遇到gcc版本和linux核版本是否一致的問題. (centos的話可以直接yum安裝顯卡驅(qū)動)

3. 下載一份Pytorch源碼. 這里會遇到pytorch子倉無法下載的問題(里面有Google的代碼倉無法直連, 需要魔法上網(wǎng), 并且把git的http/socks代理設(shè)置為梯子的本地端口號).

4. 源譯編譯一份debug版本的Pytorch. 這里會遇到cmake版本問題. 為了減少問題, 盡量把編譯腳本中能禁用的選項全部禁用掉. (DEBUG=1 USE_DISTRIBUTED=0 USE_MKLDNN=0 BUILD_TEST=0 USE_FBGEMM=0 USE_NNPACK=0 USE_QNNPACK=0 USE_XNNPACK=0 USE_NCCL=0 python3 setup.py develop
   )

5. 源碼編譯一個和gcc匹配的gdb, 這里可能遇到動態(tài)庫版本不一致的問題(比如說libstdc++.so.6).

6. 給你的linux系統(tǒng)安裝圖形界面. 再安裝一個linux版本的vscode和chrome(方便隨時查資料). 這一步可以不做, 但是為了方便代碼跳轉(zhuǎn)和閱讀的話, 最好還是在linux的圖形界面上搞.

可能是因為我的linux是centos7.9, 反正遇到的問題比較多. 也許用ubuntu18.04系統(tǒng)問題會少一點.

如果你對linux上軟件的源碼編譯安裝, 軟鏈接替換等操作不熟悉的話, 這個過程就會比較艱難, 以上每一件事兒都可能卡住半天甚至一天的時間. 大概需要一個星期才能搞定所有問題.

如果你有鈔能力的話, 以上每一個步驟你都可以在淘寶上找運維裝機服務(wù)來解決, 估計500塊錢之內(nèi)可以解決.

正文

CPU上是怎么調(diào)用過來的 跟另一篇博客大同小異, 但因為版本不同(我的torch版本為: 2.1.0a0+gitdc4a253 ), 代碼邏輯有所變化.

首先是從python_variable_methods這個文件中, 這個文件主要做的應(yīng)該是c++和python互操作的功能, 就是調(diào)用python對象中加法綁定的c++函數(shù).

1-python_variable_methods.png

然后進入THPVariable_add方法, 這個進行一個操作數(shù)的類型判斷, 調(diào)用2565行的Tensor的add方法.

2-THPVariable_add.png

進入ops的add_Tensor方法

3-TensorBody.png

調(diào)用939行創(chuàng)建的op的call方法.

4-Operators_2.png

這里是比較關(guān)鍵的角色出場了, Dispatcher, 所有的算子應(yīng)該都是注冊到這個Dispatcher中的.

5-TypedOperatorHandle.png

在Dispatcher的調(diào)度下, 根據(jù)632行op的dispatchKeySet找到對應(yīng)的op. 然后進行調(diào)用.

6-Dispatcher__call.png

這部分代碼看的云里霧里的, 應(yīng)該都是一些為了讓代碼簡潔而寫的封裝代碼吧

7-KernelFunction__call.png

8-callUnboxedKernelFunction.png

9-wrap_kernel_functor_unboxed_.png

10-WrapFunctionIntoFunctor_.png

經(jīng)過一系列云里霧里的繞繞, 進入到這個wrapper_CUDA_add_Tensor.cu這個文件, 這個就是cuda代碼的文件了.
在op.impl方法中進行實際的調(diào)用

11-wrapper_CUDA_add_Tensor.png

對于Tensor加法來說, 這個op.impl就是332行的代碼, 這里調(diào)用add_kernel方法

12-UfuncCUDA_add.png

67行這里是一個宏寫法的switch語句, 下面一大片代碼只會走其中一個分支. 因為是宏的寫法, debug進不去. 也沒關(guān)系直接看case就好了. 不管哪個case都會調(diào)用gpu_kernel這個方法

13-add_kernel.png

下面是gpu_kernel的具體實現(xiàn), 調(diào)用了gpu_kernel_impl方法

14-gpu_kernel.png

接下來是gpu_kernel_impl的實現(xiàn), 可以看到是調(diào)用了launch_vectorized_kernel方法.

15-gpu_kernel_impl.png

再往下就會調(diào)用到核函數(shù)了.

2023-05-10 14-03-06 的屏幕截圖.png

可以看到這里就已經(jīng)走到了調(diào)用核函數(shù)的位置了.

核函數(shù).png

從注釋可以看出 該部分代碼會判斷一下block的大小, 來決定是否對數(shù)據(jù)進行展開.

該核函數(shù)調(diào)用了另一個在device上調(diào)用執(zhí)行的方法

image.png

policy_t的注釋這樣寫的:

image.png

這段代碼看起來就是根據(jù)Tensor的數(shù)據(jù)類型不同來按照不同的方式加載到args數(shù)組中, 然后在GPU上調(diào)用傳入的elementwise的f函數(shù), 并把args作為入?yún)? 這樣Tensor的加法就計算完成了.