以下面這個例子作為教程，實現功能是element-wise add；

（pytorch中想呼叫cuda模組，還是用另外使用C編寫介面指令碼）

第一步：cuda程式設計的原始檔和標頭檔案

// mathutil_cuda_kernel.cu
// 標頭檔案，最後一個是cuda特有的
#include <curand.h>
#include <stdio.h>
#include <math.h>
#include <float.h>
#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 >= size) return;
  int j = i % x; i = i / x;
  int k = i % y;
  a[IDX2D(j,k,y)] += b[k];
}


// 這個函式是與c語言函式連結的介面函式
void broadcast_sum_cuda(float *a,cudaStream_t stream)
{
  int size = x * y;
  cudaError_t err;
  
  // 上面定義的函式
  broadcast_sum_kernel<<<cuda_gridsize(size),BLOCK,stream>>>(a,b,x,size);

  err = cudaGetLastError();
  if (cudaSuccess != err)
  {
    fprintf(stderr,"CUDA kernel failed : %s\n",cudaGetErrorString(err));
    exit(-1);
  }
}

#ifndef _MATHUTIL_CUDA_KERNEL
#define _MATHUTIL_CUDA_KERNEL

#define IDX2D(i,j,dj) (dj * i + j)
#define IDX3D(i,dj,dk) (IDX2D(IDX2D(i,dj),dk))

#define BLOCK 512
#define MAX_STREAMS 512

#ifdef __cplusplus
extern "C" {
#endif

void broadcast_sum_cuda(float *a,cudaStream_t stream);

#ifdef __cplusplus
}
#endif

#endif

第二步：C程式設計的原始檔和標頭檔案（介面函式）

// mathutil_cuda.c
// THC是pytorch底層GPU庫
#include <THC/THC.h>
#include "mathutil_cuda_kernel.h"

extern THCState *state;

int broadcast_sum(THCudaTensor *a_tensor,THCudaTensor *b_tensor,int y)
{
  float *a = THCudaTensor_data(state,a_tensor);
  float *b = THCudaTensor_data(state,b_tensor);
  cudaStream_t stream = THCState_getCurrentStream(state);

  // 這裡呼叫之前在cuda中編寫的介面函式
  broadcast_sum_cuda(a,stream);

  return 1;
}

int broadcast_sum(THCudaTensor *a_tensor,int y);

第三步：編譯，先編譯cuda模組，再編譯介面函式模組（不能放在一起同時編譯）

nvcc -c -o mathutil_cuda_kernel.cu.o mathutil_cuda_kernel.cu -x cu -Xcompiler -fPIC -arch=sm_52

import os
import torch
from torch.utils.ffi import create_extension

this_file = os.path.dirname(__file__)

sources = []
headers = []
defines = []
with_cuda = False

if torch.cuda.is_available():
  print('Including CUDA code.')
  sources += ['src/mathutil_cuda.c']
  headers += ['src/mathutil_cuda.h']
  defines += [('WITH_CUDA',None)]
  with_cuda = True

this_file = os.path.dirname(os.path.realpath(__file__))

extra_objects = ['src/mathutil_cuda_kernel.cu.o']  # 這裡是編譯好後的.o檔案位置
extra_objects = [os.path.join(this_file,fname) for fname in extra_objects]


ffi = create_extension(
  '_ext.cuda_util',headers=headers,sources=sources,define_macros=defines,relative_to=__file__,with_cuda=with_cuda,extra_objects=extra_objects
)

if __name__ == '__main__':
  ffi.build()

第四步：呼叫cuda模組

from _ext import cuda_util #從對應路徑中呼叫編譯好的模組

a = torch.randn(3,5).cuda()
b = torch.randn(3,1).cuda()
mathutil.broadcast_sum(a,*map(int,a.size()))

# 上面等價於下面的效果：

a = torch.randn(3,5)
b = torch.randn(3,1)
a += b

以上就是本文的全部內容，希望對大家的學習有所幫助，也希望大家多多支援我們。