Compiling an Extension with CUDA files

longcw: <SCODE>#include <THC/THC.h>

extern THCState *state;

int my_lib_add_forward_cuda(THCudaTensor *input, THCudaTensor *output)
{
  float * pinput = THCudaTensor_data(state, input);
  float * poutput = THCudaTensor_data(state, output);
  for(...)
  {
    poutput[i] = do_something(pinput[i]);
  }
  return 1;
}
<ECODE>

smth: No, you cannot access CUDA memory like this.

longcw: Thank you for your reply.

longcw: 

apaszke: 

longcw: Thank you very much. It is helpful.

lupoglaz: Hello, I have a question about compiling an extension with cuda files. I get the following error: <SCODE>distutils.errors.UnknownFileError: unknown file type '.cu'
<ECODE> when I include .cu files in the sources: <SCODE>import os
import torch
import glob
from torch.utils.ffi import create_extension

this_file = os.path.dirname(__file__)

sources = ['../Library/~.cpp',
            '../Library/~.cpp',
            '../Library/~.cu']
headers = ['../Library/~.h']
here = os.path.abspath(os.path.dirname(__file__))
lib_dir = os.path.join(here, '..', 'Library')
include_dirs = [
    os.path.join(lib_dir, '~'),
    os.path.join(lib_dir, 'Math'),
]
defines = [('WITH_CUDA', None)]
with_cuda = True

ffi = create_extension(
    '_CUDA.~',
    headers=headers,
    sources=sources,
    define_macros=defines,
    relative_to=__file__,
    with_cuda=with_cuda,
    include_dirs = include_dirs,
    extra_compile_args=["-fopenmp"]
)

if __name__ == '__main__':
    ffi.build()
    from _CUDA import ~
    print ~.__dict__
<ECODE> Am I doing something wrong?

apaszke: 

longcw: <SCODE>#!/usr/bin/env bash

CUDA_PATH=/usr/local/cuda/

cd layers/reorg/src
echo "Compiling reorg layer kernels by nvcc..."
nvcc -c -o reorg_cuda_kernel.cu.o reorg_cuda_kernel.cu -x cu -Xcompiler -fPIC -arch=sm_52

cd ../
python build.py
<ECODE> <SCODE>import os
import torch
from torch.utils.ffi import create_extension


sources = ['src/reorg_cpu.c']
headers = ['src/reorg_cpu.h']
defines = []
with_cuda = False

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

this_file = os.path.dirname(os.path.realpath(__file__))
# print(this_file)
extra_objects = ['src/reorg_cuda_kernel.cu.o']
extra_objects = [os.path.join(this_file, fname) for fname in extra_objects]

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

if __name__ == '__main__':
    ffi.build()
<ECODE> Hope it helps you.

apaszke: Sorry for the inconvenience, we’re going to be fixing that in the future.

lupoglaz: <SCODE>CUDA_ADD_LIBRARY(~ STATIC c~CUDAKernels.cu)
TARGET_LINK_LIBRARIES(~)
<ECODE> and then link them as an extra link argument: <SCODE>...
extra_link_args=['~.a']
ffi = create_extension(
    '~',
    headers=headers,
    sources=sources,
    define_macros=defines,
    relative_to=__file__,
    with_cuda=with_cuda,
    include_dirs = include_dirs,
    extra_compile_args=["-fopenmp"],
    extra_link_args = extra_link_args
)
...
<ECODE> This workaround does work with a static library. I had problems trying to link a shared library this way (haven’t tried enough though).

ThibaultGROUEIX: 

chen_dongdong: 

Jiang_He: 

GriffinLiang: