Skip to content

Stokes JIT: rank-0 compile + CSE before printing + configurable flags (#547) #612

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
bknight1 wants to merge 1 commit into
base: development
Choose a base branch
from
Open

Stokes JIT: rank-0 compile + CSE before printing + configurable flags (#547) #612

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
80 changes: 49 additions & 31 deletions src/underworld3/cython/petsc_generic_snes_solvers.pyx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(<whole block matrix>, 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)
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down
187 changes: 138 additions & 49 deletions src/underworld3/utilities/_jitextension.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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(
Expand All @@ -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,
)
Comment on lines +1216 to +1221

codeguys = []
# Create a `setup.py`
setup_py_str = """
Expand All @@ -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))
Expand All @@ -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])

Expand Down Expand Up @@ -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))
Expand Down
Loading