check_decomposition.py

Go to the documentation of this file.

#!/usr/bin/env python
"""Inspect how a grid + rank count decomposes, and how deep the pressure multigrid gets on it.
 
Answers, without spending any GPU time, the questions that decide distributed multigrid performance:
 
  * how balanced is the ORB partition (max block / min block)?
  * is a direction split that must not be (e.g. the wall-normal one in a channel)?
  * how many multigrid levels does the hierarchy ACTUALLY reach, and where does each axis stop?
  * does the coarse-first decomposition help here, and at what depth?
 
Run it before committing a grid or a rank count to a large job. It needs only an MPI build of
`flow` (the OpenMP one is fine, no GPU), and `mpirun --oversubscribe` happily runs more ranks than
there are cores because nothing here is timed.
 
    # one grid, several rank counts, both decomposition modes:
    python check_decomposition.py --grid 480,80,160 --levels 6 --np 1,2,4,8,12,16,24 --mode 0,coarse
 
    # a whole weak-scaling ladder, ORB only (no solver allocation -> any size, instantly):
    python check_decomposition.py --grid 960,160,320 --np 1 --orb-only
    python check_decomposition.py --grid 3072,512,1024 --np 32 --orb-only
 
    # what the hierarchy looks like level by level:
    python check_decomposition.py --grid 1508,240,503 --levels 5 --np 4 --verbose
 
    # the ladder after diagnostics.rebalance_by_weights(w), and the alignment it picks:
    python check_decomposition.py --grid 96,96,96 --levels 8 --np 4,8 --predict --weights w.npy
 
`--orb-only` skips constructing a Solver, so it reports the partition (balance, splits) for grids far
too large to allocate on a host — but it cannot report the achieved level count, which needs the real
MG init. Without it, the grid must fit in host memory (a few hundred M cells at most).
 
Requires PYTHONPATH to point at an MPI-enabled flow build, e.g.
    PYTHONPATH=$PWD/build_mpi_omp python scripts/check_decomposition.py ...
"""
import argparse
import os
import re
import subprocess
import sys
 
HERE = os.path.abspath(__file__)
 
 
def halvings(d):
    """How many times an axis can be halved: the multigrid's per-axis depth budget."""
    n = 0
    while d % 2 == 0 and d // 2 >= 2:
        d //= 2
        n += 1
    return n
 
 
# ---- child: runs under mpirun, prints one RESULT line ------------------------------------------
def child(args):
    from mpi4py import MPI
 
    w = MPI.COMM_WORLD
    import numpy as np
 
    from peclet import flow
 
    gx, gy, gz = args.grid
    origin, size = flow.mpi_block(gx, gy, gz, levels=args.decomp_levels)
    lnx, lny, lnz = size
    cells = lnx * lny * lnz
    hi = w.allreduce(cells, op=MPI.MAX)
    lo = w.allreduce(cells, op=MPI.MIN)
    split = [w.allreduce(1 if size[k] < (gx, gy, gz)[k] else 0, op=MPI.SUM) for k in range(3)]
 
    if not args.orb_only:
        # Constructing the Solver is what builds the MG hierarchy; with PECLET_FLOW_MG_DEBUG=1 it
        # prints the level table, which the parent parses for the achieved depth.
        s = flow.Solver(lnx, lny, lnz)
        s.set_decomposition(args.decomp_levels)  # must match the mpi_block() call above
        s.init_mpi(gx, gy, gz)
        s.set_rho(1.0)
        s.set_mu(0.1)
        s.set_dt(0.02)
        s.set_pressure_multigrid(True, args.levels)
        s.set_pressure_pcg(True, 200, 1e-6)
        if args.walls:
            s.set_domain_bc("-y", "wall")
            s.set_domain_bc("+y", "wall")
        s.set_pressure_geometry(np.asfortranarray(np.full((lnx, lny, lnz), 1e30)))
 
    if w.rank == 0:
        print(
            f"RESULT np={w.size} block={lnx}x{lny}x{lnz} imbalance={hi / lo:.3f} "
            f"split=({split[0]},{split[1]},{split[2]})",
            flush=True,
        )
    MPI.Finalize()
    os._exit(0)
 
 
