numpy.linalg.inv very slow for a stack of 3x3 arrays

[bhaveshshrimali]

numpy.linalg.inv seems to be significantly slower than a hard-coded version for a simple test case, a (1000,4,3,3,) array.
Any pointers as to what could be going off and possible remedies, if any, to make it faster. See the following MWE:

Reproducing code example:

from numpy.linalg import inv as npinv
from numpy import zeros_like, einsum, random
from timeit import timeit

# Helper function for determinant
def vdet(A):
    detA = zeros_like(A[0, 0])
    detA = A[0, 0] * (A[1, 1] * A[2, 2] - A[1, 2] * A[2, 1]) -\
           A[0, 1] * (A[2, 2] * A[1, 0] - A[2, 0] * A[1, 2]) +\
        A[0, 2] * (A[1, 0] * A[2, 1] - A[2, 0] * A[1, 1])
    return detA

# Another function for computing inverse
# using the Cayley-Hamilton corollary
def finv(A):
    detA = vdet(A)
    I1 = einsum("ii...", A)
    I2 = -0.5 * (einsum("ik...,ki...", A, A) - I1**2)
    Asq = einsum("ik...,kj...->ij...", A, A)
    eye = zeros_like(A)
    eye[0, 0] = 1.
    eye[1, 1] = 1.
    eye[2,2] = 1.
    return 1./detA * (Asq - I1 * A + I2 * eye)

# Hard coded inverse
def hdinv(A):
    invA = zeros_like(A)
    detA = vdet(A)

    invA[0, 0] = (-A[1, 2] * A[2, 1] +
                  A[1, 1] * A[2, 2]) / detA
    invA[1, 0] = (A[1, 2] * A[2, 0] -
                  A[1, 0] * A[2, 2]) / detA
    invA[2, 0] = (-A[1, 1] * A[2, 0] +
                  A[1, 0] * A[2, 1]) / detA
    invA[0, 1] = (A[0, 2] * A[2, 1] -
                  A[0, 1] * A[2, 2]) / detA
    invA[1, 1] = (-A[0, 2] * A[2, 0] +
                  A[0, 0] * A[2, 2]) / detA
    invA[2, 1] = (A[0, 1] * A[2, 0] -
                  A[0, 0] * A[2, 1]) / detA
    invA[0, 2] = (-A[0, 2] * A[1, 1] +
                  A[0, 1] * A[1, 2]) / detA
    invA[1, 2] = (A[0, 2] * A[1, 0] -
                  A[0, 0] * A[1, 2]) / detA
    invA[2, 2] = (-A[0, 1] * A[1, 0] +
                  A[0, 0] * A[1, 1]) / detA
    return invA

Timings

F = random.random((3,3,1000,4))
F2 = einsum("ij...->...ij", F)

%timeit hdinv(F) # 371 µs ± 27 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit finv(F) # 5.35 ms ± 661 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit npinv(F2) # 79.2 ms ± 12.9 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Numpy/Python version information:

python version

3.7.6 | packaged by conda-forge | (default, Jun 1 2020, 18:57:50) \n[GCC 7.5.0]

numpy version

1.19.0

[eric-wieser]

For what it's worth, the implementation is here:

numpy/numpy/linalg/umath_linalg.c.src

Lines 1692 to 1724 in 97d2db4

	static void
	@TYPE@_inv(char **args, npy_intp const *dimensions, npy_intp const *steps,
	void *NPY_UNUSED(func))
	{
	GESV_PARAMS_t params;
	fortran_int n;
	int error_occurred = get_fp_invalid_and_clear();
	INIT_OUTER_LOOP_2
	n = (fortran_int)dimensions[0];
	if (init_@lapack_func@(&params, n, n)) {
	LINEARIZE_DATA_t a_in, r_out;
	init_linearize_data(&a_in, n, n, steps[1], steps[0]);
	init_linearize_data(&r_out, n, n, steps[3], steps[2]);
	BEGIN_OUTER_LOOP_2
	int not_ok;
	linearize_@TYPE@_matrix(params.A, args[0], &a_in);
	identity_@TYPE@_matrix(params.B, n);
	not_ok = call_@lapack_func@(&params);
	if (!not_ok) {
	delinearize_@TYPE@_matrix(args[1], params.B, &r_out);
	} else {
	error_occurred = 1;
	nan_@TYPE@_matrix(args[1], &r_out);
	}
	END_OUTER_LOOP
	release_@lapack_func@(&params);
	}
	set_fp_invalid_or_clear(error_occurred);
	}

It looks like we're doing unecessary copies with linearize_@TYPE@_matrix in cases where we don't need to. Edit: We likely do need to, because lapack dgesv operates in-place.

Do you know how you installed numpy? Is it using an accelerated linalg backend?

[bhaveshshrimali]

Here is the output of numpy.show_config()

