DI: Reverse dependency structure #2301

wsmoses · 2025-02-08T17:09:19Z

@gdalle this will enable the DI ext to more properly touch internals if need be

github-actions · 2025-02-08T17:10:16Z

ext/EnzymeDIExt/EnzymeDIExt.jl

+include("reverse_onearg.jl")
+include("reverse_twoarg.jl")
+
+end # module


Suggested change

end # module

end # module

github-actions · 2025-02-08T17:10:16Z

ext/EnzymeDIExt/forward_onearg.jl

+    f::F,
+    ::AutoEnzyme{<:Union{ForwardMode,Nothing}},
+    x,
+    tx::NTuple,
+    contexts::Vararg{DI.Context,C},
+) where {F,C}
+    return DI.NoPushforwardPrep()
+end
+
+function DI.value_and_pushforward(
+    f::F,
+    ::DI.NoPushforwardPrep,
+    backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},


Suggested change

f::F,

::AutoEnzyme{<:Union{ForwardMode,Nothing}},

x,

tx::NTuple,

contexts::Vararg{DI.Context,C},

) where {F,C}

return DI.NoPushforwardPrep()

end

function DI.value_and_pushforward(

f::F,

::DI.NoPushforwardPrep,

backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},

f::F,

::AutoEnzyme{<:Union{ForwardMode, Nothing}},

x,

tx::NTuple,

contexts::Vararg{DI.Context, C},

) where {F, C}

f::F,

::DI.NoPushforwardPrep,

backend::AutoEnzyme{<:Union{ForwardMode, Nothing}},

x,

tx::NTuple{1},

contexts::Vararg{DI.Context, C},

) where {F, C}

github-actions · 2025-02-08T17:10:17Z

ext/EnzymeDIExt/forward_onearg.jl

+    f::F,
+    ::DI.NoPushforwardPrep,
+    backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
+    x,
+    tx::NTuple{B},
+    contexts::Vararg{DI.Context,C},
+) where {F,B,C}


Suggested change

f::F,

::DI.NoPushforwardPrep,

backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},

x,

tx::NTuple{B},

contexts::Vararg{DI.Context,C},

) where {F,B,C}

f::F,

::DI.NoPushforwardPrep,

backend::AutoEnzyme{<:Union{ForwardMode, Nothing}},

x,

tx::NTuple{B},

contexts::Vararg{DI.Context, C},

) where {F, B, C}

github-actions · 2025-02-08T17:10:17Z

ext/EnzymeDIExt/forward_onearg.jl

+    f::F,
+    ::DI.NoPushforwardPrep,
+    backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
+    x,
+    tx::NTuple{1},
+    contexts::Vararg{DI.Context,C},
+) where {F,C}


Suggested change

f::F,

::DI.NoPushforwardPrep,

backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},

x,

tx::NTuple{1},

contexts::Vararg{DI.Context,C},

) where {F,C}

f::F,

::DI.NoPushforwardPrep,

backend::AutoEnzyme{<:Union{ForwardMode, Nothing}},

x,

tx::NTuple{1},

contexts::Vararg{DI.Context, C},

) where {F, C}

github-actions · 2025-02-08T17:10:17Z

ext/EnzymeDIExt/forward_onearg.jl

+    f::F,
+    ::DI.NoPushforwardPrep,
+    backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
+    x,
+    tx::NTuple{B},
+    contexts::Vararg{DI.Context,C},
+) where {F,B,C}


Suggested change

f::F,

::DI.NoPushforwardPrep,

backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},

x,

tx::NTuple{B},

contexts::Vararg{DI.Context,C},

) where {F,B,C}

f::F,

::DI.NoPushforwardPrep,

backend::AutoEnzyme{<:Union{ForwardMode, Nothing}},

x,

tx::NTuple{B},

contexts::Vararg{DI.Context, C},

) where {F, B, C}

github-actions · 2025-02-08T17:10:20Z

ext/EnzymeDIExt/utils.jl

