准备工作:

要对Pytorch的c++代码进行比较方便的debug, 需要先做完以下几件事儿:

1. 购买NVIDIA显卡装在自己的电脑上. (需要会拆装显卡, 连接电源线之类的操作)

2. 安装nvidia显卡驱动, 安装cuda, 安装cudnn. 这里会遇到显卡驱动版本/显卡本身/cuda/cudnn版本是否匹配的问题. 安装过程中会遇到gcc版本和linux核版本是否一致的问题. (centos的话可以直接yum安装显卡驱动)

3. 下载一份Pytorch源码. 这里会遇到pytorch子仓无法下载的问题(里面有Google的代码仓无法直连, 需要魔法上网, 并且把git的http/socks代理设置为梯子的本地端口号).

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, 这里可能遇到动态库版本不一致的问题(比如说libstdc++.so.6).

6. 给你的linux系统安装图形界面. 再安装一个linux版本的vscode和chrome(方便随时查资料). 这一步可以不做, 但是为了方便代码跳转和阅读的话, 最好还是在linux的图形界面上搞.

可能是因为我的linux是centos7.9, 反正遇到的问题比较多. 也许用ubuntu18.04系统问题会少一点.

如果你对linux上软件的源码编译安装, 软链接替换等操作不熟悉的话, 这个过程就会比较艰难, 以上每一件事儿都可能卡住半天甚至一天的时间. 大概需要一个星期才能搞定所有问题.

如果你有钞能力的话, 以上每一个步骤你都可以在淘宝上找运维装机服务来解决, 估计500块钱之内可以解决.

正文

CPU上是怎么调用过来的 跟另一篇博客大同小异, 但因为版本不同(我的torch版本为: 2.1.0a0+gitdc4a253 ), 代码逻辑有所变化.

首先是从python_variable_methods这个文件中, 这个文件主要做的应该是c++和python互操作的功能, 就是调用python对象中加法绑定的c++函数.

1-python_variable_methods.png

然后进入THPVariable_add方法, 这个进行一个操作数的类型判断, 调用2565行的Tensor的add方法.

2-THPVariable_add.png

进入ops的add_Tensor方法

3-TensorBody.png

调用939行创建的op的call方法.

4-Operators_2.png

这里是比较关键的角色出场了, Dispatcher, 所有的算子应该都是注册到这个Dispatcher中的.

5-TypedOperatorHandle.png

在Dispatcher的调度下, 根据632行op的dispatchKeySet找到对应的op. 然后进行调用.

6-Dispatcher__call.png

这部分代码看的云里雾里的, 应该都是一些为了让代码简洁而写的封装代码吧

7-KernelFunction__call.png

8-callUnboxedKernelFunction.png

9-wrap_kernel_functor_unboxed_.png

10-WrapFunctionIntoFunctor_.png

经过一系列云里雾里的绕绕, 进入到这个wrapper_CUDA_add_Tensor.cu这个文件, 这个就是cuda代码的文件了.
在op.impl方法中进行实际的调用

11-wrapper_CUDA_add_Tensor.png

对于Tensor加法来说, 这个op.impl就是332行的代码, 这里调用add_kernel方法

12-UfuncCUDA_add.png

67行这里是一个宏写法的switch语句, 下面一大片代码只会走其中一个分支. 因为是宏的写法, debug进不去. 也没关系直接看case就好了. 不管哪个case都会调用gpu_kernel这个方法

13-add_kernel.png

下面是gpu_kernel的具体实现, 调用了gpu_kernel_impl方法

14-gpu_kernel.png

接下来是gpu_kernel_impl的实现, 可以看到是调用了launch_vectorized_kernel方法.

15-gpu_kernel_impl.png

再往下就会调用到核函数了.

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

可以看到这里就已经走到了调用核函数的位置了.

核函数.png

从注释可以看出 该部分代码会判断一下block的大小, 来决定是否对数据进行展开.

该核函数调用了另一个在device上调用执行的方法

image.png

policy_t的注释这样写的:

image.png

这段代码看起来就是根据Tensor的数据类型不同来按照不同的方式加载到args数组中, 然后在GPU上调用传入的elementwise的f函数, 并把args作为入参. 这样Tensor的加法就计算完成了.