main.jmd

# Setup
This file will act as a guide through optimising a simple monte carlo problem, and running the code at different scales. 

To begin, let's just load in some packages that will be useful throughout.

```julia
using Revise
using BenchmarkTools
using Plots
using Random
Random.seed!(1234)
```

# Monte Carlo Problem

Let's introduce a monte carlo problem, which will be what we will be optimising to run at different scales today. This problem will be just a simple random walk, which can be plotted:

```julia
begin
    plt = plot(cumsum(randn(100, 10000), dims=2)', legend=false, c=:blue, alpha=0.1)
    xlabel!(plt, "Time")
    ylabel!(plt, "x")
    return plt
end
```

Let's say that we only care about the final point of the simulation, ``x_T``, where ``T`` is the number of steps taken. We want to sample this ``n`` times. We can write an easy function which calculates this array:

```julia
function simple_monte_carlo(n, T)
    x = zeros(n)
    for i in eachindex(x)
        for t in 1:T
            x[i] += randn()
        end
    end
    return x
end
```

We can time this with `BenchmarkTools.jl`:

```julia
using BenchmarkTools
n = 2^14;
T = 100;
@btime simple_monte_carlo(n, T)
```

## Comparison to Python

We can already see that this is very fast. Out of curiosity, let's translate this into Python, and see how fast this is. First of all, let's setup our ability to call Python with `PyCall.jl`:

```julia
using PyCall
using Pkg
Pkg.build("PyCall")
```

From here, let's write the same function:

```julia
using PyCall
py"""
import random
def simple_monte_carlo(n, T):
    x = [0.0 for _ in range(n)]
    for i in range(n):
        for t in range(T):
            x[i] += random.gauss(0.0,1.0)
    return x
"""
py_simple_monte_carlo_fn = py"simple_monte_carlo"
@btime py_simple_monte_carlo_fn($n, $T)
```

We can benchmark it in the same way:

```julia
@btime simple_monte_carlo($n, $T)
@btime py_simple_monte_carlo_fn($n, $T)
```

Python took around `48.427 ms`, whereas the native Julia implementation took only `298.683 μs`, which is `162` times slower. But this is to be expected. Nobody expects Python to be fast, and this is not how one achieves speed in Python. Usually, one writes these operations as array operations, and delegates the performance to a fast library like `numpy`. Let's try that in both cases, by first adding `numpy`:

```julia
using Conda
Conda.add("numpy")
Pkg.build("PyCall")
```

Then write the function:

```julia
py"""
import numpy as np
def simple_monte_carlo_np(n, T):
    x = np.zeros(n)
    for t in range(T):
        x += np.random.randn(n)

    return x
"""
py_simple_monte_carlo_np_fn = py"simple_monte_carlo_np"
```

Now, let's benchmark:
```julia
@btime py_simple_monte_carlo_np_fn($n, $T)
```

Now this is much better, only `2.119 ms` now, but this is still around `7` times slower than our basic Julia version, which really goes to show you that the compiler is very good and can optimise your code a lot. Let's see if we can take this array approach and use a few techniques to speed up our code:

```julia
function simple_monte_carlo_optimised!(x, T)
    @inbounds for i in eachindex(x)
        x_i = zero(eltype(x))
        for _ in 1:T
            x_i += randn()
        end
        x[i] = x_i
    end
end
```

The `!` is called a "bang" and it is just notation which means that the first argument input is modified by the function. We have moved responsibility for generating `x` away from the function. This means that this function can be called many times and not allocate any memory. Also we have used the `@inbounds` macro to turn off bounds checking on the array, since we are confident that we will not index outside of it, this can be a big performance boost. Let's see what happens:

```julia
x = zeros(n); # Preallocate the array
@btime simple_monte_carlo_optimised!($x, $T)
```

Now we are at `210.704 μs`! This is around `30%` faster now, just by changing a few easy things. Also this function does not allocate.
## Multithreading

We can make this even faster now, by using multithreading, since most CPUs have multiple cores, how do we do this? In Julia, multithreading support is built into the standard libraries, so it is very easy to add:

```julia
function simple_monte_carlo_threaded!(x, T)
    @inbounds Threads.@threads for i in eachindex(x)
        x_i = zero(eltype(x))
        for _ in 1:T
            x_i += randn()
        end
        x[i] = x_i
    end
end
```

We didn't have to change much of the code, just put the `Threads.@threads` macro in front of our for loop, and it will parallelise that loop. Let's test:

```julia
Threads.nthreads()
@btime simple_monte_carlo_threaded!($x, $T)
```

