Fix slow pprof for Alloc profiles: Use integer keys; stop stringifying StackFrames (it's expensive)

.github/workflows/CI.yml

@@ -24,6 +24,8 @@ jobs:
           - x64
     steps:
       - uses: actions/checkout@v3
+        with:
+          fetch-depth: 0  # To support cross-branch testing in regression tests
       - uses: julia-actions/setup-julia@v1
         with:
           version: ${{ matrix.version }}

Project.toml

@@ -25,7 +25,9 @@ julia = "1.6"
 pprof_jll = "0.1, 1"
 
 [extras]
+Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"
+InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test"]
+test = ["InteractiveUtils", "Revise", "Test"]

src/Allocs.jl

@@ -65,9 +65,10 @@ function pprof(alloc_profile::Profile.Allocs.AllocResults = Profile.Allocs.fetch
     functions = Vector{Function}()
 
     # NOTE: It's a bug to use the actual StackFrame itself as a key in a dictionary, since
-    # different StackFrames can compare the same sometimes! 🙀 So we use its string
+    # different StackFrames can compare the same sometimes! 🙀 So we manually compute an ID
     # representation as the key. See: https://github.com/JuliaPerf/PProf.jl/issues/69
-    locs_map  = Dict{String, UInt64}()
+    # (We previously stringified the function, but that is very slow.)
+    locs_map  = Dict{Any, UInt64}()
     locations = Vector{Location}()
     samples = Vector{Sample}()
 
@@ -81,13 +82,17 @@ function pprof(alloc_profile::Profile.Allocs.AllocResults = Profile.Allocs.fetch
 
     function maybe_add_location(frame::StackFrame)::UInt64
         # See: https://github.com/JuliaPerf/PProf.jl/issues/69
-        frame_key = string(frame)
+        function_key = method_instance_id(frame)
+        line_number = frame.line
+        # For allocations profile, we uniquely identify each frame position by
+        # (function,line_number) pairs. That identifies a position in the code.
+        frame_key = (function_key, line_number)
         return get!(locs_map, frame_key) do
             loc_id = UInt64(length(locations) + 1)
 
             # Extract info from the location frame
-            (function_name, file_name, line_number) =
-                string(frame.func), string(frame.file), frame.line
+            (function_name, file_name) =
+                string(frame.func), string(frame.file)
 
             # Use a unique function id for the frame:
             function_key = method_instance_id(frame)

src/PProf.jl

@@ -33,13 +33,7 @@ NOTE: We must use Int64 throughout this package (regardless of system word-size)
 proto file specifies 64-bit integers.
 """
 function _enter!(dict::OrderedDict{T, Int64}, key::T) where T
-    if haskey(dict, key)
-        return dict[key]
-    else
-        l = Int64(length(dict))
-        dict[key] = l
-        return l
-    end
+    return get!(dict, key, Int64(length(dict)))
 end
 
 using Base.StackTraces: StackFrame

test/golden/regression_test.jl

@@ -0,0 +1,98 @@
+# This file will test that your changes produce the same profiles as it did on main.
+# It can be a bit hard to tell if a change is going to affect the profiles or not, so this
+# file can act as a sanity test.
+# It requires you to have committed your changes before running.
+
+# Here is an example of what a failure looks like, where I purposefully reintroduced the bug
+# fixed in #70. You can see that the bug is that some method instances have been incorrectly
+# replaced with other methods of the same function:
+#=
+Main binary filename not available.
+Alloc Profiles golden testing: Test Failed at /Users/nathandaly/.julia/dev/PProf/test/golden/regression_test.jl:25
+  Expression: nodiff
+Stacktrace:
+ [1] macro expansion
+   @ ~/builds/julia-1.8/usr/share/julia/stdlib/v1.8/Test/src/Test.jl:464 [inlined]
+ [2] compare_profiles(name::String)
+   @ Main ~/.julia/dev/PProf/test/golden/regression_test.jl:25
+┌ Info: Got diff!:
+│ Type: allocs
+│ Showing nodes accounting for 0, 0% of 1355 total
+│       flat  flat%   sum%        cum   cum%
+│          0     0%     0%          1 0.074%  Core.Compiler.IdDict{Int64, Int64}(::Core.Compiler.Generator{Core.Compiler.Iterators.Filter{Core.Compiler.var\"#375#382\", Core.Compiler.Pairs{Int64, Int64, Core.Compiler.LinearIndices{1, Tuple{Core.Compiler.OneTo{Int64}}}, Vector{Int64}}}, Core.Compiler.var\"#374#381\"})
+│          0     0%     0%         -1 0.074%  Core.Compiler.IdDict{Union{Core.Compiler.NewSSAValue, Core.Compiler.OldSSAValue, Core.SSAValue}, Int64}(::Core.Compiler.Generator{Core.Compiler.Iterators.Enumerate{Vector{Union{Core.Compiler.NewSSAValue, Core.Compiler.OldSSAValue, Core.SSAValue}}}, Core.Compiler.var\"#419#421\"})
+│          0     0%     0%         -1 0.074%  _collect(::Type{Int64}, ::Core.Compiler.Generator{Core.Compiler.Iterators.Filter{Core.Compiler.var\"#377#384\"{Core.Compiler.IdDict{Int64, Int64}}, Vector{Int64}}, Core.Compiler.var\"#376#383\"{Int64, Core.Compiler.IdDict{Int64, Int64}, Vector{Int64}}}, ::Core.Compiler.SizeUnknown)
+│          0     0%     0%          1 0.074%  _collect(::Type{Int64}, ::Core.Compiler.Generator{Core.Compiler.Iterators.Filter{Core.Compiler.var\"#379#386\"{Core.Compiler.IdDict{Int64, Int64}}, Vector{Int64}}, Core.Compiler.var\"#378#385\"{Int64, Core.Compiler.IdDict{Int64, Int64}, Vector{Int64}}}, ::Core.Compiler.SizeUnknown)
+│          0     0%     0%          1 0.074%  anymap(::Core.Compiler.var\"#261#262\", ::Vector{Any})
+│          0     0%     0%         -1 0.074%  anymap(::typeof(Core.Compiler.widenconditional), ::Vector{Any})
+│          0     0%     0%         -1 0.074%  append!(::Vector{Any}, ::Vector{Any})
+│          0     0%     0%          1 0.074%  append!(::Vector{Core.Compiler.BasicBlock}, ::Vector{Core.Compiler.BasicBlock})
+│          0     0%     0%         11  0.81%  argextype(::Any, ::Core.Compiler.IRCode, ::Vector{Any}, ::Vector{Any})
+│          0     0%     0%        -11  0.81%  argextype(::Any, ::Core.Compiler.IncrementalCompact, ::Vector{Any}, ::Vector{Any})
+│          0     0%     0%          8  0.59%  copy!(::Vector{Core.Compiler.DomTreeNode}, ::Vector{Core.Compiler.DomTreeNode})
+│          0     0%     0%         -8  0.59%  copy!(::Vector{Int64}, ::Vector{Int64})
+│          0     0%     0%         -1 0.074%  setindex!(::Core.Compiler.IdDict{Any, Union{Nothing, Core.Compiler.LiftedValue}}, ::Any, ::Any)
+│          0     0%     0%          6  0.44%  setindex!(::Core.Compiler.IdDict{Core.Compiler.MethodMatchKey, Union{Core.Compiler.Missing, Core.Compiler.MethodMatchResult}}, ::Any, ::Any)
+│          0     0%     0%         -3  0.22%  setindex!(::Core.Compiler.IdDict{Int64, Int64}, ::Any, ::Any)
+│          0     0%     0%         -1 0.074%  setindex!(::Core.Compiler.IdDict{Union{Core.Compiler.NewSSAValue, Core.Compiler.OldSSAValue, Core.SSAValue}, Any}, ::Any, ::Any)
+│          0     0%     0%         -1 0.074%  setindex!(::Core.Compiler.IdDict{Union{Core.Compiler.NewSSAValue, Core.Compiler.OldSSAValue, Core.SSAValue}, Int64}, ::Any, ::Any)
+│          0     0%     0%          3  0.22%  sort
+└          0     0%     0%         -3  0.22%  sort##kw
+Test Summary:                 | Pass  Fail  Total   Time
+Alloc Profiles golden testing |    1     2      3  19.8s
+ERROR: LoadError: Some tests did not pass: 1 passed, 2 failed, 0 errored, 0 broken.
+in expression starting at /Users/nathandaly/.julia/dev/PProf/test/golden/regression_test.jl:32
+=#
+
+using Profile, PProf
+using Test, Revise
+using InteractiveUtils: peakflops
+
+# Make sure that there's no local diffs (grep returns -1 if empty)
+@assert !success(pipeline(`git status --porcelain=v1`,`grep -v '^??'`))
+
+dir1 = mktempdir()
+dir2 = mktempdir()
+
+function compare_profiles(name)
+    run(`git checkout main`)
+    Revise.revise()
+    @time @eval PProf.Allocs.pprof(web=false, out="$dir1/$($name).pb.gz")
+    run(`git checkout -`)
+    Revise.revise()
+    @time @eval PProf.Allocs.pprof(web=false, out="$dir2/$($name).pb.gz")
+    # Assert that there's no diff:
+    lines = readlines(`$(PProf.pprof_jll.pprof()) -top -diff_base=$dir1/$name.pb.gz $dir2/$name.pb.gz`)
+    nodiff = (lines[end] === "      flat  flat%   sum%        cum   cum%")
+    @test nodiff
+    if !nodiff
+        @info "Got diff!: \n$(join(lines, "\n"))"
+    end
+    #@assert success(`diff $dir1/$name.pb.gz $dir2/$name.pb.gz`) "DIFF: $name"
+end
+
+@testset "Alloc Profiles golden testing" begin
+
+    Profile.Allocs.clear(); peakflops(); Profile.Allocs.@profile sample_rate=1 begin
+        for _ in 1:10 peakflops() end
+    end
+    compare_profiles("1")
+
+
+    my_function() = length(collect(Iterators.flatten(Any[1:500 for _ in 1:1000])))
+
+    Profile.Allocs.clear(); my_function(); Profile.Allocs.@profile sample_rate=0.0001 begin
+        my_function()
+    end
+
+    compare_profiles("2")
+
+    # Compilation
+    @eval my_function2() = length(collect(Iterators.flatten(Any[1:5 for _ in 1:10])))
+    @eval Profile.Allocs.clear(); Profile.Allocs.@profile sample_rate=0.1 begin
+        @eval (while false end; my_function2())
+    end
+
+    compare_profiles("3")
+
+end

test/runtests.jl

@@ -15,4 +15,8 @@ end
     @testset "Allocs.jl" begin
         include("Allocs.jl")
     end
+
+    @testset "Regression tests" begin
+        include("golden/regression_test.jl")
+    end
 end