Merge pull request SciML#3356 from SciML/fix-wrapfun-iip-inference

Fix wrapfun_iip type inference with FunctionWrappersWrappers v1.0+

Author: ChrisRackauckas

Project.toml

@@ -124,7 +124,7 @@ FastClosures = "0.3.2"
 FillArrays = "1.13"
 FiniteDiff = "2.27"
 ForwardDiff = "0.10.38, 1"
-FunctionWrappersWrappers = "0.1.3, 1"
+FunctionWrappersWrappers = "1"
 InteractiveUtils = "1.9"
 JLArrays = "0.2, 0.3"
 LineSearches = "7.4"

lib/DiffEqBase/Project.toml

@@ -87,7 +87,7 @@ FastPower = "1.1"
 FlexUnits = "0.4"
 ForwardDiff = "0.10, 1"
 FunctionWrappers = "1.0"
-FunctionWrappersWrappers = "0.1, 1"
+FunctionWrappersWrappers = "1"
 GTPSA = "1.4"
 GeneralizedGenerated = "0.3"
 InteractiveUtils = "1.9"

lib/DiffEqBase/ext/DiffEqBaseForwardDiffExt.jl

@@ -61,6 +61,27 @@ function wrapfun_oop(ff, inputs::Tuple = ())
     )
 end
 
