mirror of
https://github.com/DifferentiableUniverseInitiative/JaxPM.git
synced 2025-04-14 07:00:55 +00:00
add optimised bench script
This commit is contained in:
Wassim KABALAN
parent
4f508b7cb6
commit
5f6d42eaeb
1 changed files with 201 additions and 0 deletions
201
scripts/bench_pm.py
Normal file
201
scripts/bench_pm.py
Normal file
|
|
@ -0,0 +1,201 @@
|
|||
import os
|
||||
|
||||
os.environ["EQX_ON_ERROR"] = "nan" # avoid an allgather caused by diffrax
|
||||
import jax
|
||||
|
||||
jax.distributed.initialize()
|
||||
|
||||
rank = jax.process_index()
|
||||
size = jax.process_count()
|
||||
|
||||
import argparse
|
||||
import jax.numpy as jnp
|
||||
import jax_cosmo as jc
|
||||
from jaxpm.painting import cic_paint_dx
|
||||
from jaxpm.pm import linear_field, lpt, make_ode_fn
|
||||
from diffrax import diffeqsolve, ODETerm, Dopri5, LeapfrogMidpoint, SaveAt, PIDController
|
||||
import numpy as np
|
||||
from jax.experimental import mesh_utils
|
||||
from jax.sharding import Mesh, PartitionSpec as P, NamedSharding
|
||||
from jaxpm.kernels import interpolate_power_spectrum
|
||||
import time
|
||||
from cupy.cuda.nvtx import RangePush, RangePop
|
||||
from jax.experimental.multihost_utils import sync_global_devices
|
||||
|
||||
|
||||
def chrono_fun(fun, *args):
|
||||
start = time.perf_counter()
|
||||
out = fun(*args).block_until_ready()
|
||||
end = time.perf_counter()
|
||||
return out, end - start
|
||||
|
||||
|
||||
def run_simulation(mesh_shape,
|
||||
box_size,
|
||||
halo_size,
|
||||
solver_choice,
|
||||
iterations,
|
||||
pdims=None):
|
||||
|
||||
@jax.jit
|
||||
def simulate(omega_c, sigma8):
|
||||
# Create a small function to generate the matter power spectrum
|
||||
k = jnp.logspace(-4, 1, 128)
|
||||
pk = jc.power.linear_matter_power(
|
||||
jc.Planck15(Omega_c=omega_c, sigma8=sigma8), k)
|
||||
pk_fn = lambda x: interpolate_power_spectrum(x, k, pk)
|
||||
|
||||
# Create initial conditions
|
||||
initial_conditions = linear_field(mesh_shape,
|
||||
box_size,
|
||||
pk_fn,
|
||||
seed=jax.random.PRNGKey(0))
|
||||
|
||||
# Create particles
|
||||
cosmo = jc.Planck15(Omega_c=omega_c, sigma8=sigma8)
|
||||
dx, p, _ = lpt(cosmo, initial_conditions, 0.1, halo_size=halo_size)
|
||||
|
||||
# Evolve the simulation forward
|
||||
ode_fn = make_ode_fn(mesh_shape, halo_size=halo_size)
|
||||
term = ODETerm(
|
||||
lambda t, state, args: jnp.stack(ode_fn(state, t, args), axis=0))
|
||||
|
||||
if solver_choice == "Dopri5":
|
||||
solver = Dopri5()
|
||||
elif solver_choice == "LeapfrogMidpoint":
|
||||
solver = LeapfrogMidpoint()
|
||||
else:
|
||||
raise ValueError(
|
||||
"Invalid solver choice. Use 'Dopri5' or 'LeapfrogMidpoint'.")
|
||||
|
||||
stepsize_controller = PIDController(rtol=1e-4, atol=1e-4)
|
||||
res = diffeqsolve(term,
|
||||
solver,
|
||||
t0=0.1,
|
||||
t1=1.,
|
||||
dt0=0.01,
|
||||
y0=jnp.stack([dx, p], axis=0),
|
||||
args=cosmo,
|
||||
saveat=SaveAt(t1=True),
|
||||
stepsize_controller=stepsize_controller)
|
||||
|
||||
# Return the simulation volume at requested
|
||||
state = res.ys[-1]
|
||||
final_field = cic_paint_dx(state[0], halo_size=halo_size)
|
||||
|
||||
return final_field, res.stats
|
||||
|
||||
def run():
|
||||
# Warm start
|
||||
times = []
|
||||
RangePush("warmup")
|
||||
final_field, stats, warmup_time = chrono_fun(simulate, 0.32, 0.8)
|
||||
RangePop()
|
||||
sync_global_devices("warmup")
|
||||
for i in range(iterations):
|
||||
RangePush(f"sim iter {i}")
|
||||
final_field, stats, sim_time = chrono_fun(simulate, 0.32, 0.8)
|
||||
RangePop()
|
||||
times.append(sim_time)
|
||||
return stats, warmup_time, times, final_field
|
||||
|
||||
if jax.device_count() > 1 and pdims:
|
||||
devices = mesh_utils.create_device_mesh(pdims)
|
||||
mesh = Mesh(devices.T, axis_names=('x', 'y'))
|
||||
with mesh:
|
||||
# Warm start
|
||||
stats, warmup_time, times, final_field = run()
|
||||
else:
|
||||
stats, warmup_time, times, final_field = run()
|
||||
|
||||
return stats, warmup_time, times, final_field
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description='JAX Cosmo Simulation Benchmark')
|
||||
parser.add_argument('-m',
|
||||
'--mesh_size',
|
||||
type=int,
|
||||
help='Mesh size',
|
||||
required=True)
|
||||
parser.add_argument('-b',
|
||||
'--box_size',
|
||||
type=float,
|
||||
help='Box size',
|
||||
required=True)
|
||||
parser.add_argument('-p',
|
||||
'--pdims',
|
||||
type=str,
|
||||
help='Processor dimensions',
|
||||
default=None)
|
||||
parser.add_argument('-h',
|
||||
'--halo_size',
|
||||
type=int,
|
||||
help='Halo size',
|
||||
required=True)
|
||||
parser.add_argument('-s',
|
||||
'--solver',
|
||||
type=str,
|
||||
help='Solver',
|
||||
choices=["Dopri5", "LeapfrogMidpoint"],
|
||||
required=True)
|
||||
parser.add_argument('-i',
|
||||
'--iterations',
|
||||
type=int,
|
||||
help='Number of iterations',
|
||||
default=10)
|
||||
parser.add_argument('-o',
|
||||
'--output_path',
|
||||
type=str,
|
||||
help='Output path',
|
||||
default=".")
|
||||
parser.add_argument('-f',
|
||||
'--save_fields',
|
||||
action='store_true',
|
||||
help='Save fields')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
mesh_size = args.mesh_size
|
||||
box_size = [args.box_size] * 3
|
||||
halo_size = args.halo_size
|
||||
solver_choice = args.solver
|
||||
iterations = args.iterations
|
||||
output_path = args.output_path
|
||||
os.makedirs(output_path, exist_ok=True)
|
||||
|
||||
if args.pdims:
|
||||
pdims = tuple(map(int, args.pdims.split("x")))
|
||||
else:
|
||||
pdims = None
|
||||
|
||||
mesh_shape = [mesh_size] * 3
|
||||
|
||||
stats, warmup_time, times, final_field = run_simulation(mesh_shape,
|
||||
box_size,
|
||||
halo_size,
|
||||
solver_choice,
|
||||
iterations,
|
||||
pdims=pdims)
|
||||
|
||||
# Save the final field
|
||||
if args.save_fields:
|
||||
nb_gpus = jax.device_count()
|
||||
field_folder = f"{output_path}/final_field/{nb_gpus}/{mesh_size}_{box_size[0]}/{solver_choice}/{halo_size}"
|
||||
os.makedirs(field_folder, exist_ok=True)
|
||||
np.save(f'{field_folder}/final_field_{rank}.npy',
|
||||
final_field.addressable_data(0))
|
||||
|
||||
# Write benchmark results to CSV
|
||||
# RANK SIZE MESHSIZE BOX HALO SOLVER NUM_STEPS JITTIME MIN MAX MEAN STD
|
||||
times = np.array(times)
|
||||
with open(f"{output_path}/jax_pm_benchmark.csv", 'a') as f:
|
||||
f.write(
|
||||
f"{rank},{size},{mesh_size},{box_size[0]},{halo_size},{solver_choice},{iterations},{warmup_time},{np.min(times)},{np.max(times)},{np.mean(times)},{np.std(times)}\n"
|
||||
)
|
||||
|
||||
print(f"Finished! Warmup time: {warmup_time:.4f} seconds")
|
||||
print(f"mean times: {np.mean(times):.4f}")
|
||||
print(f"Stats")
|
||||
Loading…
Add table
Reference in a new issue