bpart: Don't invalidate for changes that inference can't see

Compiler/src/abstractinterpretation.jl

@@ -3633,18 +3633,20 @@ scan_partitions(query::Function, interp, g::GlobalRef, wwr::WorldWithRange) =
 abstract_load_all_consistent_leaf_partitions(interp, g::GlobalRef, wwr::WorldWithRange) =
     scan_leaf_partitions(abstract_eval_partition_load, interp, g, wwr)
 
+function abstract_eval_globalref_partition(interp, binding::Core.Binding, partition::Core.BindingPartition)
+    # For inference purposes, we don't particularly care which global binding we end up loading, we only
+    # care about its type. However, we would still like to terminate the world range for the particular
+    # binding we end up reaching such that codegen can emit a simpler pointer load.
+    Pair{RTEffects, Union{Nothing, Core.Binding}}(
+        abstract_eval_partition_load(interp, partition),
+        binding_kind(partition) in (PARTITION_KIND_GLOBAL, PARTITION_KIND_DECLARED) ? binding : nothing)
+end
+
 function abstract_eval_globalref(interp, g::GlobalRef, saw_latestworld::Bool, sv::AbsIntState)
     if saw_latestworld
         return RTEffects(Any, Any, generic_getglobal_effects)
     end
-    (valid_worlds, (ret, binding_if_global)) = scan_leaf_partitions(interp, g, sv.world) do interp, binding, partition
-        # For inference purposes, we don't particularly care which global binding we end up loading, we only
-        # care about its type. However, we would still like to terminate the world range for the particular
-        # binding we end up reaching such that codegen can emit a simpler pointer load.
-        Pair{RTEffects, Union{Nothing, Core.Binding}}(
-            abstract_eval_partition_load(interp, partition),
-            binding_kind(partition) in (PARTITION_KIND_GLOBAL, PARTITION_KIND_DECLARED) ? binding : nothing)
-    end
+    (valid_worlds, (ret, binding_if_global)) = scan_leaf_partitions(abstract_eval_globalref_partition, interp, g, sv.world)
     update_valid_age!(sv, valid_worlds)
     if ret.rt !== Union{} && ret.exct === UndefVarError && binding_if_global !== nothing && InferenceParams(interp).assume_bindings_static
         if isdefined(binding_if_global, :value)

Compiler/src/utilities.jl

@@ -129,6 +129,25 @@ function retrieve_code_info(mi::MethodInstance, world::UInt)
         else
             c = copy(src::CodeInfo)
         end
+        if !def.did_scan_source
+            # This scan must happen:
+            #   1. After method definition
+            #   2. Before any code instances that may have relied on information
+            #      from implicit GlobalRefs for this method are added to the cache
+            #   3. Preferably while the IR is already uncompressed
+            #   4. As late as possible, as early adding of the backedges may cause
+            #      spurious invalidations.
+            #
+            # At the moment we do so here, because
+            #  1. It's reasonably late
+            #  2. It has easy access to the uncompressed IR
+            #  3. We necessarily pass through here before relying on any
+            #     information obtained from implicit GlobalRefs.
+            #
+            # However, the exact placement of this scan is not as important as
+            # long as the above conditions are met.
+            ccall(:jl_scan_method_source_now, Cvoid, (Any, Any), def, c)
+        end
     end
     if c isa CodeInfo
         c.parent = mi

base/expr.jl

@@ -1352,6 +1352,10 @@ function make_atomic(order, ex)
                 op = :+
             elseif ex.head === :(-=)
                 op = :-
+            elseif ex.head === :(|=)
+                op = :|
+            elseif ex.head === :(&=)
+                op = :&
             elseif @isdefined string
                 shead = string(ex.head)
                 if endswith(shead, '=')

base/invalidation.jl

