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.

Dynamic Kernel Selection

The kernel used by the Python APIs can be changed at runtime as well. 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.

Kernels must be created from the underlying native representation — you cannot directly pass an instance of csr.CSR.

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: numba.experimental.jitclass.base._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.