fix: avoid duplicate function calls with FiniteDiff

[gdalle]

Meant to fix SciML/OrdinaryDiffEq.jl#4162 by making one function call inside DI.prepare_pushforward_samepoint and then reusing it inside FiniteDiff.finite_difference_jvp

[codecov]

Codecov Report

❌ Patch coverage is 18.18182% with 36 lines in your changes missing coverage. Please review.
✅ Project coverage is 88.79%. Comparing base (fab82f3) to head (c5e8e89).

Files with missing lines	Patch %	Lines
...xt/DifferentiationInterfaceFiniteDiffExt/twoarg.jl	0.00%	28 Missing ⚠️
...xt/DifferentiationInterfaceFiniteDiffExt/onearg.jl	50.00%	8 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (fab82f3) and HEAD (c5e8e89). Click for more details.

HEAD has 45 uploads less than BASE

Flag	BASE (fab82f3)	HEAD (c5e8e89)
DIT	12	4
DI	55	18

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1054      +/-   ##
==========================================
- Coverage   97.66%   88.79%   -8.88%     
==========================================
  Files         143      140       -3     
  Lines        8269     8235      -34     
==========================================
- Hits         8076     7312     -764     
- Misses        193      923     +730     

Flag	Coverage Δ
DI	88.45% <18.18%> (-9.82%)	⬇️
DIT	89.69% <ø> (-6.35%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[ChrisRackauckas]

This isn't quite it. The problem is that we have already computed f(x), so we want to reuse it in computations. I would expect the function to take in fx = f(x) under the contract that it is the correct x.

[gdalle]

That's exactly what this function does: it precomputes f(x) and stores it inside the prep while promising that prep will only ever be used at the same point x afterwards. What am I missing?

[ChrisRackauckas]

We have already computed f(x) outside of DI as the end point of an interpolation. We now want to give DI f(x) and then have it only compute f(x+eps) for the finitediff jvp

[gdalle]

I don't think that is easily done with the current DI API. What this PR does is provide you with a preparation function that computes f(x) for you, and tells DI.pushforward to reuse it. So it does waste one f call overall if you already have f(x) before calling DI.prepare_pushforward_same_point, but it then saves one f call per execution of DI.pushforward with the resulting prep object.
Can you point me to where in OrdinaryDiffEq we'd have to insert this?

[ChrisRackauckas]

For a Rosenbrock method for example, it would have already calculated this f(x) here: https://github.com/SciML/OrdinaryDiffEq.jl/blob/master/lib/OrdinaryDiffEqRosenbrock/src/rosenbrock_perform_step.jl#L842. It would go into the reinit of A https://github.com/SciML/OrdinaryDiffEq.jl/blob/4bf7627b75e881a14613f1b708f106e4243a882c/lib/OrdinaryDiffEqDifferentiation/src/linsolve_utils.jl#L59-L63 since A is a matrix-free operator for the J*v operation, used in a Jacobian-Free Newton Krylov (GMRES) method to solve the Newton Jx=b without building J, so it's just represented as an operator https://github.com/SciML/OrdinaryDiffEq.jl/blob/4bf7627b75e881a14613f1b708f106e4243a882c/lib/OrdinaryDiffEqDifferentiation/src/operators.jl#L28-L73 that calls the prebuilt cache. What we previously had was a way to declare that f(x) was already calculated via a boolean state, and then fx would just need to be kept updated outside the differentiation system under the assumption the bool was set correctly.

I noticed this was missing because the pre-DI code had this specialized with FiniteDiff.jl, that's why its cache build has this fx as something able to be passed, but then with the new solver method SciML/OrdinaryDiffEq.jl#4017 (comment) the benchmarks showed that our implementation was taking 1 extra evaluation per-step, leading to like a 20-25% performance reduction, and I tracked it down to being due to the change to DI which removed this handling.

[ChrisRackauckas]

The Jacobian code uses it https://github.com/SciML/OrdinaryDiffEq.jl/blob/4bf7627b75e881a14613f1b708f106e4243a882c/lib/OrdinaryDiffEqDifferentiation/src/derivative_wrappers.jl#L376-L381. It passes in uf: does DI prevent the f(u)=fu call with that? I assume it probably has 1 extra f-eval as well from this, though that's probably harder to ever really matter since at 10x10 that's a 9% regression 🤷 ehh).

But I think the jvp-operator code did not get as much love during the update to DI. It's my mistake as we didn't have tests that were strictly guaranteeing the right number of f calls and non-allocations. It's not really the biggest deal because generally you only use Jacobian-Free Newton Krylov for very large systems, and so if the system is large then these allocations are "relatively" small, but after making a few other things 2x-3x faster this was now showing up as the last 20-25%.

So yes, in a very concrete DI framing, what I am looking for is a way to give uf = f(u) in prepare!_pushforward so that I can change this operator code to reinitialize at new u in a non-allocating way, but also pass along the fact that I already have calculated f(u) in another part of the code.

[gdalle]

Ok so conceptually your workflow would look something like this?

prep0 = prepare_pushforward(f, backend, x0), (dx0,))
for i in loop_over_x  # ODE solve
    # your code computes yi = f(xi) somehow
    prep1 = prepare!_pushforward_same_point(f, prep0, backend, xi, (dxi,); y=yi)
    for j in loop_over_dx  # iterative linear solve based on the JVPs at xi
        pushforward(f, prep0, backend, xi, (dxi_j,))
    end
end

In which case I'd have to add:

• the function prepare!_pushforward_same_point
• the kwarg y in there

[ChrisRackauckas]

Precisely

[gdalle]

That will cost you one karaoke duet per feature

[ChrisRackauckas]

What's your duet style? Are you like a Chester to my Mike, or is this like a "me and my shadow" kind of situation? I've also got a pretty high voice, so we could also go off the deep end with summer nights or empire state of mind.