Tag CodeInstance for external abstract interpreter

External abstract interpreters need to be able to take part in
invalidation, as well as cross-session inference result caching.

MethodInstance.cache contains a set of CodeInstances which hold the
inference cache results. Instead of haveing external abstract
interpreters create detached CodeInstances we attach them the the
MethodInstance they were created from. This solves the invalidation
problem and the cross-session cache problem in one go, but requires
that we tag/mark the CodeInstance with a token to remember which
abstract interpreter created them.

`nothing` is reserved for the native interpreter and all other
interpreters can use values that compare with `jl_egal`.

To perform a cache-lookup we need the current abstract interpreter
token.

Co-authored-by: Shuhei Kadowaki <40514306+aviatesk@users.noreply.github.com>

base/boot.jl

@@ -483,13 +483,13 @@ eval(Core, quote
 end)
 
 function CodeInstance(
-    mi::MethodInstance, @nospecialize(rettype), @nospecialize(exctype), @nospecialize(inferred_const),
+    mi::MethodInstance, owner, @nospecialize(rettype), @nospecialize(exctype), @nospecialize(inferred_const),
     @nospecialize(inferred), const_flags::Int32, min_world::UInt, max_world::UInt,
     ipo_effects::UInt32, effects::UInt32, @nospecialize(analysis_results),
     relocatability::UInt8)
     return ccall(:jl_new_codeinst, Ref{CodeInstance},
-        (Any, Any, Any, Any, Any, Int32, UInt, UInt, UInt32, UInt32, Any, UInt8),
-        mi, rettype, exctype, inferred_const, inferred, const_flags, min_world, max_world,
+        (Any, Any, Any, Any, Any, Any, Int32, UInt, UInt, UInt32, UInt32, Any, UInt8),
+        mi, owner, rettype, exctype, inferred_const, inferred, const_flags, min_world, max_world,
         ipo_effects, effects, analysis_results,
         relocatability)
 end

base/compiler/cicache.jl

@@ -7,14 +7,17 @@ Internally, each `MethodInstance` keep a unique global cache of code instances
 that have been created for the given method instance, stratified by world age
 ranges. This struct abstracts over access to this cache.
 """
-struct InternalCodeCache end
+struct InternalCodeCache
+    owner::Any # `jl_egal` is used for comparison
+end
 
 function setindex!(cache::InternalCodeCache, ci::CodeInstance, mi::MethodInstance)
+    @assert ci.owner === cache.owner
     ccall(:jl_mi_cache_insert, Cvoid, (Any, Any), mi, ci)
     return cache
 end
 
-const GLOBAL_CI_CACHE = InternalCodeCache()
+const GLOBAL_CI_CACHE = InternalCodeCache(nothing)
 
 struct WorldRange
     min_world::UInt
@@ -49,11 +52,11 @@ WorldView(wvc::WorldView, wr::WorldRange) = WorldView(wvc.cache, wr)
 WorldView(wvc::WorldView, args...) = WorldView(wvc.cache, args...)
 
 function haskey(wvc::WorldView{InternalCodeCache}, mi::MethodInstance)
-    return ccall(:jl_rettype_inferred, Any, (Any, UInt, UInt), mi, first(wvc.worlds), last(wvc.worlds)) !== nothing
+    return ccall(:jl_rettype_inferred, Any, (Any, Any, UInt, UInt), wvc.cache.owner, mi, first(wvc.worlds), last(wvc.worlds)) !== nothing
 end
 
 function get(wvc::WorldView{InternalCodeCache}, mi::MethodInstance, default)
-    r = ccall(:jl_rettype_inferred, Any, (Any, UInt, UInt), mi, first(wvc.worlds), last(wvc.worlds))
+    r = ccall(:jl_rettype_inferred, Any, (Any, Any, UInt, UInt), wvc.cache.owner, mi, first(wvc.worlds), last(wvc.worlds))
     if r === nothing
         return default
     end

base/compiler/typeinfer.jl

@@ -326,7 +326,7 @@ function CodeInstance(interp::AbstractInterpreter, result::InferenceResult,
         end
     end
     # relocatability = isa(inferred_result, String) ? inferred_result[end] : UInt8(0)
-    return CodeInstance(result.linfo,
+    return CodeInstance(result.linfo, cache_owner(interp),
         widenconst(result_type), widenconst(result.exc_result), rettype_const, inferred_result,
         const_flags, first(valid_worlds), last(valid_worlds),
         # TODO: Actually do something with non-IPO effects

base/compiler/types.jl

@@ -14,6 +14,7 @@ the following methods to satisfy the `AbstractInterpreter` API requirement:
 - `OptimizationParams(interp::NewInterpreter)` - return an `OptimizationParams` instance
 - `get_world_counter(interp::NewInterpreter)` - return the world age for this interpreter
 - `get_inference_cache(interp::NewInterpreter)` - return the local inference cache
+- `cache_owner(interp::NewInterpreter)` - return the owner of any new cache entries
 - `code_cache(interp::NewInterpreter)` - return the global inference cache
 """
 :(AbstractInterpreter)