# ---- parent: spawn one mpirun per (mode, np) and tabulate ---------------------------------------
MODES = {"0": ("aligned", "0"), "aligned": ("aligned", "0"), "legacy": ("aligned", "0")}
 
 
def mode_spec(tok, levels):
    """'0'/'aligned' -> the legacy aligned ORB; 'coarse'/'coarse-first'/<int> -> coarse-first."""
    if tok in MODES:
        return MODES[tok]
    if tok in ("coarse", "coarse-first", "cf"):
        return ("coarse-first", str(levels))
    return (f"coarse-first({tok})", tok)
 
 
def predict(args):
    """Pure-function pre-flight: the hierarchy init_mpi would build, per rank count, with and
    without telescoping. No mpirun -- this is what makes np=1536 answerable on a laptop."""
    from peclet import flow
 
    gx, gy, gz = args.grid
    for np_ in args.np:
        for tele in (False, True):
            if args.weights is not None:
                # The hierarchy after diagnostics.rebalance_by_weights(w): its aligned weighted
                # ORB replaces the set_decomposition partition, and the call returns the alignment.
                rows, align = flow.predict_hierarchy(gx, gy, gz, np_, args.levels, tele,
                                                     weights=args.weights)
                if not tele:  # the partition, hence the alignment, is the same either way
                    print(f"\n{gx}x{gy}x{gz} np={np_}: rebalance_by_weights -> aligned weighted "
                          f"ORB, align {align} (a = {align.bit_length() - 1})")
            else:
                # --decomp-levels must reach predict_hierarchy: it is a pure function with no
                # process state to read since flow ad917b1, so omitting it predicts the ALIGNED
                # hierarchy.
                rows = flow.predict_hierarchy(gx, gy, gz, np_, args.levels, tele,
                                              decomposition_levels=args.decomp_levels)
            cg = rows[-1][0]
            print(f"\n{gx}x{gy}x{gz} np={np_} levels={args.levels} telescope={'ON' if tele else 'off'}"
                  f" -> {len(rows)} levels, coarsest {cg[0]}x{cg[1]}x{cg[2]}, "
                  f"max extent {max(cg)}, bottom on {rows[-1][1]} rank(s)")
            for L, (g, ranks, blk, rat, t) in enumerate(rows):
                print(f"   L{L} global {g[0]:5d}x{g[1]:5d}x{g[2]:5d}  ranks {ranks:5d}  block0 "
                      f"{blk[0]:4d}x{blk[1]:4d}x{blk[2]:4d}  ratio({rat[0]},{rat[1]},{rat[2]})"
                      f"{'  -> TELESCOPE' if t else ''}")
 
 