@@ -113,15 +113,35 @@ function invalidate_method_for_globalref!(gr::GlobalRef, method::Method, invalid
     end
 end
 
-function invalidate_code_for_globalref!(b::Core.Binding, invalidated_bpart::Core.BindingPartition, new_bpart::Union{Core.BindingPartition, Nothing}, new_max_world::UInt)
+export_affecting_partition_flags(bpart::Core.BindingPartition) =
+    ((bpart.kind & PARTITION_MASK_KIND) == PARTITION_KIND_GUARD,
+     (bpart.kind & PARTITION_FLAG_EXPORTED) != 0,
+     (bpart.kind & PARTITION_FLAG_DEPRECATED) != 0)
+
+function invalidate_code_for_globalref!(b::Core.Binding, invalidated_bpart::Core.BindingPartition, new_bpart::Core.BindingPartition, new_max_world::UInt)
     gr = b.globalref
-    if !is_some_guard(binding_kind(invalidated_bpart))
-        # TODO: We may want to invalidate for these anyway, since they have performance implications
-        foreach_module_mtable(gr.mod, new_max_world) do mt::Core.MethodTable
-            for method in MethodList(mt)
-                invalidate_method_for_globalref!(gr, method, invalidated_bpart, new_max_world)
+
+    (_, (ib, ibpart)) = Compiler.walk_binding_partition(b, invalidated_bpart, new_max_world)
+    (_, (nb, nbpart)) = Compiler.walk_binding_partition(b, new_bpart, new_max_world+1)
+
+    # abstract_eval_globalref_partition is the maximum amount of information that inference
+    # reads from a binding partition. If this information does not change - we do not need to
+    # invalidate any code that inference created, because we know that the result will not change.
+    need_to_invalidate_code =
+        Compiler.abstract_eval_globalref_partition(nothing, ib, ibpart) !==
+        Compiler.abstract_eval_globalref_partition(nothing, nb, nbpart)
+
+    need_to_invalidate_export = export_affecting_partition_flags(invalidated_bpart) !==
+                                export_affecting_partition_flags(new_bpart)
+
+    if need_to_invalidate_code
+        if (b.flags & BINDING_FLAG_ANY_IMPLICIT_EDGES) != 0
+            foreach_module_mtable(gr.mod, new_max_world) do mt::Core.MethodTable
+                for method in MethodList(mt)
+                    invalidate_method_for_globalref!(gr, method, invalidated_bpart, new_max_world)
+                end
+                return true
             end
-            return true
         end
         if isdefined(b, :backedges)
             for edge in b.backedges
@@ -133,14 +153,15 @@ function invalidate_code_for_globalref!(b::Core.Binding, invalidated_bpart::Core
                     latest_bpart.max_world == typemax(UInt) || continue
                     is_some_imported(binding_kind(latest_bpart)) || continue
                     partition_restriction(latest_bpart) === b || continue
-                    invalidate_code_for_globalref!(edge, latest_bpart, nothing, new_max_world)
+                    invalidate_code_for_globalref!(edge, latest_bpart, latest_bpart, new_max_world)
                 else
                     invalidate_method_for_globalref!(gr, edge::Method, invalidated_bpart, new_max_world)
                 end
             end
         end
     end
-    if (invalidated_bpart.kind & PARTITION_FLAG_EXPORTED != 0) || (new_bpart !== nothing && (new_bpart.kind & PARTITION_FLAG_EXPORTED != 0))
+
+    if need_to_invalidate_code || need_to_invalidate_export
         # This binding was exported - we need to check all modules that `using` us to see if they
         # have a binding that is affected by this change.
         usings_backedges = ccall(:jl_get_module_usings_backedges, Any, (Any,), gr.mod)
@@ -152,8 +173,12 @@ function invalidate_code_for_globalref!(b::Core.Binding, invalidated_bpart::Core
                 latest_bpart = user_binding.partitions
                 latest_bpart.max_world == typemax(UInt) || continue
                 binding_kind(latest_bpart) in (PARTITION_KIND_IMPLICIT, PARTITION_KIND_FAILED, PARTITION_KIND_GUARD) || continue
-                @atomic :release latest_bpart.max_world = new_max_world
-                invalidate_code_for_globalref!(convert(Core.Binding, user_binding), latest_bpart, nothing, new_max_world)
+                new_bpart = need_to_invalidate_export ?
+                    ccall(:jl_maybe_reresolve_implicit, Any, (Any, Any, Csize_t), user_binding, latest_bpart, new_max_world) :
+                    latest_bpart
+                if need_to_invalidate_code || new_bpart !== latest_bpart
+                    invalidate_code_for_globalref!(convert(Core.Binding, user_binding), latest_bpart, new_bpart, new_max_world)
+                end
             end
         end
     end
@@ -166,7 +191,11 @@ gr_needs_backedge_in_module(gr::GlobalRef, mod::Module) = gr.mod !== mod
 # N.B.: This needs to match jl_maybe_add_binding_backedge
 function maybe_add_binding_backedge!(b::Core.Binding, edge::Union{Method, CodeInstance})
     method = isa(edge, Method) ? edge : edge.def.def::Method
-    gr_needs_backedge_in_module(b.globalref, method.module) || return
+    methmod = method.module
+    if !gr_needs_backedge_in_module(b.globalref, methmod)
+        @atomic :acquire_release b.flags |= BINDING_FLAG_ANY_IMPLICIT_EDGES
+        return
+    end
     if !isdefined(b, :backedges)
         b.backedges = Any[]
     end

base/runtime_internals.jl

@@ -216,6 +216,8 @@ const PARTITION_FLAG_DEPWARN      = 0x40
 const PARTITION_MASK_KIND         = 0x0f
 const PARTITION_MASK_FLAG         = 0xf0
 
+const BINDING_FLAG_ANY_IMPLICIT_EDGES = 0x8
+
 is_defined_const_binding(kind::UInt8) = (kind == PARTITION_KIND_CONST || kind == PARTITION_KIND_CONST_IMPORT || kind == PARTITION_KIND_BACKDATED_CONST)
 is_some_const_binding(kind::UInt8) = (is_defined_const_binding(kind) || kind == PARTITION_KIND_UNDEF_CONST)
 is_some_imported(kind::UInt8) = (kind == PARTITION_KIND_IMPLICIT || kind == PARTITION_KIND_EXPLICIT || kind == PARTITION_KIND_IMPORTED)

src/gf.c

@@ -1839,7 +1839,7 @@ JL_DLLEXPORT jl_value_t *jl_debug_method_invalidation(int state)
     return jl_nothing;
 }
 
-static void _invalidate_backedges(jl_method_instance_t *replaced_mi, size_t max_world, int depth);
+static void _invalidate_backedges(jl_method_instance_t *replaced_mi, jl_code_instance_t *replaced_ci, size_t max_world, int depth);
 
 // recursively invalidate cached methods that had an edge to a replaced method
 static void invalidate_code_instance(jl_code_instance_t *replaced, size_t max_world, int depth)
@@ -1864,7 +1864,7 @@ static void invalidate_code_instance(jl_code_instance_t *replaced, size_t max_wo
     }
     assert(jl_atomic_load_relaxed(&replaced->max_world) <= max_world);
     // recurse to all backedges to update their valid range also
-    _invalidate_backedges(replaced_mi, max_world, depth + 1);
+    _invalidate_backedges(replaced_mi, replaced, max_world, depth + 1);
     JL_UNLOCK(&replaced_mi->def.method->writelock);
 }
 
@@ -1873,7 +1873,7 @@ JL_DLLEXPORT void jl_invalidate_code_instance(jl_code_instance_t *replaced, size
     invalidate_code_instance(replaced, max_world, 1);
 }
 
-static void _invalidate_backedges(jl_method_instance_t *replaced_mi, size_t max_world, int depth) {
+static void _invalidate_backedges(jl_method_instance_t *replaced_mi, jl_code_instance_t *replaced_ci, size_t max_world, int depth) {
     jl_array_t *backedges = replaced_mi->backedges;
     if (backedges) {
         // invalidate callers (if any)
@@ -1884,6 +1884,18 @@ static void _invalidate_backedges(jl_method_instance_t *replaced_mi, size_t max_
         while (i < l) {
             i = get_next_edge(backedges, i, NULL, &replaced);
             JL_GC_PROMISE_ROOTED(replaced); // propagated by get_next_edge from backedges
+            if (replaced_ci) {
+                // If we're invalidating a particular codeinstance, only invalidate
+                // this backedge it actually has an edge for our codeinstance.
+                jl_svec_t *edges = jl_atomic_load_relaxed(&replaced->edges);
+                for (size_t j = 0; j < jl_svec_len(edges); ++j) {
+                    jl_value_t *edge = jl_svecref(edges, j);
+                    if (edge == (jl_value_t*)replaced_mi || edge == (jl_value_t*)replaced_ci)
+                        goto found;
+                }
+                continue;
+            found:;
+            }
             invalidate_code_instance(replaced, max_world, depth);
         }
         JL_GC_POP();
@@ -1894,7 +1906,7 @@ static void _invalidate_backedges(jl_method_instance_t *replaced_mi, size_t max_
 static void invalidate_backedges(jl_method_instance_t *replaced_mi, size_t max_world, const char *why)
 {
     JL_LOCK(&replaced_mi->def.method->writelock);
-    _invalidate_backedges(replaced_mi, max_world, 1);
+    _invalidate_backedges(replaced_mi, NULL, max_world, 1);
     JL_UNLOCK(&replaced_mi->def.method->writelock);
     if (why && _jl_debug_method_invalidation) {
         jl_array_ptr_1d_push(_jl_debug_method_invalidation, (jl_value_t*)replaced_mi);

src/jltypes.c

@@ -3275,7 +3275,7 @@ void jl_init_types(void) JL_GC_DISABLED
                         jl_svec(5, jl_any_type/*jl_globalref_type*/, jl_any_type, jl_binding_partition_type,
                                    jl_any_type, jl_uint8_type),
                         jl_emptysvec, 0, 1, 0);
-    const static uint32_t binding_atomicfields[] = { 0x0005 }; // Set fields 2, 3 as atomic
+    const static uint32_t binding_atomicfields[] = { 0x0016 }; // Set fields 2, 3, 5 as atomic
     jl_binding_type->name->atomicfields = binding_atomicfields;
     const static uint32_t binding_constfields[] = { 0x0001 }; // Set fields 1 as constant
     jl_binding_type->name->constfields = binding_constfields;
@@ -3539,7 +3539,7 @@ void jl_init_types(void) JL_GC_DISABLED
     jl_method_type =
         jl_new_datatype(jl_symbol("Method"), core,
                         jl_any_type, jl_emptysvec,
-                        jl_perm_symsvec(31,
+                        jl_perm_symsvec(32,
                             "name",
                             "module",
                             "file",
@@ -3568,10 +3568,11 @@ void jl_init_types(void) JL_GC_DISABLED
                             "isva",
                             "is_for_opaque_closure",
                             "nospecializeinfer",
+                            "did_scan_source",
                             "constprop",
                             "max_varargs",
                             "purity"),
-                        jl_svec(31,
+                        jl_svec(32,
                             jl_symbol_type,
                             jl_module_type,
                             jl_symbol_type,
@@ -3600,6 +3601,7 @@ void jl_init_types(void) JL_GC_DISABLED
                             jl_bool_type,
                             jl_bool_type,
                             jl_bool_type,
+                            jl_bool_type,
                             jl_uint8_type,
                             jl_uint8_type,
                             jl_uint16_type),

src/julia.h

@@ -375,6 +375,8 @@ typedef struct _jl_method_t {
     uint8_t isva;
     uint8_t is_for_opaque_closure;
     uint8_t nospecializeinfer;
+    _Atomic(uint8_t) did_scan_source;
+
     // uint8 settings
     uint8_t constprop;      // 0x00 = use heuristic; 0x01 = aggressive; 0x02 = none
     uint8_t max_varargs;    // 0xFF = use heuristic; otherwise, max # of args to expand
@@ -751,7 +753,10 @@ enum jl_binding_flags {
     BINDING_FLAG_DID_PRINT_BACKDATE_ADMONITION        = 0x1,
     BINDING_FLAG_DID_PRINT_IMPLICIT_IMPORT_ADMONITION = 0x2,
     // `export` is tracked in partitions, but sets this as well
-    BINDING_FLAG_PUBLICP                              = 0x4
+    BINDING_FLAG_PUBLICP                              = 0x4,
+    // Set if any methods defined in this module implicitly reference
+    // this binding. If not, invalidation is optimized.
+    BINDING_FLAG_ANY_IMPLICIT_EDGES                   = 0x8
 };
 
 typedef struct _jl_binding_t {

src/method.c

@@ -39,6 +39,28 @@ static void check_c_types(const char *where, jl_value_t *rt, jl_value_t *at)
     }
 }
 
+JL_DLLEXPORT void jl_scan_method_source_now(jl_method_t *m, jl_value_t *src)
+{
+    if (!jl_atomic_load_relaxed(&m->did_scan_source)) {
+        jl_code_info_t *code = NULL;
+        JL_GC_PUSH1(&code);
+        if (!jl_is_code_info(src))
+            code = jl_uncompress_ir(m, NULL, src);
+        else
+            code = (jl_code_info_t*)src;
+        jl_array_t *stmts = code->code;
+        size_t i, l = jl_array_nrows(stmts);
+        for (i = 0; i < l; i++) {
+            jl_value_t *stmt = jl_array_ptr_ref(stmts, i);
+            if (jl_is_globalref(stmt)) {
+                jl_maybe_add_binding_backedge((jl_globalref_t*)stmt, m->module, (jl_value_t*)m);
+            }
+        }
+        jl_atomic_store_relaxed(&m->did_scan_source, 1);
+        JL_GC_POP();
+    }
+}
+
 // Resolve references to non-locally-defined variables to become references to global
 // variables in `module` (unless the rvalue is one of the type parameters in `sparam_vals`).
 static jl_value_t *resolve_definition_effects(jl_value_t *expr, jl_module_t *module, jl_svec_t *sparam_vals, jl_value_t *binding_edge,
@@ -47,10 +69,7 @@ static jl_value_t *resolve_definition_effects(jl_value_t *expr, jl_module_t *mod
     if (jl_is_symbol(expr)) {
         jl_error("Found raw symbol in code returned from lowering. Expected all symbols to have been resolved to GlobalRef or slots.");
     }
-    if (jl_is_globalref(expr)) {
-        jl_maybe_add_binding_backedge((jl_globalref_t*)expr, module, binding_edge);
-        return expr;
-    }
+
     if (!jl_is_expr(expr)) {
         return expr;
     }
@@ -973,6 +992,7 @@ JL_DLLEXPORT jl_method_t *jl_new_method_uninit(jl_module_t *module)
     jl_atomic_store_relaxed(&m->deleted_world, 1);
     m->is_for_opaque_closure = 0;
     m->nospecializeinfer = 0;
+    jl_atomic_store_relaxed(&m->did_scan_source, 0);
     m->constprop = 0;
     m->purity.bits = 0;
     m->max_varargs = UINT8_MAX;