Zoom HPC with Julia: the n-body problem as an example

Last updated: January 12, 2023

Zoom session 3 — part 1

Time:   2:30 pm Pacific Time.
Topic:   In this session, we will solve the n-body problem in Julia.

Sequential script

Julia script

seq.jl 

##############################################################################
# Sequential implementation
##############################################################################

using BenchmarkTools

n = 10
x = 2rand(n, 3) .- 1
v = zeros(n, 3)

nout = 100
courant = 1e-3
eps = 0.001

v[:, 1] = 2 .* x[:, 2]
v[:, 2] = .- x[:, 1] ./ x[:, 2] .* v[:, 1]

xnew, vnew = zeros(n, 3), zeros(n, 3)
vhalf = zeros(n, 3)

step = Int(3e4)

history = zeros(n, 3, Int(step/nout)+1)
history[:, :, 1] = x

@btime for t in 1:step
    force = zeros(n, 3)
    tmin = 1.e10
    for i in 1:n
        for j in 1:n
            if i != j
                distSquared = sum((x[i, :] - x[j, :]).^2)
                force[i, :] -= (x[i, :] - x[j, :]) / distSquared^1.5
                tmin = min(tmin, sqrt((distSquared+eps) / sum((v[i, :] - v[j, :]).^2)))
            end
        end
    end
    dt = tmin * courant
    # dt = tmin
    for i in n
        xnew[i, :] .= x[i, :] .+ v[i, :].*dt
        vnew[i, :] .= v[i, :] .+ force[i, :].*dt
    end
    global    x = xnew
    global    v = vnew
    if (t+1)%nout == 0
        history[:, :, trunc(Int, t/nout)+1] = x
    end
end

sbatch script

seq.sh 

#!/bin/bash
#SBATCH --time=5
#SBATCH --mem=2G
#SBATCH --output=%j.out
#SBATCH --error=%j.err

echo Sequential: running non parallel script
julia seq.jl

Shared memory script

Julia script

sha.jl 

##############################################################################
# Shared memory implementation
##############################################################################

using Base.Threads
using BenchmarkTools

n = 10
x = 2rand(n, 3) .- 1
v = zeros(n, 3)

nout = 100
courant = 1e-3
eps = 0.001

v[:, 1] = 2 .* x[:, 2]
v[:, 2] = .- x[:, 1] ./ x[:, 2] .* v[:, 1]

xnew, vnew = zeros(n, 3), zeros(n, 3)
vhalf = zeros(n, 3)

step = Int(3e4)

history = zeros(n, 3, Int(step/nout)+1)
history[:, :, 1] = x

@btime @threads for t in 1:step
    force = zeros(n, 3)
    tmin = 1.e10
    for i in 1:n
        for j in 1:n
            if i != j
                distSquared = sum((x[i, :] - x[j, :]).^2)
                force[i, :] -= (x[i, :] - x[j, :]) / distSquared^1.5
                tmin = min(tmin, sqrt((distSquared+eps) / sum((v[i, :] - v[j, :]).^2)))
            end
        end
    end
    dt = tmin * courant
    # dt = tmin
    for i in n
        xnew[i, :] .= x[i, :] .+ v[i, :].*dt
        vnew[i, :] .= v[i, :] .+ force[i, :].*dt
    end
    global    x = xnew
    global    v = vnew
    if (t+1)%nout == 0
        history[:, :, trunc(Int, t/nout)+1] = x
    end
end

sbatch script

sha.sh 

#!/bin/bash
#SBATCH --time=5
#SBATCH --cpus-per-task=4
#SBATCH --mem=8G
#SBATCH --output=%j.out
#SBATCH --error=%j.err

echo Shared memory: running parallel script on $SLURM_CPUS_PER_TASK cores
julia -t 4 sha.jl

Distributed memory script

Julia script

dis.jl

sbatch script

dis.sh 

#!/bin/bash
#SBATCH --time=10
#SBATCH --nodes=2
#SBATCH --cpus-per-task=4
#SBATCH --mem=8G
#SBATCH --output=%j.out
#SBATCH --error=%j.err

echo Distributed memory: running parallel script on 2 nodes and $SLURM_CPUS_PER_TASK cores per node
julia -p 2 -t 4 dis.jl

Comments & questions