Download all files from this example.
3D Benchmark
This example benchmarks SLIM 3D in an idealized domain using realistic numerical parameters. The benchmark workflow is split across several files:
benchmark.sh(this file) : launcher script for the full benchmark sweep
mesh.geo: idealized mesh geometry
meshsplit.py: mesh extrusion and partitioning
benchmark.py: simulation setup, timing, and SQLite output
Results from benchmark.py are written to benchmarks.sqlite.
Launcher script
This script generates a family of meshes with different horizontal resolutions, partitions each mesh, and runs the benchmark for several vertical resolutions.
It can be run concurrently by multiple jobs, provided that each job uses a different number of MPI ranks. All benchmark results are collected in the same SQLite database.
set -euo pipefail
# Change this to your device name, e.g. "A100", "MI250X", "CPU", etc.
devname="DEVICE_NAME"
# Select the command to run the Python scripts. Uncomment the one you want to use.
cmd=(python)
# cmd=(srun python)
# cmd=(mpirun -np 8 python)
mkdir -p meshes
for scale in 1.600e+01 1.345e+01 1.131e+01 9.514e+00 8.000e+00 6.727e+00 5.657e+00 4.757e+00 4.000e+00 3.364e+00 2.828e+00 2.378e+00 2.000e+00 1.682e+00 1.414e+00 1.189e+00 1.000e+00 8.409e-01 7.071e-01 5.946e-01 5.000e-01 4.204e-01 3.536e-01 2.973e-01 2.500e-01 2.102e-01 1.768e-01 1.487e-01 1.250e-01 1.051e-01 8.839e-02 7.433e-02 6.250e-02 5.256e-02 4.419e-02 3.716e-02 3.125e-02 2.628e-02; do
outfile="meshes/mesh${scale}.msh"
if [[ ! -f "$outfile" ]]; then
echo "Generating $outfile"
gmsh -2 mesh.geo -clscale "$scale" -o "$outfile" > /dev/null
fi
"${cmd[@]}" meshsplit.py -m "$scale"
for nz in 16 32; do # vertical resolutions to test
echo "Running scale=$scale, nz=$nz, devname=$devname"
"${cmd[@]}" benchmark.py -nz "$nz" -m "$scale" -devname "$devname"
done
done
Download this section.
Mesh file
Simple Gmsh geometry defining a rectangular domain discretized with unstructured triangles.
lx = 160e3;
ly = 500e3;
lc = 4e3;
Point(1) = {0,0,0,lc};
Point(2) = {lx,0,0,lc};
Point(3) = {lx,ly,0,lc};
Point(4) = {0,ly,0,lc};
Line(1) = {1,2};
Line(2) = {2,3};
Line(3) = {3,4};
Line(4) = {4,1};
Curve Loop(1) = {1, 2, 3, 4};
Plane Surface(1) = {1};
Physical Line("Wall") = {1,2,3,4};
Physical Surface("Domain") = {1};
Mesh.Algorithm=5;
Download this section.
Mesh splitting
This script partitions a benchmark mesh if it has not already been split. It is normally launched automatically by benchmark.sh rather than run directly.
The partitioning is only performed if at least one expected partition file is missing for the current MPI size.
import argparse
import os
parser = argparse.ArgumentParser()
parser.add_argument("-m", dest="scale", type=str, help="The scale")
scale = parser.parse_args().scale
from slim4 import slim2d, mpi
already_exists = os.path.exists(f"meshes/mesh{scale}_split/{mpi.size}/{mpi.rank}.msh")
if mpi.min_idx(already_exists) == 0:
slim2d.partition_mesh(
f"meshes/mesh{scale}.msh", f"meshes/mesh{scale}_split/{mpi.size}/${{rank}}.msh"
)
else:
if mpi.rank == 0:
print(f"mesh{scale}.msh already partitioned in {mpi.size} pieces")
Download this section.
Benchmark run
This script runs a short 3D SLIM benchmark on a partitioned idealized mesh and writes the timing results to benchmarks.sqlite.
It is normally launched automatically by benchmark.sh rather than run directly.
import argparse
import sqlite3
import json
import socket
import numpy as np
from slim4 import mpi, scalar, slim, slim2d, slim3d, pre3d
from slim4.slim import si_format
from time import perf_counter as clock
def parse_arguments():
parser = argparse.ArgumentParser()
parser.add_argument("-nz", dest="nz", type=int, help="An nz value")
parser.add_argument("-m", dest="scale", type=str, help="The scale")
parser.add_argument("-devname", dest="devname", type=str, help="Device name")
args = parser.parse_args()
return args.nz, args.scale, args.devname
We parse the benchmark configuration and load the partitioned mesh.
nz, scale, devname = parse_arguments()
ftype = "fp32" if scalar == np.float32 else "fp64"
mesh2d = slim2d.Mesh(f"meshes/mesh{scale}_split/{mpi.size}/{mpi.rank}.msh")
n_tri_tot = mpi.sum_idx(mesh2d.n_triangles)
if mpi.rank == 0:
print(f" Ntri = {n_tri_tot}")
We extrude the 2D mesh and initialize the 3D hydrodynamic problem.
bath2d = slim.DataStatic.create_constant(mesh2d, 1000)
mesh, z0, _bath2d = pre3d.extrude(
mesh2d, bath2d, pre3d.VMapProportional(np.ones(nz, dtype=scalar))
)
X = mesh.create_coordinate_vector(z0)
g = 9.80616
rho0 = 1000
problem = slim3d.Problem(mesh, g, bath2d, rho0)
sw3d = problem.problemUV
ts_equation = problem.add_tracer(2, "temp_sal", limiter_type="past_bound")
n_elem = mesh.z_layers[-1]
problem.set_wall_boundary("Wall")
We define idealized physical parameters and turbulence closures representative of realistic ocean-model configurations.
gamma = slim.DataStatic.create_constant(mesh2d, [0.01])
cor = slim.DataStatic.create_constant(mesh2d, 1.2e-4)
smago_coeff = slim.DataStatic.create_constant(mesh, value=0.02)
okubo_coeff = slim.DataStatic.create_constant(mesh, value=2e-4)
sw3d.set_dissipation(gamma, log_profile=False)
sw3d.set_coriolis(cor)
sw3d.set_horizontal_viscosity(smago_coeff, "smago")
for name, problem_tracer in problem.problem_tracer.items():
problem_tracer.set_horizontal_diffusivity(okubo_coeff, "okubo")
problem.add_turbulence()
The initial temperature field is initialized with a baroclinic front and a localized perturbation to trigger instabilities.
temp_sal = ts_equation.new_vector()
temperature = temp_sal.get_sub(0)
salinity = temp_sal.get_sub(1)
problem.set_state_equation_linear([(temperature, -0.2), (salinity, 0.1)])
z = slim.ScalarArrayDeviceMPI.create_from_host(z0)
uv = problem.problemUV.new_vector()
temp_bot = 10.1
temp_top = 13.1
z_bot = -975
d_temp = 1.2
y_0 = 250e3
y_a = 40e3
k = 3
Lx = 160e3
d_y = 40e3
x, y, _ = X
temp_background = temp_bot + ((temp_top - temp_bot) * (z_bot - z0[0]) / z_bot)
front_center = y_0 - y_a * np.sin(2 * np.pi * k * x / Lx)
perturbation_center = y_0 - y_a / 2 * np.sin(np.pi * ((x - 110e3) / 20e3))
temp_perturbation = 0.3 * (1 - ((y - perturbation_center) / (d_y / 2)))
temp_init = np.zeros((6, n_elem))
mask_hot = y > front_center + d_y
temp_init[mask_hot] = temp_background[mask_hot]
mask_cold = y < front_center
temp_init[mask_cold] = (temp_background - d_temp)[mask_cold]
mask_transition = ~mask_hot & ~mask_cold
temp_init[mask_transition] = (
temp_background - d_temp * (1 - ((y - front_center) / d_y))
)[mask_transition]
mask_perturbation = (
(x > 110e3)
& (x < 130e3)
& (y > perturbation_center - (d_y / 2))
& (y < perturbation_center + (d_y / 2))
)
temp_init[mask_perturbation] += temp_perturbation[mask_perturbation]
temp_init = temp_init.reshape((1, 6, -1))
temperature.set_from_host(temp_init)
salinity.fill(35.0)
We compute stable time steps and configure the benchmark duration.
dt3d, dt2d = problem.compute_cfl()
dt3d = min(dt3d, 10 * dt2d) # realistic value for M : 10-15
dt = dt3d
ddt = dt2d
M = dt / ddt
t = 0.0
nstep = 100 # for benchmark
Run the benchmark and record iteration timings.
ts = np.zeros(nstep + 1)
tic = clock()
for i in range(nstep):
t = slim3d.iterate_midpoint(problem, t, dt, M, 1, z, uv, [temp_sal])
slim.synchronize()
ts[i + 1] = clock() - tic
if (i + 1) > 10 and ts[i + 1] > 10:
break
ts = ts[: i + 2]
if mpi.max_idx(uv.has_nan()):
raise ValueError("NaN values in velocity")
Timing statistics are written to a shared SQLite database for later analysis.
if mpi.rank == 0:
dts = np.diff(ts[1:]) # skip first value to avoid cold start
print(f" min dt {si_format(dts.min(), 4, ' ')}s")
print(f" median dt {si_format(np.median(dts), 4, ' ')}s")
print(f" mean dt {si_format(dts.mean(), 4, ' ')}s")
conn = sqlite3.connect("benchmarks.sqlite", timeout=60)
conn.execute("PRAGMA journal_mode=WAL")
conn.execute("PRAGMA busy_timeout=60000")
conn.execute("""
CREATE TABLE IF NOT EXISTS benchmark (
id INTEGER PRIMARY KEY AUTOINCREMENT,
hostname TEXT NOT NULL,
devname TEXT NOT NULL,
dtype TEXT NOT NULL,
n_triangles INTEGER NOT NULL,
n_layers INTEGER NOT NULL,
mpi_size INTEGER NOT NULL,
timing_json TEXT NOT NULL,
created_at TEXT DEFAULT CURRENT_TIMESTAMP
)
""")
conn.execute(
"""
INSERT INTO benchmark
(hostname, devname, dtype, n_triangles, n_layers, mpi_size, timing_json)
VALUES (?, ?, ?, ?, ?, ?, ?)
""",
(
socket.gethostname(),
devname,
ftype,
n_tri_tot,
nz,
mpi.size,
json.dumps(dts.tolist()),
),
)
conn.commit()
conn.close()
Download this section.