diff --git a/DifferentiationInterface/CHANGELOG.md b/DifferentiationInterface/CHANGELOG.md index c4aa193b5..e75526fbf 100644 --- a/DifferentiationInterface/CHANGELOG.md +++ b/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 diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/utils.jl b/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/utils.jl index f53218249..c31832e4e 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/utils.jl +++ b/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) diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/utils.jl b/DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/utils.jl index e0d5b2dfe..e4df681cf 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/utils.jl +++ 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( diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceHyperHessiansExt/DifferentiationInterfaceHyperHessiansExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceHyperHessiansExt/DifferentiationInterfaceHyperHessiansExt.jl index ba18fad9e..b6426527f 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceHyperHessiansExt/DifferentiationInterfaceHyperHessiansExt.jl +++ b/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) diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl index 2c48fc981..f450aaadc 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl +++ b/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 diff --git a/DifferentiationInterface/src/misc/simple_finite_diff.jl b/DifferentiationInterface/src/misc/simple_finite_diff.jl index 36433ff93..3c73c1db3 100644 --- a/DifferentiationInterface/src/misc/simple_finite_diff.jl +++ b/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( diff --git a/DifferentiationInterface/src/utils/batchsize.jl b/DifferentiationInterface/src/utils/batchsize.jl index f484a11db..64c4288b4 100644 --- a/DifferentiationInterface/src/utils/batchsize.jl +++ b/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) diff --git a/DifferentiationInterface/test/Back/HyperHessians/test.jl b/DifferentiationInterface/test/Back/HyperHessians/test.jl index ea7479b33..33f5a7242 100644 --- a/DifferentiationInterface/test/Back/HyperHessians/test.jl +++ b/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( diff --git a/DifferentiationInterface/test/Core/Internals/batchsize.jl b/DifferentiationInterface/test/Core/Internals/batchsize.jl index a97ddbbaa..5ee75cef4 100644 --- a/DifferentiationInterface/test/Core/Internals/batchsize.jl +++ b/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}