Simplify: remove IsInactive, always use *_nongen

Author: danielwe

lib/EnzymeCore/src/EnzymeCore.jl

@@ -506,60 +506,35 @@ function autodiff_thunk end
 function autodiff_deferred_thunk end
 
 """
-    make_zero(
-        prev::T, ::Val{copy_if_inactive}=Val(false), ::Val{runtime_inactive}=Val(false)
-    )::T
-    make_zero(
-        ::Type{T},
-        seen::IdDict,
-        prev::T,
-        ::Val{copy_if_inactive}=Val(false),
-        ::Val{runtime_inactive}=Val(false),
-    )::T
+    make_zero(prev::T, ::Val{copy_if_inactive}=Val(false))::T
+    make_zero(::Type{T}, seen::IdDict, prev::T, ::Val{copy_if_inactive}=Val(false))::T
 
 Recursively make a copy of the value `prev::T` in which all differentiable values are
 zeroed. The argument `copy_if_inactive` specifies what to do if the type `T` or any
 of its constituent parts is guaranteed to be inactive (non-differentiable): reuse `prev`s
 instance (the default) or make a copy.
 
-The argument `runtime_inactive` specifies whether each constituent type is checked for being
-guaranteed inactive at runtime for every call to `make_zero`, or if this can be checked once
-at compile-time and reused across multiple calls to `make_zero` and related functions (the
-default). Runtime checks are necessary to pick up recently added methods to
-`EnzymeRules.inactive_type`, but may incur a significant performance penalty and is usually
-not needed unless `EnzymeRules.inactive_type` is extended interactively for types that have
-previously been passed to `make_zero` or related functions.
-
 Extending this method for custom types is rarely needed. If you implement a new type, such
-as a GPU array type, for which `make_zero` should directly invoke `zero` rather than
-iterate/broadcast when the eltype is scalar, it is sufficient to implement `Base.zero` and
-make sure your type subtypes `DenseArray`. (If subtyping `DenseArray` is not appropriate,
-extend [`EnzymeCore.isvectortype`](@ref) directly instead.)
+as a GPU array type, on which `make_zero` should directly invoke `zero` when the eltype is
+scalar, it is sufficient to implement `Base.zero` and make sure your type subtypes
+`DenseArray`. (If subtyping `DenseArray` is not appropriate, extend
+[`EnzymeCore.isvectortype`](@ref) instead.)
 """
 function make_zero end
 
 """
-    make_zero!(val::T, [seen::IdDict], ::Val{runtime_inactive}=Val(false))::Nothing
+    make_zero!(val::T, [seen::IdDict])::Nothing
 
 Recursively set a variable's differentiable values to zero. Only applicable for types `T`
 that are mutable or hold all differentiable values in mutable storage (e.g.,
 `Tuple{Vector{Float64}}` qualifies but `Tuple{Float64}` does not). The recursion skips over
 parts of `val` that are guaranteed to be inactive.
 
-The argument `runtime_inactive` specifies whether each constituent type is checked for being
-guaranteed inactive at runtime for every call to `make_zero!`, or if this can be checked once
-at compile-time and reused across multiple calls to `make_zero!` and related functions (the
-default). Runtime checks are necessary to pick up recently added methods to
-`EnzymeRules.inactive_type`, but may incur a significant performance penalty and is usually
-not needed unless `EnzymeRules.inactive_type` is extended interactively for types that have
-previously been passed to `make_zero!` or related functions.
-
 Extending this method for custom types is rarely needed. If you implement a new mutable
-type, such as a GPU array type, for which `make_zero!` should directly invoke
-`fill!(x, false)` rather than iterate/broadcast when the eltype is scalar, it is sufficient
-to implement `Base.zero`, `Base.fill!`, and make sure your type subtypes `DenseArray`. (If
-subtyping `DenseArray` is not appropriate, extend [`EnzymeCore.isvectortype`](@ref) directly
-instead.)
+type, such as a GPU array type, on which `make_zero!` should directly invoke
+`fill!(x, false)` when the eltype is scalar, it is sufficient to implement `Base.zero`,
+`Base.fill!`, and make sure your type subtypes `DenseArray`. (If subtyping `DenseArray` is
+not appropriate, extend [`EnzymeCore.isvectortype`](@ref) instead.)
 """
 function make_zero! end
 
@@ -569,7 +544,7 @@ function make_zero! end
 Trait defining types whose values should be considered leaf nodes when [`make_zero`](@ref)
 and [`make_zero!`](@ref) recurse through an object.
 
-By default, `isvectortype(T) == true` for `T` such that `isscalartype(T) == true` or
+By default, `isvectortype(T) == true` when `isscalartype(T) == true` or when
 `T <: DenseArray{U}` where `U` is a bitstype and `isscalartype(U) == true`.
 
 A new vector type, such as a GPU array type, should normally subtype `DenseArray` and
@@ -607,7 +582,7 @@ in-place mutation through any Julia API; `isscalartype(BigFloat) == true` ensure
 `make_zero!` will not try to mutate `BigFloat` values.[^BigFloat]
 
 By default, `isscalartype(T) == true` and `isscalartype(Complex{T}) == true` for concrete
-`T <: AbstractFloat`.
+types `T <: AbstractFloat`.
 
 A hypothetical new real number type with Enzyme support should usually subtype
 `AbstractFloat` and inherit the `isscalartype` trait that way. If this is not appropriate,

src/analyses/activity.jl

@@ -414,7 +414,7 @@ Base.@assume_effects :removable :foldable :nothrow @inline function guaranteed_c
     return res
 end
 
