How do I call BLAS/Lapack in a nopython block?

Numba Support: How do I do ...?
1

I would like to call LAPACK dgetrf from a nopython block.

In the numba repository, there are wrappers of LAPACK and BLAS functions at numba/_lapack.c. I am looking at numba/numba/np/linalg.py (particularly the implementation of inv) for examples of how to use these wrappers, but I am not able to replicate the results of their code.

Here is a small snippet that computes r which is of Type Signature. I am not sure how to actually call the function numba_xxgetrf. I am also not sure how to wrap this code block in a nopython block. Attempting to place all the code below in a nopython block gives an error message that _LAPACK() is not recognized.

import numpy as np
import numba.np.linalg as la

A = np.ones((10,10))
numba_xxgetrf = la._LAPACK().numba_xxgetrf(A.dtype)
n = A.shape[-1]
acpy = np.asfortranarray(A) # convert to Fortran array to call Lapack
ipiv = np.empty(n, dtype=F_INT_nptype) # pivoting vector
r = numba_xxgetrf(‘d’, n, n, acpy.ctypes, n, ipiv.ctypes)

2

Hi @annayesy,

I think the issue you are probably running in to is that Numba can’t work out how to resolve complex module calls like la._LAPACK().numba_xxgetrf(A.dtype) within compiled code, also that this is asking for typing information A.dtype to be present in compiled code. To deal with this sort of thing it’s typical to resolve the complicated lookup based on Numba’s type system outside of the compiled function and then just reference that in the compiled function itself.

Couple of ways to make a dgetrf call:

  1. Do something like this that uses a load of Numba’s internals from numba.np.linalg (be aware that this is not public API!) and also makes use of the public numba.extending API for registering an overload of a dummy function my_dgetrf_call, @overload reference guide is here. This should “just work”:

    from numba import njit, types
from numba.np.linalg import (ensure_lapack, _check_finite_matrix,
                            _check_linalg_matrix, _LAPACK,
                            _copy_to_fortran_order, F_INT_nptype,
                            _inv_err_handler, get_blas_kind,
                            _dummy_liveness_func)
from numba.extending import overload
import numpy as np

# This is a mock function to overload
def my_dgetrf_call(a):
    pass

@overload(my_dgetrf_call)
def ol_my_dgetrf_call(a):

    if not isinstance(a, types.Array):
        return None # needs to be an array

    if not isinstance(a.dtype, types.Float):
        if a.dtype.bitwidth != 64:
            return None # needs to be an array

    ensure_lapack()

    _check_linalg_matrix(a, "my_dgetrf_call")

    numba_xxgetrf = _LAPACK().numba_xxgetrf(a.dtype)

    kind = ord(get_blas_kind(a.dtype, "my_dgetrf_call"))

    def impl(a):
        n = a.shape[-1]
        if a.shape[-2] != n:
            msg = "Last 2 dimensions of the array must be square."
            raise np.linalg.LinAlgError(msg)

        _check_finite_matrix(a)

        acpy = _copy_to_fortran_order(a)

        ipiv = np.empty(n, dtype=F_INT_nptype)

        r = numba_xxgetrf(kind, n, n, acpy.ctypes, n, ipiv.ctypes)
        _inv_err_handler(r)

        # help liveness analysis
        _dummy_liveness_func([acpy.size, ipiv.size])
        return acpy, ipiv

    return impl


@njit
def foo(a):
    return my_dgetrf_call(a)

np.random.seed(0)
n = 4
x = np.random.random((n, n))
a = np.dot(x, x.T)

print(foo(a))

  2. There’s a demo of how to do this sort of thing using public APIs for LAPACK’s zgees here: How to include LAPACK functions with complex inputs · Issue #5301 · numba/numba · GitHub, hopefully that’s enough to go on to allow you to adapt something for your use case.

Hope this helps?

3

Thank you, Stuart!!
As you said, solution (1) “just works,” and it works great!

These are super useful resources. I’ve been trying to reproduce the numba.np.linalg code for a few hours, and what you said about the Numba type system look-up makes sense.

Thank you again for the quick and very useful answer!

4

@annayesy no problem, glad you have something working now :slight_smile: