Skip to content

Caching Guide

jaxamg has multiple caching layers with different goals. This page focuses on the two caches users typically configure in scripts.

Overview

  1. Metadata cache (Python, jaxamg/cache.py)

    • Caches metadata, including sparsity/coloring info for operators and MPI-related data.
    • Main goal: avoid recomputing pre-processing work and make JIT usage easier.

  2. AmgX resource cache (C++, _amgx_*)

    • Controlled by JAXAMG_CACHE_SIZE.
    • Caches native AmgX handles (matrix/vector/solver/resources).
    • Main goal: avoid repeated native setup and improve solve throughput.

There is also an internal primitive cache in jaxamg.py used automatically by the library. It usually does not require user tuning, so it is not a focus here.

Metadata cache

with_cache(A, ...)

  • Main entry point for metadata caching: attach optional metadata to A once, then reuse A across repeated solves.
  • A primary goal is to make JIT workflows easier by keeping static/precomputed metadata outside traced solve code.
  • This is object-level metadata attachment, not native AmgX-handle caching.

Before JIT, call cache_coloring(..., dtype=...) for each solve precision. Recompute the colouring if new nonzero entries appear.

When to use each option:

  • coloring=...

    • For callable operators: reuse the sparsity pattern and colouring instead of discovering them on every solve. Pass the result of cache_coloring(...).
    • Discovery traces the operator's jaxpr (Hill & Dalle, 2025), else probes it with basis vectors, at its current values: entries that are zero there are dropped. Reuse the colouring only while that pattern covers every parameter value you use.

  • pattern=...

    • Declare all possible couplings with jaxamg.pattern(rows, cols, shape), including entries that are zero initially but may appear later.
    • Pass the same declaration to new callable instances as parameters change. It covers all precisions used with those operators.
    • The declaration copies its inputs into immutable arrays. pattern= and coloring= are mutually exclusive. Attaching new colouring replaces the old declaration and any discovered colourings.

  • mpi=...

    • This reuses MPI metadata: counts, communicator pointer, config string and exchange plans.
    • Reuse it only with the same communicator, partition layout, and CSR structure (including entry order). Values may change.
    • In practice, pass the result of cache_mpi_metadata(...) into with_cache(...).

  • is_symmetric=True

    • This allows the backward pass to skip transpose-related work for symmetric systems.
    • Set it only when the matrix is truly symmetric and remains symmetric.
    • You can set it directly in with_cache(...).
    • It also lets transpose_nullspace default to nullspace.

  • nullspace=... / transpose_nullspace=...

    • Defaults for the arguments of the same name in jaxamg.solve(...); see Singular systems.
    • In MPI mode also prepare the cached config with cache_mpi_metadata(..., singular=True).

Reusing a declared pattern

Create a pattern once outside JIT and reuse it as coefficients change:

p = jaxamg.pattern(rows, cols, shape)

@jax.jit
def solve_at(coefficients, b):
    op = jaxamg.with_cache(lambda x: apply_operator(coefficients, x), pattern=p)
    return jaxamg.solve(op, b)[0]

Single-process solves reuse the pattern's layout while recomputing matrix values. Reuse the same Pattern object; a new equivalent pattern has its own cache.

If your transpose callback materializes an operator, reuse a separate pattern for it. MPI, sharded solves, and colouring tuples do not use this layout cache.

Supplying column colours

Pass both colors and n_colors to use an existing colouring:

p = jaxamg.pattern(rows, cols, shape, colors=colors, n_colors=n_colors)

colors has one integer per input column. Used columns need colours in [0, n_colors); unused columns may use -1. Columns sharing a row must have different colours. Duplicate entries are merged before validation; invalid colours are rejected, not repaired or renumbered.

Pattern describes the operator's input space, regardless of GPU count: (n_rows, n_columns) locally, (n_local, n_global) for an MPI row block, or (n_local, n_local + n_ghost) for a halo operator. Halo colours only need to be valid within each rank. For operators that communicate on global vectors, supply owned-column colours to global_operator instead.

Native AmgX resource cache

Set with environment variable:

export JAXAMG_CACHE_SIZE=2 # Default is 1

Behavior (two modes):

  • 0: isolated mode (no resource caching)
    • Create/destroy native resources every call.
    • Best for debugging behavior and cache isolation.
  • Positive values (default: 1): cache-enabled mode
    • Reuses native resources through an LRU cache for improved performance.
    • Larger values enable multi-entry reuse when alternating among multiple matrix structures/configs, including cases where the forward pass uses A and the gradient/backward pass uses a structurally different A^T.

