BufferizationStage IR amplification for runtime-coefficient Hamiltonians under @qjit

Start by running the minimal reproducible example with the stated PennyLane, Catalyst, and JAX versions, then inspect the per-stage IR produced by keep_intermediate=True, especially BufferizationStage and MLIRToLLVMDialectConversion. The issue names no source files or tests and proposes several possible directions; done would require an agreed approach that materially reduces runtime-coefficient IR amplification or provides the proposed pre-compilation size estimate without the reported memory escalation.

Summary

When using qml.TrotterProduct with JAX-traced runtime coefficients inside @qjit, the BufferizationStage of the MLIR compilation pipeline produces a large IR size amplification. This is the dominant contributor to extreme compiler memory consumption (~19 GB compiler-subprocess RSS) for even small molecules.

With fixed (Python float) coefficients, BufferizationStage achieves ~1.0x amplification (constant-folding eliminates redundant IR). With runtime (JAX array) coefficients, the symbolic representation prevents constant-folding; for an H2 reference workload the IR grows from ~3.9 MB to ~60 MB at this single stage. The subsequent MLIRToLLVMDialectConversion then amplifies a further ~4.5x, yielding a total IR growth of ~62x (3.9 MB → 241 MB).

This makes runtime-parameterized Hamiltonians impractical for systems beyond the smallest molecules: H2 (15 Pauli terms) already consumes ~21 GB total peak; H3O+ (193 terms) extrapolates well beyond commodity 32 GB hardware.

The runtime-coefficient code path is the workaround community members are pointed to (cf. #1464) and is what allows electrostatic-embedding workflows (Monte Carlo solvation, VQE-style coefficient updates) to avoid full per-step recompilation. The memory cost at this single stage currently caps how large that workaround can scale.

Related: #1464 (qml.dot with runtime coefficients), #1507 (Hamiltonian aggregate-type lowering)

Minimal reproducible example

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

# --- Hamiltonian construction (numpy path, gives 15-term H2) ---
# qml.qchem.molecular_hamiltonian expects atomic units (Bohr).
# Note: passing the geometry as a JAX array currently disables term aggregation
# in molecular_hamiltonian and yields 47 SProd operands instead of 15. We use
# numpy here to keep the workload aligned with the measurements below.
ANGSTROM_TO_BOHR = 1.8897259886
symbols = ["H", "H"]
coords_angstrom = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.74]])
coords_bohr = (coords_angstrom * ANGSTROM_TO_BOHR).flatten()

H, n_qubits = qml.qchem.molecular_hamiltonian(
    symbols=symbols,
    coordinates=coords_bohr,
    mapping="jordan_wigner",
    active_electrons=2,
    active_orbitals=2,
)
assert len(H.operands) == 15  # H2 STO-3G / JW with active space (2e, 2o)

# coeffs is a JAX array, so it becomes a runtime tracer at compile time —
# this is the configuration that triggers the BufferizationStage amplification.
coeffs = jnp.array([float(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
total_wires = n_qubits + n_est
dev = qml.device("lightning.qubit", wires=total_wires)

@qml.qjit(keep_intermediate=True)  # dump per-stage IR for diagnosis
@qml.qnode(dev)
def qpe_circuit(runtime_coeffs):
    H_rt = qml.dot(runtime_coeffs, ops)  # `ops` captured as compile-time closure
    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,  # required to admit JAX-tracer coeffs
                )
            ),
            control=n_qubits + k,
        )
    return qml.probs(wires=range(n_qubits, total_wires))

# Compilation consumes ~19 GB compiler-subprocess RSS (~21 GB total peak)
# for just 15 Pauli terms. Per-stage IR is written under the keep_intermediate
# directory; BufferizationStage shows the dominant ~9.9x growth.
result = qpe_circuit(coeffs)

Measured data

Environment: PennyLane 0.44.0, Catalyst 0.14.0, JAX 0.7.1, Linux x86_64, 30.4 GB RAM.
Measurement note: numbers below were collected with keep_intermediate=True to obtain per-stage IR sizes; production runs without keep_intermediate show similar end-to-end memory and time.

Per-stage IR size (H2, n_est=4, n_trotter=10, 15 Pauli terms)

End-to-end compilation metrics

System scaling (dynamic mode)

Impact

This is currently a blocking issue for any quantum-chemistry workflow that requires runtime-parameterized Hamiltonians under Catalyst:

1. QM/MM workflows: Monte Carlo solvation dynamics requires updating Hamiltonian coefficients at each step (electrostatic embedding). Without runtime parameterization the alternative is full per-step recompilation (~67 s/step in our setup).
2. VQE / variational methods: any workflow where Hamiltonian coefficients depend on classical parameters.
3. System-size ceiling: even H2 (15 Pauli terms) consumes ~21 GB. H3O+ (193 terms) is infeasible on commodity hardware.

Suggested directions

1. Sparse / structured IR representation: detect that runtime coefficients multiply structurally-identical operator blocks and factor them at BufferizationStage rather than expanding each symbolically.
2. Lazy lowering: defer constant-folding-dependent optimizations to JIT time when actual coefficient values are available.
3. Pre-compilation IR-size estimate: expose an API that estimates IR size before compilation so users can guard against OOM.
4. Operator fusion: combine repeated TrotterProduct blocks that share structure but differ only in the time parameter.