PennyLaneAI / PennyLaneAI/catalyst
No persistent cross-session @qjit compilation cache — every Python session recompiles from scratch
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 234
- Forks
- 84
- Avg merge
- 2d 15h
- Merged PRs (30d)
- 66
Description
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
- Development friction: 220s compilation on every session restart makes interactive development extremely slow.
- CI cost: Test suites that exercise @qjit circuits pay full compilation cost on every run.
- Workaround fragility: Our IR caching solution depends on debug-only internal APIs.
- User adoption barrier: New users who encounter multi-minute startup times may abandon Catalyst.
Related Issues
- #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.
- #1592 Improve
replace_ir/compile_executabletests — confirmsreplace_iris debug-only API, supporting our concern that the workaround is fragile.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start at the @qjit compilation path and compare its in-session caching with the workaround's catalyst.debug.replace_ir and compiled_circuit.jit_compile APIs, plus related issue #1592. Define cache identity and invalidation for reusable artifacts across fresh Python processes; done means an equivalent circuit avoids the full compilation pipeline after restart without incorrect reuse.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100