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:
If the environment variable
CSR_KERNEL
is set, its value is used as the kernel name.If the MKL kernel is compiled and MKL is available, it is used.
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 byget_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_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.