MPI cache capacity applies per communicator. Call clear_solver_cache() on every rank together; inconsistent caches cause the solve to fail on all ranks. Shared MPI resources are separate for each communicator, device, and transport mode; finalize() releases them.

Solver setup reuse

When the cache hits (same sparsity structure and config as a previous solve), the matrix values are updated via AMGX_matrix_replace_coefficients and the solver setup is repeated against the new values via AMGX_solver_resetup. resetup reuses the cached AmgX solver/matrix objects, device allocations, and fine-level matrix coloring established during the first solve, avoiding the resource-creation and matrix-upload overhead of a cold start. The AMG hierarchy itself is rebuilt against the new values by default; deeper reuse can be enabled via the structure_reuse_levels AmgX config parameter.

This setup-reuse path is substantially cheaper than a cold start while remaining correct for any change in coefficient values, making cache reuse safe for workloads where coefficients change between solves (including optimization and time-stepping).

For repeated solves where reusing the existing AMG hierarchy is acceptable, pass reuse="hierarchy" to jaxamg.solve(...). This still refreshes the fine-level matrix values, but skips AMGX_solver_resetup on cache hits. It can reduce the per-solve cost further, but the reused hierarchy may be less effective for the new values; if coefficients change enough, convergence may require more iterations or become less robust. If that happens, call jaxamg.clear_solver_cache() to force the next solve to build a fresh hierarchy.

For fixed operators, reuse="matrix" reuses both the matrix and hierarchy. Inside jit, it probes the operator only on cache misses. Declare the operator (a callable) with with_cache(A, pattern=p); operators sharing the pattern object p share the cached matrix. With make_sharded_solver(..., reuse="matrix"), each solver has its own cached matrix instead.

Each matrix uses one cache entry, so increase JAXAMG_CACHE_SIZE when alternating between solves. This mode requires unchanged operator values and does not support differentiation.

Checkpointing (jax.checkpoint)

Every solve updates the solver cache, which reuse="hierarchy" solves read, so by default a solve, single-GPU or sharded, declares a side effect to JAX. Under jax.checkpoint (JAX 0.11.1 or later; older versions reject the effect) JAX then runs each solve once and keeps its output for the backward pass: the solves' memory grows as without checkpointing.

has_side_effect=False (in solve and make_sharded_solver) declares a solve pure. Under jax.checkpoint a single-GPU or sharded solve is then recomputed instead of kept, which bounds that memory on any JAX version, and JAX may drop an unused solve or merge identical ones. Declare it only when no reuse="hierarchy" solve depends on this solve's setup, or when such a solve may use whichever setup is cached: a recomputed reuse="hierarchy" solve uses the currently cached hierarchy, which may differ from the one used in the forward pass, so its result varies within the solve tolerance. Per-rank MPI solves keep mpi4jax's ordered effect either way, which jax.checkpoint rejects (Effects not supported in partial-eval of checkpoint/remat: [OrderedMPIEffect]): they cannot be checkpointed.

Cache inspection

Use jaxamg.get_solver_cache_info() for inspecting current solver cache state, which includes:

  • Current size/capacity for both native caches (single_gpu, mpi)
  • Per-entry summaries (dimensions, mode, config, hashes)
  • isolated_mode flag

Clearing caches and cleanup

import jaxamg

jaxamg.clear_solver_cache()  # Clears C++ AmgX handle cache
jaxamg.finalize()            # Clears caches/resources and tears down native state

When to call clear_solver_cache()

In typical workloads — including optimization with changing coefficients — calling clear_solver_cache() is not required. Cache hits automatically refresh the solver against current values via AMGX_solver_resetup, so correctness is maintained without any user intervention.

Reasons you might still want to call it explicitly:

  • Free GPU memory between unrelated solve series (e.g. before moving on to a problem with a different shape or configuration).
  • Force a fresh AMGX_solver_setup if structure_reuse_levels > 0 is set in the AmgX config and the reused coarsening becomes a poor fit for the new values (not relevant with the default config, where resetup already rebuilds the hierarchy).
  • Debugging or reproducing first-solve behavior.

Sparsity-pattern changes do not require an explicit clear: the cache key includes a structural hash, so a different sparsity pattern produces a cache miss and triggers a full setup automatically.

Notes

  • clear_solver_cache() targets native C++ AmgX resources.
  • Metadata attached via with_cache(...) remains on Python objects until those objects are replaced or discarded.
  • For MPI mode, explicit finalize() during teardown can help avoid shutdown-time resource warnings.