Why does JET give these "runtime dispatch detected" lines for some LDLFactorizations calls in Tulip

I’m playing around with linear programming using Tulip.jl and BigFloat, and I wanted to decrease the time necessary for finding each LP solution, so it seemed like JET.jl might be helpful to investigate possible performance issues. No issues are reported for my code, but there seem to be some in Tulip.
Two lines in particular that are flagged by JET.jl for runtime dispatch confuse me.
These two functions are at the end of Tulip’s src/KKT/LDLFactorizations/ldlfact.jl file:

function update!(kkt::LDLFactSolver{T,K2}, θ, regP, regD) where{T}
    m, n = kkt.m, kkt.n

    # Sanity checks
    length(θ)  == n || throw(DimensionMismatch(
        "length(θ)=$(length(θ)) but KKT solver has n=$n."
    ))
    length(regP) == n || throw(DimensionMismatch(
        "length(regP)=$(length(regP)) but KKT solver has n=$n"
    ))
    length(regD) == m || throw(DimensionMismatch(
        "length(regD)=$(length(regD)) but KKT solver has m=$m"
    ))

    copyto!(kkt.θ, θ)
    copyto!(kkt.regP, regP)
    copyto!(kkt.regD, regD)

    # Update KKT matrix
    # K is stored as upper-triangular, and only its diagonal is changed
    @inbounds for j in 1:kkt.n
        k = kkt.K.colptr[1+j] - 1
        kkt.K.nzval[k] = -kkt.θ[j] - regP[j]
    end
    @inbounds for i in 1:kkt.m
        k = kkt.K.colptr[1+kkt.n+i] - 1
        kkt.K.nzval[k] = regD[i]
    end

    # Update factorization
    try
        LDLF.ldl_factorize!(Symmetric(kkt.K), kkt.F)
    catch err
        isa(err, LDLF.SQDException) && throw(PosDefException(-1))
        rethrow(err)
    end

    return nothing
end

function solve!(dx, dy, kkt::LDLFactSolver{T,K2}, ξp, ξd) where{T}
    m, n = kkt.m, kkt.n

    # Setup right-hand side
    @views copyto!(kkt.ξ[1:n], ξd)
    @views copyto!(kkt.ξ[(n+1):end], ξp)

    # Solve augmented system
    # CHOLMOD doesn't have in-place solve, so this line will allocate
    LDLF.ldiv!(kkt.F, kkt.ξ)

    # Recover dx, dy
    @views copyto!(dx, kkt.ξ[1:n])
    @views copyto!(dy, kkt.ξ[(n+1):end])

    # TODO: iterative refinement
    return nothing
end

JET detects runtime dispatch for the LDLF.ldiv! call and for the LDLF.ldl_factorize! call:

││┌ @ /home/nsajko/.julia/environments/v1.8/dev/Tulip/src/KKT/LDLFactorizations/ldlfact.jl:112 %228(%242, %245)
│││ runtime dispatch detected: %228::typeof(LDLFactorizations.ldl_factorize!)(%242::LinearAlgebra.Symmetric{BigFloat, SparseArrays.SparseMatrixCSC{BigFloat, Int64}}, %245::LDLFactorizations.LDLFactorization{BigFloat})
││└────────────────────────────────────────────────────────────────────────────────────────────
││┌ @ /home/nsajko/.julia/environments/v1.8/dev/Tulip/src/KKT/LDLFactorizations/ldlfact.jl:126 %143(%144, %145)
│││ runtime dispatch detected: %143::typeof(LinearAlgebra.ldiv!)(%144::LDLFactorizations.LDLFactorization{BigFloat}, %145::Vector{BigFloat})
││└────────────────────────────────────────────────────────────────────────────────────────────

Two things confuse me here:

  1. Why does JET say %143::typeof(LinearAlgebra.ldiv!) instead of just LinearAlgebra.ldiv!? Likewise with ldl_factorize!.
  2. Why the runtime dispatch at all, the types indicated by JET for the call are neither abstract (I think) nor unions?

How to reproduce this (if someone is interested)

The function I was interested in optimizing is ipm_optimize! in IPM/HSD/HSD.jl, with this type signature: ipm_optimize!(hsd::HSD{T}, params::IPMOptions{T}) where{T}. To run JET successfully I needed the exact types of the arguments of ipm_optimize! when called by my code, so I put something like this at the beginning of the function:

println(typeof(hsd))
println(typeof(params))

Then I used this information for defining these two type alias constants:

const HSD = Tulip.HSD{BigFloat, Vector{BigFloat}, BitVector, SparseArrays.SparseMatrixCSC{BigFloat, Int64}, Tulip.KKT.TlpLDLFactorizations.LDLFactSolver{BigFloat, Tulip.KKT.K2}}
const PAR = Tulip.IPMOptions{BigFloat}

Then it’s possible to call JET.report_opt(Tulip.ipm_optimize!, Tuple{HSD, PAR}) (I think that’s the correct way to do it).

I should also mention that I modified some parts of Tulip to remove @timeit macros and debugging branches of the code. Also, in a few places I replaced try x catch y end with just x. This will both conceivably help with my performance issues and does away with a lot of JET noise.

In my limited experience dynamic dispatch is accepted by package authors if the runtime cost is low in comparison to the computation executed. To check this I probably would use a profiler. If this would indicate a performance problem due to dynamic dispatch then filing an issue against the responsible package would make sense.

1 Like

This is interesting :thinking:
I’ve always had the (mis?) conception that using BigFloat in Tulip was going to be terribly slow, because that’s the way things are with BigFloat. If there were any compiler-related inefficiencies, I would have expected them to be negligible anyway.

As for JET: I have very little (virtually zero) experience with it, so I can’t really comment on why that line is flagged. I ran the code you shared (Julia1.7, Tulip@master), and JET also reported hundreds of issues/runtime dispatch that were mostly related to broadcast and the use of @printf

For instance, I saw a large number of such messages from JET:

││┌ @ broadcast.jl:516 Base.Broadcast.length(%157)
│││ runtime dispatch detected: Base.Broadcast.length(%157::Base.OneTo{Int64})

@nsajko: I you want to open an issue in Tulip.jl to track this, I will happily help the best I can.

1 Like

Some frontend issue. This printing is “correct” since LinearAlgebra.ldiv!::typeof(LinearAlgebra.ldiv!) holds true. But I agree with this is better printed as LinearAlgebra.ldiv!(...) instead. If anyone interested in working on this issue, here is code we want to improve: JET.jl/inferenceerrorreport.jl at 071d3108bb20fd15571f7294aad28c99e20e798a · aviatesk/JET.jl · GitHub

I’m not familiar with LDLFactorizations codebase, but it seems like LDLFactorizations.LDLFactorization{BigFloat} is not really concrete (there seems to be missing type parameters)?

help?> LDLFactorizations.LDLFactorization
  No documentation found.

  Summary
  ≡≡≡≡≡≡≡≡≡

  mutable struct LDLFactorizations.LDLFactorization{T<:Real, Ti<:Integer, Tn<:Integer, Tp<:Integer}

This seems to be a bug of JET (as the call seems to be fully concrete and should be resolved at compile time). Could you please provide me a reproducer?

3 Likes

Here you go:

As far as I can tell, the LDLFactorization struct has all its fields typed / parametrized. If there’s a missing type it’s probably in the Tulip code. I’ll look into this.

2 Likes