Type Inference Issues in IterativeSolvers with LinearMaps

I was trying to improve the performance of a program I use that makes heavy use of implicit time stepping, which results in me solving linear systems of equations ‘Ax=b’. Because of things that have to do with my specific application, it would be very convenient for me to be able to only need to be able to compute the action of A on x, and use an iterative solver to find the solution. When profiling my program using ProfileView.jl, there were red bars indicating type inference issues in the call to gmres! from IterativeSolvers.jl.

Below I have constructed a minimal working example. This example simply solves the following linear system of equations (5,000 times, so that the sampler actually captures the program well), but with the left hand matrix implement as a LinearMap

[2.0  0.0][x1]   [1.0]
[0.0  2.0][x2] = [2.0]

Below is the implementation.

import LinearMaps
import IterativeSolvers

function left_matrix_operation(result::AbstractVector{Float64}, 
        v::AbstractVector{Float64})::Vector{Float64}
    return result .= 2 .* v
end

function main()
    v_left::Vector{Float64} = zeros(2)
    v_right::Vector{Float64} = [1.0, 2.0]
    

    left_matrix_map = LinearMaps.LinearMap{Float64}(
        left_matrix_operation, 2, 2, ismutating=true
    )

    for i in 1:5000
        IterativeSolvers.gmres!(v_left, left_matrix_map, v_right)
    end
end

Here is the profile flame graph using ProfileView.jl (@profview main())

I see that there is a type inference issue (or some other run time issue) in part of the call to gmres! from the IterativeSolvers.jl package. Using Cthulu.jl to descend on the call, I get the following:

IterativeSolvers.ArnoldiDecomp(A::matT, order::Int64, T::Ty
pe) where matT @ IterativeSolvers ~/.julia/packages/IterativeSolvers/rhYBz/src/gmres.jl:11
11 (ArnoldiDecomp(A::matT, order::Int, T::Type::Type) where {matT})::IterativeSolvers.ArnoldiDecomp{_A, LinearMaps.FunctionMap{Float64, typeof(left_matrix_operation), Nothing, true}} where _A = ArnoldiDecomp{T::Type, matT}::Type{IterativeSolvers.ArnoldiDecomp{_A, LinearMaps.FunctionMap{Float64, typeof(left_matrix_operation), Nothing, true}}} where _A(
12     A,
13     zeros(T::Type, size(A, 1), order + 1)::Matrix,
14     zeros(T::Type, order + 1, order)::Matrix
15 )
Select a call to descend into or ↩ to ascend. [q]uit. [b]ookmark.
Toggles: [w]arn, [h]ide type-stable statements, [t]ype annotations, [s]yntax highlight for Source/LLVM/Native, [j]ump to source always.
Show: [S]ource code, [A]ST, [T]yped code, [L]LVM IR, [N]ative code
Actions: [E]dit source code, [R]evise and redisplay
 • size(A, 1)
   %3 = +(::Int64,::Int64)::Int64
   zeros(T::Type, size(A, 1), order + 1)
   %5 = +(::Int64,::Int64)::Int64
   zeros(T::Type, order + 1, order)
   ArnoldiDecomp{T::Type, matT}::Type{IterativeSolvers.ArnoldiDecomp{_A, LinearMaps.FunctionMap{Float64, typeof(left_matrix_operation), Nothing, true}}} where _A( A, zeros(T::Type, size(A, 1), o… 
   ↩

And here’s a picture, so you can see the red highlights:

I haven’t taken the time to fully understand the gmres code, but it looks like in ArnoldiDecomp an array of zeros is allocated, which I figure is almost certainly supposed to be an array of floating point numbers, but for some reason the compiler can’t figure that out?

Does anyone have advice on how to make the type known, and does anyone have any idea whether it would actually impact performance significantly? I am fairly new to using these profiling tools.

Pretty sure this is
https://docs.julialang.org/en/v1/manual/performance-tips/#Be-aware-of-when-Julia-avoids-specializing

1 Like

That looks like it probably is it. Thanks!