Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Import needs_primal from Reactant.jl #2021

Merged
merged 8 commits into from
Oct 30, 2024
Merged

Import needs_primal from Reactant.jl #2021

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
3415d5d
Implement `needs_primal`
mofeing Oct 28, 2024
efd2925
Test `needs_primal`
mofeing Oct 28, 2024
4ebf7a8
Export `needs_primal`
mofeing Oct 29, 2024
7b6e394
Bump versions
mofeing Oct 29, 2024
2eaaf0e
Move `needs_primal` tests EnzymeCore
mofeing Oct 30, 2024
cede9c7
Use tilde version specifier to also support v0.8.4
mofeing Oct 30, 2024
9597c03
Move `WithPrimal`, `NoPrimal` tests to EnzymeCore
mofeing Oct 30, 2024
053f627
Merge branch 'main' into needs_primal
wsmoses Oct 30, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ EnzymeStaticArraysExt = "StaticArrays"
BFloat16s = "0.2, 0.3, 0.4, 0.5"
CEnum = "0.4, 0.5"
ChainRulesCore = "1"
EnzymeCore = "0.8.4"
Copy link
Member

Choose a reason for hiding this comment

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

can the compat be set to either 0.8.4 or 0.8.5 ?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

do you mean to support both versions? we can put the ~ specifier in the version

Copy link
Member

Choose a reason for hiding this comment

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

yup so "0.8.4, 0.8.5"

EnzymeCore = "0.8.4, 0.8.5"
Enzyme_jll = "0.0.158"
GPUCompiler = "0.21, 0.22, 0.23, 0.24, 0.25, 0.26, 0.27, 1"
LLVM = "6.1, 7, 8, 9"
Expand Down
2 changes: 1 addition & 1 deletion lib/EnzymeCore/Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "EnzymeCore"
uuid = "f151be2c-9106-41f4-ab19-57ee4f262869"
authors = ["William Moses <wmoses@mit.edu>", "Valentin Churavy <vchuravy@mit.edu>"]
version = "0.8.4"
version = "0.8.5"

[compat]
Adapt = "3, 4"
Expand Down
14 changes: 14 additions & 0 deletions lib/EnzymeCore/src/EnzymeCore.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ export MixedDuplicated, BatchMixedDuplicated
export DefaultABI, FFIABI, InlineABI, NonGenABI
export BatchDuplicatedFunc
export within_autodiff
export needs_primal

function batch_size end

Expand Down Expand Up @@ -351,6 +352,15 @@ Return a new mode which excludes the primal value.
"""
@inline NoPrimal(::ReverseMode{ReturnPrimal,RuntimeActivity,ABI,Holomorphic,ErrIfFuncWritten}) where {ReturnPrimal,RuntimeActivity,ABI,Holomorphic,ErrIfFuncWritten} = ReverseMode{false,RuntimeActivity,ABI,Holomorphic,ErrIfFuncWritten}()

"""
needs_primal(::Mode)
needs_primal(::Type{Mode})

