[WIP] Add unoptimized compilation path

Author: serenity4

src/DAECompiler.jl

@@ -28,6 +28,7 @@ module DAECompiler
     include("transform/state_selection.jl")
     include("transform/common.jl")
     include("transform/runtime.jl")
+    include("transform/unoptimized.jl")
     include("transform/tearing/schedule.jl")
     include("transform/codegen/dae_factory.jl")
     include("transform/codegen/ode_factory.jl")

src/interface.jl

@@ -41,20 +41,23 @@ function factory_gen(@nospecialize(fT), settings::Settings, world::UInt = Base.g
     end
 
     # Select differential and algebraic states
-    ret = top_level_state_selection!(tstate)
-
-    if isa(ret, UncompilableIPOResult)
-        return Base.generated_body_to_codeinfo(
-            Expr(:lambda, Any[:var"#self", :settings, :f], Expr(:block, Expr(:return, Expr(:call, throw, ret.error)))),
-            @__MODULE__, false)
-    end
-    (diff_key, init_key) = ret
-
-    if settings.mode in (DAE, DAENoInit, ODE, ODENoInit)
-        tearing_schedule!(tstate, ci, diff_key, world, settings)
-    end
-    if settings.mode in (InitUncompress, DAE, ODE)
-        tearing_schedule!(tstate, ci, init_key, world, settings)
+    if settings.skip_optimizations
+        diff_key = torn_cache_key(tstate, settings)
+        init_key = nothing
+    else
+        ret = top_level_state_selection!(tstate)
+        if isa(ret, UncompilableIPOResult)
+            return Base.generated_body_to_codeinfo(
+                Expr(:lambda, Any[:var"#self", :settings, :f], Expr(:block, Expr(:return, Expr(:call, throw, ret.error)))),
+                @__MODULE__, false)
+        end
+        (diff_key, init_key) = ret
+        if settings.mode in (DAE, DAENoInit, ODE, ODENoInit)
+            tearing_schedule!(tstate, ci, diff_key, world, settings)
+        end
+        if settings.mode in (InitUncompress, DAE, ODE)
+            tearing_schedule!(tstate, ci, init_key, world, settings)
+        end
     end
 
     # Generate the IR implementation of `factory`, returning the DAEFunction/ODEFunction

src/problem_interface.jl

@@ -25,8 +25,9 @@ function DAECProblem(f, init::Union{Vector, Tuple{Vararg{Pair}}}, tspan::Tuple{R
                      force_inline_all=false,
                      insert_stmt_debuginfo=false,
                      insert_ssa_debuginfo=false,
+                     skip_optimizations=false,
                      kwargs...)
-    settings = Settings(; force_inline_all, insert_stmt_debuginfo, insert_ssa_debuginfo)
+    settings = Settings(; force_inline_all, insert_stmt_debuginfo, insert_ssa_debuginfo, skip_optimizations)
     DAECProblem(f, init, guesses, tspan, kwargs, settings, missing, nothing, nothing)
 end
 
@@ -35,13 +36,14 @@ function DAECProblem(f, tspan::Tuple{Real, Real} = (0., 1.);
                      force_inline_all=false,
                      insert_stmt_debuginfo=false,
                      insert_ssa_debuginfo=false,
+                     skip_optimizations=false,
                      kwargs...)
-    settings = Settings(; force_inline_all, insert_stmt_debuginfo, insert_ssa_debuginfo)
+    settings = Settings(; force_inline_all, insert_stmt_debuginfo, insert_ssa_debuginfo, skip_optimizations)
     DAECProblem(f, nothing, guesses, tspan, kwargs, settings, missing, nothing, nothing)
 end
 
 function DiffEqBase.get_concrete_problem(prob::DAECProblem, isadaptive; kwargs...)
-    settings = Settings(; mode=prob.init === nothing ? DAE : DAENoInit, prob.settings.force_inline_all, prob.settings.insert_stmt_debuginfo, prob.settings.insert_ssa_debuginfo)
+    settings = Settings(; mode=prob.init === nothing ? DAE : DAENoInit, prob.settings.force_inline_all, prob.settings.insert_stmt_debuginfo, prob.settings.insert_ssa_debuginfo, prob.settings.skip_optimizations)
     (daef, differential_vars) = factory(Val(settings), prob.f)
 
     u0 = zeros(length(differential_vars))
@@ -77,8 +79,9 @@ function ODECProblem(f, init::Union{Vector, Tuple{Vararg{Pair}}}, tspan::Tuple{R
                      force_inline_all=false,
                      insert_stmt_debuginfo=false,
                      insert_ssa_debuginfo=false,
+                     skip_optimizations=false,
                      kwargs...)
-    settings = Settings(; force_inline_all, insert_stmt_debuginfo, insert_ssa_debuginfo)
+    settings = Settings(; force_inline_all, insert_stmt_debuginfo, insert_ssa_debuginfo, skip_optimizations)
     ODECProblem(f, init, guesses, tspan, kwargs, settings, missing, nothing)
 end
 
@@ -87,13 +90,14 @@ function ODECProblem(f, tspan::Tuple{Real, Real} = (0., 1.);
                      force_inline_all=false,
                      insert_stmt_debuginfo=false,
                      insert_ssa_debuginfo=false,
+                     skip_optimizations=false,
                      kwargs...)
-    settings = Settings(; force_inline_all, insert_stmt_debuginfo, insert_ssa_debuginfo)
+    settings = Settings(; force_inline_all, insert_stmt_debuginfo, insert_ssa_debuginfo, skip_optimizations)
     ODECProblem(f, nothing, guesses, tspan, kwargs, settings, missing, nothing)
 end
 
 function DiffEqBase.get_concrete_problem(prob::ODECProblem, isadaptive; kwargs...)
-    settings = Settings(; mode=prob.init === nothing ? ODE : ODENoInit, prob.settings.force_inline_all, prob.settings.insert_stmt_debuginfo, prob.settings.insert_ssa_debuginfo)
+    settings = Settings(; mode=prob.init === nothing ? ODE : ODENoInit, prob.settings.force_inline_all, prob.settings.insert_stmt_debuginfo, prob.settings.insert_ssa_debuginfo, prob.settings.skip_optimizations)
     (odef, n) = factory(Val(settings), prob.f)
 
     u0 = zeros(n)

src/reflection.jl

@@ -28,7 +28,7 @@ code_ad_by_type(@nospecialize(tt::Type); kwargs...) =
 
 function code_structure_by_type(@nospecialize(tt::Type); world::UInt = Base.tls_world_age(), result = false, matched = false, mode = DAE, force_inline_all = false, insert_stmt_debuginfo = false, insert_ssa_debuginfo = false, kwargs...)
   ci = _code_ad_by_type(tt; world, kwargs...)
-  settings = Settings(; mode, force_inline_all, insert_stmt_debuginfo, insert_ssa_debuginfo)
+  settings = Settings(; mode, force_inline_all, insert_stmt_debuginfo, insert_ssa_debuginfo, skip_optimizations)
   _result = structural_analysis!(ci, world, settings)
   isa(_result, UncompilableIPOResult) && throw(_result.error)
   !matched && return result ? _result : _result.ir

src/settings.jl

@@ -14,5 +14,6 @@ struct Settings
     force_inline_all::Bool
     insert_stmt_debuginfo::Bool
     insert_ssa_debuginfo::Bool
+    skip_optimizations::Bool
 end
-Settings(; mode::GenerationMode=DAE, force_inline_all::Bool=false, insert_stmt_debuginfo::Bool=false, insert_ssa_debuginfo::Bool=false) = Settings(mode, force_inline_all, insert_stmt_debuginfo, insert_ssa_debuginfo)
+Settings(; mode::GenerationMode=DAE, force_inline_all::Bool=false, insert_stmt_debuginfo::Bool=false, insert_ssa_debuginfo::Bool=false, skip_optimizations::Bool = false) = Settings(mode, force_inline_all, insert_stmt_debuginfo, insert_ssa_debuginfo, skip_optimizations)

src/transform/codegen/dae_factory.jl

@@ -39,101 +39,51 @@ function make_daefunction(f, initf)
     DAEFunction(f; initialization_data = SciMLBase.OverrideInitData(NonlinearProblem((args...)->nothing, nothing, nothing), nothing, initf, nothing, nothing, Val{false}()))
 end
 
-"""
-    dae_factory_gen(ci, key)
-
-Generate the `factory` function for CodeInstance `ci`, returning a DAEFunction.
-The resulting function is roughly:
-
-```
-function factory(settings, f)
-    # Run all parts of `f` that do not depend on state
-    state_invariant_pieces = f_state_invariant()
-    f! = %new_opaque_closure(f_rhs, state_invariant_pieces)
-    DAEFunction(f!), differential_vars
+function continuous_variables(state::TransformationState)
+    filter(var -> varkind(state, var) == Intrinsics.Continuous, 1:length(state.result.var_to_diff))
 end
-```
-
-"""
-function dae_factory_gen(state::TransformationState, ci::CodeInstance, key::TornCacheKey, world::UInt, settings::Settings, init_key::Union{TornCacheKey, Nothing})
-    result = state.result
-    torn_ci = find_matching_ci(ci->isa(ci.owner, TornIRSpec) && ci.owner.key == key, ci.def, world)
-    torn_ir = torn_ci.inferred
-
-    (;ir_sicm) = torn_ir
 
-    ir_factory = copy(ci.inferred.ir)
-    pushfirst!(ir_factory.argtypes, Settings)
-    pushfirst!(ir_factory.argtypes, typeof(factory))
-    compact = IncrementalCompact(ir_factory)
-
-    local line
-    if ir_sicm !== nothing
-        sicm_ci = find_matching_ci(ci->isa(ci.owner, SICMSpec) && ci.owner.key == key, ci.def, world)
-        @assert sicm_ci !== nothing
-
-        line = result.ir[SSAValue(1)][:line]
-        param_list = flatten_parameter!(Compiler.fallback_lattice, compact, ci.inferred.ir.argtypes[1:end], argn->Argument(2+argn), line, settings)
-        sicm = @insert_instruction_here compact line settings invoke(param_list, sicm_ci)::Tuple
-    else
-        sicm = ()
-    end
-
-    argt = Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}, SciMLBase.NullParameters, Float64}
-
-    daef_ci = rhs_finish!(state, ci, key, world, settings, 1)
+const SCIML_ABI = Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}, SciMLBase.NullParameters, Float64}
 
