Merge pull request #4438 from SciML/as/codegen-both

AayushSabharwal · web-flow · commit 915dbf6eeee9 · 2026-04-10T14:52:45.000+05:30
feat: [AI] add `Both` and use it as default `iip` specifier
diff --git a/lib/ModelingToolkitBase/src/ModelingToolkitBase.jl b/lib/ModelingToolkitBase/src/ModelingToolkitBase.jl
@@ -351,6 +351,7 @@ const set_scalar_metadata = setmetadata
 @public renamespace, namespace_equations
 @public check_mutable_cache, store_to_mutable_cache!, should_invalidate_mutable_cache_entry
 @public convert_bindings_for_time_independent_system, get_w
+@public Both
 
 for prop in [SYS_PROPS; [:continuous_events, :discrete_events]]
     getter = Symbol(:get_, prop)
diff --git a/lib/ModelingToolkitBase/src/problems/bvproblem.jl b/lib/ModelingToolkitBase/src/problems/bvproblem.jl
@@ -9,6 +9,7 @@
     check_compatibility && check_compatible_system(BVProblem, sys)
     isnothing(callback) || error("BVP solvers do not support callbacks.")
 
+    _iip = resolve_iip(iip, op)
     dvs = unknowns(sys)
     op = to_varmap(op, dvs)
     # Systems without algebraic equations should use both fixed values + guesses
@@ -17,7 +18,7 @@
 
     fode, u0,
         p = process_SciMLProblem(
-        ODEFunction{iip, spec}, sys, _op; guesses,
+        ODEFunction{_iip, spec}, sys, _op; guesses,
         t = tspan !== nothing ? tspan[1] : tspan, check_compatibility = false, cse,
         checkbounds, time_dependent_init = false, expression, kwargs...
     )
@@ -39,7 +40,7 @@
     f_prototype = n_controls > 0 ? zeros(eltype(u0), length(dvs) - n_controls) : nothing
     bcresid_prototype = zeros(eltype(u0), length(u0_idxs) + length(constraints(sys)))
 
-    bvpfn = BVPFunction{iip}(fode, fbc; cost = fcost, f_prototype, bcresid_prototype)
+    bvpfn = BVPFunction{_iip}(fode, fbc; cost = fcost, f_prototype, bcresid_prototype)
 
     if (length(constraints(sys)) + length(op) > length(dvs))
         @warn "The BVProblem is overdetermined. The total number of conditions (# constraints + # fixed initial values given by op) exceeds the total number of states. The BVP solvers will default to doing a nonlinear least-squares optimization."
@@ -48,7 +49,7 @@
     kwargs = process_kwargs(sys; expression, tspan, kwargs...)
     args = (; bvpfn, u0, tspan, p)
 
-    return maybe_codegen_scimlproblem(expression, BVProblem{iip}, args; kwargs...)
+    return maybe_codegen_scimlproblem(expression, BVProblem{_iip}, args; kwargs...)
 end
 
 function check_compatible_system(T::Type{BVProblem}, sys::System)
diff --git a/lib/ModelingToolkitBase/src/problems/daeproblem.jl b/lib/ModelingToolkitBase/src/problems/daeproblem.jl
@@ -75,10 +75,11 @@ end
     check_complete(sys, DAEProblem)
     check_compatibility && check_compatible_system(DAEProblem, sys)
 
+    _iip = resolve_iip(iip, op)
     f, du0,
         u0,
         p = process_SciMLProblem(
-        DAEFunction{iip, spec}, sys, op;
+        DAEFunction{_iip, spec}, sys, op;
         t = tspan !== nothing ? tspan[1] : tspan, check_length, eval_expression,
         eval_module, check_compatibility, implicit_dae = true, expression, kwargs...
     )
@@ -95,5 +96,5 @@ end
     args = (; f, du0, u0, tspan, p)
     kwargs = (; differential_vars, kwargs...)
 
-    return maybe_codegen_scimlproblem(expression, DAEProblem{iip}, args; kwargs...)
+    return maybe_codegen_scimlproblem(expression, DAEProblem{_iip}, args; kwargs...)
 end
