Parameter update triggers full JIT recompilation despite JITConstant cache mechanism

[lmoresi]

Summary

Updating a UWexpression parameter value (e.g. V_top.sym = new_value or Parameters.dt_elastic = dt) triggers full JIT recompilation on the next solve, even though the _JITConstant cache mechanism is designed to handle value-only changes via PetscDS constants.

Impact: 89% of wall time spent in JIT recompilation for the VE square-wave benchmark (3.7s actual PETSc solve vs 459s JIT for 99 steps on a 128-element mesh).

Root Cause

Two interacting issues:

1. ExpressionDescriptor.__set__ (line 326-330 of _api_tools.py) unconditionally calls obj._reset() or sets obj.is_setup = False after every .sym update.

2. SolverBaseClass._build() is triggered whenever is_setup = False, and it calls getext() which does full JIT compilation — even when the JIT cache key hasn't changed.

The _JITConstant mechanism in getext() correctly substitutes constant UWexpressions with constants[i] placeholders in the cache key (so the key is stable across value changes). But _build() never checks whether the cache key is the same — it just recompiles unconditionally when is_setup = False.

Profiling Evidence

VE square-wave benchmark (99 steps, 16×8 mesh, BDF-2):

Component	Time (s)	%
getext (JIT)	459	89%
SolverBaseClass._build	458	89%
SNESSolve (PETSc)	3.7	1%
global_evaluate	37	7%

getext called 302 times for 99 steps (~3 per step).
_build called 399 times (~4 per step).

Fix

_build() should check the JIT cache key before calling getext(). If the key matches an existing compiled extension, skip recompilation and only update the PetscDS constants array values.

Alternatively: _reset() could distinguish between "structure changed" (needs recompile) and "value changed" (needs constants update only), and set a flag accordingly.

Reproduction

import underworld3 as uw
from underworld3.function import expression
import sympy

mesh = uw.meshing.StructuredQuadBox(elementRes=(16, 8), ...)
v = uw.discretisation.MeshVariable("U", mesh, mesh.dim, degree=2)
p = uw.discretisation.MeshVariable("P", mesh, 1, degree=1)

stokes = uw.systems.VE_Stokes(mesh, velocityField=v, pressureField=p, order=2)
stokes.constitutive_model = uw.constitutive_models.ViscoElasticPlasticFlowModel
stokes.constitutive_model.Parameters.shear_viscosity_0 = 1.0
stokes.constitutive_model.Parameters.shear_modulus = 1.0

V_top = expression("V_top", 0.5, "Top BC")
stokes.add_dirichlet_bc((V_top, 0.0), "Top")

uw.timing.start()
for step in range(10):
    V_top.sym = sympy.Float(0.5 * (-1)**step)  # Toggle sign
    stokes.constitutive_model.Parameters.dt_elastic = 0.02
    stokes.solve(zero_init_guess=False, timestep=0.02)

uw.timing.print_summary()
# getext will show ~30 calls for 10 steps instead of 1

Related

• Persistent JIT cache: avoid recompilation across sessions #121 (persistent JIT cache) would also help but doesn't fix the per-session recompilation
• The VE solver updates dt_elastic every step (line 1216 of solvers.py), triggering this on every single timestep

[lmoresi]

Fixed by PR #129 (JIT cache keyed on generated C source).