Record factory as an edge for all processed CodeInstances

Author: serenity4

src/analysis/ADAnalyzer.jl

@@ -88,12 +88,23 @@ end
     error(lazy"Could not find single target method for `$sig`")
 end
 
-function ad_typeinf(world, tt; force_inline_all=false)
-    @assert !force_inline_all
-    interp = ADAnalyzer(;world)
+function get_method_instance(@nospecialize(tt), world)
     match = Base._methods_by_ftype(tt, 1, world)
     isempty(match) && single_match_error(tt)
     match = only(match)
     mi = Compiler.specialize_method(match)
+end
+
+function ad_typeinf(world, tt; force_inline_all=false, edges=nothing)
+    @assert !force_inline_all
+    interp = ADAnalyzer(;world)
+    mi = get_method_instance(tt, world)
     ci = Compiler.typeinf_ext(interp, mi, Compiler.SOURCE_MODE_ABI)
+    if edges !== nothing
+        prev = @atomic ci.edges
+        # XXX: Should we return the extended edges and use them in the other CodeInstances?
+        @atomic ci.edges = Core.svec(prev..., edges...)
+        Compiler.store_backedges(ci, edges)
+    end
+    ci
 end

src/analysis/refiner.jl

@@ -49,7 +49,9 @@ Compiler.cache_owner(::StructuralRefiner) = StructureCache()
     end
 
     callee_codeinst = invokee
-    callee_result = structural_analysis!(callee_codeinst, Compiler.get_inference_world(interp))
+    # XXX: We should propagate edges down here too.
+    edges = Core.svec()
+    callee_result = structural_analysis!(callee_codeinst, Compiler.get_inference_world(interp), edges)
 
     if isa(callee_result, UncompilableIPOResult) || isa(callee_result.extended_rt, Const) || isa(callee_result.extended_rt, Type)
         # If this is uncompilable, we will be notfiying the user in the outer loop - here we just ignore it
@@ -375,4 +377,4 @@ is_any_incidence(@nospecialize args...) = any(@nospecialize(x)->isa(x, Incidence
         f::Any, bargtypes::Vector{Any}, sv::Union{AbsIntState,Nothing})
 
     return rt
-end
+end

src/analysis/structural.jl

@@ -16,21 +16,21 @@ function find_matching_ci(predicate, mi::MethodInstance, world::UInt)
     return nothing
 end
 
-function structural_analysis!(ci::CodeInstance, world::UInt)
+function structural_analysis!(ci::CodeInstance, world::UInt, edges::SimpleVector)
     # Check if we have aleady done this work - if so return the cached result
     result_ci = find_matching_ci(ci->ci.owner == StructureCache(), ci.def, world)
     if result_ci !== nothing
         return result_ci.inferred
     end
 
-    result = _structural_analysis!(ci, world)
+    result = _structural_analysis!(ci, world, edges)
     # TODO: The world bounds might have been narrowed
-    cache_dae_ci!(ci, result, nothing, nothing, StructureCache())
+    cache_dae_ci!(ci, result, nothing, nothing, StructureCache(), edges)
 
     return result
 end
 
-function _structural_analysis!(ci::CodeInstance, world::UInt)
+function _structural_analysis!(ci::CodeInstance, world::UInt, edges::SimpleVector)
     # Variables
     var_to_diff = DiffGraph(0)
     varclassification = VarEqClassification[]
@@ -291,7 +291,7 @@ function _structural_analysis!(ci::CodeInstance, world::UInt)
             (; result, mapping) = info
         else
             callee_codeinst = stmt.args[1]
-            result = structural_analysis!(callee_codeinst, Compiler.get_inference_world(refiner))
+            result = structural_analysis!(callee_codeinst, Compiler.get_inference_world(refiner), edges)
 
             if isa(result, UncompilableIPOResult)
                 # TODO: Stack trace?

src/interface.jl

@@ -22,13 +22,15 @@ This is the compile-time entry point for DAECompiler code generation. It drives
 """
 function factory_gen(world::UInt, source::Method, @nospecialize(_gen), settings, @nospecialize(fT))
     settings = settings.parameters[1]
+    factory_mi = get_method_instance(Tuple{typeof(factory),Val{settings},typeof(fT)}, world)
+    edges = Core.svec(factory_mi)
 
     # First, perform ordinary type inference, under the assumption that we may need to AD
     # parts of the function later.
-    ci = ad_typeinf(world, Tuple{fT}; force_inline_all=settings.force_inline_all)
+    ci = ad_typeinf(world, Tuple{fT}; force_inline_all=settings.force_inline_all, edges)
 
     # Perform or lookup DAECompiler specific analysis for this system.
-    result = structural_analysis!(ci, world)
+    result = structural_analysis!(ci, world, edges)
 
     if isa(result, UncompilableIPOResult)
         return Base.generated_body_to_codeinfo(
@@ -42,19 +44,19 @@ function factory_gen(world::UInt, source::Method, @nospecialize(_gen), settings,
     (diff_key, init_key) = top_level_state_selection!(tstate)
 
     if settings.mode in (DAE, DAENoInit, ODE, ODENoInit)
-        tearing_schedule!(tstate, ci, diff_key, world)
+        tearing_schedule!(tstate, ci, diff_key, world, edges)
     end
     if settings.mode in (InitUncompress, DAE, ODE)
-        tearing_schedule!(tstate, ci, init_key, world)
+        tearing_schedule!(tstate, ci, init_key, world, edges)
     end
 
     # Generate the IR implementation of `factory`, returning the DAEFunction/ODEFunction
     if settings.mode in (DAE, DAENoInit)
-        ir_factory = dae_factory_gen(tstate, ci, diff_key, world, settings.mode == DAE ? init_key : nothing)
+        ir_factory = dae_factory_gen(tstate, ci, diff_key, world, edges, settings.mode == DAE ? init_key : nothing)
     elseif settings.mode in (ODE, ODENoInit)
-        ir_factory = ode_factory_gen(tstate, ci, diff_key, world, settings.mode == ODE ? init_key : nothing)
+        ir_factory = ode_factory_gen(tstate, ci, diff_key, world, edges, settings.mode == ODE ? init_key : nothing)
     elseif settings.mode == InitUncompress
-        ir_factory = init_uncompress_gen(result, ci, init_key, diff_key, world)
+        ir_factory = init_uncompress_gen(result, ci, init_key, diff_key, world, edges)
     else
         return :(error("Unknown generation mode: $(settings.mode)"))
     end

src/transform/codegen/dae_factory.jl

@@ -60,7 +60,7 @@ end
 ```
 
 """
-function dae_factory_gen(state::TransformationState, ci::CodeInstance, key::TornCacheKey, world::UInt, init_key::Union{TornCacheKey, Nothing})
+function dae_factory_gen(state::TransformationState, ci::CodeInstance, key::TornCacheKey, world::UInt, edges::SimpleVector, 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
@@ -88,7 +88,7 @@ function dae_factory_gen(state::TransformationState, ci::CodeInstance, key::Torn
 
     argt = Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}, SciMLBase.NullParameters, Float64}
 
-    daef_ci = rhs_finish!(state, ci, key, world, 1)
+    daef_ci = rhs_finish!(state, ci, key, world, 1, edges)
 
     # Create a small opaque closure to adapt from SciML ABI to our own internal
     # ABI
@@ -188,7 +188,7 @@ function dae_factory_gen(state::TransformationState, ci::CodeInstance, key::Torn
     differential_states = Bool[v in key.diff_states for v in all_states]
 
     if init_key !== nothing
-        initf = init_uncompress_gen!(compact, result, ci, init_key, key, world)
+        initf = init_uncompress_gen!(compact, result, ci, init_key, key, world, edges)
         daef = insert_node_here!(compact, NewInstruction(Expr(:call, make_daefunction, new_oc, initf),
             DAEFunction, line), true)
     else

src/transform/codegen/init_factory.jl

@@ -1,19 +1,19 @@
 
-function init_uncompress_gen(result::DAEIPOResult, ci::CodeInstance, init_key::TornCacheKey, diff_key::TornCacheKey, world::UInt)
+function init_uncompress_gen(result::DAEIPOResult, ci::CodeInstance, init_key::TornCacheKey, diff_key::TornCacheKey, world::UInt, edges::SimpleVector)
     ir_factory = copy(result.ir)
     pushfirst!(ir_factory.argtypes, Settings)
     pushfirst!(ir_factory.argtypes, typeof(factory))
     compact = IncrementalCompact(ir_factory)
 
-    new_oc = init_uncompress_gen!(compact, result, ci, init_key, diff_key, world)
+    new_oc = init_uncompress_gen!(compact, result, ci, init_key, diff_key, world, edges)
     insert_node_here!(compact, NewInstruction(ReturnNode(new_oc), Core.OpaqueClosure, result.ir[SSAValue(1)][:line]), true)
 
     ir_factory = Compiler.finish(compact)
 
     return ir_factory
 end
 
-function init_uncompress_gen!(compact::Compiler.IncrementalCompact, result::DAEIPOResult, ci::CodeInstance, init_key::TornCacheKey, diff_key::TornCacheKey, world::UInt)
+function init_uncompress_gen!(compact::Compiler.IncrementalCompact, result::DAEIPOResult, ci::CodeInstance, init_key::TornCacheKey, diff_key::TornCacheKey, world::UInt, edges::SimpleVector)
     torn_ci = find_matching_ci(ci->isa(ci.owner, TornIRSpec) && ci.owner.key == init_key, ci.def, world)
     @assert torn_ci !== nothing
     torn_ir = torn_ci.inferred
@@ -35,7 +35,7 @@ function init_uncompress_gen!(compact::Compiler.IncrementalCompact, result::DAEI
 
     # (nlsol,)
     argt = Tuple{Any}
-    daef_ci = gen_init_uncompress!(result, ci, init_key, diff_key, world, 1)
+    daef_ci = gen_init_uncompress!(result, ci, init_key, diff_key, world, 1, edges)
 
     # Create a small opaque closure to adapt from SciML ABI to our own internal
     # ABI
@@ -102,4 +102,4 @@ function init_uncompress_gen!(compact::Compiler.IncrementalCompact, result::DAEI
         argt, Vector{Float64}, Vector{Float64}, true, oc_source_method, sicm), Core.OpaqueClosure, line), true)
 
     return new_oc
-end
+end

src/transform/codegen/init_uncompress.jl

@@ -9,10 +9,10 @@ struct InitUncompressSpec
     ordinal::Int
 end
 
-function gen_init_uncompress!(result::DAEIPOResult, ci::CodeInstance, init_key::TornCacheKey, diff_key::TornCacheKey, world::UInt, ordinal::Int, indexT=Int)
+function gen_init_uncompress!(result::DAEIPOResult, ci::CodeInstance, init_key::TornCacheKey, diff_key::TornCacheKey, world::UInt, ordinal::Int, edges::SimpleVector, indexT=Int)
     structure = make_structure_from_ipo(result)
     tstate = TransformationState(result, structure, copy(result.total_incidence))
-    return gen_init_uncompress!(tstate, ci, init_key, diff_key, world, ordinal, indexT)
+    return gen_init_uncompress!(tstate, ci, init_key, diff_key, world, ordinal, edges, indexT)
 end
 
 function gen_init_uncompress!(
@@ -22,6 +22,7 @@ function gen_init_uncompress!(
         diff_key::TornCacheKey,
         world::UInt,
         ordinal::Int,
+        edges::SimpleVector,
         indexT=Int)
 
     (; result, structure) = state
@@ -97,8 +98,8 @@ function gen_init_uncompress!(
                 spec_data = stmt.args[1]
                 callee_key = stmt.args[1][2]
                 callee_ordinal = stmt.args[1][end]::Int
-                callee_result = structural_analysis!(callee_ci, world)
-                callee_daef_ci = rhs_finish!(callee_result, callee_ci, callee_key, world, callee_ordinal)
+                callee_result = structural_analysis!(callee_ci, world, edges)
+                callee_daef_ci = rhs_finish!(callee_result, callee_ci, callee_key, world, callee_ordinal, edges)
                 # Allocate a continuous block of variables for all callee alg and diff states
 
                 empty!(stmt.args)
@@ -154,7 +155,8 @@ function gen_init_uncompress!(
         src = ir_to_src(ir)
 
         abi = Tuple{Tuple, Tuple, (VectorViewType for _ in arg_range)..., Vector{Float64}, Float64}
-        daef_ci = cache_dae_ci!(ci, src, src.debuginfo, abi, InitUncompressSpec(init_key, diff_key, ir_ordinal))
+        daef_ci = cache_dae_ci!(ci, src, src.debuginfo, abi, InitUncompressSpec(init_key, diff_key, ir_ordinal), edges)
+
         ccall(:jl_add_codeinst_to_jit, Cvoid, (Any, Any), daef_ci, src)
 
         push!(cis, daef_ci)

src/transform/codegen/ode_factory.jl

@@ -52,7 +52,7 @@ end
 ```
 
 """
-function ode_factory_gen(state::TransformationState, ci::CodeInstance, key::TornCacheKey, world::UInt, init_key::Union{TornCacheKey, Nothing})
+function ode_factory_gen(state::TransformationState, ci::CodeInstance, key::TornCacheKey, world::UInt, edges::SimpleVector, 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
@@ -75,7 +75,7 @@ function ode_factory_gen(state::TransformationState, ci::CodeInstance, key::Torn
         sicm = ()
     end
 
-    odef_ci = rhs_finish!(state, ci, key, world, 1)
+    odef_ci = rhs_finish!(state, ci, key, world, 1, edges)
 
     # Create a small opaque closure to adapt from SciML ABI to our own internal ABI
 
@@ -140,7 +140,7 @@ function ode_factory_gen(state::TransformationState, ci::CodeInstance, key::Torn
     nd = numstates[AssignedDiff] + numstates[UnassignedDiff]
     na = numstates[Algebraic] + numstates[AlgebraicDerivative]
     mass_matrix = na == 0 ? GlobalRef(LinearAlgebra, :I) : @insert_node_here compact line generate_ode_mass_matrix(nd, na)::Matrix{Float64}
-    initf = init_key !== nothing ? init_uncompress_gen!(compact, result, ci, init_key, key, world) : nothing
+    initf = init_key !== nothing ? init_uncompress_gen!(compact, result, ci, init_key, key, world, edges) : nothing
     odef = @insert_node_here compact line make_odefunction(new_oc, mass_matrix, initf)::ODEFunction true
 
     odef_and_n = @insert_node_here compact line tuple(odef, nd + na)::Tuple true

src/transform/codegen/rhs.jl

@@ -63,10 +63,10 @@ function compute_slot_ranges(info::MappingInfo, callee_key, var_assignment, eq_a
     return state_ranges
 end
 
-function rhs_finish!(result::DAEIPOResult, ci::CodeInstance, key::TornCacheKey, world::UInt, ordinal::Int, indexT=Int)
+function rhs_finish!(result::DAEIPOResult, ci::CodeInstance, key::TornCacheKey, world::UInt, ordinal::Int, edges::SimpleVector, indexT=Int)
     structure = make_structure_from_ipo(result)
     tstate = TransformationState(result, structure, copy(result.total_incidence))
-    return rhs_finish!(tstate, ci, key, world, ordinal, indexT)
+    return rhs_finish!(tstate, ci, key, world, ordinal, edges, indexT)
 end
 
 function rhs_finish!(
@@ -75,6 +75,7 @@ function rhs_finish!(
     key::TornCacheKey,
     world::UInt,
     ordinal::Int,
+    edges::SimpleVector,
     indexT=Int)
 
     (; result, structure) = state
@@ -145,8 +146,8 @@ function rhs_finish!(
                 spec_data = stmt.args[1]
                 callee_key = spec_data[2]
                 callee_ordinal = spec_data[end]::Int
-                callee_result = structural_analysis!(callee_ci, world)
-                callee_daef_ci = rhs_finish!(callee_result, callee_ci, callee_key, world, callee_ordinal)
+                callee_result = structural_analysis!(callee_ci, world, edges)
+                callee_daef_ci = rhs_finish!(callee_result, callee_ci, callee_key, world, callee_ordinal, edges)
                 # Allocate a continuous block of variables for all callee alg and diff states
 
                 empty!(stmt.args)
@@ -218,7 +219,7 @@ function rhs_finish!(
         src = ir_to_src(ir)
 
         abi = Tuple{Tuple, Tuple, (VectorViewType for _ in arg_range)..., Float64}
-        daef_ci = cache_dae_ci!(ci, src, src.debuginfo, abi, RHSSpec(key, ir_ordinal))
+        daef_ci = cache_dae_ci!(ci, src, src.debuginfo, abi, RHSSpec(key, ir_ordinal), edges)
         ccall(:jl_add_codeinst_to_jit, Cvoid, (Any, Any), daef_ci, src)
 
         push!(cis, daef_ci)

src/transform/common.jl

@@ -63,10 +63,11 @@ function ir_to_src(ir::IRCode)
     return src
 end
 
-function cache_dae_ci!(old_ci, src, debuginfo, abi, owner)
+function cache_dae_ci!(old_ci, src, debuginfo, abi, owner, edges::SimpleVector)
     daef_ci = CodeInstance(abi === nothing ? old_ci.def : Core.ABIOverride(abi, old_ci.def), owner, Tuple, Union{}, nothing, src, Int32(0),
         UInt(1)#=ci.min_world=#, old_ci.max_world, old_ci.ipo_purity_bits,
-        nothing, debuginfo, Compiler.empty_edges)
+        nothing, debuginfo, edges)
+    Compiler.store_backedges(daef_ci, edges)
     ccall(:jl_mi_cache_insert, Cvoid, (Any, Any), old_ci.def, daef_ci)
     return daef_ci
 end