refactor: mark u argument in generated functions where appropriate

lib/ModelingToolkitBase/src/ModelingToolkitBase.jl

@@ -386,6 +386,7 @@ function __init__()
     SU.hashcons(unwrap(ASSERTION_LOG_VARIABLE), true)
     SU.hashcons(DDE_AT_IDX_SYM, true)
     SU.hashcons(DDE_DELAY_SYM, true)
+    SU.hashcons(MTKUNKNOWNS_ARG, true)
     return nothing
 end
 

lib/ModelingToolkitBase/src/inputoutput.jl

@@ -313,7 +313,7 @@ function generate_control_function(
         args = (ddvs, args...)
     end
     f = build_function_wrapper(
-        sys, rhss, args...; p_start = 3 + implicit_dae,
+        sys, rhss, args...; u_arg = 1 + Int(implicit_dae), p_start = 3 + implicit_dae,
         p_end = length(p) + 2 + implicit_dae, kwargs...
     )
     f = eval_or_rgf.(f; eval_expression, eval_module)

lib/ModelingToolkitBase/src/systems/abstractsystem.jl

@@ -40,6 +40,7 @@ function generate_custom_function(
             dvs,
             p...,
             get_iv(sys);
+            u_arg = 1,
             kwargs...,
             expression = Val{true}
         )
@@ -48,6 +49,7 @@ function generate_custom_function(
             sys, exprs,
             dvs,
             p...;
+            u_arg = 1,
             kwargs...,
             expression = Val{true}
         )
@@ -74,6 +76,7 @@ function wrap_assignments(isscalar, assignments; let_block = false)
 end
 
 const MTKPARAMETERS_ARG = SSym(:___mtkparameters___; type = Vector{Vector{Any}}, shape = SymbolicUtils.Unknown(1))
+const MTKUNKNOWNS_ARG = SSym(:___mtkunknowns___; type = Vector{Real}, shape = SymbolicUtils.Unknown(1))
 
 """
     $(TYPEDSIGNATURES)

lib/ModelingToolkitBase/src/systems/callbacks.jl

@@ -1002,7 +1002,7 @@ Base.@nospecializeinfer function compile_condition(
     end
 
     fs = build_function_wrapper(
-        sys, condit, u, p..., t; kwargs..., cse = false
+        sys, condit, u, p..., t; u_arg = 1, kwargs..., cse = false
     )
     fs = GeneratedFunctionWrapper{(2, 3, is_split(sys))}(
         Val{false}, fs...; eval_expression, eval_module
@@ -1451,13 +1451,13 @@ Base.@nospecializeinfer function compile_explicit_affect(
     u_up,
         u_up! = build_function_wrapper(
         sys, (@view rhss[is_u]), dvs, _ps..., t;
-        wrap_code = add_integrator_header(sys, integ, :u),
+        u_arg = 1, wrap_code = add_integrator_header(sys, integ, :u),
         outputidxs = u_idxs, wrap_mtkparameters, iip_config = (false, true)
     )
     p_up,
         p_up! = build_function_wrapper(
         sys, (@view rhss[is_p]), dvs, _ps..., t;
-        wrap_code = add_integrator_header(sys, integ, :p),
+        u_arg = 1, wrap_code = add_integrator_header(sys, integ, :p),
         outputidxs = p_idxs, wrap_mtkparameters, iip_config = (false, true)
     )
 

lib/ModelingToolkitBase/src/systems/codegen.jl

@@ -108,8 +108,9 @@ function generate_rhs(
         p_start += 1
     end
 
+    u_arg = scalar ? -1 : (implicit_dae ? 2 : 1)
     res = build_function_wrapper(
-        sys, rhss, args...; p_start, extra_assignments,
+        sys, rhss, args...; p_start, extra_assignments, u_arg,
         expression = Val{true}, expression_module = eval_module, kwargs...
     )
     nargs = length(args) - length(p) + 1
@@ -147,7 +148,7 @@ function generate_diffusion_function(
         eqs = vec(eqs)
     end
     p = reorder_parameters(sys, ps)
-    res = build_function_wrapper(sys, eqs, dvs, p..., get_iv(sys); kwargs...)
+    res = build_function_wrapper(sys, eqs, dvs, p..., get_iv(sys); u_arg = 1, kwargs...)
     if expression == Val{true}
         return res
     end
@@ -262,7 +263,7 @@ function generate_jacobian(
         nargs = 3
     end
     res = build_function_wrapper(
-        sys, jac, args...; wrap_code, expression = Val{true},
+        sys, jac, args...; wrap_code, u_arg = 1, expression = Val{true},
         expression_module = eval_module, checkbounds, kwargs...
     )
     return maybe_compile_function(
@@ -309,6 +310,7 @@ function generate_tgrad(
         dvs,
         p...,
         get_iv(sys);
+        u_arg = 1,
         expression = Val{true},
         expression_module = eval_module,
         kwargs...
@@ -392,7 +394,7 @@ function generate_W(
     p = reorder_parameters(sys, ps)
     res = build_function_wrapper(
         sys, W, dvs, p..., W_GAMMA, t; wrap_code,
-        p_end = 1 + length(p), checkbounds, kwargs...
+        u_arg = 1, p_end = 1 + length(p), checkbounds, kwargs...
     )
     return maybe_compile_function(
         expression, wrap_gfw, (2, 4, is_split(sys)), res; eval_expression, eval_module
@@ -432,7 +434,7 @@ function generate_dae_jacobian(
     p = reorder_parameters(sys, ps)
     res = build_function_wrapper(
         sys, jac, derivatives, dvs, p..., W_GAMMA, t;
-        p_start = 3, p_end = 2 + length(p), kwargs...
+        u_arg = 2, p_start = 3, p_end = 2 + length(p), kwargs...
     )
     return maybe_compile_function(
         expression, wrap_gfw, (3, 5, is_split(sys)), res; eval_expression, eval_module
@@ -694,8 +696,9 @@ function generate_cost(
         args = (dvs, ps...)
         nargs = 2
     end
+    u_arg = is_time_dependent(sys) ? -1 : 1
     res = build_function_wrapper(
-        sys, obj, args...; expression = Val{true}, p_start, p_end, wrap_delays,
+        sys, obj, args...; expression = Val{true}, p_start, p_end, wrap_delays, u_arg,
         histfn = (p, t) -> BVP_SOLUTION(t), histfn_symbolic = BVP_SOLUTION, kwargs...
     )[1]
     if expression == Val{true}
@@ -788,7 +791,7 @@ function generate_cost_gradient(
     dvs = unknowns(sys)
     ps = reorder_parameters(sys)
     exprs = calculate_cost_gradient(sys; simplify)
-    res = build_function_wrapper(sys, exprs, dvs, ps...; expression = Val{true}, kwargs...)
+    res = build_function_wrapper(sys, exprs, dvs, ps...; u_arg = 1, expression = Val{true}, kwargs...)
     return maybe_compile_function(
         expression, wrap_gfw, (2, 2, is_split(sys)), res; eval_expression, eval_module
     )
@@ -847,7 +850,7 @@ function generate_cost_hessian(
     if sparse
         sparsity = similar(exprs, Float64)
     end
-    res = build_function_wrapper(sys, exprs, dvs, ps...; expression = Val{true}, kwargs...)
+    res = build_function_wrapper(sys, exprs, dvs, ps...; u_arg = 1, expression = Val{true}, kwargs...)
     fn = maybe_compile_function(
         expression, wrap_gfw, (2, 2, is_split(sys)), res; eval_expression, eval_module
     )
@@ -879,7 +882,7 @@ function generate_cons(
     cons = canonical_constraints(sys)
     dvs = unknowns(sys)
     ps = reorder_parameters(sys)
-    res = build_function_wrapper(sys, cons, dvs, ps...; expression = Val{true}, kwargs...)
+    res = build_function_wrapper(sys, cons, dvs, ps...; u_arg = 1, expression = Val{true}, kwargs...)
     return maybe_compile_function(
         expression, wrap_gfw, (2, 2, is_split(sys)), res; eval_expression, eval_module
     )
@@ -936,7 +939,7 @@ function generate_constraint_jacobian(
         sparsity = calculate_constraint_jacobian(
         sys; simplify, sparse, return_sparsity = true
     )
-    res = build_function_wrapper(sys, jac, dvs, ps...; expression = Val{true}, kwargs...)
+    res = build_function_wrapper(sys, jac, dvs, ps...; u_arg = 1, expression = Val{true}, kwargs...)
     fn = maybe_compile_function(
         expression, wrap_gfw, (2, 2, is_split(sys)), res; eval_expression, eval_module
     )
@@ -995,7 +998,7 @@ function generate_constraint_hessian(
         sparsity = calculate_constraint_hessian(
         sys; simplify, sparse, return_sparsity = true
     )
-    res = build_function_wrapper(sys, hess, dvs, ps...; expression = Val{true}, kwargs...)
+    res = build_function_wrapper(sys, hess, dvs, ps...; u_arg = 1, expression = Val{true}, kwargs...)
     fn = maybe_compile_function(
         expression, wrap_gfw, (2, 2, is_split(sys)), res; eval_expression, eval_module
     )
@@ -1048,7 +1051,7 @@ function generate_control_jacobian(
     ps = parameters(sys; initial_parameters = true)
     jac = calculate_control_jacobian(sys; simplify = simplify, sparse = sparse)
     p = reorder_parameters(sys, ps)
-    res = build_function_wrapper(sys, jac, dvs, p..., get_iv(sys); kwargs...)
+    res = build_function_wrapper(sys, jac, dvs, p..., get_iv(sys); u_arg = 1, kwargs...)
     return maybe_compile_function(
         expression, wrap_gfw, (2, 3, is_split(sys)), res; eval_expression, eval_module
     )
@@ -1059,6 +1062,7 @@ function generate_rate_function(js::System, rate)
     return build_function_wrapper(
         js, rate, unknowns(js), p...,
         get_iv(js),
+        u_arg = 1,
         expression = Val{true},
         iip_config = (true, false),
     )[1]

lib/ModelingToolkitBase/src/systems/codegen_utils.jl

@@ -256,6 +256,10 @@ generated functions, and `args` are the arguments.
   `MTKParameters` object are present. These are collapsed into a single argument and
   destructured inside the function. `p_start` must also be provided for non-split systems
   since it is used by `wrap_delays`.
+- `u_arg`: The index in `args` of the argument corresponding to `unknowns(sys)` (the `u`
+  vector). If `-1` (the default), the u vector is not treated specially. Otherwise, the
+  argument must be a `Vector` and is wrapped in a `DestructuredArgs` with the common
+  identifier `MTKUNKNOWNS_ARG`, giving it the predictable name `___mtkunknowns___`.
 - `compress_args`: A list of argument ranges that end before `p_start`.
   Each range will be compressed into a single argument to the function. For example,
   If there are 5 elements in `args` and `compress_args = [2:3]`, then the generated function
@@ -296,7 +300,7 @@ All other keyword arguments are forwarded to `build_function`.
 Base.@nospecializeinfer function build_function_wrapper(
         sys::AbstractSystem, @nospecialize(expr), @nospecialize(args...); p_start = 2,
         p_end = is_time_dependent(sys) ? length(args) - 1 : length(args), compress_args = UnitRange{Int}[],
-        non_standard_param_layout = false,
+        non_standard_param_layout = false, u_arg::Integer = -1,
         wrap_delays = is_dde(sys), histfn = DDE_HISTORY_FUN, histfn_symbolic = histfn, wrap_code = identity,
         add_observed = true, obsidxs_to_use = nothing,
         create_bindings = false, output_type = nothing, mkarray = nothing,
@@ -307,6 +311,18 @@ Base.@nospecializeinfer function build_function_wrapper(
     obs = observed(sys)
     args = Vector{Any}(collect(args))
     assignments = Assignment[]
+
+    if u_arg != -1
+        args[u_arg] isa AbstractVector ||
+            throw(ArgumentError("argument at u_arg = $u_arg must be a Vector, got $(typeof(args[u_arg]))"))
+    end
+
+    u_argument_name = if u_arg == -1
+        generated_argument_name
+    else
+        push!(assignments, Assignment(generated_argument_name(u_arg), :___mtkunknowns___))
+        i -> i == u_arg ? :___mtkunknowns___ : generated_argument_name(i)
+    end
     # turn delayed unknowns into calls to the history function
     if wrap_delays
         param_arg = is_split(sys) ? MTKPARAMETERS_ARG : generated_argument_name(p_start)
@@ -343,7 +359,7 @@ Base.@nospecializeinfer function build_function_wrapper(
 
     # assignments for reconstructing scalarized array symbolics
     if non_standard_param_layout
-        append!(assignments, array_variable_assignments(args...; filter_vars = required_arrvars))
+        append!(assignments, array_variable_assignments(args...; filter_vars = required_arrvars, argument_name = u_argument_name))
     else
         cached = check_mutable_cache(sys, ParameterArrayAssignments, ParameterArrayAssignments, nothing)
         if cached isa ParameterArrayAssignments
@@ -357,14 +373,16 @@ Base.@nospecializeinfer function build_function_wrapper(
                 param_var_to_arridxs; buffer_offset = p_start - 1, filter_vars = required_arrvars
             )
         )
-        other_assigns = array_variable_assignments(args...; ignore_vars = keys(param_var_to_arridxs), filter_vars = required_arrvars)
+        other_assigns = array_variable_assignments(args...; ignore_vars = keys(param_var_to_arridxs), filter_vars = required_arrvars, argument_name = u_argument_name)
         append!(assignments, other_assigns)
     end
     append!(assignments, extra_assignments)
 
     for (i, arg) in enumerate(args)
         # Make sure to use the proper names for arguments
-        args[i] = if symbolic_type(arg) == NotSymbolic() && arg isa AbstractArray
+        args[i] = if u_arg != -1 && i == u_arg
+            DestructuredArgs(arg, MTKUNKNOWNS_ARG; create_bindings)
+        elseif symbolic_type(arg) == NotSymbolic() && arg isa AbstractArray
             DestructuredArgs(arg, generated_argument_name(i); create_bindings)
         else
             arg

src/systems/codegen.jl

@@ -334,18 +334,18 @@ function generate_semiquadratic_functions(
     end
 
     f1_iip = build_function_wrapper(
-        sys, nothing, Symbolics.DEFAULT_OUTSYM, dvs, ps..., iv; p_start = 3,
+        sys, nothing, Symbolics.DEFAULT_OUTSYM, dvs, ps..., iv; u_arg = 2, p_start = 3,
         extra_assignments = f1_iip_ir, expression = Val{true}, iip_config = (true, false), kwargs...
     )[1]
     f2_iip = build_function_wrapper(
-        sys, nothing, Symbolics.DEFAULT_OUTSYM, dvs, ps..., iv; p_start = 3,
+        sys, nothing, Symbolics.DEFAULT_OUTSYM, dvs, ps..., iv; u_arg = 2, p_start = 3,
         extra_assignments = f2_iip_ir, expression = Val{true}, iip_config = (true, false), kwargs...
     )[1]
     f1_oop = build_function_wrapper(
-        sys, f1_expr, dvs, ps..., iv; expression = Val{true}, iip_config = (true, false), kwargs...
+        sys, f1_expr, dvs, ps..., iv; u_arg = 1, expression = Val{true}, iip_config = (true, false), kwargs...
     )[1]
     f2_oop = build_function_wrapper(
-        sys, f2_expr, dvs, ps..., iv; expression = Val{true}, iip_config = (true, false), kwargs...
+        sys, f2_expr, dvs, ps..., iv; u_arg = 1, expression = Val{true}, iip_config = (true, false), kwargs...
     )[1]
 
     f1 = maybe_compile_function(
@@ -511,11 +511,11 @@ function generate_semiquadratic_jacobian(
     oop_expr = length(terms) == 1 ? only(terms) : term(+, terms...)
 
     j_iip, _ = build_function_wrapper(
-        sys, nothing, Symbolics.DEFAULT_OUTSYM, dvs, ps..., iv; p_start = 3,
+        sys, nothing, Symbolics.DEFAULT_OUTSYM, dvs, ps..., iv; u_arg = 2, p_start = 3,
         extra_assignments = iip_ir, expression = Val{true}, iip_config = (true, false), kwargs...
     )
     j_oop, _ = build_function_wrapper(
-        sys, oop_expr, dvs, ps..., iv; expression = Val{true}, iip_config = (true, false), kwargs...
+        sys, oop_expr, dvs, ps..., iv; u_arg = 1, expression = Val{true}, iip_config = (true, false), kwargs...
     )
     return maybe_compile_function(
         expression, wrap_gfw, (2, 3, is_split(sys)),