Skip to content

Commit

Permalink
fix: better handling of Enzyme split mode + run Enzyme tests on 1.11 (#…
Browse files Browse the repository at this point in the history 
…654)

* Test Enzyme on 1.11

* Remove matrix inputs

* Explicit EnzymeCore dep

* Skip step 6

* Custom diff implem

* Even more manual diff

* Filter sparse scenarios
  • Loading branch information
@gdalle
gdalle authored Dec 5, 2024
1 parent 3c6e5dc commit 3abcee1
Show file tree
Hide file tree
Showing 5 changed files with 91 additions and 77 deletions.
4 changes: 1 addition & 3 deletions .github/workflows/Test.yml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ jobs:
actions: write
contents: read
strategy:
fail-fast: false # TODO: toggle
fail-fast: true # TODO: toggle
matrix:
version:
- "1.10"
Expand Down Expand Up @@ -57,8 +57,6 @@ jobs:
# version: "1.10"
- version: "1"
group: Back/ChainRules
- version: "1"
group: Back/Enzyme
env:
JULIA_DI_TEST_GROUP: ${{ matrix.group }}
steps:
Expand Down
7 changes: 4 additions & 3 deletions DifferentiationInterface/Project.toml
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "DifferentiationInterface"
uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
authors = ["Guillaume Dalle", "Adrian Hill"]
version = "0.6.25"
version = "0.6.26"

[deps]
ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
Expand All @@ -12,6 +12,7 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
Diffractor = "9f5e2b26-1114-432f-b630-d3fe2085c51c"
Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
FastDifferentiation = "eb9bf01b-bf85-4b60-bf87-ee5de06c00be"
FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
Expand All @@ -29,7 +30,7 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
[extensions]
DifferentiationInterfaceChainRulesCoreExt = "ChainRulesCore"
DifferentiationInterfaceDiffractorExt = "Diffractor"
DifferentiationInterfaceEnzymeExt = "Enzyme"
DifferentiationInterfaceEnzymeExt = ["EnzymeCore", "Enzyme"]
DifferentiationInterfaceFastDifferentiationExt = "FastDifferentiation"
DifferentiationInterfaceFiniteDiffExt = "FiniteDiff"
DifferentiationInterfaceFiniteDifferencesExt = "FiniteDifferences"
Expand All @@ -49,7 +50,7 @@ ADTypes = "1.9.0"
ChainRulesCore = "1.23.0"
DiffResults = "1.1.0"
Diffractor = "=0.2.6"
Enzyme = "0.13.6"
Enzyme = "0.13.17"
ExplicitImports = "1.10.1"
FastDifferentiation = "0.4.1"
FiniteDiff = "2.23.1"
Expand Down
7 changes: 5 additions & 2 deletions ...ationInterface/ext/DifferentiationInterfaceEnzymeExt/DifferentiationInterfaceEnzymeExt.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,12 @@ module DifferentiationInterfaceEnzymeExt
using ADTypes: ADTypes, AutoEnzyme
using Base: Fix1
import DifferentiationInterface as DI
using Enzyme:
using EnzymeCore:
Active,
Annotation,
BatchDuplicated,
BatchMixedDuplicated,
Combined,
Const,
Duplicated,
DuplicatedNoNeed,
Expand All @@ -25,7 +26,9 @@ using Enzyme:
ReverseSplitWidth,
ReverseSplitWithPrimal,
ReverseWithPrimal,
WithPrimal,
Split,
WithPrimal
using Enzyme:
autodiff,
autodiff_thunk,
create_shadows,
Expand Down
9 changes: 6 additions & 3 deletions DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/utils.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,9 +58,12 @@ reverse_noprimal(::AutoEnzyme{Nothing}) = Reverse
reverse_withprimal(backend::AutoEnzyme{<:ReverseMode}) = WithPrimal(backend.mode)
reverse_withprimal(::AutoEnzyme{Nothing}) = ReverseWithPrimal

function reverse_split_withprimal(backend::AutoEnzyme)
mode = ReverseSplitWithPrimal
return set_err(mode, backend)
function reverse_split_withprimal(backend::AutoEnzyme{<:ReverseMode})
return set_err(WithPrimal(Split(backend.mode)), backend)
end

function reverse_split_withprimal(backend::AutoEnzyme{Nothing})
return set_err(ReverseSplitWithPrimal, backend)
end

set_err(mode::Mode, ::AutoEnzyme{<:Any,Nothing}) = EnzymeCore.set_err_if_func_written(mode)
Expand Down
141 changes: 75 additions & 66 deletions DifferentiationInterface/test/Back/Enzyme/test.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,12 +13,20 @@ check_no_implicit_imports(DifferentiationInterface)

LOGGING = get(ENV, "CI", "false") == "false"

function remove_matrix_inputs(scens::Vector{<:Scenario}) # TODO: remove
if VERSION < v"1.11"
return scens
else
# for https://github.com/EnzymeAD/Enzyme.jl/issues/2071
return filter(s -> s.x isa Union{Number,AbstractVector}, scens)
end
end

backends = [
AutoEnzyme(; mode=nothing),
AutoEnzyme(; mode=Enzyme.Forward),
AutoEnzyme(; mode=Enzyme.Reverse),
AutoEnzyme(; mode=Enzyme.Forward, function_annotation=Enzyme.Const),
AutoEnzyme(; mode=Enzyme.Reverse, function_annotation=Enzyme.Const),
AutoEnzyme(; mode=nothing, function_annotation=Enzyme.Const),
]