@@ -404,6 +405,7 @@ InferenceParams(interp::NativeInterpreter) = interp.inf_params
 OptimizationParams(interp::NativeInterpreter) = interp.opt_params
 get_world_counter(interp::NativeInterpreter) = interp.world
 get_inference_cache(interp::NativeInterpreter) = interp.inf_cache
+cache_owner(interp::NativeInterpreter) = nothing
 code_cache(interp::NativeInterpreter) = WorldView(GLOBAL_CI_CACHE, get_world_counter(interp))
 
 """

base/compiler/utilities.jl

@@ -321,25 +321,6 @@ function iterate(iter::BackedgeIterator, i::Int=1)
     return BackedgePair(item, backedges[i+1]::MethodInstance), i+2            # `invoke` calls
 end
 
-"""
-    add_invalidation_callback!(callback, mi::MethodInstance)
-
-Register `callback` to be triggered upon the invalidation of `mi`.
-`callback` should a function taking two arguments, `callback(replaced::MethodInstance, max_world::UInt32)`,
-and it will be recursively invoked on `MethodInstance`s within the invalidation graph.
-"""
-function add_invalidation_callback!(@nospecialize(callback), mi::MethodInstance)
-    if !isdefined(mi, :callbacks)
-        callbacks = mi.callbacks = Any[callback]
-    else
-        callbacks = mi.callbacks::Vector{Any}
-        if !any(@nospecialize(cb)->cb===callback, callbacks)
-            push!(callbacks, callback)
-        end
-    end
-    return callbacks
-end
-
 #########
 # types #
 #########

base/reflection.jl

@@ -1331,6 +1331,15 @@ function method_instances(@nospecialize(f), @nospecialize(t), world::UInt)
     return results
 end
 
+function method_instance(@nospecialize(f), @nospecialize(t);
+                         world=Base.get_world_counter(), method_table=nothing)
+    tt = signature_type(f, t)
+    mi = ccall(:jl_method_lookup_by_tt, Any,
+                (Any, Csize_t, Any),
+                tt, world, method_table)
+    return mi::Union{Nothing, MethodInstance}
+end
+
 default_debug_info_kind() = unsafe_load(cglobal(:jl_default_debug_info_kind, Cint))
 
 # this type mirrors jl_cgparams_t (documented in julia.h)

doc/src/devdocs/ast.md

@@ -627,6 +627,10 @@ for important details on how to modify these fields safely.
 
     The `MethodInstance` that this cache entry is derived from.
 
+  * `owner`
+
+    A token that represents the owner of this `CodeInstance`. Will use `jl_egal` to match.
+
 
   * `rettype`/`rettype_const`
 

doc/src/devdocs/locks.md

@@ -155,6 +155,7 @@ MethodInstance/CodeInstance updates : Method->writelock, codegen lock
 >       * specTypes
 >       * sparam_vals
 >       * def
+>       * owner
 
 >   * These are set by `jl_type_infer` (while holding codegen lock):
 >       * cache

src/builtins.c

@@ -2378,6 +2378,7 @@ jl_fptr_args_t jl_get_builtin_fptr(jl_datatype_t *dt)
     jl_typemap_entry_t *entry = (jl_typemap_entry_t*)jl_atomic_load_relaxed(&dt->name->mt->defs);
     jl_method_instance_t *mi = jl_atomic_load_relaxed(&entry->func.method->unspecialized);
     jl_code_instance_t *ci = jl_atomic_load_relaxed(&mi->cache);
+    assert(ci->owner == jl_nothing);
     return jl_atomic_load_relaxed(&ci->specptr.fptr1);
 }