Universal Functions#

Numba provides a set of decorators to create NumPy universal functions-like routines that are JIT compiled. Although, a close analog to NumPy universal functions Numba’s @vectorize are not fully compatible with a regular NumPy ufunc.. Refer Creating NumPy universal functions for details.

Numba-dpex only supports numba.vectorize decorator and not yet the numba.guvectorize decorator. Another present limitation is that numba-dpex ufunc kernels cannot invoke a numba_dpex.kernel function. .. Ongoing work is in progress to address these limitations.

Example 1: Basic Usage#

Full example can be found at numba_dpex/examples/vectorize.py.

@vectorize(nopython=True)
def ufunc_kernel(x, y):
    return x + y

def test_njit():
    N = 10
    dtype = np.float64

    A = np.arange(N, dtype=dtype)
    B = np.arange(N, dtype=dtype) * 10

    # Use the environment variable SYCL_DEVICE_FILTER to change the default device.
    # See https://github.com/intel/llvm/blob/sycl/sycl/doc/EnvironmentVariables.md#sycl_device_filter.
    device = dpctl.select_default_device()
    print("Using device ...")
    device.print_device_info()

    with dpctl.device_context(device):
        C = ufunc_kernel(A, B)

    print(C)

    print("Done...")

Example 2: Calling numba.vectorize inside a numba_dpex.kernel#

Full example can be found at numba_dpex/examples/blacksholes_njit.py.

Note

numba.cuda requires target='cuda' parameter for numba.vectorize and numba.guvectorize functions. Numba-dpex eschews the target parameter for @vectorize and infers the target from the dpctl.device_context in which the numba.vectorize function is called.

Full Examples#

  • numba_dpex/examples/vectorize.py

  • numba_dpex/examples/blacksholes_njit.py