Kernels

More sophisticated CSR operations, such as multiplications, are provided by a kernel. Numba supports multiple kernels; at import time, a default kernel is loaded and is available as csr.kernel. Numba-compiled functions using the CSR kernel should access it through this module.

There are currently three kernels available:

numba

This kernel implements the operations directly with Numba-compiled functions.

mkl

This kernel uses the Intel Math Kernel Library (MKL) to implement the operations.

scipy

This kernel is not intended for general use, and only exists to make testing and benchmarking easier. It uses SciPy’s sparse matrices to implement matrix math, and cannot be used from Numba.

The default kernel is automatically selected when csr is imported. It is selected as follows:

  1. If the environment variable CSR_KERNEL is set, its value is used as the kernel name.

  2. If the MKL kernel is compiled and MKL is available, it is used.

  3. Otherwise, the Numba kernel is used.

The PyPI packages do not include the MKL kernel. CSR only supports MKL when used with Conda; the conda-forge packages for Intel platforms include the MKL kernel.

More kernels may be added in the future; for example, it may be useful to add a CUDA kernel, or if an optimized sparse matrix package for 64-bit ARM becomes available.

There are two ways to access the kernel. One is to import csr.kernel and use the functions directly; this uses the kernel selected at import time and cannot be dynamically changed. Numba-optimized functions must use this access method; pure Python code generally doesn’t need to access kernels directly, but may use either this method or dynamic selection (described in the next section).

Dynamic Kernel Selection

The kernel used by the Python APIs can be changed at runtime as well. The Python implementations of the various methods provided by csr.CSR use this API to access the active kernel, so changing the kernel with set_kernel() will change the kernel for pure Python code using CSR and its methods, but will not change the code used for Numba-based code. Dynamic kernel selection is primarily useful for testing and benchmarking.

Note

This does not change the kernel exposed as csr.kernel, which is used by Numba-compiled client functions.

csr.kernels.get_kernel(name=None)

Get a kernel.

Parameters:

name (str or None) – The name of the kernel. If None, returns the current default kernel.

csr.kernels.set_kernel(name)

Set the default kernel. It is very rare to need to use this — letting CSR select its default kernel, or configuring the kernel through the CSR_KERNEL environment variable, is the best option for the vast majority of applications. This is here primarily to enable test code to switch kernels.

This function does not change the kernel exposed by importing csr.kernel, which is typically used by compiled Numba functions. It only changes the kernel returned by get_kernel() and used by the pure-Python APIs,

Parameters:

name (str) – The name of the kernel.

csr.kernels.use_kernel(name)

Context manager to run code with a specified (thread-local) kernel. It calls set_kernel(), and restores the previously-active kernel when the context exits.

Kernel Interface

The csr.kernel module exposes the kernel interface. These same functions are available on any kernel, including those returned by csr.kernels.get_kernel().

Handles

The kernel interface is built on opaque handles: a csr._CSR needs to be converted to a handle with to_handle(), and subsequent operations use that handle to access the matrix. Handles must be explicitly released, or they will generally leak memory.

csr.kernel.to_handle(csr)

Convert a native CSR to a handle. The caller must arrange for the CSR last at least as long as the handle. The handle must be explicitly released.

Handles are opaque as far as callers are concerned.

csr.kernel.from_handle(h)

Convert a handle to a CSR. The handle may be released after this is called.

csr.kernel.release_handle(h)

Release a handle.

Manipulation

csr.kernel.order_columns(h)

Sort matrix rows in increasing column order.

Multiplication

csr.kernel.mult_vec(h: CSR, v)
csr.kernel.mult_ab(a_h, b_h)

Multiply matrices A and B.

Parameters:
  • a_h – the handle of matrix A

  • b_h – the handle of matrix B

Returns:

the handle of the product; it must be released when no longer needed.

csr.kernel.mult_abt(a_h, b_h)

Multiply matrices A and B^T.

Parameters:
  • a_h – the handle of matrix A

  • b_h – the handle of matrix B

Returns:

the handle of the product; it must be released when no longer needed.

Additional Requirements

There are additional requirements for kernel implementations:

csr.kernel.max_nnz: int

This attribute stores the maximum number of non-zero entries supported by this kernel. It is not exposed on the kernel module, but is exposed on dynamically-obtained kernels, and must be provided by any kernel implementation. It is used by csr._CSR to automatically handle matrices that are too large for a particular kernel when possible.