Optimizing Julia code

This session provides an introduction to optimizing Julia code. The examples are developed with Julia v1.9.3. You can download all files from the summer school website:

This website renders the content of optimizing_julia.jl.

First, we install all required packages

import Pkg
Pkg.activate(@__DIR__)
Pkg.instantiate()

The markdown file is created from the source code using Literate.jl. You can create the markdown file via

using Literate
Literate.markdown(joinpath(@__DIR__, "optimizing_julia.jl");
                  flavor = Literate.CommonMarkFlavor())

Basics

The Julia manual contains basic information about performance. In particular, you should be familiar with

https://docs.julialang.org/en/v1/manual/performance-tips/
https://docs.julialang.org/en/v1/manual/style-guide/

if you care about performance.

Example: A linear advection simulation

This is a classical partial differential equation (PDE) simulation setup. If you are not familiar with it, just ignore what’s going on - but we need an example. This one is similar to several problematic cases I have seen in practice.

First, we generate initial data and store it in a file.

using HDF5
x = range(-1.0, 1.0, length = 1000)
dx = step(x)
h5open("initial_data.h5", "w") do io
    u0 = sinpi.(x)
    write_dataset(io, "x", collect(x))
    write_dataset(io, "u0", u0)
end

Next, we write our simulation code as a function - as we have learned from the performance tips!

function simulate()
    x, u0 = h5open("initial_data.h5", "r") do io
        read_dataset(io, "x"), read_dataset(io, "u0")
    end

    t = 0
    u = u0
    while t < t_end
        u = u + dt * rhs(u)
        t = t + dt
    end

    return t, x, u, u0
end

function rhs(u)
    du = similar(u)
    for i in eachindex(u)
        im1 = i == firstindex(u) ? lastindex(u) : (i - 1)
        du[i] = -(u[i] - u[im1]) / dx
    end
    return du
end

Now, we can define our parameters, run a simulation, and plot the results.

using Plots, LaTeXStrings

t_end = 2.5
dt = 0.9 * dx
t, x, u, u0 = simulate()
plot(x, u0; label = L"u_0", xguide = L"x")
plot!(x, u; label = L"u")

Maybe you can already spot problems in the code above. Anyway, before you begin optimizing some code, you should test it and make sure it’s correct. Let’s assume this is the case here. You should write tests making sure that the code continues to work as expected while we optimize it. We will not do this here right now.

Profiling

First, we need to measure the performance of our code to see whether it’s already reasonable. For this, you can use @time - or better BenchmarkTools.jl.

using BenchmarkTools
@benchmark simulate()

This shows quite a lot of allocations - typically a bad sign if you are working with numerical simulations.

There are also profiling tools available in Julia. Some of them are already loaded in the Visual Studio Code extension for Julia. If you prefer working in the REPL, you can install ProfileView.jl (@profview simulate()) and PProf.jl (pprof() after creating a profile via @profview).

@profview simulate()

Task: Optimize the code!

You can find one improved version in the file optimizing_julia_solution.jl.

Introduction to generic Julia code and AD

One of the strengths of Julia is that you can use quite a few modern tools like AD. However, you need to learn Julia a bit to use everything efficiently. Here, we concentrate on ForwardDiff.jl.

using ForwardDiff
using Statistics

function foo0(x)
    vector = zeros(typeof(x), 100)
    foo0!(vector, x)
end

function foo0!(vector, scalar)
    for i in eachindex(vector)
        vector[i] = atan(i * scalar) / π
    end

    for _ in 1:1000
        for i in eachindex(vector)
            vector[i] = cos(vector[i])
        end
    end

    return mean(vector)
end

let x = 2 * π
    @show foo0(x)
    @show ForwardDiff.derivative(foo0, x)
    @benchmark foo0($x)
end

The code above is basically a fixed point iteration. By doing some analysis, one could figure out that it should converge to the solution x of x == cos(x), the “Dottie number”. See https://en.wikipedia.org/wiki/Dottie_number

using SimpleNonlinearSolve

function dottie_number()
    f(u, p) = cos(u) - u
    bounds = (0.0, 2.0)
    prob = IntervalNonlinearProblem(f, bounds)
    sol = solve(prob, ITP())
    return sol.u
end

dottie_number()

Next, we introduce a custom struct. Can you spot the problems?

struct MyData1
    scalar
    vector
end

function foo1(x)
    data = MyData1(x, zeros(typeof(x), 100))
    foo1!(data)
end

function foo1!(data)
    (; scalar, vector) = data

    for i in eachindex(vector)
        vector[i] = atan(i * scalar) / π
    end

    for _ in 1:1000
        for i in eachindex(vector)
            vector[i] = cos(vector[i])
        end
    end

    return mean(vector)
end

let x = 2 * π
    @show foo1(x)
    @show ForwardDiff.derivative(foo1, x)
    @benchmark foo1($x)
end

We can fix type instabilities by introducing types explicitly. But we lose the possibility to use AD this way!

struct MyData2
    scalar::Float64
    vector::Vector{Float64}
end

function foo2(x)
    data = MyData2(x, zeros(typeof(x), 100))
    foo1!(data)
end

let x = 2 * π
    @show foo2(x)
    @benchmark foo2($x)
end

let x = 2 * π
    ForwardDiff.derivative(foo2, x)
end

Thus, we need parametric types!

struct MyData3{T}
    scalar::T
    vector::Vector{T}
end

function foo3(x)
    data = MyData3(x, zeros(typeof(x), 100))
    foo1!(data)
end

let x = 2 * π
    @show foo3(x)
    @show ForwardDiff.derivative(foo3, x)
    @benchmark foo3($x)
end

This page was generated using Literate.jl.