-    # Create a small opaque closure to adapt from SciML ABI to our own internal
-    # ABI
+function sciml_to_internal_abi!(ir::IRCode, state::TransformationState, internal_ci::CodeInstance, key::TornCacheKey, var_eq_matching, settings::Settings)
+    (; result, structure) = state
 
     numstates = zeros(Int, Int(LastEquationStateKind))
-
-    all_states = Int[]
-    for var = 1:length(result.var_to_diff)
-        varkind(state, var) == Intrinsics.Continuous || continue
+    for var in continuous_variables(state)
         kind = classify_var(result.var_to_diff, key, var)
         kind == nothing && continue
         numstates[kind] += 1
-        (kind != AlgebraicDerivative) && push!(all_states, var)
     end
 
-    ir_oc = copy(ci.inferred.ir)
-    empty!(ir_oc.argtypes)
-    push!(ir_oc.argtypes, Tuple)
-    push!(ir_oc.argtypes, Vector{Float64})
-    push!(ir_oc.argtypes, Vector{Float64})
-    push!(ir_oc.argtypes, Vector{Float64})
-    push!(ir_oc.argtypes, SciMLBase.NullParameters)
-    push!(ir_oc.argtypes, Float64)
+    empty!(ir.argtypes)
+    push!(ir.argtypes, Tuple) # opaque closure captures
+    append!(ir.argtypes, fieldtypes(SCIML_ABI))
 
