diff --git a/src/underworld3/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index c42bbb47..f61b98de 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -7842,41 +7842,51 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): L = self.Unknowns.L Nc = dim - # Normalise the F0 Jacobian source to a flat list of dim scalar - # entries — F0 arrives as (1, dim) or (dim, 1) depending on how the - # template/bodyforce was written (Array indexing is strict). + # Normalise the F0 Jacobian source to a flat (dim,) Array — F0 arrives + # as (1, dim) or (dim, 1) depending on how the template/bodyforce was + # written (Array indexing is strict). _f0_flat = sympy.Array(F0_jac).reshape(dim) - f0_jac_list = [_f0_flat[c] for c in range(dim)] + + # Jacobian blocks are built with MATRIX-LEVEL diffs: one + # ``sympy.diff(, var)`` call per derivative + # variable instead of per-entry ``diff`` loops. Measured ~1.7x faster + # on large (monster-viscosity) fluxes — the batched elementwise diff + # reuses the derivative work across the block's shared subexpressions — + # and produces bit-identical entries (verified against the loop form). + # The flat PETSc [fc, gc, df, dg] layout is preserved by assigning each + # variable's matrix derivative into its slots below. # uu_G0[fc, gc] = dF0[fc] / dU[gc] G0 = sympy.zeros(Nc, Nc) - for fc in range(Nc): - for gc in range(Nc): - G0[fc, gc] = sympy.diff(f0_jac_list[fc], U_list[gc]) + for gc in range(Nc): + dF0_dU = sympy.diff(_f0_flat, U_list[gc]) + for fc in range(Nc): + G0[fc, gc] = dF0_dU[fc] # uu_G1[fc*Nc + gc, dg] = dF0[fc] / dL[gc, dg] G1 = sympy.zeros(Nc * Nc, dim) - for fc in range(Nc): - for gc in range(Nc): - for dg in range(dim): - G1[fc * Nc + gc, dg] = sympy.diff(f0_jac_list[fc], L[gc, dg]) + for gc in range(Nc): + for dg in range(dim): + dF0_dL = sympy.diff(_f0_flat, L[gc, dg]) + for fc in range(Nc): + G1[fc * Nc + gc, dg] = dF0_dL[fc] # uu_G2[fc*Nc + gc, df] = dF1[fc, df] / dU[gc] G2 = sympy.zeros(Nc * Nc, dim) - for fc in range(Nc): - for gc in range(Nc): + for gc in range(Nc): + dF1_dU = sympy.diff(F1_for_jac, U_list[gc]) + for fc in range(Nc): for df in range(dim): - G2[fc * Nc + gc, df] = sympy.diff(F1_for_jac[fc, df], U_list[gc]) + G2[fc * Nc + gc, df] = dF1_dU[fc, df] # uu_G3[fc*Nc + gc, df*dim + dg] = dF1[fc, df] / dL[gc, dg] G3 = sympy.zeros(Nc * Nc, dim * dim) - for fc in range(Nc): - for gc in range(Nc): - for df in range(dim): - for dg in range(dim): - G3[fc * Nc + gc, df * dim + dg] = sympy.diff( - F1_for_jac[fc, df], L[gc, dg] - ) + for gc in range(Nc): + for dg in range(dim): + dF1_dL = sympy.diff(F1_for_jac, L[gc, dg]) + for fc in range(Nc): + for df in range(dim): + G3[fc * Nc + gc, df * dim + dg] = dF1_dL[fc, df] self._uu_G0 = sympy.ImmutableMatrix(G0) self._uu_G1 = sympy.ImmutableMatrix(G1) @@ -7892,27 +7902,31 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): # up_G0[fc, 0] = dF0[fc] / dp G0 = sympy.zeros(dim, 1) + dF0_dp = sympy.diff(_f0_flat, p_scalar) for fc in range(dim): - G0[fc, 0] = sympy.diff(f0_jac_list[fc], p_scalar) + G0[fc, 0] = dF0_dp[fc] # up_G1[fc, dg] = dF0[fc] / d(dp/dx_dg) G1 = sympy.zeros(dim, dim) - for fc in range(dim): - for dg in range(dim): - G1[fc, dg] = sympy.diff(f0_jac_list[fc], Gp[0, dg]) + for dg in range(dim): + dF0_dGp = sympy.diff(_f0_flat, Gp[0, dg]) + for fc in range(dim): + G1[fc, dg] = dF0_dGp[fc] # up_G2[fc, df] = dF1[fc, df] / dp G2 = sympy.zeros(dim, dim) + dF1_dp = sympy.diff(F1_for_jac, p_scalar) for fc in range(dim): for df in range(dim): - G2[fc, df] = sympy.diff(F1_for_jac[fc, df], p_scalar) + G2[fc, df] = dF1_dp[fc, df] # up_G3[fc*dim + df, dg] = dF1[fc, df] / d(dp/dx_dg) G3 = sympy.zeros(dim * dim, dim) - for fc in range(dim): - for df in range(dim): - for dg in range(dim): - G3[fc * dim + df, dg] = sympy.diff(F1_for_jac[fc, df], Gp[0, dg]) + for dg in range(dim): + dF1_dGp = sympy.diff(F1_for_jac, Gp[0, dg]) + for fc in range(dim): + for df in range(dim): + G3[fc * dim + df, dg] = dF1_dGp[fc, df] self._up_G0 = sympy.ImmutableMatrix(G0) # zero in stokes tests self._up_G1 = sympy.ImmutableMatrix(G1) # zero in stokes tests @@ -8185,7 +8199,11 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): verbose=verbose, debug=debug, debug_name=debug_name, - cache=False, + # Disk cache + rank-0-only compile: under MPI only rank 0 invokes + # cc and publishes to the shared cache dir; the other ranks load + # the compiled module. Without this, every rank compiles its own + # copy of the (potentially enormous) pointwise module. + cache=True, ) self.compiled_extensions = _getext_result.ptrobj self.ext_dict = _getext_result.fn_dicts diff --git a/src/underworld3/utilities/_jitextension.py b/src/underworld3/utilities/_jitextension.py index 01e1a058..7d63ae18 100644 --- a/src/underworld3/utilities/_jitextension.py +++ b/src/underworld3/utilities/_jitextension.py @@ -680,6 +680,7 @@ def getext( module = _ext_dict[source_hash] else: from underworld3.utilities import _jit_cache as _jc + from mpi4py import MPI as _MPI # Disk lookup is cheap and rank-local — every rank checks # independently. UW assumes a shared filesystem for the cache dir. @@ -689,7 +690,23 @@ def getext( if module is not None and verbose and underworld3.mpi.rank == 0: print(f"JIT compiled module cached (disk) ... {source_hash}", flush=True) - if module is None: + # COLLECTIVE decision. The disk lookup above is rank-local — its own + # comment says so — and the branch below contains a Barrier, so the two + # must not be joined by a rank-local predicate: one rank finding the + # published module while another does not leaves the second waiting at + # a barrier nobody else enters. On the shared filesystems this path + # exists for, that divergence is ordinary (attribute caching, metadata + # lag, a write not yet visible), and it is what the barrier is meant to + # manage rather than something it can assume away. + # + # `needs_compile` is a global OR: if ANY rank lacks the module, every + # rank enters the branch and reaches the barrier. Ranks that already + # have it keep it and simply take part. + needs_compile = underworld3.mpi.comm.allreduce( + module is None, op=_MPI.LOR + ) + + if needs_compile: # Cold compile path. With MPI, only rank 0 invokes cc and # publishes; the other ranks barrier-wait and then load the # freshly-published .so from disk. This avoids the N× cc @@ -699,10 +716,17 @@ def getext( # cache is disabled (no way to share a build), in which # case the wasted cc cost is on the user's chosen path. multi_rank = underworld3.mpi.size > 1 - disk_enabled = cache and _jc.get_cache_dir() is not None + # Also collective, and for the same reason: a rank whose cache + # directory cannot be resolved would otherwise take the `else` + # branch — which has no barrier — while its peers wait in one. A + # global AND makes every rank fall back together, at the cost of + # each compiling locally, which is the safe direction. + disk_enabled = underworld3.mpi.comm.allreduce( + cache and _jc.get_cache_dir() is not None, op=_MPI.LAND + ) if multi_rank and disk_enabled: - if underworld3.mpi.rank == 0: + if underworld3.mpi.rank == 0 and module is None: if verbose: print( f"JIT compiling new module on rank 0 ... {source_hash}", @@ -719,7 +743,7 @@ def getext( # Synchronisation point: rank 0 has now published the # entry; other ranks may load it. underworld3.mpi.comm.Barrier() - if underworld3.mpi.rank != 0: + if module is None: module = _jc.load_module( source_hash, real_modname, constants_manifest ) @@ -735,17 +759,18 @@ def getext( print( f"JIT compiling new module ... {source_hash}", flush=True ) - module, tmpdir = compile_and_load( - real_modname, codeguys_final, verbose=verbose - ) - if verbose and underworld3.mpi.rank == 0: - # Tests in test_0004_pointwise_fns parse this exact prefix - # to find the per-call build directory. - print(f"Location of compiled module: {tmpdir}", flush=True) - if cache: - _jc.store_module( - source_hash, real_modname, tmpdir, constants_manifest + if module is None: + module, tmpdir = compile_and_load( + real_modname, codeguys_final, verbose=verbose ) + if verbose and underworld3.mpi.rank == 0: + # Tests in test_0004_pointwise_fns parse this exact + # prefix to find the per-call build directory. + print(f"Location of compiled module: {tmpdir}", flush=True) + if cache: + _jc.store_module( + source_hash, real_modname, tmpdir, constants_manifest + ) if cache: _ext_dict[source_hash] = module @@ -990,54 +1015,71 @@ def _basescalar_ccode(self, printer): # Save original for debugging fn_original = fn - # Three-phase lowering (issue #302 — must match prepare_for_cache_key): - # Phase 1: reveal constants nested inside other UWexpressions, so the - # substitution below can reach them. A top-level xreplace - # missed constants inside template-wrapped parameters and - # baked them as C literals while the manifest listed them. - fn = _reveal_constants(fn) - - # A truly-constant atom the manifest does NOT know about would be - # silently folded to a literal in phase 3 — the manifest and the - # C source must never disagree (issue #302). + # --- Gate the UW lowering (issue #302 pipeline) on the presence of + # UW-expression atoms. Plain-sympy components — the derivative + # blocks, which dominate the expression size — have no UW atoms, so + # the reveal / validate / xreplace / unwrap pipeline (≈5 full + # traversals per component) would be pure overhead: skip it entirely + # when there is nothing to lower. from underworld3.function.expressions import UWexpression as _UWexpr - if constants_subs_map is not None and hasattr(fn, 'atoms'): - unmanifested = [ - a.name for a in _stable_sorted(fn.atoms(sympy.Symbol)) - if isinstance(a, _UWexpr) - and _is_truly_constant(a, _UWexpr) - and a not in constants_subs_map - ] - if unmanifested: - raise RuntimeError( - f"JIT constants manifest is incomplete: constant expression(s) " - f"{unmanifested} appear in a kernel but have no constants[] " - f"slot — they would be baked into the C source (issue #302)." - ) + from underworld3.function.expressions import UWDerivativeExpression as _UWderiv + + _needs_lowering = ( + isinstance(fn, (_UWexpr, _UWderiv)) + # `has` is a bare traversal (no atom-set build) — the atoms() + # form built a set of every node, which cost seconds per 100k-node + # Jacobian component (measured ~20 s on a large collision model). + or (hasattr(fn, "has") and fn.has(_UWexpr)) + or not isinstance(fn, (sympy.MatrixBase, sympy.MatrixExpr)) + ) + if _needs_lowering: + # Phase 1: reveal constants nested inside other UWexpressions, so + # the substitution below can reach them. A top-level + # xreplace missed constants inside template-wrapped + # parameters and baked them as C literals while the + # manifest listed them. + fn = _reveal_constants(fn) + + # A truly-constant atom the manifest does NOT know about would be + # silently folded to a literal in phase 3 — the manifest and the + # C source must never disagree (issue #302). + if constants_subs_map is not None and hasattr(fn, 'atoms'): + unmanifested = [ + a.name for a in _stable_sorted(fn.atoms(sympy.Symbol)) + if isinstance(a, _UWexpr) + and _is_truly_constant(a, _UWexpr) + and a not in constants_subs_map + ] + if unmanifested: + raise RuntimeError( + f"JIT constants manifest is incomplete: constant expression(s) " + f"{unmanifested} appear in a kernel but have no constants[] " + f"slot — they would be baked into the C source (issue #302)." + ) - # Phase 2: Substitute constant UWexpressions with _JITConstant symbols - # These survive into C code as constants[i] - if constants_subs_map and fn is not None: - try: - fn = fn.xreplace(constants_subs_map) if hasattr(fn, 'xreplace') else fn - except Exception: - pass - - # Phase 3: Unwrap remaining non-constant UWexpressions to numerical values - fn = underworld3.function.expressions.unwrap(fn, keep_constants=False, return_self=False) - - # A manifested constant surviving to here bypassed its constants[] - # slot and is about to be baked — refuse rather than freeze the - # parameter silently (issue #302). - if constants_subs_map and hasattr(fn, 'atoms'): - baked = [a.name for a in _stable_sorted(fn.atoms(sympy.Symbol)) - if a in constants_subs_map] - if baked: - raise RuntimeError( - f"Manifested constant(s) {baked} were not routed through " - f"constants[] and would be baked into the C source " - f"(issue #302)." - ) + # Phase 2: Substitute constant UWexpressions with _JITConstant symbols + # These survive into C code as constants[i] + if constants_subs_map and fn is not None: + try: + fn = fn.xreplace(constants_subs_map) if hasattr(fn, 'xreplace') else fn + except Exception: + pass + + # Phase 3: Unwrap remaining non-constant UWexpressions to numerical values + fn = underworld3.function.expressions.unwrap(fn, keep_constants=False, return_self=False) + + # A manifested constant surviving to here bypassed its constants[] + # slot and is about to be baked — refuse rather than freeze the + # parameter silently (issue #302). + if constants_subs_map and hasattr(fn, 'atoms'): + baked = [a.name for a in _stable_sorted(fn.atoms(sympy.Symbol)) + if a in constants_subs_map] + if baked: + raise RuntimeError( + f"Manifested constant(s) {baked} were not routed through " + f"constants[] and would be baked into the C source " + f"(issue #302)." + ) if isinstance(fn, sympy.vector.Vector): fn = fn.to_matrix(mesh.N)[0 : mesh.dim, 0] @@ -1114,7 +1156,58 @@ def _basescalar_ccode(self, printer): print(f" - {sym} (type: {type(sym).__name__}, _ccodestr: {getattr(sym, '_ccodestr', 'N/A')})") out = sympy.MatrixSymbol("out", *fn.shape) - eqn = ("eqn_" + str(index), printer.doprint(fn, out)) + + # CSE before printing: shared subexpressions become ``double xN = ...;`` + # temps evaluated in dependency order, so the generated C — and hence + # the codegen time, gcc memory/time, and .so size — collapses on large + # expressions (measured: monster Jacobian output ~460k nodes -> ~30k). + # Semantics-preserving: temps are exact aliases of repeated + # subexpressions, so the generated kernel evaluates identical values. + # Opt in with UW_JIT_CSE=1 (default is off to preserve original behavior). + if os.environ.get("UW_JIT_CSE") in ("1", "true", "True", "yes", "YES"): + from sympy.simplify.cse_main import cse + from sympy.vector.scalar import BaseScalar + + _repl, _red = cse([fn]) + if _repl: + # cse may mint NEW coordinate instances (BaseScalar / + # UWCoordinate wrappers) that lack the mesh-set _ccodestr; + # recover it from their _id (same scheme as the + # COORDINATE SYMBOL RECOVERY above). + def _patch_coords(expr): + for _sym in set(expr.free_symbols): + _target = getattr(_sym, "_original_base_scalar", _sym) + if isinstance(_target, BaseScalar) and not hasattr( + _target, "_ccodestr" + ): + _idx = _target._id[0] + _sys = str(_target._id[1]) + _target._ccodestr = ( + f"petsc_n[{_idx}]" + if "Gamma" in _sys + else f"petsc_x[{_idx}]" + ) + + for _t_sym, _t_expr in _repl: + _patch_coords(_t_expr) + _patch_coords(_red[0]) + + _temp_code = "\n".join( + "double {} = {};".format( + printer.doprint(t_sym), printer.doprint(t_expr) + ) + for t_sym, t_expr in _repl + ) + _red_code = printer.doprint(_red[0], out) + if _red_code.startswith("// Not supported in C:"): + eqn = ("eqn_" + str(index), _red_code) + else: + eqn = ("eqn_" + str(index), _temp_code + "\n" + _red_code) + else: + eqn = ("eqn_" + str(index), printer.doprint(fn, out)) + else: + eqn = ("eqn_" + str(index), printer.doprint(fn, out)) + if eqn[1].startswith("// Not supported in C:"): spliteqn = eqn[1].split("\n") raise RuntimeError( @@ -1131,6 +1224,27 @@ def _basescalar_ccode(self, printer): MODNAME = "fn_ptr_ext_" + str(name) + # JIT compile flags for the generated kernels. Default keeps -O3 (kernel + # runtime speed) but adds -g0 to drop the debug info that the base Python + # CFLAGS injects via sysconfig -- pure overhead, and a memory hog on huge + # expressions, for these generated kernels. For very large expressions + # whose gcc -O3 compile is slow or OOM-killed, set UW3_JIT_CFLAGS to a + # lower optimisation level, e.g. UW3_JIT_CFLAGS="-O1 -g0". -std=c99 is + # always prepended (the generated code relies on it). + _default_jit_cflags = ["-O3", "-g0"] + _jit_cflags_env = os.environ.get("UW3_JIT_CFLAGS") + extra_compile_args = ( + ["-std=c99", *_jit_cflags_env.split()] + if _jit_cflags_env is not None + else ["-std=c99", *_default_jit_cflags] + ) + if verbose: + print( + f"JIT compile flags: {extra_compile_args}" + f"{' (from UW3_JIT_CFLAGS)' if _jit_cflags_env is not None else ' (default)'}", + flush=True, + ) + codeguys = [] # Create a `setup.py` setup_py_str = """ @@ -1148,7 +1262,7 @@ def _basescalar_ccode(self, printer): library_dirs={LIBDIRS}, runtime_library_dirs={LIBDIRS}, libraries={LIBFILES}, - extra_compile_args=['-std=c99','-O3'], + extra_compile_args={EXTRA_COMPILE_ARGS}, extra_link_args=[] )] setup(ext_modules=cythonize(ext_mods)) @@ -1157,6 +1271,7 @@ def _basescalar_ccode(self, printer): HEADERS=list(_stable_sorted(underworld3._incdirs.keys())), LIBDIRS=list(_stable_sorted(underworld3._libdirs.keys())), LIBFILES=list(_stable_sorted(underworld3._libfiles.keys())), + EXTRA_COMPILE_ARGS=extra_compile_args, ) codeguys.append(["setup.py", setup_py_str]) @@ -1199,7 +1314,6 @@ def _basescalar_ccode(self, printer): import string import random - import os if not "UW_JITNAME" in os.environ: randstr = "".join(random.choices(string.ascii_uppercase, k=5))