270 lines
11 KiB
Python
270 lines
11 KiB
Python
"""Performance baseline for minimal host-interference LBM stepping.
|
|
|
|
This script builds a temporary config, runs warmup + measured batches, and
|
|
reports MLUPS under a FluidX3D-like benchmark mindset:
|
|
|
|
- keep the main loop on GPU (`stepper.step`)
|
|
- avoid host downloads by default
|
|
- selectively enable host-touch paths to quantify overhead
|
|
- sweep `inlet.scheme` to check stability-sensitive combinations
|
|
|
|
Usage::
|
|
python tests/run_perf_baseline.py
|
|
python tests/run_perf_baseline.py --lattice-model D2Q9 --nx 384 --ny 192 --steps 30000
|
|
python tests/run_perf_baseline.py --macro-every 500 --ddf-every 1000
|
|
python tests/run_perf_baseline.py --with-cylinder --obs-every 50
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import time
|
|
from typing import Any, Dict, List
|
|
|
|
import pycuda.driver as cuda
|
|
|
|
_REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
|
|
_DEFAULT_LBM = os.path.join(_REPO, "src", "CelerisLab", "configs", "config_lbm.json")
|
|
|
|
|
|
def _load_json(path: str) -> dict:
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
|
|
|
|
def _write_json(path: str, payload: dict) -> None:
|
|
with open(path, "w", encoding="utf-8") as f:
|
|
json.dump(payload, f, indent=2)
|
|
|
|
|
|
def _build_lbm_cfg(base: dict, args: argparse.Namespace) -> dict:
|
|
cfg = json.loads(json.dumps(base))
|
|
cfg["grid"]["lattice_model"] = args.lattice_model
|
|
cfg["grid"]["nx"] = int(args.nx)
|
|
cfg["grid"]["ny"] = int(args.ny)
|
|
cfg["grid"]["nz"] = int(args.nz)
|
|
|
|
cfg["physics"]["viscosity"] = float(args.viscosity)
|
|
cfg["physics"]["velocity"] = float(args.velocity)
|
|
cfg["physics"]["rho"] = float(args.rho)
|
|
|
|
cfg["method"]["collision"] = str(args.collision).upper()
|
|
cfg["method"]["streaming"] = str(args.streaming)
|
|
cfg["method"]["store_precision"] = str(args.store_precision).upper()
|
|
cfg["method"]["ddf_shifting"] = bool(args.ddf_shifting)
|
|
cfg["method"]["les"]["enabled"] = bool(args.enable_les)
|
|
cfg["method"]["outlet"]["mode"] = str(args.outlet_mode)
|
|
cfg["method"]["inlet"]["profile"] = str(args.inlet_profile)
|
|
# Expose inlet scheme as a benchmark axis; useful when stability depends
|
|
# on collision/inlet coupling.
|
|
cfg["method"]["inlet"]["scheme"] = str(args.inlet_scheme)
|
|
cfg["method"]["y_wall_bc"] = str(args.y_wall_bc)
|
|
|
|
cfg["cuda"]["threads_per_block"] = int(args.threads_per_block)
|
|
cfg["cuda"]["compute_capability"] = str(args.compute_capability)
|
|
return cfg
|
|
|
|
|
|
def _build_body_cfg(args: argparse.Namespace) -> dict:
|
|
if not args.with_cylinder:
|
|
return {"objects": []}
|
|
cx = 0.5 * float(args.nx)
|
|
cy = 0.5 * float(args.ny)
|
|
radius = max(2.0, min(float(args.nx), float(args.ny)) * 0.08)
|
|
return {
|
|
"objects": [
|
|
{
|
|
"type": "cylinder",
|
|
"center": [cx, cy],
|
|
"radius": radius,
|
|
}
|
|
]
|
|
}
|
|
|
|
|
|
def _maybe_probe_macroscopic(sim: Any, step: int, every: int, repeat: int) -> None:
|
|
if every > 0 and step % every == 0:
|
|
for _ in range(max(1, int(repeat))):
|
|
_ = sim.get_macroscopic()
|
|
|
|
|
|
def _maybe_probe_ddf(sim: Any, step: int, every: int, repeat: int) -> None:
|
|
if every > 0 and step % every == 0:
|
|
for _ in range(max(1, int(repeat))):
|
|
_ = sim.get_ddf()
|
|
|
|
|
|
def _maybe_checkpoint(sim: Any, step: int, every: int, out_dir: str) -> None:
|
|
if every > 0 and step % every == 0:
|
|
path = os.path.join(out_dir, f"checkpoint_step_{step:09d}.h5")
|
|
sim.save_checkpoint(path)
|
|
|
|
|
|
def _maybe_probe_obs(sim: Any, stream: cuda.Stream, step: int, every: int) -> None:
|
|
if every <= 0 or step % every != 0:
|
|
return
|
|
if sim.bodies.count <= 0:
|
|
return
|
|
sim.bodies.download_obs_full_async(stream)
|
|
stream.synchronize()
|
|
_ = sim.bodies.read_force(0)
|
|
|
|
|
|
def run(args: argparse.Namespace) -> Dict[str, Any]:
|
|
if not os.path.isfile(_DEFAULT_LBM):
|
|
raise FileNotFoundError(f"Base config missing: {_DEFAULT_LBM}")
|
|
|
|
base = _load_json(_DEFAULT_LBM)
|
|
lbm_cfg = _build_lbm_cfg(base, args)
|
|
body_cfg = _build_body_cfg(args)
|
|
|
|
tmpd = tempfile.mkdtemp(prefix="celeris_perf_baseline_")
|
|
lbm_path = os.path.join(tmpd, "config_lbm.json")
|
|
body_path = os.path.join(tmpd, "config_body.json")
|
|
ckpt_dir = os.path.join(tmpd, "checkpoints")
|
|
os.makedirs(ckpt_dir, exist_ok=True)
|
|
_write_json(lbm_path, lbm_cfg)
|
|
_write_json(body_path, body_cfg)
|
|
|
|
from CelerisLab import Simulation # noqa: WPS433
|
|
|
|
sim = Simulation(lbm_config_path=lbm_path, body_config_path=body_path, device_id=args.device_id)
|
|
sim.initialize()
|
|
|
|
stream = cuda.Stream()
|
|
total_cells = int(args.nx) * int(args.ny) * int(args.nz)
|
|
|
|
# Warmup before measurement window.
|
|
warmup_done = 0
|
|
while warmup_done < int(args.warmup_steps):
|
|
chunk = min(int(args.batch_steps), int(args.warmup_steps) - warmup_done)
|
|
sim.stepper.step(chunk, action_gpu=sim.bodies.action_gpu, obs_gpu=sim.bodies.obs_gpu, stream=stream)
|
|
warmup_done += chunk
|
|
stream.synchronize()
|
|
|
|
measured_batch_s: List[float] = []
|
|
measured_steps = int(args.steps)
|
|
done = 0
|
|
t0 = time.perf_counter()
|
|
while done < measured_steps:
|
|
chunk = min(int(args.batch_steps), measured_steps - done)
|
|
step_start = time.perf_counter()
|
|
sim.stepper.step(chunk, action_gpu=sim.bodies.action_gpu, obs_gpu=sim.bodies.obs_gpu, stream=stream)
|
|
stream.synchronize()
|
|
step_end = time.perf_counter()
|
|
measured_batch_s.append(step_end - step_start)
|
|
|
|
done += chunk
|
|
global_step = sim.stepper.step_count
|
|
_maybe_probe_macroscopic(sim, global_step, int(args.macro_every), int(args.macro_repeat))
|
|
_maybe_probe_ddf(sim, global_step, int(args.ddf_every), int(args.ddf_repeat))
|
|
_maybe_checkpoint(sim, global_step, int(args.checkpoint_every), ckpt_dir)
|
|
_maybe_probe_obs(sim, stream, global_step, int(args.obs_every))
|
|
|
|
stream.synchronize()
|
|
elapsed_s = time.perf_counter() - t0
|
|
|
|
# Optional final sanity readback (outside core timing path by default).
|
|
if args.final_macro_snapshot:
|
|
_ = sim.get_macroscopic()
|
|
|
|
sim.close()
|
|
|
|
mlups = (total_cells * measured_steps) / max(elapsed_s, 1e-12) / 1.0e6
|
|
batch_ms = [1000.0 * x for x in measured_batch_s]
|
|
batch_ms_sorted = sorted(batch_ms)
|
|
p50 = batch_ms_sorted[len(batch_ms_sorted) // 2] if batch_ms_sorted else 0.0
|
|
p90 = batch_ms_sorted[min(len(batch_ms_sorted) - 1, int(0.9 * (len(batch_ms_sorted) - 1)))] if batch_ms_sorted else 0.0
|
|
|
|
return {
|
|
"benchmark": "celerislab_stepper_baseline",
|
|
"device_id": int(args.device_id),
|
|
"lattice_model": args.lattice_model,
|
|
"grid": {"nx": int(args.nx), "ny": int(args.ny), "nz": int(args.nz)},
|
|
"collision": str(args.collision).upper(),
|
|
"inlet_scheme": str(args.inlet_scheme),
|
|
"streaming": str(args.streaming),
|
|
"store_precision": str(args.store_precision).upper(),
|
|
"steps": measured_steps,
|
|
"warmup_steps": int(args.warmup_steps),
|
|
"batch_steps": int(args.batch_steps),
|
|
"mlups": float(mlups),
|
|
"elapsed_s": float(elapsed_s),
|
|
"batch_ms_p50": float(p50),
|
|
"batch_ms_p90": float(p90),
|
|
"overhead_switches": {
|
|
"macro_every": int(args.macro_every),
|
|
"macro_repeat": int(args.macro_repeat),
|
|
"ddf_every": int(args.ddf_every),
|
|
"ddf_repeat": int(args.ddf_repeat),
|
|
"checkpoint_every": int(args.checkpoint_every),
|
|
"obs_every": int(args.obs_every),
|
|
"with_cylinder": bool(args.with_cylinder),
|
|
"final_macro_snapshot": bool(args.final_macro_snapshot),
|
|
},
|
|
}
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
ap = argparse.ArgumentParser(description="CelerisLab pure-step performance baseline")
|
|
ap.add_argument("--device-id", type=int, default=0)
|
|
ap.add_argument("--lattice-model", choices=("D2Q9", "D3Q19"), default="D3Q19")
|
|
ap.add_argument("--nx", type=int, default=256)
|
|
ap.add_argument("--ny", type=int, default=256)
|
|
ap.add_argument("--nz", type=int, default=256)
|
|
ap.add_argument("--steps", type=int, default=3000)
|
|
ap.add_argument("--warmup-steps", type=int, default=400)
|
|
ap.add_argument("--batch-steps", type=int, default=100)
|
|
|
|
ap.add_argument("--collision", choices=("SRT", "TRT", "MRT"), default="SRT")
|
|
ap.add_argument("--streaming", choices=("double_buffer", "esopull"), default="double_buffer")
|
|
ap.add_argument("--store-precision", choices=("FP32", "FP16S"), default="FP32")
|
|
ap.add_argument("--ddf-shifting", action="store_true")
|
|
ap.add_argument("--enable-les", action="store_true")
|
|
ap.add_argument("--outlet-mode", choices=("neq_extrap", "zero_gradient", "blended"), default="neq_extrap")
|
|
ap.add_argument("--inlet-profile", choices=("uniform", "parabolic"), default="uniform")
|
|
ap.add_argument(
|
|
"--inlet-scheme",
|
|
choices=("zou_he_local", "channel_stabilized", "equilibrium", "regularized"),
|
|
default="zou_he_local",
|
|
)
|
|
ap.add_argument("--y-wall-bc", choices=("bounce_back", "free_slip"), default="bounce_back")
|
|
|
|
ap.add_argument("--viscosity", type=float, default=0.0035)
|
|
ap.add_argument("--velocity", type=float, default=0.03)
|
|
ap.add_argument("--rho", type=float, default=1.0)
|
|
ap.add_argument("--threads-per-block", type=int, default=256)
|
|
ap.add_argument("--compute-capability", type=str, default="auto")
|
|
|
|
# Overhead attribution toggles.
|
|
ap.add_argument("--macro-every", type=int, default=0, help="Call get_macroscopic() every N steps (0=off)")
|
|
ap.add_argument("--macro-repeat", type=int, default=1, help="Repeat get_macroscopic() calls per probe step")
|
|
ap.add_argument("--ddf-every", type=int, default=0, help="Call get_ddf() every N steps (0=off)")
|
|
ap.add_argument("--ddf-repeat", type=int, default=1, help="Repeat get_ddf() calls per probe step")
|
|
ap.add_argument("--checkpoint-every", type=int, default=0, help="Save checkpoint every N steps (0=off)")
|
|
ap.add_argument("--obs-every", type=int, default=0, help="Download object obs every N steps (0=off)")
|
|
ap.add_argument("--with-cylinder", action="store_true", help="Inject one cylinder object (needed for obs probes)")
|
|
ap.add_argument("--final-macro-snapshot", action="store_true")
|
|
|
|
ap.add_argument("--json-out", type=str, default="", help="Optional path to save metrics JSON")
|
|
return ap.parse_args()
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
result = run(args)
|
|
print(json.dumps(result, indent=2))
|
|
if args.json_out.strip():
|
|
out = os.path.abspath(args.json_out)
|
|
os.makedirs(os.path.dirname(out), exist_ok=True)
|
|
_write_json(out, result)
|
|
print(f"Wrote: {out}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|