Skip flattening in unoptimized context, WIP on IPO support

Author: serenity4

src/analysis/structural.jl

@@ -357,18 +357,19 @@ function _structural_analysis!(ci::CodeInstance, world::UInt, settings::Settings
             opaque_eligible = false
         end
 
-        # Rewrite to flattened ABI
-        compact[SSAValue(i)] = nothing
-        compact.result_idx -= 1
-        new_args = _flatten_parameter!(Compiler.optimizer_lattice(refiner), compact, callee_codeinst.inferred.ir.argtypes, arg->stmt.args[arg+1], line, settings)
-
-        new_call = insert_instruction_here!(compact, settings, @__SOURCE__,
-                NewInstruction(Expr(:invoke, (StructuralSSARef(compact.result_idx), callee_codeinst), new_args...), stmtype, info, line, stmtflags))
-        compact.ssa_rename[compact.idx - 1] = new_call
+        if !settings.skip_optimizations
+            # Rewrite to flattened ABI
+            compact[SSAValue(i)] = nothing
+            compact.result_idx -= 1
+            new_args = _flatten_parameter!(Compiler.optimizer_lattice(refiner), compact, callee_codeinst.inferred.ir.argtypes, arg->stmt.args[arg+1], line, settings)
+            new_call = insert_instruction_here!(compact, settings, @__SOURCE__,
+            NewInstruction(Expr(:invoke, (StructuralSSARef(compact.result_idx), callee_codeinst), new_args...), stmtype, info, line, stmtflags))
+            compact.ssa_rename[compact.idx - 1] = new_call
+        end
 
         cms = CallerMappingState(result, refiner.var_to_diff, refiner.varclassification, refiner.varkinds, eqclassification, eqkinds)
-        err = add_internal_equations_to_structure!(refiner, cms, total_incidence, eq_callee_mapping, StructuralSSARef(new_call.id),
-            result, mapping)
+        err = add_internal_equations_to_structure!(refiner, cms, total_incidence, eq_callee_mapping,
+            StructuralSSARef(i), result, mapping)
         if err !== true
             return UncompilableIPOResult(warnings, UnsupportedIRException(err, ir))
         end

src/transform/reconstruct.jl

@@ -53,7 +53,7 @@ end
 function expand_residuals(f, residuals, u, du, t)
     result = @code_structure result=true f()
     structure = make_structure_from_ipo(result)
-    state = TransformationState(result, structure, copy(result.total_incidence))
+    state = TransformationState(result, structure)
     key, _ = top_level_state_selection!(state)
     return expand_residuals(state, key, residuals, u, du, t)
 end

src/transform/unoptimized.jl

@@ -18,7 +18,8 @@ function rhs_finish_noopt!(
     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]
     append!(empty!(ir.argtypes), argtypes)
-    captures, vars, out, du, u, out_indices, du_indics, u_indices, t = Argument.(eachindex(slotnames))
+    captures, vars, out, du, u, out_indices, du_indices, u_indices, t = Argument.(eachindex(slotnames))
+    # TODO: use `out_indices`, `du_indices`, `u_indices`
     @assert length(slotnames) == length(ir.argtypes)
 
     equations = Pair{SSAValue, Eq}[]
@@ -45,17 +46,27 @@ function rhs_finish_noopt!(
             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, variable, compact)
+        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
-        elseif is_known_invoke_or_call(stmt, sim_time, compact)
+        elseif is_known_invoke_or_call(stmt, Intrinsics.sim_time, compact)
             inst[:stmt] = t
-        # TODO: process flattened variables
-        # TODO: process other intrinsics (epsilon, etc)
-        # else
-        #     replace_if_intrinsic!(compact, settings, ssaidx, nothing, nothing, nothing, t, nothing)
+        elseif is_known_invoke_or_call(stmt, Intrinsics.epsilon, compact)
+            inst[:stmt] = 0.0
+        elseif isexpr(stmt, :invoke)
+            @sshow stmt
+            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)
         end
         type = inst[:type]
         if isa(type, Incidence) || isa(type, Eq)

test/validation.jl

@@ -11,7 +11,7 @@ const *ᵢ = Core.Intrinsics.mul_float
 const +ᵢ = Core.Intrinsics.add_float
 const -ᵢ = Core.Intrinsics.sub_float
 
-@noinline function f()
+function f()
     x₁ = continuous() # selected
     x₂ = continuous() # selected
     x₃ = continuous() # algebraic, optimized away
@@ -22,12 +22,42 @@ const -ᵢ = Core.Intrinsics.sub_float
     always!(x₄ *ᵢ x₄ -ᵢ ddt(x₁))
 end
 
+@noinline function onecall!()
+    x = continuous()
+    always!(ddt(x) - x)
+end
+
+function twocall!()
+    onecall!(); onecall!();
+    return nothing
+end
+
 @testset "Validation" begin
     refresh() # TODO: remove before merge
-    f()
+
+    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 = [3.0, 1.0, 100.0, 4.0]
     du = [3.0, 0.0, 0.0, 0.0]
     residuals, expanded_residuals = compute_residual_vectors(f, u, du; t = 1.0)
     @test residuals ≈ [0.0, -3.0, 97.0, 13.0]
     @test residuals ≈ expanded_residuals
+
+    # IPO
+
+    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]
+    @test residuals ≈ expanded_residuals
 end;