Add invalidation edges transitively, don't pass them around

Author: serenity4

src/analysis/refiner.jl

@@ -10,7 +10,6 @@ struct StructuralRefiner <: Compiler.AbstractInterpreter
     var_to_diff::DiffGraph
     varkinds::Vector{Intrinsics.VarKind}
     varclassification::Vector{VarEqClassification}
-    edges::SimpleVector
 end
 
 struct StructureCache; end
@@ -50,7 +49,7 @@ Compiler.cache_owner(::StructuralRefiner) = StructureCache()
     end
 
     callee_codeinst = invokee
-    callee_result = structural_analysis!(callee_codeinst, Compiler.get_inference_world(interp), interp.edges)
+    callee_result = structural_analysis!(callee_codeinst, Compiler.get_inference_world(interp))
 
     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

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, edges::SimpleVector)
+function structural_analysis!(ci::CodeInstance, world::UInt)
     # 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, edges)
+    result = _structural_analysis!(ci, world)
     # TODO: The world bounds might have been narrowed
-    cache_dae_ci!(ci, result, nothing, nothing, StructureCache(), edges)
+    cache_dae_ci!(ci, result, nothing, nothing, StructureCache())
 
     return result
 end
 
-function _structural_analysis!(ci::CodeInstance, world::UInt, edges::SimpleVector)
+function _structural_analysis!(ci::CodeInstance, world::UInt)
     # Variables
     var_to_diff = DiffGraph(0)
     varclassification = VarEqClassification[]
@@ -66,7 +66,7 @@ function _structural_analysis!(ci::CodeInstance, world::UInt, edges::SimpleVecto
     ir = copy(ci.inferred.ir)
 
     # Allocate variable and equation numbers of any incoming arguments
-    refiner = StructuralRefiner(world, var_to_diff, varkinds, varclassification, Core.svec(ci))
+    refiner = StructuralRefiner(world, var_to_diff, varkinds, varclassification)
     argtypes = Any[make_argument_lattice_elem(Compiler.typeinf_lattice(refiner), Argument(i), argt, add_variable!, add_equation!, add_scope!) for (i, argt) in enumerate(ir.argtypes)]
     nexternalvars = length(var_to_diff)
     nexternaleqs = length(eqssas)
@@ -291,7 +291,7 @@ function _structural_analysis!(ci::CodeInstance, world::UInt, edges::SimpleVecto
             (; result, mapping) = info
         else
             callee_codeinst = stmt.args[1]
-            result = structural_analysis!(callee_codeinst, Compiler.get_inference_world(refiner), edges)
+            result = structural_analysis!(callee_codeinst, Compiler.get_inference_world(refiner))
 
             if isa(result, UncompilableIPOResult)
                 # TODO: Stack trace?

src/interface.jl

@@ -27,10 +27,9 @@ function factory_gen(world::UInt, source::Method, @nospecialize(_gen), settings,
     # 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, edges=Core.svec(factory_mi))
-    edges = Core.svec(ci)
 
     # Perform or lookup DAECompiler specific analysis for this system.
-    result = structural_analysis!(ci, world, edges)
+    result = structural_analysis!(ci, world)
 
     if isa(result, UncompilableIPOResult)
         return Base.generated_body_to_codeinfo(
@@ -44,19 +43,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, edges)
+        tearing_schedule!(tstate, ci, diff_key, world)
     end
     if settings.mode in (InitUncompress, DAE, ODE)
-        tearing_schedule!(tstate, ci, init_key, world, edges)
+        tearing_schedule!(tstate, ci, init_key, world)
     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, edges, settings.mode == DAE ? init_key : nothing)
+        ir_factory = dae_factory_gen(tstate, ci, diff_key, world, settings.mode == DAE ? init_key : nothing)
     elseif settings.mode in (ODE, ODENoInit)
-        ir_factory = ode_factory_gen(tstate, ci, diff_key, world, edges, settings.mode == ODE ? init_key : nothing)
+        ir_factory = ode_factory_gen(tstate, ci, diff_key, world, settings.mode == ODE ? init_key : nothing)
     elseif settings.mode == InitUncompress
-        ir_factory = init_uncompress_gen(result, ci, init_key, diff_key, world, edges)
+        ir_factory = init_uncompress_gen(result, ci, init_key, diff_key, world)
     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, edges::SimpleVector, init_key::Union{TornCacheKey, Nothing})
+function dae_factory_gen(state::TransformationState, ci::CodeInstance, key::TornCacheKey, world::UInt, 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, edges)
+    daef_ci = rhs_finish!(state, ci, key, world, 1)
 
     # 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, edges)
+        initf = init_uncompress_gen!(compact, result, ci, init_key, key, world)
         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, edges::SimpleVector)