duplicated_backends = [
Expand All @@ -33,27 +41,25 @@ duplicated_backends = [
end
end;

## First order

test_differentiation(backends, default_scenarios(); excluded=SECOND_ORDER, logging=LOGGING);

test_differentiation(
backends[1:3],
default_scenarios(; include_normal=false, include_constantified=true);
excluded=SECOND_ORDER,
logging=LOGGING,
);

#=
# TODO: reactivate closurified tests once Enzyme#2056 is fixed
test_differentiation(
duplicated_backends,
default_scenarios(; include_normal=false, include_closurified=true);
excluded=SECOND_ORDER,
logging=LOGGING,
);
=#
@testset "First order" begin
test_differentiation(
backends, default_scenarios(); excluded=SECOND_ORDER, logging=LOGGING
)

test_differentiation(
backends[1:3],
default_scenarios(; include_normal=false, include_constantified=true);
excluded=SECOND_ORDER,
logging=LOGGING,
)

test_differentiation(
duplicated_backends,
default_scenarios(; include_normal=false, include_closurified=true);
excluded=SECOND_ORDER,
logging=LOGGING,
)
end

#=
# TODO: reactivate type stability tests
Expand All @@ -68,50 +74,53 @@ test_differentiation(
);
=#

## Second order

test_differentiation(
AutoEnzyme(),
default_scenarios(; include_constantified=true);
excluded=FIRST_ORDER,
logging=LOGGING,
);

test_differentiation(
AutoEnzyme(; mode=Enzyme.Forward);
excluded=vcat(FIRST_ORDER, [:hessian, :hvp]),
logging=LOGGING,
);

test_differentiation(
AutoEnzyme(; mode=Enzyme.Reverse);
excluded=vcat(FIRST_ORDER, [:second_derivative]),
logging=LOGGING,
);

test_differentiation(
SecondOrder(AutoEnzyme(; mode=Enzyme.Reverse), AutoEnzyme(; mode=Enzyme.Forward));
logging=LOGGING,
);

## Sparse
@testset "Second order" begin
test_differentiation(
[
AutoEnzyme(),
SecondOrder(
AutoEnzyme(; mode=Enzyme.Reverse), AutoEnzyme(; mode=Enzyme.Forward)
),
],
remove_matrix_inputs(default_scenarios(; include_constantified=true));
excluded=FIRST_ORDER,
logging=LOGGING,
)

test_differentiation(
AutoEnzyme(; mode=Enzyme.Forward);
excluded=vcat(FIRST_ORDER, [:hessian, :hvp]),
logging=LOGGING,
)
end

test_differentiation(
MyAutoSparse.(AutoEnzyme(; function_annotation=Enzyme.Const)),
sparse_scenarios();
sparsity=true,
logging=LOGGING,
);
@testset "Sparse" begin
test_differentiation(
MyAutoSparse.(AutoEnzyme(; function_annotation=Enzyme.Const)),
if VERSION < v"1.11"
sparse_scenarios()
else
filter(sparse_scenarios()) do s
# for https://github.com/EnzymeAD/Enzyme.jl/issues/2168
(s.x isa AbstractVector) &&
(s.f != DIT.sumdiffcube) &&
(s.f != DIT.sumdiffcube_mat)
end
end;
sparsity=true,
logging=LOGGING,
)
end

##
@testset "Static" begin
filtered_static_scenarios = filter(static_scenarios()) do s
DIT.operator_place(s) == :out && DIT.function_place(s) == :out
end

filtered_static_scenarios = filter(static_scenarios()) do s
DIT.operator_place(s) == :out && DIT.function_place(s) == :out
test_differentiation(
[AutoEnzyme(; mode=Enzyme.Forward), AutoEnzyme(; mode=Enzyme.Reverse)],
filtered_static_scenarios;
excluded=SECOND_ORDER,
logging=LOGGING,
)
end

test_differentiation(
[AutoEnzyme(; mode=Enzyme.Forward), AutoEnzyme(; mode=Enzyme.Reverse)],
filtered_static_scenarios;
excluded=SECOND_ORDER,
logging=LOGGING,
)

2 comments on commit 3abcee1

@gdalle
Copy link
Member Author

@gdalle gdalle commented on 3abcee1 Dec 5, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@JuliaRegistrator register subdir=DifferentiationInterface

@JuliaRegistrator
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Registration pull request created: JuliaRegistries/General/120770

Tip: Release Notes

Did you know you can add release notes too? Just add markdown formatted text underneath the comment after the text
"Release notes:" and it will be added to the registry PR, and if TagBot is installed it will also be added to the
release that TagBot creates. i.e.

@JuliaRegistrator register

Release notes:

## Breaking changes

- blah

To add them here just re-invoke and the PR will be updated.

Tagging

After the above pull request is merged, it is recommended that a tag is created on this repository for the registered package version.

This will be done automatically if the Julia TagBot GitHub Action is installed, or can be done manually through the github interface, or via:

git tag -a DifferentiationInterface-v0.6.26 -m "<description of version>" 3abcee154396c8c5c36d88b7ee643821e627bd99
git push origin DifferentiationInterface-v0.6.26

Please sign in to comment.