Skip to content

Commit

Permalink
Add hook in inference recursion resolution for external AbstractInter…
Browse files Browse the repository at this point in the history
…preter (#36401)

This extends hookability to the same-frame comparison in inference's
recursion cycle detection. The case I ran into that made this necessary
is a recursive, nested AD transform. In this case, inference must detect
if two frames have different orders of derivatives (e.g. the primitive
for `-`, again calls `-`; the external interpreter makes sure that
inference results for these end up in different caches).
  • Loading branch information
Keno authored Jun 26, 2020
1 parent 063525f commit ab520c7
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 5 deletions.
7 changes: 6 additions & 1 deletion base/compiler/inferencestate.jl
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,10 @@ mutable struct InferenceState
# `max_valid`, to be used in inlining
matching_methods_cache::IdDict{Any, Tuple{Any, UInt, UInt}}

# The interpreter that created this inference state. Not looked at by
# NativeInterpreter. But other interpreters may use this to detect cycles
interp::AbstractInterpreter

# src is assumed to be a newly-allocated CodeInfo, that can be modified in-place to contain intermediate results
function InferenceState(result::InferenceResult, src::CodeInfo,
cached::Bool, interp::AbstractInterpreter)
Expand Down Expand Up @@ -107,7 +111,8 @@ mutable struct InferenceState
Vector{InferenceState}(), # callers_in_cycle
#=parent=#nothing,
cached, false, false, false,
IdDict{Any, Tuple{Any, UInt, UInt}}())
IdDict{Any, Tuple{Any, UInt, UInt}}(),
interp)
result.result = frame
cached && push!(get_inference_cache(interp), result)
return frame
Expand Down
12 changes: 8 additions & 4 deletions base/compiler/typeinfer.jl
Original file line number Diff line number Diff line change
Expand Up @@ -439,21 +439,25 @@ function merge_call_chain!(parent::InferenceState, ancestor::InferenceState, chi
end
end

function is_same_frame(interp::AbstractInterpreter, linfo::MethodInstance, frame::InferenceState)
return linfo === frame.linfo
end

# Walk through `linfo`'s upstream call chain, starting at `parent`. If a parent
# frame matching `linfo` is encountered, then there is a cycle in the call graph
# (i.e. `linfo` is a descendant callee of itself). Upon encountering this cycle,
# we "resolve" it by merging the call chain, which entails unioning each intermediary
# frame's `callers_in_cycle` field and adding the appropriate backedges. Finally,
# we return `linfo`'s pre-existing frame. If no cycles are found, `nothing` is
# returned instead.
function resolve_call_cycle!(linfo::MethodInstance, parent::InferenceState)
function resolve_call_cycle!(interp::AbstractInterpreter, linfo::MethodInstance, parent::InferenceState)
frame = parent
uncached = false
limited = false
while isa(frame, InferenceState)
uncached |= !frame.cached # ensure we never add an uncached frame to a cycle
limited |= frame.limited
if frame.linfo === linfo
if is_same_frame(interp, linfo, 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)
Expand All @@ -465,7 +469,7 @@ function resolve_call_cycle!(linfo::MethodInstance, parent::InferenceState)
return frame
end
for caller in frame.callers_in_cycle
if caller.linfo === linfo
if is_same_frame(interp, linfo, caller)
if uncached
poison_callstack(parent, frame, false)
return true
Expand Down Expand Up @@ -496,7 +500,7 @@ function typeinf_edge(interp::AbstractInterpreter, method::Method, @nospecialize
# (if we asked resolve_call_cyle, it might instead detect that there is a cycle that it can't merge)
frame = false
else
frame = resolve_call_cycle!(mi, caller)
frame = resolve_call_cycle!(interp, mi, caller)
end
if frame === false
# completely new
Expand Down

1 comment on commit ab520c7

@nanosoldier
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Executing the daily package evaluation, I will reply here when finished:

@nanosoldier runtests(ALL, isdaily = true)

Please sign in to comment.