+function init_uncompress_gen(result::DAEIPOResult, ci::CodeInstance, init_key::TornCacheKey, diff_key::TornCacheKey, world::UInt)
     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, edges)
+    new_oc = init_uncompress_gen!(compact, result, ci, init_key, diff_key, world)
     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, edges::SimpleVector)
+function init_uncompress_gen!(compact::Compiler.IncrementalCompact, result::DAEIPOResult, ci::CodeInstance, init_key::TornCacheKey, diff_key::TornCacheKey, world::UInt)
     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, edges)
+    daef_ci = gen_init_uncompress!(result, ci, init_key, diff_key, world, 1)
 
     # Create a small opaque closure to adapt from SciML ABI to our own internal
     # ABI

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, edges::SimpleVector, indexT=Int)
+function gen_init_uncompress!(result::DAEIPOResult, ci::CodeInstance, init_key::TornCacheKey, diff_key::TornCacheKey, world::UInt, ordinal::Int, 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, edges, indexT)
+    return gen_init_uncompress!(tstate, ci, init_key, diff_key, world, ordinal, indexT)
 end
 
 function gen_init_uncompress!(
@@ -22,7 +22,6 @@ function gen_init_uncompress!(
         diff_key::TornCacheKey,
         world::UInt,
         ordinal::Int,
-        edges::SimpleVector,
         indexT=Int)
 
     (; result, structure) = state
@@ -98,8 +97,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, edges)
-                callee_daef_ci = rhs_finish!(callee_result, callee_ci, callee_key, world, callee_ordinal, edges)
+                callee_result = structural_analysis!(callee_ci, world)
+                callee_daef_ci = rhs_finish!(callee_result, callee_ci, callee_key, world, callee_ordinal)
                 # Allocate a continuous block of variables for all callee alg and diff states
 
                 empty!(stmt.args)
@@ -155,7 +154,7 @@ 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), edges)
+        daef_ci = cache_dae_ci!(ci, src, src.debuginfo, abi, InitUncompressSpec(init_key, diff_key, ir_ordinal))
         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, edges::SimpleVector, init_key::Union{TornCacheKey, Nothing})
+function ode_factory_gen(state::TransformationState, ci::CodeInstance, key::TornCacheKey, world::UInt, 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, edges)
+    odef_ci = rhs_finish!(state, ci, key, world, 1)
 
     # 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, edges) : nothing
+    initf = init_key !== nothing ? init_uncompress_gen!(compact, result, ci, init_key, key, world) : 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, edges::SimpleVector, indexT=Int)
+function rhs_finish!(result::DAEIPOResult, ci::CodeInstance, key::TornCacheKey, world::UInt, ordinal::Int, indexT=Int)
     structure = make_structure_from_ipo(result)
     tstate = TransformationState(result, structure, copy(result.total_incidence))
