[Compiler] fix some cycle_fix_limited usage (#57545)

Recompute some O(n) things a bit less on every statement. Fix an assert
that ensured LimitedAccuracy does not accidentally get preserved when it
should have been deleted: the (local) cache should not contain things
that are marked as dead (RIP), as that was leading to much code not
getting cached when it logically should have. Simplify the computation
of frame_parent when the cycle_parent isn't needed.

Author: vtjnash

Compiler/src/abstractinterpretation.jl

@@ -286,19 +286,12 @@ function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(fun
                 state.rettype = Any
             end
             # if from_interprocedural added any pclimitations to the set inherited from the arguments,
-            # some of those may be part of our cycles, so those can be deleted now
-            # TODO: and those might need to be deleted later too if the cycle grows to include them?
             if isa(sv, InferenceState)
                 # TODO (#48913) implement a proper recursion handling for irinterp:
-                # This works just because currently the `:terminate` condition guarantees that
-                # irinterp doesn't fail into unresolved cycles, but it's not a good solution.
+                # This works most of the time just because currently the `:terminate` condition often guarantees that
+                # irinterp doesn't fail into unresolved cycles, but it is not a good (or working) solution.
                 # We should revisit this once we have a better story for handling cycles in irinterp.
-                if !isempty(sv.pclimitations) # remove self, if present
-                    delete!(sv.pclimitations, sv)
-                    for caller in callers_in_cycle(sv)
-                        delete!(sv.pclimitations, caller)
-                    end
-                end
+                delete!(sv.pclimitations, sv) # remove self, if present
             end
         else
             # there is unanalyzed candidate, widen type and effects to the top
@@ -775,7 +768,7 @@ function edge_matches_sv(interp::AbstractInterpreter, frame::AbsIntState,
         # check in the cycle list first
         # all items in here are considered mutual parents of all others
         if !any(p::AbsIntState->matches_sv(p, sv), callers_in_cycle(frame))
-            let parent = frame_parent(frame)
+            let parent = cycle_parent(frame)
                 parent === nothing && return false
                 (is_cached(parent) || frame_parent(parent) !== nothing) || return false
                 matches_sv(parent, sv) || return false
@@ -1379,6 +1372,7 @@ function const_prop_call(interp::AbstractInterpreter,
         inf_result.result = concrete_eval_result.rt
         inf_result.ipo_effects = concrete_eval_result.effects
     end
+    typ = inf_result.result
     return const_prop_result(inf_result)
 end
 

Compiler/src/inferenceresult.jl

@@ -183,7 +183,8 @@ function cache_lookup(𝕃::AbstractLattice, mi::MethodInstance, given_argtypes:
     method = mi.def::Method
     nargtypes = length(given_argtypes)
     for cached_result in cache
-        cached_result.linfo === mi || @goto next_cache
+        cached_result.tombstone && continue # ignore deleted entries (due to LimitedAccuracy)
+        cached_result.linfo === mi || continue
         cache_argtypes = cached_result.argtypes
         @assert length(cache_argtypes) == nargtypes "invalid `cache_argtypes` for `mi`"
         cache_overridden_by_const = cached_result.overridden_by_const::BitVector

Compiler/src/inferencestate.jl

@@ -292,7 +292,7 @@ mutable struct InferenceState
 
     # IPO tracking of in-process work, shared with all frames given AbstractInterpreter
     callstack #::Vector{AbsIntState}
-    parentid::Int # index into callstack of the parent frame that originally added this frame (call frame_parent to extract the current parent of the SCC)
+    parentid::Int # index into callstack of the parent frame that originally added this frame (call cycle_parent to extract the current parent of the SCC)
     frameid::Int # index into callstack at which this object is found (or zero, if this is not a cached frame and has no parent)
     cycleid::Int # index into the callstack of the topmost frame in the cycle (all frames in the same cycle share the same cycleid)
 
@@ -908,14 +908,17 @@ function frame_module(sv::AbsIntState)
     return def.module
 end
 
-function frame_parent(sv::InferenceState)
+frame_parent(sv::AbsIntState) = sv.parentid == 0 ? nothing : (sv.callstack::Vector{AbsIntState})[sv.parentid]
+
+function cycle_parent(sv::InferenceState)
     sv.parentid == 0 && return nothing
     callstack = sv.callstack::Vector{AbsIntState}
     sv = callstack[sv.cycleid]::InferenceState
     sv.parentid == 0 && return nothing
     return callstack[sv.parentid]
 end
-frame_parent(sv::IRInterpretationState) = sv.parentid == 0 ? nothing : (sv.callstack::Vector{AbsIntState})[sv.parentid]
+cycle_parent(sv::IRInterpretationState) = frame_parent(sv)
+
 
 # add the orphan child to the parent and the parent to the child
 function assign_parentchild!(child::InferenceState, parent::AbsIntState)

Compiler/src/ssair/irinterp.jl

@@ -1,7 +1,7 @@
 # This file is a part of Julia. License is MIT: https://julialang.org/license
 
 function collect_limitations!(@nospecialize(typ), ::IRInterpretationState)
-    @assert !isa(typ, LimitedAccuracy) "irinterp is unable to handle heavy recursion"
+    @assert !isa(typ, LimitedAccuracy) "irinterp is unable to handle heavy recursion correctly"
     return typ
 end
 
@@ -212,6 +212,7 @@ function reprocess_instruction!(interp::AbstractInterpreter, inst::Instruction,
     else
         rt = argextype(stmt, irsv.ir)
     end
+    @assert !(rt isa LimitedAccuracy)
     if rt !== nothing
         if has_flag(inst, IR_FLAG_UNUSED)
             # Don't bother checking the type if we know it's unused

Compiler/src/typeinfer.jl

@@ -195,7 +195,7 @@ function finish!(interp::AbstractInterpreter, mi::MethodInstance, ci::CodeInstan
 end
 
 function finish_nocycle(::AbstractInterpreter, frame::InferenceState)
-    finishinfer!(frame, frame.interp)
+    finishinfer!(frame, frame.interp, frame.cycleid)
     opt = frame.result.src
     if opt isa OptimizationState # implies `may_optimize(caller.interp) === true`
         optimize(frame.interp, opt, frame.result)
@@ -232,7 +232,7 @@ function finish_cycle(::AbstractInterpreter, frames::Vector{AbsIntState}, cyclei
     for frameid = cycleid:length(frames)
         caller = frames[frameid]::InferenceState
         adjust_cycle_frame!(caller, cycle_valid_worlds, cycle_valid_effects)
-        finishinfer!(caller, caller.interp)
+        finishinfer!(caller, caller.interp, cycleid)
     end
     for frameid = cycleid:length(frames)
         caller = frames[frameid]::InferenceState
@@ -312,26 +312,21 @@ function cache_result!(interp::AbstractInterpreter, result::InferenceResult, ci:
     return true
 end
 
-function cycle_fix_limited(@nospecialize(typ), sv::InferenceState)
+function cycle_fix_limited(@nospecialize(typ), sv::InferenceState, cycleid::Int)
     if typ isa LimitedAccuracy
-        if sv.parentid === 0
-            # we might have introduced a limit marker, but we should know it must be sv and other callers_in_cycle
-            #@assert !isempty(callers_in_cycle(sv))
-            #  FIXME: this assert fails, appearing to indicate there is a bug in filtering this list earlier.
-            #  In particular (during doctests for example), during inference of
-            #  show(Base.IOContext{Base.GenericIOBuffer{Memory{UInt8}}}, Base.Multimedia.MIME{:var"text/plain"}, LinearAlgebra.BunchKaufman{Float64, Array{Float64, 2}, Array{Int64, 1}})
-            #  we observed one of the ssavaluetypes here to be Core.Compiler.LimitedAccuracy(typ=Any, causes=Core.Compiler.IdSet(getproperty(LinearAlgebra.BunchKaufman{Float64, Array{Float64, 2}, Array{Int64, 1}}, Symbol)))
-            return typ.typ
-        end
-        causes = copy(typ.causes)
-        delete!(causes, sv)
-        for caller in callers_in_cycle(sv)
-            delete!(causes, caller)
-        end
-        if isempty(causes)
-            return typ.typ
+        frames = sv.callstack::Vector{AbsIntState}
+        causes = typ.causes
+        for frameid = cycleid:length(frames)
+            caller = frames[frameid]::InferenceState
+            caller in causes || continue
+            causes === typ.causes && (causes = copy(causes))
+            pop!(causes, caller)
+            if isempty(causes)
+                return typ.typ
+            end
         end
-        if length(causes) != length(typ.causes)
+        @assert sv.parentid != 0
+        if causes !== typ.causes
             return LimitedAccuracy(typ.typ, causes)
         end
     end
@@ -439,20 +434,23 @@ const empty_edges = Core.svec()
 
 # inference completed on `me`
 # update the MethodInstance
-function finishinfer!(me::InferenceState, interp::AbstractInterpreter)
+function finishinfer!(me::InferenceState, interp::AbstractInterpreter, cycleid::Int)
     # prepare to run optimization passes on fulltree
     @assert isempty(me.ip)
     # inspect whether our inference had a limited result accuracy,
     # else it may be suitable to cache
-    bestguess = me.bestguess = cycle_fix_limited(me.bestguess, me)
-    exc_bestguess = me.exc_bestguess = cycle_fix_limited(me.exc_bestguess, me)
+    bestguess = me.bestguess = cycle_fix_limited(me.bestguess, me, cycleid)
+    exc_bestguess = me.exc_bestguess = cycle_fix_limited(me.exc_bestguess, me, cycleid)
     limited_ret = bestguess isa LimitedAccuracy || exc_bestguess isa LimitedAccuracy
     limited_src = false
-    if !limited_ret
+    if limited_ret
+        @assert me.parentid != 0
+    else
         gt = me.ssavaluetypes
         for j = 1:length(gt)
-            gt[j] = gtj = cycle_fix_limited(gt[j], me)
-            if gtj isa LimitedAccuracy && me.parentid != 0
+            gt[j] = gtj = cycle_fix_limited(gt[j], me, cycleid)
+            if gtj isa LimitedAccuracy
+                @assert me.parentid != 0
                 limited_src = true
                 break
             end
@@ -474,6 +472,7 @@ function finishinfer!(me::InferenceState, interp::AbstractInterpreter)
         # a parent may be cached still, but not this intermediate work:
         # we can throw everything else away now
         result.src = nothing
+        result.tombstone = true
         me.cache_mode = CACHE_MODE_NULL
         set_inlineable!(me.src, false)
     elseif limited_src
@@ -714,7 +713,7 @@ function merge_call_chain!(::AbstractInterpreter, parent::InferenceState, child:
         add_cycle_backedge!(parent, child)
         parent.cycleid === ancestorid && break
         child = parent
-        parent = frame_parent(child)::InferenceState
+        parent = cycle_parent(child)::InferenceState
     end
     # ensure that walking the callstack has the same cycleid (DAG)
     for frameid = reverse(ancestorid:length(frames))
@@ -750,7 +749,7 @@ end
 # returned instead.
 function resolve_call_cycle!(interp::AbstractInterpreter, mi::MethodInstance, parent::AbsIntState)
     # TODO (#48913) implement a proper recursion handling for irinterp:
-    # This works currently just because the irinterp code doesn't get used much with
+    # This works most of the time currently just because the irinterp code doesn't get used much with
     # `@assume_effects`, so it never sees a cycle normally, but that may not be a sustainable solution.
     parent isa InferenceState || return false
     frames = parent.callstack::Vector{AbsIntState}
@@ -762,7 +761,7 @@ function resolve_call_cycle!(interp::AbstractInterpreter, mi::MethodInstance, pa
         if is_same_frame(interp, mi, frame)
             if uncached
                 # our attempt to speculate into a constant call lead to an undesired self-cycle
-                # that cannot be converged: poison our call-stack (up to the discovered duplicate frame)
+                # that cannot be converged: if necessary, poison our call-stack (up to the discovered duplicate frame)
                 # with the limited flag and abort (set return type to Any) now
                 poison_callstack!(parent, frame)
                 return true

Compiler/src/types.jl

@@ -110,6 +110,7 @@ mutable struct InferenceResult
     effects::Effects                  # if optimization is finished
     analysis_results::AnalysisResults # AnalysisResults with e.g. result::ArgEscapeCache if optimized, otherwise NULL_ANALYSIS_RESULTS
     is_src_volatile::Bool             # `src` has been cached globally as the compressed format already, allowing `src` to be used destructively
+    tombstone::Bool
 
     #=== uninitialized fields ===#
     ci::CodeInstance                  # CodeInstance if this result may be added to the cache
@@ -120,7 +121,7 @@ mutable struct InferenceResult
         ipo_effects = effects = Effects()
         analysis_results = NULL_ANALYSIS_RESULTS
         return new(mi, argtypes, overridden_by_const, result, exc_result, src,
-            valid_worlds, ipo_effects, effects, analysis_results, #=is_src_volatile=#false)
+            valid_worlds, ipo_effects, effects, analysis_results, #=is_src_volatile=#false, false)
     end
 end
 function InferenceResult(mi::MethodInstance, 𝕃::AbstractLattice=fallback_lattice)