Suggestions for performance improvement in Julia - performance

Question

I'm making my first effort to move from Matlab to Julia and have found my code to improve by ~3x but still think there is more to come, I'm not using any global variables in the function and have preallocated all the arrays used (I think?). If there was any thoughts on how it could be sped up even further it would be greatly appreciated, I'll fully convert even at the current improvement I think!
function word_sim(tau::Int, omega::Int, mu::Float64)
# inserts a word in position (tau+1), at each point creates a new word with prob mu
# otherwise randomly chooses a previously used. Runs the program until time omega
words = zeros(Int32, 1, omega) # to store the words
tests = rand(1,omega) # will compare mu to these
words[1] = 1; # initialize the words
next_word = 2 # will be the next word used
words[tau+1] = omega + 1; # max possible word so insert that at time tau
innovates = mu .> tests; # when we'll make a new word
for i = 2:tau # simulate the process
if innovates[i] == 1 # innovate
words[i] = next_word
next_word = next_word + 1
else # copy
words[i] = words[rand(1:(i-1))]
end
end
# force the word we're interested in
for i = (tau+2):omega
if innovates[i] == 1 # innovate
words[i] = next_word
next_word = next_word + 1
else # copy
words[i] = words[rand(1:(i-1))]
end
end
result = sum(words .== (omega + 1)); # count how many times our word occurred
return result
end
and when I run it with these values it takes ~.26 seconds on my PC
using Statistics
#time begin
nsim = 10^3;
omega = 100;
seed = [0:1:(omega-1);];
mu = 0.01;
results = zeros(Float64, 1, length(seed));
pops = zeros(Int64, 1, nsim);
for tau in seed
for jj = 1:nsim
pops[jj] = word_sim(tau, omega, mu);
end
results[tau+1] = mean(pops);
end
end
Or perhaps I'd be better writing the code in C++? Julia was my first reaction as I've heard rave reviews about its syntax, which to be honest is fantastic!
Any comments greatly appreciated.

Answer 1

A 3x speedup is a nice start, but it turns out there are a few more things you can do to improve performance significantly!
As a starting point, using your example posted above in Julia 1.6.1, I get
0.301665 seconds (798.10 k allocations: 164.778 MiB, 12.70% gc time)
That's a lot of allocations, and a fair amount of garbage collector ("gc") time, so it seems we're producing a fair amount of garbage here. Some of the culprits are lines like
tests = rand(1,omega) # will compare mu to these
or
innovates = mu .> tests; # when we'll make a new word
In languages like Matlab or Python, pre-calculating these things whole-vector-at-a-time can be good for performance, but in Julia it's generally not really necessary, and can even hurt because each of these lines is causing a brand new array to be allocated. If we remove these and just generate our tests on the fly, we can avoid these allocations. One other line that allocates in here is
result = sum(words .== (omega + 1))
where you first build a whole new array before taking the sum of it. You could avoid this by writing it as a for loop (even though this may feel wrong coming from Matlab, it's quite fast in Julia). Or, to keep it as a one-liner, use either count or sum with a function that does the comparison as the first argument
result = count(x->(x == omega+1), words)
(in this example, just using an anonymous function x->(x == omega+1)).
Adding up these changes so far then
function word_sim(tau::Int, omega::Int, mu::Float64)
# inserts a word in position (tau+1), at each point creates a new word with prob mu
# otherwise randomly chooses a previously used. Runs the program until time omega
words = zeros(Int32, 1, omega) # to store the words
words[1] = 1; # initialize the words
next_word = 2 # will be the next word used
words[tau+1] = omega + 1; # max possible word so insert that at time tau
for i = 2:tau # simulate the process
if mu > rand() # innovate
words[i] = next_word
next_word = next_word + 1
else # copy
words[i] = words[rand(1:(i-1))]
end
end
# force the word we're interested in
for i = (tau+2):omega
if mu > rand() # innovate
words[i] = next_word
next_word = next_word + 1
else # copy
words[i] = words[rand(1:(i-1))]
end
end
result = count(x->(x == omega+1), words) # count how many times our word occurred
return result
end
Using the same timing code, this now brings us down to
0.177766 seconds (298.10 k allocations: 51.863 MiB, 13.01% gc time)
So about half the time and half the allocations. There's still more though!
First, let's move the allocation of the words array outside of the word_sim function and instead make an in-place version of that function. We can also speed things up a adding an #inbounds to the tight for loops.
function word_sim!(words::AbstractArray, tau::Int, omega::Int, mu::Float64)
# inserts a word in position (tau+1), at each point creates a new word with prob mu
# otherwise randomly chooses a previously used. Runs the program until time omega
fill!(words, 0) # Probably not necessary actually, but I haven't spent enough time looking at the code to be sure
words[1] = 1; # initialize the words
next_word = 2 # will be the next word used
words[tau+1] = omega + 1; # max possible word so insert that at time tau
#inbounds for i = 2:tau # simulate the process
if mu > rand() # innovate
words[i] = next_word
next_word = next_word + 1
else # copy
words[i] = words[rand(1:(i-1))]
end
end
# force the word we're interested in
#inbounds for i = (tau+2):omega
if mu > rand() # innovate
words[i] = next_word
next_word = next_word + 1
else # copy
words[i] = words[rand(1:(i-1))]
end
end
result = count(x->(x == omega+1), words) # count how many times our word occurred
return result
end
In-place functions that modify one of their input arguments are usually denoted by a ! at the end of their name by convention in Julia, hence the new function name.
Since we have to modify the timing code a bit to pre-allocate words now, let's also take the opportunity to put that timing code into a function to avoid any globals in the timing.
function run_word_sim()
nsim = 10^3
omega = 100
seed = [0:1:(omega-1);]
mu = 0.01
results = zeros(Float64, 1, length(seed))
pops = zeros(Int64, 1, nsim)
words = zeros(Int32, 1, omega) # to store the words
for tau in seed
for jj = 1:nsim
pops[jj] = word_sim!(words, tau, omega, mu)
end
results[tau+1] = mean(pops)
end
return results
end
Then get the most accurate timing results (and optionally some useful plots and statistics) we can use the BenchmarkTools package and its #btime or #benchmark macros
julia> using BenchmarkTools
julia> #btime run_word_sim()
124.178 ms (4 allocations: 10.17 KiB)
or
So, almost another 3x speedup, and reduced allocations and memory usage (by four or five orders of magnitude) down to only the four arrays used in the timing code (seed, results, pops and words).
For the absolute maximum performance, you could possibly go even farther with LoopVectorization.jl and its #turbo macro, though it would likely require a change in algorithm since these loops depend on previous state, so don't appear to be compatible with loop re-ordering. You could turn the count into a for loop and #turbo that for a slight additional speedup though.
There are also other options for potentially faster random number generation, such as VectorizedRNG.jl as discussed in the discourse thread linked in the comments. While allocating a new vector of random numbers on each call of word_sim is likely not optimal, RNG is generally faster when you can generate a lot of random numbers at once, so passing a pre-allocated buffer of random numbers to word_sim! and filling that in-place with rand! as provided by either the Random stdlib or VectorizedRNG could yield a significant additional speedup.
Some of the tricks and rules of thumb used in this answer are discussed more generally in https://github.com/brenhinkeller/JuliaAdviceForMatlabProgrammers, along with a few other general Matlab -> Julia tips.