|
| 1 | +# TODO: Keep this in a separate module? |
| 2 | +module Allocs |
| 3 | + |
| 4 | +# Most of this file was copied from the PProf.jl package, and then adapted to |
| 5 | +# export a profile of the heap profile data from this package. |
| 6 | +# This code is pretty hacky, and I could probably do a better job re-using |
| 7 | +# logic from the PProf package, but :shrug:. |
| 8 | + |
| 9 | + |
| 10 | +import Profile # For Profile.Allocs structures |
| 11 | + |
| 12 | +# Import the PProf generated protobuf types from the PProf package: |
| 13 | +import PProf |
| 14 | +using PProf.perftools.profiles: ValueType, Sample, Function, Location, Line, Label |
| 15 | +using PProf: _enter! |
| 16 | +const PProfile = PProf.perftools.profiles.Profile |
| 17 | +using Base.StackTraces: StackFrame |
| 18 | + |
| 19 | +using PProf.ProtoBuf |
| 20 | +using PProf.OrderedCollections |
| 21 | + |
| 22 | +# input: e.g. "maybe_handle_const_call! at ./compiler/ssair/inlining.jl:1243" |
| 23 | +function parse_location(loc_str::String) |
| 24 | + at_split = split(loc_str, " at "; limit=2) |
| 25 | + function_name = at_split[1] |
| 26 | + file_and_line = at_split[2] |
| 27 | + colon_split = split(file_and_line, ":") |
| 28 | + file = colon_split[1] |
| 29 | + line = parse(Int, colon_split[2]) |
| 30 | + |
| 31 | + return (;function_name = function_name, file = file, line = line) |
| 32 | +end |
| 33 | + |
| 34 | +function to_pprof(alloc_profile::Profile.Allocs.AllocResults |
| 35 | + ; |
| 36 | + web::Bool = true, |
| 37 | + webhost::AbstractString = "localhost", |
| 38 | + webport::Integer = 62261, # Use a different port than PProf (chosen via rand(33333:99999)) |
| 39 | + out::AbstractString = "alloc-profile.pb.gz", |
| 40 | + from_c::Bool = true, |
| 41 | + drop_frames::Union{Nothing, AbstractString} = nothing, |
| 42 | + keep_frames::Union{Nothing, AbstractString} = nothing, |
| 43 | + ui_relative_percentages::Bool = true, |
| 44 | + # TODO: decide how to name this: |
| 45 | + aggregate_by_type::Bool = true, |
| 46 | + ) |
| 47 | + period = UInt64(0x1) |
| 48 | + |
| 49 | + @assert !isempty(basename(out)) "`out=` must specify a file path to write to. Got unexpected: '$out'" |
| 50 | + if !endswith(out, ".pb.gz") |
| 51 | + out = "$out.pb.gz" |
| 52 | + @info "Writing output to $out" |
| 53 | + end |
| 54 | + |
| 55 | + string_table = OrderedDict{AbstractString, Int64}() |
| 56 | + enter!(string) = _enter!(string_table, PProf._escape_name_for_pprof(string)) |
| 57 | + enter!(::Nothing) = _enter!(string_table, "nothing") |
| 58 | + ValueType!(_type, unit) = ValueType(_type = enter!(_type), unit = enter!(unit)) |
| 59 | + |
| 60 | + # Setup: |
| 61 | + enter!("") # NOTE: pprof requires first entry to be "" |
| 62 | + |
| 63 | + funcs_map = Dict{String, UInt64}() |
| 64 | + functions = Vector{Function}() |
| 65 | + |
| 66 | + locs_map = Dict{StackFrame, UInt64}() |
| 67 | + locations = Vector{Location}() |
| 68 | + |
| 69 | + sample_type = [ |
| 70 | + ValueType!("allocs", "count"), # Mandatory |
| 71 | + ValueType!("size", "bytes") |
| 72 | + ] |
| 73 | + |
| 74 | + prof = PProfile( |
| 75 | + sample = [], location = [], _function = [], |
| 76 | + mapping = [], string_table = [], |
| 77 | + sample_type = sample_type, default_sample_type = 2, # size |
| 78 | + period = period, period_type = ValueType!("heap", "bytes") |
| 79 | + ) |
| 80 | + |
| 81 | + if drop_frames !== nothing |
| 82 | + prof.drop_frames = enter!(drop_frames) |
| 83 | + end |
| 84 | + if keep_frames !== nothing |
| 85 | + prof.keep_frames = enter!(keep_frames) |
| 86 | + end |
| 87 | + |
| 88 | + function maybe_add_location(frame::StackFrame)::UInt64 |
| 89 | + return get!(locs_map, frame) do |
| 90 | + loc_id = UInt64(length(locations) + 1) |
| 91 | + |
| 92 | + # Extract info from the location frame |
| 93 | + (function_name, file_name, line_number) = |
| 94 | + string(frame.func), string(frame.file), frame.line |
| 95 | + |
| 96 | + ## Decode the IP into information about this stack frame |
| 97 | + #if (!from_c && location_from_c) |
| 98 | + # continue |
| 99 | + #end |
| 100 | + |
| 101 | + function_id = get!(funcs_map, function_name) do |
| 102 | + func_id = UInt64(length(functions) + 1) |
| 103 | + |
| 104 | + # Store the function in our functions dict |
| 105 | + funcProto = Function() |
| 106 | + funcProto.id = func_id |
| 107 | + file = function_name |
| 108 | + simple_name = function_name |
| 109 | + # TODO: Get full name with arguments from profile data |
| 110 | + local full_name_with_args |
| 111 | + # WEIRD TRICK: By entering a separate copy of the string (with a |
| 112 | + # different string id) for the name and system_name, pprof will use |
| 113 | + # the supplied `name` *verbatim*, without pruning off the arguments. |
| 114 | + # So even when full_signatures == false, we want to generate two `enter!` ids. |
| 115 | + funcProto.system_name = enter!(simple_name) |
| 116 | + #if full_signatures |
| 117 | + # funcProto.name = enter!(full_name_with_args) |
| 118 | + #else |
| 119 | + funcProto.name = enter!(simple_name) |
| 120 | + #end |
| 121 | + file = Base.find_source_file(file_name) |
| 122 | + file = file !== nothing ? file : file_name |
| 123 | + funcProto.filename = enter!(file) |
| 124 | + push!(functions, funcProto) |
| 125 | + |
| 126 | + return func_id |
| 127 | + end |
| 128 | + |
| 129 | + locationProto = Location(;id = loc_id, |
| 130 | + line=[Line(function_id = function_id, line = line_number)]) |
| 131 | + push!(locations, locationProto) |
| 132 | + |
| 133 | + return loc_id |
| 134 | + end |
| 135 | + end |
| 136 | + |
| 137 | + function construct_location_for_type(typename) |
| 138 | + # TODO: Lol something less hacky than this: |
| 139 | + return maybe_add_location(StackFrame("Alloc: $(typename)", "nothing", 0)) |
| 140 | + end |
| 141 | + |
| 142 | + for sample in alloc_profile.allocs # convert the sample.stack to vector of location_ids |
| 143 | + # for each location in the sample.stack, if it's the first time seeing it, |
| 144 | + # we also enter that location into the locations table |
| 145 | + location_ids = UInt64[ |
| 146 | + maybe_add_location(location) |
| 147 | + for location in sample.stacktrace |
| 148 | + ] |
| 149 | + |
| 150 | + if aggregate_by_type |
| 151 | + # Add location_id for the type: |
| 152 | + pushfirst!(location_ids, construct_location_for_type(sample.type)) |
| 153 | + end |
| 154 | + |
| 155 | + # report the value: allocs = 1 (count) |
| 156 | + # report the value: size (bytes) |
| 157 | + value = [ |
| 158 | + 1, # allocs |
| 159 | + sample.size, # bytes |
| 160 | + ] |
| 161 | + # TODO: Consider reporting a label? (Dangly thingy) |
| 162 | + |
| 163 | + labels = Label[ |
| 164 | + Label(key = enter!("bytes"), num = sample.size, num_unit = enter!("bytes")), |
| 165 | + ] |
| 166 | + if !aggregate_by_type |
| 167 | + push!(labels, Label(key = enter!("type"), str = enter!(sample.type))) |
| 168 | + end |
| 169 | + |
| 170 | + push!(prof.sample, Sample(;location_id = location_ids, value = value, label = labels)) |
| 171 | + end |
| 172 | + |
| 173 | + |
| 174 | + # Build Profile |
| 175 | + prof.string_table = collect(keys(string_table)) |
| 176 | + # If from_c=false funcs and locs should NOT contain C functions |
| 177 | + prof._function = functions |
| 178 | + prof.location = locations |
| 179 | + |
| 180 | + # Write to disk |
| 181 | + open(out, "w") do io |
| 182 | + writeproto(io, prof) |
| 183 | + end |
| 184 | + |
| 185 | + if web |
| 186 | + PProf.refresh(webhost = webhost, webport = webport, file = out, |
| 187 | + ui_relative_percentages = ui_relative_percentages, |
| 188 | + ) |
| 189 | + end |
| 190 | + |
| 191 | + out |
| 192 | +end |
| 193 | + |
| 194 | +end # module Allocs |
0 commit comments