Performance asymptotics for sparsity detection

[j-fu]

Hi, I am testing sparsity detection via DifferentiationInterface against Symbolics.jl+SparseDiffTools.jl, and I see a regression of asymptotic performance for sparse jacobian calculation away from the O(n) behavior of Symbolics.jl+SparseDiffTools.jl. May be I miss some options to tune this in the right way.

The test involves essentially a tridiagonal operator, something like a nonlinear scalar 1D PDE discretized with finite differences. Below I provide links to the pluto notebook behind the image (with a smaller version of the problem by default) and rendered versions.

Links:

• Pluto/Julia
• html
• pdf

[gdalle]

Hi @j-fu, thanks for trying out DI, and for running cool experiments!

In general, why do you benchmark tdetect + tassemble? The whole idea of Jacobian "preparation" is that you only perform it once, outside of the hot loop, and then you reuse the results as long as your sparsity pattern doesn't change. Any reason why that wouldn't work here? Note that the remark also applies for the SparseDiffTools pipeline.
That doesn't mean detection and coloring time isn't relevant, but it might not be your true bottleneck if you use DI as intended.

For DI+Symbolics, I'm not sure why we see the deviation at the end. One possible reason might be that you include a new, more sophisticated coloring type, with the IncidenceDegree() ordering which SparseDiffTools definitely doesn't bother with. Does it give you fewer colors in your case? Can you try it again with NaturalOrder(), to see if the small deviation disappears?

For DI+SCT, I think it has to do with the way we store sets inside sparsity pattern detection. It's always a tradeoff that depends on the size and density of the problem. SCT comes with a way to switch set representations, but it isn't public yet so the interface is a little nasty. Could you try a comparison between

sparsity_detector = SCT.TracerSparsityDetector()  # the default

and the alternative

sparsity_detector = SCT.TracerSparsityDetector(;
    gradient_tracer_type=SCT.GradientTracer{
        SCT.IndexSetGradientPattern{Int,Set{Int}}
    }
)

to see if it speeds things up? Essentially, the default index set storage is a BitSet, which works very well in small dimension but becomes burdensome in high dimension. Replacing it with a Set{Int} could give you the speedup you need.

Ping @adrhill, we may want to think about making that interface more palatable, especially since we use it in our paper.

[j-fu]

In general, why do you benchmark tdetect + tassemble? The whole idea of Jacobian "preparation" is that you only perform it once, outside of the hot loop, and then you reuse the results as long as your sparsity pattern doesn't change. Any reason why that wouldn't work here? Note that the remark also applies for the SparseDiffTools pipeline.
That doesn't mean detection and coloring time isn't relevant, but it might not be your true bottleneck if you use DI as intended.

Sure, in fact, this is the way I am using this - detect sparsity once and and use this for further assembly. The use case in VoronoiFVM is a "generic operator" which occasionally is used if a problem does not fit into the standard "nodal sources/sinks/reactions + edgefluxes" graph based paradigm, where I assemble things from the pieces by using a linked-list based intermediate sparse matrix structure for new entries.

And those models I am aware of, where the generic operator is used, in the moment have rather small sizes.
Nevertheless I think, if O(n) is possible here, one should strive for it. (The problem here is 1D, it would be interesting to test higher dimensions as well, I put this on my list...)

[j-fu]

A quick test gives no significantly different picture with NaturalOrder and the GradientTracer enabled:

[gdalle]

Very strangely, my own benchmarking (code below) finds that the difference is not due to sparsity detection, but instead to the wider Jacobian preparation.
What's weird about that is: outside of sparsity detection, the rest of the preparation pipeline is exactly the same regardless of whether SCT or Symbolics was used as the detector!

using ADTypes
using CairoMakie
using Chairmarks
import DifferentiationInterface as DI
using ForwardDiff: ForwardDiff
using LinearAlgebra
import SparseConnectivityTracer as SCT
import SparseMatrixColorings as SMC
using Symbolics: Symbolics
using Test

const detector_sct_default = SCT.TracerSparsityDetector()
const detector_sct_set = SCT.TracerSparsityDetector(;
    gradient_tracer_type=SCT.GradientTracer{SCT.IndexSetGradientPattern{Int,Set{Int}}}
)
const detector_symb = Symbolics.SymbolicsSparsityDetector()

function MyAutoSparse(sparsity_detector)
    return AutoSparse(
        AutoForwardDiff();
        sparsity_detector,
        coloring_algorithm=SMC.GreedyColoringAlgorithm(),
    )
end

function f!(y, u)
    n = length(u)
    h = 1 / (n - 1)
    y[1] = (u[1]^2 - 1) * 0.5 * h
    for i in 2:(n - 1)
        y[i] = (u[i]^2 - 1.0) * h
    end
    y[end] = (u[end]^2 - 1) * 0.5 * h

    for i in 1:(n - 1)
        du = (u[i + 1]^2 - u[i]^2) / h
        y[i + 1] += du
        y[i] -= du
    end
    return nothing
