De-duplicate edges in typeinfer instead of gf.c (#58117)

Without this PR, my system spends ~1.07 seconds just running
`store_backedges` when doing `using CairoMakie`

With this change, that drops to `0.641 seconds`

That's still not fast enough for me, but we do call this function
236,092 times so maybe it's understandable.

Author: topolarity

Compiler/src/typeinfer.jl

@@ -687,12 +687,17 @@ function finishinfer!(me::InferenceState, interp::AbstractInterpreter, cycleid::
     nothing
 end
 
-# record the backedges
-function store_backedges(caller::CodeInstance, edges::SimpleVector)
-    isa(caller.def.def, Method) || return # don't add backedges to toplevel method instance
-    i = 1
-    while true
-        i > length(edges) && return nothing
+# Iterate a series of back-edges that need registering, based on the provided forward edge list.
+# Back-edges are returned as (invokesig, item), where the item is a Binding, MethodInstance, or
+# MethodTable.
+struct ForwardToBackedgeIterator
+    forward_edges::SimpleVector
+end
+
+function Base.iterate(it::ForwardToBackedgeIterator, i::Int = 1)
+    edges = it.forward_edges
+    i > length(edges) && return nothing
+    while i ≤ length(edges)
         item = edges[i]
         if item isa Int
             i += 2
@@ -702,34 +707,55 @@ function store_backedges(caller::CodeInstance, edges::SimpleVector)
             i += 1
             continue
         elseif isa(item, Core.Binding)
-            i += 1
-            maybe_add_binding_backedge!(item, caller)
-            continue
+            return ((nothing, item), i + 1)
         end
         if isa(item, CodeInstance)
-            item = item.def
-        end
-        if isa(item, MethodInstance) # regular dispatch
-            ccall(:jl_method_instance_add_backedge, Cvoid, (Any, Any, Any), item, nothing, caller)
-            i += 1
+            item = get_ci_mi(item)
+            return ((nothing, item), i + 1)
+        elseif isa(item, MethodInstance) # regular dispatch
+            return ((nothing, item), i + 1)
         else
+            invokesig = item
             callee = edges[i+1]
-            if isa(callee, MethodTable) # abstract dispatch (legacy style edges)
-                ccall(:jl_method_table_add_backedge, Cvoid, (Any, Any, Any), callee, item, caller)
-                i += 2
-                continue
-            elseif isa(callee, Method)
-                # ignore `Method`-edges (from e.g. failed `abstract_call_method`)
-                i += 2
-                continue
-            # `invoke` edge
-            elseif isa(callee, CodeInstance)
-                callee = get_ci_mi(callee)
+            isa(callee, Method) && (i += 2; continue) # ignore `Method`-edges (from e.g. failed `abstract_call_method`)
+            if isa(callee, MethodTable)
+                # abstract dispatch (legacy style edges)
+                return ((invokesig, callee), i + 2)
             else
-                callee = callee::MethodInstance
+                # `invoke` edge
+                callee = isa(callee, CodeInstance) ? get_ci_mi(callee) : callee::MethodInstance
+                return ((invokesig, callee), i + 2)
+            end
+        end
+    end
+    return nothing
+end
+
+# record the backedges
+function store_backedges(caller::CodeInstance, edges::SimpleVector)
+    isa(caller.def.def, Method) || return # don't add backedges to toplevel method instance
+
+    backedges = ForwardToBackedgeIterator(edges)
+    for (i, (invokesig, item)) in enumerate(backedges)
+        # check for any duplicate edges we've already registered
+        duplicate_found = false
+        for (i′, (invokesig′, item′)) in enumerate(backedges)
+            i == i′ && break
+            if item′ === item && invokesig′ == invokesig
+                duplicate_found = true
+                break
+            end
+        end
+
+        if !duplicate_found
+            if item isa Core.Binding
+                maybe_add_binding_backedge!(item, caller)
+            elseif item isa MethodTable
+                ccall(:jl_method_table_add_backedge, Cvoid, (Any, Any, Any), item, invokesig, caller)
+            else
+                item::MethodInstance
+                ccall(:jl_method_instance_add_backedge, Cvoid, (Any, Any, Any), item, invokesig, caller)
             end
-            ccall(:jl_method_instance_add_backedge, Cvoid, (Any, Any, Any), callee, item, caller)
-            i += 2
         end
     end
     nothing

Compiler/test/invalidation.jl

@@ -177,25 +177,24 @@ begin
     end
 
     # Verify that adding the backedge again does not actually add a new backedge
-    let mi1 = Base.method_instance(deduped_caller1, (Int,)),
-        mi2 = Base.method_instance(deduped_caller2, (Int,)),
-        ci1 = mi1.cache
-        ci2 = mi2.cache
+    let mi = Base.method_instance(deduped_caller1, (Int,)),
+        ci = mi.cache
 
         callee_mi = Base.method_instance(deduped_callee, (Int,))
 
         # Inference should have added the callers to the callee's backedges
-        @test ci1 in callee_mi.backedges
-        @test ci2 in callee_mi.backedges
+        @test ci in callee_mi.backedges
 
+        # In practice, inference will never end up calling `store_backedges`
+        # twice on the same CodeInstance like this - we only need to check
+        # that de-duplication works for a single invocation
         N = length(callee_mi.backedges)
-        Core.Compiler.store_backedges(ci1, Core.svec(callee_mi))
-        Core.Compiler.store_backedges(ci2, Core.svec(callee_mi))
+        Core.Compiler.store_backedges(ci, Core.svec(callee_mi, callee_mi))
         N′ = length(callee_mi.backedges)
 
-        # The number of backedges should not be affected by an additional store,
-        # since de-duplication should have noticed the edge is already tracked
-        @test N == N′
+        # A single `store_backedges` invocation should de-duplicate any of the
+        # edges it is adding.
+        @test N′ - N == 1
     end
 end
 

src/gf.c

@@ -1997,7 +1997,6 @@ JL_DLLEXPORT void jl_method_instance_add_backedge(jl_method_instance_t *callee,
     assert(invokesig == NULL || jl_is_type(invokesig));
     JL_LOCK(&callee->def.method->writelock);
     if (jl_atomic_load_relaxed(&allow_new_worlds)) {
-        int found = 0;
         jl_array_t *backedges = jl_mi_get_backedges(callee);
         // TODO: use jl_cache_type_(invokesig) like cache_method does to save memory
         if (!backedges) {
@@ -2006,25 +2005,7 @@ JL_DLLEXPORT void jl_method_instance_add_backedge(jl_method_instance_t *callee,
             callee->backedges = backedges;
             jl_gc_wb(callee, backedges);
         }
-        else {
-            size_t i = 0, l = jl_array_nrows(backedges);
-            for (i = 0; i < l; i++) {
-                // optimized version of while (i < l) i = get_next_edge(callee->backedges, i, &invokeTypes, &mi);
-                jl_value_t *ciedge = jl_array_ptr_ref(backedges, i);
-                if (ciedge != (jl_value_t*)caller)
-                    continue;
-                jl_value_t *invokeTypes = i > 0 ? jl_array_ptr_ref(backedges, i - 1) : NULL;
-                if (invokeTypes && jl_is_code_instance(invokeTypes))
-                    invokeTypes = NULL;
-                if ((invokesig == NULL && invokeTypes == NULL) ||
-                    (invokesig && invokeTypes && jl_types_equal(invokesig, invokeTypes))) {
-                    found = 1;
-                    break;
-                }
-            }
-        }
-        if (!found)
-            push_edge(backedges, invokesig, caller);
+        push_edge(backedges, invokesig, caller);
     }
     JL_UNLOCK(&callee->def.method->writelock);
 }
@@ -2047,14 +2028,6 @@ JL_DLLEXPORT void jl_method_table_add_backedge(jl_methtable_t *mt, jl_value_t *t
         else {
             // check if the edge is already present and avoid adding a duplicate
             size_t i, l = jl_array_nrows(mt->backedges);
-            for (i = 1; i < l; i += 2) {
-                if (jl_array_ptr_ref(mt->backedges, i) == (jl_value_t*)caller) {
-                    if (jl_types_equal(jl_array_ptr_ref(mt->backedges, i - 1), typ)) {
-                        JL_UNLOCK(&mt->writelock);
-                        return;
-                    }
-                }
-            }
             // reuse an already cached instance of this type, if possible
             // TODO: use jl_cache_type_(tt) like cache_method does, instead of this linear scan?
             for (i = 1; i < l; i += 2) {

src/staticdata.c

@@ -785,6 +785,12 @@ static void jl_insert_into_serialization_queue(jl_serializer_state *s, jl_value_
         }
         goto done_fields; // for now
     }
+    if (s->incremental && jl_is_mtable(v)) {
+        jl_methtable_t *mt = (jl_methtable_t *)v;
+        // Any back-edges will be re-validated and added by staticdata.jl, so
+        // drop them from the image here
+        record_field_change((jl_value_t**)&mt->backedges, NULL);
+    }
     if (jl_is_method_instance(v)) {
         jl_method_instance_t *mi = (jl_method_instance_t*)v;
         if (s->incremental) {
@@ -800,12 +806,14 @@ static void jl_insert_into_serialization_queue(jl_serializer_state *s, jl_value_
                 // we only need 3 specific fields of this (the rest are restored afterward, if valid)
                 // in particular, cache is repopulated by jl_mi_cache_insert for all foreign function,
                 // so must not be present here
-                record_field_change((jl_value_t**)&mi->backedges, NULL);
                 record_field_change((jl_value_t**)&mi->cache, NULL);
             }
             else {
                 assert(!needs_recaching(v, s->query_cache));
             }
+            // Any back-edges will be re-validated and added by staticdata.jl, so
+            // drop them from the image here
+            record_field_change((jl_value_t**)&mi->backedges, NULL);
             // n.b. opaque closures cannot be inspected and relied upon like a
             // normal method since they can get improperly introduced by generated
             // functions, so if they appeared at all, we will probably serialize