-    return rhs_finish!(tstate, ci, key, world, ordinal, edges, indexT)
+    return rhs_finish!(tstate, ci, key, world, ordinal, indexT)
 end
 
 function rhs_finish!(
@@ -75,7 +75,6 @@ function rhs_finish!(
     key::TornCacheKey,
     world::UInt,
     ordinal::Int,
-    edges::SimpleVector,
     indexT=Int)
 
     (; result, structure) = state
@@ -146,8 +145,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, edges)
-                callee_daef_ci = rhs_finish!(callee_result, callee_ci, callee_key, world, callee_ordinal, edges)
+                callee_result = structural_analysis!(callee_ci, world)
+                callee_daef_ci = rhs_finish!(callee_result, callee_ci, callee_key, world, callee_ordinal)
                 # Allocate a continuous block of variables for all callee alg and diff states
 
                 empty!(stmt.args)
@@ -219,7 +218,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), edges)
+        daef_ci = cache_dae_ci!(ci, src, src.debuginfo, abi, RHSSpec(key, ir_ordinal))
         ccall(:jl_add_codeinst_to_jit, Cvoid, (Any, Any), daef_ci, src)
 
         push!(cis, daef_ci)

src/transform/common.jl

@@ -63,7 +63,8 @@ function ir_to_src(ir::IRCode)
     return src
 end
 
-function cache_dae_ci!(old_ci, src, debuginfo, abi, owner, edges::SimpleVector)
+function cache_dae_ci!(old_ci, src, debuginfo, abi, owner)
+    edges = Core.svec(old_ci)
     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, edges)

src/transform/tearing/schedule.jl

@@ -536,13 +536,13 @@ function matching_for_key(result::DAEIPOResult, key::TornCacheKey, structure = m
     return var_eq_matching
 end
 
-function tearing_schedule!(result::DAEIPOResult, ci::CodeInstance, key::TornCacheKey, world::UInt, edges::SimpleVector)
+function tearing_schedule!(result::DAEIPOResult, ci::CodeInstance, key::TornCacheKey, world::UInt)
     structure = make_structure_from_ipo(result)
     tstate = TransformationState(result, structure, copy(result.total_incidence))
-    return tearing_schedule!(tstate, ci, key, world, edges)
+    return tearing_schedule!(tstate, ci, key, world)
 end
 
-function tearing_schedule!(state::TransformationState, ci::CodeInstance, key::TornCacheKey, world::UInt, edges::SimpleVector)
+function tearing_schedule!(state::TransformationState, ci::CodeInstance, key::TornCacheKey, world::UInt)
     result_ci = find_matching_ci(ci->isa(ci.owner, SICMSpec) && ci.owner.key == key, ci.def, world)
     if result_ci !== nothing
         return result_ci
@@ -709,8 +709,8 @@ function tearing_schedule!(state::TransformationState, ci::CodeInstance, key::To
                 if isa(callee_codeinst, MethodInstance)
                     callee_codeinst = Compiler.get(Compiler.code_cache(interp), callee_codeinst, nothing)
                 end
-                callee_result = structural_analysis!(callee_codeinst, world, edges)
-                callee_sicm_ci = tearing_schedule!(callee_result, callee_codeinst, callee_key, world, edges)
+                callee_result = structural_analysis!(callee_codeinst, world)
+                callee_sicm_ci = tearing_schedule!(callee_result, callee_codeinst, callee_key, world)
 
                 inst[:type] = Any
                 inst[:flag] = UInt32(0)
@@ -1029,10 +1029,10 @@ function tearing_schedule!(state::TransformationState, ci::CodeInstance, key::To
         debuginfo = src.debuginfo
     end
 
-    sicm_ci = cache_dae_ci!(ci, src, debuginfo, sig, SICMSpec(key), edges)
+    sicm_ci = cache_dae_ci!(ci, src, debuginfo, sig, SICMSpec(key))
     ccall(:jl_add_codeinst_to_jit, Cvoid, (Any, Any), sicm_ci, src)
 
-    torn_ci = cache_dae_ci!(ci, TornIR(ir_sicm, irs), nothing, sig, TornIRSpec(key), edges)
+    torn_ci = cache_dae_ci!(ci, TornIR(ir_sicm, irs), nothing, sig, TornIRSpec(key))
 
     return sicm_ci
 end