-    oc_compact = IncrementalCompact(ir_oc)
+    compact = IncrementalCompact(ir)
 
     # Zero the output
-    line = ir_oc[SSAValue(1)][:line]
-    @insert_instruction_here oc_compact line settings zero!(Argument(2))::VectorViewType
+    line = ir[SSAValue(1)][:line]
+    @insert_instruction_here compact line settings zero!(Argument(2))::VectorViewType
 
     # out_du_mm, out_eq, in_u_mm, in_u_unassgn, in_du_unassgn, in_alg
     nassgn = numstates[AssignedDiff]
     ntotalstates = numstates[AssignedDiff] + numstates[UnassignedDiff] + numstates[Algebraic]
-    out_du_mm = @insert_instruction_here oc_compact line settings view(Argument(2), 1:nassgn)::VectorViewType
-    out_eq = @insert_instruction_here oc_compact line settings view(Argument(2), (nassgn+1):ntotalstates)::VectorViewType
+    out_du_mm = @insert_instruction_here compact line settings view(Argument(2), 1:nassgn)::VectorViewType
+    out_eq = @insert_instruction_here compact line settings view(Argument(2), (nassgn+1):ntotalstates)::VectorViewType
 
-    (in_du_assgn, in_du_unassgn) = sciml_dae_split_du!(oc_compact, line, settings, Argument(3), numstates)
-    (in_u_mm, in_u_unassgn, in_alg) = sciml_dae_split_u!(oc_compact, line, settings, Argument(4), numstates)
+    (in_du_assgn, in_du_unassgn) = sciml_dae_split_du!(compact, line, settings, Argument(3), numstates)
+    (in_u_mm, in_u_unassgn, in_alg) = sciml_dae_split_u!(compact, line, settings, Argument(4), numstates)
 
     # Call DAECompiler-generated RHS with internal ABI
