Created
September 19, 2016 14:14
-
-
Save sklam/702af16db009caf66dc4bff2c190ebd7 to your computer and use it in GitHub Desktop.
Numba using NVCC device function
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
""" | |
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) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
/* | |
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; | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment