Mark codeinst->inferred as atomic and rescope type inference lock

base/compiler/abstractinterpretation.jl

@@ -1076,9 +1076,16 @@ function const_prop_methodinstance_heuristic(
             return false
         else
             code = get(code_cache(interp), mi, nothing)
-            if isdefined(code, :inferred) && inlining_policy(
-                    interp, code.inferred, IR_FLAG_NULL, mi, argtypes) !== nothing
-                return true
+            if isdefined(code, :inferred)
+                if isa(code, CodeInstance)
+                    inferred = @atomic :monotonic code.inferred
+                else
+                    inferred = code.inferred
+                end
+                if inlining_policy(
+                    interp, inferred, IR_FLAG_NULL, mi, argtypes) !== nothing
+                    return true
+                end
             end
         end
     end

base/compiler/ssair/inlining.jl

@@ -1,5 +1,7 @@
 # This file is a part of Julia. License is MIT: https://julialang.org/license
 
+const depends = :monotonic
+
 @nospecialize
 
 struct Signature
@@ -826,7 +828,7 @@ function resolve_todo(todo::InliningTodo, state::InliningState, flag::UInt8)
                 et !== nothing && push!(et, mi)
                 return ConstantCase(quoted(code.rettype_const))
             else
-                src = code.inferred
+                src = @atomic depends code.inferred
             end
             effects = decode_effects(code.ipo_purity_bits)
         else

base/compiler/typeinfer.jl

@@ -1,9 +1,10 @@
 # This file is a part of Julia. License is MIT: https://julialang.org/license
 
 # Tracking of newly-inferred MethodInstances during precompilation
-const track_newly_inferred = RefValue{Bool}(false)
 const newly_inferred = MethodInstance[]
 
+const depends = :monotonic
+
 # build (and start inferring) the inference frame for the top-level MethodInstance
 function typeinf(interp::AbstractInterpreter, result::InferenceResult, cache::Symbol)
     frame = InferenceState(result, cache, interp)
@@ -205,9 +206,11 @@ const __measure_typeinf__ = fill(false)
 # Wrapper around _typeinf that optionally records the exclusive time for each invocation.
 function typeinf(interp::AbstractInterpreter, frame::InferenceState)
     if __measure_typeinf__[]
+        ccall(:jl_typeinf_begin, Cvoid, ())
         Timings.enter_new_timer(frame)
         v = _typeinf(interp, frame)
         Timings.exit_current_timer(frame)
+        ccall(:jl_typeinf_end, Cvoid, ())
         return v
     else
         return _typeinf(interp, frame)
@@ -394,10 +397,10 @@ function cache_result!(interp::AbstractInterpreter, result::InferenceResult)
     if !already_inferred
         inferred_result = transform_result_for_cache(interp, linfo, valid_worlds, result.src, result.ipo_effects)
         code_cache(interp)[linfo] = CodeInstance(result, inferred_result, valid_worlds)
-        if track_newly_inferred[]
-            m = linfo.def
-            if isa(m, Method)
-                m.module != Core && push!(newly_inferred, linfo)
+        m = linfo.def
+        if isa(m, Method)
+            if m.module != Core
+                ccall(:jl_log_inferred, Cvoid, (Any,), linfo)
             end
         end
     end
@@ -841,7 +844,8 @@ function typeinf_edge(interp::AbstractInterpreter, method::Method, @nospecialize
     mi = specialize_method(method, atype, sparams)::MethodInstance
     code = get(code_cache(interp), mi, nothing)
     if code isa CodeInstance # return existing rettype if the code is already inferred
-        if code.inferred === nothing && is_stmt_inline(get_curr_ssaflag(caller))
+        inferred = @atomic depends code.inferred
+        if inferred === nothing && is_stmt_inline(get_curr_ssaflag(caller))
             # we already inferred this edge before and decided to discard the inferred code,
             # nevertheless we re-infer it here again and keep it around in the local cache
             # since the inliner will request to use it later
@@ -921,60 +925,52 @@ end
 # compute an inferred frame
 function typeinf_frame(interp::AbstractInterpreter, method::Method, @nospecialize(atype), sparams::SimpleVector, run_optimizer::Bool)
     mi = specialize_method(method, atype, sparams)::MethodInstance
-    ccall(:jl_typeinf_begin, Cvoid, ())
     result = InferenceResult(mi)
     frame = InferenceState(result, run_optimizer ? :global : :no, interp)
     frame === nothing && return nothing
     typeinf(interp, frame)
-    ccall(:jl_typeinf_end, Cvoid, ())
     return frame
 end
 
 # compute (and cache) an inferred AST and return type
 function typeinf_ext(interp::AbstractInterpreter, mi::MethodInstance)
     method = mi.def::Method
-    for i = 1:2 # test-and-lock-and-test
-        i == 2 && ccall(:jl_typeinf_begin, Cvoid, ())
-        code = get(code_cache(interp), mi, nothing)
-        if code isa CodeInstance
-            # see if this code already exists in the cache
-            inf = code.inferred
-            if use_const_api(code)
-                i == 2 && ccall(:jl_typeinf_end, Cvoid, ())
-                tree = ccall(:jl_new_code_info_uninit, Ref{CodeInfo}, ())
-                rettype_const = code.rettype_const
-                tree.code = Any[ ReturnNode(quoted(rettype_const)) ]
-                nargs = Int(method.nargs)
-                tree.slotnames = ccall(:jl_uncompress_argnames, Vector{Symbol}, (Any,), method.slot_syms)
-                tree.slotflags = fill(IR_FLAG_NULL, nargs)
-                tree.ssavaluetypes = 1
-                tree.codelocs = Int32[1]
-                tree.linetable = [LineInfoNode(method.module, method.name, method.file, method.line, Int32(0))]
-                tree.inferred = true
-                tree.ssaflags = UInt8[0]
-                tree.pure = true
-                tree.inlineable = true
-                tree.parent = mi
-                tree.rettype = Core.Typeof(rettype_const)
-                tree.min_world = code.min_world
-                tree.max_world = code.max_world
-                return tree
-            elseif isa(inf, CodeInfo)
-                i == 2 && ccall(:jl_typeinf_end, Cvoid, ())
-                if !(inf.min_world == code.min_world &&
-                     inf.max_world == code.max_world &&
-                     inf.rettype === code.rettype)
-                    inf = copy(inf)
-                    inf.min_world = code.min_world
-                    inf.max_world = code.max_world
-                    inf.rettype = code.rettype
-                end
-                return inf
-            elseif isa(inf, Vector{UInt8})
-                i == 2 && ccall(:jl_typeinf_end, Cvoid, ())
-                inf = _uncompressed_ir(code, inf)
-                return inf
+    code = get(code_cache(interp), mi, nothing)
+    if code isa CodeInstance
+        # see if this code already exists in the cache
+        inf = @atomic depends code.inferred
+        if use_const_api(code)
+            tree = ccall(:jl_new_code_info_uninit, Ref{CodeInfo}, ())
+            rettype_const = code.rettype_const
+            tree.code = Any[ ReturnNode(quoted(rettype_const)) ]
+            nargs = Int(method.nargs)
+            tree.slotnames = ccall(:jl_uncompress_argnames, Vector{Symbol}, (Any,), method.slot_syms)
+            tree.slotflags = fill(IR_FLAG_NULL, nargs)
+            tree.ssavaluetypes = 1
+            tree.codelocs = Int32[1]
+            tree.linetable = [LineInfoNode(method.module, method.name, method.file, method.line, Int32(0))]
+            tree.inferred = true
+            tree.ssaflags = UInt8[0]
+            tree.pure = true
+            tree.inlineable = true
+            tree.parent = mi
+            tree.rettype = Core.Typeof(rettype_const)
+            tree.min_world = code.min_world
+            tree.max_world = code.max_world
+            return tree
+        elseif isa(inf, CodeInfo)
+            if !(inf.min_world == code.min_world &&
+                    inf.max_world == code.max_world &&
+                    inf.rettype === code.rettype)
+                inf = copy(inf)
+                inf.min_world = code.min_world
+                inf.max_world = code.max_world
+                inf.rettype = code.rettype
             end
+            return inf
+        elseif isa(inf, Vector{UInt8})
+            inf = _uncompressed_ir(code, inf)
+            return inf
         end
     end
     if ccall(:jl_get_module_infer, Cint, (Any,), method.module) == 0 && !generating_sysimg()
@@ -984,7 +980,6 @@ function typeinf_ext(interp::AbstractInterpreter, mi::MethodInstance)
     frame = InferenceState(InferenceResult(mi), #=cache=#:global, interp)
     frame === nothing && return nothing
     typeinf(interp, frame)
-    ccall(:jl_typeinf_end, Cvoid, ())
     frame.src.inferred || return nothing
     return frame.src
 end
@@ -995,18 +990,13 @@ function typeinf_type(interp::AbstractInterpreter, method::Method, @nospecialize
         return Union{} # don't ask: it does weird and unnecessary things, if it occurs during bootstrap
     end
     mi = specialize_method(method, atype, sparams)::MethodInstance
-    for i = 1:2 # test-and-lock-and-test
-        i == 2 && ccall(:jl_typeinf_begin, Cvoid, ())
-        code = get(code_cache(interp), mi, nothing)
-        if code isa CodeInstance
-            # see if this rettype already exists in the cache
-            i == 2 && ccall(:jl_typeinf_end, Cvoid, ())
-            return code.rettype
-        end
+    code = get(code_cache(interp), mi, nothing)
+    if code isa CodeInstance
+        # see if this rettype already exists in the cache
+        return code.rettype
     end
     result = InferenceResult(mi)
     typeinf(interp, result, :global)
-    ccall(:jl_typeinf_end, Cvoid, ())
     result.result isa InferenceState && return nothing
     return widenconst(ignorelimited(result.result))
 end
@@ -1019,17 +1009,13 @@ function typeinf_ext_toplevel(interp::AbstractInterpreter, linfo::MethodInstance
         src = typeinf_ext(interp, linfo)
     else
         src = linfo.uninferred::CodeInfo
+        # toplevel lambda - infer directly
         if !src.inferred
-            # toplevel lambda - infer directly
-            ccall(:jl_typeinf_begin, Cvoid, ())
-            if !src.inferred
-                result = InferenceResult(linfo)
-                frame = InferenceState(result, src, #=cache=#:global, interp)
-                typeinf(interp, frame)
-                @assert frame.inferred # TODO: deal with this better
-                src = frame.src
-            end
-            ccall(:jl_typeinf_end, Cvoid, ())
+            result = InferenceResult(linfo)
+            frame = InferenceState(result, src, #=cache=#:global, interp)
+            typeinf(interp, frame)
+            @assert frame.inferred # TODO: deal with this better
+            src = frame.src
         end
     end
     return src

base/loading.jl

@@ -1470,17 +1470,16 @@ function include_package_for_output(pkg::PkgId, input::String, depot_path::Vecto
         task_local_storage()[:SOURCE_PATH] = source
     end
 
-    Core.Compiler.track_newly_inferred.x = true
+    ccall(:jl_set_newly_inferred, Cvoid, (Any,), Core.Compiler.newly_inferred)
     try
         Base.include(Base.__toplevel__, input)
     catch ex
         precompilableerror(ex) || rethrow()
         @debug "Aborting `create_expr_cache'" exception=(ErrorException("Declaration of __precompile__(false) not allowed"), catch_backtrace())
         exit(125) # we define status = 125 means PrecompileableError
     finally
-        Core.Compiler.track_newly_inferred.x = false
+        ccall(:jl_untrack_newly_inferred, Cvoid, ())
     end
-    ccall(:jl_set_newly_inferred, Cvoid, (Any,), Core.Compiler.newly_inferred)
 end
 
 const PRECOMPILE_TRACE_COMPILE = Ref{String}()

src/aotcompile.cpp

@@ -217,7 +217,7 @@ static void jl_ci_cache_lookup(const jl_cgparams_t &cgparams, jl_method_instance
     jl_code_instance_t *codeinst = NULL;
     if (ci != jl_nothing) {
         codeinst = (jl_code_instance_t*)ci;
-        *src_out = (jl_code_info_t*)codeinst->inferred;
+        *src_out = (jl_code_info_t*)jl_atomic_load_relaxed(&codeinst->inferred);
         jl_method_t *def = codeinst->def->def.method;
         if ((jl_value_t*)*src_out == jl_nothing)
             *src_out = NULL;
@@ -232,8 +232,8 @@ static void jl_ci_cache_lookup(const jl_cgparams_t &cgparams, jl_method_instance
             *src_out = jl_type_infer(mi, world, 0);
             if (*src_out) {
                 codeinst = jl_get_method_inferred(mi, (*src_out)->rettype, (*src_out)->min_world, (*src_out)->max_world);
-                if ((*src_out)->inferred && !codeinst->inferred)
-                    codeinst->inferred = jl_nothing;
+                if ((*src_out)->inferred && !jl_atomic_load_relaxed(&codeinst->inferred))
+                    jl_atomic_store_relaxed(&codeinst->inferred, jl_nothing);
             }
         }
     }
@@ -1010,7 +1010,7 @@ void *jl_get_llvmf_defn_impl(jl_method_instance_t *mi, size_t world, char getwra
         jl_value_t *ci = jl_rettype_inferred(mi, world, world);
         if (ci != jl_nothing) {
             jl_code_instance_t *codeinst = (jl_code_instance_t*)ci;
-            src = (jl_code_info_t*)codeinst->inferred;
+            src = (jl_code_info_t*)jl_atomic_load_relaxed(&codeinst->inferred);
             if ((jl_value_t*)src != jl_nothing && !jl_is_code_info(src) && jl_is_method(mi->def.method))
                 src = jl_uncompress_ir(mi->def.method, codeinst, (jl_array_t*)src);
             jlrettype = codeinst->rettype;

src/codegen.cpp

@@ -4793,13 +4793,15 @@ static std::pair<Function*, Function*> get_oc_function(jl_codectx_t &ctx, jl_met
     jl_method_instance_t *mi = jl_specializations_get_linfo(closure_method, sigtype, jl_emptysvec);
     jl_code_instance_t *ci = (jl_code_instance_t*)jl_rettype_inferred(mi, ctx.world, ctx.world);
 
-    if (ci == NULL || (jl_value_t*)ci == jl_nothing || ci->inferred == NULL || ci->inferred == jl_nothing) {
+    auto inferred = jl_atomic_load_relaxed(&ci->inferred);
+
+    if (ci == NULL || (jl_value_t*)ci == jl_nothing || inferred == NULL || inferred == jl_nothing) {
         JL_GC_POP();
         return std::make_pair((Function*)NULL, (Function*)NULL);
     }
     ++EmittedOpaqueClosureFunctions;
 
-    ir = jl_uncompress_ir(closure_method, ci, (jl_array_t*)ci->inferred);
+    ir = jl_uncompress_ir(closure_method, ci, (jl_array_t*)inferred);
 
     // TODO: Emit this inline and outline it late using LLVM's coroutine support.
     orc::ThreadSafeModule closure_m = jl_create_llvm_module(
@@ -8084,7 +8086,7 @@ jl_llvm_functions_t jl_emit_codeinst(
     JL_TIMING(CODEGEN);
     JL_GC_PUSH1(&src);
     if (!src) {
-        src = (jl_code_info_t*)codeinst->inferred;
+        src = (jl_code_info_t*)jl_atomic_load_relaxed(&codeinst->inferred);
         jl_method_t *def = codeinst->def->def.method;
         if (src && (jl_value_t*)src != jl_nothing && jl_is_method(def))
             src = jl_uncompress_ir(def, codeinst, (jl_array_t*)src);
@@ -8118,33 +8120,33 @@ jl_llvm_functions_t jl_emit_codeinst(
         if (// don't alter `inferred` when the code is not directly being used
             params.world &&
             // don't change inferred state
-            codeinst->inferred) {
+            jl_atomic_load_relaxed(&codeinst->inferred)) {
             jl_method_t *def = codeinst->def->def.method;
             if (// keep code when keeping everything
                 !(JL_DELETE_NON_INLINEABLE) ||
                 // aggressively keep code when debugging level >= 2
                 jl_options.debug_level > 1) {
                 // update the stored code
-                if (codeinst->inferred != (jl_value_t*)src) {
+                if (jl_atomic_load_relaxed(&codeinst->inferred) != (jl_value_t*)src) {
                     if (jl_is_method(def)) {
                         src = (jl_code_info_t*)jl_compress_ir(def, src);
                         assert(jl_typeis(src, jl_array_uint8_type));
                         codeinst->relocatability = ((uint8_t*)jl_array_data(src))[jl_array_len(src)-1];
                     }
-                    codeinst->inferred = (jl_value_t*)src;
+                    jl_atomic_store_relaxed(&codeinst->inferred, (jl_value_t*)src);
                     jl_gc_wb(codeinst, src);
                 }
             }
             else if (// don't delete toplevel code
                      jl_is_method(def) &&
                      // and there is something to delete (test this before calling jl_ir_flag_inlineable)
-                     codeinst->inferred != jl_nothing &&
+                     jl_atomic_load_relaxed(&codeinst->inferred) != jl_nothing &&
                      // don't delete inlineable code, unless it is constant
-                     (codeinst->invoke == jl_fptr_const_return_addr || !jl_ir_flag_inlineable((jl_array_t*)codeinst->inferred)) &&
+                     (codeinst->invoke == jl_fptr_const_return_addr || !jl_ir_flag_inlineable((jl_array_t*)jl_atomic_load_relaxed(&codeinst->inferred))) &&
                      // don't delete code when generating a precompile file
                      !(params.imaging || jl_options.incremental)) {
                 // if not inlineable, code won't be needed again
-                codeinst->inferred = jl_nothing;
+                jl_atomic_store_relaxed(&codeinst->inferred, jl_nothing);
             }
         }
     }
@@ -8197,7 +8199,7 @@ void jl_compile_workqueue(
                 // Reinfer the function. The JIT came along and removed the inferred
                 // method body. See #34993
                 if (policy != CompilationPolicy::Default &&
-                    codeinst->inferred && codeinst->inferred == jl_nothing) {
+                    jl_atomic_load_relaxed(&codeinst->inferred) && jl_atomic_load_relaxed(&codeinst->inferred) == jl_nothing) {
                     src = jl_type_infer(codeinst->def, jl_atomic_load_acquire(&jl_world_counter), 0);
                     if (src) {
                         orc::ThreadSafeModule result_m =