-    oc_sicm = @insert_instruction_here oc_compact line settings getfield(Argument(1), 1)::Core.OpaqueClosure
+    oc_sicm = @insert_instruction_here compact line settings getfield(Argument(1), 1)::Core.OpaqueClosure
 
     # N.B: The ordering of arguments should match the ordering in the StateKind enum
-    @insert_instruction_here oc_compact line settings (:invoke)(daef_ci, oc_sicm, (), in_u_mm, in_u_unassgn, in_du_unassgn, in_alg, out_du_mm, out_eq, Argument(6))::Nothing
-
-    # TODO: We should not have to recompute this here
-    var_eq_matching = matching_for_key(state, key)
-    (slot_assignments, var_assignment, eq_assignment) = assign_slots(state, key, var_eq_matching)
+    @insert_instruction_here compact line settings (:invoke)(internal_ci, oc_sicm, (), in_u_mm, in_u_unassgn, in_du_unassgn, in_alg, out_du_mm, out_eq, Argument(6))::Nothing
 
     # Manually apply mass matrix and implicit equations between selected states
-    for v = 1:ndsts(state.structure.graph)
-        vdiff = state.structure.var_to_diff[v]
+    (_, var_assignment, _) = assign_slots(state, key, var_eq_matching)
+    for v = 1:ndsts(structure.graph)
+        vdiff = structure.var_to_diff[v]
         vdiff === nothing && continue
 
         if var_eq_matching[v] !== SelectedState() || var_eq_matching[vdiff] !== SelectedState()
@@ -146,22 +96,81 @@ function dae_factory_gen(state::TransformationState, ci::CodeInstance, key::Torn
         @assert kind == AssignedDiff
         @assert dkind in (AssignedDiff, UnassignedDiff)
 
-        v_val = @insert_instruction_here oc_compact line settings getindex(dkind == AssignedDiff ? in_u_mm : in_u_unassgn, dslot)::Any
-        @insert_instruction_here oc_compact line settings setindex!(out_du_mm, v_val, slot)::Any
+        v_val = @insert_instruction_here compact line settings getindex(dkind == AssignedDiff ? in_u_mm : in_u_unassgn, dslot)::Any
+        @insert_instruction_here compact line settings setindex!(out_du_mm, v_val, slot)::Any
     end
 
-    bc = @insert_instruction_here oc_compact line settings Base.Broadcast.broadcasted(-, out_du_mm, in_du_assgn)::Any
-    @insert_instruction_here oc_compact line settings Base.Broadcast.materialize!(out_du_mm, bc)::Nothing
+    bc = @insert_instruction_here compact line settings Base.Broadcast.broadcasted(-, out_du_mm, in_du_assgn)::Any
+    @insert_instruction_here compact line settings Base.Broadcast.materialize!(out_du_mm, bc)::Nothing
 
     # Return
-    @insert_instruction_here oc_compact line settings (return nothing)::Union{}
+    @insert_instruction_here compact line settings (return nothing)::Union{}
 
-    ir_oc = Compiler.finish(oc_compact)
-    maybe_rewrite_debuginfo!(ir_oc, settings)
-    resize!(ir_oc.cfg.blocks, 1)
-    empty!(ir_oc.cfg.blocks[1].succs)
-    Compiler.verify_ir(ir_oc)
-    oc = Core.OpaqueClosure(ir_oc)
+    ir = Compiler.finish(compact)
+    maybe_rewrite_debuginfo!(ir, settings)
+    resize!(ir.cfg.blocks, 1)
+    empty!(ir.cfg.blocks[1].succs)
+    Compiler.verify_ir(ir)
+
+    @async @eval Main begin
+        interface_ir = $ir
+    end
+
+    return Core.OpaqueClosure(ir; slotnames = [:captures, :out, :du, :u, :p, :t])
+end
+
+"""
+    dae_factory_gen(ci, key)
+
+Generate the `factory` function for CodeInstance `ci`, returning a DAEFunction.
+The resulting function is roughly:
+
+```
+function factory(settings, f)
+    # Run all parts of `f` that do not depend on state
+    state_invariant_pieces = f_state_invariant()
+    f! = %new_opaque_closure(f_rhs, state_invariant_pieces)
+    DAEFunction(f!), differential_vars
+end
+```
+
+"""
+function dae_factory_gen(state::TransformationState, ci::CodeInstance, key::TornCacheKey, world::UInt, settings::Settings, init_key::Union{TornCacheKey, Nothing})
+    result = state.result
+    # TODO: We should not have to recompute this here
+
+    ir_factory = copy(ci.inferred.ir)
+    pushfirst!(ir_factory.argtypes, Settings)
+    pushfirst!(ir_factory.argtypes, typeof(factory))
+    compact = IncrementalCompact(ir_factory)
+
+    # Create a small opaque closure to adapt from SciML ABI to our own internal ABI
+    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)
+        oc = sciml_to_internal_abi_noopt!(copy(ci.inferred.ir), state, daef_ci, settings)
+    else
+        var_eq_matching = matching_for_key(state, key)
+
+        torn_ci = find_matching_ci(ci->isa(ci.owner, TornIRSpec) && ci.owner.key == key, ci.def, world)
+        torn_ir = torn_ci.inferred
+
+        (; ir_sicm) = torn_ir
+
+        local line
+        if ir_sicm !== nothing
+            sicm_ci = find_matching_ci(ci->isa(ci.owner, SICMSpec) && ci.owner.key == key, ci.def, world)
+            @assert sicm_ci !== nothing
+
+            line = result.ir[SSAValue(1)][:line]
+            param_list = flatten_parameter!(Compiler.fallback_lattice, compact, ci.inferred.ir.argtypes[1:end], argn->Argument(2+argn), line, settings)
+            sicm = @insert_instruction_here compact line settings invoke(param_list, sicm_ci)::Tuple
+        end
+
+        daef_ci = rhs_finish!(state, ci, key, world, settings, 1)
+        oc = sciml_to_internal_abi!(copy(ci.inferred.ir), state, daef_ci, key, var_eq_matching, settings)
+    end
 
     line = result.ir[SSAValue(1)][:line]
 
@@ -173,6 +182,7 @@ function dae_factory_gen(state::TransformationState, ci::CodeInstance, key::Torn
 
     new_oc = @insert_instruction_here compact line settings (:new_opaque_closure)(argt, Union{}, Nothing, true, oc_source_method, sicm)::Core.OpaqueClosure true
 
+    all_states = filter(var -> classify_var(result, key, var) != AlgebraicDerivative, continuous_variables(state))
     differential_states = Bool[v in key.diff_states for v in all_states]
 
     if init_key !== nothing
@@ -192,6 +202,6 @@ function dae_factory_gen(state::TransformationState, ci::CodeInstance, key::Torn
     empty!(ir_factory.cfg.blocks[1].succs)
     Compiler.verify_ir(ir_factory)
 
-    slotnames = [[:factory, :settings]; Symbol.(:arg, 1:(length(ir_factory.argtypes) - 2))]
+    slotnames = [:factory, :settings, :f]
     return ir_factory, slotnames
 end