blas_mkl_info:
    libraries = ['blas', 'cblas', 'lapack', 'pthread', 'blas', 'cblas', 'lapack']
    library_dirs = ['/srv/conda/envs/notebook/lib']
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ['/srv/conda/envs/notebook/include']
blas_opt_info:
    libraries = ['blas', 'cblas', 'lapack', 'pthread', 'blas', 'cblas', 'lapack', 'blas', 'cblas', 'lapack']
    library_dirs = ['/srv/conda/envs/notebook/lib']
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ['/srv/conda/envs/notebook/include']
lapack_mkl_info:
    libraries = ['blas', 'cblas', 'lapack', 'pthread', 'blas', 'cblas', 'lapack']
    library_dirs = ['/srv/conda/envs/notebook/lib']
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ['/srv/conda/envs/notebook/include']
lapack_opt_info:
    libraries = ['blas', 'cblas', 'lapack', 'pthread', 'blas', 'cblas', 'lapack', 'blas', 'cblas', 'lapack']
    library_dirs = ['/srv/conda/envs/notebook/lib']
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ['/srv/conda/envs/notebook/include']

[eric-wieser]

It looks like we could compute inv more efficiently as dgetrf followed by dgetri instead of calling dgesv...

[thrasibule]

This is what I get on my machine:

In [43]: %timeit npinv(F2[0,0,:,:])
10.7 µs ± 309 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [44]: %timeit hdinv(F[:,:,0,0])
20.1 µs ± 316 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [45]: %timeit npinv(F2)
13.4 ms ± 495 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [46]: %timeit hdinv(F)
152 µs ± 1.38 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

so for a single 3x3 matrix npinv is actually twice as fast as hdinv. But for the block matrix, npinv scale linearly (13ms ~12000 x 10.7micros) whereas hdinv is only 7.5 times slower for 12000 matrices... There is something magical going on.

[bhaveshshrimali]

so for a single 3x3 matrix npinv is actually twice as fast as hdinv.

Interesting, because on my machine for a simple 3x3 setup I get

F = random.random((3,3))

%timeit npinv(F)
44.3 µs ± 9.72 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit hdinv(F)
31.1 µs ± 4 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit finv(F)
36.6 µs ± 782 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Roughly the same order as I posted above but not as fast as yours.

[thrasibule]

You can actually try dgetri by using inv from scipy which uses it:

In [26]: F = np.random.rand(3,3)

In [27]: from scipy.linalg import inv as spinv

In [28]: %timeit spinv(F, check_finite=False)
6.78 µs ± 61.6 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

which is like 25% faster on my machine but there is no miracle. Even by rewriting the implicit loop in npinv in cython, which would get rid of the python overhead, it will get maybe twice as fast but will never get as fast as hdinv.

Thinking about it more I realized what you're doing in hdinv is really quite clever, because all the operations are vectorized on vectors of size 12000. The algorithm is quite different from what npinv is doing (basically you're reordering the loops), which is why it runs so much faster.

[bhaveshshrimali]

Indeed spinv runs about twice as fast compared to npinv on a single 3x3 input.

As an aside, Numba does use the implementations in dgetrf and dgetri when overloading numpy.linalg.inv but it still seems to operate only on 2D arrays so writing the remaining loops explicitly does make it faster but still not as fast as hdinv I think.

from numba import jit, prange

@jit(cache=True, nopython=True, nogil=True)
def vecinv(A):
    invA = zeros_like(A)
    for i in range(A.shape[0]):
        for j in range(A.shape[1]):
            invA[i,j,:, :] = npinv(A[i,j,:,:])
    return invA

F = random.random((3,3,1000,4))
F2 = einsum("ij...->...ij", F)

%timeit vecinv(F2)
4.61 ms ± 237 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit hdinv(F)
248 µs ± 67.5 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit npinv(F2)
106 ms ± 1.27 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

[thrasibule]

Very interesting, using numba to rewrite the loop is a good idea. If I replace npinv by a hand written version (basically you hdinv but for 3x3 matrices), I get the following timings:

@jit('float64[:,:](float64[:,:])', cache=True, nopython=True, nogil=True)
def fast_inverse(A):
    inv = np.empty_like(A)
    a = A[0, 0]
    b = A[0, 1]
    c = A[0, 2]
    d = A[1, 0]
    e = A[1, 1]
    f = A[1, 2]
    g = A[2, 0]
    h = A[2, 1]
    i = A[2, 2]

    inv[0, 0] = (e * i - f * h)
    inv[1, 0] = -(d* i - f * g)
    inv[2, 0] = (d * h - e * g)
    inv_det = 1 / (a * inv[0, 0] + b * inv[1, 0] + c * inv[2, 0])

    inv[0, 0] *= inv_det
    inv[0, 1] = - inv_det * (b * i - c * h)
    inv[0, 2] = inv_det * (b * f - c * e)
    inv[1, 0] *= inv_det
    inv[1, 1] = inv_det * (a * i - c * g)
    inv[1, 2] = -inv_det * (a * f - c * d)
    inv[2, 0] *= inv_det
    inv[2, 1] = -inv_det * (a * h - b * g)
    inv[2, 2] = inv_det * (a * e - b * d)
    return inv

%timeit vecinv(F2)
351 µs ± 13.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

so on my machine just twice as slow as hdinv. Basically I think the conclusion is: np.linalg inv is pretty slow for 3x3 matrix, and your hdinv is optimal in a way because of the loop reordering.

[ilayn]

If you make your hdinv robust (singularity checks etc.) that would also get slower. You are basically avoiding the LAPACK calls and some overhead for possibly underflowed zero determinants (apart from calling dgetri slowdown).

Alternative would be flattening arrays to 2D and solving a banded system.