Trim allocations in dual linear problem

[hersle]

I am differentiating ODEs with ForwardDiff in my package.

using SymBoltz, ProfileCanvas
...
@time ForwardDiff.jacobian(logP, logθ) # ForwardDiff-differentiates ~20 stiff ODEs with ~50x50 Jacobians
@profview_allocs ForwardDiff.jacobian(logP, logθ)

  1.537143 seconds (1.17 M allocations: 906.513 MiB, 18.37% gc time)

There are a lot of allocations from nodual_value(val) and partial_vals(val) in LinearSolve's ForwardDiff extension. They set some cache field with a new reallocated array every time. This seems unnecessary to me.

In this PR I attempt to instead use map! and a custom setproperty! dispatch to set the elements nodual_value(val[i]) and partial_vals(val[i]) without having to reallocate the arrays. This gives me a much shorter runtime and lower memory usage:

  0.823024 seconds (218.36 k allocations: 18.415 MiB)

Now the LinearSolve allocations in the first flamegraph are gone.

Checklist

• Appropriate tests were added
• Any code changes were done in a way that does not break public API
• All documentation related to code changes were updated
• The new code follows the
  contributor guidelines, in particular the SciML Style Guide and
  COLPRAC.
• Any new documentation only uses public API

Additional context

Add any other context about the problem here.

[hersle]

Sorry for not providing a full self-isolated test/example here, I ] dev LinearSolve while working on my package.

Let me know if I should make one. The issue here was somewhat similar to SciML/OrdinaryDiffEq.jl#2837, but with a bigger ODE (~50x50 Jacobian here, instead of 3x3 in that issue).

[ChrisRackauckas]

are these needed?

[hersle]

Not required, but they gave me a non-negligible speedup, as I saw checkbounds showing up in profiling.

[hersle]

I added it so the user should not suffer a performance penalty from library internals. I can remove them if you want to keep it safe.

This function is just shuffling data around. The optimal solution would be to avoid it altogether, but I am not sure if it's easily possible.

[ChrisRackauckas]

Is this used? Also, it's a pun so probably should be changed.

[hersle]

Yes, it's used here:

LinearSolve.jl/ext/LinearSolveForwardDiffExt.jl

Line 315 in 3c47844

setproperty!(dc.linear_cache, sym, val, nodual_value) # maps nodual_value.(val) without allocating

Hmm, what do you mean by it's a pun? 😅

[hersle]

It's a "trick" to call setproperty! while doing an in-place of the array.

[ChrisRackauckas]

But why use the name of setproperty? It's not actually a setproperty! call

[hersle]

Yeah that's true. After understanding everything better, I came up with a simpler version that just uses standard setproperty! calls and new in-place nodual_value! and partial_vals! functions.

[hersle]

I fixed the tests in test/forwarddiff_overloads.jl locally, so CI is worth another try.

The failures seemed related to nested duals, so I added a fallback that reintroduces the old-ish behavior in that case.

[hersle]

I have to move this before the solve! to make this test pass. Is that fine or am I overlooking something?

[ChrisRackauckas]

No that's not good.

[hersle]

Sorry this was never a problem, my local state messed up the testing. Removing this.

[hersle]

@ChrisRackauckas Refactored to use just standard setproperty! calls and new in-place nodual_value! and partial_vals! functions. I think this is cleaner. Local test passes:

julia> include("test/forwarddiff_overloads.jl")
Test Passed

[hersle]

Hold on, there is a small type instability in

Seems related to the (somewhat complex) branching/recursion on sym in setproperty!.

[hersle]

I reworked the branching in setproperty! a bit to make the if-else path simpler. Tests still pass and the type instability is gone:

Hopefully that's it.