diff --git a/lib/ModelingToolkitBase/src/problems/ddeproblem.jl b/lib/ModelingToolkitBase/src/problems/ddeproblem.jl
@@ -52,9 +52,10 @@ end
     check_complete(sys, DDEProblem)
     check_compatibility && check_compatible_system(DDEProblem, sys)
 
+    _iip = resolve_iip(iip, op)
     f, u0,
         p = process_SciMLProblem(
-        DDEFunction{iip, spec}, sys, op;
+        DDEFunction{_iip, spec}, sys, op;
         t = tspan !== nothing ? tspan[1] : tspan, check_length, cse, checkbounds,
         eval_expression, eval_module, check_compatibility, symbolic_u0 = true,
         expression, u0_constructor, kwargs...
@@ -80,7 +81,7 @@ end
     )
     args = (; f, u0, h, tspan, p)
 
-    return maybe_codegen_scimlproblem(expression, DDEProblem{iip}, args; kwargs...)
+    return maybe_codegen_scimlproblem(expression, DDEProblem{_iip}, args; kwargs...)
 end
 
 function check_compatible_system(T::Union{Type{DDEFunction}, Type{DDEProblem}}, sys::System)
diff --git a/lib/ModelingToolkitBase/src/problems/discreteproblem.jl b/lib/ModelingToolkitBase/src/problems/discreteproblem.jl
@@ -48,12 +48,13 @@ end
     check_complete(sys, DiscreteProblem)
     check_compatibility && check_compatible_system(DiscreteProblem, sys)
 
+    _iip = resolve_iip(iip, op)
     dvs = unknowns(sys)
     op = to_varmap(op, dvs)
     add_toterms!(op; replace = true)
     f, u0,
         p = process_SciMLProblem(
-        DiscreteFunction{iip, spec}, sys, op;
+        DiscreteFunction{_iip, spec}, sys, op;
         t = tspan !== nothing ? tspan[1] : tspan, check_compatibility, expression,
         kwargs...
     )
@@ -67,7 +68,7 @@ end
     kwargs = process_kwargs(sys; kwargs...)
     args = (; f, u0, tspan, p)
 
-    return maybe_codegen_scimlproblem(expression, DiscreteProblem{iip}, args; kwargs...)
+    return maybe_codegen_scimlproblem(expression, DiscreteProblem{_iip}, args; kwargs...)
 end
 
 function check_compatible_system(
diff --git a/lib/ModelingToolkitBase/src/problems/implicitdiscreteproblem.jl b/lib/ModelingToolkitBase/src/problems/implicitdiscreteproblem.jl
@@ -55,20 +55,21 @@ end
     check_complete(sys, ImplicitDiscreteProblem)
     check_compatibility && check_compatible_system(ImplicitDiscreteProblem, sys)
 
+    _iip = resolve_iip(iip, op)
     dvs = unknowns(sys)
     op = to_varmap(op, dvs)
     add_toterms!(op; replace = true)
     f, u0,
         p = process_SciMLProblem(
-        ImplicitDiscreteFunction{iip, spec}, sys, op;
+        ImplicitDiscreteFunction{_iip, spec}, sys, op;
         t = tspan !== nothing ? tspan[1] : tspan, check_compatibility,
         expression, kwargs...
     )
 
     kwargs = process_kwargs(sys; kwargs...)
     args = (; f, u0, tspan, p)
     return maybe_codegen_scimlproblem(
-        expression, ImplicitDiscreteProblem{iip}, args; kwargs...
+        expression, ImplicitDiscreteProblem{_iip}, args; kwargs...
     )
 end
 
diff --git a/lib/ModelingToolkitBase/src/problems/initializationproblem.jl b/lib/ModelingToolkitBase/src/problems/initializationproblem.jl
@@ -39,6 +39,7 @@ All other keyword arguments are forwarded to the wrapped problem constructor.
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `mtkcompile` on the system before creating an `ODEProblem`")
     end
+    _iip = resolve_iip(iip, op)
     if !fast_path
         op = build_operating_point(sys, op)
         fast_path = true
@@ -122,7 +123,7 @@ All other keyword arguments are forwarded to the wrapped problem constructor.
         sys, isys; warn_initialize_determined,
         kwargs...
     )
