Fixes related to FunctionWrappers and wrapped number types

src/analyses/activity.jl

@@ -97,10 +97,10 @@ Base.@nospecializeinfer @inline function is_mutable_array(@nospecialize(T::Type)
     if T isa DataType
         if hasproperty(T, :name) && hasproperty(T.name, :module)
             mod = T.name.module
-            if string(mod) == "Reactant" && (T.name.name == :ConcretePJRTArray || T.name.name == :ConcreteIFRTArray || T.name.name == :TracedRArray)
+            if nameof(mod) === :Reactant && (T.name.name == :ConcretePJRTArray || T.name.name == :ConcreteIFRTArray || T.name.name == :TracedRArray)
                 return true
             end
-            if string(mod) == "CUDA" && (T.name.name == :CuRefValue || T.name.name == :CuPtr)
+            if nameof(mod) === :CUDA && (T.name.name == :CuRefValue || T.name.name == :CuPtr)
                 return true
             end
         end
@@ -109,15 +109,17 @@ Base.@nospecializeinfer @inline function is_mutable_array(@nospecialize(T::Type)
 end
 
 Base.@nospecializeinfer @inline function is_wrapped_number(@nospecialize(T::Type))
-    if T isa UnionAll
-        return is_wrapped_number(T.body)
-    end
     while T isa UnionAll
         T = T.body
     end
     if T isa DataType && hasproperty(T, :name) && hasproperty(T.name, :module)
         mod = T.name.module
-        if string(mod) == "Reactant" && (T.name.name == :ConcretePJRTNumber || T.name.name == :ConcreteIFRTNumber || T.name.name == :TracedRNumber)
+        if nameof(mod) === :Reactant && T.name.name in (
+                :ConcretePJRTNumber, :ConcreteIFRTNumber, :TracedRNumber,
+                :TracedRInteger, :TracedRFloat, :TracedRComplex,
+                :ConcretePJRTInteger, :ConcretePJRTFloat, :ConcretePJRTComplex,
+                :ConcreteIFRTInteger, :ConcreteIFRTFloat, :ConcreteIFRTComplex,
+            )
             return true
         end
     end
@@ -162,14 +164,9 @@ Base.@nospecializeinfer @inline function active_reg_inner(
         )
     end
 
-    if T <: BigFloat
-        return DupState
-    end
-
-    if T <: AbstractFloat
-        return ActiveState
-    end
-
+    # Wrapped numbers are mutable references to their underlying value; they
+    # may subtype AbstractFloat (or Integer) and must be classified before the
+    # raw float/integer fast paths.
     if is_wrapped_number(T)
         if justActive
             return AnyState
@@ -193,6 +190,14 @@ Base.@nospecializeinfer @inline function active_reg_inner(
         end
     end
 
+    if T <: BigFloat
+        return DupState
+    end
+
+    if T <: AbstractFloat
+        return ActiveState
+    end
+
     if is_mutable_array(T)
         if justActive
             return AnyState

src/compiler/interpreter.jl

@@ -1308,21 +1308,26 @@ function abstract_call_known(
     let ft = Core.Compiler.widenconst(argtypes[1])
         if ft isa DataType && ft.name.name == Symbol("#do_ccall") && string(ft.name.module) == "FunctionWrappers"
             ft2 = Core.Compiler.widenconst(argtypes[2])
-            Ret = ft2.parameters[1]
-            ArgsT = ft2.parameters[2]
-            arginfo2 = ArgInfo(
-                fargs isa Nothing ? nothing :
-                [:(FuncWrapperRewriter{Ret, ArgsT}), fargs[2:end]...],
-                [Core.Const(FuncWrapperRewriter{Ret, ArgsT}()), argtypes[2:end]...],
-            )
-            return Base.@invoke abstract_call_known(
-                interp::AbstractInterpreter,
-                FuncWrapperRewriter{Ret, ArgsT}(),
-                arginfo2::ArgInfo,
-                si::StmtInfo,
-                sv::AbsIntState,
-                max_methods::Int,
-            )
+            # The FunctionWrapper type may be only partially known during
+            # abstract interpretation (a UnionAll), rewrite only fully
+            # determined calls:
+            if ft2 isa DataType && length(ft2.parameters) >= 2
+                Ret = ft2.parameters[1]
+                ArgsT = ft2.parameters[2]
+                arginfo2 = ArgInfo(
+                    fargs isa Nothing ? nothing :
+                    [:(FuncWrapperRewriter{Ret, ArgsT}), fargs[2:end]...],
+                    [Core.Const(FuncWrapperRewriter{Ret, ArgsT}()), argtypes[2:end]...],
+                )
+                return Base.@invoke abstract_call_known(
+                    interp::AbstractInterpreter,
+                    FuncWrapperRewriter{Ret, ArgsT}(),
+                    arginfo2::ArgInfo,
+                    si::StmtInfo,
+                    sv::AbsIntState,
+                    max_methods::Int,
+                )
+            end
         end
     end
 

test/core/activity.jl

@@ -43,3 +43,59 @@ end
 
     @test Enzyme.Compiler.active_reg(Tuple, Base.get_world_counter(), AbstractIsMixed=true, justActive=true) == Enzyme.Compiler.MixedState
 end
+
+# Reactant number and array wrappers are recognized by module/type name in
+# `is_wrapped_number`/`is_mutable_array`; mock the names so classification is
+# testable without a Reactant dependency.
+module Reactant
+    mutable struct TracedRNumber{T} <: Number
+        v::T
+    end
+    mutable struct TracedRInteger{T} <: Integer
+        v::T
+    end
+    mutable struct TracedRFloat{T} <: AbstractFloat
+        v::T
+    end
+    mutable struct TracedRComplex{T} <: Number
+        v::T
+    end
+    mutable struct ConcretePJRTFloat{T, D} <: AbstractFloat
+        v::T
+    end
+    mutable struct ConcreteIFRTInteger{T} <: Integer
+        v::T
+    end
+    mutable struct TracedRArray{T, N} <: AbstractArray{T, N} end
+end
+
+module NotReactant
+    mutable struct TracedRFloat{T} <: AbstractFloat
+        v::T
+    end
+end
+
+@testset "Wrapped number activity" begin
+    @test Enzyme.Compiler.is_wrapped_number(Reactant.TracedRFloat{Float64})
+    @test Enzyme.Compiler.is_wrapped_number(Reactant.ConcretePJRTFloat{Float64, 1})
+    @test !Enzyme.Compiler.is_wrapped_number(NotReactant.TracedRFloat{Float64})
+    @test !Enzyme.Compiler.is_wrapped_number(Float64)
+
+    # wrapped numbers subtyping AbstractFloat/Integer must classify as mutable
+    # wrappers of their unwrapped type, not hit the raw float/integer fast paths
+    @test Enzyme.Compiler.active_reg(Reactant.TracedRFloat{Float64}, Base.get_world_counter()) == Enzyme.Compiler.DupState
+    @test Enzyme.Compiler.active_reg(Reactant.TracedRFloat{Float64}, Base.get_world_counter(), justActive = true) == Enzyme.Compiler.AnyState
+    @test Enzyme.Compiler.active_reg(Reactant.TracedRInteger{Int}, Base.get_world_counter()) == Enzyme.Compiler.AnyState
+    @test Enzyme.Compiler.active_reg(Reactant.TracedRComplex{ComplexF64}, Base.get_world_counter()) == Enzyme.Compiler.DupState
+    @test Enzyme.Compiler.active_reg(Reactant.TracedRNumber{Float64}, Base.get_world_counter()) == Enzyme.Compiler.DupState
+    @test Enzyme.Compiler.active_reg(Reactant.ConcretePJRTFloat{Float64, 1}, Base.get_world_counter()) == Enzyme.Compiler.DupState
+    @test Enzyme.Compiler.active_reg(Reactant.ConcreteIFRTInteger{Int}, Base.get_world_counter()) == Enzyme.Compiler.AnyState
+    @test Enzyme.Compiler.active_reg(Tuple{Reactant.TracedRFloat{Float64}, Float64}, Base.get_world_counter()) == Enzyme.Compiler.MixedState
+
+    # same-named types outside the Reactant module use the plain rules
+    @test Enzyme.Compiler.active_reg(NotReactant.TracedRFloat{Float64}, Base.get_world_counter()) == Enzyme.Compiler.ActiveState
+
+    @test Enzyme.Compiler.is_mutable_array(Reactant.TracedRArray{Float64, 1})
+    @test Enzyme.Compiler.active_reg(Reactant.TracedRArray{Float64, 1}, Base.get_world_counter()) == Enzyme.Compiler.DupState
+    @test Enzyme.Compiler.active_reg(Reactant.TracedRArray{Int, 1}, Base.get_world_counter()) == Enzyme.Compiler.AnyState
+end

test/ext/functionwrappers.jl

@@ -18,3 +18,25 @@ end
     res = Enzyme.gradient(Enzyme.set_runtime_activity(Enzyme.Reverse), wrapper_loss, [2.0])
     @test res[1] ≈ [-1.0]
 end
+
+# The FunctionWrapper type visible to abstract interpretation may be a
+# UnionAll when the wrapper comes from a container with a partially known
+# element type. Compilation must not fail on such calls, even if they are
+# never reached at runtime:
+function partially_known_wrapper(x::Float64, take_wrapper::Bool, fws::Vector{FunctionWrapper{Float64}})
+    if take_wrapper
+        return fws[1](x)::Float64
+    else
+        return sin(x)
+    end
+end
+
+@testset "FunctionWrappers with partially known wrapper type" begin
+    fws = Vector{FunctionWrapper{Float64}}(undef, 1)
+    fws[1] = FunctionWrapper{Float64,Tuple{Float64}}(cos)
+    res = Enzyme.autodiff(
+        Enzyme.set_runtime_activity(Enzyme.Reverse), partially_known_wrapper,
+        Enzyme.Active, Enzyme.Active(0.5), Enzyme.Const(false), Enzyme.Const(fws)
+    )
+    @test res[1][1] ≈ cos(0.5)
+end