Currently, I have `8` threads and we were able to get `43.062 μs`. This is around a `5` times speedup, showing that threading does have some overhead. This is less noticable when the workload is much larger.

For now, let's take a look at the array approach:

```julia
function simple_monte_carlo_array!(x, T, cache=similar(x))
    fill!(x, zero(eltype(x))) # Fill the array x with zeros
    for _ in 1:T
        randn!(cache) # Put random numbers in the cache array
        x .+= cache # Add the random numbers element-wise to x
    end
end
```

Now we have a cache for the random numbers. This can be created before, or optionally, created on the fly:

```julia
x = zeros(n);
cache = similar(x);
@btime simple_monte_carlo_array!($x, $T, $cache)
```

Notice that this is our fastest implementation so far, without using threads. This approach is not allocating, as it will be the final change to our code for the rest of this walkthrough. A quick note - `similar` is a function which allocates an array which is the same size and type as the input, but does not try to initialise any of the values.

## Multiprocessing

In Julia, you can start and connect multiple processing using the `Distributed.jl` standard library. One can start Julia with multiple processes in the command line as well:
```bash
julia -p 8
```

One can also add processes during runtime with:
```julia
using Distributed
addprocs(8; exeflags=["--project"])
```

One can see the number of processes available with:
```julia
workers()
```

We can also use `DistributedArrays.jl` to parallelise our previous function to a much larger sample, we use the `@everywhere` macro to run it on all processes.

```julia
@everywhere using DistributedArrays
n = 10^7 # 10 million
T = 100
x_dist = dzeros(n);
x_cache_dist = dzeros(n);
```
Since `DistributedArrays.jl` is very bare-bones, we will have to wrap our code to get this working, but due to multiple dispatch, we can still use our original function!

```julia
@everywhere using Random
@everywhere function simple_monte_carlo_array!(x::AbstractArray, T, cache::AbstractArray)
    fill!(x, zero(eltype(x))) # Fill the array x with zeros
    for _ in 1:T
        randn!(cache) # Put random numbers in the cache array
        x .+= cache # Add the random numbers element-wise to x
    end
    nothing
end
@everywhere function simple_monte_carlo_array!(x::DArray, T, cache::DArray)
    if myid() == 1 # Main process (not a worker)
        SPMD.spmd(simple_monte_carlo_array!, x, T, cache; pids=workers())
    else
        simple_monte_carlo_array!(localpart(x), T, localpart(cache))
    end
end
```

Now we can use our previous function, but changing the inputs:
```julia
simple_monte_carlo_array!(x_dist, T, x_cache_dist)
```

This takes `332.710 ms` for `8` process, let's see what the original time is on this case:

```julia
x = Array(x_dist);
cache = similar(x);
@btime simple_monte_carlo_array!($x, $T, $cache)
```

Now we measure the time to be `1.874 s`, which means by moving to a distributed array we have a `6.6` times speedup. Considering we only have `8` processes, and the workload is still small, this is pretty good.

## GPUs

Luckily, Julia has a whole ecosystem of packages designed to write native Julia code which can run on the GPU. CUDA is NVIDIA's proprietary SDK for writing GPU code for their hardware. The bindings for this SDK are available in `CUDA.jl`, which is the most developed amongst the GPU packages. In order to get our code running on the GPU, all we need to do is just convert our input arrays into CUDA arrays:

```julia
using CUDA
x_large_gpu = cu(x_large);
x_large_cache_gpu = similar(x_large_gpu);
```

Now to benchmark, we need to use the `CUDA.@sync` macro to make sure that our timing is accurate:

```julia
@btime CUDA.@sync simple_monte_carlo_array!($x_large_gpu, $T, $x_large_cache_gpu)
```

We can see that our original code was compiled into GPU kernels that run on our graphics card. Here, the execution only took around `33.396ms`, which is around `60` times faster than a single core CPU execution, and all we had to do was move our input arrays to the GPU!

## Benchmarks

As a final part to this, we should benchmark each implementation for different parameters of ``n`` and ``T``. We can start by making sure they are all using the same datatypes. 

```julia
py"""
import numpy as np
def simple_monte_carlo_np_f32(n, T):
    x = np.zeros(n, dtype=np.float32)
    for t in range(T):
        x += np.random.randn(n)
    
    return None
"""
py_simple_monte_carlo_np_f32_fn = py"simple_monte_carlo_np_f32"
```

