Ctrl+K

Writing Data Parallel Kernels#

Kernel Declaration#

A kernel function is a device function that is meant to be called from host code, where a device can be any SYCL supported device such as a GPU, CPU, or an FPGA. The main characteristics of a kernel function are:

  • Scalars must be passed as an array. Kernels operate with dpnp array arguments only. If your want a scalar argument, then represent it as 0-dimensional dpnp array.

Note

Please refer to Data Parallel Extension for Numpy* to learn more about dpnp.

  • Kernels cannot explicitly return a value. All result data must be written to dpnp array passed as a function’s argument.

Here is an example of a kernel that computes sum of two vectors a and b. Arguments are two input vectors a and b and one output vector c for storing the result of vector summation:

EXAMPLE: Data parallel kernel implementing the vector sum a+b#
import numba_dpex as ndpx

# Data parallel kernel implementing vector sum
@ndpx.kernel
def kernel_vector_sum(a, b, c):
    i = ndpx.get_global_id(0)
    c[i] = a[i] + b[i]

Kernel Invocation#

The kernel launch parameter syntax for specifying global and local sizes are similar to SYCL’s range and ndrange classes. The global and local sizes need to be specified with numba_dpex’s Range and NdRange classes.

For example, below is a following kernel that computes a sum of two vectors:

EXAMPLE: A vector sum kernel#
import numba_dpex as ndpx

# Data parallel kernel implementing vector sum
@ndpx.kernel
def kernel_vector_sum(a, b, c):
    i = ndpx.get_global_id(0)
    c[i] = a[i] + b[i]

If the global size parameter is needed to run, it could be like this (where global_size is an int):

EXAMPLE: A vector sum kernel with a global size/range#
import numba_dpex as ndpx

# Utility function for printing and testing
def driver(a, b, c, global_size):
    kernel_vector_sum[ndpx.Range(global_size)](a, b, c)
    a_np = dpnp.asnumpy(a)  # Copy dpnp array a to NumPy array a_np
    b_np = dpnp.asnumpy(b)  # Copy dpnp array b to NumPy array b_np
    c_np = dpnp.asnumpy(c)  # Copy dpnp array c to NumPy array c_np
    testing.assert_equal(c_np, a_np + b_np)

If both local and global ranges are needed, they can be specified using two instances of Range inside an NdRange object. For example, below is a kernel to compute pair-wise Euclidean distances of n-dimensional data points:

EXAMPLE: A kernel to compute pair-wise Euclidean distances#
import numba_dpex as ndpx

@ndpx.kernel
def pairwise_distance(X, D, xshape0, xshape1):
    """
    An Euclidean pairwise distance computation implemented as
    a ``kernel`` function.
    """
    idx = ndpx.get_global_id(0)

    d0 = X[idx, 0] - X[idx, 0]
    # for i in range(xshape0):
    for j in range(X.shape[0]):
        d = d0
        for k in range(X.shape[1]):
            tmp = X[idx, k] - X[j, k]
            d += tmp * tmp
        D[idx, j] = sqrt(d)

Now the local and global sizes can be specified as follows (here both args.n and args.l are int):

EXAMPLE: A kernel to compute pair-wise Euclidean distances with a global and a local size/range#
import numba_dpex as ndpx

# Global work size is equal to the number of points
global_size = ndpx.Range(args.n)
# Local Work size is optional
local_size = ndpx.Range(args.l)

def driver():
    # measure running time
    times = list()
    for repeat in range(args.r):
        start = time()
        pairwise_distance[ndpx.NdRange(global_size, local_size)](
            X, D, X.shape[0], X.shape[1]
        )
        end = time()

        total_time = end - start
        times.append(total_time)

    return times

Kernel Indexing Functions#

In data parallel kernel programming all work items are enumerated and accessed by their index. You will use numba_dpex.get_global_id() function to get the index of a current work item from the kernel. The total number of work items can be determined by calling numba_dpex.get_global_size() function.

The work group size can be determined by calling numba_dpex.get_local_size() function. Work items in the current work group are accessed by calling numba_dpex.get_local_id().

The total number of work groups are determined by calling numba_dpex.get_num_groups() function. The current work group index is obtained by calling numba_dpex.get_group_id() function.

On this page