一个简单的torch-cuda demo

写一个简单CUDA demo的尝试1

功能

对于一维tensor a, b，实现实现a[i]+b[i]。

参考: https://zhuanlan.zhihu.com/p/595851188

需要的代码文件

test.py: 测试文件
AddMlp/: 算子仓库

• add_mlp_cuda.cu: 算子核心操作，包括一个cuda计算kernel和一个launcher来启动cuda kernel的并行计算
• add_mlp.cpp: laucher的函数声明，对封装.cuda中的算子并为python提供调用的方法接口（如forward函数）
• setup.py: 封装.cu和.cpp，为python提供调用接口的算子实例/类
• addMlp.py: 在python端调用C接口并封装、并实现backward方法
• init.py: addMlp.py的init文件

封装/调用链条

.cu里面的add_mlp_kernel -> .cu里面的add_mlp_kernel -> .cpp 里面的add_mlp_launcher定义及add_mlp_gpu并绑定forward方法->setup.py绑定.cu和.cpp->addMlp.py中AddMLP类调用add_mlp.forward并绑定为add_mlp_op->init.py绑定->test.py调用add_mlp_op

打包与安装

因为cuda cpp的code和python是靠setup.py连接起来的，所以必须要用pip install .安装后才能使用。值得注意的是setup.py里的name只是包名(例如scikit-learn)，import的时候是import带有__init__.py的那个folder的名字，比如这里就是from AddMlp import add_mlp_op.

backward

cuda backward的输入是上一步的gradient，这一步的gradient怎么求是要自己推好了写在backward函数里的，比如是常数那就直接返回常数。

代码

https://github.com/JiuTongBro/CudaDemo1