PennyLaneAI / PennyLaneAI/catalyst

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

Open
#2,760 0 comments 0 reactions 0 assignees View on GitHub

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

  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

  • #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_executable tests — confirms replace_ir is debug-only API, supporting our concern that the workaround is fragile.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.