-Base.@assume_effects :removable :foldable :nothrow @inline function guaranteed_const_nongen(::Type{T}, world)::Bool where {T}
+Base.@assume_effects :removable :foldable :nothrow @inline function guaranteed_const_nongen(::Type{T}, world=nothing)::Bool where {T}
     rt = active_reg_inner(T, (), world)
     res = rt == AnyState
     return res
@@ -427,7 +427,7 @@ Base.@assume_effects :removable :foldable :nothrow @inline function guaranteed_n
     return rt == Enzyme.Compiler.AnyState || rt == Enzyme.Compiler.DupState
 end
 
-Base.@assume_effects :removable :foldable :nothrow @inline function guaranteed_nonactive_nongen(::Type{T}, world)::Bool where {T}
+Base.@assume_effects :removable :foldable :nothrow @inline function guaranteed_nonactive_nongen(::Type{T}, world=nothing)::Bool where {T}
     rt = Enzyme.Compiler.active_reg_inner(T, (), world)
     return rt == Enzyme.Compiler.AnyState || rt == Enzyme.Compiler.DupState
 end

src/typeutils/recursive_add.jl

@@ -1,7 +1,7 @@
 using .RecursiveMaps: RecursiveMaps, recursive_map, recursive_map!
 
 """
-    recursive_add(x::T, y::T, f=identity, forcelhs=guaranteed_const)
+    recursive_add(x::T, y::T, f=identity, forcelhs=guaranteed_const_nongen)
 
 !!! warning
     Internal function, documented for developer convenience but not covered by semver API
@@ -18,7 +18,9 @@ The optional argument `forcelhs` takes a callable such that if `forcelhs(S) == t
 such that `zi = xi + f(yi)` applies to differentiable values, while `zi = xi` applies to
 non-differentiable values.
 """
-function recursive_add(x::T, y::T, f::F=identity, forcelhs::L=guaranteed_const) where {T,F,L}
+function recursive_add(
+    x::T, y::T, f::F=identity, forcelhs::L=guaranteed_const_nongen
+) where {T,F,L}
     function addf(xi::S, yi::S) where {S}
         @assert EnzymeCore.isvectortype(S)
         return ((xi + f(yi))::S,)
@@ -27,8 +29,7 @@ function recursive_add(x::T, y::T, f::F=identity, forcelhs::L=guaranteed_const)
 end
 
 """
-    accumulate_seen!(f, seen::IdDict, ::Val{runtime_inactive}=Val(false))
-    accumulate_seen!(f, seen::IdDict, isinactivetype::RecursiveMaps.IsInactive)
+    accumulate_seen!(f, seen::IdDict)
 
 !!! warning
     Internal function, documented for developer convenience but not covered by semver API
@@ -42,35 +43,17 @@ vector type instance mappping to another object of the same type and structure.
 returned value.
 
 The recursion stops at instances of types that are themselves cached by `make_zero`
-(`recursive_map`), as these objects should have their own entries in `seen`.
-
-Inactive objects that would be shared/copied rather than zeroed by `make_zero` are skipped.
-If the optional `::Val{runtime_inactive}` argument was passed to `make_zero`, the same value
-should be passed to `accumulate_seen` for consistency. If needed, a custom
-`RecursiveMaps.IsInactive` instance can be provided instead.
+(`recursive_map`), as these objects should have their own entries in `seen`. The recursion
+also stops at inactive objects that not be zeroed by `make_zero`.
 """
-function accumulate_seen!(
-    f::F, seen::IdDict, ::Val{runtime_inactive}=Val(false)
-) where {F,runtime_inactive}
-    accumulate_seen!(f, seen, RecursiveMaps.IsInactive{runtime_inactive}())
-    return nothing
-end
-
-function accumulate_seen!(
-    f::F, seen::IdDict, isinactivetype::RecursiveMaps.IsInactive
-) where {F}
-    isinactivetype_or_seen = RecursiveMaps.IsInactive(
-        isinactivetype, RecursiveMaps.iscachedtype
-    )
+function accumulate_seen!(f::F, seen::IdDict) where {F}
     for (k, v) in seen
-        _accumulate_seen_item!(f, k, v, isinactivetype, isinactivetype_or_seen)
+        _accumulate_seen_item!(f, k, v)
     end
     return nothing
 end
 
-function _accumulate_seen_item!(
-    f::F, k::T, v::T, isinactivetype, isinactivetype_or_seen
-) where {F,T}
+function _accumulate_seen_item!(f::F, k::T, v::T) where {F,T}
     function addf!!(ki::S, vi::S) where {S}
         @assert EnzymeCore.isvectortype(S)
         return ((ki .+ f.(vi))::S,)
@@ -82,10 +65,13 @@ function _accumulate_seen_item!(
         ki .+= f.(vi)
         return (ki::S,)
     end
-    RecursiveMaps.check_nonactive(T, isinactivetype)
-    if !isinactivetype(T)
+    @inline function isinactive_or_cachedtype(::Type{T}) where {T}
+        return guaranteed_const_nongen(T) || RecursiveMaps.iscachedtype(T)
+    end
+    RecursiveMaps.check_nonactive(T)
+    if !guaranteed_const_nongen(T)
         newks = RecursiveMaps.recursive_map_inner(
-            nothing, addf!!, (k,), (k, v), Val(false), isinactivetype_or_seen
+            nothing, addf!!, (k,), (k, v), Val(false), isinactive_or_cachedtype
         )
         @assert only(newks) === k
     end