diff --git a/src/analyses/activity.jl b/src/analyses/activity.jl index 4abe36e514..a85c066c90 100644 --- a/src/analyses/activity.jl +++ b/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 diff --git a/src/compiler/interpreter.jl b/src/compiler/interpreter.jl index a04d384a31..12540dbf4c 100644 --- a/src/compiler/interpreter.jl +++ b/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 diff --git a/test/core/activity.jl b/test/core/activity.jl index cc17b51bac..85e278466a 100644 --- a/test/core/activity.jl +++ b/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 diff --git a/test/ext/functionwrappers.jl b/test/ext/functionwrappers.jl index 05d11f6f8e..d79fe31090 100644 --- a/test/ext/functionwrappers.jl +++ b/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