Merge pull request #2613 from SebastianM-C/speed

Speed up `replace` and `canonicalize` by improving type stability

Author: ChrisRackauckas

Project.toml

@@ -119,6 +119,7 @@ DelayDiffEq = "bcd4f6db-9728-5f36-b5f7-82caef46ccdb"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
 Ipopt_jll = "9cc047cb-c261-5740-88fc-0cf96f7bdcc7"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 ModelingToolkitStandardLibrary = "16a59e39-deab-5bd0-87e4-056b12336739"
 NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
@@ -137,4 +138,4 @@ Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["AmplNLWriter", "BenchmarkTools", "ControlSystemsBase", "DelayDiffEq", "NonlinearSolve", "ForwardDiff", "Ipopt", "Ipopt_jll", "ModelingToolkitStandardLibrary", "Optimization", "OptimizationOptimJL", "OptimizationMOI", "Random", "ReferenceTests", "SafeTestsets", "StableRNGs", "Statistics", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials", "StochasticDelayDiffEq", "Pkg"]
+test = ["AmplNLWriter", "BenchmarkTools", "ControlSystemsBase", "DelayDiffEq", "NonlinearSolve", "ForwardDiff", "Ipopt", "Ipopt_jll", "ModelingToolkitStandardLibrary", "Optimization", "OptimizationOptimJL", "OptimizationMOI", "Random", "ReferenceTests", "SafeTestsets", "StableRNGs", "Statistics", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials", "StochasticDelayDiffEq", "Pkg", "JET"]

src/systems/parameter_buffer.jl

@@ -155,35 +155,53 @@ function MTKParameters(
 end
 
 function buffer_to_arraypartition(buf)
-    return ArrayPartition(Tuple(eltype(v) <: AbstractArray ? buffer_to_arraypartition(v) :
-                                v for v in buf))
+    return ArrayPartition(ntuple(i -> _buffer_to_arrp_helper(buf[i]), Val(length(buf))))
 end
 
-function split_into_buffers(raw::AbstractArray, buf; recurse = true)
-    idx = 1
-    function _helper(buf_v; recurse = true)
-        if eltype(buf_v) <: AbstractArray && recurse
-            return _helper.(buf_v; recurse = false)
-        else
-            res = reshape(raw[idx:(idx + length(buf_v) - 1)], size(buf_v))
-            idx += length(buf_v)
-            return res
-        end
-    end
-    return Tuple(_helper(buf_v; recurse) for buf_v in buf)
+_buffer_to_arrp_helper(v::T) where {T} = _buffer_to_arrp_helper(eltype(T), v)
+_buffer_to_arrp_helper(::Type{<:AbstractArray}, v) = buffer_to_arraypartition(v)
+_buffer_to_arrp_helper(::Any, v) = v
+
+function _split_helper(buf_v::T, recurse, raw, idx) where {T}
+    _split_helper(eltype(T), buf_v, recurse, raw, idx)
+end
+
+function _split_helper(::Type{<:AbstractArray}, buf_v, ::Val{true}, raw, idx)
+    map(b -> _split_helper(eltype(b), b, Val(false), raw, idx), buf_v)
+end
+
+function _split_helper(::Type{<:AbstractArray}, buf_v, ::Val{false}, raw, idx)
+    _split_helper((), buf_v, (), raw, idx)
+end
+
+function _split_helper(_, buf_v, _, raw, idx)
+    res = reshape(raw[idx[]:(idx[] + length(buf_v) - 1)], size(buf_v))
+    idx[] += length(buf_v)
+    return res
+end
+
+function split_into_buffers(raw::AbstractArray, buf, recurse = Val(true))
+    idx = Ref(1)
+    ntuple(i -> _split_helper(buf[i], recurse, raw, idx), Val(length(buf)))
+end
+
+function _update_tuple_helper(buf_v::T, raw, idx) where {T}
+    _update_tuple_helper(eltype(T), buf_v, raw, idx)
+end
+
+function _update_tuple_helper(::Type{<:AbstractArray}, buf_v, raw, idx)
+    ntuple(i -> _update_tuple_helper(buf_v[i], raw, idx), Val(length(buf_v)))
+end
+
+function _update_tuple_helper(::Any, buf_v, raw, idx)
+    copyto!(buf_v, view(raw, idx[]:(idx[] + length(buf_v) - 1)))
+    idx[] += length(buf_v)
+    return nothing
 end
 
 function update_tuple_of_buffers(raw::AbstractArray, buf)
-    idx = 1
-    function _helper(buf_v)
-        if eltype(buf_v) <: AbstractArray
-            _helper.(buf_v)
-        else
-            copyto!(buf_v, view(raw, idx:(idx + length(buf_v) - 1)))
-            idx += length(buf_v)
-        end
-    end
-    _helper.(buf)
+    idx = Ref(1)
+    ntuple(i -> _update_tuple_helper(buf[i], raw, idx), Val(length(buf)))
 end
 
 SciMLStructures.isscimlstructure(::MTKParameters) = true
@@ -213,7 +231,7 @@ for (Portion, field) in [(SciMLStructures.Tunable, :tunable)
         @set! p.$field = split_into_buffers(newvals, p.$field)
         if p.dependent_update_oop !== nothing
             raw = p.dependent_update_oop(p...)
-            @set! p.dependent = split_into_buffers(raw, p.dependent; recurse = false)
+            @set! p.dependent = split_into_buffers(raw, p.dependent, Val(false))
         end
         p
     end

test/mtkparameters.jl

@@ -4,6 +4,7 @@ using SymbolicIndexingInterface
 using SciMLStructures: SciMLStructures, canonicalize, Tunable, Discrete, Constants
 using OrdinaryDiffEq
 using ForwardDiff
+using JET
 
 @parameters a b c d::Integer e[1:3] f[1:3, 1:3]::Int g::Vector{AbstractFloat} h::String
 @named sys = ODESystem(
@@ -136,6 +137,77 @@ ps = [p => 1.0] # Value for `d` is missing
 @test_throws ModelingToolkit.MissingVariablesError ODEProblem(sys, u0, tspan, ps)
 @test_nowarn ODEProblem(sys, u0, tspan, [ps..., d => 1.0])
 
+# JET tests
+
+# scalar parameters only
+function level1()
+    @parameters p1=0.5 [tunable = true] p2=1 [tunable = true] p3=3 [tunable = false] p4=3 [tunable = true] y0=1
+    @variables x(t)=2 y(t)=y0
+    D = Differential(t)
+
+    eqs = [D(x) ~ p1 * x - p2 * x * y
+           D(y) ~ -p3 * y + p4 * x * y]
+
+    sys = structural_simplify(complete(ODESystem(
+        eqs, t, tspan = (0, 3.0), name = :sys, parameter_dependencies = [y0 => 2p4])))
+    prob = ODEProblem{true, SciMLBase.FullSpecialize}(sys)
+end
+
+# scalar and vector parameters
+function level2()
+    @parameters p1=0.5 [tunable = true] (p23[1:2]=[1, 3.0]) [tunable = true] p4=3 [tunable = false] y0=1
+    @variables x(t)=2 y(t)=y0
+    D = Differential(t)
+
+    eqs = [D(x) ~ p1 * x - p23[1] * x * y
+           D(y) ~ -p23[2] * y + p4 * x * y]
+
+    sys = structural_simplify(complete(ODESystem(
+        eqs, t, tspan = (0, 3.0), name = :sys, parameter_dependencies = [y0 => 2p4])))
+    prob = ODEProblem{true, SciMLBase.FullSpecialize}(sys)
+end
+
+# scalar and vector parameters with different scalar types
+function level3()
+    @parameters p1=0.5 [tunable = true] (p23[1:2]=[1, 3.0]) [tunable = true] p4::Int=3 [tunable = true] y0::Int=1
+    @variables x(t)=2 y(t)=y0
+    D = Differential(t)
+
+    eqs = [D(x) ~ p1 * x - p23[1] * x * y
+           D(y) ~ -p23[2] * y + p4 * x * y]
+
+    sys = structural_simplify(complete(ODESystem(
+        eqs, t, tspan = (0, 3.0), name = :sys, parameter_dependencies = [y0 => 2p4])))
+    prob = ODEProblem{true, SciMLBase.FullSpecialize}(sys)
+end
+
+@testset "level$i" for (i, prob) in enumerate([level1(), level2(), level3()])
+    ps = prob.p
+    @testset "Type stability of $portion" for portion in [
+        Tunable(), Discrete(), Constants()]
+        @test_call canonicalize(portion, ps)
+        # @inferred canonicalize(portion, ps)
+        broken = (i ∈ [2, 3] && portion == Tunable())
+
+        # broken because the size of a vector of vectors can't be determined at compile time
+        @test_opt broken=broken target_modules=(ModelingToolkit,) canonicalize(
+            portion, ps)
+
+        buffer, repack, alias = canonicalize(portion, ps)
+
+        @test_call SciMLStructures.replace(portion, ps, ones(length(buffer)))
+        @inferred SciMLStructures.replace(portion, ps, ones(length(buffer)))
+        @test_opt target_modules=(ModelingToolkit,) SciMLStructures.replace(
+            portion, ps, ones(length(buffer)))
+
+        @test_call target_modules=(ModelingToolkit,) SciMLStructures.replace!(
+            portion, ps, ones(length(buffer)))
+        @inferred SciMLStructures.replace!(portion, ps, ones(length(buffer)))
+        @test_opt target_modules=(ModelingToolkit,) SciMLStructures.replace!(
+            portion, ps, ones(length(buffer)))
+    end
+end
+
 # Issue#2642
 @parameters α β γ δ
 @variables x(t) y(t)