Expand SnoopCompileCore.@snoopi to report _all MethodInstances_ that are inferred during a single inference call.

[NHDaly]

Before this PR, it reported only the entry MethodInstance:

julia> module M
           h(a::Array) = a[1]::Integer + 1
           g(y::Integer, x) = h(Any[y]) + Int(x)
       end;

julia> SnoopCompileCore.@snoopi begin
           M.g(2,3.0)
       end
1-element Array{Tuple{Float64,Core.MethodInstance},1}:
 (0.0011429786682128906, MethodInstance for g(::Int64, ::Float64))

After this PR, we want to be able to get all the MethodInstances that are (recursively) inferred during the inference to the top method instance.

Currently, it now looks like this:

julia> SnoopCompileCore.@snoopi begin
           M.g(2,3.0)
       end
1-element Vector{Tuple{Float64,Vector{Core.MethodInstance}}}:
 (0.001068115234375, [MethodInstance for g(::Int64, ::Float64), MethodInstance for h(::Vector{Any})])

But we're very open to other ideas for how best to represent this! :)
I think it might be super cool to return this as a DAG rather than a vector, so that we could maybe then work our way up from the leaves timing code_typed on each of them one at a time to try to understand which function had the most exclusive time?

---

This PR is basically applying the same trick we did in #37136 to get the full set of MethodInstances being compiled, rather than just the entry method instance.

[NHDaly]

I'm super not sure where the best places actually are to do the logging from within Julia. I just poked around putting print statements in a bunch of places until I found a spot that seems(?) reasonable.

I'd super appreciate help / advice / suggestions on how best to do this! ❤️! Thanks in advance ~~ hopefully this is useful to others! :)

[timholy]

I haven't had a chance to think about this in detail, but you're not the first person to wish for this. 👍

[Keno]

This might be better done with a custom AbstractInterpreter, so you can also control caching. Otherwise the results of this will be wildly unstable, depending on what happens to be in the global cache.

[NHDaly]

This might be better done with a custom AbstractInterpreter, so you can also control caching. Otherwise the results of this will be wildly unstable, depending on what happens to be in the global cache.

Thanks @Keno. If I'm not mistaken, isn't what you're describing already true of the behavior of @snoopi? I think @snoopi is already going to return different answers depending on what's in the global cache, and I think this PR will behave the same way. Is that what you're referring to?
e.g. this?:

julia> f(x) = x+3
f (generic function with 1 method)

julia> SnoopCompileCore.@snoopi begin
           f(2)
       end
1-element Array{Tuple{Float64,Core.MethodInstance},1}:
 (0.0002541542053222656, MethodInstance for f(::Int64))

julia> SnoopCompileCore.@snoopi begin
           f(2)
       end
0-element Array{Tuple{Float64,Core.MethodInstance},1}

julia> 

If so, I think this is actually what we want, because we want to know for a given run exactly what ended up taking time in the compiler, not necessarily what functions are most expensive to compile from scratch (though that would also be useful to know).

Is the above what you were referring to? Or were you referring to something else? Thanks!

[NHDaly]

Question:
Is it expected that julia would be calling infer() on the same target function (even with the same caller) multiple times, inside typeinf_edge()?

Unless I've messed up my data collection, which is of course possible, in the example I just recorded it seems like there are hundreds of exact duplicate edges. I've added the logging inside typeinf_edge, on line 541, right before the call to infer():

julia/base/compiler/typeinfer.jl

Lines 525 to 545 in be69119

	if frame === false
	# completely new
	lock_mi_inference(interp, mi)
	result = InferenceResult(mi)
	frame = InferenceState(result, #=cached=#true, interp) # always use the cache for edge targets
	if frame === nothing
	# can't get the source for this, so we know nothing
	unlock_mi_inference(interp, mi)
	return Any, nothing
	end
	if caller.cached || caller.limited # don't involve uncached functions in cycle resolution
	frame.parent = caller
	end
	if __collect_inference_callees__[]
	#Core.println("edge!")
	#Core.println(__inference_callee_edges4__)
	push!(__inference_callee_edges4__[end][2], (caller.linfo, frame.linfo))
	end
	typeinf(interp, frame)
	update_valid_age!(frame, caller)
	return widenconst_bestguess(frame.bestguess), frame.inferred ? mi : nothing

That's inside the body of typeinf_edge, whose signature is here:

julia/base/compiler/typeinfer.jl

Line 503 in be69119

function typeinf_edge(interp::AbstractInterpreter, method::Method, @nospecialize(atypes), sparams::SimpleVector, caller::InferenceState)

And when I collected the data, I found that within a single top-level call to typeinf_ext_toplevel(), we hit duplicate edges many times. For example, this edge shows up 144 times:

julia> [e for e in edges if last(e) == sorted_targets[1]]
144-element Vector{Pair{Core.MethodInstance,Core.MethodInstance}}:
 MethodInstance for (::Base.var"#printstyled##kw")(::NamedTuple{(:color,),Tuple{Symbol}}, ::typeof(printstyled), ::String) => MethodInstance for #printstyled#779(::Bool, ::Symbol, ::typeof(printstyled), ::String)
 MethodInstance for (::Base.var"#printstyled##kw")(::NamedTuple{(:color,),Tuple{Symbol}}, ::typeof(printstyled), ::String) => MethodInstance for #printstyled#779(::Bool, ::Symbol, ::typeof(printstyled), ::String)
 MethodInstance for (::Base.var"#printstyled##kw")(::NamedTuple{(:color,),Tuple{Symbol}}, ::typeof(printstyled), ::String) => MethodInstance for #printstyled#779(::Bool, ::Symbol, ::typeof(printstyled), ::String)
 MethodInstance for (::Base.var"#printstyled##kw")(::NamedTuple{(:color,),Tuple{Symbol}}, ::typeof(printstyled), ::String) => MethodInstance for #printstyled#779(::Bool, ::Symbol, ::typeof(printstyled), ::String)
 MethodInstance for (::Base.var"#printstyled##kw")(::NamedTuple{(:color,),Tuple{Symbol}}, ::typeof(printstyled), ::String) => MethodInstance for #printstyled#779(::Bool, ::Symbol, ::typeof(printstyled), ::String)
 MethodInstance for (::Base.var"#printstyled##kw")(::NamedTuple{(:color,),Tuple{Symbol}}, ::typeof(printstyled), ::String) => MethodInstance for #printstyled#779(::Bool, ::Symbol, ::typeof(printstyled), ::String)
 MethodInstance for (::Base.var"#printstyled##kw")(::NamedTuple{(:color,),Tuple{Symbol}}, ::typeof(printstyled), ::String) => MethodInstance for #printstyled#779(::Bool, ::Symbol, ::typeof(printstyled), ::String)
 MethodInstance for (::Base.var"#printstyled##kw")(::NamedTuple{(:color,),Tuple{Symbol}}, ::typeof(printstyled), ::String) => MethodInstance for #printstyled#779(::Bool, ::Symbol, ::typeof(printstyled), ::String)
 MethodInstance for (::Base.var"#printstyled##kw")(::NamedTuple{(:color,),Tuple{Symbol}}, ::typeof(printstyled), ::String) => MethodInstance for #printstyled#779(::Bool, ::Symbol, ::typeof(printstyled), ::String)
 MethodInstance for (::Base.var"#printstyled##kw")(::NamedTuple{(:color,),Tuple{Symbol}}, ::typeof(printstyled), ::String) => MethodInstance for #printstyled#779(::Bool, ::Symbol, ::typeof(printstyled), ::String)
 MethodInstance for (::Base.var"#printstyled##kw")(::NamedTuple{(:color,),Tuple{Symbol}}, ::typeof(printstyled), ::String) => MethodInstance for #printstyled#779(::Bool, ::Symbol, ::typeof(printstyled), ::String)
...

julia> unique([e for e in edges if last(e) == sorted_targets[1]])
1-element Vector{Pair{Core.MethodInstance,Core.MethodInstance}}:
 MethodInstance for (::Base.var"#printstyled##kw")(::NamedTuple{(:color,),Tuple{Symbol}}, ::typeof(printstyled), ::String) => MethodInstance for #printstyled#779(::Bool, ::Symbol, ::typeof(printstyled), ::String)

Out of the 6,212 edges I collected for one top-level function, there are only 145 unique edges, and only 118 unique MethodInstances.

This also means that there are some MethodInstances being inferred from multiple parents. (There are 21 such cases, mostly with 2 callers, four with 3, and one with 5.)

---

Is this expected? Based on reading the comments and the code, I thought that this should only be called if mi had not yet been inferred and wasn't part of a cycle, which to me means that there shouldn't be duplicates?

Am I misunderstanding this?

Thanks! :) I figured I'd mention it here in case this is a bug / performance regression that could easily be fixed! :)

[NHDaly]

And, actually, more evidence for the above:

For the method instance in question, @snoopi reported 2.6 seconds to infer in a cold process, but when I instead summed the time for code_typed() on each inner MethodInstance encountered above (also on a cold process), starting from the leaves working all the way back up in a sorted DAG, the total time was only 1.01 seconds.

So I'm thinking that there's maybe a problem in the inference code, where it's not detecting duplicate nodes correctly, and it's redundantly inferring the same MethodInstances multiple times within the same inferable set!

@JeffBezanson: Maybe this is also something we can take a look at together tomorrow when we meet up? :)

EDIT: Although, hmm, maybe not: is code_typed not doing everything inference does? When I just ran code_typed() in a cold process, on the top-level MethodInstance above, the total time was also only 0.16 seconds. So I'm now thinking code_typed is just not the right thing to be timing to time what inference does?

---

EDIT EDIT: On a bigger example, where the original @snoopi reported 73 secs, summing individual code_typeds takes only 2.3 seconds. I recorded 29,135 (caller, frame) edges, which deduplicates down to 284 unique edges over 228 unique MethodInstances.