Performance issue with mixed mode Jacobian

[ErikQQY]

In the performance comparison between forward, reverse, and mixed mode Jacobian, when we are using Enzyme as the backend, the result is quite satisfying:

3×2 DataFrame
 Row │ mode                     time       
     │ Abstract…                Float64    
─────┼─────────────────────────────────────
   1 │ ForwardMode()            0.00721971
   2 │ ReverseMode()            0.0138332
   3 │ ForwardAndReverseMode()  9.5529e-5

however, if we use ForwardDiff.jl and ReverseDiff.jl, the difference is not that significant compared with only using ForwardDiff.jl:

using DifferentiationInterface, DifferentiationInterfaceTest
using SparseConnectivityTracer, SparseMatrixColorings
using SparseArrays, StableRNGs
using ForwardDiff: ForwardDiff
using ReverseDiff: ReverseDiff

function MyAutoSparse(backend)
    sparsity_detector = TracerSparsityDetector()
    coloring_algorithm = GreedyColoringAlgorithm(RandomOrder(StableRNG(0)))
    return AutoSparse(backend; sparsity_detector, coloring_algorithm)
end

forward = MyAutoSparse(AutoForwardDiff())
reverse = MyAutoSparse(AutoReverseDiff())
mixed = MyAutoSparse(
    MixedMode(AutoForwardDiff(), AutoReverseDiff())
)

function f!(y, x)
    copyto!(y, x)  # diagonal of Jacobian
    y[1] += sum(x)  # first row of Jacobian
    return y .+= x[1]  # first column of Jacobian
end

n = 1000
x, y = ones(n), zeros(n);
J = spdiagm(0 => ones(n))
J[1, :] .= 1
J[:, 1] .= 1
J[1, 1]  = 3
scenario = Scenario{:jacobian,:in}(f!, y, x; res1=J)

data = benchmark_differentiation([forward, reverse, mixed], [scenario]; logging=true)
data[!, :mode] = DifferentiationInterface.mode.(data[!, :backend])
data = data[data[!, :operator] .== :jacobian!, [:mode, :time]]

3×2 DataFrame
 Row │ mode                     time       
     │ Abstract…                Float64    
─────┼─────────────────────────────────────
   1 │ ForwardMode()            0.00320217
   2 │ ReverseMode()            0.707737
   3 │ ForwardAndReverseMode()  0.00196804

So is the advantage of mixed mode Jacobian only available on Enzyme?

[gdalle]

First, the good news: Enzyme performing well seems to indicate that the sparse autodiff pipeline is not the bottleneck for your specific use case.
Now the bad news: using ReverseDiff for sparse Jacobians is unlikely to be efficient anytime soon. The reason is that it doesn't provide me with a tape-able pullback (#321), so the tape is re-recorded for every single pullback.

A few things about benchmarking though: at the moment, random ordering is more performant because of limitations on the SparseMatrixColorings side, which I hope to solve in the coming weeks. But just because you set the seed at the beginning doesn't mean the RNG will always be seeded the same way. See gdalle/SparseMatrixColorings.jl#160, which I'll try to fix today.

[ErikQQY]

OK, so the mixed mode is best suited for Enzyme.jl for now.

Another related question, IIUC, since the res in Scenario is used for the reference Jacobian prototype of the final result, so we didn’t tell the backend to exploit the known banded pattern of the banded part of f!, the question is that if the mixed mode is being able to auto detect which part use forward and which part use reverse, or the above example is not exploiting the known sparsity pattern?

[gdalle]

The sparsity pattern that gets used is the one detected by the sparsity_detector inside AutoSparse, not the one from the provided res1 (which is only used to pass jac = similar(res1) to the in-place jacobian! operator).
So the bidirectional coloring gets a generic sparse matrix, and then it decides the split between forward and reverse however it wants. If you have additional knowledge about the problem (like there is one block you want to do with forward and the rest with reverse) then you will need to essentially recode the sparse Jacobian manually.