+# Construct FunctionWrappersWrapper bypassing the convenience constructor.
+# The convenience constructor's `map` doesn't infer when the callable has many
+# type parameters (e.g. ODEFunction with 20+), because the FunctionWrapper
+# constructor in the separately-precompiled FunctionWrappers package can't be
+# traced through for complex types. Using Type{A} dispatch binds the arglist
+# types as type parameters, making FW{Nothing, A}(vff) fully inferrable.
+function _make_fww(
+        @nospecialize(vff),
+        ::Type{A1}, ::Type{A2}, ::Type{A3}, ::Type{A4},
+    ) where {A1, A2, A3, A4}
+    FW = FunctionWrappersWrappers.FunctionWrappers.FunctionWrapper
+    fwt = (
+        FW{Nothing, A1}(vff), FW{Nothing, A2}(vff),
+        FW{Nothing, A3}(vff), FW{Nothing, A4}(vff),
+    )
+    cs = FunctionWrappersWrappers.SingleCacheStorage()
+    return FunctionWrappersWrappers.FunctionWrappersWrapper{
+        typeof(fwt), FunctionWrappersWrappers.AllowNonIsBits, typeof(cs),
+    }(fwt, cs)
+end
+
 function wrapfun_iip(
         ff,
         inputs::Tuple{T1, T2, T3, T4}
@@ -71,16 +92,13 @@ function wrapfun_iip(
     dualT2 = ArrayInterface.promote_eltype(T2, dualT)
     dualT4 = dualgen(promote_type(T, T4))
 
-    iip_arglists = (
+    return _make_fww(
+        Void(ff),
         Tuple{T1, T2, T3, T4},
         Tuple{dualT1, dualT2, T3, T4},
         Tuple{dualT1, T2, T3, dualT4},
-        Tuple{dualT1, dualT2, T3, dualT4},
+        Tuple{dualT1, dualT2, T3, dualT4}
     )
-
-    iip_returnlists = ntuple(x -> Nothing, 4)
-
-    return FunctionWrappersWrappers.FunctionWrappersWrapper(Void(ff), iip_arglists, iip_returnlists)
 end
 
 # 3-arg version: compile FunctionWrapper variants with the specified chunk size.
@@ -103,45 +121,15 @@ function wrapfun_iip(
     dualT1_time = ArrayInterface.promote_eltype(T1, dualT_time)
     dualT4_time = dualgen(promote_type(T, T4))
 
-    iip_arglists = (
-        Tuple{T1, T2, T3, T4},                                  # plain
-        Tuple{dualT1_jac, dualT2_jac, T3, T4},                  # Jacobian (u dual, chunk=CS)
-        Tuple{dualT1_time, T2, T3, dualT4_time},                # time derivative (chunk=1)
-        Tuple{dualT1_jac, dualT2_jac, T3, dualT4_time},         # both
+    return _make_fww(
+        Void(ff),
+        Tuple{T1, T2, T3, T4},
+        Tuple{dualT1_jac, dualT2_jac, T3, T4},
+        Tuple{dualT1_time, T2, T3, dualT4_time},
+        Tuple{dualT1_jac, dualT2_jac, T3, dualT4_time}
     )
-
-    iip_returnlists = ntuple(x -> Nothing, 4)
-
-    return FunctionWrappersWrappers.FunctionWrappersWrapper(Void(ff), iip_arglists, iip_returnlists)
 end
 
-const iip_arglists_default = (
-    Tuple{
-        Vector{Float64}, Vector{Float64}, Vector{Float64},
-        Float64,
-    },
-    Tuple{
-        Vector{Float64}, Vector{Float64},
-        SciMLBase.NullParameters,
-        Float64,
-    },
-    Tuple{Vector{dualT}, Vector{Float64}, Vector{Float64}, dualT},
-    Tuple{Vector{dualT}, Vector{dualT}, Vector{Float64}, dualT},
-    Tuple{Vector{dualT}, Vector{dualT}, Vector{Float64}, Float64},
-    Tuple{
-        Vector{dualT}, Vector{dualT}, SciMLBase.NullParameters,
-        Float64,
-    },
-    Tuple{
-        Vector{dualT}, Vector{Float64},
-        SciMLBase.NullParameters, dualT,
-    },
-)
-const iip_returnlists_default = ntuple(x -> Nothing, length(iip_arglists_default))
-
-function wrapfun_iip(@nospecialize(ff))
-    return FunctionWrappersWrappers.FunctionWrappersWrapper(Void(ff), iip_arglists_default, iip_returnlists_default)
-end
 
 function promote_tspan(u0::AbstractArray{<:ForwardDiff.Dual}, p, tspan, prob, kwargs)
     if (haskey(kwargs, :callback) && has_continuous_callback(kwargs[:callback])) ||

lib/DiffEqBase/test/downstream/inference.jl

@@ -17,7 +17,7 @@ sol = solve(prob, Tsit5(), save_idxs = 1)
 
 prob = ODEProblem{true, SciMLBase.FullSpecialize}(lorenz, u0, tspan)
 
-@inferred SciMLBase.wrapfun_iip(prob.f)
+@inferred SciMLBase.wrapfun_iip(prob.f, (u0, u0, Float64[], tspan[1]))
 @inferred remake(prob, u0 = [1.0; 0.0; 0.0])
 @inferred remake(prob, u0 = Float32[1.0; 0.0; 0.0])
 @test_broken @inferred(solve(prob, Tsit5(), u0 = Float32[1.0; 0.0; 0.0])) ==
@@ -29,7 +29,8 @@ prob = ODEProblem(lorenz, Float32[1.0; 0.0; 0.0], tspan)
     solve(prob, Tsit5(), u0 = [1.0; 0.0; 0.0])
 remake(prob, u0 = [1.0; 0.0; 0.0])
 
-@inferred SciMLBase.wrapfun_iip(prob.f)
+u0_32 = Float32[1.0; 0.0; 0.0]
+@inferred SciMLBase.wrapfun_iip(prob.f, (u0_32, u0_32, Float32[], tspan[1]))
 @test_broken @inferred(
     ODEFunction{
         isinplace(prob), SciMLBase.FunctionWrapperSpecialize,

lib/OrdinaryDiffEqCore/Project.toml

@@ -57,7 +57,7 @@ Accessors = "0.1.36"
 ConcreteStructs = "0.2"
 StaticArraysCore = "1.4.3"
 SciMLStructures = "1.7"
-FunctionWrappersWrappers = "0.1, 1"
+FunctionWrappersWrappers = "1"
 FastBroadcast = "1.3"
 Random = "<0.0.1, 1"
 DiffEqDevTools = "2.44.4"

lib/OrdinaryDiffEqDifferentiation/Project.toml

@@ -12,7 +12,7 @@ DiffEqBase = {path = "../DiffEqBase"}
 [compat]
 Pkg = "1"
 ForwardDiff = "0.10.38, 1"
-FunctionWrappersWrappers = "0.1, 1"
+FunctionWrappersWrappers = "1"
 FastBroadcast = "1.3"
 Random = "<0.0.1, 1"
 DiffEqDevTools = "2.44.4"