Correctly compute state/equation indices for IPO

This includes a refactor that considers mapping to states,
instead of individual mappings into u/du vectors

Author: serenity4

src/transform/codegen/dae_factory.jl

@@ -147,7 +147,7 @@ function dae_factory_gen(state::TransformationState, ci::CodeInstance, key::Unio
     argt = Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}, SciMLBase.NullParameters, Float64}
     sicm = ()
     if settings.skip_optimizations
-        daef_ci = rhs_finish_noopt!(state, ci, key, world, settings, 1)
+        daef_ci = rhs_finish_noopt!(state, ci, key, world, settings)
         oc = sciml_to_internal_abi_noopt!(copy(ci.inferred.ir), state, daef_ci, settings)
     else
         # TODO: We should not have to recompute this here

src/transform/codegen/rhs.jl

@@ -1,5 +1,4 @@
 const VectorViewType = SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int}}, true}
-const VectorIntViewType = SubArray{Int, 1, Vector{Int}, Tuple{UnitRange{Int}}, true}
 
 """
     struct RHSSpec

src/transform/common.jl

@@ -116,7 +116,7 @@ function replace_call!(ir::Union{IRCode,IncrementalCompact}, idx::SSAValue, @nos
     ir[idx][:type] = Any
     ir[idx][:info] = Compiler.NoCallInfo()
     ir[idx][:flag] |= Compiler.IR_FLAG_REFINED
-    return new_call
+    return idx
 end
 
 function maybe_insert_debuginfo!(compact::IncrementalCompact, settings::Settings, source::LineNumberNode, previous = nothing, i = compact.result_idx)
@@ -129,6 +129,7 @@ function maybe_insert_debuginfo!(debuginfo::DebugInfoStream, settings::Settings,
 end
 
 function insert_debuginfo!(debuginfo::DebugInfoStream, i::Integer, source::LineNumberNode, previous)
+    prev_edge_index = prev_edge_line = nothing
     if previous !== nothing && isa(previous, Tuple)
         prev_edge_index, prev_edge_line = previous[2], previous[3]
     end

src/transform/reconstruct.jl

@@ -7,47 +7,35 @@ function expand_residuals(state::TransformationState, key::TornCacheKey, compres
     # (slot_assignments, var_assignment, _) = assign_slots(state, key, var_eq_matching)
     for (eq, incidence) in enumerate(result.total_incidence)
         if !is_var_part_known_linear(incidence)
-            sign = get_linear_residual_sign(result, eq)
-            push!(expanded, sign * compressed[i])
+            push!(expanded, compressed[i])
             i += 1
             continue
         end
 
         residual = 0.0
-        sign = get_linear_residual_sign(result, eq)
         for (coeff, var) in zip(nonzeros(incidence.row), rowvals(incidence.row))
             var -= 1
             if var == 0
                 value = t
             else
-                vint = invview(structure.var_to_diff)[var]
-                (slot, source) = vint === nothing ? (var, u) : (vint, du)
-                vint !== nothing && (i += 1)
+                is_diff = is_differential_variable(structure, var)
+                source = ifelse(is_diff, du, u)
+                slot = count(v -> is_differential_variable(structure, v) == is_diff, 1:var)
+                eq == var && !is_diff && (i += 1)
                 value = source[slot]
             end
             residual += value * coeff
         end
         constant_term = isa(incidence.typ, Const) ? incidence.typ.val::Float64 : 0.0
-        push!(expanded, constant_term + sign * residual)
+        push!(expanded, constant_term + residual)
     end
     return expanded
 end
 
-function get_linear_residual_sign(result::DAEIPOResult, eq::Int)
-    # If a linear solved term appears with a positive coefficient,
-    # the residual will be taken as the negative of the value provided to `always!`.
-    # For example: ẋ₁ - x₁x₂ = 0
-    #                    -ẋ₁ = -x₁x₂
-    #                     ẋ₁ = -x₁x₂/-1
-    #                     ẋ₁ = x₁x₂
-    #                      0 = x₁x₂ - ẋ₁   <-- residual
-    incidence = result.total_incidence[eq]
-    for (coeff, var) in zip(nonzeros(incidence.row), rowvals(incidence.row))
-        var - 1 === eq || continue
-        isa(coeff, Float64) || continue
-        return coeff ≤ 0 ? -1 : 1
-    end
-    return 1
+function is_differential_variable(structure::DAESystemStructure, var)
+    structure.var_to_diff[var] !== nothing && return false
+    return invview(structure.var_to_diff)[var] !== nothing && return true
+    @assert false
 end
 
 function expand_residuals(f, residuals, u, du, t)
@@ -71,6 +59,7 @@ function extract_removed_states(state::TransformationState, key::TornCacheKey, t
         vint === nothing || key.diff_states === nothing || !in(vint, key.diff_states) || continue
         push!(removed_states, var)
     end
+    # @sshow removed_states
     return removed_states
 end
 

src/transform/unoptimized.jl

@@ -6,67 +6,90 @@ function rhs_finish_noopt!(
     key::UnoptimizedKey,
     world::UInt,
     settings::Settings,
-    indexT=Int)
+    equation_to_residual_mapping = 1:length(state.structure.eq_to_diff),
+    variable_to_state_mapping = map_variables_to_states(state))
 
     (; result, structure) = state
-    result_ci = find_matching_ci(ci -> ci.inferred === key, ci.def, world)
+    result_ci = find_matching_ci(ci -> ci.owner === key, ci.def, world)
     if result_ci !== nothing
         return result_ci
     end
 
     ir = copy(result.ir)
-    slotnames = [:captures, :vars, :out, :du, :u, :out_indices, :du_indices, :u_indices, :t]
-    argtypes = [Tuple, Vector{Float64}, Vector{Float64}, Vector{Float64}, Vector{Float64}, VectorIntViewType, VectorIntViewType, VectorIntViewType, Float64]
+    # TODO: use original function arguments too
+    slotnames = [:captures, :out, :du, :u, :residuals, :states, :t]
+    argtypes = [Tuple, Vector{Float64}, Vector{Float64}, Vector{Float64}, Vector{Int}, Vector{Int}, Float64]
     append!(empty!(ir.argtypes), argtypes)
-    captures, vars, out, du, u, out_indices, du_indices, u_indices, t = Argument.(eachindex(slotnames))
-    # TODO: use `out_indices`, `du_indices`, `u_indices`
+    captures, out, du, u, residuals, states, t = Argument.(eachindex(slotnames))
     @assert length(slotnames) == length(ir.argtypes)
 
     equations = Pair{SSAValue, Eq}[]
+    callee_to_caller_eq_map = invert_eq_callee_mapping(result.eq_callee_mapping)
     compact = IncrementalCompact(ir)
 
-    for ((_, i), _) in compact
+    # @sshow equation_to_residual_mapping variable_to_state_mapping
+
+    for ((old, i), _) in compact
         ssaidx = SSAValue(i)
         inst = compact[ssaidx]
         stmt = inst[:stmt]
         type = inst[:type]
         line = inst[:line]
 
-        if is_known_invoke(stmt, Intrinsics.equation, compact)
+        if i == 1
+            # @insert_instruction_here(compact, nothing, settings, println("Residuals: ", residuals)::Any)
+            # @insert_instruction_here(compact, nothing, settings, println("States: ", states)::Any)
+            replace_uses!(compact, (old, inst) => SSAValue(3))
+        end
+
+        if is_known_invoke_or_call(stmt, Intrinsics.variable, compact)
+            var = idnum(type)
+            index = @insert_instruction_here(compact, line, settings, getindex(states, var)::Int)
+            value = @insert_instruction_here(compact, line, settings, getindex(u, index)::Float64)
+            replace_uses!(compact, (old, inst) => value)
+            # @insert_instruction_here(compact, line, settings, println("Variable (", var, "): ", value)::Float64)
+        elseif is_known_invoke(stmt, Intrinsics.ddt, compact)
+            var = idnum(type)
+            index = @insert_instruction_here(compact, line, settings, getindex(states, var)::Int)
+            value = @insert_instruction_here(compact, line, settings, getindex(du, index)::Float64)
+            replace_uses!(compact, (old, inst) => value)
+            # @insert_instruction_here(compact, line, settings, println("Variable derivative (", var, " := ", invview(structure.var_to_diff)[var], "′): ", value)::Any)
+        elseif is_known_invoke(stmt, Intrinsics.equation, compact)
             push!(equations, ssaidx => type::Eq)
             inst[:stmt] = nothing
-        elseif is_known_invoke(stmt, Intrinsics.ddt, compact)
-            var = invview(structure.var_to_diff)[idnum(type)]
-            getdu = Expr(:call, getindex, du, var)
-            replace_call!(compact, ssaidx, getdu, settings, @__SOURCE__)
-            inst[:type] = Float64
         elseif is_equation_call(stmt, compact)
             callee, value = stmt.args[2], stmt.args[3]
             i = findfirst(x -> first(x) == callee, equations)::Int
             eq = last(equations[i])
-            call = Expr(:call, setindex!, out, value, eq.id)
-            replace_call!(compact, ssaidx, call, settings, @__SOURCE__)
-        elseif is_known_invoke_or_call(stmt, Intrinsics.variable, compact)
-            var = idnum(type)
-            call = Expr(:call, getindex, u, var)
-            replace_call!(compact, ssaidx, call, settings, @__SOURCE__)
-            inst[:type] = Float64
+            index = @insert_instruction_here(compact, line, settings, getindex(residuals, eq.id)::Int)
+            ret = @insert_instruction_here(compact, line, settings, setindex!(out, value, index)::Any)
+            replace_uses!(compact, (old, inst) => ret)
+            # @insert_instruction_here(compact, line, settings, println("Residuals (index = ", index, ", value = ", value, "): ", residuals)::Any)
         elseif is_known_invoke_or_call(stmt, Intrinsics.sim_time, compact)
             inst[:stmt] = t
         elseif is_known_invoke_or_call(stmt, Intrinsics.epsilon, compact)
             inst[:stmt] = 0.0
         elseif isexpr(stmt, :invoke)
-            @sshow stmt
+            info = inst[:info]::MappingInfo
             callee_ci, callee_f = stmt.args[1]::CodeInstance, stmt.args[2]
             callee_result = structural_analysis!(callee_ci, world, settings)
             callee_structure = make_structure_from_ipo(callee_result)
             callee_state = TransformationState(callee_result, callee_structure)
-            callee_daef_ci = rhs_finish_noopt!(callee_state, callee_ci, UnoptimizedKey(), world, settings)
-            callee_captures = ()
-            # TODO: compute indices into `u`/`du`/`out`
-            empty!(stmt.args)
-            push!(stmt.args, callee_daef_ci, callee_captures, vars,
-                out, du, u, out_indices, du_indices, u_indices, t)
+
+            callee_residuals = equation_to_residual_mapping[callee_to_caller_eq_map[StructuralSSARef(old)]]
+            caller_variables = idnum.(info.mapping.var_coeffs)
+            callee_states = variable_to_state_mapping[caller_variables]
+
+            callee_daef_ci = rhs_finish_noopt!(callee_state, callee_ci, UnoptimizedKey(), world, settings, callee_residuals, callee_states)
+            call = @insert_instruction_here(compact, line, settings, (:invoke)(callee_daef_ci, callee_f,
+                out,
+                du,
+                u,
+                @insert_instruction_here(compact, line, settings, getindex(Int, callee_residuals...)::Vector{Int}),
+                @insert_instruction_here(compact, line, settings, getindex(Int, callee_states...)::Vector{Int}),
+                t)::type)
+            # TODO: add `stmt.args[3:end]`
+            replace_uses!(compact, (old, inst) => call)
         end
         type = inst[:type]
         if isa(type, Incidence) || isa(type, Eq)
@@ -75,9 +98,28 @@ function rhs_finish_noopt!(
     end
 
     daef_ci = rhs_finish_ir!(Compiler.finish(compact), ci, settings, key, slotnames)
+    # @sshow daef_ci.inferred
     return daef_ci
 end
 
+function map_variables_to_states(state::TransformationState)
+    (; structure) = state
+    diff_to_var = invview(structure.var_to_diff)
+    states = Int[]
+    prev_state = 0
+    for var in continuous_variables(state)
+        ref = is_differential_variable(structure, var) ? diff_to_var[var] : var
+        state = @something(get(states, ref, nothing), prev_state += 1)
+        push!(states, state)
+    end
+    return states
+end
+
+function replace_uses!(compact, ((old, inst), new))
+    inst[:stmt] = nothing
+    compact.ssa_rename[old] = new
+end
+
 function sciml_to_internal_abi_noopt!(ir::IRCode, state::TransformationState, internal_ci::CodeInstance, settings::Settings)
     slotnames = [:captures, :out, :du, :u, :p, :t]
     captures, out, du, u, p, t = Argument.(eachindex(slotnames))
@@ -90,14 +132,11 @@ function sciml_to_internal_abi_noopt!(ir::IRCode, state::TransformationState, in
     line = ir[SSAValue(1)][:line]
 
     internal_oc = @insert_instruction_here(compact, line, settings, getfield(captures, 1)::Core.OpaqueClosure)
-    # TODO: Compute proper indices.
     neqs = length(state.structure.eq_to_diff)
-    out_indices = @insert_instruction_here(compact, line, settings, view(out, 1:neqs)::VectorIntViewType)
-    du_indices = @insert_instruction_here(compact, line, settings, view(du, 1:neqs)::VectorIntViewType)
-    u_indices = @insert_instruction_here(compact, line, settings, view(u, 1:neqs)::VectorIntViewType)
-    # TODO: Provide actual external variables.
-    vars = @insert_instruction_here(compact, line, settings, getindex(Float64)::Vector{Float64})
-    @insert_instruction_here(compact, line, settings, (:invoke)(internal_ci, internal_oc, vars, out, du, u, out_indices, du_indices, u_indices, t)::Nothing)
+    nvars = length(state.structure.var_to_diff)
+    residuals = @insert_instruction_here(compact, line, settings, getindex(Int, 1:neqs...)::Vector{Int})
+    states = @insert_instruction_here(compact, line, settings, getindex(Int, map_variables_to_states(state)...)::Vector{Int})
+    @insert_instruction_here(compact, line, settings, (:invoke)(internal_ci, internal_oc, out, du, u, residuals, states, t)::Nothing)
     @insert_instruction_here(compact, line, settings, (return nothing)::Union{})
 
     ir = Compiler.finish(compact)

src/utils.jl

@@ -162,15 +162,19 @@ end
     @sshow stmt
     @sshow length(ir.stmts) typeof(val)
 
-Drop-in replacement for `@show`, but using `jl_safe_printf` to avoid task switches.
+Drop-in replacement for `@show`, but using `Core.println` to avoid task switches.
 
 This directly prints to C stdout; `stdout` redirects won't have any effect.
 """
 macro sshow(exs...)
     blk = Expr(:block)
     for ex in exs
