From 8978ad8bd25b3fc7f47ba96d90cd6af02db7ace3 Mon Sep 17 00:00:00 2001 From: Ben Knight <55677727+bknight1@users.noreply.github.com> Date: Fri, 14 Aug 2026 10:02:28 +0800 Subject: [PATCH 1/3] Stokes JIT: rank-0 compile + CSE before printing + configurable flags (#547) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four fixes for the HPC OOM in the runtime JIT compilation of Stokes pointwise functions (gcc: fatal error: Killed signal terminated program cc1 under mpirun -np N): 1. SNES_Stokes_SaddlePt._setup_pointwise_functions now calls getext with cache=True (the only call site that opted out), activating the existing rank-0-only compile + cross-rank disk handoff. Previously every rank compiled its own copy of the (potentially enormous) pointwise module, exhausting node memory (20x -> 1x concurrent compilers). 2. CSE before printing in generate_c_source: shared subexpressions become double xN = ...; temps evaluated in dependency order, shrinking the generated C ~63x on large rheologies (1.45 MB -> 23 KB header) and gcc peak RSS from 1.32 GB to 62-74 MB — comfortably inside a 900 MB per-rank cap. Semantics-preserving (temps are exact aliases); recovers _ccodestr on the new coordinate instances cse can mint; UW_JIT_NOCSE=1 escapes. 3. JIT compile flags configurable via UW3_JIT_CFLAGS (e.g. "-O1 -g0"); default adds -g0 to drop the debug info sysconfig injects. The env var is the memory lever for constrained HPC nodes. 4. Matrix-level sympy.diff for the Stokes uu/up Jacobian blocks instead of per-entry loops (~1.7x faster derivatives, bit-identical entries, flat PETSc [fc,gc,df,dg] layout preserved). Verified: 132 tests pass (incl. the FD-oracle Jacobian layout test, JIT determinism, constants routing) plus the CSE-off path; MPI rank-0-only compile and disk-cache reuse confirmed under mpirun -np 2. Underworld development team with AI support from Claude Code --- .../cython/petsc_generic_snes_solvers.pyx | 80 +++++--- src/underworld3/utilities/_jitextension.py | 187 +++++++++++++----- 2 files changed, 187 insertions(+), 80 deletions(-) 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..2b388ac0 100644 --- a/src/underworld3/utilities/_jitextension.py +++ b/src/underworld3/utilities/_jitextension.py @@ -990,54 +990,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 +1131,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 out with UW_JIT_NOCSE=1 if a pathological case regresses. + if os.environ.get("UW_JIT_NOCSE") not in ("1", "true", "True"): + 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 +1199,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 +1237,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 +1246,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 +1289,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)) From 9b8b3cd7edc8b549a8de21825fed59674664aaef Mon Sep 17 00:00:00 2001 From: lmoresi Date: Wed, 19 Aug 2026 20:00:45 +1000 Subject: [PATCH 2/3] JIT cache gate: decide collectively whether a compile is needed MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The disk-cache lookup is rank-local — its own comment says so — and the Barrier that follows sits inside `if module is None:`. Joining the two by a rank-local predicate means a rank that finds the published module skips the barrier while a rank that does not enters it alone and waits forever. That divergence is ordinary on the shared filesystems this path exists for: attribute caching, metadata lag, or a write not yet visible to every rank just after rank 0 publishes. It is what the barrier is meant to manage rather than something it can assume away. A second route is `disk_enabled`, which is `cache and get_cache_dir() is not None` — a rank whose cache directory cannot be resolved takes the else branch, which has no barrier, while its peers wait in one. Both predicates are now reduced. `needs_compile` is a global OR, so if any rank lacks the module every rank enters the branch and reaches the barrier; `disk_enabled` is a global AND, so a rank without a cache directory makes all ranks fall back together to local compiles, which is the safe direction. Ranks that already hold the module keep it: the rank-0 compile, the post-barrier load, and the no-disk fallback are each guarded on `module is None`. Demonstrated rather than argued. Monkeypatching the disk lookup so rank 0 sees the cache and the others do not — the shape NFS or Lustre lag produces — deadlocks this PR as it stands: at np=4 all four ranks enter the solve, none returns, mpirun times out at rc=241. With the reduction, rc=0 on four consecutive runs. This is not a defect introduced here: the gate predates the PR. What the PR changes is the exposure. Stokes was the one call site passing cache=False, so it always took the no-barrier branch; switching it to cache=True is the point of the change, and it puts the main solver path onto the branch with the conditional collective, on exactly the filesystems #547 is about. Full ./uw test 1556 passed; the four JIT test files 14 passed. Underworld development team with AI support from Claude Code --- src/underworld3/utilities/_jitextension.py | 53 ++++++++++++++++------ 1 file changed, 39 insertions(+), 14 deletions(-) diff --git a/src/underworld3/utilities/_jitextension.py b/src/underworld3/utilities/_jitextension.py index 2b388ac0..d38af739 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 From 630ed8d23fa586534f0f340df54a12b927d5b4da Mon Sep 17 00:00:00 2001 From: Ben Knight <55677727+bknight1@users.noreply.github.com> Date: Fri, 21 Aug 2026 08:27:20 +0800 Subject: [PATCH 3/3] JIT: make SymPy CSE opt-in via UW_JIT_CSE=1 (default off) Change the CSE preprocessing step in generate_c_source from opt-out (UW_JIT_NOCSE=1) to opt-in (UW_JIT_CSE=1), so that default code generation behavior and output are preserved verbatim. Users on memory-constrained HPC nodes running very large/complex rheologies can set UW_JIT_CSE=1 to activate CSE header shrinking. Underworld development team with AI support from Claude Code --- src/underworld3/utilities/_jitextension.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/underworld3/utilities/_jitextension.py b/src/underworld3/utilities/_jitextension.py index d38af739..7d63ae18 100644 --- a/src/underworld3/utilities/_jitextension.py +++ b/src/underworld3/utilities/_jitextension.py @@ -1163,8 +1163,8 @@ def _basescalar_ccode(self, printer): # 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 out with UW_JIT_NOCSE=1 if a pathological case regresses. - if os.environ.get("UW_JIT_NOCSE") not in ("1", "true", "True"): + # 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