```julia
function simple_monte_carlo_array_block_threaded!(x, T, cache=similar(x))
    n = length(x) # Number of elements
    n_threads = Threads.nthreads() # Number of threads/cores
    block_size = Int(ceil(n / n_threads)) # Elements per core
    Threads.@threads for i = 0:n_threads-1 # Loop over each core
        start_i = i * block_size + 1
        end_i = min(n, (i+1)*block_size)
        indices = start_i:end_i
        simple_monte_carlo_array!(view(x, indices), T, view(cache,indices))
    end
end
```

```julia
function gather_benchmarks(n, T, element_type=Float32)
    x = zeros(element_type, n)
    cache = similar(x)
    x_dist = distribute(x)
    x_dist_cache = similar(x_dist)
    x_gpu = cu(x)
    x_gpu_cache = similar(x_gpu)

    timings = Dict{Symbol, Float64}();

    timings[:array] = @belapsed simple_monte_carlo_array!($x, $T, $cache)
    timings[:numpy] = element_type == Float32 ? (@belapsed py_simple_monte_carlo_np_f32_fn($n, $T)) : (@belapsed py_simple_monte_carlo_np_fn($n, $T))
    timings[:threaded] = @belapsed simple_monte_carlo_threaded!($x, $T)
    timings[:threaded_blocked] = @belapsed simple_monte_carlo_array_block_threaded!($x, $T, $cache)
    timings[:distributed] = @belapsed simple_monte_carlo_array!($x_dist, $T, $x_dist_cache)
    timings[:gpu] = @belapsed CUDA.@sync simple_monte_carlo_array!($x_gpu, $T, $x_gpu_cache)

    return timings
end
using Random
function gather_benchmarks(n, T)
    x_gpu = CUDA.zeros(n)
    x_gpu_cache = similar(x_gpu)

    timings = Dict{Symbol, Float64}();
    timings[:gpu] = @belapsed CUDA.@sync simple_monte_carlo_array!($x_gpu, $T, $x_gpu_cache) seconds=30

    return timings
end
```

```julia
function plot_bar_chart(timings)
    times = collect(values(results))
    labels = string.(keys(timings))

    bar(1:length(times), times, xticks=(collect(1:length(times)), labels), legend=false, yscale=:log10)
    ylabel!("Time (s)")
end
```

```julia
function gather_results()
    T = 100
    ns = Int.( 2 .^ (3:2:24))

    return (T=T, ns=ns, timings=gather_benchmarks.(ns, T));
end
```

```julia
function compare_speedups(results; ref_label=:array, num_workers=Threads.nthreads(), exclude_symbols=[])
    T = results.T
    ns = results.ns
    timings = results.timings

    labels = collect(keys(timings[begin]))
    ref_times = (x->x[ref_label]).(timings)
    markers = [:circle, :diamond, :utriangle, :star4, :dtriangle]

    plt = plot(ns, (x->num_workers).(ns); linestyle=:dot, markershape=:none, label="# Cores", lw=3, c=:black)
    plt = plot!(plt, ns, (x->1.0).(ns); linestyle=:dash, markershape=:none, label="Array", lw=2, c=:black)
    colours = [:darkred, :purple, :green, :lightblue, :orange, :pink, :brown]
    for (i, label) in enumerate(sort(labels))
        if label == ref_label || label in exclude_symbols
            continue
        end
        l_timings = (x->x[label]).(timings)
        rel_speedup = ref_times ./ l_timings
        plot!(plt, ns, rel_speedup, label=string(label), xscale=:log10, yscale=:log10, markershape=markers[(i-1)%length(markers)+1], lw=2, markersize=6, yformatter=:plain, c=colours[(i-1)%length(colours)+1])
    end
    xlabel!(plt, "n")
    ylabel!(plt, "Relative Speedup")
    xticks!(plt, collect(10 .^ (Int(ceil(log10(minimum(ns)))):Int(ceil(log10(maximum(ns)))))))
    y_ticks = [0.01, 0.1, 1, 10, 100]
    yticks!(plt, (y_ticks, (x-> x % 1 == 0 ? string(Int(x)) : string(x)).(y_ticks)))
    
    plot!(plt; legend=:bottomright, legend_foreground_color=nothing, dpi=300)

    return plt
end
```

Let's see what these results look like:

```julia
using BSON
if isfile("results.bson")
    BSON.@load "results.bson" results
else
    results = gather_results()
    BSON.@save "results.bson" results
end
```

Let's plot these results:
```julia
compare_speedups(results)
```

The GPU is around ``90`` times faster than the array implementation for around ``n=10^7``.