Address review comments, fix CodeInstances not being invalidated

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)
@@ -95,16 +102,9 @@ function get_method_instance(@nospecialize(tt), world)
     mi = Compiler.specialize_method(match)
 end
 
-function ad_typeinf(world, tt; force_inline_all=false, edges=nothing)
+function ad_typeinf(world, tt; force_inline_all=false, edges=Compiler.empty_edges)
     @assert !force_inline_all
-    interp = ADAnalyzer(;world)
+    interp = ADAnalyzer(; world, edges)
     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

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

src/analysis/structural.jl

@@ -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)
+    refiner = StructuralRefiner(world, var_to_diff, varkinds, varclassification, Core.svec(ci))
     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)

src/interface.jl

@@ -23,11 +23,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)
-    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, edges)
+    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)

src/transform/common.jl

@@ -67,11 +67,19 @@ 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, edges)
-    Compiler.store_backedges(daef_ci, edges)
+    add_backedges_to_callees(daef_ci, edges)
     ccall(:jl_mi_cache_insert, Cvoid, (Any, Any), old_ci.def, daef_ci)
     return daef_ci
 end
 
+# Equivalent to `Compiler.store_backedges` in our case, but we allow `caller.def.def` to not be a `Method`.
+function add_backedges_to_callees(caller::CodeInstance, edges::SimpleVector)
+    for edge in edges
+        isa(edge, CodeInstance) && (edge = edge.def)
+        ccall(:jl_method_instance_add_backedge, Cvoid, (Any, Any, Any), edge::MethodInstance, nothing, caller)
+    end
+end
+
 function replace_call!(ir::Union{IRCode,IncrementalCompact}, idx::SSAValue, new_call::Expr)
     @assert !isa(ir[idx][:inst], PhiNode)
     ir[idx][:inst] = new_call

test/invalidation.jl

@@ -30,8 +30,8 @@ trivialeq!() = always!(ddt(continuous()))
   @test ci !== nothing
 
   cached = get_cached_code_instances(mi)
-  @test getproperty.(cached, :max_world) == fill(typemax(UInt), length(cached))
   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))
 
@@ -43,4 +43,4 @@ trivialeq!() = always!(ddt(continuous()))
   @test ci !== nothing
 
   @test length(get_cached_code_instances(mi)) == 2 * length(cached)
-end
+end;