Utilities for splitting and unsplitting mode objects

lib/EnzymeCore/src/EnzymeCore.jl

@@ -384,7 +384,7 @@ Subtype of [`Mode`](@ref) for split reverse mode differentiation, to use in [`au
     - [`set_abi`](@ref)
     - [`ReverseSplitModified`](@ref), [`ReverseSplitWidth`](@ref)
 """
-struct ReverseModeSplit{ReturnPrimal,ReturnShadow,Width,RuntimeActivity,ModifiedBetween,ABI,Holomorphic,ErrIfFuncWritten,ShadowInit} <: Mode{ABI, ErrIfFuncWritten,RuntimeActivity} end
+struct ReverseModeSplit{ReturnPrimal,ReturnShadow,RuntimeActivity,Width,ModifiedBetween,ABI,Holomorphic,ErrIfFuncWritten,ShadowInit} <: Mode{ABI, ErrIfFuncWritten,RuntimeActivity} end
 
 """
     const ReverseSplitNoPrimal
@@ -578,4 +578,99 @@ Return a new mode with its [`ABI`](@ref) set to the chosen type.
 """
 function set_abi end
 
+"""
+    split_mode(
+        ::ReverseMode, [::Val{ReturnShadow}, ::Val{Width}, ::Val{ModifiedBetween}, ::Val{ShadowInit}]
+    )
+
+Turn a [`ReverseMode`](@ref) object into a [`ReverseModeSplit`](@ref) object while preserving as many of the settings as possible.
+The rest of the settings can be configured with optional positional arguments of `Val` type.
+
+This function acts as the identity on a [`ReverseModeSplit`](@ref).
+
+See also [`combined_mode`](@ref).
+"""
+function split_mode(
+    ::ReverseMode{
+        ReturnPrimal,
+        RuntimeActivity,
+        ABI,
+        Holomorphic,
+        ErrIfFuncWritten
+    },
+    ::Val{ReturnShadow}=Val(true),
+    ::Val{Width}=Val(0),
+    ::Val{ModifiedBetween}=Val(true),
+    ::Val{ShadowInit}=Val(false),
+) where {
+    ReturnPrimal,
+    ReturnShadow,
+    RuntimeActivity,
+    Width,
+    ModifiedBetween,
+    ABI,
+    Holomorphic,
+    ErrIfFuncWritten,
+    ShadowInit
+}
+    mode_split = ReverseModeSplit{
+        ReturnPrimal,
+        ReturnShadow,
+        RuntimeActivity,
+        Width,
+        ModifiedBetween,
+        ABI,
+        Holomorphic,
+        ErrIfFuncWritten,
+        ShadowInit
+    }()
+    return mode_split
+end
+
+split_mode(mode::ReverseModeSplit, args...) = mode
+
+"""
+    combined_mode(::ReverseMode)
+
+Turn a [`ReverseModeSplit`](@ref) object into a [`ReverseMode`](@ref) object while preserving as many of the settings as possible.
+
+This function acts as the identity on a [`ReverseMode`](@ref).
+
+See also [`split_mode`](@ref).
+"""
+function combined_mode(
+    ::ReverseModeSplit{
+        ReturnPrimal,
+        ReturnShadow,
+        RuntimeActivity,
+        Width,
+        ModifiedBetween,
+        ABI,
+        Holomorphic,
+        ErrIfFuncWritten,
+        ShadowInit
+    }
+) where {
+    ReturnPrimal,
+    ReturnShadow,
+    RuntimeActivity,
+    Width,
+    ModifiedBetween,
+    ABI,
+    Holomorphic,
+    ErrIfFuncWritten,
+    ShadowInit
+}
+    mode_unsplit = ReverseMode{
+        ReturnPrimal,
+        RuntimeActivity,
+        ABI,
+        Holomorphic,
+        ErrIfFuncWritten
+    }()
+    return mode_unsplit
+end
+
+combined_mode(mode::ReverseMode) = mode
+
 end # module EnzymeCore

lib/EnzymeCore/test/mode_modification.jl

@@ -0,0 +1,23 @@
+using EnzymeCore
+using EnzymeCore: InlineABI, ReverseModeSplit, split_mode, combined_mode, set_runtime_activity, set_err_if_func_written, set_abi
+using Test
+
+@testset "Split / unsplit mode" begin
+    @test split_mode(Reverse) == ReverseSplitNoPrimal
+    @test split_mode(ReverseWithPrimal) == ReverseSplitWithPrimal
+    @test split_mode(ReverseSplitNoPrimal) == ReverseSplitNoPrimal
+    @test split_mode(ReverseSplitWithPrimal) == ReverseSplitWithPrimal
+
+    @test split_mode(set_runtime_activity(Reverse)) == set_runtime_activity(ReverseSplitNoPrimal)
+    @test split_mode(set_err_if_func_written(Reverse)) == set_err_if_func_written(ReverseSplitNoPrimal)
+
+    @test split_mode(Reverse, Val(:ReturnShadow), Val(:Width), Val(:ModifiedBetween), Val(:ShadowInit)) == ReverseModeSplit{false,:ReturnShadow,false,:Width,:ModifiedBetween,EnzymeCore.DefaultABI,false,false,:ShadowInit}()
+
+    @test combined_mode(Reverse) == Reverse
+    @test combined_mode(ReverseWithPrimal) == ReverseWithPrimal
+    @test combined_mode(ReverseSplitNoPrimal) == Reverse
+    @test combined_mode(ReverseSplitWithPrimal) == ReverseWithPrimal
+
+    @test combined_mode(set_runtime_activity(ReverseSplitNoPrimal)) == set_runtime_activity(Reverse)
+    @test combined_mode(set_err_if_func_written(ReverseSplitNoPrimal)) == set_err_if_func_written(Reverse)
+end

lib/EnzymeCore/test/runtests.jl

@@ -25,4 +25,6 @@ y = @view data[3:end]
 @test_throws ErrorException Duplicated(d, y)
 
 @test_throws ErrorException Active(data)
-@test_throws ErrorException Active(d)
+@test_throws ErrorException Active(d)
+
+include("mode_modification.jl")