JET integration

[timholy]

This allows one to generate reports on both the callee and caller of an inference trigger.

[timholy]

Demo:

julia> f(c) = sum(c[1])
f (generic function with 1 method)

julia> c = Any[Any[1,2,3]]
1-element Vector{Any}:
 Any[1, 2, 3]

julia> using SnoopCompile
[ Info: Precompiling SnoopCompile [aa65fe97-06da-5843-b5b1-d5d13cad87d2]

julia> tinf = @snoopi_deep f(c)
InferenceTimingNode: 0.032447/0.037888 on Core.Compiler.Timings.ROOT() with 2 direct children

julia> using JET

julia> @report_call f(c)
No errors !
(Any, 0)

julia> itrigs = inference_triggers(tinf)
2-element Vector{InferenceTrigger}:
 Inference triggered to call f(::Vector{Any}) from eval (./boot.jl:373) inlined into REPL.eval_user_input(::Any, ::REPL.REPLBackend) (/home/tim/src/julia-master/usr/share/julia/stdlib/v1.8/REPL/src/REPL.jl:150)
 Inference triggered to call sum(::Vector{Any}) from f (./REPL[1]:1) with specialization f(::Vector{Any})

julia> report_callee(itrigs[1])
No errors !
(Any, 0)

julia> report_callee(itrigs[2])
═════ 1 possible error found ═════
┌ @ reducedim.jl:889 Base.#sum#738(Base.:, Base.pairs(Core.NamedTuple()), #self#, a)
│┌ @ reducedim.jl:889 Base._sum(a, dims)
││┌ @ reducedim.jl:893 Base.#_sum#740(Base.pairs(Core.NamedTuple()), #self#, a, _3)
│││┌ @ reducedim.jl:893 Base._sum(Base.identity, a, Base.:)
││││┌ @ reducedim.jl:894 Base.#_sum#741(Base.pairs(Core.NamedTuple()), #self#, f, a, _4)
│││││┌ @ reducedim.jl:894 Base.mapreduce(f, Base.add_sum, a)
││││││┌ @ reducedim.jl:322 Base.#mapreduce#731(Base.:, Base._InitialValue(), #self#, f, op, A)
│││││││┌ @ reducedim.jl:322 Base._mapreduce_dim(f, op, init, A, dims)
││││││││┌ @ reducedim.jl:330 Base._mapreduce(f, op, Base.IndexStyle(A), A)
│││││││││┌ @ reduce.jl:402 Base.mapreduce_empty_iter(f, op, A, Base.IteratorEltype(A))
││││││││││┌ @ reduce.jl:353 Base.reduce_empty_iter(Base.MappingRF(f, op), itr, ItrEltype)
│││││││││││┌ @ reduce.jl:357 Base.reduce_empty(op, Base.eltype(itr))
││││││││││││┌ @ reduce.jl:331 Base.mapreduce_empty(Base.getproperty(op, :f), Base.getproperty(op, :rf), _)
│││││││││││││┌ @ reduce.jl:345 Base.reduce_empty(op, T)
││││││││││││││┌ @ reduce.jl:322 Base.reduce_empty(Base.+, _)
│││││││││││││││┌ @ reduce.jl:313 Base.zero(_)
││││││││││││││││┌ @ missing.jl:106 Base.throw(Base.MethodError(Base.zero, Core.tuple(Base.Any)))
│││││││││││││││││ MethodError: no method matching zero(::Type{Any})
││││││││││││││││└──────────────────
(Any, 1)

julia> report_caller(itrigs[2])
No errors !
(Any, 0)

CC @aviatesk

I might leave this up long enough for folks to experiment with it before merging. An alternative approach is to split out @snoopi_deep and a small portion of its infrastructure into a new package, ProfileInference.jl, which could then be imported by both SnoopCompile & JET. @aviatesk, I'm perfectly happy with either approach, please let me know what you prefer.

[codecov]

Codecov Report

Merging #253 (0ef71df) into master (2ff6c69) will increase coverage by 0.12%.
The diff coverage is 90.90%.

❗ Current head 0ef71df differs from pull request most recent head 9469a1c. Consider uploading reports for the commit 9469a1c to get more accurate results

@@            Coverage Diff             @@
##           master     #253      +/-   ##
==========================================
+ Coverage   86.82%   86.95%   +0.12%     
==========================================
  Files          16       16              
  Lines        2012     2023      +11     
==========================================
+ Hits         1747     1759      +12     
+ Misses        265      264       -1     

Impacted Files	Coverage Δ
SnoopCompileCore/src/snoopi_deep.jl	92.00% <ø> (ø)
src/SnoopCompile.jl	100.00% <ø> (ø)
src/parcel_snoopi_deep.jl	89.77% <90.90%> (+0.22%)	⬆️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 2ff6c69...9469a1c. Read the comment docs.

[aviatesk]

Great ! I think this would be a simple but very nice first step toward analyzing Julia programs using both static/dynamic information.

An alternative approach is to split out @snoopi_deep and a small portion of its infrastructure into a new package, ProfileInference.jl, which could then be imported by both SnoopCompile & JET

Hmm, I'm still not too sure how SnoopCompile and JET want to use "ProfileInference". My naive understanding is, "ProfileInference" (or SnoopCompile if we go with the current approach) will provide user-facing utilities on top of functionalities of SnoopCompile and JET, and so there is no need for those two packages to load "ProfileInference" ?

[timholy]

One option would be possible to extract:

• https://github.com/timholy/SnoopCompile.jl/blob/master/SnoopCompileCore/src/snoopi_deep.jl

• SnoopCompile.jl/src/parcel_snoopi_deep.jl

  Lines 17 to 44 in 1abf0a2

  	Core.MethodInstance(mi_info::InferenceFrameInfo) = mi_info.mi
  	Core.MethodInstance(t::InferenceTiming) = MethodInstance(t.mi_info)
  	Core.MethodInstance(t::InferenceTimingNode) = MethodInstance(t.mi_timing)
  	Core.Method(x::InferenceNode) = MethodInstance(x).def::Method # deliberately throw an error if this is a module
  	isROOT(mi::MethodInstance) = mi === Core.Compiler.Timings.ROOTmi
  	isROOT(m::Method) = m === Core.Compiler.Timings.ROOTmi.def
  	isROOT(mi_info::InferenceNode) = isROOT(MethodInstance(mi_info))
  	isROOT(node::InferenceTimingNode) = isROOT(node.mi_timing)
  	# Record instruction pointers we've already looked up (performance optimization)
  	const lookups = Dict{Union{UInt, Core.Compiler.InterpreterIP}, Vector{StackTraces.StackFrame}}()
  	lookups_key(ip) = ip
  	lookups_key(ip::Ptr{Nothing}) = UInt(ip)
  	# These should be in SnoopCompileCore, except that it promises not to specialize Base methods
  	Base.show(io::IO, t::InferenceTiming) = (print(io, "InferenceTiming: "); _show(io, t))
  	function _show(io::IO, t::InferenceTiming)
  	print(io, @sprintf("%8.6f", exclusive(t)), "/", @sprintf("%8.6f", inclusive(t)), " on ")
  	print(io, stripifi(t.mi_info))
  	end
  	function Base.show(io::IO, node::InferenceTimingNode)
  	print(io, "InferenceTimingNode: ")
  	_show(io, node.mi_timing)
  	print(io, " with ", string(length(node.children)), " direct children")
  	end

and put it in a standalone package. It might need a few more things, but if you just want to capture the raw data and then write your own handling, that's the minimum. But it seems likely that you'd soon want

SnoopCompile.jl/src/parcel_snoopi_deep.jl

Lines 590 to 877 in 1abf0a2

	## Analysis of inference triggers
	"""
	InferenceTrigger(callee::MethodInstance, callerframes::Vector{StackFrame}, btidx::Int, bt)
	Organize information about the "triggers" of inference. `callee` is the `MethodInstance` requiring inference,
	`callerframes`, `btidx` and `bt` contain information about the caller.
	`callerframes` are the frame(s) of call site that triggered inference; it's a `Vector{StackFrame}`, rather than a
	single `StackFrame`, due to the possibility that the caller was inlined into something else, in which case the first entry
	is the direct caller and the last entry corresponds to the MethodInstance into which it was ultimately inlined.
	`btidx` is the index in `bt`, the backtrace collected upon entry into inference, corresponding to `callerframes`.
	`InferenceTrigger`s are created by calling [`inference_triggers`](@ref).
	See also: [`callerinstance`](@ref) and [`callingframe`](@ref).
	"""
	struct InferenceTrigger
	node::InferenceTimingNode
	callerframes::Vector{StackTraces.StackFrame}
	btidx::Int # callerframes = StackTraces.lookup(bt[btidx])
	end
	function Base.show(io::IO, itrig::InferenceTrigger)
	print(io, "Inference triggered to call ")
	printstyled(io, stripmi(MethodInstance(itrig.node)); color=:yellow)
	if !isempty(itrig.callerframes)
	sf = first(itrig.callerframes)
	print(io, " from ")
	printstyled(io, sf.func; color=:red, bold=true)
	print(io, " (", sf.file, ':', sf.line, ')')
	caller = itrig.callerframes[end].linfo
	if isa(caller, MethodInstance)
	length(itrig.callerframes) == 1 ? print(io, " with specialization ") : print(io, " inlined into ")
	printstyled(io, stripmi(caller); color=:blue)
	if length(itrig.callerframes) > 1
	sf = itrig.callerframes[end]
	print(io, " (", sf.file, ':', sf.line, ')')
	end
	elseif isa(caller, Core.CodeInfo)
	print(io, " called from toplevel code ", caller)
	end
	else
	print(io, " called from toplevel")
	end
	end
	"""
	mi = callerinstance(itrig::InferenceTrigger)
	Return the MethodInstance `mi` of the caller in the selected stackframe in `itrig`.
	"""
	callerinstance(itrig::InferenceTrigger) = itrig.callerframes[end].linfo
	function callerinstances(itrigs::AbstractVector{InferenceTrigger})
	callers = Set{MethodInstance}()
	for itrig in itrigs
	!isempty(itrig.callerframes) && push!(callers, callerinstance(itrig))
	end
	return callers
	end
	function callermodule(itrig::InferenceTrigger)
	if !isempty(itrig.callerframes)
	m = callerinstance(itrig).def
	return isa(m, Module) ? m : m.module
	end
	return nothing
	end
	# Select the next (caller) frame that's a Julia (as opposed to C) frame; returns the stackframe and its index in bt, or nothing
	function next_julia_frame(bt, idx, Δ=1; methodinstanceonly::Bool=true, methodonly::Bool=true)
	while 1 <= idx+Δ <= length(bt)
	ip = lookups_key(bt[idx+=Δ])
	sfs = get!(()->Base.StackTraces.lookup(ip), lookups, ip)
	sf = sfs[end]
	sf.from_c && continue
	mi = sf.linfo
	methodinstanceonly && (isa(mi, Core.MethodInstance) || continue)
	if isa(mi, MethodInstance)
	m = mi.def
	methodonly && (isa(m, Method) || continue)
	# Exclude frames that are in Core.Compiler
	isa(m, Method) && m.module === Core.Compiler && continue
	end
	return sfs, idx
	end
	return nothing
	end
	SnoopCompileCore.exclusive(itrig::InferenceTrigger) = exclusive(itrig.node)
	SnoopCompileCore.inclusive(itrig::InferenceTrigger) = inclusive(itrig.node)
	StackTraces.stacktrace(itrig::InferenceTrigger) = stacktrace(itrig.node.bt)
	isprecompilable(itrig::InferenceTrigger) = isprecompilable(MethodInstance(itrig.node))
	"""
	itrigs = inference_triggers(tinf::InferenceTimingNode; exclude_toplevel=true)
	Collect the "triggers" of inference, each a fresh entry into inference via a call dispatched at runtime.
	All the entries in `itrigs` are previously uninferred, or are freshly-inferred for specific constant inputs.
	`exclude_toplevel` determines whether calls made from the REPL, `include`, or test suites are excluded.
	# Example
	We'll use [`SnoopCompile.itrigs_demo`](@ref), which runs `@snoopi_deep` on a workload designed to yield reproducible results:
	```jldoctest triggers; setup=:(using SnoopCompile), filter=r"([0-9]*\\.?[0-9]+([eE][-+]?[0-9]+)?|.*/deep_demos\\.jl:\\d+|WARNING: replacing module ItrigDemo\\.\\n)"
	julia> tinf = SnoopCompile.itrigs_demo()
	InferenceTimingNode: 0.004490576/0.004711168 on InferenceFrameInfo for Core.Compiler.Timings.ROOT() with 2 direct children
	julia> itrigs = inference_triggers(tinf)
	2-element Vector{InferenceTrigger}:
	Inference triggered to call MethodInstance for double(::UInt8) from calldouble1 (/pathto/SnoopCompile/src/deep_demos.jl:86) inlined into MethodInstance for calldouble2(::Vector{Vector{Any}}) (/pathto/SnoopCompile/src/deep_demos.jl:87)
	Inference triggered to call MethodInstance for double(::Float64) from calldouble1 (/pathto/SnoopCompile/src/deep_demos.jl:86) inlined into MethodInstance for calldouble2(::Vector{Vector{Any}}) (/pathto/SnoopCompile/src/deep_demos.jl:87)
	```
	```
	julia> edit(itrigs[1]) # opens an editor at the spot in the caller
	julia> ascend(itrigs[2]) # use Cthulhu to inspect the stacktrace (caller is the second item in the trace)
	Choose a call for analysis (q to quit):
	> double(::Float64)
	calldouble1 at /pathto/SnoopCompile/src/deep_demos.jl:86 => calldouble2(::Vector{Vector{Any}}) at /pathto/SnoopCompile/src/deep_demos.jl:87
	calleach(::Vector{Vector{Vector{Any}}}) at /pathto/SnoopCompile/src/deep_demos.jl:88
	...
	```
	"""
	function inference_triggers(tinf::InferenceTimingNode; exclude_toplevel::Bool=true)
	function first_julia_frame(bt)
	ret = next_julia_frame(bt, 1)
	if ret === nothing
	return StackTraces.StackFrame[], 0
	end
	return ret
	end
	itrigs = map(tinf.children) do child
	bt = child.bt
	bt === nothing && throw(ArgumentError("it seems you've supplied a child node, but backtraces are collected only at the entrance to inference"))
	InferenceTrigger(child, first_julia_frame(bt)...)
	end
	if exclude_toplevel
	filter!(maybe_internal, itrigs)
	end
	return itrigs
	end
	function maybe_internal(itrig::InferenceTrigger)
	for sf in itrig.callerframes
	linfo = sf.linfo
	if isa(linfo, MethodInstance)
	m = linfo.def
	if isa(m, Method)
	if m.module === Base
	m.name === :include_string && return false
	m.name === :_include_from_serialized && return false
	m.name === :return_types && return false # from `@inferred`
	end
	m.name === :eval && return false
	end
	end
	match(rextest, string(sf.file)) !== nothing && return false
	end
	return true
	end
	"""
	itrigcaller = callingframe(itrig::InferenceTrigger)
	"Step out" one layer of the stacktrace, referencing the caller of the current frame of `itrig`.
	You can retrieve the proximal trigger of inference with `InferenceTrigger(itrigcaller)`.
	# Example
	We collect data using the [`SnoopCompile.itrigs_demo`](@ref):
	```julia
	julia> itrig = inference_triggers(SnoopCompile.itrigs_demo())[1]
	Inference triggered to call MethodInstance for double(::UInt8) from calldouble1 (/pathto/SnoopCompile/src/parcel_snoopi_deep.jl:762) inlined into MethodInstance for calldouble2(::Vector{Vector{Any}}) (/pathto/SnoopCompile/src/parcel_snoopi_deep.jl:763)
	julia> itrigcaller = callingframe(itrig)
	Inference triggered to call MethodInstance for double(::UInt8) from calleach (/pathto/SnoopCompile/src/parcel_snoopi_deep.jl:764) with specialization MethodInstance for calleach(::Vector{Vector{Vector{Any}}})
	```
	"""
	function callingframe(itrig::InferenceTrigger)
	idx = itrig.btidx
	if idx < length(itrig.node.bt)
	ret = next_julia_frame(itrig.node.bt, idx)
	if ret !== nothing
	return InferenceTrigger(itrig.node, ret...)
	end
	end
	return InferenceTrigger(itrig.node, StackTraces.StackFrame[], length(itrig.node.bt)+1)
	end
	"""
	itrig0 = InferenceTrigger(itrig::InferenceTrigger)
	Reset an inference trigger to point to the stackframe that triggered inference.
	This can be useful to undo the actions of [`callingframe`](@ref) and [`skiphigherorder`](@ref).
	"""
	InferenceTrigger(itrig::InferenceTrigger) = InferenceTrigger(itrig.node, next_julia_frame(itrig.node.bt, 1)...)
	"""
	itrignew = skiphigherorder(itrig; exact::Bool=false)
	Attempt to skip over frames of higher-order functions that take the callee as a function-argument.
	This can be useful if you're analyzing inference triggers for an entire package and would prefer to assign
	triggers to package-code rather than Base functions like `map!`, `broadcast`, etc.
	# Example
	We collect data using the [`SnoopCompile.itrigs_higherorder_demo`](@ref):
	```julia
	julia> itrig = inference_triggers(SnoopCompile.itrigs_higherorder_demo())[1]
	Inference triggered to call MethodInstance for double(::Float64) from mymap! (/pathto/SnoopCompile/src/parcel_snoopi_deep.jl:706) with specialization MethodInstance for mymap!(::typeof(SnoopCompile.ItrigHigherOrderDemo.double), ::Vector{Any}, ::Vector{Any})
	julia> callingframe(itrig) # step out one (non-inlined) frame
	Inference triggered to call MethodInstance for double(::Float64) from mymap (/pathto/SnoopCompile/src/parcel_snoopi_deep.jl:710) with specialization MethodInstance for mymap(::typeof(SnoopCompile.ItrigHigherOrderDemo.double), ::Vector{Any})
	julia> skiphigherorder(itrig) # step out to frame that doesn't have `double` as a function-argument
	Inference triggered to call MethodInstance for double(::Float64) from callmymap (/pathto/SnoopCompile/src/parcel_snoopi_deep.jl:711) with specialization MethodInstance for callmymap(::Vector{Any})
	```
	!!! warn
	By default `skiphigherorder` is conservative, and insists on being sure that it's the callee being passed to the higher-order function.
	Higher-order functions that do not get specialized (e.g., with `::Function` argument types) will not be skipped over.
	You can pass `exact=false` to allow `::Function` to also be passed over, but keep in mind that this may falsely skip some frames.
	"""
	function skiphigherorder(itrig::InferenceTrigger; exact::Bool=true)
	ft = Base.unwrap_unionall(Base.unwrap_unionall(MethodInstance(itrig.node).specTypes).parameters[1])
	sfs, idx = itrig.callerframes, itrig.btidx
	while idx < length(itrig.node.bt)
	if !isempty(sfs)
	callermi = sfs[end].linfo
	if !hasparameter(callermi.specTypes, ft, exact)
	return InferenceTrigger(itrig.node, sfs, idx)
	end
	end
	ret = next_julia_frame(itrig.node.bt, idx)
	ret === nothing && return InferenceTrigger(itrig.node, sfs, idx)
	sfs, idx = ret
	end
	return itrig
	end
	function hasparameter(@nospecialize(typ), @nospecialize(ft), exact::Bool)
	isa(typ, Type) || return false
	typ = Base.unwrap_unionall(typ)
	typ === ft && return true
	exact || (typ === Function && return true)
	typ === Union{} && return false
	if isa(typ, Union)
	hasparameter(typ.a, ft, exact) && return true
	hasparameter(typ.b, ft, exact) && return true
	return false
	end
	for p in typ.parameters
	hasparameter(p, ft, exact) && return true
	end
	return false
	end
	"""
	ncallees, ncallers = diversity(itrigs::AbstractVector{InferenceTrigger})
	Count the number of distinct MethodInstances among the callees and callers, respectively, among the triggers in `itrigs`.
	"""
	function diversity(itrigs)
	# Analyze caller => callee argument type diversity
	callees, callers, ncextra = Set{MethodInstance}(), Set{MethodInstance}(), 0
	for itrig in itrigs
	push!(callees, MethodInstance(itrig.node))
	caller = itrig.callerframes[end].linfo
	if isa(caller, MethodInstance)
	push!(callers, caller)
	else
	ncextra += 1
	end
	end
	return length(callees), length(callers) + ncextra
	end

, and at that point it might make more sense to just have SnoopCompile be a wrapper. It's already a snoop/Cthulhu/AbstractTrees integration package, so this would not be weird.

Anyway, the point is, let me know what you think makes the most sense.

[aviatesk]

Thanks so much for your explanation.

Okay, so to me it seems more natural to have this integration in SnoopCompile. From JET's viewpoint, it's much simpler if it just focuses on its (half-)static analysis feature, while it just makes sense to have external package (SnoopCompile) that combines JET's static analysis and runtime information.
If you think you want to make SnoopCompile focus on "inference timing" tools, then it also makes sense to have a separate package that is dedicated to execute a program with collecting entry points for static analysis, but as far as I understand, SnoopCompile is already such a package that combines both static and runtime information, right ?

Just a random thought, in the long run, we may want another, more automated way to combine static analysis and runtime execution, e.g. we can use JuliaInterpreter/Cassette/compiler-plugin to automatically invoke "static" analysis provided by SnoopCompile/JET/whatsoever in the middle of program execution and collect some information. But also, we know (basically-)interactive inspections like SnoopCompile/Cthulhu are the most accurate and efficient tools we have at this moment.