Skip to content

Commit 7d8b5e6

Browse files
committed
Skip parameter args when rebuilding non-parameter array variables
`build_function_wrapper`` reconstructs scalarized array variables in two passes: the parameters (whose layout is cached and reused) and everything else (unknowns, cache vars, the independent variable). The second pass walked the parameter args as well, running split_indexed_var on every parameter symbol and then throwing the result away because parameters are already handled by the cache.
1 parent 6745da8 commit 7d8b5e6

1 file changed

Lines changed: 10 additions & 6 deletions

File tree

lib/ModelingToolkitBase/src/systems/codegen_utils.jl

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -182,11 +182,11 @@ function generated_argument_name(i::Int)
182182
return Symbol(:__mtk_arg_, i)
183183
end
184184

185-
function compute_array_variable_buffer_idxs(@nospecialize(args); ignore_vars = Set{SymbolicT}())
186-
return _compute_array_variable_buffer_idxs(args isa Vector ? args : collect(args), ignore_vars)
185+
function compute_array_variable_buffer_idxs(@nospecialize(args); ignore_vars = Set{SymbolicT}(), ignore_arg_idxs = nothing)
186+
return _compute_array_variable_buffer_idxs(args isa Vector ? args : collect(args), ignore_vars, ignore_arg_idxs)
187187
end
188188

189-
function _compute_array_variable_buffer_idxs(args::Vector, ignore_vars)
189+
function _compute_array_variable_buffer_idxs(args::Vector, ignore_vars, ignore_arg_idxs = nothing)
190190
# map array symbolic to an identically sized array where each element is (buffer_idx, idx_in_buffer)
191191
var_to_arridxs = Dict{SymbolicT, Vector{Tuple{Int, Int}}}()
192192
for (i, arg) in enumerate(args)
@@ -195,6 +195,10 @@ function _compute_array_variable_buffer_idxs(args::Vector, ignore_vars)
195195
# scalarized array symbolic. This works because the only non-array element
196196
# is the independent variable
197197
arg isa Vector{SymbolicT} || continue
198+
# entire arg-vectors whose decomposition is already accounted for (e.g. the
199+
# parameter slice, handled via the cached `param_var_to_arridxs`) are skipped
200+
# here so we never re-run `split_indexed_var`/`get_stable_index` on them.
201+
ignore_arg_idxs === nothing || i in ignore_arg_idxs && continue
198202

199203
# go through symbolics
200204
for (j, var) in enumerate(arg)
@@ -291,9 +295,9 @@ reconstruct array variables if they are present scalarized in `args`.
291295
"""
292296
function array_variable_assignments(
293297
@nospecialize(args...); ignore_vars = Set{SymbolicT}(), filter_vars = nothing,
294-
argument_name = generated_argument_name, buffer_offset = 0
298+
argument_name = generated_argument_name, buffer_offset = 0, ignore_arg_idxs = nothing
295299
)
296-
var_to_arridxs = compute_array_variable_buffer_idxs(args; ignore_vars)
300+
var_to_arridxs = compute_array_variable_buffer_idxs(args; ignore_vars, ignore_arg_idxs)
297301
return array_variable_buffer_idxs_to_assignments(var_to_arridxs; argument_name, buffer_offset, filter_vars)
298302
end
299303

@@ -555,7 +559,7 @@ Base.@nospecializeinfer function build_function_wrapper(
555559
param_var_to_arridxs; buffer_offset = p_start - 1, filter_vars = required_arrvars
556560
)
557561
)
558-
other_assigns = array_variable_assignments(args...; ignore_vars = keys(param_var_to_arridxs), filter_vars = required_arrvars, argument_name = u_argument_name)
562+
other_assigns = array_variable_assignments(args...; ignore_vars = keys(param_var_to_arridxs), filter_vars = required_arrvars, argument_name = u_argument_name, ignore_arg_idxs = p_start:p_end)
559563
append!(assignments, other_assigns)
560564
end
561565
append!(assignments, extra_assignments)

0 commit comments

Comments
 (0)