perf: check mutability of array before preallocating dual buffer (#619)

* perf: check mutability of array before preallocating ForwardDiff dual buffer

* Add allocation testing

Author: gdalle

DifferentiationInterface/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterface"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.6.22"
+version = "0.6.23"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/DifferentiationInterfaceForwardDiffExt.jl

@@ -25,7 +25,8 @@ using DifferentiationInterface:
     outer,
     shuffled_gradient,
     unwrap,
-    with_contexts
+    with_contexts,
+    ismutable_array
 import ForwardDiff.DiffResults as DR
 using ForwardDiff.DiffResults:
     DiffResults, DiffResult, GradientResult, HessianResult, MutableDiffResult

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/onearg.jl

@@ -68,7 +68,11 @@ function DI.prepare_pushforward(
     f::F, backend::AutoForwardDiff, x, tx::NTuple, contexts::Vararg{Context,C}
 ) where {F,C}
     T = tag_type(f, backend, x)
-    xdual_tmp = make_dual_similar(T, x, tx)
+    if ismutable_array(x)
+        xdual_tmp = make_dual_similar(T, x, tx)
+    else
+        xdual_tmp = nothing
+    end
     return ForwardDiffOneArgPushforwardPrep{T,typeof(xdual_tmp)}(xdual_tmp)
 end
 
@@ -92,8 +96,12 @@ function compute_ydual_onearg(
     tx::NTuple{B},
     contexts::Vararg{Context,C},
 ) where {F,T,B,C}
-    (; xdual_tmp) = prep
-    make_dual!(T, xdual_tmp, x, tx)
+    if ismutable_array(x)
+        make_dual!(T, prep.xdual_tmp, x, tx)
+        xdual_tmp = prep.xdual_tmp
+    else
+        xdual_tmp = make_dual(T, x, tx)
+    end
     contexts_dual = translate(T, Val(B), contexts...)
     ydual = f(xdual_tmp, contexts_dual...)
     return ydual

DifferentiationInterface/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl

@@ -13,6 +13,8 @@ function DI.stack_vec_row(t::NTuple{B,<:StaticArray}) where {B}
     return vcat(transpose.(map(vec, t))...)
 end
 
+DI.ismutable_array(::Type{<:SArray}) = false
+
 function DI.BatchSizeSettings(::AutoForwardDiff{nothing}, x::StaticArray)
     return BatchSizeSettings{length(x),true,true}(length(x))
 end

DifferentiationInterface/src/utils/linalg.jl

@@ -1,2 +1,5 @@
 stack_vec_col(t::NTuple) = stack(vec, t; dims=2)
 stack_vec_row(t::NTuple) = stack(vec, t; dims=1)
+
+ismutable_array(::Type) = true
+ismutable_array(x) = ismutable_array(typeof(x))

DifferentiationInterface/test/Back/ForwardDiff/test.jl

@@ -3,6 +3,7 @@ Pkg.add("ForwardDiff")
 
 using ComponentArrays: ComponentArrays
 using DifferentiationInterface, DifferentiationInterfaceTest
+import DifferentiationInterfaceTest as DIT
 using ForwardDiff: ForwardDiff
 using StaticArrays: StaticArrays
 using Test
@@ -65,3 +66,19 @@ test_differentiation(
 ## Static
 
 test_differentiation(AutoForwardDiff(), static_scenarios(); logging=LOGGING)
+
+@testset verbose = true "No allocations on StaticArrays" begin
+    filtered_static_scenarios = filter(static_scenarios(; include_batchified=false)) do scen
+        DIT.function_place(scen) == :out && DIT.operator_place(scen) == :out
+    end
+    data = benchmark_differentiation(
+        AutoForwardDiff(),
+        filtered_static_scenarios;
+        benchmark=:prepared,
+        excluded=[:hessian, :pullback],  # TODO: figure this out
+        logging=LOGGING,
+    )
+    @testset "$(row[:scenario])" for row in eachrow(data)
+        @test row[:allocs] == 0
+    end
+end;