-    TProb{iip}(isys, op; kwargs..., build_initializeprob = false, is_initializeprob = true)
+    TProb{_iip}(isys, op; kwargs..., build_initializeprob = false, is_initializeprob = true)
 end
 
 function overdetermined_initialization_message(neqs::Integer, nunknown::Integer, extra::AbstractString)
diff --git a/lib/ModelingToolkitBase/src/problems/nonlinearproblem.jl b/lib/ModelingToolkitBase/src/problems/nonlinearproblem.jl
@@ -71,17 +71,18 @@ end
     end
     check_compatibility && check_compatible_system(NonlinearProblem, sys)
 
+    _iip = resolve_iip(iip, op)
     f, u0,
         p = process_SciMLProblem(
-        NonlinearFunction{iip, spec}, sys, op;
+        NonlinearFunction{_iip, spec}, sys, op;
         check_length, check_compatibility, expression, kwargs...
     )
 
     kwargs = process_kwargs(sys; kwargs...)
     ptype = getmetadata(sys, ProblemTypeCtx, StandardNonlinearProblem())
     args = (; f, u0, p, ptype)
 
-    return maybe_codegen_scimlproblem(expression, NonlinearProblem{iip}, args; kwargs...)
+    return maybe_codegen_scimlproblem(expression, NonlinearProblem{_iip}, args; kwargs...)
 end
 
 @fallback_iip_specialize function SciMLBase.NonlinearLeastSquaresProblem{iip, spec}(
@@ -91,17 +92,18 @@ end
     check_complete(sys, NonlinearLeastSquaresProblem)
     check_compatibility && check_compatible_system(NonlinearLeastSquaresProblem, sys)
 
+    _iip = resolve_iip(iip, op)
     f, u0,
         p = process_SciMLProblem(
-        NonlinearFunction{iip}, sys, op;
+        NonlinearFunction{_iip}, sys, op;
         check_length, expression, kwargs...
     )
 
     kwargs = process_kwargs(sys; kwargs...)
     args = (; f, u0, p)
 
     return maybe_codegen_scimlproblem(
-        expression, NonlinearLeastSquaresProblem{iip}, args; kwargs...
+        expression, NonlinearLeastSquaresProblem{_iip}, args; kwargs...
     )
 end
 
diff --git a/lib/ModelingToolkitBase/src/problems/odeproblem.jl b/lib/ModelingToolkitBase/src/problems/odeproblem.jl
@@ -107,9 +107,10 @@ Base.@nospecializeinfer @fallback_iip_specialize function SciMLBase.ODEProblem{i
     check_complete(sys, ODEProblem)
     check_compatibility && check_compatible_system(ODEProblem, sys)
 
+    _iip = resolve_iip(iip, op)
     f, u0,
         p = process_SciMLProblem(
-        ODEFunction{iip, spec}, sys, op;
+        ODEFunction{_iip, spec}, sys, op;
         t = tspan !== nothing ? tspan[1] : tspan, check_length, eval_expression,
         eval_module, expression, check_compatibility, kwargs...
     )
@@ -120,7 +121,7 @@ Base.@nospecializeinfer @fallback_iip_specialize function SciMLBase.ODEProblem{i
 
     ptype = getmetadata(sys, ProblemTypeCtx, StandardODEProblem())
     args = (; f, u0, tspan, p, ptype)
-    maybe_codegen_scimlproblem(expression, ODEProblem{iip}, args; kwargs...)
+    maybe_codegen_scimlproblem(expression, ODEProblem{_iip}, args; kwargs...)
 end
 
 @fallback_iip_specialize function DiffEqBase.SteadyStateProblem{iip, spec}(
@@ -130,17 +131,18 @@ end
     check_complete(sys, SteadyStateProblem)
     check_compatibility && check_compatible_system(SteadyStateProblem, sys)
 
+    _iip = resolve_iip(iip, op)
     f, u0,
         p = process_SciMLProblem(
-        ODEFunction{iip}, sys, op;
+        ODEFunction{_iip}, sys, op;
         steady_state = true, check_length, check_compatibility, expression,
         is_steadystateprob = true, kwargs...
     )
 
     kwargs = process_kwargs(sys; expression, tspan = (0, Inf), kwargs...)
     args = (; f, u0, p)
 
-    maybe_codegen_scimlproblem(expression, SteadyStateProblem{iip}, args; kwargs...)
+    maybe_codegen_scimlproblem(expression, SteadyStateProblem{_iip}, args; kwargs...)
 end
 
 function check_compatible_system(
diff --git a/lib/ModelingToolkitBase/src/problems/sddeproblem.jl b/lib/ModelingToolkitBase/src/problems/sddeproblem.jl
@@ -56,9 +56,10 @@ end
     check_complete(sys, SDDEProblem)
     check_compatibility && check_compatible_system(SDDEProblem, sys)
 
+    _iip = resolve_iip(iip, op)
     f, u0,
         p = process_SciMLProblem(
-        SDDEFunction{iip, spec}, sys, op;
+        SDDEFunction{_iip, spec}, sys, op;
         t = tspan !== nothing ? tspan[1] : tspan, check_length, cse, checkbounds,
         eval_expression, eval_module, check_compatibility, sparse, symbolic_u0 = true,
         expression, u0_constructor, kwargs...
@@ -90,7 +91,7 @@ end
     args = (; f, g, u0, h, tspan, p)
     kwargs = (; noise, noise_rate_prototype, kwargs...)
 
-    return maybe_codegen_scimlproblem(expression, SDDEProblem{iip}, args; kwargs...)
+    return maybe_codegen_scimlproblem(expression, SDDEProblem{_iip}, args; kwargs...)
 end
 
 function check_compatible_system(
diff --git a/lib/ModelingToolkitBase/src/problems/sdeproblem.jl b/lib/ModelingToolkitBase/src/problems/sdeproblem.jl
@@ -84,9 +84,10 @@ end
     check_complete(sys, SDEProblem)
     check_compatibility && check_compatible_system(SDEProblem, sys)
 
+    _iip = resolve_iip(iip, op)
     f, u0,
         p = process_SciMLProblem(
-        SDEFunction{iip, spec}, sys, op;
+        SDEFunction{_iip, spec}, sys, op;
         t = tspan !== nothing ? tspan[1] : tspan, check_length, eval_expression,
         eval_module, check_compatibility, sparse, expression, kwargs...
     )
@@ -113,7 +114,7 @@ end
     args = (; f, u0, tspan, p)
     kwargs = (; noise, noise_rate_prototype, kwargs...)
 
-    return maybe_codegen_scimlproblem(expression, SDEProblem{iip}, args; kwargs...)
+    return maybe_codegen_scimlproblem(expression, SDEProblem{_iip}, args; kwargs...)
 end
 
 function check_compatible_system(T::Union{Type{SDEFunction}, Type{SDEProblem}}, sys::System)
diff --git a/lib/ModelingToolkitBase/src/systems/problem_utils.jl b/lib/ModelingToolkitBase/src/systems/problem_utils.jl
@@ -1802,12 +1802,32 @@ function SymbolicTstops(
     end
 end
 
+"""
+    Both
+
+Sentinel type used as the `iip` type parameter in problem constructors when the caller
+did not explicitly specify in-place vs. out-of-place behavior. The actual value is
+resolved at construction time: `iip = false` when the operating-point `op` is a
+`StaticArray`, and `iip = true` otherwise.
+"""
+struct Both end
+
+"""
+    resolve_iip(iip, op)
+
+Resolve the `iip` type parameter for a problem constructor. When `iip` is `Both`,
+returns `false` if `op` is a `StaticArray` and `true` otherwise. For any other value
+of `iip`, returns `iip` unchanged.
+"""
+resolve_iip(iip, @nospecialize(op)) = iip
+resolve_iip(::Type{Both}, @nospecialize(op)) = !(op isa StaticArray)
+
 """
     $(TYPEDSIGNATURES)
 
 Macro for writing problem/function constructors. Expects a function definition with type
 parameters for `iip` and `specialize`. Generates fallbacks with
-`specialize = SciMLBase.AutoSpecialize` and `iip = true`.
+`specialize = SciMLBase.AutoSpecialize` and `iip = Both` (resolved at construction time).
 """
 # Unwrap `@nospecialize(arg)` to get the underlying argument expression.
 # Returns the argument unchanged if not wrapped in @nospecialize.
@@ -1889,32 +1909,18 @@ macro fallback_iip_specialize(ex)
     fnwhere_iip = Expr(:where, fncall_iip, where_args[1])
     fn_iip = Expr(:function, fnwhere_iip, callexpr_iip)
 
-    # `ODEProblem{true}(call_args...)`
-    callexpr_base = Expr(:call, Expr(:curly, fnname_name, true), call_args...)
+    # Problem constructors default to `Both` (iip resolved at construction time).
+    # Function constructors keep `true` as the iip default.
+    is_problem = occursin("Problem", string(fnname_name))
+    default_iip = is_problem ? Both : true
+    callexpr_base = Expr(:call, Expr(:curly, fnname_name, default_iip), call_args...)
     # `ODEProblem(sig_args...)` - use sig_args for fallback signature
     fncall_base = Expr(:call, fnname_name, sig_args...)
     fn_base = Expr(:function, fncall_base, callexpr_base)
 
-    # Handle case when this is a problem constructor and `u0map` is a `StaticArray`,
-    # where `iip` should default to `false`.
+    # The StaticArray-specific fallback is no longer needed: `Both` defers the
+    # iip decision to the body of the fully-parameterized method.
     fn_sarr = nothing
-    if occursin("Problem", string(fnname_name))
-        # sig_args should at least contain an argument for the `u0map`
-        @assert length(sig_args) > 2
-        u0_arg = sig_args[3]
-        # should not have a type-annotation
-        @assert !Meta.isexpr(u0_arg, :(::))
-        if Meta.isexpr(u0_arg, :kw)
-            argname, default = u0_arg.args
-            u0_arg = Expr(:kw, Expr(:(::), argname, StaticArray), default)
-        else
-            u0_arg = Expr(:(::), u0_arg, StaticArray)
-        end
-
-        callexpr_sarr = Expr(:call, Expr(:curly, fnname_name, false), call_args...)
-        fncall_sarr = Expr(:call, fnname_name, sig_args[1], sig_args[2], u0_arg, sig_args[4:end]...)
-        fn_sarr = Expr(:function, fncall_sarr, callexpr_sarr)
-    end
     return quote
         $fn_base
         $fn_sarr
diff --git a/src/problems/semilinearodeproblem.jl b/src/problems/semilinearodeproblem.jl
@@ -125,17 +125,18 @@ end
     @set! sys.guesses = guess
     @set! sys.initial_conditions = defs
 
+    _iip = resolve_iip(iip, op)
     f, u0,
         p = process_SciMLProblem(
-        SemilinearODEFunction{iip, spec}, sys, op;
+        SemilinearODEFunction{_iip, spec}, sys, op;
         t = tspan !== nothing ? tspan[1] : tspan, expression, check_compatibility,
         semiquadratic_form, sparse, u0_eltype, stiff_linear, stiff_quadratic, stiff_nonlinear, jac, kwargs...
     )
 
     kwargs = process_kwargs(sys; expression, callback, kwargs...)
 
     args = (; f, u0, tspan, p)
-    maybe_codegen_scimlproblem(expression, SplitODEProblem{iip}, args; kwargs...)
+    maybe_codegen_scimlproblem(expression, SplitODEProblem{_iip}, args; kwargs...)
 end
 
 """
