Expose changes to Julia via @snoopi_deep to report timing profile breakdown of _all MethodInstances_ that are inferred during type inference

.github/workflows/ci.yml

@@ -11,7 +11,7 @@ jobs:
     name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
     runs-on: ${{ matrix.os }}
     strategy:
-      fail-fast: true
+      fail-fast: false
       matrix:
         version:
           - '1.0'

Project.toml

@@ -5,18 +5,21 @@ version = "2.1.2"
 
 [deps]
 Cthulhu = "f68482b8-f384-11e8-15f7-abe071a5a75f"
+FlameGraphs = "08572546-2f56-4bcf-ba4e-bab62c3a3f89"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 SnoopCompileCore = "e2b509da-e806-4183-be48-004708413034"
 
 [compat]
 Cthulhu = "1.2"
+FlameGraphs = "0.2"
 OrderedCollections = "1"
 SnoopCompileCore = "~2.1.2"
 julia = "1"
 
 [extras]
+AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
 ColorTypes = "3da002f7-5984-5a60-b8a6-cbb66c0b333f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 FixedPointNumbers = "53c48c17-4a7d-5ca2-90c5-79b7896eea93"
@@ -28,4 +31,4 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["ColorTypes", "Documenter", "FixedPointNumbers", "InteractiveUtils", "JLD", "MethodAnalysis", "Pkg", "SparseArrays", "Test"]
+test = ["AbstractTrees", "ColorTypes", "Documenter", "FixedPointNumbers", "InteractiveUtils", "JLD", "MethodAnalysis", "Pkg", "SparseArrays", "Test"]

SnoopCompileCore/src/SnoopCompileCore.jl

@@ -11,6 +11,9 @@ if VERSION >= v"1.2.0-DEV.573"
     include("snoopi.jl")
 end
 
+if VERSION >= v"1.6.0-DEV.1190"  # https://github.com/JuliaLang/julia/pull/37749
+    include("snoopi_deep.jl")
+end
 
 if VERSION >= v"1.6.0-DEV.154"
     include("snoopr.jl")

SnoopCompileCore/src/snoopi.jl

@@ -1,4 +1,4 @@
-export @snoopi
+export @snoopi, @snoopi_deep
 
 const __inf_timing__ = Tuple{Float64,MethodInstance}[]
 
@@ -51,7 +51,7 @@ end
     inf_timing = @snoopi tmin=0.0 commands
 
 Execute `commands` while snooping on inference. Returns an array of `(t, linfo)`
-tuples, where `t` is the amount of time spent infering `linfo` (a `MethodInstance`).
+tuples, where `t` is the amount of time spent inferring `linfo` (a `MethodInstance`).
 
 Methods that take less time than `tmin` will not be reported.
 """
@@ -99,5 +99,6 @@ function __init__()
     end
     precompile(start_timing, ())
     precompile(stop_timing, ())
+
     nothing
 end

SnoopCompileCore/src/snoopi_deep.jl

@@ -0,0 +1,72 @@
+function start_deep_timing()
+    Core.Compiler.Timings.reset_timings()
+    Core.Compiler.__set_measure_typeinf(true)
+end
+function stop_deep_timing()
+    Core.Compiler.__set_measure_typeinf(false)
+    Core.Compiler.Timings.close_current_timer()
+end
+
+function finish_snoopi_deep()
+    return Core.Compiler.Timings._timings[1]
+end
+
+function _snoopi_deep(cmd::Expr)
+    return quote
+        start_deep_timing()
+        try
+            $(esc(cmd))
+        finally
+            stop_deep_timing()
+        end
+        finish_snoopi_deep()
+    end
+end
+
+"""
+    timing_tree = @snoopi_deep commands
+
+Produce a profile of julia's type inference, containing the amount of time spent inferring
+every `MethodInstance` processed while executing `commands`.
+
+The top-level node in this profile tree is `ROOT`, which contains the time spent _not_ in
+julia's type inference (codegen, llvm_opt, runtime, etc).
+
+To make use of these results, see the processing functions in SnoopCompile:
+    - [`SnoopCompile.flatten_times(timing_tree)`](@ref)
+    - [`SnoopCompile.flamegraph(timing_tree)`](@ref)
+
+# Examples
+```julia
+julia> timing = @snoopi_deep begin
+           @eval sort(rand(100))  # Evaluate some code and profile julia's type inference
+       end;
+
+julia> using SnoopCompile, ProfileView
+
+julia> times = flatten_times(timing, tmin_secs=0.001)
+4-element Vector{Any}:
+ 0.001088448 => Core.Compiler.Timings.InferenceFrameInfo(MethodInstance for fpsort!(...
+ 0.001618478 => Core.Compiler.Timings.InferenceFrameInfo(MethodInstance for rand!(...
+ 0.002289655 => Core.Compiler.Timings.InferenceFrameInfo(MethodInstance for _rand_max383!(...
+ 0.093143594 => Core.Compiler.Timings.InferenceFrameInfo(MethodInstance for ROOT(), ...
+
+julia> fg = flamegraph(timing)
+Node(FlameGraphs.NodeData(ROOT() at typeinfer.jl:70, 0x00, 0:15355670))
+
+julia> ProfileView.view(fg);  # Display the FlameGraph in a package that supports it
+
+julia> fg = flamegraph(timing; tmin_secs=0.0001)  # Skip very tiny frames
+Node(FlameGraphs.NodeData(ROOT() at typeinfer.jl:70, 0x00, 0:15355670))
+```
+"""
+macro snoopi_deep(cmd)
+    return _snoopi_deep(cmd)
+end
+
+# These are okay to come at the top-level because we're only measuring inference, and
+# inference results will be cached in a `.ji` file.
+@assert precompile(Core.Compiler.Timings.reset_timings, ())
+@assert precompile(start_deep_timing, ())
+@assert precompile(stop_deep_timing, ())
+@assert precompile(finish_snoopi_deep, ())

src/SnoopCompile.jl

@@ -3,6 +3,9 @@ module SnoopCompile
 using SnoopCompileCore
 export @snoopc
 isdefined(SnoopCompileCore, Symbol("@snoopi")) && export @snoopi
+if isdefined(SnoopCompileCore, Symbol("@snoopi_deep"))
+    export @snoopi_deep, flamegraph, flatten_times
+end
 if isdefined(SnoopCompileCore, Symbol("@snoopr"))
     export @snoopr, uinvalidated, invalidation_trees, filtermod, findcaller, ascend
 end
@@ -24,6 +27,10 @@ if VERSION >= v"1.2.0-DEV.573"
     include("parcel_snoopi.jl")
 end
 
+if VERSION >= v"1.6.0-DEV.1190"  # https://github.com/JuliaLang/julia/pull/37749
+    include("parcel_snoopi_deep.jl")
+end
+
 if isdefined(SnoopCompileCore, Symbol("@snoopr"))
     include("invalidations.jl")
 end

src/parcel_snoopi_deep.jl

@@ -0,0 +1,147 @@
+import FlameGraphs
+
+using Base.StackTraces: StackFrame
+using FlameGraphs.LeftChildRightSiblingTrees: Node, addchild
+using Core.Compiler.Timings: Timing
+
+const flamegraph = FlameGraphs.flamegraph  # For re-export
+
+"""
+    flatten_times(timing::Core.Compiler.Timings.Timing; tmin_secs = 0.0)
+
+Flatten the execution graph of Timings returned from `@snoopi_deep` into a Vector of pairs,
+with the exclusive time for each invocation of type inference, skipping any frames that
+took less than `tmin_secs` seconds. Results are sorted by time.
+"""
+function flatten_times(timing::Core.Compiler.Timings.Timing; tmin_secs = 0.0)
+    out = Pair{Float64,Core.Compiler.Timings.InferenceFrameInfo}[]
+    frontier = [timing]
+    while !isempty(frontier)
+        t = popfirst!(frontier)
+        exclusive_time = (t.time / 1e9)
+        if exclusive_time >= tmin_secs
+            push!(out, exclusive_time => t.mi_info)
+        end
+        append!(frontier, t.children)
+    end
+    return sort(out; by=tl->tl[1])
+end
+
+struct InclusiveTiming
+    mi_info::Core.Compiler.Timings.InferenceFrameInfo
+    inclusive_time::UInt64
+    start_time::UInt64
+    children::Vector{InclusiveTiming}
+end
+
+inclusive_time(t::InclusiveTiming) = t.inclusive_time
+
+function build_inclusive_times(t::Timing)
+    child_times = InclusiveTiming[
+        build_inclusive_times(child)
+        for child in t.children
+    ]
+    incl_time = t.time + sum(inclusive_time, child_times; init=UInt64(0))
+    return InclusiveTiming(t.mi_info, incl_time, t.start_time, child_times)
+end
+
+"""
+    flamegraph(t::Core.Compiler.Timings.Timing; tmin_secs=0.0)
+    flamegraph(t::SnoopCompile.InclusiveTiming; tmin_secs=0.0)
+
+Convert the call tree of inference timings returned from `@snoopi_deep` into a FlameGraph.
+Returns a FlameGraphs.FlameGraph structure that represents the timing trace recorded for
+type inference.
+
+Frames that take less than `tmin_secs` seconds of _inclusive time_ will not be included
+in the resultant FlameGraph (meaning total time including it and all of its children).
+This can be helpful if you have a very big profile, to save on processing time.
+
+# Examples
+```julia
+julia> timing = @snoopi_deep begin
+           @eval sort(rand(100))  # Evaluate some code and profile julia's type inference
+       end;
+
+julia> fg = flamegraph(timing)
+Node(FlameGraphs.NodeData(ROOT() at typeinfer.jl:70, 0x00, 0:15355670))
+
+julia> ProfileView.view(fg);  # Display the FlameGraph in a package that supports it
+
+julia> fg = flamegraph(timing; tmin_secs=0.0001)  # Skip very tiny frames
+Node(FlameGraphs.NodeData(ROOT() at typeinfer.jl:70, 0x00, 0:15355670))
+```
+
+NOTE: This function must touch every frame in the provided `Timing` to build inclusive
+timing information (`InclusiveTiming`). If you have a very large profile, and you plan to
+call this function multiple times (say with different values for `tmin_secs`), you can save
+some intermediate time by first calling [`SnoopCompile.build_inclusive_times(t)`](@ref), only once,
+and then passing in the `InclusiveTiming` object for all subsequent calls.
+"""
+function FlameGraphs.flamegraph(t::Timing; tmin_secs = 0.0)
+    it = build_inclusive_times(t)
+    flamegraph(it; tmin_secs=tmin_secs)
+end
+
+function FlameGraphs.flamegraph(to::InclusiveTiming; tmin_secs = 0.0)
+    tmin_ns = UInt64(round(tmin_secs * 1e9))
+
+    # Compute a "root" frame for the top-level node, to cover the whole profile
+    node_data = _flamegraph_frame(to, to.start_time; toplevel=true)
+    root = Node(node_data)
+    return _build_flamegraph!(root, to, to.start_time, tmin_ns)
+end
+function _build_flamegraph!(root, to::InclusiveTiming, start_ns, tmin_ns)
+    for child in to.children
+        if child.inclusive_time > tmin_ns
+            node_data = _flamegraph_frame(child, start_ns; toplevel=false)
+            node = addchild(root, node_data)
+            _build_flamegraph!(node, child, start_ns, tmin_ns)
+        end
+    end
+    return root
+end
+
+function frame_name(mi_info::Core.Compiler.Timings.InferenceFrameInfo)
+    frame_name(mi_info.mi::Core.Compiler.MethodInstance)
+end
+function frame_name(mi::Core.Compiler.MethodInstance)
+    frame_name(mi.def.name, mi.specTypes)
+end
+# Special printing for Type Tuples so they're less ugly in the FlameGraph
+function frame_name(name, ::Type{TT}) where TT<:Tuple
+    io = IOBuffer()
+    Base.show_tuple_as_call(io, name, TT)
+    v = String(take!(io))
+    return v
+end
+
+# NOTE: The "root" node doesn't cover the whole profile, because it's only the _complement_
+# of the inference times (so it's missing the _overhead_ from the measurement).
+# SO we need to manually create a root node that covers the whole thing.
+function max_end_time(t::InclusiveTiming)
+    # It's possible that t is already the longest-reaching node.
+    t_end = UInt64(t.start_time + t.inclusive_time)
+    # It's also possible that the last child extends past the end of t. (I think this is
+    # possible because of the small unmeasured overhead in computing these measurements.)
+    last_node = length(t.children) > 0 ? t.children[end] : t
+    child_end = last_node.start_time + last_node.inclusive_time
+    # Return the maximum end time to make sure the top node covers the entire graph.
+    return max(t_end, child_end)
+end
+
+# Make a flat frame for this Timing
+function _flamegraph_frame(to::InclusiveTiming, start_ns; toplevel)
+    mi = to.mi_info.mi
+    tt = Symbol(frame_name(to.mi_info))
+    sf = StackFrame(tt, mi.def.file, mi.def.line, mi, false, false, UInt64(0x0))
+    status = 0x0  # "default" status -- See FlameGraphs.jl
+    start = to.start_time - start_ns
+    if toplevel
+        # Compute a range over the whole profile for the top node.
+        range = Int(start) : Int(max_end_time(to) - start_ns)
+    else
+        range = Int(start) : Int(start + to.inclusive_time)
+    end
+    return FlameGraphs.NodeData(sf, status, range)
+end

test/runtests.jl

@@ -1,11 +1,18 @@
+using Test
+
 if VERSION >= v"1.2.0-DEV.573"
     include("snoopi.jl")
 end
 
+if VERSION >= v"1.6.0-DEV.1190"  # https://github.com/JuliaLang/julia/pull/37749
+    @testset "snoopi_deep" begin
+        include("snoopi_deep.jl")
+    end
+end
+
 using SnoopCompile
 using JLD
 using SparseArrays
-using Test
 
 # issue #26
 logfile = joinpath(tempdir(), "anon.log")

test/snoopi_deep.jl

@@ -0,0 +1,58 @@
+using SnoopCompile
+using SnoopCompile.SnoopCompileCore
+using Test
+
+using AbstractTrees  # For FlameGraphs tests
+
+@testset "@snoopi_deep" begin
+    # WARMUP (to compile all the small, reachable methods)
+    @eval module M  # Example with some functions that include type instability
+        i(x) = x+5
+        h(a::Array) = i(a[1]::Integer) + 2
+        g(y::Integer) = h(Any[y])
+    end
+    M.g(2)  # Warmup all deeply reachable functions
+
+    # Redefine the module, so the snoop will only show these functions:
+    @eval module M  # Example with some functions that include type instability
+        i(x) = x+5
+        h(a::Array) = i(a[1]::Integer) + 2
+        g(y::Integer) = h(Any[y])
+    end
+
+    timing = SnoopCompileCore.@snoopi_deep begin
+        M.g(2)
+    end
+    times = SnoopCompile.flatten_times(timing)
+    @test length(times) == 5  # ROOT, g(...), h(...), i(::Integer), i(::Int)
+    names = [mi_info.mi.def.name for (time, mi_info) in times]
+    @test sort(names) == [:ROOT, :g, :h, :i, :i]
+
+    longest_frame_time = times[end][1]
+    @test length(SnoopCompile.flatten_times(timing, tmin_secs=longest_frame_time)) == 1
+end
+
+@testset "flamegraph_export" begin
+    @eval module M  # Take another timing
+        i(x) = x+5
+        h(a::Array) = i(a[1]::Integer) + 2
+        g(y::Integer) = h(Any[y])
+    end
+
+    timing = SnoopCompileCore.@snoopi_deep begin
+        M.g(2)
+    end
+    times = SnoopCompile.flatten_times(timing)
+
+    fg = SnoopCompile.flamegraph(timing)
+    @test length(collect(AbstractTrees.PreOrderDFS(fg))) == 5
+    # Test that the span covers the whole tree.
+    for leaf in AbstractTrees.PreOrderDFS(fg)
+        @test leaf.data.span.start in fg.data.span
+        @test leaf.data.span.stop in fg.data.span
+    end
+
+    cutoff_bottom_frame = (times[1][1] + times[2][1]) / 2
+    fg2 = SnoopCompile.flamegraph(timing, tmin_secs = cutoff_bottom_frame)
+    @test length(collect(AbstractTrees.PreOrderDFS(fg2))) == (length(collect(AbstractTrees.PreOrderDFS(fg))) - 1)
+end