pytorch 如何调用cuda_pytorch中使用cuda扩展的实现示例

以下面这个例子作为教程，实现功能是element-wise add；

(pytorch中想调用cuda模块，还是用另外使用C编写接口脚本)

第一步：cuda编程的源文件和头文件

// mathutil_cuda_kernel.cu

// 头文件，最后一个是cuda特有的

#include

#include

#include

#include

#include "mathutil_cuda_kernel.h"

// 获取GPU线程通道信息

dim3 cuda_gridsize(int n)

{

int k = (n - 1) / BLOCK + 1;

int x = k;

int y = 1;

if(x > 65535) {

x = ceil(sqrt(k));

y = (n - 1) / (x * BLOCK) + 1;

}

dim3 d(x, y, 1);

return d;

}

// 这个函数是cuda执行函数，可以看到细化到了每一个元素

__global__ void broadcast_sum_kernel(float *a, float *b, int x, int y, int size)

{

int i = (blockIdx.x + blockIdx.y * gridDim.x) * blockDim.x + threadIdx.x;

if(i >= s