From 9a80842018deadeceb3f41754e37e409eba2df9d Mon Sep 17 00:00:00 2001 From: Louis Moresi Date: Wed, 25 Mar 2026 16:14:31 +1100 Subject: [PATCH 1/2] Introduce JITCallbackSet for structured JIT cache keys Refactor the JIT compilation pipeline to use a JITCallbackSet dataclass that groups the five PETSc callback lists (residual, bcs, jacobian, bd_residual, bd_jacobian) into a single structured container. This addresses the root cause of the cache-collision bug (PR #92) at an architectural level: the flat tuple hash that lost callback role information is replaced by a structured signature that preserves which slot each expression belongs to. Changes: - Add JITCallbackSet dataclass with flat(), signature(), map(), counts - Extract _structural_expand() as a module-level function (was inline) - Refactor getext() to accept JITCallbackSet (with backward compat) - Refactor _createext() to accept JITCallbackSet - Update all 6 call sites: 3 solvers (Scalar, Vector, Stokes) and 3 integrals (Integral, Integral._evaluate_integral, BdIntegral) - Include PR #92 regression tests (spherical shell cache collision) Incorporates the fix from PR #92 (gthyagi) which identified the bug and added the regression tests. Test results: 374 passed, 7 skipped, 1 xfailed (level_1 suite) Underworld development team with AI support from Claude Code --- .../cython/petsc_generic_snes_solvers.pyx | 72 ++-- src/underworld3/cython/petsc_maths.pyx | 12 +- src/underworld3/utilities/_jitextension.py | 318 +++++++++--------- tests/test_0004_pointwise_fns.py | 38 ++- 4 files changed, 232 insertions(+), 208 deletions(-) diff --git a/src/underworld3/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index 34cd3a328..e830fdc42 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -9,7 +9,7 @@ from petsc4py import PETSc import underworld3 import underworld3 as uw -from underworld3.utilities._jitextension import getext +from underworld3.utilities._jitextension import getext, JITCallbackSet import underworld3.timing as timing from underworld3.utilities._api_tools import uw_object @@ -1542,15 +1542,19 @@ class SNES_Scalar(SolverBaseClass): print(f"Scalar SNES: Jacobians complete, now compile", flush=True) prim_field_list = [self.u] - _getext_result = getext(self.mesh, - tuple(fns_residual), - tuple(fns_jacobian), - [x.fn for x in self.essential_bcs], - tuple(fns_bd_residual), - tuple(fns_bd_jacobian), - primary_field_list=prim_field_list, - verbose=verbose, - debug=debug,) + _getext_result = getext( + self.mesh, + JITCallbackSet( + residual=tuple(fns_residual), + bcs=tuple(x.fn for x in self.essential_bcs), + jacobian=tuple(fns_jacobian), + bd_residual=tuple(fns_bd_residual), + bd_jacobian=tuple(fns_bd_jacobian), + ), + prim_field_list, + verbose=verbose, + debug=debug, + ) self.compiled_extensions = _getext_result.ptrobj self.ext_dict = _getext_result.fn_dicts self.constants_manifest = _getext_result.constants_manifest @@ -2291,15 +2295,19 @@ class SNES_Vector(SolverBaseClass): # note also that the order here is important. prim_field_list = [self.u,] - _getext_result = getext(self.mesh, - tuple(fns_residual), - tuple(fns_jacobian), - [x.fn for x in self.essential_bcs], - tuple(fns_bd_residual), - tuple(fns_bd_jacobian), - primary_field_list=prim_field_list, - verbose=verbose, - debug=debug,) + _getext_result = getext( + self.mesh, + JITCallbackSet( + residual=tuple(fns_residual), + bcs=tuple(x.fn for x in self.essential_bcs), + jacobian=tuple(fns_jacobian), + bd_residual=tuple(fns_bd_residual), + bd_jacobian=tuple(fns_bd_jacobian), + ), + prim_field_list, + verbose=verbose, + debug=debug, + ) self.compiled_extensions = _getext_result.ptrobj self.ext_dict = _getext_result.fn_dicts self.constants_manifest = _getext_result.constants_manifest @@ -3669,17 +3677,21 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): print(f"Stokes: Jacobians complete, now compile", flush=True) prim_field_list = [self.u, self.p] - _getext_result = getext(self.mesh, - tuple(fns_residual), - tuple(fns_jacobian), - [x.fn for x in self.essential_bcs], - tuple(fns_bd_residual), - tuple(fns_bd_jacobian), - primary_field_list=prim_field_list, - verbose=verbose, - debug=debug, - debug_name=debug_name, - cache=False) + _getext_result = getext( + self.mesh, + JITCallbackSet( + residual=tuple(fns_residual), + bcs=tuple(x.fn for x in self.essential_bcs), + jacobian=tuple(fns_jacobian), + bd_residual=tuple(fns_bd_residual), + bd_jacobian=tuple(fns_bd_jacobian), + ), + prim_field_list, + verbose=verbose, + debug=debug, + debug_name=debug_name, + cache=False, + ) self.compiled_extensions = _getext_result.ptrobj self.ext_dict = _getext_result.fn_dicts self.constants_manifest = _getext_result.constants_manifest diff --git a/src/underworld3/cython/petsc_maths.pyx b/src/underworld3/cython/petsc_maths.pyx index fe2d19f8e..26e2f828b 100644 --- a/src/underworld3/cython/petsc_maths.pyx +++ b/src/underworld3/cython/petsc_maths.pyx @@ -3,7 +3,7 @@ import sympy import underworld3 import underworld3.timing as timing -from underworld3.utilities._jitextension import getext +from underworld3.utilities._jitextension import getext, JITCallbackSet from petsc4py import PETSc @@ -89,7 +89,8 @@ class Integral: self.dm = self.mesh.dm # .clone() mesh=self.mesh - _getext_result = getext(self.mesh, [self.fn,], [], [], [], [], self.mesh.vars.values(), verbose=verbose) + _getext_result = getext(self.mesh, JITCallbackSet(residual=(self.fn,)), + self.mesh.vars.values(), verbose=verbose) cdef PtrContainer ext = _getext_result.ptrobj # Pull out vec for variables, and go ahead with the integral @@ -273,7 +274,8 @@ class CellWiseIntegral: elif isinstance(self.fn, sympy.vector.Dyadic): raise RuntimeError("Integral evaluation for Dyadic integrands not supported.") - cdef PtrContainer ext = getext(self.mesh, [self.fn,], [], [], [], [], self.mesh.vars.values()).ptrobj + cdef PtrContainer ext = getext(self.mesh, JITCallbackSet(residual=(self.fn,)), + self.mesh.vars.values()).ptrobj # Pull out vec for variables, and go ahead with the integral self.mesh.update_lvec() @@ -388,8 +390,8 @@ class BdIntegral: # Compile integrand using the boundary residual slot (includes petsc_n[] in signature) _getext_result = getext( - self.mesh, [], [], [], [self.fn,], [], self.mesh.vars.values(), verbose=verbose - ) + self.mesh, JITCallbackSet(bd_residual=(self.fn,)), + self.mesh.vars.values(), verbose=verbose) cdef PtrContainer ext = _getext_result.ptrobj # Prepare the solution vector diff --git a/src/underworld3/utilities/_jitextension.py b/src/underworld3/utilities/_jitextension.py index 45af5636a..b8b9d5c03 100644 --- a/src/underworld3/utilities/_jitextension.py +++ b/src/underworld3/utilities/_jitextension.py @@ -1,11 +1,11 @@ -from typing import List +from typing import List, Optional, Tuple import subprocess from xmlrpc.client import boolean import sympy import underworld3 import underworld3.timing as timing -from typing import Optional from collections import namedtuple +from dataclasses import dataclass ## This is not required in sympy >= 1.9 @@ -33,6 +33,108 @@ _ext_dict = {} +# ============================================================================ +# JIT Callback Set +# ============================================================================ +# +# Groups the five callback lists that PETSc requires for pointwise functions. +# Using a structured container prevents cache-key collisions between callback +# roles (e.g. volume residual vs boundary residual) that share the same +# symbolic form. +# ============================================================================ + +@dataclass(frozen=True) +class JITCallbackSet: + """Immutable container for the five PETSc pointwise callback lists. + + Each slot holds a tuple of SymPy expressions for one callback role. + The structured representation ensures that cache keys preserve which + role each expression belongs to, preventing the collision bug where + ``Integral(fn=1)`` and ``BdIntegral(fn=1)`` would share a cached module. + + Parameters + ---------- + residual : tuple + Volume residual expressions (F0, F1 for each field). + bcs : tuple + Essential boundary condition expressions. + jacobian : tuple + Jacobian expressions (G0, G1, G2, G3 for each field pair). + bd_residual : tuple + Boundary residual expressions (includes ``petsc_n[]`` access). + bd_jacobian : tuple + Boundary Jacobian expressions. + """ + residual: tuple = () + bcs: tuple = () + jacobian: tuple = () + bd_residual: tuple = () + bd_jacobian: tuple = () + + def flat(self) -> tuple: + """Concatenate all slots into a single ordered tuple. + + The ordering (residual, bcs, jacobian, bd_residual, bd_jacobian) + matches what ``_createext()`` expects. + """ + return self.residual + self.bcs + self.jacobian + self.bd_residual + self.bd_jacobian + + def signature(self) -> tuple: + """Hashable key that preserves callback role separation. + + Two callback sets with the same expressions in different roles + will produce different signatures. + """ + return (self.residual, self.bcs, self.jacobian, self.bd_residual, self.bd_jacobian) + + def map(self, fn) -> 'JITCallbackSet': + """Apply *fn* to every expression in every slot, returning a new set.""" + return JITCallbackSet( + residual=tuple(fn(f) for f in self.residual), + bcs=tuple(fn(f) for f in self.bcs), + jacobian=tuple(fn(f) for f in self.jacobian), + bd_residual=tuple(fn(f) for f in self.bd_residual), + bd_jacobian=tuple(fn(f) for f in self.bd_jacobian), + ) + + @property + def counts(self): + """Lengths of each slot, for ``_createext()`` offset calculation.""" + return (len(self.residual), len(self.bcs), len(self.jacobian), + len(self.bd_residual), len(self.bd_jacobian)) + + +def prepare_for_cache_key(fn, constants_subs_map): + """Prepare a single expression for JIT cache hashing. + + Two-phase process: + 1. Substitute constant UWexpressions with ``_JITConstant`` placeholders + so that changing a constant's *value* does not invalidate the cache. + 2. Unwrap remaining (non-constant) UWexpressions to pure SymPy so the + hash is deterministic. + + Parameters + ---------- + fn : sympy expression or None + The expression to expand. + constants_subs_map : dict or None + Mapping from UWexpression symbols to ``_JITConstant`` placeholders. + """ + # Phase 1: Substitute constants with _JITConstant placeholders + if constants_subs_map and fn is not None: + try: + fn_structural = fn.xreplace(constants_subs_map) if hasattr(fn, "xreplace") else fn + except Exception: + fn_structural = fn + else: + fn_structural = fn + + # Phase 2: Unwrap remaining (non-constant) expressions + return underworld3.function.expressions.unwrap( + fn_structural, keep_constants=False, return_self=False + ) + + # ============================================================================ # JIT Constants Support # ============================================================================ @@ -275,20 +377,25 @@ def debugging_text_bd(randstr, fn, fn_type, eqn_no): @timing.routine_timer_decorator def getext( mesh, - fns_residual, - fns_jacobian, - fns_bcs, - fns_bd_residual, - fns_bd_jacobian, + callbacks: JITCallbackSet, primary_field_list, verbose=False, debug=False, debug_name=None, cache=True, ): - """ - Check if we've already created an equivalent extension - and use if available. + """Compile (or retrieve cached) JIT extension for PETSc pointwise functions. + + Parameters + ---------- + mesh : Mesh + Supporting mesh for coordinate system and variable information. + callbacks : JITCallbackSet + Callback expressions grouped by PETSc role (residual, bcs, jacobian, + bd_residual, bd_jacobian). + primary_field_list : iterable + Variables that map to PETSc primary arrays (``petsc_u[]``). + All others map to auxiliary arrays (``petsc_a[]``). Returns ------- @@ -302,65 +409,17 @@ def getext( time_s = time.time() primary_field_list = tuple(primary_field_list) - raw_fns_residual = tuple(fns_residual) - raw_fns_bcs = tuple(fns_bcs) - raw_fns_jacobian = tuple(fns_jacobian) - raw_fns_bd_residual = tuple(fns_bd_residual) - raw_fns_bd_jacobian = tuple(fns_bd_jacobian) - - raw_fns = ( - raw_fns_residual - + raw_fns_bcs - + raw_fns_jacobian - + raw_fns_bd_residual - + raw_fns_bd_jacobian - ) - # Extract constant UWexpressions that will go through constants[] array - constants_manifest, constants_subs_map = _extract_constants(raw_fns, mesh) + constants_manifest, constants_subs_map = _extract_constants(callbacks.flat(), mesh) # Build structurally-expanded functions for cache hashing. # Constants are replaced with placeholder symbols (value-independent), # so changing a constant value won't cause a cache miss. - def _structural_expand(fn): - # Phase 1: Substitute constants with _JITConstant placeholders - if constants_subs_map and fn is not None: - try: - fn_structural = fn.xreplace(constants_subs_map) if hasattr(fn, "xreplace") else fn - except Exception: - fn_structural = fn - else: - fn_structural = fn - - # Phase 2: Unwrap remaining (non-constant) expressions - return underworld3.function.expressions.unwrap( - fn_structural, keep_constants=False, return_self=False - ) - - expanded_fns_residual = tuple(_structural_expand(fn) for fn in raw_fns_residual) - expanded_fns_bcs = tuple(_structural_expand(fn) for fn in raw_fns_bcs) - expanded_fns_jacobian = tuple(_structural_expand(fn) for fn in raw_fns_jacobian) - expanded_fns_bd_residual = tuple(_structural_expand(fn) for fn in raw_fns_bd_residual) - expanded_fns_bd_jacobian = tuple(_structural_expand(fn) for fn in raw_fns_bd_jacobian) - - fns = ( - expanded_fns_residual - + expanded_fns_bcs - + expanded_fns_jacobian - + expanded_fns_bd_residual - + expanded_fns_bd_jacobian - ) - fns_signature = ( - expanded_fns_residual, - expanded_fns_bcs, - expanded_fns_jacobian, - expanded_fns_bd_residual, - expanded_fns_bd_jacobian, - ) + expanded = callbacks.map(lambda fn: prepare_for_cache_key(fn, constants_subs_map)) if debug and underworld3.mpi.rank == 0: print(f"Expanded functions for compilation:") - for i, fn in enumerate(fns): + for i, fn in enumerate(expanded.flat()): print(f"{i}: {fn}") if constants_manifest: print(f"Constants manifest ({len(constants_manifest)} entries):") @@ -378,13 +437,13 @@ def _structural_expand(fn): # unique modules. jitname += "_" + str(len(_ext_dict.keys())) - else: # Else name from a structured hash — function role/signature must be preserved. + else: # Name from structured hash — function role must be preserved. primary_field_signature = tuple( (getattr(field, "field_id", None), getattr(field, "clean_name", None)) for field in primary_field_list ) jitname = abs( - hash((mesh, fns_signature, tuple(mesh.vars.keys()), primary_field_signature)) + hash((mesh, expanded.signature(), tuple(mesh.vars.keys()), primary_field_signature)) ) # Create the module if not in dictionary @@ -392,11 +451,7 @@ def _structural_expand(fn): _createext( jitname, mesh, - fns_residual, - fns_bcs, - fns_jacobian, - fns_bd_residual, - fns_bd_jacobian, + callbacks, primary_field_list, constants_subs_map=constants_subs_map, verbose=verbose, @@ -410,25 +465,11 @@ def _structural_expand(fn): module = _ext_dict[jitname] ptrobj = module.getptrobj() - i_res = {} - for index, fn in enumerate(fns_residual): - i_res[fn] = index - - i_ebc = {} - for index, fn in enumerate(fns_bcs): - i_ebc[fn] = index - - i_jac = {} - for index, fn in enumerate(fns_jacobian): - i_jac[fn] = index - - i_bd_res = {} - for index, fn in enumerate(fns_bd_residual): - i_bd_res[fn] = index - - i_bd_jac = {} - for index, fn in enumerate(fns_bd_jacobian): - i_bd_jac[fn] = index + i_res = {fn: i for i, fn in enumerate(callbacks.residual)} + i_ebc = {fn: i for i, fn in enumerate(callbacks.bcs)} + i_jac = {fn: i for i, fn in enumerate(callbacks.jacobian)} + i_bd_res = {fn: i for i, fn in enumerate(callbacks.bd_residual)} + i_bd_jac = {fn: i for i, fn in enumerate(callbacks.bd_jacobian)} extn_fn_dict = namedtuple( "Functions", @@ -444,74 +485,37 @@ def _structural_expand(fn): def _createext( name: str, mesh: underworld3.discretisation.Mesh, - fns_residual: List[sympy.Basic], - fns_bcs: List[sympy.Basic], - fns_jacobian: List[sympy.Basic], - fns_bd_residual: List[sympy.Basic], - fns_bd_jacobian: List[sympy.Basic], - primary_field_list: List[underworld3.discretisation.MeshVariable], + callbacks: JITCallbackSet, + primary_field_list, constants_subs_map: Optional[dict] = None, verbose: Optional[bool] = False, debug: Optional[bool] = False, debug_name=None, ): - """ - This creates the required extension which houses the JIT - fn pointer for PETSc. + """Create the JIT extension module with PETSc function pointers. Note that it is not possible to replace loaded shared libraries in Python, so we instead create a new extension for each new function. - We hash the functions and create a dictionary of the generated extensions - to avoid redundantly creating new extensions. - - Params - ------ - name: - Name for the extension. It will be prepended with "fn_ptr_ext_" - mesh: - Supporting mesh. It is used to get coordinate system and variable - information. - fns_residual: - List of system's residual sympy functions for which JIT equivalents - will be generated. - fns_jacobian: - List of system's Jacobian sympy functions for which JIT equivalents - will be generated. - fns_bcs: - List of system's boundary condition sympy functions for which JIT equivalents - will be generated. - fns_bd_residual: - List of system's boundary integral sympy functions for which JIT equivalents - will be generated. - fns_bd_jacobian: - List of system's boundary integral jacobian sympy functions for which JIT equivalents - will be generated. - primary_field_list - List of variables that will map from petsc primary variable arrays. All - other variables will be obtained from the mesh object and will be mapped to - petsc auxiliary variable arrays. Note that *all* the variables in the - calling system's corresponding `PetscDM` must be included in this list. - They must also be ordered according to their `field_id`. - + Parameters + ---------- + name : str + Name for the extension. It will be prepended with "fn_ptr_ext_". + mesh : Mesh + Supporting mesh for coordinate system and variable information. + callbacks : JITCallbackSet + Structured callback expressions grouped by role. + primary_field_list : list + Variables that map to PETSc primary variable arrays (``petsc_u[]``). + All other variables map to auxiliary arrays (``petsc_a[]``). + Must be ordered by ``field_id``. """ from sympy import symbols, Eq, MatrixSymbol from underworld3 import VarType - # Note that the order here is important. - fns = ( - tuple(fns_residual) - + tuple(fns_bcs) - + tuple(fns_jacobian) - + tuple(fns_bd_residual) - + tuple(fns_bd_jacobian) - ) - - count_residual_sig = len(fns_residual) - count_bc_sig = len(fns_bcs) - count_jacobian_sig = len(fns_jacobian) - count_bd_residual_sig = len(fns_bd_residual) - count_bd_jacobian_sig = len(fns_bd_jacobian) + fns = callbacks.flat() + count_residual_sig, count_bc_sig, count_jacobian_sig, \ + count_bd_residual_sig, count_bd_jacobian_sig = callbacks.counts # `_ccode` patching def ccode_patch_fns(varlist, prefix_str): @@ -932,39 +936,39 @@ def _basescalar_ccode(self, printer): clsguy.fns_bd_residual = malloc({}*sizeof(PetscDSBdResidualFn)) clsguy.fns_bd_jacobian = malloc({}*sizeof(PetscDSBdJacobianFn)) """.format( - len(fns_residual), - len(fns_bcs), - len(fns_jacobian), - len(fns_bd_residual), - len(fns_bd_jacobian), + count_residual_sig, + count_bc_sig, + count_jacobian_sig, + count_bd_residual_sig, + count_bd_jacobian_sig, ) eqn_count = 0 - for index, eqn in enumerate(eqns[eqn_count : eqn_count + len(fns_residual)]): + for index, eqn in enumerate(eqns[eqn_count : eqn_count + count_residual_sig]): pyx_str += " clsguy.fns_residual[{}] = {}_petsc_{}\n".format(index, randstr, eqn[0]) eqn_count += 1 residual_equations = (0, eqn_count) - for index, eqn in enumerate(eqns[eqn_count : eqn_count + len(fns_bcs)]): + for index, eqn in enumerate(eqns[eqn_count : eqn_count + count_bc_sig]): pyx_str += " clsguy.fns_bcs[{}] = {}_petsc_{}\n".format(index, randstr, eqn[0]) eqn_count += 1 boundary_equations = (residual_equations[1], eqn_count) - for index, eqn in enumerate(eqns[eqn_count : eqn_count + len(fns_jacobian)]): + for index, eqn in enumerate(eqns[eqn_count : eqn_count + count_jacobian_sig]): pyx_str += " clsguy.fns_jacobian[{}] = {}_petsc_{}\n".format(index, randstr, eqn[0]) eqn_count += 1 jacobian_equations = (boundary_equations[1], eqn_count) - for index, eqn in enumerate(eqns[eqn_count : eqn_count + len(fns_bd_residual)]): + for index, eqn in enumerate(eqns[eqn_count : eqn_count + count_bd_residual_sig]): pyx_str += " clsguy.fns_bd_residual[{}] = {}_petsc_{}\n".format(index, randstr, eqn[0]) eqn_count += 1 boundary_residual_equations = (jacobian_equations[1], eqn_count) - for index, eqn in enumerate(eqns[eqn_count : eqn_count + len(fns_bd_jacobian)]): + for index, eqn in enumerate(eqns[eqn_count : eqn_count + count_bd_jacobian_sig]): pyx_str += " clsguy.fns_bd_jacobian[{}] = {}_petsc_{}\n".format(index, randstr, eqn[0]) eqn_count += 1 @@ -1061,23 +1065,23 @@ def load_dynamic(name, path, file=None): flush=True, ) print( - f"{randstr} {len(fns_residual):5d} residuals: {residual_equations[0]}:{residual_equations[1]}", + f"{randstr} {count_residual_sig:5d} residuals: {residual_equations[0]}:{residual_equations[1]}", flush=True, ) print( - f"{randstr} {len(fns_bcs):5d} boundaries: {boundary_equations[0]}:{boundary_equations[1]}", + f"{randstr} {count_bc_sig:5d} boundaries: {boundary_equations[0]}:{boundary_equations[1]}", flush=True, ) print( - f"{randstr} {len(fns_jacobian):5d} jacobians: {jacobian_equations[0]}:{jacobian_equations[1]}", + f"{randstr} {count_jacobian_sig:5d} jacobians: {jacobian_equations[0]}:{jacobian_equations[1]}", flush=True, ) print( - f"{randstr} {len(fns_bd_residual):5d} boundary_res: {boundary_residual_equations[0]}:{boundary_residual_equations[1]}", + f"{randstr} {count_bd_residual_sig:5d} boundary_res: {boundary_residual_equations[0]}:{boundary_residual_equations[1]}", flush=True, ) print( - f"{randstr} {len(fns_bd_jacobian):5d} boundary_jac: {boundary_jacobian_equations[0]}:{boundary_jacobian_equations[1]}", + f"{randstr} {count_bd_jacobian_sig:5d} boundary_jac: {boundary_jacobian_equations[0]}:{boundary_jacobian_equations[1]}", flush=True, ) diff --git a/tests/test_0004_pointwise_fns.py b/tests/test_0004_pointwise_fns.py index 7c5f3b60d..63c2570d2 100644 --- a/tests/test_0004_pointwise_fns.py +++ b/tests/test_0004_pointwise_fns.py @@ -14,7 +14,7 @@ import numpy as np import sympy -from underworld3.utilities._jitextension import getext +from underworld3.utilities._jitextension import getext, JITCallbackSet # build a small mesh - we'll load up a simple problem and then see what functions are loaded @@ -64,11 +64,13 @@ def test_getext_simple(): with uw.utilities.CaptureStdout(split=True) as captured_setup_solver: _getext_result = getext( mesh, - [res_fn, res_fn], - [jac_fn], - [bc_fn], - [bd_res_fn], - [bd_jac_fn], + JITCallbackSet( + residual=(res_fn, res_fn), + bcs=(bc_fn,), + jacobian=(jac_fn,), + bd_residual=(bd_res_fn,), + bd_jacobian=(bd_jac_fn,), + ), mesh.vars.values(), verbose=True, debug=True, @@ -109,11 +111,13 @@ def test_getext_sympy_fns(): with uw.utilities.CaptureStdout(split=True) as captured_setup_solver: _getext_result = getext( mesh, - [res_fn, res_fn], - [jac_fn], - [bc_fn], - [bd_res_fn], - [bd_jac_fn], + JITCallbackSet( + residual=(res_fn, res_fn), + bcs=(bc_fn,), + jacobian=(jac_fn,), + bd_residual=(bd_res_fn,), + bd_jacobian=(bd_jac_fn,), + ), mesh.vars.values(), verbose=True, debug=True, @@ -163,11 +167,13 @@ def test_getext_meshVar(): with uw.utilities.CaptureStdout(split=True) as captured_setup_solver: _getext_result = getext( mesh, - [res_fn, res_fn], - [jac_fn], - [bc_fn], - [bd_res_fn], - [bd_jac_fn], + JITCallbackSet( + residual=(res_fn, res_fn), + bcs=(bc_fn,), + jacobian=(jac_fn,), + bd_residual=(bd_res_fn,), + bd_jacobian=(bd_jac_fn,), + ), mesh.vars.values(), verbose=True, debug=True, From a3f203c300913ff2eb29edbc0efb66a0c407c800 Mon Sep 17 00:00:00 2001 From: Louis Moresi Date: Wed, 25 Mar 2026 22:36:07 +1100 Subject: [PATCH 2/2] Add per-function JIT cache to avoid redundant recompilation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The existing JIT cache keys the entire bundle of functions (residuals + Jacobians + BCs) as one hash. If any single function changes, everything recompiles. This is wasteful for SNES_Tensor_Projection which loops over tensor components, changing only the residual RHS while the Jacobian (identity) never changes. New per-function cache: each compiled C function is cached individually by its structural hash + signature type. When getext() is called: - Fast path: whole-bundle hash hit → return immediately (unchanged) - Slow path: check per-function hashes, compile only new functions, assemble PtrContainer by copying cached function pointers PtrContainer gains allocate() and copy_*_from() methods to support cross-module pointer assembly. Also fixes elastic_dt → dt_elastic bug in VE_Stokes.solve(). Includes timing/profiling test scripts for JIT benchmarking. Underworld development team with AI support from Claude Code --- src/underworld3/cython/petsc_types.pxd | 13 +- src/underworld3/cython/petsc_types.pyx | 31 +++- src/underworld3/systems/solvers.py | 2 +- src/underworld3/utilities/_jitextension.py | 179 ++++++++++++++++++--- tests/minimal_vep_timing.py | 89 ++++++++++ tests/profile_jit_phases.py | 163 +++++++++++++++++++ 6 files changed, 453 insertions(+), 24 deletions(-) create mode 100644 tests/minimal_vep_timing.py create mode 100644 tests/profile_jit_phases.py diff --git a/src/underworld3/cython/petsc_types.pxd b/src/underworld3/cython/petsc_types.pxd index 7d2509666..07ce18876 100644 --- a/src/underworld3/cython/petsc_types.pxd +++ b/src/underworld3/cython/petsc_types.pxd @@ -22,21 +22,21 @@ ctypedef void(*PetscDSResidualFn)(PetscInt, PetscInt, PetscInt, ctypedef void (*PetscDSJacobianFn)(PetscInt, PetscInt, PetscInt, const PetscInt*, const PetscInt*, const PetscScalar*, const PetscScalar*, const PetscScalar*, const PetscInt*, const PetscInt*, const PetscScalar*, const PetscScalar*, const PetscScalar*, - PetscReal, PetscReal, const PetscReal*, PetscInt, const PetscScalar*, + PetscReal, PetscReal, const PetscReal*, PetscInt, const PetscScalar*, PetscScalar*) ctypedef void(*PetscDSBdResidualFn)( PetscInt, PetscInt, PetscInt, const PetscInt*, const PetscInt*, const PetscScalar*, const PetscScalar*, const PetscScalar*, const PetscInt*, const PetscInt*, const PetscScalar*, const PetscScalar*, const PetscScalar*, - PetscReal, const PetscReal*,const PetscReal*, PetscInt, const PetscScalar*, + PetscReal, const PetscReal*,const PetscReal*, PetscInt, const PetscScalar*, PetscScalar* ) ctypedef void (*PetscDSBdJacobianFn)( PetscInt, PetscInt, PetscInt, const PetscInt*, const PetscInt*, const PetscScalar*, const PetscScalar*, const PetscScalar*, const PetscInt*, const PetscInt*, const PetscScalar*, const PetscScalar*, const PetscScalar*, - PetscReal, PetscReal, const PetscReal*, const PetscReal*, PetscInt, + PetscReal, PetscReal, const PetscReal*, const PetscReal*, PetscInt, const PetscScalar*, PetscScalar* ) @@ -47,4 +47,9 @@ cdef class PtrContainer: cdef PetscDSBdResidualFn* fns_bd_residual cdef PetscDSBdJacobianFn* fns_bd_jacobian - + cpdef allocate(self, int n_res, int n_bcs, int n_jac, int n_bd_res, int n_bd_jac) + cpdef copy_residual_from(self, int dst, PtrContainer src, int src_idx) + cpdef copy_bcs_from(self, int dst, PtrContainer src, int src_idx) + cpdef copy_jacobian_from(self, int dst, PtrContainer src, int src_idx) + cpdef copy_bd_residual_from(self, int dst, PtrContainer src, int src_idx) + cpdef copy_bd_jacobian_from(self, int dst, PtrContainer src, int src_idx) diff --git a/src/underworld3/cython/petsc_types.pyx b/src/underworld3/cython/petsc_types.pyx index 01a95da62..609dad2a6 100644 --- a/src/underworld3/cython/petsc_types.pyx +++ b/src/underworld3/cython/petsc_types.pyx @@ -1,2 +1,31 @@ +from libc.stdlib cimport malloc + cdef class PtrContainer: - pass \ No newline at end of file + + cpdef allocate(self, int n_res, int n_bcs, int n_jac, int n_bd_res, int n_bd_jac): + """Allocate function pointer arrays of the given sizes.""" + self.fns_residual = malloc(n_res * sizeof(PetscDSResidualFn)) + self.fns_bcs = malloc(n_bcs * sizeof(PetscDSResidualFn)) + self.fns_jacobian = malloc(n_jac * sizeof(PetscDSJacobianFn)) + self.fns_bd_residual = malloc(n_bd_res * sizeof(PetscDSBdResidualFn)) + self.fns_bd_jacobian = malloc(n_bd_jac * sizeof(PetscDSBdJacobianFn)) + + cpdef copy_residual_from(self, int dst, PtrContainer src, int src_idx): + """Copy a residual function pointer from another container.""" + self.fns_residual[dst] = src.fns_residual[src_idx] + + cpdef copy_bcs_from(self, int dst, PtrContainer src, int src_idx): + """Copy a BC function pointer from another container.""" + self.fns_bcs[dst] = src.fns_bcs[src_idx] + + cpdef copy_jacobian_from(self, int dst, PtrContainer src, int src_idx): + """Copy a Jacobian function pointer from another container.""" + self.fns_jacobian[dst] = src.fns_jacobian[src_idx] + + cpdef copy_bd_residual_from(self, int dst, PtrContainer src, int src_idx): + """Copy a boundary residual function pointer from another container.""" + self.fns_bd_residual[dst] = src.fns_bd_residual[src_idx] + + cpdef copy_bd_jacobian_from(self, int dst, PtrContainer src, int src_idx): + """Copy a boundary Jacobian function pointer from another container.""" + self.fns_bd_jacobian[dst] = src.fns_bd_jacobian[src_idx] diff --git a/src/underworld3/systems/solvers.py b/src/underworld3/systems/solvers.py index 24caecef1..550632306 100644 --- a/src/underworld3/systems/solvers.py +++ b/src/underworld3/systems/solvers.py @@ -1275,7 +1275,7 @@ def solve( timestep = self.delta_t.sym if timestep != self.delta_t: - self._constitutive_model.Parameters.elastic_dt = timestep # this will force an initialisation because the functions need to be updated + self._constitutive_model.Parameters.dt_elastic = timestep # this will force an initialisation because the functions need to be updated if _force_setup: self.is_setup = False diff --git a/src/underworld3/utilities/_jitextension.py b/src/underworld3/utilities/_jitextension.py index b8b9d5c03..c2d737bcb 100644 --- a/src/underworld3/utilities/_jitextension.py +++ b/src/underworld3/utilities/_jitextension.py @@ -32,6 +32,21 @@ _ext_dict = {} +# Per-function cache: maps (fn_hash, sig_type) -> _CachedFn +_fn_cache = {} + + +@dataclass +class _CachedFn: + """Registry entry for a single cached compiled function pointer.""" + ptr_container: object # PtrContainer from the .so that compiled this fn + sig_type: str # "residual" | "jacobian" | "bcs" | "bd_residual" | "bd_jacobian" + index: int # index within that container's sig_type array + + +# Signature type names used by the per-function cache +_SIG_TYPES = ("residual", "bcs", "jacobian", "bd_residual", "bd_jacobian") + # ============================================================================ # JIT Callback Set @@ -428,6 +443,11 @@ def getext( import os + primary_field_signature = tuple( + (getattr(field, "field_id", None), getattr(field, "clean_name", None)) + for field in primary_field_list + ) + if debug_name is not None: jitname = debug_name @@ -438,33 +458,155 @@ def getext( jitname += "_" + str(len(_ext_dict.keys())) else: # Name from structured hash — function role must be preserved. - primary_field_signature = tuple( - (getattr(field, "field_id", None), getattr(field, "clean_name", None)) - for field in primary_field_list - ) jitname = abs( hash((mesh, expanded.signature(), tuple(mesh.vars.keys()), primary_field_signature)) ) - # Create the module if not in dictionary - if jitname not in _ext_dict.keys() or not cache: + # ── Fast path: whole-bundle cache hit ────────────────────────────────── + if jitname in _ext_dict and cache: + if verbose and underworld3.mpi.rank == 0: + print(f"JIT compiled module cached ... {jitname} ", flush=True) + + module = _ext_dict[jitname] + ptrobj = module.getptrobj() + + i_res = {fn: i for i, fn in enumerate(callbacks.residual)} + i_ebc = {fn: i for i, fn in enumerate(callbacks.bcs)} + i_jac = {fn: i for i, fn in enumerate(callbacks.jacobian)} + i_bd_res = {fn: i for i, fn in enumerate(callbacks.bd_residual)} + i_bd_jac = {fn: i for i, fn in enumerate(callbacks.bd_jacobian)} + + extn_fn_dict = namedtuple( + "Functions", ["res", "jac", "ebc", "bd_res", "bd_jac"], + ) + return _GextResult( + ptrobj, + extn_fn_dict(i_res, i_jac, i_ebc, i_bd_res, i_bd_jac), + constants_manifest, + ) + + # ── Per-function cache: check which individual functions are cached ── + # Build a hashable key for the constants manifest so it's part of + # every per-function hash (ensures constants[i] means the same thing). + constants_manifest_key = tuple( + (str(expr), idx) for idx, expr in constants_manifest + ) + + # Compute per-function hashes for each expression, grouped by sig_type + fn_hashes = {} # maps (sig_type, slot_index) -> hash + for sig_type, slot_fns in zip( + _SIG_TYPES, + [expanded.residual, expanded.bcs, expanded.jacobian, + expanded.bd_residual, expanded.bd_jacobian], + ): + for i, fn_expanded in enumerate(slot_fns): + fn_hashes[(sig_type, i)] = abs( + hash((mesh, fn_expanded, tuple(mesh.vars.keys()), + sig_type, constants_manifest_key, primary_field_signature)) + ) + + # Partition into cached vs new + cached_hits = {} # (sig_type, slot_index) -> _CachedFn + new_needed = {} # (sig_type, slot_index) -> original expression + for sig_type, slot_fns in zip( + _SIG_TYPES, + [callbacks.residual, callbacks.bcs, callbacks.jacobian, + callbacks.bd_residual, callbacks.bd_jacobian], + ): + for i, fn_orig in enumerate(slot_fns): + key = (sig_type, i) + fn_hash = fn_hashes[key] + cache_key = (fn_hash, sig_type) + if cache_key in _fn_cache and cache: + cached_hits[key] = _fn_cache[cache_key] + else: + new_needed[key] = fn_orig + + n_hits = len(cached_hits) + n_new = len(new_needed) + + if verbose and underworld3.mpi.rank == 0: + total = n_hits + n_new + print(f"Per-function cache: {n_hits}/{total} hits, {n_new} new", flush=True) + + # ── Compile new functions ─────────────────────────────────────────── + new_ptr = None + new_indices = {} # (sig_type, slot_index) -> index in new_ptr's arrays + + if n_new > 0: + # Build a JITCallbackSet containing ONLY the new functions, + # preserving order within each sig_type. + new_by_type = {st: [] for st in _SIG_TYPES} + new_slot_map = {st: [] for st in _SIG_TYPES} # tracks original slot indices + for (sig_type, slot_idx), fn_orig in sorted(new_needed.items()): + new_by_type[sig_type].append(fn_orig) + new_slot_map[sig_type].append(slot_idx) + + new_callbacks = JITCallbackSet( + residual=tuple(new_by_type["residual"]), + bcs=tuple(new_by_type["bcs"]), + jacobian=tuple(new_by_type["jacobian"]), + bd_residual=tuple(new_by_type["bd_residual"]), + bd_jacobian=tuple(new_by_type["bd_jacobian"]), + ) + + # Compile the new functions + new_jitname = abs(hash((jitname, "partial", n_new, time.time()))) _createext( - jitname, + new_jitname, mesh, - callbacks, + new_callbacks, primary_field_list, constants_subs_map=constants_subs_map, verbose=verbose, debug=debug, debug_name=debug_name, ) - else: - if verbose and underworld3.mpi.rank == 0: - print(f"JIT compiled module cached ... {jitname} ", flush=True) - module = _ext_dict[jitname] - ptrobj = module.getptrobj() + new_module = _ext_dict[new_jitname] + new_ptr = new_module.getptrobj() + + # Register each new function in the per-function cache + for sig_type in _SIG_TYPES: + for local_idx, slot_idx in enumerate(new_slot_map[sig_type]): + key = (sig_type, slot_idx) + fn_hash = fn_hashes[key] + cache_key = (fn_hash, sig_type) + entry = _CachedFn( + ptr_container=new_ptr, + sig_type=sig_type, + index=local_idx, + ) + _fn_cache[cache_key] = entry + cached_hits[key] = entry + + # Also register the full bundle in _ext_dict if ALL functions were new + # (common case: first compile of a solver) + if n_new > 0 and n_hits == 0: + _ext_dict[jitname] = _ext_dict[new_jitname] + + # ── Assemble PtrContainer from cached function pointers ───────────── + from underworld3.cython.petsc_types import PtrContainer + + result_ptr = PtrContainer() + counts = callbacks.counts + result_ptr.allocate(*counts) + + _copy_methods = { + "residual": result_ptr.copy_residual_from, + "bcs": result_ptr.copy_bcs_from, + "jacobian": result_ptr.copy_jacobian_from, + "bd_residual": result_ptr.copy_bd_residual_from, + "bd_jacobian": result_ptr.copy_bd_jacobian_from, + } + for sig_type, n_fns in zip(_SIG_TYPES, counts): + copy_fn = _copy_methods[sig_type] + for slot_idx in range(n_fns): + entry = cached_hits[(sig_type, slot_idx)] + copy_fn(slot_idx, entry.ptr_container, entry.index) + + # ── Build fn_dicts (unchanged from original) ──────────────────────── i_res = {fn: i for i, fn in enumerate(callbacks.residual)} i_ebc = {fn: i for i, fn in enumerate(callbacks.bcs)} i_jac = {fn: i for i, fn in enumerate(callbacks.jacobian)} @@ -472,13 +614,14 @@ def getext( i_bd_jac = {fn: i for i, fn in enumerate(callbacks.bd_jacobian)} extn_fn_dict = namedtuple( - "Functions", - ["res", "jac", "ebc", "bd_res", "bd_jac"], + "Functions", ["res", "jac", "ebc", "bd_res", "bd_jac"], ) - extensions_functions_dicts = extn_fn_dict(i_res, i_jac, i_ebc, i_bd_res, i_bd_jac) - - return _GextResult(ptrobj, extensions_functions_dicts, constants_manifest) + return _GextResult( + result_ptr, + extn_fn_dict(i_res, i_jac, i_ebc, i_bd_res, i_bd_jac), + constants_manifest, + ) @timing.routine_timer_decorator diff --git a/tests/minimal_vep_timing.py b/tests/minimal_vep_timing.py new file mode 100644 index 000000000..e03f2937b --- /dev/null +++ b/tests/minimal_vep_timing.py @@ -0,0 +1,89 @@ +"""Minimal VEP timing test — isolate where time is spent. + +Run with: pixi run -e amr-dev python tests/minimal_vep_timing.py +""" + +import time +import sympy +import underworld3 as uw +from underworld3.systems import VE_Stokes + +t0 = time.time() + +# --- Mesh --- +mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), + cellSize=1.0 / 8, qdegree=3, +) +print(f"Mesh: {time.time() - t0:.1f}s") + +# --- Variables --- +v = uw.discretisation.MeshVariable("U", mesh, 2, degree=2, vtype=uw.VarType.VECTOR) +p = uw.discretisation.MeshVariable("P", mesh, 1, degree=1, continuous=True, + vtype=uw.VarType.SCALAR) +print(f"Variables: {time.time() - t0:.1f}s") + +# --- Solver + constitutive model --- +stokes = VE_Stokes(mesh, velocityField=v, pressureField=p, order=1) +stokes.constitutive_model = uw.constitutive_models.ViscoElasticPlasticFlowModel +stokes.constitutive_model.Parameters.shear_viscosity_0 = 1.0 +stokes.constitutive_model.Parameters.shear_modulus = 1.0 +stokes.constitutive_model.Parameters.shear_viscosity_min = 1.0e-3 +stokes.constitutive_model.Parameters.strainrate_inv_II_min = 1.0e-10 +stokes.saddle_preconditioner = 1.0 +stokes.tolerance = 1.0e-4 +print(f"Solver setup: {time.time() - t0:.1f}s") + +# --- Yield stress: Piecewise (fault layer) --- +x, y = mesh.X +tau_y = sympy.Piecewise( + (0.3, (y >= 0.45) & (y <= 0.55)), + (1.0e6, True), +) +stokes.constitutive_model.Parameters.yield_stress = tau_y +print(f"Yield stress set: {time.time() - t0:.1f}s") + +# --- BCs --- +stokes.add_essential_bc(sympy.Matrix([0.5, 0.0]), "Top") +stokes.add_essential_bc(sympy.Matrix([0.0, 0.0]), "Bottom") +stokes.add_essential_bc((sympy.oo, 0.0), "Left") +stokes.add_essential_bc((sympy.oo, 0.0), "Right") +stokes.bodyforce = sympy.Matrix([0.0, 0.0]) +stokes.petsc_options["ksp_type"] = "fgmres" +print(f"BCs set: {time.time() - t0:.1f}s") + +# --- First solve (includes JIT compilation) --- +# Set dt_elastic explicitly to work around elastic_dt alias bug in VE_Stokes.solve() +t1 = time.time() +stokes.solve(timestep=0.02, zero_init_guess=True) +print(f"First solve (incl JIT): {time.time() - t1:.1f}s") + +# --- Second solve (cached JIT) --- +t2 = time.time() +stokes.solve(timestep=0.02, zero_init_guess=False) +print(f"Second solve (cached): {time.time() - t2:.1f}s") + +# --- Compare: pure VE (no yield) --- +stokes2 = VE_Stokes(mesh, velocityField=v, pressureField=p, order=1) +stokes2.constitutive_model = uw.constitutive_models.ViscoElasticPlasticFlowModel +stokes2.constitutive_model.Parameters.shear_viscosity_0 = 1.0 +stokes2.constitutive_model.Parameters.shear_modulus = 1.0 +# yield_stress defaults to sympy.oo — pure VE +stokes2.saddle_preconditioner = 1.0 +stokes2.tolerance = 1.0e-4 +stokes2.add_essential_bc(sympy.Matrix([0.5, 0.0]), "Top") +stokes2.add_essential_bc(sympy.Matrix([0.0, 0.0]), "Bottom") +stokes2.add_essential_bc((sympy.oo, 0.0), "Left") +stokes2.add_essential_bc((sympy.oo, 0.0), "Right") +stokes2.bodyforce = sympy.Matrix([0.0, 0.0]) +stokes2.petsc_options["ksp_type"] = "fgmres" + +t3 = time.time() +stokes2.solve(timestep=0.02, zero_init_guess=True) +print(f"Pure VE first solve (incl JIT): {time.time() - t3:.1f}s") + +t4 = time.time() +stokes2.solve(timestep=0.02, zero_init_guess=False) +print(f"Pure VE second solve (cached): {time.time() - t4:.1f}s") + +print(f"\nTotal: {time.time() - t0:.1f}s") diff --git a/tests/profile_jit_phases.py b/tests/profile_jit_phases.py new file mode 100644 index 000000000..ecd12eba0 --- /dev/null +++ b/tests/profile_jit_phases.py @@ -0,0 +1,163 @@ +"""Profile JIT compilation phases — isolate where time is spent. + +Instruments: sympy derivatives, expression unwrapping, hashing, C code generation, +Cython compilation, and the actual PETSc solve. + +Run with: pixi run -e default python tests/profile_jit_phases.py +""" + +import time +import sympy +import underworld3 as uw +from underworld3.systems import VE_Stokes + +# ── Setup (fast) ────────────────────────────────────────────────────────────── + +mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), + cellSize=1.0 / 8, qdegree=3, +) + +v = uw.discretisation.MeshVariable("U", mesh, 2, degree=2, vtype=uw.VarType.VECTOR) +p = uw.discretisation.MeshVariable("P", mesh, 1, degree=1, continuous=True, + vtype=uw.VarType.SCALAR) + +stokes = VE_Stokes(mesh, velocityField=v, pressureField=p, order=1) +stokes.constitutive_model = uw.constitutive_models.ViscoElasticPlasticFlowModel +stokes.constitutive_model.Parameters.shear_viscosity_0 = 1.0 +stokes.constitutive_model.Parameters.shear_modulus = 1.0 +stokes.constitutive_model.Parameters.shear_viscosity_min = 1.0e-3 +stokes.constitutive_model.Parameters.strainrate_inv_II_min = 1.0e-10 +stokes.saddle_preconditioner = 1.0 +stokes.tolerance = 1.0e-4 + +x, y = mesh.X +tau_y = sympy.Piecewise( + (0.3, (y >= 0.45) & (y <= 0.55)), + (1.0e6, True), +) +stokes.constitutive_model.Parameters.yield_stress = tau_y + +stokes.add_essential_bc(sympy.Matrix([0.5, 0.0]), "Top") +stokes.add_essential_bc(sympy.Matrix([0.0, 0.0]), "Bottom") +stokes.add_essential_bc((sympy.oo, 0.0), "Left") +stokes.add_essential_bc((sympy.oo, 0.0), "Right") +stokes.bodyforce = sympy.Matrix([0.0, 0.0]) +stokes.petsc_options["ksp_type"] = "fgmres" + +print("Setup complete.\n") + +# ── Phase 1: Sympy derivative computation ───────────────────────────────────── + +dim = mesh.dim + +# Get residual terms (these are already built by the constitutive model) +F0 = sympy.Array(stokes.F0.sym) +F1 = sympy.Array(stokes.F1.sym) +PF0 = sympy.Array(stokes.PF0.sym) + +print(f"F0 has {len(F0.free_symbols)} free symbols, {sum(1 for _ in sympy.preorder_traversal(sympy.Matrix(F0)))} nodes") +print(f"F1 has {len(F1.free_symbols)} free symbols, {sum(1 for _ in sympy.preorder_traversal(sympy.Matrix(F1)))} nodes") + +t0 = time.time() +sympy.core.cache.clear_cache() +t_cache_clear = time.time() - t0 +print(f"\nsympy.core.cache.clear_cache(): {t_cache_clear:.3f}s") + +# UU block derivatives +t0 = time.time() +G0 = sympy.derive_by_array(F0, stokes.u.sym) +t_uu_g0 = time.time() - t0 + +t0 = time.time() +G1 = sympy.derive_by_array(F0, stokes.Unknowns.L) +t_uu_g1 = time.time() - t0 + +t0 = time.time() +G2 = sympy.derive_by_array(F1, stokes.u.sym) +t_uu_g2 = time.time() - t0 + +t0 = time.time() +G3 = sympy.derive_by_array(F1, stokes.Unknowns.L) +t_uu_g3 = time.time() - t0 + +print(f"\nderive_by_array (UU block):") +print(f" G0 (dF0/dU): {t_uu_g0:.3f}s") +print(f" G1 (dF0/dL): {t_uu_g1:.3f}s") +print(f" G2 (dF1/dU): {t_uu_g2:.3f}s") +print(f" G3 (dF1/dL): {t_uu_g3:.3f}s") +print(f" Total UU: {t_uu_g0+t_uu_g1+t_uu_g2+t_uu_g3:.3f}s") + +# UP block +t0 = time.time() +sympy.derive_by_array(F0, stokes.p.sym) +sympy.derive_by_array(F0, stokes._G) +sympy.derive_by_array(F1, stokes.p.sym) +sympy.derive_by_array(F1, stokes._G) +t_up = time.time() - t0 +print(f"\nderive_by_array (UP block): {t_up:.3f}s") + +# PU block +t0 = time.time() +sympy.derive_by_array(PF0, stokes.u.sym) +sympy.derive_by_array(PF0, stokes.Unknowns.L) +t_pu = time.time() - t0 +print(f"derive_by_array (PU block): {t_pu:.3f}s") + +print(f"\nTotal derivative computation: {t_uu_g0+t_uu_g1+t_uu_g2+t_uu_g3+t_up+t_pu:.3f}s") + +# ── Phase 2: getext (unwrap + hash + compile) ──────────────────────────────── + +# Now let the solver do its full setup and time it +print(f"\n--- Full solver._setup_pointwise_functions + getext ---") +stokes.is_setup = False +stokes.constitutive_model._solver_is_setup = False +stokes.DFDt.psi_fn = stokes.constitutive_model.flux.T + +t0 = time.time() +stokes._setup_pointwise_functions(verbose=True) +t_setup_pw = time.time() - t0 +print(f"_setup_pointwise_functions: {t_setup_pw:.3f}s") + +t0 = time.time() +stokes._setup_discretisation(verbose=True) +t_setup_disc = time.time() - t0 +print(f"_setup_discretisation: {t_setup_disc:.3f}s") + +t0 = time.time() +stokes._setup_solver(verbose=True) +t_setup_solver = time.time() - t0 +print(f"_setup_solver: {t_setup_solver:.3f}s") + +# ── Phase 3: DFDt update + actual solve ─────────────────────────────────────── + +print(f"\n--- DFDt + solve ---") +t0 = time.time() +stokes.DFDt.update_pre_solve(0.02, verbose=True) +t_dfdt_pre = time.time() - t0 +print(f"DFDt.update_pre_solve: {t_dfdt_pre:.3f}s") + +t0 = time.time() +# Call the parent (Stokes) solve directly to skip VE_Stokes overhead +from underworld3.systems.solvers import SNES_Stokes_SaddlePt +SNES_Stokes_SaddlePt.solve(stokes, zero_init_guess=True, _force_setup=False, verbose=True) +t_solve = time.time() - t0 +print(f"PETSc SNES solve: {t_solve:.3f}s") + +t0 = time.time() +stokes.DFDt.update_post_solve(0.02, verbose=True) +t_dfdt_post = time.time() - t0 +print(f"DFDt.update_post_solve: {t_dfdt_post:.3f}s") + +# ── Summary ─────────────────────────────────────────────────────────────────── + +print(f"\n{'='*50}") +print(f"SUMMARY") +print(f"{'='*50}") +print(f"Derivative computation: {t_uu_g0+t_uu_g1+t_uu_g2+t_uu_g3+t_up+t_pu:.1f}s") +print(f"_setup_pointwise_functions: {t_setup_pw:.1f}s (includes derivatives + getext)") +print(f"_setup_discretisation: {t_setup_disc:.1f}s") +print(f"_setup_solver: {t_setup_solver:.1f}s") +print(f"DFDt pre-solve: {t_dfdt_pre:.1f}s") +print(f"PETSc solve: {t_solve:.1f}s") +print(f"DFDt post-solve: {t_dfdt_post:.1f}s")