Invalidate caches on refresh() (#24)

* Record `factory` as an edge for all processed `CodeInstance`s

* Address review comments, fix `CodeInstance`s not being invalidated

* Add invalidation edges transitively, don't pass them around

* Simplify backedge storage when we know we have only 1 edge

Author: serenity4

src/analysis/ADAnalyzer.jl

@@ -4,7 +4,7 @@ using Compiler
 using Diffractor
 using Core: SimpleVector, CodeInstance, Const
 using Compiler: ArgInfo, StmtInfo, AbstractInterpreter, InferenceParams, OptimizationParams,
-    AbsIntState, CallInfo, InferenceResult
+    AbsIntState, CallInfo, InferenceResult, InferenceState
 
 struct ADCache; end
 
@@ -17,10 +17,12 @@ AD'd using Diffractor.
 struct ADAnalyzer <: Compiler.AbstractInterpreter
     world::UInt
     inf_cache::Vector{Compiler.InferenceResult}
+    edges::SimpleVector # additional edges
     function ADAnalyzer(;
             world::UInt = Base.get_world_counter(),
-            inf_cache::Vector{Compiler.InferenceResult} = Compiler.InferenceResult[])
-        new(world, inf_cache)
+            inf_cache::Vector{Compiler.InferenceResult} = Compiler.InferenceResult[],
+            edges = Compiler.empty_edges)
+        new(world, inf_cache, edges)
     end
 end
 
@@ -60,6 +62,11 @@ struct AnalyzedSource
     inline_cost::Compiler.InlineCostType
 end
 
+@override function Compiler.result_edges(interp::ADAnalyzer, caller::InferenceState)
+    edges = @invoke Compiler.result_edges(interp::AbstractInterpreter, caller::InferenceState)
+    Core.svec(edges..., interp.edges...)
+end
+
 @override function Compiler.transform_result_for_cache(interp::ADAnalyzer, result::InferenceResult, edges::SimpleVector)
     ir = result.src.optresult.ir
     params = Compiler.OptimizationParams(interp)
@@ -88,12 +95,16 @@ 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=Compiler.empty_edges)
+    @assert !force_inline_all
+    interp = ADAnalyzer(; world, edges)
+    mi = get_method_instance(tt, world)
     ci = Compiler.typeinf_ext(interp, mi, Compiler.SOURCE_MODE_ABI)
 end

src/analysis/refiner.jl

@@ -415,4 +415,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/interface.jl

@@ -22,10 +22,11 @@ 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)
 
     # 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=Core.svec(factory_mi))
 
     # Perform or lookup DAECompiler specific analysis for this system.
     result = structural_analysis!(ci, world)

src/transform/codegen/init_factory.jl

@@ -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/common.jl

@@ -64,10 +64,13 @@ function ir_to_src(ir::IRCode)
 end
 
 function cache_dae_ci!(old_ci, src, debuginfo, abi, owner)
+    mi = old_ci.def
+    edges = Core.svec(mi)
     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)
-    ccall(:jl_mi_cache_insert, Cvoid, (Any, Any), old_ci.def, daef_ci)
+        nothing, debuginfo, edges)
+    ccall(:jl_method_instance_add_backedge, Cvoid, (Any, Any, Any), mi, nothing, daef_ci)
+    ccall(:jl_mi_cache_insert, Cvoid, (Any, Any), mi, daef_ci)
     return daef_ci
 end
 

test/invalidation.jl

@@ -0,0 +1,46 @@
+using Test
+using SciMLBase
+using Sundials
+using Compiler
+using Core.IR
+using DAECompiler
+using DAECompiler: get_method_instance, refresh
+using DAECompiler.Intrinsics
+
+function get_cached_code_instances(mi::MethodInstance)
+  ci = mi.cache
+  cached = CodeInstance[ci]
+  while isdefined(ci, :next)
+    push!(cached, ci.next)
+    ci = ci.next
+  end
+  return cached
+end
+
+trivialeq!() = always!(ddt(continuous()))
+
+@testset "Invalidation" begin
+  mi = get_method_instance(Tuple{typeof(trivialeq!)}, Base.get_world_counter())
+
+  ci = DAECompiler.find_matching_ci(ci->ci.owner == DAECompiler.StructureCache(), mi, Base.get_world_counter())
+  @test ci === nothing
+
+  solve(DAECProblem(trivialeq!, (1,) .=> 1.), IDA())
+  ci = DAECompiler.find_matching_ci(ci->ci.owner == DAECompiler.StructureCache(), mi, Base.get_world_counter())
+  @test ci !== nothing
+
+  cached = get_cached_code_instances(mi)
+  world_before = Base.get_world_counter()
+  @test getproperty.(cached, :max_world) == fill(typemax(UInt), length(cached))
+  refresh()
+  @test getproperty.(cached, :max_world) == fill(world_before, length(cached))
+
+  ci = DAECompiler.find_matching_ci(ci->ci.owner == DAECompiler.StructureCache(), mi, Base.get_world_counter())
+  @test ci === nothing
+
+  solve(DAECProblem(trivialeq!, (1,) .=> 1.), IDA())
+  ci = DAECompiler.find_matching_ci(ci->ci.owner == DAECompiler.StructureCache(), mi, Base.get_world_counter())
+  @test ci !== nothing
+
+  @test length(get_cached_code_instances(mi)) == 2 * length(cached)
+end;

test/runtests.jl

@@ -2,3 +2,4 @@ include("basic.jl")
 include("ipo.jl")
 include("ssrm.jl")
 include("regression.jl")
+include("invalidation.jl")