Feature Request: Add function for in-place eigendecomposition

[longisland-icetea]

Hi maintainers,

I'd like to propose adding an in-place eigendecomposition function, eigen_inplace! (as new function or overload of eigen!), to LinearAlgebra.jl. This would allow users to compute eigenvalues (and optionally eigenvectors) without allocating memory for eigenvalues, which can be especially useful for large amount of small matrices.

Below is a minimal working example that computes the eigenvalues and overwrites a Hermitian matrix A with its eigenvectors using the MKL routine on Hermitian matrix (imitating syev! in

LinearAlgebra.jl/src/lapack.jl

Line 5543 in 28ee87e

Base.@constprop :none function _syev!(jobz::AbstractChar, uplo::AbstractChar, A::AbstractMatrix{$elty})

):

function eigen_inplace!(A::AbstractMatrix{ComplexF64}, W::AbstractArray{Float64,1})
    n = size(A, 1)
    work = Vector{ComplexF64}(undef, 1)
    lwork = Int32(-1)
    rwork = Vector{Float64}(undef, max(1, 3n - 2))
    info = Ref{Int32}()
    for i = 1:2  # first call returns lwork as work[1]
        ccall((:zheev_, libmkl), Cvoid,
            (Ref{UInt8}, Ref{UInt8}, Ref{Int32}, Ptr{ComplexF64}, Ref{Int32},
                Ptr{Float64}, Ptr{ComplexF64}, Ref{Int32}, Ptr{Float64}, Ptr{Int32},
                Clong, Clong),
            'V', 'U', n, A, stride(A, 2), W, work, lwork, rwork, info, 1, 1)
        if i == 1
            lwork = Int32(real(work[1]))
            resize!(work, lwork)
        end
    end
    (W, A)
end

Ideally, a more general and robust version would:

Support both real and complex Hermitian/symmetric matrices,
Allow choosing whether to compute eigenvectors,
Be compatible with different BLAS/LAPACK providers (not just MKL, I manually config libmkl because I don't know how libblastrampoline work),

I can't do more as I'm not familiar with julia programing.

Thanks for your time!

[araujoms]

I don't think that would be worth it, because the eigenvectors don't need much memory, work, rwork, and iwork use much more. And to pre-allocate those you need the low-level interface anyway.

[longisland-icetea]

I don't think that would be worth it, because the eigenvectors don't need much memory, work, rwork, and iwork use much more. And to pre-allocate those you need the low-level interface anyway.

rwork should be larger but I have test in my project and got about 40% improvement, here is the minimal test code, which didn't touch the low-level interface :

using LinearAlgebra, MKL, Random
function get_blas_lib_names()
    config = BLAS.get_config()
    libs = getfield(config, :loaded_libs)
    if isempty(libs)
        return String[]
    else
        return collect(libs)[1].libname
    end
end
const libmkl = get_blas_lib_names()
BLAS.set_num_threads(1)

const n = 8
const num_matrices = 1000000

function eigen_inplace!(A::AbstractMatrix{ComplexF64}, W::AbstractArray{Float64,1})
    n = size(A, 1)
    work = Vector{ComplexF64}(undef, 1)
    lwork = Int32(-1)
    rwork = Vector{Float64}(undef, max(1, 3n - 2))
    info = Ref{Int32}()
    for i = 1:2  # first call returns lwork as work[1]
        ccall((:zheev_, libmkl), Cvoid,
            (Ref{UInt8}, Ref{UInt8}, Ref{Int32}, Ptr{ComplexF64}, Ref{Int32},
                Ptr{Float64}, Ptr{ComplexF64}, Ref{Int32}, Ptr{Float64}, Ptr{Int32},
                Clong, Clong),
            'V', 'U', n, A, stride(A, 2), W, work, lwork, rwork, info, 1, 1)
        if i == 1
            lwork = Int32(real(work[1]))
            resize!(work, lwork)
        end
    end
    (W, A)
end

function test_eigen()
    matrices = zeros(ComplexF64, n, n, num_matrices)
    for i in 1:num_matrices
        @views matrices[:, :, i] = Hermitian(randn(ComplexF64, n, n))
    end
    @time for i in 1:num_matrices
        @views eigen!(matrices[:, :, i])
    end
end

function test_eigen_inplace()
    matrices = zeros(ComplexF64, n, n, num_matrices)
    for i in 1:num_matrices
        @views matrices[:, :, i] = Hermitian(randn(ComplexF64, n, n))
    end
    eigenvalues = zeros(Float64, n, num_matrices)
    @time for i in 1:num_matrices
        @views eigen_inplace!(matrices[:, :, i], eigenvalues[:, i])
    end
end

function test_eigen_performance()
    println("test eigen_inplace:")
    @time test_eigen_inplace()
    println("test eigen:")
    @time test_eigen()
end

test_eigen_performance()

in my machine (12700KF 32GB RAM) the result is:

test eigen_inplace:
5.454230 seconds (4.00 M allocations: 534.058 MiB, 0.26% gc time)
6.248530 seconds (6.01 M allocations: 2.593 GiB, 1.34% gc time, 0.13% compilation time)
test eigen:
8.218828 seconds (19.00 M allocations: 5.141 GiB, 1.59% gc time)
9.068301 seconds (21.00 M allocations: 7.153 GiB, 2.97% gc time)

for smaller matrix the improvement will be larger.

[araujoms]

eigen! uses by default syevr, you are using syev in your code.

[longisland-icetea]

eigen! uses by default syevr, you are using syev in your code.

I see. Thanks.

[kshyatt]

FWIW with LinearAlgebra moving out of the sys image, you could try out https://github.com/QuantumKitHub/MatrixAlgebraKit.jl and its "YALAPACK" which allows you to supply input matrices in cases like these