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.
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.
-
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.
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.
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.