-        push!(blk.args, :(Core.println($(sprint(Base.show_unquoted,ex)*" = "),
-                                  repr(begin local value = $(esc(ex)) end))))
+        push!(blk.args, quote
+            value = $(esc(ex))
+            Core.print($(sprint(Base.show_unquoted, ex)))
+            Core.print(" = ")
+            Core.println(sprint(print, value, context = :color => true))
+        end)
     end
     isempty(exs) || push!(blk.args, :value)
     return blk

test/validation.jl

@@ -49,15 +49,15 @@ end
 
     # IPO
 
+    u = [2.0]
+    du = [3.0]
     residuals, expanded_residuals = compute_residual_vectors(() -> onecall!(), u, du; t = 1.0)
     @test residuals ≈ [1.0]
     @test residuals ≈ expanded_residuals
 
     u = [2.0, 4.0]
     du = [3.0, 7.0]
-    # ERROR: BoundsError: attempt to access 2-element Vector{Float64} at index [3]
-    # (for `var = 3`)
-    refresh(); residuals, expanded_residuals = compute_residual_vectors(twocall!, u, du; t = 1.0)
-    @test residuals ≈ [1.0]
+    residuals, expanded_residuals = compute_residual_vectors(twocall!, u, du; t = 1.0)
+    @test residuals ≈ [1.0, 3.0]
     @test residuals ≈ expanded_residuals
 end;