Profiling Julia code with BenchmarkTools and Profile: find the slow function

The problem

The analysis on this page follows 2,000 random walkers in a box, molecules or cells alike, and computes after every step their mean squared displacement, the average squared distance from where each started. One run takes well under a second on the machine that built this page, but a sweep of a thousand runs turns every variant into minutes. A rewrite has to wait for two answers: which function eats the time, and whether the fix is faster with the same result. Profiling gives the first, with the Profile standard library sampling where the program is as it runs. BenchmarkTools gives the second. Swap in your own code for the walkers, cut into functions the way the next section explains.

The code

These functions stand in for your analysis. The timing and the profile both report per function call, and the profile names every line by its function, so the work belongs in functions that take inputs and return results, as step, reflect, and msd do. Do not overwrite the slow function with its fix: msd_fast is a new name beside msd, because the last lines of the recipe run both and compare them. simulate takes either one as its first argument, with the slow msd as the default (msd = msd is not a typo), and the seeded Xoshiro(seed) gives both the same walks.

using Random

step(pos, s, rng) = pos .+ s .* randn(rng, size(pos))
reflect(pos, box) = box .- abs.(box .- abs.(pos))     # bounce off the walls at 0 and at box

function msd(pos, start)
    total = 0.0
    for i in axes(pos, 1)                             # one walker at a time
        d = pos[i, :] .- start[i, :]                  # each slice is a new array
        total += sum(d .^ 2)
    end
    return total / size(pos, 1)
end
function msd_fast(pos, start)                         # the fix: numbers instead of slices
    total = 0.0
    for i in axes(pos, 1)
        dx = pos[i, 1] - start[i, 1]
        dy = pos[i, 2] - start[i, 2]
        total += dx * dx + dy * dy
    end
    return total / size(pos, 1)
end

function simulate(msd = msd; n_walkers = 2_000, n_steps = 2_000, s = 0.01, box = 1.0, seed = 1)
    rng = Xoshiro(seed)
    pos = start = box .* rand(rng, n_walkers, 2)
    curve = zeros(n_steps)
    for t in 1:n_steps
        pos = reflect(step(pos, s, rng), box)
        curve[t] = msd(pos, start)                    # the mean squared displacement after each step
    end
    return curve
end

pos, start = rand(Xoshiro(0), 2_000, 2), rand(Xoshiro(1), 2_000, 2)   # inputs for timing msd alone
nothing

The recipe opens with the first look most people take, two @time calls of msd_fast. The first call of a Julia function compiles it, and @time says how much of its time went into compiling. The warm-up line runs both simulations once, so that @profile measures the run and not the compiler. @profile notes where the program is about every millisecond, once Profile.clear() has emptied any earlier samples. Profile.print(format = :flat) lists every line that turned up, with its count, and the keywords sort the rows, drop those under 50 samples, and keep only the thread that ran your code. display prints the full @benchmark report in the middle of a cell. Your own call goes on the line marked # <- your code here.

using BenchmarkTools, Profile, Printf

# ---- a first look with @time: the first call compiles
@time msd_fast(pos, start)
@time msd_fast(pos, start)
simulate(msd); simulate(msd_fast)                     # warm up before profiling

# ---- profile one run
Profile.clear()                                       # empty the old samples
@profile simulate(msd)                                # <- your code here
Profile.print(format = :flat, sortedby = :count, mincount = 50, threads = Threads.threadid())

# ---- time the slow function alone, and its fix
display(@benchmark msd($pos, $start))
@btime msd($pos, $start)
@btime msd_fast($pos, $start)

# ---- profile again after the fix: ten runs, enough samples
Profile.clear()
@profile for _ in 1:10; simulate(msd_fast); end
Profile.print(format = :flat, sortedby = :count, mincount = 50, threads = Threads.threadid())

