Cache the jitted RHS/Jacobian so solve() does not recompile every call

[timlichtenberg]

Summary

The JAX RHS and Jacobian are rebuilt and re-wrapped in fresh jax.jit objects on every solve() call, so the in-memory trace/compile cache never hits and each solve pays a full re-trace (and, without the persistent compilation cache, a full XLA re-compile). In a coupled PROTEUS run this shows up as

[ INFO  ] JAX RHS first call (JIT compile complete)
[ INFO  ] JAX Jacobian first call (JIT compile complete)

on every solve rather than once per process. It surfaced while profiling PROTEUS #706 (Super-Earth structure and thermal evolution), where a lot of wall-time in an Aragog+AGNI+MORS run is spent in this initialisation stage. This is aragog-side and independent of that PR.

Mechanism

The factory build_jax_rhs_and_jacobian defines a fresh closure _rhs_nondim and wraps it on each call:

aragog/src/aragog/solver/cvode_jax.py

Lines 200 to 202 in f016737

	rhs_jit = jax.jit(_rhs_nondim)
	# jacrev is fine for square Jacobians; jacfwd would also work
	jac_jit = jax.jit(jax.jacrev(_rhs_nondim, argnums=1))

That closure captures concrete JAX arrays as trace constants (state_scale_jax, rhs_scale_jax, heating_jax, and the EOS/phase/mesh/boundary pytrees; see the scale setup and the args tuple).

The factory is called from inside EntropySolver.solve(), so a new pair of jitted wrappers is created on every solve. JAX keys its in-memory cache on the identity of the wrapped callable plus input avals and captured constants, so a fresh closure object every solve is a guaranteed miss: JAX re-traces (re-runs the Python and the jacrev transform) and re-lowers each time.

In-memory cache versus persistent cache

The inline comment at entropy_solver.py#L2926-L2929 claims the trace cache is hit after the first call within a process lifetime. That holds only when the persistent on-disk XLA compilation cache is enabled, which caches the compiled executable by computation hash independent of the Python object identity. PROTEUS sets that cache on its Slurm grid path via JAX_COMPILATION_CACHE_DIR and JAX_COMPILATION_CACHE_MAX_SIZE, so grid runs amortise the XLA re-compile to a disk load; the per-solve re-trace cost still applies. Without the persistent cache, as in a plain local run, every solve pays a full re-trace and re-compile. The comment should be corrected to reflect this.

The log is not proof of recompilation

The first call (JIT compile complete) messages are gated by the per-factory-instance info dict, which is re-initialised each factory call, so they fire once per solve even when the persistent cache served the executable:

aragog/src/aragog/solver/cvode_jax.py

Lines 204 to 233 in f016737

	info = {
	'rhs_calls': 0,
	'jac_calls': 0,
	'first_rhs_compile_done': False,
	'first_jac_compile_done': False,
	}
	def rhs_fn(t_nd, y_nd, ydot_nd):
	"""scikits.odes RHS function: fills ydot in-place."""
	try:
	result = rhs_jit(float(t_nd), jnp.asarray(y_nd))
	ydot_nd[:] = np.asarray(result)
	info['rhs_calls'] += 1
	if not info['first_rhs_compile_done']:
	info['first_rhs_compile_done'] = True
	logger.info('JAX RHS first call (JIT compile complete)')
	return 0
	except Exception as exc:
	logger.error('JAX RHS failed: %s', exc)
	return 1
	def jacfn(t_nd, y_nd, fy_nd, J, user_data=None):
	"""scikits.odes Jacobian function: fills J in-place."""
	try:
	jac = jac_jit(float(t_nd), jnp.asarray(y_nd))
	J[...] = np.asarray(jac)
	info['jac_calls'] += 1
	if not info['first_jac_compile_done']:
	info['first_jac_compile_done'] = True
	logger.info('JAX Jacobian first call (JIT compile complete)')

The guaranteed cost from the code structure is the re-trace; the re-compile is guaranteed only when the persistent cache is absent. The log alone does not size the problem.

Impact assessment

• Least affected: PROTEUS production grids, which run through the Slurm path and can enable the persistent compilation cache, so they pay the re-trace but not the re-compile.
• Most affected: local single runs (developer machines, notebooks) with no persistent cache, which pay a full re-trace and re-compile on every solve. This is what was measured in the PROTEUS #706 profiling.
• Worsened by solve count: the per-cell step caps introduced on the Super-Earth branch (aragog #20, released as 26.07.04) multiply the number of solve() calls, so they multiply this overhead.
• Correctness: none. This is a pure performance concern; the numerics are unchanged.

Please measure before fixing

Before committing to the fix, measure the per-solve cost so we can size it:

• Time a representative multi-solve loop (for example the Aragog+AGNI+MORS fixed-mesh case) two ways: with the persistent XLA compilation cache enabled, and with it disabled.
• Separate the re-trace cost from the re-compile cost (for example compare wall-time per solve against a build-once prototype, or inspect JAX compile counters / XLA logs), so we know how much comes from tracing versus lowering.
• Report the per-solve overhead and how it scales with the number of solves.

Fix direction

Build the jitted RHS/Jacobian pair once per solver instance and reuse it across solve() calls, rebuilding only when a genuinely static input changes (state shape, core_bc mode, EOS table identity). This is not a plain memoise: the quantities that change between solves are currently captured as trace constants, in particular the nondim scales (which depend on the current entropy magnitude and change every solve) and, in the Zalmoxis-coupled case, mesh_arrays (regenerated on every structure re-solve). Reusing a stale jitted pair would silently apply a stale mesh or stale scaling. The clean approach is to pass the varying quantities (the nondim scales, and for coupled runs the mesh) as traced arguments so the compiled artifact is invariant across solves and the in-memory cache hits on matching shapes. The fixed-mesh case (dummy interior structure) is the simplest win: only the scales vary there, so making the scales arguments is enough.

Related

• Origin of the Option Z JAX CVODE path: aragog #12.
• Step caps that multiply the solve count: aragog #20 (released 26.07.04).
• Earlier JAX-vs-numpy solver work: aragog #15.
• Profiling context: PROTEUS #706. The PROTEUS side has its own separate slowness fix in that PR (a redundant structure re-solve on the coupled path); this issue is the aragog-side contribution and is independent of it, so it does not need to land before #706.