Add Readout Error Mitigation (REM)

[volodymyr-hq]

Context:

Readout Error Mitigation (REM) is a probabilistic QEM technique aimed at mitigating the readout/measurement noise (not to be confused with quantum noise). This is achieved by running two additional calibration circuits: all-zeroes (an empty circuit) and all-ones (a circuit with a Pauli-X gate on each qubit), then using the resulting bitstring distribution after measurement to calculate per-qubit confusion matrices. Finally, these confusion matrices and the resulting bitstring distribution after measurement for the mitigatee circuit are used to solve a linear system and obtain mitigated bitstring distribution. Haiqu's aim as part of the CODE project is to add REM support to the Catalyst compiler.

Description of the Change:

1. Frontend
   We introduced the public mitigate_with_rem wrapper. This function hooks into the Catalyst compilation pipeline by wrapping target QNodes, binding a new rem_p JAX primitive, and routing the resulting user and calibration samples through JAX-compiled post-processing functions. We tried to follow the format/naming conventions established the existing ZNE Mitigation implementation where possible.

2. JAX-Compilable Post-Processing
   Top-level JAX-compilable post-processing functions were added to calculate per-qubit confusion matrices and to solve the linear system to correct noisy samples. A trace-time path dispatch switches from full histograms (bounded by 2^n_qubits x 2^n_qubits) to a sort-based encoding proposed by Nation et al., which is bounded by num_shots x num_shots, when 2^n_qubits exceeds shot count. This optimization allows the program to scale past 25 qubits. Currently, the SampleMP, CountsMP and ProbsMP are supported.

3. Mitigation Dialect and Lowering Pass
   New RemOp (mitigation.rem) operation was added to the Mitigation dialect. The lowering pass rewrites a single calling operation into the original callee alongside two inlined calibration circuits (all-zeros and all-ones states). The quantum.device attributes are forwarded to the calibration circuits, automatically propagating the noise probability to both calibration circuits.

4. Calibration Caching Support
   A boolean runCalibration attribute in the mitigation.rem operation that turns on or off the emission of operations for calibration circuits. When set to False, the lowering pass skips generating calibration circuits, which is useful for per-device confusion matrix caching.

Benefits:

• Extends catalysts support for QEM technique (only ZNE is supported as of now)
• Is scalable for large qubits counts (not bound by exponential growth of 2^num_qubits for SampleMP), which is beneficial for on-hardware execution and large scale simulation.
• Is backend-agnostic
• Research such as Quantum Error Correction on Error-mitigated Physical Qubits by Jeon et al. shows that it is possible to combine QEM techniques with QEC. Since Catalyst already offers QEC dialects, it would make sense to support such pipelines in Catalyst.

Possible Drawbacks:

• No support for observables, since Catalyst does not support injecting any classical post-processing between shot sampling and observable calculation. Instead, each backend is expected implement support for each observable type. In the future, It is possible to circumvent this by implementing each observable in JAX and adding a pass that replaces every observable MP by SampleMP + classical post-processing. This was however out of scope for this project.
• Currently, no supported device in Catalyst has native support for readout noise channels. Therefore, the tests that we added can only test REM on an identity noise channel, as well as separately testing the JAX post-processing functions. We updated the qrack.simulator C++ backend to support simulated readout noise, work is being done to get it merged into pennylane-qrack upstream. For now it is possible to test with our forked version available here: https://github.com/volodymyr-hq/pennylane-qrack/tree/readout_noise_channel. Adding pennylane-qrack as a dependency for the tests seems unreasonable.
• The JAX post-processing functions to calibrate confusion matrices and solve for mitigated results are added at JAX trace-time, not as part of the lower-mitigation pass that replaces mitigation.rem. This breaks the abstraction to some extent, and it would be better if these functions were inserted into the program as part of the lowering pass. However, we did not find a reasonably simple way to implement this, since the post-processing functions rely on linear algebra functions (best represented in StableHLO), which is lowered at an earlier point in the pipeline. We are open to suggestions if there is a better way to do this.
• The returned type after mitigation for SampleMP is not a tensor of all sampled bitstrings, but a tuple of two tensors: one with a sorted list of unique sampled bitstrings, and another with their respective counts. This differs from what is usually expected from a kernel with a SampleMP measurement process, but makes sense for mitigation purposes. Unrolling the counts into samples is not optimal, since that would require probabilistic re-sampling of shots.
• Confusion matrix caching has to be done explicitly by the user.
• This PR modifies ‎mlir/lib/Mitigation/Transforms/CMakeLists.txt to set "-Wno-deprecated-declarations" since Rem.cpp still uses the legacy rewriter.create<OpTy>(...) builder, while upstream MLIR has deprecated it in favour of OpTy::create(rewriter, ...). This can be corrected and Rem.cpp can be ported to the new interface if needed.

[volodymyr-hq]

See also unitaryfoundation/pennylane-qrack#40 for a PR into pennylane-qrack repo with a readout noise channel support. It was used to test REM on simulated noise.