Numba using NVCC device function

cuda_extern.py

"""
Demonstrating CUDA JIT integration
"""
from __future__ import print_function
from numba import cuda
import numpy
import os

# Declare function to link to
bar = cuda.declare_device('bar', 'int32(voidptr, int32)')

# Get path to precompiled library
curdir = os.path.join(os.path.dirname(__file__))
link = os.path.join(curdir, 'cufunc.o')
print("Linking: %s", link)

# Code that uses CUDA JIT
@cuda.jit('void(int32[:], int32[:])', link=[link])
def foo(inp, out):
	i = cuda.grid(1)
	# we will pass the array pointer as void* (voidptr)
	out[i] = bar(inp[i:].ctypes.data, 2)

print(foo.ptx)


n = 5
inp = numpy.arange(n, dtype='int32')
out = numpy.zeros_like(inp)
foo[1, out.size](inp, out)

print("inp =", inp)
print("out =", out)

cufunc.cu

/*
Numba requires return value to be passed as
a pointer in the first argument.


To compile:

	nvcc -arch=sm_20 -dc cufunc.cu -o cufunc.o

*/
#include <cstdio>

extern "C" {

/* Note: numba uses a different signature.
   The following is equivalent to the numba signature
    int(voidptr, int)
   The numba return value is the first arg.
   The C return value is for error status.
*/
__device__
int bar(int* retval, int *a, int b){
	/* Fill this function with anything */
    *retval = a[0] + b;
	/* Return 0 to indicate success */
	return 0;
}

}