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 
#include 

你可能感兴趣的:(高性能计算,pytorch,人工智能,python,1024程序员节)