def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--grid", required=True, help="GNX,GNY,GNZ")
    ap.add_argument("--np", default="1", help="comma-separated rank counts")
    ap.add_argument("--levels", type=int, default=5, help="multigrid levels requested (default 5)")
    ap.add_argument("--mode", default="0", help="comma-separated: 0/aligned, coarse, or a depth")
    ap.add_argument("--walls", action="store_true", help="no-slip -y/+y (the channel's domain-BC MG path)")
    ap.add_argument("--orb-only", action="store_true", help="partition only; no Solver, so any size")
    ap.add_argument("--verbose", action="store_true", help="also print the per-level dims/ratios")
    ap.add_argument("--mpirun", default="mpirun")
    ap.add_argument("--predict", action="store_true",
                    help="use flow.predict_hierarchy (pure function, no mpirun, any np): print the "
                         "level ladder with and without coarse-level telescoping")
    ap.add_argument("--weights", default=None,
                    help="with --predict: a .npy of per-cell weights (what "
                         "diagnostics.rebalance_by_weights takes: flat x-fastest, or an (nx,ny,nz) "
                         "array) -- predict the hierarchy AFTER that rebalance and log the "
                         "alignment its aligned weighted ORB picks")
    ap.add_argument("--decomp-levels", type=int, default=0,
                    help="Solver.set_decomposition(levels) / flow.mpi_block(levels=): 0 = the "
                         "aligned ORB, >= 2 = coarse-first with that depth")
    ap.add_argument("--child", action="store_true", help=argparse.SUPPRESS)
    args = ap.parse_args()
    if args.weights is not None and not args.predict:
        ap.error("--weights needs --predict")
    if getattr(args, "predict", False):
        if isinstance(args.grid, str):
            args.grid = [int(v) for v in args.grid.split(",")]
        if isinstance(args.np, str):
            args.np = [int(v) for v in args.np.split(",")]
        if args.weights is not None:
            import numpy as np
 
            w = np.load(args.weights)
            if w.ndim == 3:
                if tuple(w.shape) != tuple(args.grid):
                    sys.exit(f"--weights: shape {w.shape} is not the grid {tuple(args.grid)}")
                w = w.ravel(order="F")  # x-fastest, as rebalance_by_weights takes it
            args.weights = np.asarray(w, dtype=np.float64).ravel().tolist()
            # The weighted ORB replaces the set_decomposition partition, so --mode does not apply.
            print("=== weighted: diagnostics.rebalance_by_weights(w) (--mode / --decomp-levels "
                  "ignored) ===")
            predict(args)
            return
        # In --predict, --mode selects the depth (it can list several). --decomp-levels is the
        # run-mode flag; honour it when the user gave it and left --mode at its default, rather
        # than silently overwriting it below.
        if args.decomp_levels and str(args.mode) == "0":
            args.mode = str(args.decomp_levels)
        for spec in (args.mode if isinstance(args.mode, list) else str(args.mode).split(",")):
            label, lv = mode_spec(spec, args.levels)
            args.decomp_levels = int(lv)
            print(f"=== decomposition mode {label} (set_decomposition({int(lv)})) ===")
            predict(args)
        return
    args.grid = [int(v) for v in args.grid.split(",")]
    if args.child:
        return child(args)
 
    gx, gy, gz = args.grid
    print(f"grid {gx}x{gy}x{gz} = {gx * gy * gz / 1e6:.1f} Mcells   "
          f"halvings x/y/z = {halvings(gx)}/{halvings(gy)}/{halvings(gz)}"
          f"{'   <-- an axis with 0 halvings NEVER coarsens' if 0 in (halvings(gx), halvings(gy), halvings(gz)) else ''}")
    print(f"multigrid levels requested: {args.levels}\n")
    print(f"{'mode':<20} {'np':>4} {'block':>18} {'imbalance':>10} {'split x,y,z':>12} {'levels':>7}")
 
    for tok in args.mode.split(","):
        label, decomp = mode_spec(tok.strip(), args.levels)
        for np_ in [int(v) for v in args.np.split(",")]:
            env = dict(os.environ, PECLET_FLOW_MG_DEBUG="1")
            cmd = [args.mpirun, "--oversubscribe", "-np", str(np_), sys.executable, HERE, "--child",
                   "--grid", ",".join(map(str, args.grid)), "--levels", str(args.levels),
                   "--decomp-levels", str(decomp)]
            if args.walls:
                cmd.append("--walls")
            if args.orb_only:
                cmd.append("--orb-only")
            try:
                out = subprocess.run(cmd, env=env, capture_output=True, text=True, timeout=900).stdout
            except subprocess.TimeoutExpired:
                print(f"{label:<20} {np_:>4}   TIMED OUT")
                continue
            res = re.search(r"RESULT np=(\d+) block=(\S+) imbalance=(\S+) split=\‍((\S+)\‍)", out)
            if not res:
                print(f"{label:<20} {np_:>4}   FAILED (rerun by hand for the error)")
                continue
            lev = re.search(r"-> (\d+) levels", out)
            cf = re.search(r"coarse-first depth (\d+) \‍(refine (\S+),", out)
            note = f"  [coarse-first depth {cf.group(1)}, refine {cf.group(2)}]" if cf else ""
            print(f"{label:<20} {res.group(1):>4} {res.group(2):>18} {res.group(3):>10} "
                  f"{res.group(4):>12} {lev.group(1) if lev else '-':>7}{note}")
            if args.verbose:
                for line in out.splitlines():
                    if line.startswith("[mg]  L"):
                        print("    " + line)
    print("\nreminders: an ODD axis never coarsens (measured 3.2x slower on one GPU);")
    print("           in MPI a level coarsens an axis only if EVERY rank's block is even on it,")
    print("           so the achievable depth is set by the per-rank block, not the global grid.")
 
 
if __name__ == "__main__":
    main()