fix: avoid double dynamic dispatch in pick_batchsize with run-time batch size

DifferentiationInterface/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased](https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.21...main)
 
+### Fixed
+
+- Speed up `pick_batchsize` when the batch size is only known at run time ([#1076](https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/1076))
+
 ## [0.7.21](https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.20...v0.7.21)
 
 ### Changed

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/utils.jl

@@ -10,7 +10,7 @@ const AnyDuplicated = Union{
 # until https://github.com/EnzymeAD/Enzyme.jl/pull/1545 is merged
 function DI.pick_batchsize(::AutoEnzyme, N::Integer)
     B = DI.reasonable_batchsize(N, 16)
-    return DI.BatchSizeSettings{B}(N)
+    return DI.BatchSizeSettings(B, N)
 end
 
 to_val(::DI.BatchSizeSettings{B}) where {B} = Val(B)

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/utils.jl

@@ -1,10 +1,10 @@
 function DI.pick_batchsize(::AutoForwardDiff{nothing}, N::Integer)
     chunksize = pickchunksize(N)
-    return DI.BatchSizeSettings{chunksize}(N)
+    return DI.BatchSizeSettings(chunksize, N)
 end
 
 function DI.pick_batchsize(::AutoForwardDiff{chunksize}, N::Integer) where {chunksize}
-    return DI.BatchSizeSettings{chunksize}(N)
+    return DI.BatchSizeSettings(chunksize, N)
 end
 
 function DI.threshold_batchsize(

DifferentiationInterface/ext/DifferentiationInterfaceHyperHessiansExt/DifferentiationInterfaceHyperHessiansExt.jl

@@ -29,12 +29,12 @@ chunk_from_backend(::DI.AutoHyperHessians{CS}, ::Integer, ::Type) where {CS} = C
 
 function DI.pick_batchsize(backend::DI.AutoHyperHessians, x::AbstractArray)
     B = chunksize(chunk_from_backend(backend, x))
-    return DI.BatchSizeSettings{B}(length(x))
+    return DI.BatchSizeSettings(B, length(x))
 end
 
 function DI.pick_batchsize(backend::DI.AutoHyperHessians, N::Integer)
     B = chunksize(chunk_from_backend(backend, N, Float64))
-    return DI.BatchSizeSettings{B}(N)
+    return DI.BatchSizeSettings(B, N)
 end
 
 function DI.threshold_batchsize(backend::DI.AutoHyperHessians{nothing}, ::Integer)

DifferentiationInterface/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl

@@ -29,11 +29,11 @@ end
 function DI.pick_batchsize(
         ::DI.AutoSimpleFiniteDiff{chunksize}, x::StaticArray
     ) where {chunksize}
-    return DI.BatchSizeSettings{chunksize}(Val(length(x)))
+    return DI.BatchSizeSettings(chunksize, length(x))
 end
 
 function DI.pick_batchsize(::AutoForwardDiff{chunksize}, x::StaticArray) where {chunksize}
-    return DI.BatchSizeSettings{chunksize}(Val(length(x)))
+    return DI.BatchSizeSettings(chunksize, length(x))
 end
 
 end

DifferentiationInterface/src/misc/simple_finite_diff.jl

@@ -22,11 +22,11 @@ inner_preparation_behavior(::AutoSimpleFiniteDiff) = PrepareInnerSimple()
 
 function pick_batchsize(::AutoSimpleFiniteDiff{nothing}, N::Integer)
     B = reasonable_batchsize(N, 12)
-    return BatchSizeSettings{B}(N)
+    return BatchSizeSettings(B, N)
 end
 
 function pick_batchsize(::AutoSimpleFiniteDiff{chunksize}, N::Integer) where {chunksize}
-    return BatchSizeSettings{chunksize}(N)
+    return BatchSizeSettings(chunksize, N)
 end
 
 function threshold_batchsize(

DifferentiationInterface/src/utils/batchsize.jl

@@ -32,14 +32,8 @@ function BatchSizeSettings{B, singlebatch, aligned}(N::Integer) where {B, single
     return BatchSizeSettings{B, singlebatch, aligned}(N, A, B_last)
 end
 
-function BatchSizeSettings{B}(::Val{N}) where {B, N}
-    singlebatch = B == N
-    aligned = (B == N == 0) || (N % B == 0)
-    return BatchSizeSettings{B, singlebatch, aligned}(N)
-end
-
-function BatchSizeSettings{B}(N::Integer) where {B}
-    # type-unstable
+function BatchSizeSettings(B::Int, N::Int)
+    # type-unstable, unless `B` and `N` are constant-propagated
     singlebatch = B == N
     aligned = (B == N == 0) || (N % B == 0)
     return BatchSizeSettings{B, singlebatch, aligned}(N)

DifferentiationInterface/test/Back/HyperHessians/test.jl

@@ -24,6 +24,15 @@ for backend in backends
     test_counterparts(backend)
 end
 
+@testset "Batch size" begin
+    @test DI.pick_batchsize(DI.AutoHyperHessians(), rand(10)) isa DI.BatchSizeSettings
+    @test DI.pick_batchsize(DI.AutoHyperHessians(), 10) isa DI.BatchSizeSettings
+    @test DI.pick_batchsize(DI.AutoHyperHessians(; chunksize = 4), rand(10)) isa
+        DI.BatchSizeSettings{4}
+    @test DI.pick_batchsize(DI.AutoHyperHessians(; chunksize = 4), 10) isa
+        DI.BatchSizeSettings{4}
+end
+
 scenarios = default_scenarios(; include_constantified = true, include_cachified = true)
 
 test_differentiation(

DifferentiationInterface/test/Core/Internals/batchsize.jl

@@ -25,6 +25,16 @@ BSS = BatchSizeSettings
     )
 end
 
+@testset "Constructor" begin
+    @test BSS(1, 0) === BSS{1, false, true}(0, 0, 0)
+    @test BSS(0, 0) === BSS{0, true, true}(0, 0, 0)
+    @test BSS(1, 1) === BSS{1, true, true}(1, 1, 0)
+    @test BSS(2, 2) === BSS{2, true, true}(2, 1, 0)
+    @test BSS(4, 8) === BSS{4, false, true}(8, 2, 0)
+    @test BSS(8, 13) === BSS{8, false, false}(13, 2, 5)
+    @test_throws ArgumentError BSS(3, 2)
+end
+
 @testset "SimpleFiniteDiff (adaptive)" begin
     @test (pick_batchsize(AutoSimpleFiniteDiff(), zeros(0))) isa BSS{1, false, true}
     @test (pick_batchsize(AutoSimpleFiniteDiff(), zeros(2))) isa BSS{2, true, true}