pytorch / extension-ffi

Examples of C extensions for PyTorch
258 stars 70 forks source link

CUDA kernel #4

Closed ikostrikov closed 7 years ago

ikostrikov commented 7 years ago

When I try to add a simple cuda kernel, for example:

__global__ void fill_array(int* array, int value, int length) {
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    if (x < length) {
    array[x] = value;
    }
}

It produces the following error: /home/kostrikov/extension-ffi/script/src/my_lib_cuda.c: In function ‘fill_array’: /home/kostrikov/extension-ffi/script/src/my_lib_cuda.c:8:10: error: ‘blockIdx’ undeclared (first use in this function) int x = blockIdx.x blockDim.x + threadIdx.x; ^ /home/kostrikov/extension-ffi/script/src/my_lib_cuda.c:8:10: note: each undeclared identifier is reported only once for each function it appears in /home/kostrikov/extension-ffi/script/src/my_lib_cuda.c:8:23: error: ‘blockDim’ undeclared (first use in this function) int x = blockIdx.x blockDim.x + threadIdx.x; ^ /home/kostrikov/extension-ffi/script/src/my_lib_cuda.c:8:36: error: ‘threadIdx’ undeclared (first use in this function) int x = blockIdx.x * blockDim.x + threadIdx.x;

jekbradbury commented 7 years ago

A CUDA kernel must be in a .cu file and compiled with nvcc.

ikostrikov commented 7 years ago

Thanks!

Is there an example how to integrate it with ffi?

jekbradbury commented 7 years ago

I haven't seen one; also unless you're building a complex C extension this is probably easier https://gist.github.com/szagoruyko/dccce13465df1542621b728fcc15df53

ikostrikov commented 7 years ago

Yes, I think this one will work.

Thanks a lot for the link!

zhanghang1989 commented 7 years ago

Hi, just in case someone may need it. I have an example using Torch generic files with CUDA backend in this repo https://github.com/zhanghang1989/PyTorch-Encoding