InternalCodeCache update

Author: Keno

base/boot.jl

@@ -501,9 +501,8 @@ struct LineInfoNode # legacy support for aiding Serializer.deserialize of old IR
     LineInfoNode(mod::Module, @nospecialize(method), file::Symbol, line::Int32, inlined_at::Int32) = new(mod, method, file, line, inlined_at)
 end
 
-
 function CodeInstance(
-    mi::MethodInstance, @nospecialize(rettype), @nospecialize(exctype), @nospecialize(inferred_const),
+    mi::MethodSpecialization, @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, edges::DebugInfo)
@@ -649,12 +648,12 @@ Symbol(s::Symbol) = s
 # module providing the IR object model
 module IR
 
-export CodeInfo, MethodInstance, CodeInstance, GotoNode, GotoIfNot, ReturnNode,
+export CodeInfo, MethodSpecialization, MethodInstance, CodeInstance, GotoNode, GotoIfNot, ReturnNode,
     NewvarNode, SSAValue, SlotNumber, Argument,
     PiNode, PhiNode, PhiCNode, UpsilonNode, DebugInfo,
     Const, PartialStruct, InterConditional, EnterNode
 
-using Core: CodeInfo, MethodInstance, CodeInstance, GotoNode, GotoIfNot, ReturnNode,
+using Core: CodeInfo, MethodSpecialization, MethodInstance, CodeInstance, GotoNode, GotoIfNot, ReturnNode,
     NewvarNode, SSAValue, SlotNumber, Argument,
     PiNode, PhiNode, PhiCNode, UpsilonNode, DebugInfo,
     Const, PartialStruct, InterConditional, EnterNode
@@ -1006,6 +1005,9 @@ const check_top_bit = check_sign_bit
 EnterNode(old::EnterNode, new_dest::Int) = isdefined(old, :scope) ?
     EnterNode(new_dest, old.scope) : EnterNode(new_dest)
 
+eval(Core, :((MS::Type{<:MethodSpecialization})(def::Union{Method, Module, MethodSpecialization}, abi::Type{<:Tuple}) =
+    $(Expr(:new, :MS, :def, :abi))))
+
 include(Core, "optimized_generics.jl")
 
 ccall(:jl_set_istopmod, Cvoid, (Any, Bool), Core, true)

base/compiler/cicache.jl

@@ -3,15 +3,35 @@
 """
     struct InternalCodeCache
 
-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.
+The internal code cache is keyed on type specializations, represented by
+MethodSpecialization{DefaultSpec} aka MethodInstance. External abstract
+interpreters may use this same structure by using a different `Spec` for their
+`MethodSpecialization{Spec}`. `InternalCodeCache` will match such specializations
+by type. Additionally, it is possible to specialize methods on properties other
+than types, but this requires custom caching logic. `InternalCodeCache` currently
+only supports type-based specialization.
 """
 struct InternalCodeCache
+    mitype::DataType # <: MethodSpecialization, but stored as DataType for efficient ===
+    InternalCodeCache(T::Type{<:MethodSpecialization}) =
+        new(T)
 end
+InternalCodeCache() = InternalCodeCache(MethodInstance)
 
 function setindex!(cache::InternalCodeCache, ci::CodeInstance, mi::MethodInstance)
-    ccall(:jl_mi_cache_insert, Cvoid, (Any, Any), mi, ci)
+    ms::MethodSpecialization = mi
+    while typeof(ms) !== cache.mitype
+        if !isdefined(ms, :next)
+            # No specialization for this spec. Try to allocate it now.
+            newms = cache.mitype(mi.def, mi.specTypes)
+            if @atomiconce :sequentially_consistent (ms.next = newms)
+                ms = newms
+                break
+            end
+        end
+        ms = @atomic :acquire ms.next
+    end
+    ccall(:jl_mi_cache_insert, Cvoid, (Any, Any), ms, ci)
     return cache
 end
 
@@ -48,11 +68,21 @@ 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
+    ms::MethodSpecialization = mi
+    while typeof(ms) !== wvc.cache.mitype
+        isdefined(ms, :next) || return false
+        ms = ms.next
+    end
+    return ccall(:jl_rettype_inferred, Any, (Any, UInt, UInt), ms, 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))
+    ms::MethodSpecialization = mi
+    while typeof(ms) !== wvc.cache.mitype
+        isdefined(ms, :next) || return default
+        ms = ms.next
+    end
+    r = ccall(:jl_rettype_inferred, Any, (Any, UInt, UInt), ms, first(wvc.worlds), last(wvc.worlds))
     if r === nothing
         return default
     end

base/compiler/compiler.jl

@@ -6,11 +6,13 @@ using Core.Intrinsics, Core.IR
 
 import Core: print, println, show, write, unsafe_write, stdout, stderr,
              _apply_iterate, svec, apply_type, Builtin, IntrinsicFunction,
+             MethodSpecialization,
              MethodInstance, CodeInstance, MethodTable, MethodMatch, PartialOpaque,
              TypeofVararg
 
 const getproperty = Core.getfield
 const setproperty! = Core.setfield!
+const setpropertyonce! = Core.setfieldonce!
 const swapproperty! = Core.swapfield!
 const modifyproperty! = Core.modifyfield!
 const replaceproperty! = Core.replacefield!

src/codegen.cpp

@@ -5162,7 +5162,7 @@ static jl_cgval_t emit_invoke(jl_codectx_t &ctx, const jl_cgval_t &lival, ArrayR
     jl_cgval_t result;
     if (lival.constant) {
         jl_method_instance_t *mi = (jl_method_instance_t*)lival.constant;
-        assert(jl_is_method_instance(mi));
+        assert(jl_is_method_specialization(mi));
         if (mi == ctx.linfo) {
             // handle self-recursion specially
             jl_returninfo_t::CallingConv cc = jl_returninfo_t::CallingConv::Boxed;

src/jltypes.c

@@ -3332,7 +3332,7 @@ void jl_init_types(void) JL_GC_DISABLED
         (jl_unionall_t*)jl_meth_spec_type->name->wrapper;
     // These fields should be constant, but Serialization wants to mutate them in initialization
     //const static uint32_t method_instance_constfields[1] = { 0x00000007 }; // (1<<0)|(1<<1);
-    const static uint32_t method_instance_atomicfields[1] = { 0x0000008 }; // (1<<3)
+    const static uint32_t method_instance_atomicfields[1] = { 0x0000018 }; // (1<<3)|(1<<4)
     //Fields 3 and 4 must be protected by method->write_lock, and thus all operations on jl_method_instance_t are threadsafe. TODO: except inInference
     //jl_method_instance_type->name->constfields = method_instance_constfields;
     jl_meth_spec_type->name->atomicfields = method_instance_atomicfields;

src/julia.h

@@ -1528,6 +1528,7 @@ static inline int jl_field_isconst(jl_datatype_t *st, int i) JL_NOTSAFEPOINT
 #define jl_is_newvarnode(v)  jl_typetagis(v,jl_newvarnode_type)
 #define jl_is_linenode(v)    jl_typetagis(v,jl_linenumbernode_type)
 #define jl_is_method_instance(v) jl_typetagis(v,jl_method_instance_type)
+#define jl_is_method_specialization(v) jl_isa((jl_value_t*)v,(jl_value_t*)jl_method_specialization_type)
 #define jl_is_code_instance(v) jl_typetagis(v,jl_code_instance_type)
 #define jl_is_code_info(v)   jl_typetagis(v,jl_code_info_type)
 #define jl_is_method(v)      jl_typetagis(v,jl_method_type)