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

This implements the optimizations I promised in #57602 by checking
in invalidation whether or not the information that inference sees
changes. The primary situation in which this is useful is avoiding
an invalidation for `export` flag changes or changes of which module
a binding is being imported from that do not change what the actual
binding being imported is. Base itself uses these patterns sparingly,
so the bootstrap impact over #57615 is minimal (though non-zero).

Author: Keno

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)

base/invalidation.jl

@@ -113,33 +113,55 @@ 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 (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)
+
+    (_, (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
-    end
-    if isdefined(b, :backedges)
-        for edge in b.backedges
-            if isa(edge, CodeInstance)
-                ccall(:jl_invalidate_code_instance, Cvoid, (Any, UInt), edge, new_max_world)
-            elseif isa(edge, Core.Binding)
-                isdefined(edge, :partitions) || continue
-                latest_bpart = edge.partitions
-                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)
-            else
-                invalidate_method_for_globalref!(gr, edge::Method, invalidated_bpart, new_max_world)
+        if isdefined(b, :backedges)
+            for edge in b.backedges
+                if isa(edge, CodeInstance)
+                    ccall(:jl_invalidate_code_instance, Cvoid, (Any, UInt), edge, new_max_world)
+                elseif isa(edge, Core.Binding)
+                    isdefined(edge, :partitions) || continue
+                    latest_bpart = edge.partitions
+                    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, 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)
@@ -151,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

src/module.c

@@ -106,6 +106,9 @@ void jl_check_new_binding_implicit(
             if (tempbmax_world < max_world)
                 max_world = tempbmax_world;
 
+            // N.B.: Which aspects of the partition are considered here needs to
+            // be kept in sync with `export_affecting_partition_flags` in the
+            // invalidation code.
             if ((tempbpart->kind & PARTITION_FLAG_EXPORTED) == 0)
                 continue;
 
@@ -165,6 +168,29 @@ void jl_check_new_binding_implicit(
     return;
 }
 
+JL_DLLEXPORT jl_binding_partition_t *jl_maybe_reresolve_implicit(jl_binding_t *b, size_t new_max_world)
+{
+    jl_binding_partition_t *new_bpart = new_binding_partition();
+    jl_binding_partition_t *bpart = jl_atomic_load_acquire(&b->partitions);
+    assert(bpart);
+    while (1) {
+        jl_atomic_store_relaxed(&new_bpart->next, bpart);
+        jl_gc_wb(new_bpart, bpart);
+        jl_check_new_binding_implicit(new_bpart, b, NULL, new_max_world+1);
+        JL_GC_PROMISE_ROOTED(new_bpart); // TODO: Analyzer doesn't understand MAYBE_UNROOTED properly
+        if (bpart->kind & PARTITION_FLAG_EXPORTED)
+            new_bpart->kind |= PARTITION_FLAG_EXPORTED;
+        if (new_bpart->kind == bpart->kind && new_bpart->restriction == bpart->restriction)
+            return bpart;
+        // Resolution changed, insert the new partition
+        size_t expected_max_world = ~(size_t)0;
+        if (jl_atomic_cmpswap(&bpart->max_world, &expected_max_world, new_max_world) &&
+            jl_atomic_cmpswap(&b->partitions, &bpart, new_bpart))
+            break;
+    }
+    return new_bpart;
+}
+
 STATIC_INLINE jl_binding_partition_t *jl_get_binding_partition_(jl_binding_t *b JL_PROPAGATES_ROOT, jl_value_t *parent, _Atomic(jl_binding_partition_t *)*insert, size_t world, modstack_t *st) JL_GLOBALLY_ROOTED
 {
     assert(jl_is_binding(b));
@@ -183,7 +209,7 @@ STATIC_INLINE jl_binding_partition_t *jl_get_binding_partition_(jl_binding_t *b
         if (!new_bpart)
             new_bpart = new_binding_partition();
         jl_atomic_store_relaxed(&new_bpart->next, bpart);
-        jl_gc_wb_fresh(new_bpart, bpart);
+        jl_gc_wb(new_bpart, bpart); // Not fresh the second time around the loop
         new_bpart->min_world = bpart ? jl_atomic_load_relaxed(&bpart->max_world) + 1 : 0;
         jl_atomic_store_relaxed(&new_bpart->max_world, max_world);
         JL_GC_PROMISE_ROOTED(new_bpart); // TODO: Analyzer doesn't understand MAYBE_UNROOTED properly
@@ -1112,10 +1138,12 @@ JL_DLLEXPORT void jl_module_using(jl_module_t *to, jl_module_t *from)
             jl_sym_t *var = b->globalref->name;
             jl_binding_t *tob = jl_get_module_binding(to, var, 0);
             if (tob) {
-                jl_binding_partition_t *tobpart = jl_get_binding_partition(tob, new_world);
-                enum jl_partition_kind kind = jl_binding_kind(tobpart);
-                if (jl_bkind_is_some_implicit(kind)) {
-                    jl_replace_binding_locked(tob, tobpart, NULL, PARTITION_KIND_IMPLICIT_RECOMPUTE, new_world);
+                jl_binding_partition_t *tobpart = jl_atomic_load_relaxed(&tob->partitions);
+                if (tobpart) {
+                    enum jl_partition_kind kind = jl_binding_kind(tobpart);
+                    if (jl_bkind_is_some_implicit(kind)) {
+                        jl_replace_binding_locked(tob, tobpart, NULL, PARTITION_KIND_IMPLICIT_RECOMPUTE, new_world);
+                    }
                 }
             }
         }
@@ -1380,20 +1408,24 @@ JL_DLLEXPORT jl_binding_partition_t *jl_replace_binding_locked2(jl_binding_t *b,
         jl_atomic_store_relaxed(&jl_first_image_replacement_world, new_world);
 
     assert(jl_atomic_load_relaxed(&b->partitions) == old_bpart);
-    jl_atomic_store_release(&old_bpart->max_world, new_world-1);
     jl_binding_partition_t *new_bpart = new_binding_partition();
     JL_GC_PUSH1(&new_bpart);
     new_bpart->min_world = new_world;
     if ((kind & PARTITION_MASK_KIND) == PARTITION_KIND_IMPLICIT_RECOMPUTE) {
         assert(!restriction_val);
         jl_check_new_binding_implicit(new_bpart /* callee rooted */, b, NULL, new_world);
         new_bpart->kind |= kind & PARTITION_MASK_FLAG;
+        if (new_bpart->kind == old_bpart->kind && new_bpart->restriction == old_bpart->restriction) {
+            JL_GC_POP();
+            return old_bpart;
+        }
     }
     else {
         new_bpart->kind = kind;
         new_bpart->restriction = restriction_val;
         jl_gc_wb_fresh(new_bpart, restriction_val);
     }
+    jl_atomic_store_release(&old_bpart->max_world, new_world-1);
     jl_atomic_store_relaxed(&new_bpart->next, old_bpart);
     jl_gc_wb_fresh(new_bpart, old_bpart);