+    f::F, ::AutoEnzyme{M,Nothing}, mode::Mode, ::Val{B}=Val(1)
+) where {F,M,B}
+    return f
+end
+
+@inline function get_f_and_df(
+    f::F, ::AutoEnzyme{M,<:Const}, mode::Mode, ::Val{B}=Val(1)
+) where {F,M,B}
+    return Const(f)
+end
+
+@inline function get_f_and_df(
+    f::F,
+    ::AutoEnzyme{
+        M,
+        <:Union{
+            Duplicated,
+            MixedDuplicated,


Suggested change

f::F, ::AutoEnzyme{M,Nothing}, mode::Mode, ::Val{B}=Val(1)

) where {F,M,B}

return f

end

@inline function get_f_and_df(

f::F, ::AutoEnzyme{M,<:Const}, mode::Mode, ::Val{B}=Val(1)

) where {F,M,B}

return Const(f)

end

@inline function get_f_and_df(

f::F,

::AutoEnzyme{

M,

<:Union{

Duplicated,

MixedDuplicated,

f::F, ::AutoEnzyme{M, Nothing}, mode::Mode, ::Val{B} = Val(1)

) where {F, M, B}

f::F, ::AutoEnzyme{M, <:Const}, mode::Mode, ::Val{B} = Val(1)

) where {F, M, B}

f::F,

::AutoEnzyme{

M,

<:Union{

Duplicated,

MixedDuplicated,

BatchDuplicated,

BatchMixedDuplicated,

DuplicatedNoNeed,

BatchDuplicatedNoNeed,

},

mode::Mode,

::Val{B} = Val(1),

) where {F, M, B}

github-actions · 2025-02-08T17:10:20Z

ext/EnzymeDIExt/utils.jl

+force_annotation(f::F) where {F<:Annotation} = f
+force_annotation(f::F) where {F} = Const(f)
+
+@inline function _translate(
+    ::AutoEnzyme, ::Mode, ::Val{B}, c::Union{DI.Constant,DI.BackendContext}


Suggested change

force_annotation(f::F) where {F<:Annotation} = f

force_annotation(f::F) where {F} = Const(f)

@inline function _translate(

::AutoEnzyme, ::Mode, ::Val{B}, c::Union{DI.Constant,DI.BackendContext}

force_annotation(f::F) where {F <: Annotation} = f

::AutoEnzyme, ::Mode, ::Val{B}, c::Union{DI.Constant, DI.BackendContext}

) where {B}

backend::AutoEnzyme, mode::Mode, ::Val{B}, c::DI.Cache

) where {B}

github-actions · 2025-02-08T17:10:20Z

ext/EnzymeDIExt/utils.jl

+    backend::AutoEnzyme, mode::Mode, ::Val{B}, c::DI.FunctionContext
+) where {B}
+    return force_annotation(get_f_and_df(DI.unwrap(c), backend, mode, Val(B)))
+end


Suggested change

backend::AutoEnzyme, mode::Mode, ::Val{B}, c::DI.FunctionContext

) where {B}

return force_annotation(get_f_and_df(DI.unwrap(c), backend, mode, Val(B)))

end

backend::AutoEnzyme, mode::Mode, ::Val{B}, c::DI.FunctionContext

) where {B}

backend::AutoEnzyme, mode::Mode, ::Val{B}, contexts::Vararg{DI.Context, C}

) where {B, C}

github-actions · 2025-02-08T17:10:20Z

ext/EnzymeDIExt/utils.jl

+set_err(mode::Mode, ::AutoEnzyme{<:Any,Nothing}) = EnzymeCore.set_err_if_func_written(mode)
+set_err(mode::Mode, ::AutoEnzyme{<:Any,<:Annotation}) = mode


Suggested change

set_err(mode::Mode, ::AutoEnzyme{<:Any,Nothing}) = EnzymeCore.set_err_if_func_written(mode)

set_err(mode::Mode, ::AutoEnzyme{<:Any,<:Annotation}) = mode

set_err(mode::Mode, ::AutoEnzyme{<:Any, Nothing}) = EnzymeCore.set_err_if_func_written(mode)

set_err(mode::Mode, ::AutoEnzyme{<:Any, <:Annotation}) = mode

github-actions · 2025-02-08T17:10:20Z

ext/EnzymeDIExt/utils.jl

+function annotate(::Type{BatchDuplicated{T,B}}, x, tx::NTuple{B}) where {T,B}
+    return BatchDuplicated(x, tx)
+end


Suggested change

function annotate(::Type{BatchDuplicated{T,B}}, x, tx::NTuple{B}) where {T,B}

