Calling foreign functions from Python kernels

Python kernels can call device functions written in other languages. CUDA C/C++, PTX, and binary objects (cubins, fat binaries, etc.) are directly supported; sources in other languages must be compiled to PTX first. The constituent parts of a Python kernel call to a foreign device function are:

• The device function implementation in a foreign language (e.g. CUDA C).

• A declaration of the device function in Python.

• A kernel that links with and calls the foreign function.

Device function ABI

Numba’s ABI for calling device functions defines the following prototype in C/C++:

extern "C"
__device__ int
function(
  T* return_value,
  ...
);

Components of the prototype are as follows:

• extern "C" is used to prevent name-mangling so that it is easy to declare the function in Python. It can be removed, but then the mangled name must be used in the declaration of the function in Python.

• __device__ is required to define the function as a device function.

• The return value is always of type int, and is used to signal whether a Python exception occurred. Since Python exceptions don’t occur in foreign functions, this should always be set to 0 by the callee.

• The first argument is a pointer to the return value of type T, which is allocated in the local address space [1] and passed in by the caller. If the function returns a value, the pointee should be set by the callee to store the return value.

• Subsequent arguments should match the types and order of arguments passed to the function from the Python kernel.

Functions written in other languages must compile to PTX that conforms to this prototype specification.

A function that accepts two floats and returns a float would have the following prototype:

extern "C"
__device__ int
mul_f32_f32(
  float* return_value,
  float x,
  float y
);

Notes

[1]

Care must be taken to ensure that any operations on the return value are applicable to data in the local address space. Some operations, such as atomics, cannot be performed on data in the local address space.

Declaration in Python

To declare a foreign device function in Python, use declare_device():

numba.cuda.declare_device(name, sig)

Declare the signature of a foreign function. Returns a descriptor that can be used to call the function from a Python kernel.

Parameters:

• name (str) – The name of the foreign function.

• sig – The Numba signature of the function.

The returned descriptor name need not match the name of the foreign function. For example, when:

mul = cuda.declare_device('mul_f32_f32', 'float32(float32, float32)')

is declared, calling mul(a, b) inside a kernel will translate into a call to mul_f32_f32(a, b) in the compiled code.

Passing pointers

Numba’s calling convention requires multiple values to be passed for array arguments. These include the data pointer along with shape, stride, and other information. This is incompatible with the expectations of most C/C++ functions, which generally only expect a pointer to the data. To align the calling conventions between C device code and Python kernels it is necessary to declare array arguments using C pointer types.

For example, a function with the following prototype:

numba/cuda/tests/doc_examples/ffi/functions.cu

extern "C"
__device__ int
sum_reduce(
  float* return_value,
  float* array,
  int n
);

would be declared as follows:

from test_ex_from_buffer in numba/cuda/tests/doc_examples/test_ffi.py

signature = 'float32(CPointer(float32), int32)'
sum_reduce = cuda.declare_device('sum_reduce', signature)

To obtain a pointer to array data for passing to foreign functions, use the from_buffer() method of a cffi.FFI instance. For example, a kernel using the sum_reduce function could be defined as:

from test_ex_from_buffer in numba/cuda/tests/doc_examples/test_ffi.py

import cffi
ffi = cffi.FFI()

@cuda.jit(link=[functions_cu])
def reduction_caller(result, array):
    array_ptr = ffi.from_buffer(array)
    result[()] = sum_reduce(array_ptr, len(array))

where result and array are both arrays of float32 data.

Linking and Calling functions

The link keyword argument of the @cuda.jit decorator accepts a list of file names specified by absolute path or a path relative to the current working directory. Files whose name ends in .cu will be compiled with the NVIDIA Runtime Compiler (NVRTC) and linked into the kernel as PTX; other files will be passed directly to the CUDA Linker.

For example, the following kernel calls the mul() function declared above with the implementation mul_f32_f32() in a file called functions.cu:

@cuda.jit(link=['functions.cu'])
def multiply_vectors(r, x, y):
    i = cuda.grid(1)

    if i < len(r):
        r[i] = mul(x[i], y[i])

C/C++ Support

Support for compiling and linking of CUDA C/C++ code is provided through the use of NVRTC subject to the following considerations:

• It is only available when using the NVIDIA Bindings. See NUMBA_CUDA_USE_NVIDIA_BINDING.

• A suitable version of the NVRTC library for the installed version of the NVIDIA CUDA Bindings must be available.

• The CUDA include path is assumed by default to be /usr/local/cuda/include on Linux and $env:CUDA_PATH\include on Windows. It can be modified using the environment variable NUMBA_CUDA_INCLUDE_PATH.

• The CUDA include directory will be made available to NVRTC on the include path; additional includes are not supported.

Complete Example

This example demonstrates calling a foreign function written in CUDA C to multiply pairs of numbers from two arrays.

The foreign function is written as follows:

numba/cuda/tests/doc_examples/ffi/functions.cu

// Foreign function example: multiplication of a pair of floats

extern "C" __device__ int
mul_f32_f32(
  float* return_value,
  float x,
  float y)
{
  // Compute result and store in caller-provided slot
  *return_value = x * y;

  // Signal that no Python exception occurred
  return 0;
}

The Python code and kernel are:

from test_ex_linking_cu in numba/cuda/tests/doc_examples/test_ffi.py

from numba import cuda
import numpy as np
import os

# Declaration of the foreign function
mul = cuda.declare_device('mul_f32_f32', 'float32(float32, float32)')

# Path to the source containing the foreign function
# (here assumed to be in a subdirectory called "ffi")
basedir = os.path.dirname(os.path.abspath(__file__))
functions_cu = os.path.join(basedir, 'ffi', 'functions.cu')

# Kernel that links in functions.cu and calls mul
@cuda.jit(link=[functions_cu])
def multiply_vectors(r, x, y):
    i = cuda.grid(1)

    if i < len(r):
        r[i] = mul(x[i], y[i])

# Generate random data
N = 32
np.random.seed(1)
x = np.random.rand(N).astype(np.float32)
y = np.random.rand(N).astype(np.float32)
r = np.zeros_like(x)

# Run the kernel
multiply_vectors[1, 32](r, x, y)

# Sanity check - ensure the results match those expected
np.testing.assert_array_equal(r, x * y)

Note

The example above is minimal in order to illustrate a foreign function call - it would not be expected to be particularly performant due to the small grid and light workload of the foreign function.