pytorch调用手写CUDA算子和pybind11的使用
之前介绍了很多CUDA编写算子的代码，但是一直缺乏一个好的方法来证明自己手写算子的正确性，以及希望知道自己的手写算子在时间上和pytorch的差异，这里我们需要用到pybind11这个工具，这个工具可以将我们手写的CUDA代码编译为动态库，最终形成一个崭新的module，使得pytorch可以调用我们手写的算子进行计算。下面我们先介绍一下代码结构：
代码架构
-
- softmax.cu
-
- py::array_t < < <float > > >
- request()方法
- .ptr
- 总结
- setup.py
-
- 导入模块
- 设置python包
- 定义扩展模块
- 指定构建命令
- 使用方法
- test_softmax.py
- run.sh
- tensor代替numpy
-
- softmax.cu
- test_softmax.py
softmax.cu
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>