No persistent cross-session @qjit compilation cache — every Python session recompiles from scratch

[yjmaxpayne]

Description

Catalyst's @qjit does not persist compiled artifacts across Python sessions. Each fresh import + @qjit call triggers the full MLIR -> LLVM -> machine code pipeline from scratch, even when the circuit topology (number of qubits, gates, structure) is identical to a previous session.

For our H2 QPE circuit with runtime-parameterized coefficients, this means ~220 seconds of recompilation every time a new Python session is started (or ~306s with profiling overhead). This is particularly painful during:

• Interactive development/debugging cycles
• CI/CD test suites
• Jupyter notebook restarts
• Production workflows that spawn fresh Python processes

Minimal Reproducible Example

import pennylane as qml
from jax import numpy as jnp
import time

# H2 molecule setup
symbols = ["H", "H"]
coords = jnp.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.4]])
H, n_qubits = qml.qchem.molecular_hamiltonian(symbols, coords, mapping="jordan_wigner")
coeffs = jnp.array([op.scalar for op in H.operands if isinstance(op, qml.ops.SProd)])
ops = [op.base for op in H.operands if isinstance(op, qml.ops.SProd)]

n_est = 4
n_trotter = 10
dev = qml.device("lightning.qubit", wires=n_qubits + n_est)

@qml.qjit
@qml.qnode(dev)
def qpe_circuit(runtime_coeffs):
    H_rt = qml.dot(runtime_coeffs, ops)
    for k in range(n_est):
        t = 2 ** (n_est - 1 - k)
        qml.ctrl(
            qml.adjoint(
                qml.TrotterProduct(H_rt, time=t, n=n_trotter, order=2, check_hermitian=False)
            ),
            control=n_qubits + k,
        )
    qml.adjoint(qml.QFT)(wires=range(n_qubits, n_qubits + n_est))
    return qml.probs(wires=range(n_qubits, n_qubits + n_est))

start = time.time()
result = qpe_circuit(coeffs)
elapsed = time.time() - start
print(f"First call (compilation + execution): {elapsed:.1f}s")
# Output: ~220s for H_dynamic, ~7.5s for H_fixed

start = time.time()
result2 = qpe_circuit(coeffs)
elapsed2 = time.time() - start
print(f"Second call (execution only): {elapsed2:.3f}s")
# Output: ~0.015s — compilation is cached within the same session

# But restarting Python and running the same script → another 220s

Measured Data

Environment: PennyLane 0.44.0, Catalyst 0.14.0, JAX 0.7.1

Compilation Times by Mode

Mode	First call (compile)	Subsequent calls	Session restart
H_fixed (constants)	7.5s	0.009s	7.5s again
H_dynamic (runtime params)	220s	0.015s	220s again

Our Workaround: Two-Phase IR Caching

We implemented a two-phase cache using Catalyst's internal replace_ir + jit_compile APIs:

Phase	When	Time	What
A (cache miss)	First session	~289–301s	Subprocess full @qjit(keep_intermediate=True) → save LLVM IR to disk
B (cache hit)	Subsequent sessions	~28.8–61.4s	Load IR from disk → replace_ir() → jit_compile() (skips MLIR pipeline, still does LLVM→machine code)

Net result: 4.9× speedup and 79% peak memory reduction vs full @qjit, but Phase B is still tens of seconds (the cache eliminates MLIR amplification, not LLVM→native codegen).

Cache key includes: molecule name, basis set, active space, QPE parameters, circuit style.
Cache key excludes: coordinates, n_waters, temperature (they don't affect IR topology).

Limitation: This workaround uses internal APIs (catalyst.debug.replace_ir, compiled_circuit.jit_compile) that the Catalyst team has explicitly described as "intended for debugging" (see #1592). Building production caching on debug APIs is fragile.

Impact Assessment

1. Development friction: 220s compilation on every session restart makes interactive development extremely slow.
2. CI cost: Test suites that exercise @qjit circuits pay full compilation cost on every run.
3. Workaround fragility: Our IR caching solution depends on debug-only internal APIs.
4. User adoption barrier: New users who encounter multi-minute startup times may abandon Catalyst.

Related Issues

• [MLIR] Speed up compile times #467 [MLIR] Speed up compile times — addresses single-compile MLIR throughput; orthogonal to this issue, which asks for cross-session reuse so recompilation is avoided entirely.
• Improve replace_ir / compile_executable tests #1592 Improve replace_ir / compile_executable tests — confirms replace_ir is debug-only API, supporting our concern that the workaround is fragile.