Returns `true` if the mode needs the primal value, otherwise `false`.
"""
@inline needs_primal(::ReverseMode{ReturnPrimal}) where {ReturnPrimal} = ReturnPrimal
@inline needs_primal(::Type{<:ReverseMode{ReturnPrimal}}) where {ReturnPrimal} = ReturnPrimal

"""
struct ReverseModeSplit{
ReturnPrimal,
Expand Down Expand Up @@ -424,6 +434,8 @@ Return a new instance of [`ReverseModeSplit`](@ref) mode where `Width` is set to
@inline WithPrimal(::ReverseModeSplit{ReturnPrimal,ReturnShadow,RuntimeActivity,Width,ModifiedBetween,ABI, Holomorphic, ErrIfFuncWritten, ShadowInit}) where {ReturnPrimal,ReturnShadow,RuntimeActivity,Width,ModifiedBetween,ABI, Holomorphic, ErrIfFuncWritten, ShadowInit} = ReverseModeSplit{true,ReturnShadow,RuntimeActivity,Width,ModifiedBetween,ABI, Holomorphic, ErrIfFuncWritten, ShadowInit}()
@inline NoPrimal(::ReverseModeSplit{ReturnPrimal,ReturnShadow,RuntimeActivity,Width,ModifiedBetween,ABI, Holomorphic, ErrIfFuncWritten, ShadowInit}) where {ReturnPrimal,ReturnShadow,RuntimeActivity,Width,ModifiedBetween,ABI, Holomorphic, ErrIfFuncWritten, ShadowInit} = ReverseModeSplit{false,ReturnShadow,RuntimeActivity,Width,ModifiedBetween,ABI, Holomorphic, ErrIfFuncWritten, ShadowInit}()

@inline needs_primal(::ReverseModeSplit{ReturnPrimal}) where {ReturnPrimal} = ReturnPrimal
@inline needs_primal(::Type{<:ReverseModeSplit{ReturnPrimal}}) where {ReturnPrimal} = ReturnPrimal

"""
struct ForwardMode{
Expand Down Expand Up @@ -480,6 +492,8 @@ const ForwardWithPrimal = ForwardMode{true, DefaultABI, false, false}()
@inline WithPrimal(::ForwardMode{ReturnPrimal,ABI,ErrIfFuncWritten,RuntimeActivity}) where {ReturnPrimal,ABI,ErrIfFuncWritten,RuntimeActivity} = ForwardMode{true,ABI,ErrIfFuncWritten,RuntimeActivity}()
@inline NoPrimal(::ForwardMode{ReturnPrimal,ABI,ErrIfFuncWritten,RuntimeActivity}) where {ReturnPrimal,ABI,ErrIfFuncWritten,RuntimeActivity} = ForwardMode{false,ABI,ErrIfFuncWritten,RuntimeActivity}()

@inline needs_primal(::ForwardMode{ReturnPrimal}) where {ReturnPrimal} = ReturnPrimal
@inline needs_primal(::Type{<:ForwardMode{ReturnPrimal}}) where {ReturnPrimal} = ReturnPrimal

function autodiff end
function autodiff_deferred end
Expand Down
28 changes: 28 additions & 0 deletions lib/EnzymeCore/test/runtests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,34 @@ using Test
using EnzymeCore

@testset verbose = true "EnzymeCore" begin
@testset "WithPrimal" begin
@test WithPrimal(Reverse) === ReverseWithPrimal
@test NoPrimal(Reverse) === Reverse
@test WithPrimal(ReverseWithPrimal) === ReverseWithPrimal
@test NoPrimal(ReverseWithPrimal) === Reverse

@test WithPrimal(set_runtime_activity(Reverse)) === set_runtime_activity(ReverseWithPrimal)

@test WithPrimal(Forward) === ForwardWithPrimal
@test NoPrimal(Forward) === Forward
@test WithPrimal(ForwardWithPrimal) === ForwardWithPrimal
@test NoPrimal(ForwardWithPrimal) === Forward

@test WithPrimal(ReverseSplitNoPrimal) === ReverseSplitWithPrimal
@test NoPrimal(ReverseSplitNoPrimal) === ReverseSplitNoPrimal
@test WithPrimal(ReverseSplitWithPrimal) === ReverseSplitWithPrimal
@test NoPrimal(ReverseSplitWithPrimal) === ReverseSplitNoPrimal
end

@testset "needs_primal" begin
@test needs_primal(Reverse) === false
@test needs_primal(ReverseWithPrimal) === true
@test needs_primal(Forward) === false
@test needs_primal(ForwardWithPrimal) === true
@test needs_primal(ReverseSplitNoPrimal) === false
@test needs_primal(ReverseSplitWithPrimal) === true
end

@testset "Miscellaneous" begin
include("misc.jl")
end
Expand Down
19 changes: 0 additions & 19 deletions test/runtests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3615,25 +3615,6 @@ end
@test res[2][6] ≈ 6.0
end

@testset "WithPrimal" begin
@test WithPrimal(Reverse) === ReverseWithPrimal
@test NoPrimal(Reverse) === Reverse
@test WithPrimal(ReverseWithPrimal) === ReverseWithPrimal
@test NoPrimal(ReverseWithPrimal) === Reverse

@test WithPrimal(set_runtime_activity(Reverse)) === set_runtime_activity(ReverseWithPrimal)

@test WithPrimal(Forward) === ForwardWithPrimal
@test NoPrimal(Forward) === Forward
@test WithPrimal(ForwardWithPrimal) === ForwardWithPrimal
@test NoPrimal(ForwardWithPrimal) === Forward

@test WithPrimal(ReverseSplitNoPrimal) === ReverseSplitWithPrimal
@test NoPrimal(ReverseSplitNoPrimal) === ReverseSplitNoPrimal
@test WithPrimal(ReverseSplitWithPrimal) === ReverseSplitWithPrimal
@test NoPrimal(ReverseSplitWithPrimal) === ReverseSplitNoPrimal
end

# TEST EXTENSIONS
using SpecialFunctions
@testset "SpecialFunctions ext" begin
Expand Down
Loading