diff --git a/lib/ModelingToolkitBase/src/utils.jl b/lib/ModelingToolkitBase/src/utils.jl index cc27ce015b..32cf45be44 100644 --- a/lib/ModelingToolkitBase/src/utils.jl +++ b/lib/ModelingToolkitBase/src/utils.jl @@ -1263,79 +1263,94 @@ function observed_equations_used_by( end """ - $(TYPEDSIGNATURES) - -Given an expression `expr`, return a dictionary mapping subexpressions of `expr` that do -not involve variables in `vars` to anonymous symbolic variables. Also return the modified -`expr` with the substitutions indicated by the dictionary. If `expr` is a function -of only `vars`, then all of the returned subexpressions can be precomputed. - -Note that this will only process subexpressions floating point value. Additionally, -array variables must be passed in both scalarized and non-scalarized forms in `vars`. -""" -function subexpressions_not_involving_vars(expr, vars) - expr = unwrap(expr) - vars = map(unwrap, vars) - state = Dict() - newexpr = subexpressions_not_involving_vars!(expr, vars, state) - return state, newexpr -end - -""" - $(TYPEDSIGNATURES) + $TYPEDSIGNATURES -Mutating version of `subexpressions_not_involving_vars` which writes to `state`. Only -returns the modified `expr`. +Given a list of expressions `exprs`, find all top-level subexpressions in `exprs` +that do not involve variables in `banned_vars`. "Top-level" implies that for all +such subexpressions, any parent of theirs in `exprs` will involve something in +`banned_vars`. `state` will be populated as a map from the identified subexpressions +to anonymous symbols they can be replaced by. """ -function subexpressions_not_involving_vars!(expr, vars, state::Dict{Any, Any}) - expr = unwrap(expr) - if symbolic_type(expr) == NotSymbolic() - if is_array_of_symbolics(expr) - return map(expr) do el - subexpressions_not_involving_vars!(el, vars, state) - end +function subexpressions_not_involving_vars!( + ir::SU.IRStructure{VartypeT}, exprs::AbstractArray{SymbolicT}, + banned_vars::Set{SymbolicT}, state::Dict{SymbolicT, SymbolicT} + ) + # Populate the IR first to ensure that the `RecursiveDFS` has a correctly sized + # `visited` buffer. + for x in exprs + populate_ir!(ir, x) + end + for x in banned_vars + populate_ir!(ir, x) + end + # Get the nodes that are reachable from `exprs`. We need to find the topologically + # earliest subset of these nodes that do not contain `banned_vars`. + reachable_nodes = Set{Int}() + # We only want to descent into non-atomic nodes. Otherwise e.g. the index `1` in `x[1]` + # will end up being a subexpression not involving `banned_vars`. + filtered_nbors = let ir = ir + function __filtered_nbors(graph, i) + # Use `Iterators.filter` instead of just checking `ir[i]` and returning `()` for + # type-stability. The early-exit infers as a `Union`. + is_atomic = SU.default_is_atomic(ir[i]) + # We also don't want to descend into constants - those are not worth caching. + Iterators.filter(j -> !is_atomic && !SU.isconst(ir[j]), Graphs.outneighbors(graph, i)) end - return expr - end - any(isequal(expr), vars) && return expr - iscall(expr) || return expr - symbolic_has_known_size(expr) || return expr - haskey(state, expr) && return state[expr] - op = operation(expr) - args = arguments(expr) - # if this is a `getindex` and the getindex-ed value is a `Sym` - # or it is not a called parameter - # OR - # none of `vars` are involved in `expr` - if op === getindex && (issym(args[1]) || !iscalledparameter(args[1])) || - (vs = SU.search_variables(expr); intersect!(vs, vars); isempty(vs)) - sym = gensym(:subexpr) - var = similar_variable(expr, sym) - state[expr] = var - return var - end - - if (op == (+) || op == (*)) && symbolic_type(expr) !== ArraySymbolic() - indep_args = SymbolicT[] - dep_args = SymbolicT[] - for arg in args - _vs = SU.search_variables(arg) - intersect!(_vs, vars) - if !isempty(_vs) - push!(dep_args, subexpressions_not_involving_vars!(arg, vars, state)) - else - push!(indep_args, arg) - end + end + rdfs = SU.RecursiveDFS( + ir.dependency_graph; neighbors_fn = filtered_nbors, + on_exit = Base.Fix1(push!, reachable_nodes) + ) + for x in exprs + for idx in ir.weak_definitions[x] + rdfs(idx) + end + end + # We want to retain the reachability information for later + unbanned_subexprs = copy(reachable_nodes) + # Walk through the usages of `banned_vars` that are in `unbanned_subexprs`, and remove + # from the candidates any expression we encounter. The remaining vertices are + # ones that do not use `banned_vars`. + unbanned_nbors_fn = let unbanned_subexprs = unbanned_subexprs + function __unbanned_nbors_fn(graph, i) + return Iterators.filter(in(unbanned_subexprs), Graphs.inneighbors(graph, i)) + end + end + rdfs = SU.RecursiveDFS( + ir.dependency_graph; + neighbors_fn = unbanned_nbors_fn, + on_exit = Base.Fix1(delete!, unbanned_subexprs) + ) + for var in banned_vars + for idx in ir.weak_definitions[var] + rdfs(idx) end - indep_term = reduce(op, indep_args; init = Int(op == (*))) - indep_term = subexpressions_not_involving_vars!(indep_term, vars, state) - dep_term = reduce(op, dep_args; init = Int(op == (*))) - return op(indep_term, dep_term) end - newargs = map(args) do arg - subexpressions_not_involving_vars!(arg, vars, state) + + # Now, the nodes we actually care about and will populate `state` with are ones + # present in `unbanned_subexprs` and are used by a node in + # `setdiff(reachable_nodes, unbanned_subexprs)`. All of `setdiff!`, the subsequent + # `filter!`, and then populating `state` will iterate over `unbanned_subexprs`. We + # might as well combine this into one loop. + for node in unbanned_subexprs + nbors = Graphs.inneighbors(ir.dependency_graph, node) + is_valid_node = false + for nbor in nbors + nbor in reachable_nodes || continue + nbor in unbanned_subexprs && continue + is_valid_node = true + break + end + is_valid_node || continue + + expr = ir[node] + haskey(state, expr) && continue + anon_sym = Symbolics.SSym( + Symbol(:__cached_, length(state)); + type = SU.symtype(expr), shape = SU.shape(expr) + ) + state[expr] = anon_sym end - return maketerm(typeof(expr), op, newargs, metadata(expr)) end """ diff --git a/src/problems/sccnonlinearproblem.jl b/src/problems/sccnonlinearproblem.jl index d1be37ecf3..4c16426d72 100644 --- a/src/problems/sccnonlinearproblem.jl +++ b/src/problems/sccnonlinearproblem.jl @@ -65,45 +65,30 @@ equations are subset, delete accordingly. Requires that `sys` is complete and fl - `available_vars`: A list of variables that the subset system should assume are precomputed or already available. Will be mutated with the unknowns and observables of the subset. This is useful for SCC decomposition. -- `prevobsidxs`: Indices of observed equations that the subset system should assume are - precomputed or already available. Will be appended with the indices of observed equations - required by this subset. """ function subset_system( sys::System, vscc::Vector{Int}, escc::Vector{Int}; - available_vars = Set{SymbolicT}(), prevobsidxs = Int[] + available_vars = Set{SymbolicT}() ) check_complete(sys, "subset_system") @assert isempty(get_systems(sys)) "`subset_system` requires a flattened system" dvs = unknowns(sys) ps = parameters(sys) - eqs = equations(sys) - obs = observed(sys) + eqs = full_equations(sys) # subset unknowns and equations _dvs = dvs[vscc] _eqs = eqs[escc] - # get observed equations required by this SCC union!(available_vars, _dvs) - obsidxs = observed_equations_used_by(sys, _eqs; available_vars) - # the ones used by previous SCCs can be precomputed into the cache - setdiff!(obsidxs, prevobsidxs) - _obs = obs[obsidxs] - _observables = SymbolicT[] - for eq in _obs - push!(_observables, eq.lhs) - end - union!(available_vars, _observables) - append!(prevobsidxs, obsidxs) subsys = ConstructionBase.setproperties( - sys; unknowns = _dvs, eqs = _eqs, observed = _obs, + sys; unknowns = _dvs, eqs = _eqs, observed = Equation[], parameter_bindings_graph = get_parameter_bindings_graph(sys), complete = true ) if get_index_cache(sys) !== nothing @set! subsys.index_cache = subset_unknowns_observed( - get_index_cache(sys), sys, _dvs, getproperty.(_obs, (:lhs,)) + get_index_cache(sys), sys, _dvs, SymbolicT[], ) end cached_param_arr_assigns = check_mutable_cache( @@ -120,29 +105,26 @@ end const BlockIdxsT = typeof(BlockVector{Int}(undef_blocks, Int[])) -mutable struct SCCDecomposition - const subsystems::Vector{System} - const var_sccs::Vector{Vector{Int}} - const eq_sccs::Vector{Vector{Int}} - const islinear::BitVector - const hints::Vector{StructuralHint.Type} +struct SCCDecomposition + subsystems::Vector{System} + var_sccs::Vector{Vector{Int}} + eq_sccs::Vector{Vector{Int}} + islinear::BitVector + hints::Vector{StructuralHint.Type} # Cache buffer types and corresponding sizes. Stored as a pair of arrays instead of a # dict to maintain a consistent order of buffers across SCCs - const cachetypes::Vector{TypeT} - const cachesizes::Vector{Int} + cachetypes::Vector{TypeT} + cachesizes::Vector{Int} # explicitfun! related information for each SCC # We need to compute buffer sizes before doing any codegen - const scc_cachevars::Vector{SCCCacheVarsExprsElT} - const scc_cacheexprs::Vector{SCCCacheVarsExprsElT} - obsidxs::BlockIdxsT - const obsidxs_for_cacheexprs::Vector{Vector{Int}} + scc_cachevars::Vector{SCCCacheVarsExprsElT} + scc_cacheexprs::Vector{SCCCacheVarsExprsElT} end function SCCDecomposition() return SCCDecomposition( System[], Vector{Int}[], Vector{Int}[], BitVector(), StructuralHint.Type[], - TypeT[], Int[], SCCCacheVarsExprsElT[], SCCCacheVarsExprsElT[], - BlockVector{Int}(undef_blocks, Int[]), Vector{Int}[] + TypeT[], Int[], SCCCacheVarsExprsElT[], SCCCacheVarsExprsElT[] ) end @@ -164,9 +146,8 @@ function SCCDecomposition( final_decomposition, active_decomposition, first(nbors), active, nbors ) end - blockpush!(active_decomposition.obsidxs, Int[]) subsys = subset_system( - sys, vscc, escc; available_vars, prevobsidxs = active_decomposition.obsidxs + sys, vscc, escc; available_vars ) push!(active_decomposition.subsystems, subsys) push!(active_decomposition.var_sccs, vscc) @@ -189,7 +170,6 @@ function SCCDecomposition( end function copy_scc!(dst::SCCDecomposition, src::SCCDecomposition, tgt::Int) - blockpush!(dst.obsidxs, copy(view(src.obsidxs, Block(tgt)))) push!(dst.subsystems, src.subsystems[tgt]) push!(dst.var_sccs, src.var_sccs[tgt]) push!(dst.eq_sccs, src.eq_sccs[tgt]) @@ -304,33 +284,27 @@ end function build_caches!(sys::System, decomposition::SCCDecomposition) banned_vars = Set{SymbolicT}() - state = Dict() - prev_obseqs = Equation[] + state = Dict{SymbolicT, SymbolicT}() + ir = get_irstructure(sys) for i in eachindex(decomposition.subsystems) empty!(banned_vars) empty!(state) subsys = decomposition.subsystems[i] union!(banned_vars, unknowns(subsys)) - union!(banned_vars, observables(subsys)) for u in unknowns(subsys) push!(banned_vars, split_indexed_var(u)[1]) end - for u in observables(subsys) - push!(banned_vars, split_indexed_var(u)[1]) - end - _obs = get_observed(subsys) - for i in eachindex(_obs) - _obs[i] = _obs[i].lhs ~ subexpressions_not_involving_vars!( - _obs[i].rhs, banned_vars, state - ) - end _eqs = get_eqs(subsys) + exprs_to_search = SymbolicT[] for i in eachindex(_eqs) - _eqs[i] = _eqs[i].lhs ~ subexpressions_not_involving_vars!( - _eqs[i].rhs, banned_vars, state - ) + push!(exprs_to_search, _eqs[i].rhs) + end + subexpressions_not_involving_vars!(ir, exprs_to_search, banned_vars, state) + subber = SU.IRSubstituter{false}(ir, state; filterer = !SU.default_is_atomic) + for i in eachindex(_eqs) + _eqs[i] = _eqs[i].lhs ~ subber(_eqs[i].rhs) end if decomposition.islinear[i] @@ -351,19 +325,6 @@ function build_caches!(sys::System, decomposition::SCCDecomposition) cacheexprs = SCCCacheVarsExprsElT() push!(decomposition.scc_cachevars, cachevars) push!(decomposition.scc_cacheexprs, cacheexprs) - # observed of previous SCCs are in the cache - # NOTE: When we get proper CSE, we can substitute these - # and then use `subexpressions_not_involving_vars!` - for eq in prev_obseqs - T = symtype(eq.lhs) - buf = get!(() -> SymbolicT[], cachevars, T) - push!(buf, eq.lhs) - - buf = get!(() -> SymbolicT[], cacheexprs, T) - push!(buf, eq.lhs) - end - - append!(prev_obseqs, observed(subsys)) for (k, v) in state k = unwrap(k) @@ -375,8 +336,6 @@ function build_caches!(sys::System, decomposition::SCCDecomposition) push!(buf, k) end all_cacheexprs = reduce(vcat, values(cacheexprs); init = SymbolicT[]) - obsidxs_for_scc_cacheexprs = observed_equations_used_by(sys, all_cacheexprs) - push!(decomposition.obsidxs_for_cacheexprs, obsidxs_for_scc_cacheexprs) # update the sizes of cache buffers for (T, buf) in cachevars idx = findfirst(isequal(T), decomposition.cachetypes) @@ -400,7 +359,6 @@ function _collapse_into!(decomposition::SCCDecomposition, i::Int, js) parent = decomposition.subsystems[i] new_eqs = copy(equations(parent)) new_dvs = copy(unknowns(parent)) - new_obs = copy(observed(parent)) cached_ab::Union{CachedLinearAb, Nothing} = if decomposition.islinear[i] calculate_A_b(parent) check_mutable_cache(parent, CachedLinearAb, CachedLinearAb, nothing)::CachedLinearAb @@ -411,7 +369,6 @@ function _collapse_into!(decomposition::SCCDecomposition, i::Int, js) cur = decomposition.subsystems[j] append!(new_eqs, equations(cur)) append!(new_dvs, unknowns(cur)) - append!(new_obs, observed(cur)) decomposition.islinear[i] &= decomposition.islinear[j] if cached_ab isa CachedLinearAb && decomposition.islinear[j] A = cached_ab.A @@ -426,21 +383,13 @@ function _collapse_into!(decomposition::SCCDecomposition, i::Int, js) append!(decomposition.var_sccs[i], decomposition.var_sccs[j]) append!(decomposition.eq_sccs[i], decomposition.eq_sccs[j]) end - unique!(new_obs) new_parent = decomposition.subsystems[i] = ConstructionBase.setproperties( - parent; eqs = new_eqs, unknowns = new_dvs, observed = new_obs + parent; eqs = new_eqs, unknowns = new_dvs, ) if cached_ab isa CachedLinearAb store_to_mutable_cache!(new_parent, CachedLinearAb, cached_ab) end - parent_obsidxs = view(decomposition.obsidxs, Block(i)) - for j in js - append!(parent_obsidxs, view(decomposition.obsidxs, Block(j))) - end - sort!(parent_obsidxs) - unique!(parent_obsidxs) - return nothing end @@ -463,7 +412,6 @@ function SCCNonlinearFunction{iip}( f = generate_rhs( subsys; expression = Val{false}, wrap_gfw = Val{true}, cachesyms, eval_expression, eval_module, - obsidxs_to_use = eachindex(observed(subsys)) ) return NonlinearFunction{iip}(f; sys = subsys) @@ -536,7 +484,6 @@ function SciMLBase.SCCNonlinearProblem{iip}( dvs = unknowns(sys) ps = parameters(sys) eqs = equations(sys) - obs = observed(sys) _, u0, p = process_SciMLProblem( EmptySciMLFunction{iip}, sys, op; eval_expression, eval_module, symbolic_u0 = true, @@ -559,7 +506,6 @@ function SciMLBase.SCCNonlinearProblem{iip}( cachevars = decomposition.scc_cachevars[i] cacheexprs = decomposition.scc_cacheexprs[i] subsys = decomposition.subsystems[i] - _prevobsidxs = decomposition.obsidxs_for_cacheexprs[i] if isempty(cachevars) push!(explicitfuns, Returns(nothing)) else @@ -670,7 +616,7 @@ function SciMLBase.SCCNonlinearProblem{iip}( new_eqs = eqs[reduce(vcat, decomposition.eq_sccs)] sys = ConstructionBase.setproperties( sys; unknowns = new_dvs, eqs = new_eqs, index_cache = subset_unknowns_observed( - get_index_cache(sys), sys, new_dvs, getproperty.(obs, (:lhs,)) + get_index_cache(sys), sys, new_dvs, SymbolicT[] ) )