Add ABI setter (#1709)

EnzymeAD · Aug 7, 2024 · d19307d · d19307d
1 parent ece399d
commit d19307d
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 8 deletions.
diff --git a/lib/EnzymeCore/src/EnzymeCore.jl b/lib/EnzymeCore/src/EnzymeCore.jl
@@ -228,6 +228,8 @@ const ReverseHolomorphicWithPrimal = ReverseMode{true,DefaultABI, true, false}()
 @inline set_err_if_func_written(::ReverseMode{ReturnPrimal,ABI,Holomorphic,ErrIfFuncWritten}) where {ReturnPrimal,ABI,Holomorphic,ErrIfFuncWritten} = ReverseMode{ReturnPrimal,ABI,Holomorphic,true}()
 @inline clear_err_if_func_written(::ReverseMode{ReturnPrimal,ABI,Holomorphic,ErrIfFuncWritten}) where {ReturnPrimal,ABI,Holomorphic,ErrIfFuncWritten} = ReverseMode{ReturnPrimal,ABI,Holomorphic,false}()
 
+@inline set_abi(::ReverseMode{ReturnPrimal,OldABI,Holomorphic,ErrIfFuncWritten}, ::Type{NewABI}) where {ReturnPrimal,OldABI,Holomorphic,ErrIfFuncWritten,NewABI<:ABI} = ReverseMode{ReturnPrimal,NewABI,Holomorphic,ErrIfFuncWritten}()
+
 """
     struct ReverseModeSplit{ReturnPrimal,ReturnShadow,Width,ModifiedBetween,ABI} <: Mode{ABI}
 
@@ -255,6 +257,8 @@ const Forward = ForwardMode{DefaultABI, false}()
 @inline set_err_if_func_written(::ForwardMode{ABI,ErrIfFuncWritten}) where {ABI,ErrIfFuncWritten} = ForwardMode{ABI,true}()
 @inline clear_err_if_func_written(::ForwardMode{ABI,ErrIfFuncWritten}) where {ABI,ErrIfFuncWritten} = ForwardMode{ABI,false}()
 
+@inline set_abi(::ForwardMode{OldABI,ErrIfFuncWritten}, ::Type{NewABI}) where {OldABI,ErrIfFuncWritten,NewABI<:ABI} = ForwardMode{NewABI,ErrIfFuncWritten}()
+
 function autodiff end
 function autodiff_deferred end
 function autodiff_thunk end

diff --git a/src/Enzyme.jl b/src/Enzyme.jl
@@ -5,8 +5,8 @@ import EnzymeCore
 import EnzymeCore: Forward, Reverse, ReverseWithPrimal, ReverseSplitNoPrimal, ReverseSplitWithPrimal, ReverseSplitModified, ReverseSplitWidth, ReverseMode, ForwardMode, ReverseHolomorphic, ReverseHolomorphicWithPrimal
 export Forward, Reverse, ReverseWithPrimal, ReverseSplitNoPrimal, ReverseSplitWithPrimal, ReverseSplitModified, ReverseSplitWidth, ReverseMode, ForwardMode, ReverseHolomorphic, ReverseHolomorphicWithPrimal
 
-import EnzymeCore: Annotation, Const, Active, Duplicated, DuplicatedNoNeed, BatchDuplicated, BatchDuplicatedNoNeed, ABI, DefaultABI, FFIABI, InlineABI, NonGenABI, set_err_if_func_written, clear_err_if_func_written
-export Annotation, Const, Active, Duplicated, DuplicatedNoNeed, BatchDuplicated, BatchDuplicatedNoNeed, DefaultABI, FFIABI, InlineABI, NonGenABI, set_err_if_func_written, clear_err_if_func_written
+import EnzymeCore: Annotation, Const, Active, Duplicated, DuplicatedNoNeed, BatchDuplicated, BatchDuplicatedNoNeed, ABI, DefaultABI, FFIABI, InlineABI, NonGenABI, set_err_if_func_written, clear_err_if_func_written, set_abi
+export Annotation, Const, Active, Duplicated, DuplicatedNoNeed, BatchDuplicated, BatchDuplicatedNoNeed, DefaultABI, FFIABI, InlineABI, NonGenABI, set_err_if_func_written, clear_err_if_func_written, set_abi
 
 import EnzymeCore: BatchDuplicatedFunc
 export BatchDuplicatedFunc

diff --git a/test/abi.jl b/test/abi.jl
@@ -9,11 +9,11 @@ using Test
     res = autodiff(Reverse, f, Const, Const(nothing))
     @test res === ((nothing,),)
 
-    res = autodiff(ReverseMode{false,NonGenABI, false, false}(), f, Const, Const(nothing))
+    res = autodiff(Enzyme.set_abi(Reverse, NonGenABI), f, Const, Const(nothing))
     @test res === ((nothing,),)
 
     @test () === autodiff(Forward, f, Const, Const(nothing))
-    @test () === autodiff(ForwardMode{NonGenABI, false}(), f, Const, Const(nothing))
+    @test () === autodiff(Enzyme.set_abi(Forward, NonGenABI), f, Const, Const(nothing))
 
     res = autodiff(Reverse, f, Const(nothing))
     @test res === ((nothing,),)
@@ -22,11 +22,11 @@ using Test
 
     res = autodiff_deferred(Reverse, f, Const(nothing))
     @test res === ((nothing,),)
-    res = autodiff_deferred(ReverseMode{false,NonGenABI, false, false}(), f, Const, Const(nothing))
+    res = autodiff_deferred(Enzyme.set_abi(Reverse, NonGenABI), f, Const, Const(nothing))
     @test res === ((nothing,),)
 
     @test () === autodiff_deferred(Forward, f, Const(nothing))
-    @test () === autodiff_deferred(ForwardMode{NonGenABI, false}(), f, Const, Const(nothing))
+    @test () === autodiff_deferred(Enzyme.set_abi(Forward, NonGenABI), f, Const, Const(nothing))
 
     # ConstType -> Type{Int}
     res = autodiff(Reverse, f, Const, Const(Int))
@@ -65,15 +65,15 @@ using Test
     _, res0 = autodiff(Reverse, unused, Active, Const(nothing), Active(2.0))[1]
     @test res0 ≈ 1.0
 
-    _, res0 = autodiff(ReverseMode{false, NonGenABI, false, false}(), unused, Active, Const(nothing), Active(2.0))[1]
+    _, res0 = autodiff(Enzyme.set_abi(Reverse, NonGenABI), unused, Active, Const(nothing), Active(2.0))[1]
     @test res0 ≈ 1.0
 
     res0, = autodiff(Forward, unused, DuplicatedNoNeed, Const(nothing), Duplicated(2.0, 1.0))
     @test res0 ≈ 1.0
     res0, = autodiff(Forward, unused, DuplicatedNoNeed, Const(nothing), DuplicatedNoNeed(2.0, 1.0))
     @test res0 ≈ 1.0
 
-    res0, = autodiff(ForwardMode{NonGenABI, false}(), unused, DuplicatedNoNeed, Const(nothing), Duplicated(2.0, 1.0))
+    res0, = autodiff(Enzyme.set_abi(Forward, NonGenABI), unused, DuplicatedNoNeed, Const(nothing), Duplicated(2.0, 1.0))
     @test res0 ≈ 1.0
 
     _, res0 = autodiff(Reverse, unused, Const(nothing), Active(2.0))[1]