Performance of computing the diagonal of V' * A *V

[stevengj]

@stevengj

Instead of multiplying A by the columns of V one by one, why not use BLAS for the matrix–matrix product Y = AV, which does exactly the same O(n^3) work but much faster? Then “manually” compute the dot products of the columns of V with the columns of Y — this is only O(n^2) , so it will asymptotically beat the BLAS matrix–matrix product.

In particular, this is about 2x faster than diagonal_elements2! as expected (because it saves one of the matrix–matrix multiplications) for your 2000 \times 2000 test case on my machine:

@views function diagonal_elements3!(diag_elem, A, V, Y)
    mul!(Y, A, V)
    for i in eachindex(diag_elem)
        diag_elem[i] = dot(V[:,i], Y[:,i])
    end
    return diag_elem
end

(No performance reason to declare the argument types, though for clarity you might declare them as diag_elem::AbstractVector, A::AbstractMatrix, V::AbstractMatrix, Y::AbstractMatrix.)