return BatchDuplicated(x, tx)

end

function annotate(::Type{BatchDuplicated{T, B}}, x, tx::NTuple{B}) where {T, B}

batchify_activity(::Type{Active{T}}, ::Val{B}) where {T, B} = Active{T}

batchify_activity(::Type{Duplicated{T}}, ::Val{B}) where {T, B} = BatchDuplicated{T, B}

github-actions · 2025-02-08T17:20:48Z

Benchmark Results

	main	`00803ee`...	main/00803ee98e24c7...
basics/overhead	5.26 ± 0.01 ns	4.64 ± 0.01 ns	1.13
time_to_load	1.1 ± 0.012 s	1.13 ± 0.034 s	0.967

Benchmark Plots

A plot of the benchmark results have been uploaded as an artifact to the workflow run for this PR.
Go to "Actions"->"Benchmark a pull request"->[the most recent run]->"Artifacts" (at the bottom).

wsmoses · 2025-02-08T18:55:30Z

ext/EnzymeDIExt/utils.jl

+    return set_err(ReverseSplitWithPrimal, backend)
+end
+
+set_err(mode::Mode, ::AutoEnzyme{<:Any,Nothing}) = EnzymeCore.set_err_if_func_written(mode)


@gdalle as discussed on slack this probably should be an extension to the set_err_if_func_written function to take an ADmode, so likely we have this in an EnzymeCoreADTypes ext?

wsmoses · 2025-02-08T18:56:54Z

ext/EnzymeDIExt/utils.jl

+forward_withprimal(backend::AutoEnzyme{<:ForwardMode}) = WithPrimal(backend.mode)
+forward_withprimal(::AutoEnzyme{Nothing}) = ForwardWithPrimal
+
+reverse_noprimal(backend::AutoEnzyme{<:ReverseMode}) = NoPrimal(backend.mode)


@gdalle similarly here we can make an ADTypes ext func for get_mode_or_default(AutoEnzyme, defaultMode)

wsmoses · 2025-02-08T18:59:05Z

ext/EnzymeDIExt/reverse_twoarg.jl

+    dy_sametype = convert(typeof(y), only(prep.ty_copy))
+    x_and_dx = Duplicated(x, dx_sametype)
+    y_and_dy = Duplicated(y, dy_sametype)
+    annotated_contexts = translate(backend, mode, Val(1), contexts...)


@gdalle I presume this can be moved into Enzyme.gradient! And have DI call that?

DI: Reverse dependency structure

00803ee

github-actions bot reviewed Feb 8, 2025

View reviewed changes

wsmoses marked this pull request as draft February 8, 2025 18:34

wsmoses commented Feb 8, 2025

View reviewed changes

danielwe mentioned this pull request Feb 16, 2025

Suggestion: remove code-suggester, print diff instead #2312

Closed

		set_err(mode::Mode, ::AutoEnzyme{<:Any,Nothing}) = EnzymeCore.set_err_if_func_written(mode)
		set_err(mode::Mode, ::AutoEnzyme{<:Any,<:Annotation}) = mode

-function annotate(::Type{BatchDuplicated{T,B}}, x, tx::NTuple{B}) where {T,B}
-    return BatchDuplicated(x, tx)
-end
+function annotate(::Type{BatchDuplicated{T, B}}, x, tx::NTuple{B}) where {T, B}
+batchify_activity(::Type{Active{T}}, ::Val{B}) where {T, B} = Active{T}
+batchify_activity(::Type{Duplicated{T}}, ::Val{B}) where {T, B} = BatchDuplicated{T, B}

DI: Reverse dependency structure #2301

Are you sure you want to change the base?

DI: Reverse dependency structure #2301

Uh oh!

Conversation

wsmoses commented Feb 8, 2025

Uh oh!

github-actions bot Feb 8, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot Feb 8, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot Feb 8, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot Feb 8, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot Feb 8, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot Feb 8, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot Feb 8, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot Feb 8, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot Feb 8, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot Feb 8, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Feb 8, 2025

Benchmark Results

Benchmark Plots

Uh oh!

wsmoses Feb 8, 2025

Choose a reason for hiding this comment

Uh oh!

wsmoses Feb 8, 2025

Choose a reason for hiding this comment

Uh oh!

wsmoses Feb 8, 2025

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants