flamegraphs

Author: timholy

src/parcel_snoop_inference.jl

@@ -1789,27 +1789,33 @@ stripifi(args...) = strip_prefix(args..., "InferenceFrameInfo for ")
 ## Flamegraph creation
 
 """
-    flamegraph(tinf::InferenceTimingNode; tmin=0.0, excluded_modules=Set([Main]), mode=nothing)
+    flamegraph(tinf::InferenceTimingNode; include_llvm=true, tmin=0.0, excluded_modules=Set([Main]), mode=nothing)
 
-Convert the call tree of inference timings returned from `@snoop_inference` into a FlameGraph.
-Returns a FlameGraphs.FlameGraph structure that represents the timing trace recorded for
-type inference.
+Convert the call tree of inference timings returned from `@snoop_inference` into
+a FlameGraph. Returns a FlameGraphs.FlameGraph structure that represents the
+timing trace recorded for type inference and, if `include_llvm` is `true`, LLVM
+compilation.
 
 Frames that take less than `tmin` 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.
+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.
 
-Non-precompilable frames are marked in reddish colors. `excluded_modules` can be used to mark methods
-defined in modules to which you cannot or do not wish to add precompiles.
+Non-precompilable frames are marked in reddish colors. `excluded_modules` can be
+used to mark methods defined in modules to which you cannot or do not wish to
+add precompiles.
 
-`mode` controls how frames are named in tools like ProfileView.
-`nothing` uses the default of just the qualified function name, whereas
-supplying `mode=Dict(method => count)` counting the number of specializations of
-each method will cause the number of specializations to be included in the frame name.
+`mode` controls how frames are named in tools like ProfileView. `nothing` uses
+the default of just the qualified function name, whereas `mode=:spec` will print
+the full specialization of each MethodInstance. Supplying a pre-populated
+`mode=Dict(method => count)` will cause `count` to be included in the frame name
+of the corresponding `method`. This can be used, for example, to display the
+number of specializations of each method.
 
 # Example
 
-We'll use [`SnoopCompile.flatten_demo`](@ref), which runs `@snoop_inference` on a workload designed to yield reproducible results:
+We'll use [`SnoopCompile.flatten_demo`](@ref), which runs `@snoop_inference` on
+a workload designed to yield reproducible results:
 
 ```jldoctest flamegraph; setup=:(using SnoopCompile), filter=r"([0-9]*\\.?[0-9]+([eE][-+]?[0-9]+)?/[0-9]*\\.?[0-9]+([eE][-+]?[0-9]+)?|at.*typeinfer\\.jl:\\d+|0:\\d+|WARNING: replacing module FlattenDemo\\.\\n)"
 julia> tinf = SnoopCompile.flatten_demo()
@@ -1823,36 +1829,44 @@ Node(FlameGraphs.NodeData(ROOT() at typeinfer.jl:75, 0x00, 0:3334431))
 julia> ProfileView.view(fg);  # Display the FlameGraph in a package that supports it
 ```
 
-You should be able to reconcile the resulting flamegraph to `print_tree(tinf)` (see [`flatten`](@ref)).
+You should be able to reconcile the resulting flamegraph to `print_tree(tinf)`
+(see [`flatten`](@ref)).
 
-The empty horizontal periods in the flamegraph correspond to times when something other than inference is running.
-The total width of the flamegraph is set from the `ROOT` node.
+The empty horizontal periods in the flamegraph correspond to times when
+something other than inference is running. The total width of the flamegraph is
+set from the `ROOT` node.
 """
-function FlameGraphs.flamegraph(tinf::InferenceTimingNode; tmin = 0.0, excluded_modules=Set([Main::Module]), mode=nothing)
+function FlameGraphs.flamegraph(tinf::InferenceTimingNode; include_llvm::Bool=true, tmin = 0.0, excluded_modules=Set([Main::Module]), mode=nothing)
     isROOT(tinf) && isempty(tinf.children) && @warn "Empty profile: no compilation was recorded."
+    duration(node) = inclusive(node) + include_llvm * reinterpret(Float16, node.ci.time_compile)
+    ctimes = pushfirst!(cumsum(duration.(tinf.children)), 0)
     io = IOBuffer()
     # Compute a "root" frame for the top-level node, to cover the whole profile
-    node_data, _ = _flamegraph_frame(io, tinf, tinf.start_time, true, excluded_modules, mode; toplevel=true)
+    node_data, _ = _flamegraph_frame(duration, io, tinf, 0, false, excluded_modules, mode; stop_secs=ctimes[end])
     root = Node(node_data)
     if !isROOT(tinf)
-        node_data, child_check_precompilable = _flamegraph_frame(io, tinf, tinf.start_time, true, excluded_modules, mode; toplevel=false)
+        node_data, child_check_precompilable = _flamegraph_frame(duration, io, tinf, 0, true, excluded_modules, mode)
         root = addchild(root, node_data)
     end
-    return _build_flamegraph!(root, io, tinf, tinf.start_time, tmin, true, excluded_modules, mode)
+    return _build_flamegraph!(duration, root, io, tinf, ctimes, tmin, true, excluded_modules, mode)
 end
-function _build_flamegraph!(root, io::IO, node::InferenceTimingNode, start_secs, tmin, check_precompilable, excluded_modules, mode)
-    for child in node.children
-        if inclusive(child) > tmin
-            node_data, child_check_precompilable = _flamegraph_frame(io, child, start_secs, check_precompilable, excluded_modules, mode; toplevel=false)
+function _build_flamegraph!(duration, root, io::IO, node::InferenceTimingNode, ctimes, tmin, check_precompilable, excluded_modules, mode)
+    start_secs = first(ctimes)
+    for (i, child) in pairs(node.children)
+        if duration(child) > tmin
+            node_data, child_check_precompilable = _flamegraph_frame(duration, io, child, ctimes[i], check_precompilable, excluded_modules, mode)
             node = addchild(root, node_data)
-            _build_flamegraph!(node, io, child, start_secs, tmin, child_check_precompilable, excluded_modules, mode)
+            if !isempty(child.children)
+                newtimes = pushfirst!(cumsum(duration.(child.children)), 0) .+ start_secs
+                _build_flamegraph!(duration, node, io, child, newtimes, tmin, child_check_precompilable, excluded_modules, mode)
+            end
         end
     end
     return root
 end
 
 # Create a profile frame for this node
-function _flamegraph_frame(io::IO, node::InferenceTimingNode, start_secs, check_precompilable::Bool, excluded_modules, mode; toplevel)
+function _flamegraph_frame(duration, io::IO, node::InferenceTimingNode, start_secs, check_precompilable::Bool, excluded_modules, mode; stop_secs=nothing)
     function func_name(mi::MethodInstance, ::Nothing)
         m = mi.def
         return isa(m, Method) ? string(m.module, '.', m.name) : string(m, '.', "thunk")
@@ -1868,20 +1882,17 @@ function _flamegraph_frame(io::IO, node::InferenceTimingNode, start_secs, check_
         end
         return str
     end
-    # function func_name(io::IO, mi_info::InferenceFrameInfo, mode)
-    #     if mode === :slots
-    #         show(io, mi_info)
-    #         str = String(take!(io))
-    #         startswith(str, "InferenceFrameInfo for ") && (str = str[length("InferenceFrameInfo for ")+1:end])
-    #         return str
-    #     elseif mode === :spec
-    #         return frame_name(io, mi_info)
-    #     else
-    #         return func_name(MethodInstance(mi_info), mode)
-    #     end
-    # end
+    function func_name(io::IO, node::InferenceTimingNode, mode)
+        if mode === :slots
+            error("unsupported in Julia 1.12")
+        elseif mode === :spec
+            return frame_name(io, MethodInstance(node))
+        else
+            return func_name(MethodInstance(node), mode)
+        end
+    end
 
-    mistr = Symbol(func_name(io, MethodInstance(node), mode))
+    mistr = Symbol(func_name(io, node, mode))
     mi = MethodInstance(node)
     m = mi.def
     sf = isa(m, Method) ? StackFrame(mistr, mi.def.file, mi.def.line, mi, false, false, UInt64(0x0)) :
@@ -1899,14 +1910,10 @@ function _flamegraph_frame(io::IO, node::InferenceTimingNode, start_secs, check_
     if hasconstprop(InferenceTiming(node))
         status |= FlameGraphs.gc_event
     end
-    start = node.start_time - start_secs
-    if toplevel
-        # Compute a range over the whole profile for the top node.
-        stop_secs = isROOT(node) ? max_end_time(node) : max_end_time(node, true)
-        range = round(Int, start*1e9) : round(Int, (stop_secs - start_secs)*1e9)
-    else
-        range = round(Int, start*1e9) : round(Int, (start + inclusive(node))*1e9)
+    if stop_secs === nothing
+        stop_secs = duration(node) + start_secs
     end
+    range = round(Int, start_secs*1e9) : round(Int, (stop_secs - start_secs)*1e9)
     return FlameGraphs.NodeData(sf, status, range), check_precompilable
 end
 

test/snoop_inference.jl

@@ -633,7 +633,7 @@ end
     fgnodes = collect(AbstractTrees.PreOrderDFS(fg))
     for tgtname in (:h, :i, :+)
         @test mapreduce(|, fgnodes; init=false) do node
-            node.data.sf.linfo.def.name == tgtname
+            SnoopCompile.methodinstance(node.data.sf.linfo).def.name == tgtname
         end
     end
     # Test that the span covers the whole tree, and check for const-prop