# ---- did the fix help, and is the answer the same?
t, t_fast = @belapsed(simulate(msd)), @belapsed(simulate(msd_fast))
@printf("one run %.3f s -> %.3f s; 1,000 runs %.1f min -> %.1f min\n", t, t_fast, 1000t / 60, 1000t_fast / 60)
println("same curve: ", simulate(msd) == simulate(msd_fast))

  0.017456 seconds (11.43 k allocations: 653.617 KiB, 99.72% compilation time)
  0.000024 seconds (1 allocation: 16 bytes)
 Count  Overhead File                            Line Function
 =====  ======== ====                            ==== ========
    60         0 @Base/boot.jl                   648 Array
    70        69 @Base/boot.jl                   588 GenericMemory
    80         0 @Base/multidimensional.jl       996 _unsafe_getindex
    80         0 @Base/abstractarray.jl          825 similar
    80         0 @Base/abstractarray.jl          836 similar
    80         0 @Base/array.jl                  407 similar
    85         0 @Base/multidimensional.jl       984 _getindex
    87         0 @Base/abstractarray.jl         1363 getindex
    92         0 In[2]                            10 msd(pos::Matrix{Float64}, start::Matrix{Float…
   123         0 In[2]                             9 msd(pos::Matrix{Float64}, start::Matrix{Float…
   131         0 @Base/broadcast.jl              229 similar
   131         0 @Base/broadcast.jl              230 similar
   131         0 @Base/abstractarray.jl          876 similar
   131         0 @Base/array.jl                  409 similar
   131         0 @Base/abstractarray.jl          877 similar
   132         0 @Base/boot.jl                   671 Array
   142       141 @Base/boot.jl                   649 Array
   143         0 @Base/broadcast.jl              924 copy
   145         0 @Base/broadcast.jl              899 materialize
   202         0 @Base/boot.jl                   661 Array
   215         0 In[2]                            30 simulate(msd::typeof(msd); n_walkers::Int64, …
   243         0 In[2]                            24 simulate(msd::Function)
   245         0 @IJulia/src/eventloop.jl         71 (::IJulia.var"#waitloop##2#waitloop##3"{IJuli…
   245         0 @IJulia/src/eventloop.jl         26 eventloop(socket::ZMQ.Socket, kernel::IJulia.…
   245         0 @IJulia/src/execute_request.jl  129 execute_request(socket::ZMQ.Socket, kernel::I…
   245         0 @Base/loading.jl               3045 include_string(mapexpr::typeof(REPL.softscope…
   245         2 @Base/boot.jl                   489 eval(m::Module, e::Any)
Total snapshots: 245. Utilization: 100% across all threads and tasks. Use the `groupby` kwarg to break down by thread and/or task.
BenchmarkTools.Trial: 10000 samples with 1 evaluation per sample.
 Range (min … max):  120.062 μs …   7.464 ms  ┊ GC (min … max):  0.00% … 96.81%
 Time  (median):     148.178 μs               ┊ GC (median):     0.00%
 Time  (mean ± σ):   239.897 μs ± 453.846 μs  ┊ GC (mean ± σ):  27.41% ± 13.83%

  █▆▄▃▂                                                         ▁
  ██████▆▆▅▄▁▅▅▃▄▄▄▄▃▃▁▃▃▃▁▄▃▃▁▁▁▃▁▄▁▁▁▁▃▄▁▃▁▁▁▁▁▁▅▃▄▄▅▅▆▆▇▇▇▇▇ █
  120 μs        Histogram: log(frequency) by time       2.98 ms <

 Memory estimate: 625.00 KiB, allocs estimate: 16000.
  120.456 μs (16000 allocations: 625.00 KiB)
  2.364 μs (0 allocations: 0 bytes)
 Count  Overhead File                            Line Function
 =====  ======== ====                            ==== ========
    51         0 @Base/broadcast.jl              617 getindex
    51         0 @Base/broadcast.jl              621 _getindex
    51         0 @Base/broadcast.jl              677 _broadcast_getindex
    51         0 @Base/broadcast.jl              704 _broadcast_getindex_evalf
    63         0 @Base/broadcast.jl              229 similar
    63         0 @Base/broadcast.jl              230 similar
    63         0 @Base/abstractarray.jl          876 similar
    63         0 @Base/array.jl                  409 similar
    63         0 @Base/abstractarray.jl          877 similar
    74         0 @Base/simdloop.jl                77 macro expansion
    74         0 @Base/broadcast.jl             1000 macro expansion
    76         0 @Base/broadcast.jl              952 copyto!
    76         0 @Base/broadcast.jl              999 copyto!
    78         0 In[2]                             4 reflect(pos::Matrix{Float64}, box::Float64)
    79         0 @Base/boot.jl                   671 Array
    79         0 @Base/boot.jl                   662 Array
    79         0 @Base/boot.jl                   652 Array
    79         0 @Base/boot.jl                   605 new_as_memoryref
    79        79 @Base/boot.jl                   588 GenericMemory
    81         0 @Random/src/normal.jl           260 randn!(rng::Xoshiro, A::Matrix{Float64})
   133         0 @Random/src/normal.jl           275 randn
   133         0 @Random/src/normal.jl           269 randn
   139         0 @Base/broadcast.jl              899 materialize
   139         0 @Base/broadcast.jl              924 copy
   195         0 In[2]                             3 step(pos::Matrix{Float64}, s::Float64, rng::X…
   273         0 In[2]                            29 simulate(msd::typeof(msd_fast); n_walkers::In…
   304         0 @IJulia/src/eventloop.jl         71 (::IJulia.var"#waitloop##2#waitloop##3"{IJuli…
   304         0 @IJulia/src/eventloop.jl         26 eventloop(socket::ZMQ.Socket, kernel::IJulia.…
   304         0 @IJulia/src/execute_request.jl  129 execute_request(socket::ZMQ.Socket, kernel::I…
   304         0 @Base/loading.jl               3045 include_string(mapexpr::typeof(REPL.softscope…
   304         0 @Base/boot.jl                   489 eval(m::Module, e::Any)
   304         0 In[4]                            20 top-level scope
   304         0 @Profile/src/Profile.jl          60 macro expansion
   304         0 In[4]                            20 macro expansion
   304         0 In[2]                            24 simulate
Total snapshots: 304. Utilization: 100% across all threads and tasks. Use the `groupby` kwarg to break down by thread and/or task.
one run 0.493 s -> 0.039 s; 1,000 runs 8.2 min -> 0.7 min
same curve: true

The page recorded this output the first time it ran the cell and shows that record on every rebuild. In your copy the output prints straight under the cell. The last line compares with ==, which demands equality to the last bit and holds here only because msd_fast adds the same numbers in the same order as msd. A correct fix that sums in another order changes the last digits and prints false, so compare such a curve with ≈ (isapprox), which accepts a relative difference of about 10⁻⁸.

The knobs

@benchmark runs the call again and again, up to 10,000 samples or 5 s by default, and reports the minimum, median, and mean with a histogram. Compare minimums: the code did the same work in every sample, so a slower one lost its time to the rest of the machine or to the garbage collector, which frees the memory of arrays no longer in use and whose share the report gives as GC. That is why the mean of msd lies well above its median. The $ hands the benchmark the values of pos and start instead of the names of two global variables, so the timing leaves out looking them up. @btime prints the same benchmark as one line, minimum and allocations, and @belapsed returns the minimum in seconds. The 16,000 per call of msd are eight per walker: four new arrays (the two slices, d, and d .^ 2) of two allocations each, the array and the block of memory that holds its numbers. The fix brings the count to zero and msd to about a fiftieth of its time. The whole run, where step and reflect remain, gets an order of magnitude faster, and so does the sweep.

In the flat report, Count is the number of samples in which the line was on the call stack, running or waiting for something it called, and Overhead the number in which the line itself was running. Overhead moves between rows from one run to the next: read Count. In[2] in the File column is the cell with the functions, in your copy your file or cell. The rows are sorted by Count, largest at the bottom, without percentages, so divide a Count by the total in the last line. Read up from the bottom, past the wrapper rows with nearly the full count: IJulia's eventloop, eval, and simulate at line 24. simulate has two rows because Julia turns a function with keyword arguments into two methods: a short one, listed at the function line with only msd::Function, that fills in the defaults, and the body, listed with the keywords. The first row of the body is line 30, where simulate calls msd, with most of the samples, and above it line 9 of msd, the slice line. Further up come rows of Julia's Base, materialize and copy in broadcast.jl and Array in boot.jl: the broadcasts .- and .^ making their new arrays, on top of the two copies the slices make. After the fix step leads, and the randn rows above it are its random numbers, already one call on the whole array.

Pitfalls

Timing one call with @time. The first @time line reads milliseconds, and @time itself reports that more than 99 % of them went into compilation. The second reads microseconds, still far above the @btime minimum of msd_fast, with one allocation the function never makes: a single call of a fast function times the clock and the global variables along with it. A sweep compiles once, not a thousand times, so neither line is what a run costs you. Use @btime or @benchmark with $, which compile first and then repeat the call, and warm up a long run before @time or @profile, as the recipe does.

Too few samples. The ten runs of the fast simulation gave a few hundred samples, one run alone would give a few dozen, and a function with 5 % of the time would then get one or two, which is noise. Profile a loop of runs until the total is in the hundreds, or shorten the interval with Profile.init(delay = 0.0001), and confirm a suspect with @btime before you rewrite it.

Optimizing what takes a small share. Before the fix, step and reflect get the samples of line 24 that line 30 does not have, fewer than msd gets and well under half the total, so making both infinitely fast would not halve the run. Begin at the first row of your own code past the wrappers. Once the leading rows are calls that each handle a whole array at a time, like step and its randn after the fix, rewriting lines has little left to give: the gain is in the method.