fix: make FiniteDiff prep non-allocating with static arrays

[dingraha]

I noticed some allocations when using FiniteDiff with StaticArrays that I eventually traced to calls to similar in the FiniteDiff extension. Adding the check for immutable input or output arrays seemed to fix this. Haven't attempted to add tests for this—wanted to get some feedback before diving in. What do you think?

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 98.10%. Comparing base (1912e3d) to head (f298e1b).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1019   +/-   ##
=======================================
  Coverage   98.09%   98.10%           
=======================================
  Files         142      143    +1     
  Lines        8239     8268   +29     
=======================================
+ Hits         8082     8111   +29     
  Misses        157      157           

Flag	Coverage Δ
DI	98.87% <100.00%> (+<0.01%)	⬆️
DIT	96.03% <ø> (ø)

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.

[gdalle]

Nice catch! The alternative would be replacing similar with copy, does that work on SArrays too?

[dingraha]

Yes, copy returns a SArray (unlike similar), but that doesn't eliminate all the allocations, I think because the FiniteDiff.JVPCache struct is mutable. Here is a short script that shows this:

using FiniteDiff
using DifferentiationInterface
using StaticArrays
function doit()
    backend = AutoFiniteDiff()
    x = @SVector [1.0, 2.0]
    tx = (2.0.*x,)
    f(x) = @. 3.0 * x
    # prep = DifferentiationInterface.prepare_pushforward(f, backend, x, tx);
    prep =  DifferentiationInterface.prepare_pushforward_nokwarg(Val(true), f, backend, x, tx)
    return prep
end

When DI.prepare_pushforward_nokwarg looks like this:

function DI.prepare_pushforward_nokwarg(
        strict::Val, f, backend::AutoFiniteDiff, x, tx::NTuple, contexts::Vararg{DI.Context, C}
    ) where {C}
    _sig = DI.signature(f, backend, x, tx, contexts...; strict)
    fc = DI.fix_tail(f, map(DI.unwrap, contexts)...)
    y = fc(x)
    cache = if x isa Number || y isa Number
        nothing
    else
        JVPCache(copy(x), y, fdtype(backend))
    end
    relstep = if isnothing(backend.relstep)
        default_relstep(fdtype(backend), eltype(x))
    else
        backend.relstep
    end
    absstep = if isnothing(backend.absstep)
        relstep
    else
        backend.absstep
    end
    dir = backend.dir
    return FiniteDiffOneArgPushforwardPrep(_sig, cache, relstep, absstep, dir)
end

I see these allocations:

julia> @allocated prep = doit()
18134560

julia> @allocated prep = doit()
96

julia> 

(I assume the first call includes the allocations associated with compiling.)

Then, if I add those ismutable_array checks:

function DI.prepare_pushforward_nokwarg(
        strict::Val, f, backend::AutoFiniteDiff, x, tx::NTuple, contexts::Vararg{DI.Context, C}
    ) where {C}
    _sig = DI.signature(f, backend, x, tx, contexts...; strict)
    fc = DI.fix_tail(f, map(DI.unwrap, contexts)...)
    y = fc(x)
    cache = if x isa Number || y isa Number || (! DI.ismutable_array(x)) || (! DI.ismutable_array(y))
        nothing
    else
        JVPCache(copy(x), y, fdtype(backend))
    end
    relstep = if isnothing(backend.relstep)
        default_relstep(fdtype(backend), eltype(x))
    else
        backend.relstep
    end
    absstep = if isnothing(backend.absstep)
        relstep
    else
        backend.absstep
    end
    dir = backend.dir
    return FiniteDiffOneArgPushforwardPrep(_sig, cache, relstep, absstep, dir)
end

I see this in the REPL:

julia> @allocated prep = doit()
751888

julia> @allocated prep = doit()
0

julia> 

Here's another example that isolates the FiniteDiff.JVPCache:

julia> function doit2()
                x = @SVector [1.0, 2.0]
                y = @SVector [2.0, 3.0]
                foo = Val(:forward)
                cache = FiniteDiff.JVPCache(x, y, foo)
                return cache
              end
doit2 (generic function with 1 method)

julia> @allocated cache = doit2()
882800

julia> @allocated cache = doit2()
48

julia> 

[gdalle]

I tried implementing an upstream fix in JuliaDiff/FiniteDiff.jl#216

[dingraha]

@gdalle Excellent, thanks! That fixes the allocations I was seeing with this small example:

julia> function doit2()
               x = @SVector [1.0, 2.0]
               y = @SVector [2.0, 3.0]
               foo = Val(:forward)
               cache = FiniteDiff.JVPCache(x, y, foo)
               return cache
             end
doit2 (generic function with 1 method)

julia> @allocated cache = doit2()
882800

julia> @allocated cache = doit2()
48

julia> 

To avoid the allocations with this example

using FiniteDiff
using DifferentiationInterface
using StaticArrays
function doit()
    backend = AutoFiniteDiff()
    x = @SVector [1.0, 2.0]
    tx = (2.0.*x,)
    f(x) = @. 3.0 * x
    # prep = DifferentiationInterface.prepare_pushforward(f, backend, x, tx);
    prep =  DifferentiationInterface.prepare_pushforward_nokwarg(Val(true), f, backend, x, tx)
    return prep
end

it looks like we'd need to replace the similar calls with copy, like you suggested. I updated the PR to reflect this.

Interestingly, it looks like FiniteDiff.finite_difference_jvp doesn't actually use the cache::JVPCache argument at all, so I guess it doesn't matter what's in there.

[gdalle]

Thanks!
Two questions:

• Any idea why you need the internal form of prepare_pushforward(Val(true), ...) to avoid allocations? Is it because passing strict as a keyword argument causes some form of dynamic dispatch?
• Can you add a small test in the FiniteDiff test file to make sure that we have no regression?

[gdalle]

@dingraha is this ready for another review?

[dingraha]

@gdalle Yes, sorry for the delay.

[gdalle]

Could it also be due to

DifferentiationInterface.jl/DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/onearg.jl

Lines 350 to 352 in 5572e56

	x1 = similar(x)
	fx = similar(y)
	fx1 = similar(y)

?

[dingraha]

IIUC this PR changes those to copy and I still see a few allocations.

[dingraha]

I think the test failures are unrelated to this PR?