end

@testset "Correctness" begin
    u_test = ones(10)
    y_test = ones(10)
    J = Tridiagonal(ones(Bool, 9), ones(Bool, 10), ones(Bool, 9))
    J1 = jacobian_sparsity(f!, y_test, u_test, detector_sct_default)
    J2 = jacobian_sparsity(f!, y_test, u_test, detector_sct_set)
    J3 = jacobian_sparsity(f!, y_test, u_test, detector_symb)
    @test J1 == J
    @test J2 == J
    @test J3 == J
    prep = DI.prepare_jacobian(f!, y_test, MyAutoSparse(detector_sct_default), u_test)
    @test SMC.ncolors(prep) == 3
    @test SMC.sparsity_pattern(prep) == J
end

function benchmark_detection(n_values, detector)
    benchs = map(n_values) do n
        yield()
        @info "Benchmarking detection" n detector
        u, y = ones(n), ones(n)
        @b jacobian_sparsity(f!, y, u, detector)
    end
    return map(b -> b.time, benchs)
end

function benchmark_preparation(n_values, detector)
    backend = MyAutoSparse(detector)
    benchs = map(n_values) do n
        yield()
        @info "Benchmarking preparation" n detector
        u, y = ones(n), ones(n)
        @b DI.prepare_jacobian(f!, y, backend, u)
    end
    return map(b -> b.time, benchs)
end

function plot_benchmarks(title, n_values; times_sct_default, times_sct_set, times_symb)
    fig = Figure()
    ax = Axis(
        fig[1, 1];
        title,
        xlabel="Input dimension",
        ylabel="Time [s]",
        xscale=log10,
        yscale=log10,
    )
    # linear regression in log scale
    X = [log10.(n_values) ones(length(n_values))]
    slope_sct_default, _ = round.(X \ log10.(times_sct_default); digits=3)
    slope_sct_set, _ = round.(X \ log10.(times_sct_set); digits=3)
    slope_symb, _ = round.(X \ log10.(times_symb); digits=3)
    scatterlines!(
        ax, n_values, times_sct_default; label="SCT (default)\nslope $slope_sct_default"
    )
    scatterlines!(ax, n_values, times_sct_set; label="SCT (set)\nslope $slope_sct_set")
    scatterlines!(ax, n_values, times_symb; label="Symbolics\nslope $slope_symb")
    axislegend(ax; position=:rb)
    return fig
end

n_values = ceil.(Int, 10 .^ (1:0.3:6))

## Detection

detection_times_sct_default = benchmark_detection(n_values, detector_sct_default)
detection_times_sct_set = benchmark_detection(n_values, detector_sct_set)
detection_times_symb = benchmark_detection(n_values, detector_symb)

fig_det = plot_benchmarks(
    "Sparsity detection alone",
    n_values;
    times_sct_default=detection_times_sct_default,
    times_sct_set=detection_times_sct_set,
    times_symb=detection_times_symb,
)
save("detection.png", fig_det)

## Preparation

preparation_times_sct_default = benchmark_preparation(n_values, detector_sct_default)
preparation_times_sct_set = benchmark_preparation(n_values, detector_sct_set)
preparation_times_symb = benchmark_preparation(n_values, detector_symb)

fig_prep = plot_benchmarks(
    "Full Jacobian preparation",
    n_values;
    times_sct_default=preparation_times_sct_default,
    times_sct_set=preparation_times_sct_set,
    times_symb=preparation_times_symb,
)
save("preparation.png", fig_prep)

[gdalle]

With the one-line fix in #777, I get this much better curve. @j-fu thanks for the report! Can you try the DI version from that branch on your whole benchmark suite?

[j-fu]

Wow what a difference!

DI delivers O(n) now, and DI+SparseConnectivityTracer is slightly better than VoronoiFVM!

Interesting what happens for 2D/3D and systems. Need to do other homework though, but put this on my list.

[gdalle]

Could you add the benchmarks without the preparation step? I'm curious to see if we can beat your handcrafted code when we amortize preparation, as we would in realistic settings

[gdalle]

And there's another fix landing in adrhill/SparseConnectivityTracer.jl#239

[adrhill]

SCT v0.6.17 contains the fix!

[gdalle]

With the default settings

SCTBackend = AutoSparse(
    AutoForwardDiff();  # any object from ADTypes
    sparsity_detector = TracerSparsityDetector(),
    coloring_algorithm = GreedyColoringAlgorithm()
)

and the upcoming patches in DI and SCT, I observe that DI+SCT competes with the best manual version, and that's including sparsity detection

Now if we exclude preparation from the timings for DI and SparseDiffTools, the manual version is completely overpowered (there might be optimizations available in the manual version too, so the plot below isn't entirely fair)