云南学校 手机网站建设北京优化网站推广
写一个简单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