From a95bf029bdb59fe4f9051040197109dd042929b3 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Sat, 16 Aug 2025 00:08:42 +0200 Subject: [PATCH 1/4] feat: add runtime activity and chunksize parameters to `AutoEnzyme` --- Project.toml | 2 +- ext/ADTypesConstructionBaseExt.jl | 10 +++-- src/dense.jl | 63 +++++++++++++++++++++++-------- test/dense.jl | 20 ++++++++++ test/misc.jl | 1 + 5 files changed, 75 insertions(+), 21 deletions(-) diff --git a/Project.toml b/Project.toml index 240038e..6860184 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "ADTypes" uuid = "47edcb42-4c32-4615-8424-f2b9edc5f35b" authors = ["Vaibhav Dixit , Guillaume Dalle and contributors"] -version = "1.17.0" +version = "1.18.0" [deps] ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" diff --git a/ext/ADTypesConstructionBaseExt.jl b/ext/ADTypesConstructionBaseExt.jl index 883dd7f..1340de6 100644 --- a/ext/ADTypesConstructionBaseExt.jl +++ b/ext/ADTypesConstructionBaseExt.jl @@ -3,12 +3,14 @@ module ADTypesConstructionBaseExt using ADTypes: AutoEnzyme, AutoForwardDiff, AutoPolyesterForwardDiff using ConstructionBase: ConstructionBase -struct InternalAutoEnzymeReconstructor{A} end +struct InternalAutoEnzymeReconstructor{A, R, C} end -InternalAutoEnzymeReconstructor{A}(mode::M) where {M, A} = AutoEnzyme{M, A}(mode) +function InternalAutoEnzymeReconstructor{A, R, C}(mode::M) where {M, A, R, C} + AutoEnzyme{M, A, R, C}(mode) +end -function ConstructionBase.constructorof(::Type{<:AutoEnzyme{M, A}}) where {M, A} - return InternalAutoEnzymeReconstructor{A} +function ConstructionBase.constructorof(::Type{<:AutoEnzyme{M, A, R, C}}) where {M, A, R, C} + return InternalAutoEnzymeReconstructor{A, R, C} end function ConstructionBase.constructorof(::Type{<:AutoForwardDiff{chunksize}}) where {chunksize} diff --git a/src/dense.jl b/src/dense.jl index 38072e6..5e523c8 100644 --- a/src/dense.jl +++ b/src/dense.jl @@ -39,7 +39,7 @@ struct AutoDiffractor <: AbstractADType end mode(::AutoDiffractor) = ForwardOrReverseMode() """ - AutoEnzyme{M,A} + AutoEnzyme{M,A,R,C} Struct used to select the [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl) backend for automatic differentiation. @@ -47,38 +47,69 @@ Defined by [ADTypes.jl](https://github.com/SciML/ADTypes.jl). # Constructors - AutoEnzyme(; mode::M=nothing, function_annotation::Type{A}=Nothing) + AutoEnzyme(; + mode::Union{EnzymeCore.Mode,ADTypes.AbstractMode,Nothing}=nothing, + function_annotation::Type{<:Union{EnzymeCore.Annotation,Nothing}}=Nothing, + runtime_activity::Union{Bool,Nothing}=nothing, + chunksize::Union{Int,Float64,Nothing}=nothing, + ) -# Type parameters + - `mode::M` determines the autodiff mode (forward or reverse). It can be: - - `A` determines how the function `f` to differentiate is passed to Enzyme. It can be: + + a mode object from EnzymeCore.jl, like `EnzymeCore.Forward` or `EnzymeCore.Reverse` (possibly modified with additional settings like runtime activity) + + a mode object from ADTypes.jl, either [`ADTypes.ForwardMode`](@ref) or [`ADTypes.ReverseMode`](@ref) + + `nothing` to choose the best mode automatically + + - `A=function_annotation` determines how the function `f` to differentiate is passed to Enzyme. It can be: + a subtype of `EnzymeCore.Annotation` (like `EnzymeCore.Const` or `EnzymeCore.Duplicated`) to enforce a given annotation - + `Nothing` to simply pass `f` and let Enzyme choose the most appropriate annotation -# Fields + + `Nothing` (the type, not the object) to simply pass `f` and let Enzyme choose the most appropriate annotation + - `R=runtime_activity` determines whether [runtime activity](https://enzymead.github.io/Enzyme.jl/stable/faq/#faq-runtime-activity) is activated by Enzyme. It can be: - - `mode::M` determines the autodiff mode (forward or reverse). It can be: + + a `Bool` to force a given setting - + an object subtyping `EnzymeCore.Mode` (like `EnzymeCore.Forward` or `EnzymeCore.Reverse`) if a specific mode is required - + `nothing` to choose the best mode automatically + + `nothing` to fall back on the setting of the `mode` object + - `C=chunksize` determines the number of derivatives evaluated simultaneously when computing operators like a Jacobian or a forward-mode gradient. It can be: + + + a positive `Int` to fix a constant chunk size + + `Inf` to pick the maximum chunk size, corresponding to the array length + + `nothing` to choose a good chunk size automatically® """ -struct AutoEnzyme{M, A} <: AbstractADType +struct AutoEnzyme{M, A, R, C} <: AbstractADType mode::M + + function AutoEnzyme{M, A, R, C}(mode::M) where {M, A, R, C} + @assert R isa Union{Nothing, Bool} + @assert C isa Union{Nothing, Int, Float64} + if C isa Int + @assert C > 0 + elseif C isa Float64 + @assert C == Inf + end + return new{M, A, R, C}(mode) + end end function AutoEnzyme(; - mode::M = nothing, function_annotation::Type{A} = Nothing) where {M, A} - return AutoEnzyme{M, A}(mode) + mode::M = nothing, + function_annotation::Type{A} = Nothing, + runtime_activity::Union{Nothing, Bool} = nothing, + chunksize::Union{Nothing, Int, Float64} = nothing +) where {M, A} + return AutoEnzyme{M, A, runtime_activity, chunksize}(mode) end mode(::AutoEnzyme) = ForwardOrReverseMode() # specialized in the extension +mode(::AutoEnzyme{<:ForwardMode}) = ForwardMode() +mode(::AutoEnzyme{<:ReverseMode}) = ReverseMode() -function Base.show(io::IO, backend::AutoEnzyme{M, A}) where {M, A} +function Base.show(io::IO, backend::AutoEnzyme{M, A, R, C}) where {M, A, R, C} print(io, AutoEnzyme, "(") - !isnothing(backend.mode) && print(io, "mode=", repr(backend.mode; context = io)) - !isnothing(backend.mode) && !(A <: Nothing) && print(io, ", ") - !(A <: Nothing) && print(io, "function_annotation=", repr(A; context = io)) + !isnothing(backend.mode) && print(io, "mode=", repr(backend.mode; context = io), ", ") + !(A <: Nothing) && print(io, "function_annotation=", repr(A; context = io), ", ") + !(R === nothing) && print(io, "runtime_activity=", repr(R; context = io), ", ") + !(C === nothing) && print(io, "chunksize=", repr(C; context = io)) print(io, ")") end diff --git a/test/dense.jl b/test/dense.jl index 307403d..ba3adb4 100644 --- a/test/dense.jl +++ b/test/dense.jl @@ -25,12 +25,17 @@ end @test mode(ad) isa ForwardOrReverseMode end +get_runtime_activity(::AutoEnzyme{M, A, R, C}) where {M, A, R, C} = R +get_chunksize(::AutoEnzyme{M, A, R, C}) where {M, A, R, C} = C + @testset "AutoEnzyme" begin ad = AutoEnzyme() @test ad isa AbstractADType @test ad isa AutoEnzyme{Nothing, Nothing} @test mode(ad) isa ForwardOrReverseMode @test ad.mode === nothing + @test get_runtime_activity(ad) === nothing + @test get_chunksize(ad) === nothing ad = AutoEnzyme(; mode = EnzymeCore.Forward) @test ad isa AbstractADType @@ -50,6 +55,21 @@ end @test ad isa AutoEnzyme{typeof(EnzymeCore.Reverse), EnzymeCore.Duplicated} @test mode(ad) isa ReverseMode @test ad.mode == EnzymeCore.Reverse + + ad = AutoEnzyme(; mode = ForwardMode(), runtime_activity = true, chunksize = Inf) + @test mode(ad) isa ForwardMode + @test get_runtime_activity(ad) + @test get_chunksize(ad) == Inf + + ad = AutoEnzyme(; mode = ReverseMode(), runtime_activity = false, chunksize = 3) + @test mode(ad) isa ReverseMode + @test !get_runtime_activity(ad) + @test get_chunksize(ad) == 3 + + @test_throws TypeError AutoEnzyme(; runtime_activity = :yes) + @test_throws TypeError AutoEnzyme(; chunksize = :big) + @test_throws AssertionError AutoEnzyme(; chunksize = 0) + @test_throws AssertionError AutoEnzyme(; chunksize = 1.3) end @testset "AutoFastDifferentiation" begin diff --git a/test/misc.jl b/test/misc.jl index 95c73bc..616cf9b 100644 --- a/test/misc.jl +++ b/test/misc.jl @@ -34,6 +34,7 @@ for backend in [ ADTypes.AutoEnzyme(mode = :forward), ADTypes.AutoEnzyme(function_annotation = Val{:forward}), ADTypes.AutoEnzyme(mode = :reverse, function_annotation = Val{:duplicated}), + ADTypes.AutoEnzyme(runtime_activity = true, chunksize = 2), ADTypes.AutoFastDifferentiation(), ADTypes.AutoFiniteDiff(), ADTypes.AutoFiniteDiff(fdtype = :fd, fdjtype = :fdj, fdhtype = :fdh), From 428bae4fd8933c8f10a532391704b6aa3ff9eb48 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Wed, 27 Aug 2025 07:55:10 +0200 Subject: [PATCH 2/4] Update src/dense.jl --- src/dense.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/dense.jl b/src/dense.jl index 5e523c8..2804f91 100644 --- a/src/dense.jl +++ b/src/dense.jl @@ -74,7 +74,7 @@ Defined by [ADTypes.jl](https://github.com/SciML/ADTypes.jl). + a positive `Int` to fix a constant chunk size + `Inf` to pick the maximum chunk size, corresponding to the array length - + `nothing` to choose a good chunk size automatically® + + `nothing` to choose a good chunk size automatically """ struct AutoEnzyme{M, A, R, C} <: AbstractADType mode::M From 8e4cd221263893b9bd8d5a69b144536fd20cf6a4 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Wed, 27 Aug 2025 15:17:58 +0200 Subject: [PATCH 3/4] Remove ADTypes mode --- src/dense.jl | 5 +---- test/dense.jl | 6 ++---- 2 files changed, 3 insertions(+), 8 deletions(-) diff --git a/src/dense.jl b/src/dense.jl index 2804f91..048eb6b 100644 --- a/src/dense.jl +++ b/src/dense.jl @@ -48,7 +48,7 @@ Defined by [ADTypes.jl](https://github.com/SciML/ADTypes.jl). # Constructors AutoEnzyme(; - mode::Union{EnzymeCore.Mode,ADTypes.AbstractMode,Nothing}=nothing, + mode::Union{EnzymeCore.Mode,Nothing}=nothing, function_annotation::Type{<:Union{EnzymeCore.Annotation,Nothing}}=Nothing, runtime_activity::Union{Bool,Nothing}=nothing, chunksize::Union{Int,Float64,Nothing}=nothing, @@ -57,7 +57,6 @@ Defined by [ADTypes.jl](https://github.com/SciML/ADTypes.jl). - `mode::M` determines the autodiff mode (forward or reverse). It can be: + a mode object from EnzymeCore.jl, like `EnzymeCore.Forward` or `EnzymeCore.Reverse` (possibly modified with additional settings like runtime activity) - + a mode object from ADTypes.jl, either [`ADTypes.ForwardMode`](@ref) or [`ADTypes.ReverseMode`](@ref) + `nothing` to choose the best mode automatically - `A=function_annotation` determines how the function `f` to differentiate is passed to Enzyme. It can be: @@ -101,8 +100,6 @@ function AutoEnzyme(; end mode(::AutoEnzyme) = ForwardOrReverseMode() # specialized in the extension -mode(::AutoEnzyme{<:ForwardMode}) = ForwardMode() -mode(::AutoEnzyme{<:ReverseMode}) = ReverseMode() function Base.show(io::IO, backend::AutoEnzyme{M, A, R, C}) where {M, A, R, C} print(io, AutoEnzyme, "(") diff --git a/test/dense.jl b/test/dense.jl index ba3adb4..88f9078 100644 --- a/test/dense.jl +++ b/test/dense.jl @@ -56,13 +56,11 @@ get_chunksize(::AutoEnzyme{M, A, R, C}) where {M, A, R, C} = C @test mode(ad) isa ReverseMode @test ad.mode == EnzymeCore.Reverse - ad = AutoEnzyme(; mode = ForwardMode(), runtime_activity = true, chunksize = Inf) - @test mode(ad) isa ForwardMode + ad = AutoEnzyme(; runtime_activity = true, chunksize = Inf) @test get_runtime_activity(ad) @test get_chunksize(ad) == Inf - ad = AutoEnzyme(; mode = ReverseMode(), runtime_activity = false, chunksize = 3) - @test mode(ad) isa ReverseMode + ad = AutoEnzyme(; runtime_activity = false, chunksize = 3) @test !get_runtime_activity(ad) @test get_chunksize(ad) == 3 From 2790b3cedd5ee2acf37b1c9a63f2378f6a7f2ab7 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Wed, 27 Aug 2025 15:35:18 +0200 Subject: [PATCH 4/4] Remove runtime activity --- ext/ADTypesConstructionBaseExt.jl | 10 +++++----- src/dense.jl | 21 ++++++--------------- test/dense.jl | 17 +++++++++-------- test/misc.jl | 2 +- 4 files changed, 21 insertions(+), 29 deletions(-) diff --git a/ext/ADTypesConstructionBaseExt.jl b/ext/ADTypesConstructionBaseExt.jl index 1340de6..9700213 100644 --- a/ext/ADTypesConstructionBaseExt.jl +++ b/ext/ADTypesConstructionBaseExt.jl @@ -3,14 +3,14 @@ module ADTypesConstructionBaseExt using ADTypes: AutoEnzyme, AutoForwardDiff, AutoPolyesterForwardDiff using ConstructionBase: ConstructionBase -struct InternalAutoEnzymeReconstructor{A, R, C} end +struct InternalAutoEnzymeReconstructor{A, C} end -function InternalAutoEnzymeReconstructor{A, R, C}(mode::M) where {M, A, R, C} - AutoEnzyme{M, A, R, C}(mode) +function InternalAutoEnzymeReconstructor{A, C}(mode::M) where {M, A, C} + AutoEnzyme{M, A, C}(mode) end -function ConstructionBase.constructorof(::Type{<:AutoEnzyme{M, A, R, C}}) where {M, A, R, C} - return InternalAutoEnzymeReconstructor{A, R, C} +function ConstructionBase.constructorof(::Type{<:AutoEnzyme{M, A, C}}) where {M, A, C} + return InternalAutoEnzymeReconstructor{A, C} end function ConstructionBase.constructorof(::Type{<:AutoForwardDiff{chunksize}}) where {chunksize} diff --git a/src/dense.jl b/src/dense.jl index 048eb6b..19b232e 100644 --- a/src/dense.jl +++ b/src/dense.jl @@ -39,7 +39,7 @@ struct AutoDiffractor <: AbstractADType end mode(::AutoDiffractor) = ForwardOrReverseMode() """ - AutoEnzyme{M,A,R,C} + AutoEnzyme{M,A,C} Struct used to select the [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl) backend for automatic differentiation. @@ -50,7 +50,6 @@ Defined by [ADTypes.jl](https://github.com/SciML/ADTypes.jl). AutoEnzyme(; mode::Union{EnzymeCore.Mode,Nothing}=nothing, function_annotation::Type{<:Union{EnzymeCore.Annotation,Nothing}}=Nothing, - runtime_activity::Union{Bool,Nothing}=nothing, chunksize::Union{Int,Float64,Nothing}=nothing, ) @@ -64,48 +63,40 @@ Defined by [ADTypes.jl](https://github.com/SciML/ADTypes.jl). + a subtype of `EnzymeCore.Annotation` (like `EnzymeCore.Const` or `EnzymeCore.Duplicated`) to enforce a given annotation + `Nothing` (the type, not the object) to simply pass `f` and let Enzyme choose the most appropriate annotation - - `R=runtime_activity` determines whether [runtime activity](https://enzymead.github.io/Enzyme.jl/stable/faq/#faq-runtime-activity) is activated by Enzyme. It can be: - - + a `Bool` to force a given setting - - + `nothing` to fall back on the setting of the `mode` object - `C=chunksize` determines the number of derivatives evaluated simultaneously when computing operators like a Jacobian or a forward-mode gradient. It can be: + a positive `Int` to fix a constant chunk size + `Inf` to pick the maximum chunk size, corresponding to the array length + `nothing` to choose a good chunk size automatically """ -struct AutoEnzyme{M, A, R, C} <: AbstractADType +struct AutoEnzyme{M, A, C} <: AbstractADType mode::M - function AutoEnzyme{M, A, R, C}(mode::M) where {M, A, R, C} - @assert R isa Union{Nothing, Bool} + function AutoEnzyme{M, A, C}(mode::M) where {M, A, C} @assert C isa Union{Nothing, Int, Float64} if C isa Int @assert C > 0 elseif C isa Float64 @assert C == Inf end - return new{M, A, R, C}(mode) + return new{M, A, C}(mode) end end function AutoEnzyme(; mode::M = nothing, function_annotation::Type{A} = Nothing, - runtime_activity::Union{Nothing, Bool} = nothing, chunksize::Union{Nothing, Int, Float64} = nothing ) where {M, A} - return AutoEnzyme{M, A, runtime_activity, chunksize}(mode) + return AutoEnzyme{M, A, chunksize}(mode) end mode(::AutoEnzyme) = ForwardOrReverseMode() # specialized in the extension -function Base.show(io::IO, backend::AutoEnzyme{M, A, R, C}) where {M, A, R, C} +function Base.show(io::IO, backend::AutoEnzyme{M, A, C}) where {M, A, C} print(io, AutoEnzyme, "(") !isnothing(backend.mode) && print(io, "mode=", repr(backend.mode; context = io), ", ") !(A <: Nothing) && print(io, "function_annotation=", repr(A; context = io), ", ") - !(R === nothing) && print(io, "runtime_activity=", repr(R; context = io), ", ") !(C === nothing) && print(io, "chunksize=", repr(C; context = io)) print(io, ")") end diff --git a/test/dense.jl b/test/dense.jl index 88f9078..022c081 100644 --- a/test/dense.jl +++ b/test/dense.jl @@ -25,8 +25,7 @@ end @test mode(ad) isa ForwardOrReverseMode end -get_runtime_activity(::AutoEnzyme{M, A, R, C}) where {M, A, R, C} = R -get_chunksize(::AutoEnzyme{M, A, R, C}) where {M, A, R, C} = C +get_chunksize(::AutoEnzyme{M, A, C}) where {M, A, C} = C @testset "AutoEnzyme" begin ad = AutoEnzyme() @@ -34,7 +33,6 @@ get_chunksize(::AutoEnzyme{M, A, R, C}) where {M, A, R, C} = C @test ad isa AutoEnzyme{Nothing, Nothing} @test mode(ad) isa ForwardOrReverseMode @test ad.mode === nothing - @test get_runtime_activity(ad) === nothing @test get_chunksize(ad) === nothing ad = AutoEnzyme(; mode = EnzymeCore.Forward) @@ -56,15 +54,18 @@ get_chunksize(::AutoEnzyme{M, A, R, C}) where {M, A, R, C} = C @test mode(ad) isa ReverseMode @test ad.mode == EnzymeCore.Reverse - ad = AutoEnzyme(; runtime_activity = true, chunksize = Inf) - @test get_runtime_activity(ad) + ad = AutoEnzyme(; chunksize = nothing) + @test get_chunksize(ad) === nothing + + ad = AutoEnzyme(; chunksize = 3) + @test get_chunksize(ad) == 3 + + ad = AutoEnzyme(; chunksize = Inf) @test get_chunksize(ad) == Inf - ad = AutoEnzyme(; runtime_activity = false, chunksize = 3) - @test !get_runtime_activity(ad) + ad = AutoEnzyme(; chunksize = 3) @test get_chunksize(ad) == 3 - @test_throws TypeError AutoEnzyme(; runtime_activity = :yes) @test_throws TypeError AutoEnzyme(; chunksize = :big) @test_throws AssertionError AutoEnzyme(; chunksize = 0) @test_throws AssertionError AutoEnzyme(; chunksize = 1.3) diff --git a/test/misc.jl b/test/misc.jl index 616cf9b..30ac9ad 100644 --- a/test/misc.jl +++ b/test/misc.jl @@ -34,7 +34,7 @@ for backend in [ ADTypes.AutoEnzyme(mode = :forward), ADTypes.AutoEnzyme(function_annotation = Val{:forward}), ADTypes.AutoEnzyme(mode = :reverse, function_annotation = Val{:duplicated}), - ADTypes.AutoEnzyme(runtime_activity = true, chunksize = 2), + ADTypes.AutoEnzyme(chunksize = 2), ADTypes.AutoFastDifferentiation(), ADTypes.AutoFiniteDiff(), ADTypes.AutoFiniteDiff(fdtype = :fd, fdjtype = :fdj, fdhtype = :fdh),