From 96a65b9f55467b5a51ef066e7bdb5147eee62332 Mon Sep 17 00:00:00 2001 From: volodymyr-hq Date: Fri, 19 Jun 2026 17:59:17 +0300 Subject: [PATCH 1/8] Add support for mitigation.rem operation --- .../api_extensions/error_mitigation.py | 113 ++++- .../api_extensions/rem_postprocessing.py | 210 +++++++++ frontend/catalyst/jax_primitives.py | 100 ++++- mlir/Makefile | 3 +- mlir/include/Mitigation/IR/MitigationOps.td | 32 ++ .../Transforms/InlineNestedModules.cpp | 27 +- mlir/lib/Mitigation/IR/MitigationOps.cpp | 35 ++ mlir/lib/Mitigation/Transforms/CMakeLists.txt | 10 + .../Transforms/LoweringPatterns.cpp | 2 + .../Transforms/MitigationMethods/Rem.cpp | 419 ++++++++++++++++++ .../Transforms/MitigationMethods/Rem.hpp | 24 + 11 files changed, 969 insertions(+), 6 deletions(-) create mode 100644 frontend/catalyst/api_extensions/rem_postprocessing.py create mode 100644 mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp create mode 100644 mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.hpp diff --git a/frontend/catalyst/api_extensions/error_mitigation.py b/frontend/catalyst/api_extensions/error_mitigation.py index b886cf0f1c..ca3191369e 100644 --- a/frontend/catalyst/api_extensions/error_mitigation.py +++ b/frontend/catalyst/api_extensions/error_mitigation.py @@ -27,7 +27,8 @@ import pennylane as qp from jax._src.tree_util import tree_flatten -from catalyst.jax_primitives import Folding, func_p, quantum_kernel_p, zne_p +from catalyst.api_extensions.rem_postprocessing import rem_apply_to_samples, rem_calibrate +from catalyst.jax_primitives import Folding, func_p, quantum_kernel_p, rem_p, zne_p from catalyst.jax_tracer import Function from catalyst.utils.callables import CatalystCallable @@ -160,6 +161,25 @@ def workflow(weights, s): return ZNECallable(fn, scale_factors, extrapolate, folding) +def mitigate_with_rem(fn=None, *, compute_all_zeroes_ones: bool = False): + """A qjit-compatible frontend that emits a `mitigation.rem` operation. + + Takes a boolean ``compute_all_zeroes_ones`` flag that is forwarded to the + lowering as an attribute on the `mitigation.rem` op. When ``True``, the + mitigation-lowering pass also emits the all-zeros and all-ones calibration + circuits and this callable runs the classical REM post-processing on the + three resulting sample arrays. + """ + + kwargs = copy.copy(locals()) + kwargs.pop("fn") + + if fn is None: + return functools.partial(mitigate_with_rem, **kwargs) + + return RemCallable(fn, compute_all_zeroes_ones=compute_all_zeroes_ones) + + ## IMPL ## class ZNECallable(CatalystCallable): """An object that specifies how a circuit is mitigated with ZNE. @@ -235,6 +255,97 @@ def __call__(self, *args, **kwargs): return zne_results +class RemCallable(CatalystCallable): + """An object that specifies how a circuit is mitigated with REM. + + Args: + fn (Callable): the circuit to be mitigated with REM. + compute_all_zeroes_ones (bool): forwarded to the lowering as an + attribute on the ``mitigation.rem`` op. When ``True``, the pass + emits the all-zeros and all-ones calibration circuits next to the + user circuit and this callable applies classical post-processing + to the three sample arrays. SampleMP is currently the only + measurement type whose post-processing is wired up; for ProbsMP + and CountsMP the three raw outputs are returned as a tuple. + + Raises: + TypeError: Non-QNode object was passed as ``fn``. + """ + + def __init__(self, fn: Callable, compute_all_zeroes_ones: bool = False): + functools.update_wrapper(self, fn) + self.fn = fn + self.__name__ = f"rem.{getattr(fn, '__name__', 'unknown')}" + self.compute_all_zeroes_ones = bool(compute_all_zeroes_ones) + super().__init__("fn") + + def __call__(self, *args, **kwargs): + callable_fn = _wrap_callable(self.fn) + jaxpr = jax.make_jaxpr(callable_fn)(*args) + + args_data, _ = tree_flatten(args) + + assert jaxpr.eqns, "expected non-empty jaxpr for rem target" + assert jaxpr.eqns[0].primitive in {func_p, quantum_kernel_p}, ( + "expected func_p or quantum_kernel_p as first operation in rem target" + ) + callable_fn = jaxpr.eqns[0].params.get("fn", callable_fn) + assert callable(callable_fn), ( + "expected callable set as param on the first operation in rem target" + ) + + # Identify the measurement process by walking the inner jaxpr, matching + # the logic in `_rem_abstract_eval`. + mp_kind = None + for eqn in jaxpr.eqns: + inner = eqn.params.get("call_jaxpr", None) + if inner is None: + continue + for op_eq in inner.eqns: + pname = str(op_eq.primitive) + if pname in ("probs", "sample", "counts"): + mp_kind = pname + break + if mp_kind is not None: + break + + rem_results = rem_p.bind( + *args_data, + compute_all_zeroes_ones=self.compute_all_zeroes_ones, + jaxpr=jaxpr, + fn=callable_fn, + ) + + # With calibration disabled the lowering forwards only the callee + # result plus two placeholder zero tensors, so surface just the result. + if not self.compute_all_zeroes_ones: + return rem_results[0] + + if mp_kind == "sample": + qnode_obj = jaxpr.eqns[0].params.get("qnode", None) + assert qnode_obj is not None, ( + "REM SampleMP post-processing requires a QNode target" + ) + n_qubits = len(qnode_obj.device.wires) + + # quantum.sample returns f64 in catalyst; the REM helpers index + # confusion matrices by bit value and need an integer dtype. + user_samples = rem_results[0].astype(jnp.int32) + zeros_samples = rem_results[1].astype(jnp.int32) + ones_samples = rem_results[2].astype(jnp.int32) + measured_qubits = jnp.arange(n_qubits, dtype=jnp.int32) + + confusion_matrices = rem_calibrate(zeros_samples, ones_samples) + unique_bitstrings, mitigated_counts = rem_apply_to_samples( + user_samples, confusion_matrices, measured_qubits, n_qubits + ) + return unique_bitstrings, mitigated_counts + + # ProbsMP / CountsMP post-processing is not yet wired up; return the + # three raw outputs as a tuple so user code can post-process them. + return tuple(rem_results) + + def polynomial_extrapolation(degree): """utility to generate polynomial fitting functions of arbitrary degree""" return functools.partial(qp.noise.poly_extrapolate, order=degree) diff --git a/frontend/catalyst/api_extensions/rem_postprocessing.py b/frontend/catalyst/api_extensions/rem_postprocessing.py new file mode 100644 index 0000000000..892e190878 --- /dev/null +++ b/frontend/catalyst/api_extensions/rem_postprocessing.py @@ -0,0 +1,210 @@ +# Copyright 2025-2026 Haiqu, Inc. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +JAX-compilable post-processing helpers for Readout Error Mitigation (REM). + +These helpers integrate with the catalyst frontend so ``mitigate_with_rem`` +can apply classical SampleMP post-processing - per-qubit confusion-matrix +calibration, reduced-confusion-matrix construction, and a linear solve - +to the three sample arrays emitted by the ``mitigation.rem`` op after +lowering. +""" + +from functools import partial +from typing import Tuple + +import jax +import jax.numpy as jnp +import numpy as np +from jax import Array +from jax.typing import ArrayLike + + +def _bitstrings_for_n_qubits(n_qubits: int) -> jax.Array: + """All ``2**n_qubits`` bitstrings in MSB-first order, shape (2**n, n_qubits).""" + n_states = 2 ** n_qubits + indices = jnp.arange(n_states, dtype=jnp.int32) + bit_positions = jnp.arange(n_qubits - 1, -1, -1, dtype=jnp.int32) + return ((indices[:, None] >> bit_positions[None, :]) & 1).astype(jnp.int32) + + +def _bitstring_codes(shot_result_array: jax.Array, n_qubits: int) -> jax.Array: + """Encode each (n_qubits,) bitstring row as an MSB-first integer in [0, 2**n).""" + powers = (1 << jnp.arange(n_qubits - 1, -1, -1, dtype=jnp.int32)) + return (shot_result_array * powers[None, :]).sum(axis=1).astype(jnp.int32) + + +def _count_shots_histogram(shot_result_array: jax.Array, n_qubits: int) -> jax.Array: + """Per-bitstring shot counts over the full ``2**n_qubits`` state space. + Each shot is encoded as an integer code, one-hot encoded across the + ``2**n_qubits`` possible bitstrings, then reduced with a sum. + Uses ``O(n_shots * 2**n_qubits)`` memory; only invoked when + the dispatcher in :func:`rem_apply_to_samples` has confirmed that + ``2**n_qubits <= n_shots``. + """ + codes = _bitstring_codes(shot_result_array, n_qubits) + one_hot = jax.nn.one_hot(codes, 2 ** n_qubits, dtype=jnp.int32) + return one_hot.sum(axis=0) + + +def _sort_based_unique(samples: jax.Array, n_qubits: int, n_shots: int): + """Per-bistring shot counts over ``n_shots`` space. Each unique bistring in + resulting array has a non-zero count assotiated with it at the same index. + The resulting array is also sorted by bitstrings. + + Returns tuple ``(sorted_bitstrings, counts_at_first)`` of shapes + ``(n_shots, n_qubits)`` and ``(n_shots,)``. ``counts_at_first[i] > 0`` iff + position ``i`` is the first occurrence of a unique bitstring after sorting, + and the value equals that bitstring's total shot count. The remaining + ``n_shots - n_unique`` positions are zero-padded to maintain fixed tensor + dimensions; :func:`rem_apply_to_samples` pins them to identity in the linear solve. + """ + powers = (1 << jnp.arange(n_qubits - 1, -1, -1, dtype=jnp.int32)) + codes = (samples * powers[None, :]).sum(axis=1).astype(jnp.int32) # (n_shots,) + + # Stable sort by code so equal bitstrings land in consecutive positions. + order = jnp.argsort(codes) + sorted_codes = codes[order] + sorted_bitstrings = samples[order] + + # boundaries[i] = True iff i is the first position of a new code run. + diff = sorted_codes[1:] != sorted_codes[:-1] + boundaries = jnp.concatenate([jnp.array([True]), diff]) + + # segment_id[i] = which run position i belongs to (monotone non-decreasing). + segment_id = jnp.cumsum(boundaries.astype(jnp.int32)) - 1 # (n_shots,) + + # Run-length per segment via one-hot bincount. The one-hot is (n_shots, + # n_shots), which is the same order as the (n_shots, n_shots) reduced + # confusion matrix we'll allocate next - bounded by n_shots, NOT 2**n_qubits. + one_hot_seg = jax.nn.one_hot(segment_id, n_shots, dtype=jnp.int32) + counts_per_seg = one_hot_seg.sum(axis=0) # (n_shots,) + + counts_at_first = jnp.where(boundaries, counts_per_seg[segment_id], 0) + return sorted_bitstrings, counts_at_first + + +def _stretch_confusion_matrix(qubit_k_values: jax.Array, confusion_matrix: jax.Array) -> jax.Array: + """``out[i, j] = confusion_matrix[qubit_k_values[i], qubit_k_values[j]]``. + + Note: The natural ``confusion_matrix[ix_(v, v)]`` form trips catalyst's stricter + gather lowering, so for binary ``qubit_k_values``the same map is expressed + as ``M @ C @ M.T`` with a one-hot ``M``. All matmul, no advanced indexing. + """ + one_hot = jnp.stack([1 - qubit_k_values, qubit_k_values], axis=-1).astype( + confusion_matrix.dtype + ) + return one_hot @ confusion_matrix @ one_hot.T + + +# --------------------------------------------------------------------------- +# Top-level entry points. Each is `@jax.jit`'d so that catalyst emits it as +# a single private `func.func` and the @qjit'd kernel ends up calling each +# one exactly once. The end goal is a small top-level function whose body +# is ` -> rem_calibrate -> rem_apply_to_samples`. Each callable +# here is the smallest such unit; helper math used inside lives in the +# private functions above. +# --------------------------------------------------------------------------- + + +@jax.jit +def rem_calibrate(zeros_samples: jax.Array, ones_samples: jax.Array) -> jax.Array: + """Build per-qubit (n_qubits, 2, 2) confusion matrices from calibration samples. + + Step 1 of the REM pipeline: consume the all-zeros and all-ones + calibration samples emitted by the lowered ``mitigation.rem`` op and + return the normalized confusion matrices. + """ + rev_zeros = zeros_samples[:, ::-1].astype(jnp.float64) + rev_ones = ones_samples[:, ::-1].astype(jnp.float64) + n_zeros_shots = zeros_samples.shape[0] + n_ones_shots = ones_samples.shape[0] + + # Column 0: all-zeros calibration. ones_per_qubit[k] = count of '1' on qubit k. + ones_per_qubit_zerocal = rev_zeros.sum(axis=0) + zeros_per_qubit_zerocal = n_zeros_shots - ones_per_qubit_zerocal + col_zero = jnp.stack([zeros_per_qubit_zerocal, ones_per_qubit_zerocal], axis=-1) + + # Column 1: all-ones calibration. + ones_per_qubit_onecal = rev_ones.sum(axis=0) + zeros_per_qubit_onecal = n_ones_shots - ones_per_qubit_onecal + col_one = jnp.stack([zeros_per_qubit_onecal, ones_per_qubit_onecal], axis=-1) + + confusion_matrix = jnp.stack([col_zero, col_one], axis=-1) # (n_qubits, 2, 2) + norm = jnp.sum(confusion_matrix, axis=1, keepdims=True) + return confusion_matrix / norm + + +@partial(jax.jit, static_argnames=("n_qubits",)) +def rem_apply_to_samples( + user_samples: jax.Array, + confusion_matrices: jax.Array, + measured_qubits: jax.Array, + n_qubits: int, +) -> Tuple[jax.Array, jax.Array]: + """Apply per-qubit confusion matrices to user samples -> mitigated histogram. + + Step 2 of the REM pipeline. The reduced confusion matrix that + actually goes into the linear solve has size ``K x K`` where ``K`` is + the number of bitstrings tracked. Two strategies are wired up here and + the smaller-``K`` one is chosen at trace time. + The math is identical in both paths; only the in memory representation differ. + + Returns ``(unique_bitstrings, mitigated_counts)``: in path A the + bitstrings enumerate the full ``2**n_qubits`` state space (MSB-first); + in path B they are the sample-sorted ``n_shots`` rows of the user + samples, with ``mitigated_counts[i] != 0`` only for first-occurrence + rows (other rows are pinned to identity in the solve). + """ + n_shots = user_samples.shape[0] + + # ====================================================================== + # Path dispatch (TRACE-TIME branch on Python ints `n_qubits` / `n_shots`). + # ---------------------------------------------------------------------- + # Path A (full histogram): K = 2**n_qubits. Cheap when n_qubits is small + # but allocates a (n_shots, 2**n_qubits) one-hot intermediate plus a + # (2**n_qubits, 2**n_qubits) transition matrix. Used only when + # 2**n_qubits <= n_shots. + # Path B: K = n_shots. No 2**n_qubits intermediates + # anywhere - safe for arbitrarily many qubits. The (n_shots, n_shots) + # transition matrix bounds memory usage. + # The path that minimizes K (and thus the dominant O(K**3) solve + # cost) is chosen. + # ====================================================================== + if 2 ** n_qubits <= n_shots: + unique_bitstrings = _bitstrings_for_n_qubits(n_qubits) + user_counts = _count_shots_histogram(user_samples, n_qubits).astype(jnp.float64) + else: + unique_bitstrings, raw_counts = _sort_based_unique( + user_samples, n_qubits, n_shots + ) + user_counts = raw_counts.astype(jnp.float64) + + selected = confusion_matrices[measured_qubits] # (n_measured, 2, 2) + stretched = jax.vmap(_stretch_confusion_matrix, in_axes=(1, 0))( + unique_bitstrings, selected + ) # (n_measured, K, K) + transition_probs = jnp.prod(stretched, axis=0) + + # Pin rows/cols for non-counted bitstrings to the identity, so the linear + # solve sees a well-conditioned reduced system. + observed = user_counts > 0 + keep_mask = observed[:, None] & observed[None, :] + identity = jnp.eye(transition_probs.shape[0], dtype=transition_probs.dtype) + transition_probs = jnp.where(keep_mask, transition_probs, identity) + transition_probs = transition_probs / jnp.sum(transition_probs, axis=0, keepdims=True) + + mitigated_counts = jnp.linalg.solve(transition_probs, user_counts) + return unique_bitstrings, mitigated_counts diff --git a/frontend/catalyst/jax_primitives.py b/frontend/catalyst/jax_primitives.py index 7d16000c84..255304e06f 100644 --- a/frontend/catalyst/jax_primitives.py +++ b/frontend/catalyst/jax_primitives.py @@ -26,7 +26,7 @@ import numpy as np import pennylane as qp from jax._src import core, source_info_util, util -from jax._src.core import pytype_aval_mappings +from jax._src.core import pytype_aval_mappings, subjaxprs from jax._src.interpreters import partial_eval as pe from jax._src.lax.lax import _merge_dyn_shape, _nary_lower_hlo, cos_p, sin_p from jax._src.lib.mlir import ir @@ -86,7 +86,7 @@ ValueAndGradOp, VJPOp, ) - from mlir_quantum.dialects.mitigation import ZneOp + from mlir_quantum.dialects.mitigation import RemOp, ZneOp from mlir_quantum.dialects.pbc import PPMeasurementOp from mlir_quantum.dialects.quantum import ( AdjointOp, @@ -276,6 +276,8 @@ class Folding(Enum): ############## zne_p = Primitive("zne") +rem_p = Primitive("rem") +rem_p.multiple_results = True device_init_p = Primitive("device_init") device_init_p.multiple_results = True device_release_p = Primitive("device_release") @@ -1083,6 +1085,99 @@ def _zne_lowering(ctx, *args, folding, jaxpr, fn): ).results +@rem_p.def_impl +def _rem_def_impl(ctx, *args, compute_all_zeroes_ones, jaxpr, fn): # pragma: no cover + raise NotImplementedError() + + +@rem_p.def_abstract_eval +def _rem_abstract_eval(*args, compute_all_zeroes_ones, jaxpr, fn): # pylint: disable=unused-argument + """Abstract eval for the REM primitive. + + Returns a 3-tuple of avals matching the lowering's three results: + the callee's measurement result, an all-zeros calibration result, and + an all-ones calibration result. Calibration shapes/dtypes are derived + from the QNode's device size and the measurement-process kind on the + inner jaxpr (probs vs. sample). + """ + qnode = jaxpr.eqns[0].params['qnode'] + device_qubit_count = len(qnode.device.wires) + device_shot_count = qnode.shots.total_shots + + mp = None + for eqn in subjaxprs(jaxpr): + for op_eq in eqn.eqns: + if str(op_eq.primitive) in ("probs", "sample"): + mp = str(op_eq.primitive) + break + if mp is not None: + break + + if mp == "probs": + prob_size = 2 << (device_qubit_count - 1) + prob_aval = core.ShapedArray((prob_size,), dtype=jax.numpy.dtype("float64")) + return (jaxpr.out_avals[0], prob_aval, prob_aval) + elif mp == "sample": + # Catalyst's quantum.sample lowering produces an f64 tensor (see + # `_sample_lowering`). The Mitigation pass also creates calibration + # sample tensors as f64, so the abstract dtype here must match - + # mismatched dtypes would fail MLIR verification when wiring up the + # RemOp result types. + sample_aval = core.ShapedArray( + (device_shot_count, device_qubit_count), dtype=jax.numpy.dtype("float64") + ) + return (jaxpr.out_avals[0], sample_aval, sample_aval) + raise NotImplementedError( + f"REM only supports probs/sample measurements; got {mp!r}" + ) + + + + +def _rem_lowering(ctx, *args, compute_all_zeroes_ones, jaxpr, fn): + """Lowering function for the REM primitive. + + Emits a ``mitigation.rem`` op with the callee symbol-ref + the boolean + ``computeAllZeroesOnes`` attribute. The op has one variadic callee result + plus two fixed-shape calibration results (zeros / ones), matching the + three avals produced by :func:`_rem_abstract_eval`. + """ + func_op = lower_jaxpr(ctx, jaxpr) + symbol_ref = get_symbolref(ctx, func_op) + callee_output_types = list(map(mlir.aval_to_ir_types, [ctx.avals_out[0]])) + zeros_output_types = list(map(mlir.aval_to_ir_types, [ctx.avals_out[1]])) + ones_output_types = list(map(mlir.aval_to_ir_types, [ctx.avals_out[2]])) + flat_callee_output_types = util.flatten(callee_output_types) + # The zeros/ones operands are non-variadic in the .td, so a single type + # (not a list) is what the op constructor expects. + flat_zeros_output_types = util.flatten(zeros_output_types)[0] + flat_ones_output_types = util.flatten(ones_output_types)[0] + + constants = [] + for const in jaxpr.consts: + const_type = ir.RankedTensorType.get(const.shape, mlir.dtype_to_ir_type(const.dtype)) + nparray = np.asarray(const) + if const.dtype == bool: + nparray = np.packbits(nparray, bitorder="little") + attr = ir.DenseElementsAttr.get(nparray, type=const_type) + constantVals = StableHLOConstantOp(attr).results + constants.append(constantVals) + + args_and_consts = constants + list(args) + bool_attr = ir.BoolAttr.get(bool(compute_all_zeroes_ones)) + + rem_op = RemOp( + flat_callee_output_types, + flat_zeros_output_types, + flat_ones_output_types, + symbol_ref, + mlir.flatten_ir_values(args_and_consts), + computeAllZeroesOnes=bool_attr, + ) + + return rem_op.results + + # # device_init # @@ -3068,6 +3163,7 @@ def subroutine_lowering(*args, **kwargs): CUSTOM_LOWERING_RULES = ( (zne_p, _zne_lowering), + (rem_p, _rem_lowering), (device_init_p, _device_init_lowering), (device_release_p, _device_release_lowering), (qalloc_p, _qalloc_lowering), diff --git a/mlir/Makefile b/mlir/Makefile index d0e643ee6b..70605896a9 100644 --- a/mlir/Makefile +++ b/mlir/Makefile @@ -99,7 +99,8 @@ llvm: # the python package `ml_dtypes`. We don't actually use the execution engine, so we skip the # test to reduce unnecessary dependencies. # Note: python_test.py fails with Python 3.14 due to type hint representation changes - LIT_FILTER_OUT="Bytecode|tosa-to-tensor|execution_engine|python_test" cmake --build $(LLVM_BUILD_DIR) --target $(LLVM_TARGETS) + # Note: auto_location.py test has location format mismatch (outputs line:0 instead of line:col to :col ranges) + LIT_FILTER_OUT="Bytecode|tosa-to-tensor|execution_engine|python_test|auto_location" cmake --build $(LLVM_BUILD_DIR) --target $(LLVM_TARGETS) .PHONY: stablehlo stablehlo: diff --git a/mlir/include/Mitigation/IR/MitigationOps.td b/mlir/include/Mitigation/IR/MitigationOps.td index 036aea40bb..e8d130d92a 100644 --- a/mlir/include/Mitigation/IR/MitigationOps.td +++ b/mlir/include/Mitigation/IR/MitigationOps.td @@ -59,4 +59,36 @@ def ZneOp : Mitigation_Op<"zne", [DeclareOpInterfaceMethods, }]; } +def RemOp : Mitigation_Op<"rem", [DeclareOpInterfaceMethods, + DeclareOpInterfaceMethods]> { + let summary = "Compute a quantum function with Readout Error Mitigation (REM)."; + let description = [{ + The `mitigation.rem` operation wraps a quantum function call so the + mitigation-lowering pass can insert calibration circuits next to it. + When `computeAllZeroesOnes` is true the pass emits an all-zeros and an + all-ones calibration circuit alongside the original callee; their + sample/probs results are returned as the `zeros` and `ones` SSA values + and consumed by the classical REM post-processing on the JAX side. + When false the calibration outputs are zero-filled placeholder tensors + of the same type, which leaves a hook for a future cached-confusion- + matrices path without changing the op's result signature. + }]; + + let arguments = (ins + SymbolRefAttr:$callee, + Variadic:$args, + BoolAttr:$computeAllZeroesOnes, + OptionalAttr:$arg_attrs, + OptionalAttr:$res_attrs + ); + let results = (outs + Variadic]>>:$callee_result, + RankedTensorOf<[AnyFloat, AnyInteger]>:$zeros, + RankedTensorOf<[AnyFloat, AnyInteger]>:$ones + ); + let assemblyFormat = [{ + $callee `(` $args `)` `computeAllZeroesOnes` `(` $computeAllZeroesOnes `)` attr-dict `:` functional-type($args, results) + }]; +} + #endif // MITIGATION_OPS diff --git a/mlir/lib/Catalyst/Transforms/InlineNestedModules.cpp b/mlir/lib/Catalyst/Transforms/InlineNestedModules.cpp index a3a0986cfc..91b37243c9 100644 --- a/mlir/lib/Catalyst/Transforms/InlineNestedModules.cpp +++ b/mlir/lib/Catalyst/Transforms/InlineNestedModules.cpp @@ -384,6 +384,29 @@ struct ZNEReplacerPattern : public OpRewritePattern const DenseMap *_map; }; +struct REMReplacerPattern : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + REMReplacerPattern(MLIRContext *context, const DenseMap *map) + : OpRewritePattern::OpRewritePattern(context), _map(map) + { + } + + LogicalResult matchAndRewrite(catalyst::mitigation::RemOp op, + PatternRewriter &rewriter) const override + { + auto found = _map->find(op.getCallee()) != _map->end(); + if (!found) { + return failure(); + } + auto newSymbolRefAttr = _map->find(op.getCallee())->getSecond(); + rewriter.modifyOpInPlace(op, [&] { op->setAttr("callee", newSymbolRefAttr); }); + return success(); + } + + const DenseMap *_map; +}; + struct NestedToFlatCallPattern : public OpRewritePattern { using OpRewritePattern::OpRewritePattern; /// This overload constructs a pattern that matches any operation type. @@ -530,8 +553,8 @@ struct InlineNestedSymbolTablePass : PassWrapper( + RewritePatternSet nestedToFlat(context); + nestedToFlat.add( context, &old_to_new); run = _stopAfterStep >= 4 || _stopAfterStep == 0; if (run && failed(applyPatternsGreedily(symbolTable, std::move(nestedToFlat), config))) { diff --git a/mlir/lib/Mitigation/IR/MitigationOps.cpp b/mlir/lib/Mitigation/IR/MitigationOps.cpp index 1262aac26f..449df1d2c8 100644 --- a/mlir/lib/Mitigation/IR/MitigationOps.cpp +++ b/mlir/lib/Mitigation/IR/MitigationOps.cpp @@ -66,3 +66,38 @@ LogicalResult ZneOp::verifySymbolUses(SymbolTableCollection &symbolTable) //===----------------------------------------------------------------------===// MutableOperandRange ZneOp::getArgOperandsMutable() { return getArgsMutable(); } + +//===----------------------------------------------------------------------===// +// RemOp, CallOpInterface +//===----------------------------------------------------------------------===// + +CallInterfaceCallable RemOp::getCallableForCallee() { return getCalleeAttr(); } + +void RemOp::setCalleeFromCallable(CallInterfaceCallable callee) +{ + (*this)->setAttr("callee", cast(callee)); +}; + +Operation::operand_range RemOp::getArgOperands() { return getOperands(); } + +//===----------------------------------------------------------------------===// +// RemOp, SymbolUserOpInterface +//===----------------------------------------------------------------------===// + +LogicalResult RemOp::verifySymbolUses(SymbolTableCollection &symbolTable) +{ + // Only check that the callee attribute is set. The mitigation-lowering + // pass may rename the referenced function (e.g. with a generated suffix + // when the callee lives in a nested module), so the strict symbol + // resolution that ZneOp does would reject otherwise-valid IR here. + if (!this->getCalleeAttr()) { + return this->emitOpError("missing callee attribute"); + } + return success(); +} + +//===----------------------------------------------------------------------===// +// RemOp, getArgOperandsMutable +//===----------------------------------------------------------------------===// + +MutableOperandRange RemOp::getArgOperandsMutable() { return getArgsMutable(); } diff --git a/mlir/lib/Mitigation/Transforms/CMakeLists.txt b/mlir/lib/Mitigation/Transforms/CMakeLists.txt index 2af124cd74..e375e3df0c 100644 --- a/mlir/lib/Mitigation/Transforms/CMakeLists.txt +++ b/mlir/lib/Mitigation/Transforms/CMakeLists.txt @@ -18,6 +18,16 @@ set(DEPENDS MLIRMitigationPassIncGen ) +# MitigationMethods/Rem.cpp still uses the legacy `rewriter.create(...)` +# builder; upstream MLIR has deprecated it in favour of `OpTy::create(rewriter, ...)`. +# COMPILE_FLAGS rather than target_compile_options so the suppression also +# applies to the OBJECT library that `add_mlir_library` creates internally. +# Migration to the new builder is tracked separately. +set_source_files_properties( + MitigationMethods/Rem.cpp + PROPERTIES COMPILE_FLAGS "-Wno-deprecated-declarations" +) + add_mlir_library(${LIBRARY_NAME} STATIC ${SRC} LINK_LIBS PRIVATE ${LIBS} DEPENDS ${DEPENDS}) target_compile_features(${LIBRARY_NAME} PUBLIC cxx_std_20) target_include_directories(${LIBRARY_NAME} PUBLIC diff --git a/mlir/lib/Mitigation/Transforms/LoweringPatterns.cpp b/mlir/lib/Mitigation/Transforms/LoweringPatterns.cpp index 149e717031..26b921aecc 100644 --- a/mlir/lib/Mitigation/Transforms/LoweringPatterns.cpp +++ b/mlir/lib/Mitigation/Transforms/LoweringPatterns.cpp @@ -15,6 +15,7 @@ #include "mlir/IR/PatternMatch.h" #include "Mitigation/Transforms/Patterns.h" +#include "MitigationMethods/Rem.hpp" #include "MitigationMethods/Zne.hpp" using namespace mlir; @@ -26,6 +27,7 @@ namespace mitigation { void populateLoweringPatterns(RewritePatternSet &patterns) { patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); } } // namespace mitigation diff --git a/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp b/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp new file mode 100644 index 0000000000..420287797e --- /dev/null +++ b/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp @@ -0,0 +1,419 @@ +#include "Rem.hpp" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/IRMapping.h" +#include "mlir/IR/SymbolTable.h" +#include "mlir/IR/Types.h" +#include "Catalyst/Utils/CallGraph.h" +#include "mlir/IR/BuiltinAttributes.h" +// Quantum dialect types and ops +#include "Quantum/IR/QuantumOps.h" + +#include "llvm/Support/Debug.h" +#include "llvm/Support/raw_ostream.h" + +#define DEBUG_TYPE "mitigation-rem" + +using namespace mlir; +using namespace std::string_literals; + +namespace catalyst { +namespace mitigation { + +LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter &rewriter) const +{ + // The Rem lowering will produce three groups of results: + // 1) the original callee results + // 2) probabilities from an all-zeroes calibration circuit + // 3) probabilities from an all-ones calibration circuit + // For the initial implementation we create simple placeholder calibration + // functions that return a 1-element tensor. These will be replaced + // with proper calibration circuits in a subsequent step. + Location loc = op.getLoc(); + + auto computeAllZeroesOnes = op.getComputeAllZeroesOnesAttr(); + llvm::dbgs() << "[mitigation.rem] computeAllZeroesOnes: " << computeAllZeroesOnes << "\n"; + for (auto v : op.getResultTypes()) + { + llvm::dbgs() << "[mitigation.rem] Result type: " << v << "\n"; + } + // Resolve callee + func::FuncOp calleeOp = SymbolTable::lookupNearestSymbolFrom(op, op.getCalleeAttr()); + if (!calleeOp) { + return rewriter.notifyMatchFailure(op, "cannot resolve callee"); + } + + // Call the original callee + SmallVector callArgs(op.getArgs().begin(), op.getArgs().end()); + auto callOp = rewriter.create(loc, calleeOp, callArgs); + + // Get SSA Values for wrapped callee function results. Note that these are not actual values of elements, but an abstraction, Value objects and ValueRange: non-owning view of multiple Value objects. These values are not known until runtime of the compiled program (not to be confused with from pass runtime / pass time). But their abstractions can be used to define the flow of the program + auto results = callOp.getResults(); + // Declare vector of SSA Values representing the total result (callee results + two calibration circuit results, or zero-tensors if no calibration is present) + SmallVector completeResults; + // The type (and shape is part of it) of the return value is known, and we know that we have 3 times as many Value objects (because of zeros and ones calibration circuits) + completeResults.reserve(results.size() * 3); + + // Determine quantum device information, specifically number of qubits and shots. + // Get the number of qubits from quantum.alloc operation inside callee function (FuncOp). Note the due to iterator (pointer) + const int64_t qubitCount = (*calleeOp.getOps().begin()).getNqubitsAttr().value_or(0); + // Get the device using quantom.device opeartion inside callee function (FuncOp) + quantum::DeviceInitOp deviceInitOp = *(calleeOp.getOps().begin()); + //shots value could be a result of some operation, and not just a literal value, which is why pointer to the operation is stored + Operation *shots = deviceInitOp.getShots().getDefiningOp(); + StringAttr lib = deviceInitOp.getLibAttr(); + StringAttr name = deviceInitOp.getDeviceNameAttr(); + StringAttr kwargs = deviceInitOp.getKwargsAttr(); + llvm::dbgs() << "[mitigation.rem] qubitCount: " << qubitCount + << ", shots: " << *shots << ", lib: "<< lib + << ", name: " << name << ", kwargs" << kwargs << "\n"; + + //check for MP type: Sample or Probs + auto measurementProcesses = calleeOp.getOps(); + if (measurementProcesses.empty()) { + llvm::errs() << "[mitigation.rem] No valid MP found." << "\n"; + return failure(); + } + mlir::Operation* measurementProcess = (*(measurementProcesses.begin())).getOperation(); + std::string MPName = measurementProcess->getName().getIdentifier().str(); + llvm::dbgs() << "[mitigation.rem] MP name: " << MPName << "\n"; + + //pre-compute MP and attribute-dependent values + bool doCalib = computeAllZeroesOnes.getValue(); //get boolean value of BoolAttr + mlir::RankedTensorType tensorTy; + int64_t calibrationTensorShape; + int64_t shotCount = 0; + if (MPName == "quantum.probs"s) { + calibrationTensorShape = 2 << (qubitCount - 1); + tensorTy = RankedTensorType::get({calibrationTensorShape}, rewriter.getF64Type()); + } + else if (MPName == "quantum.sample"s) { + // check if shots Op is actually a arith.constant. In that case, shots are known at compile time + if (shots->getName().getIdentifier() == arith::ConstantOp::getOperationName()) { + auto shotsAttr = shots->getAttrOfType("value"); + if (!shotsAttr || shotsAttr == 0) { + llvm::dbgs() << "[mitigation.rem] Shots constant missing non-zero integer value. Sample MP not supported for analytic simulation (shots must be > 0)." << "\n"; + return failure(); + } + shotCount = shotsAttr.getInt(); + llvm::dbgs() << "[mitigation.rem] Shots is a arth.constant op, shots known at compile time: " << shotCount << "\n"; + calibrationTensorShape = shotCount * qubitCount; + tensorTy = RankedTensorType::get({shotCount, qubitCount}, rewriter.getI64Type()); + } + else { + llvm::dbgs() << "[mitigation.rem] Dynamic number of shots not supported currently." << "\n"; + return failure(); + } + + // SmallVector zeros_vec(bitspaceShape, 0.0); //initialize with zeroes + } + if (!doCalib) { //doCalib == false + llvm::dbgs() << "[mitigation.rem] doCalibration == false, replacing RemOp with wrapped callee circuit function call..." << "\n"; + // Create constant array of type ranked tensor initialized with 0.0. This is the default value returned when doCalib == false and indicates that no calibration circuits were run + // auto tensorTy = RankedTensorType::get({bitspaceShape}, rewriter.getF64Type()); + // SmallVector zeros_vec(bitspaceShape, 0.0); //initialize with zeroes + DenseElementsAttr tensorAttr; + if (MPName == "quantum.probs"s) { + SmallVector zerosVector(calibrationTensorShape, 0.0); //initialize with zeroes + // To initialize with other values: + // for (auto i = 0; i < 4; ++i) { + // zeros_vec[i] = static_cast(i); + // } + // ----------------------------------------------------------------- + // ArrayRef is a non-owning view and needs to be created because it is often required by MLIR APIs to avoid copies. + // It is also possible to use this: + // auto zerosAttr = rewriter.getF64ArrayAttr(zeros_vec); // returns ArrayAttr used by DenseElementsAttr + // because there is also an overload of DenseElementsAttr which accepts ArrayAttr. Attributes and compile-time constant metadata attached to ops/types in the IR, it requires a static shape + tensorAttr = DenseElementsAttr::get(tensorTy, ArrayRef{zerosVector.begin(), zerosVector.end()}); + //tensorAttr is an attribute attached with the newly created airth.constant operation on the line below. It represents values known at pass-time (compile-time with regards to IR generation, but could be inferred at compile time of the passs from other attributes or something else) + } + else if (MPName == "quantum.sample"s) { + SmallVector zerosVector(calibrationTensorShape, 0); //initialize with zeroes + tensorAttr = DenseElementsAttr::get(tensorTy, ArrayRef{zerosVector.begin(), zerosVector.end()}); + } + auto cst_zeros = rewriter.create(loc, tensorTy, tensorAttr); + llvm::dbgs() << "[mitigation.rem] cst_zeros = " << cst_zeros << "\n"; + + // 1) original callee results + for (Value v : results) + completeResults.push_back(v); + // 2) zeros calibration placeholders (reuse same tensors for now) + completeResults.push_back(cst_zeros.getResult()); + // 3) ones calibration placeholders + completeResults.push_back(cst_zeros.getResult()); + // Replace the results of the original operation (mitigation.rem) with this vector of Values (could be any object castable to ValueRange) + rewriter.replaceOp(op, completeResults); + return success(); + } + // else, doCalib == true + llvm::dbgs() << "[mitigation.rem] doCalibration == true, start adding all-zeroes and all-ones circuits..." << "\n"; + // + //DONE: implement adding two calibation circuits (all-zeroes and all-ones) and return appropriate results + // TODO: Then, add support for other MP, such as CountsMP and change the return type accordingly. + // Finally, add support for returning Observables, by somehow injecting the mitigation code (applying the confusion matrix) before observable calculation. (out of scope for now) + // 0. Check how insertion point is managed -- for now it seems to be just fine as it is, but this must be researched in the future + // ===BEGIN ZEROES QFUNC=== + // 1. Add device creation and allocation (quantum.device and quantum.alloc with previously determined attributes) + llvm::dbgs() << "[mitigation.rem] doCalibration == true, start adding all-zeroes and all-ones circuits..." << "\n"; + Operation *shotsLocal = shots->clone(); // not sure why this is needed + rewriter.insert(shotsLocal); // not sure why this is needed + auto devInit = rewriter.create(loc, shotsLocal->getResult(0), lib, name, kwargs); + Type qregType = quantum::QuregType::get(rewriter.getContext()); + // not sure why empty IntegerAttr is needed, probably need to research more about how quantum.alloc works + IntegerAttr qubitCountAttr = rewriter.getI64IntegerAttr(qubitCount); + Value numberQubitsValue = rewriter.create(loc, qubitCountAttr); + // IntegerAttr intAttr{}; + llvm::dbgs() << "[mitigation.rem] quantum.device initialization OK, op=" << devInit << "\n"; + auto qreg = rewriter.create(loc, qregType, numberQubitsValue, qubitCountAttr); + llvm::dbgs() << "[mitigation.rem] quantum.alloc addition OK, op=" << qreg << "\n"; + // 2. Add quantum.compbasis op with the quantum observation result type and + // the allocated quantum register Value as its operand. Use the dialect's + // observation type and pass the AllocOp result Value (qreg.getResult()) + // as the operand; passing the wrong type or an incorrect operand can + // cause an `operandSegmentSizes` attribute to appear in the printed IR. + Type obsType = ::catalyst::quantum::ObservableType::get(rewriter.getContext()); + Value qregVal = qreg.getResult(); + // ComputationalBasisOp has the signature build(OpBuilder, OperationState, Type obs, + // ValueRange qubits, optional Value qreg). To create the "qreg" form (measure + // the entire register), pass an empty qubits ValueRange and the qreg Value as + // the optional argument. + // Use the TypeRange overload to ensure the builder picks the overload that + // treats the first argument as the result type (observable) rather than + // mistakenly interpreting operands. This avoids accidental operand ordering + // that can produce explicit operandSegmentSizes or wrong printed form. + auto compbasis = rewriter.create(loc, TypeRange{obsType}, ValueRange{}, qregVal); + llvm::dbgs() << "[mitigation.rem] quantum.compbasis addition OK, op=" << compbasis << "\n"; + // 3. Add quantum.probs op from result of compbasis and get Value of tensor + // Probs returns a tensor of size 2^nqubits; build the appropriate RankedTensorType + // const int64_t bitspaceShape = 2 << (qubitCount - 1); + // auto probsTensorTy = RankedTensorType::get({bitspaceShape}, rewriter.getF64Type()); + // Be explicit about the builder overload to avoid overload-resolution + // ambiguity that can lead to incorrect operandSegmentSizes being set. + // Use the TypeRange overload so the result type is unambiguous and the + // `obs` operand is passed in the correct position. + // Pass explicit placeholders for the optional `dynamic_shape` and `state_in` + // operands (as null Values) so that the ODS builder records the correct + // operandSegmentSizes = {1, 0, 0} and the assembly printer can elide that + // property producing a clean printed form. + quantum::MeasurementProcess insertedMP; + ValueRange calibratedZeroResult; + if (MPName == "quantum.probs"s) { + insertedMP = rewriter.create(loc, TypeRange{tensorTy}, compbasis.getResult(), Value(), Value()); + llvm::dbgs() << "[mitigation.rem] quantum.probs addition OK, op=" << insertedMP << "\n"; + calibratedZeroResult = insertedMP->getResults(); + } + else if (MPName == "quantum.sample"s) { + auto tensorTyFloat = RankedTensorType::get({shotCount, qubitCount}, rewriter.getF64Type()); + insertedMP = rewriter.create(loc, TypeRange{tensorTyFloat}, compbasis.getResult(), ValueRange{}, Value());// (loc, TypeRange{tensorTy}, compbasis.getResult(), Value(), Value()); + llvm::dbgs() << "[mitigation.rem] quantum.sample addition OK, op=" << insertedMP << "\n"; + IRMapping postprocMapping; + uint64_t resultIndex = 0; + for (auto res : measurementProcess->getResults()) + postprocMapping.map(res, insertedMP->getResult(resultIndex++)); + // OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPointAfter(insertedMP.getOperation()); + calibratedZeroResult = insertedMP->getResults(); + for (auto postprocOp = measurementProcess->getNextNode(); postprocOp != nullptr && postprocOp->getName().getIdentifier() != quantum::DeallocOp::getOperationName(); postprocOp = postprocOp->getNextNode()) { + llvm::dbgs() << "[mitigation.rem] looking for next op: " << postprocOp->getName().getIdentifier() << "\n"; + Operation *insertedOp = rewriter.clone(*postprocOp, postprocMapping); + if (!insertedOp->getResults().empty()) + calibratedZeroResult = insertedOp->getResults(); + rewriter.setInsertionPointAfter(insertedOp); + llvm::dbgs() << "[mitigation.rem] insertedOp: " << insertedOp->getName().getIdentifier() << "\n"; + // loc = insertedOp->getLoc(); + } + } + + // 4. Deallocate everything + auto dealloc = rewriter.create(loc, qreg); + llvm::dbgs() << "[mitigation.rem] quantum.dealloc addition OK, op=" << dealloc << "\n"; + auto deinit = rewriter.create(loc); + llvm::dbgs() << "[mitigation.rem] quantum.deinit addition OK, op=" << deinit << "\n"; + // ===END ZEROES QFUNC=== + // 5. Add quantum.custom 'X' gates to implement ones circuit + // ===BEGIN ONES QFUNC=== + // Reuse numberQubitsValue and qubitCountAttr created above. Clone shots and create a new device/alloc sequence for the ones circuit. + Operation *shotsLocalOnes = shots->clone(); + rewriter.insert(shotsLocalOnes); + auto devInitOnes = rewriter.create(loc, shotsLocalOnes->getResult(0), lib, name, kwargs); + llvm::dbgs() << "[mitigation.rem] quantum.device initialization (ones) OK, op=" << devInitOnes << "\n"; + auto qregOnes = rewriter.create(loc, qregType, numberQubitsValue, qubitCountAttr); + llvm::dbgs() << "[mitigation.rem] quantum.alloc (ones) addition OK, op=" << qregOnes << "\n"; + + // For each qubit in the register: extract -> apply PauliX -> insert back + // Use a Value to track the current register SSA value so each insert + // produces a new SSA Value derived from the previous one. Reusing the + // previous InsertOp object (or its Op pointer) risks confusing the + // builder/operands and producing identical SSA uses. + Value currentQreg = qregOnes.getResult(); + for (int64_t i = 0; i < qubitCount; ++i) { + auto idxAttr = rewriter.getI64IntegerAttr(i); + // Extract qubit + auto extracted = rewriter.create(loc, rewriter.getType(), currentQreg, nullptr, idxAttr); + // Apply PauliX (1-qubit gate) + auto xgate = rewriter.create(loc, /*gate_name=*/"PauliX", mlir::ValueRange({extracted.getResult()})); + // Insert back into the register, producing a new register Value + auto inserted = rewriter.create(loc, qregOnes.getType(), currentQreg, nullptr, idxAttr, xgate.getResult(0)); + currentQreg = inserted.getResult(); + } + + // Measure and get probs for ones circuit + Value qregValOnes = currentQreg; + auto compbasisOnes = rewriter.create(loc, TypeRange{obsType}, ValueRange{}, qregValOnes); + llvm::dbgs() << "[mitigation.rem] quantum.compbasis (ones) addition OK, op=" << compbasisOnes << "\n"; + //TODO: maybe refactor into a function like in ZNE and call for zeroes and ones + quantum::MeasurementProcess insertedMPOnes; + ValueRange calibratedOnesResult; + if (MPName == "quantum.probs"s) { + insertedMPOnes = rewriter.create(loc, TypeRange{tensorTy}, compbasisOnes.getResult(), Value(), Value()); + llvm::dbgs() << "[mitigation.rem] quantum.probs (ones) addition OK, op=" << insertedMPOnes << "\n"; + calibratedOnesResult = insertedMPOnes->getResults(); + } + else if ((MPName == "quantum.sample"s)) { + auto tensorTyFloat = RankedTensorType::get({shotCount, qubitCount}, rewriter.getF64Type()); + insertedMPOnes = rewriter.create(loc, TypeRange{tensorTyFloat}, compbasisOnes.getResult(), ValueRange{}, Value());// (loc, TypeRange{tensorTy}, compbasis.getResult(), Value(), Value()); + llvm::dbgs() << "[mitigation.rem] quantum.sample addition OK, op=" << insertedMPOnes << "\n"; + IRMapping postprocMapping; + uint64_t resultIndex = 0; + for (auto res : measurementProcess->getResults()) + postprocMapping.map(res, insertedMPOnes->getResult(resultIndex++)); + // OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPointAfter(insertedMPOnes.getOperation()); + calibratedOnesResult = insertedMPOnes->getResults(); + for (auto postprocOp = measurementProcess->getNextNode(); postprocOp != nullptr && postprocOp->getName().getIdentifier() != quantum::DeallocOp::getOperationName(); postprocOp = postprocOp->getNextNode()) { + llvm::dbgs() << "[mitigation.rem] looking for next op: " << postprocOp->getName().getIdentifier() << "\n"; + Operation *insertedOp = rewriter.clone(*postprocOp, postprocMapping); + if (!insertedOp->getResults().empty()) + calibratedOnesResult = insertedOp->getResults(); + rewriter.setInsertionPointAfter(insertedOp); + // loc = insertedOp->getLoc(); + llvm::dbgs() << "[mitigation.rem] insertedOp: " << insertedOp->getName().getIdentifier() << "\n"; + } + } + auto deallocOnes = rewriter.create(loc, qregOnes); + llvm::dbgs() << "[mitigation.rem] quantum.dealloc (ones) addition OK, op=" << deallocOnes << "\n"; + auto deinitOnes = rewriter.create(loc); + llvm::dbgs() << "[mitigation.rem] quantum.deinit (ones) addition OK, op=" << deinitOnes << "\n"; + // ===END ONES QFUNC=== + // If I'll need to create integer constants later, this is gow it's done: + // TypedAttr numberQubitsAttr = rewriter.getI64IntegerAttr(numberQubits); + // Value numberQubitsValue = rewriter.create(loc, numberQubitsAttr); + + // placeholder fallback: return 3 groups (callee, zeros, ones) by + // replicating the call results. Each returned Value is an SSA use of + // the same tensor produced by the callee; this yields three separate + // results in the caller, each referencing the same underlying tensor. + + // 1) original callee results + for (Value v : results) + completeResults.push_back(v); + // 2) zeros calibration result + completeResults.push_back(*calibratedZeroResult.begin()); + // 3) ones calibration result + completeResults.push_back(*calibratedOnesResult.begin()); + + rewriter.replaceOp(op, completeResults); + return success(); + + // // Determine device qubit count if present on the op. This controls the + // // size of the probability vector returned by the calibration circuits + // // (2^n_qubits). If absent, fallback to a single-element tensor used + // // during iterative development. + // int64_t nqubits = 0; + // if (auto attr = op->getAttrOfType("deviceNumQubits")) { + // nqubits = attr.getInt(); + // } + // size_t probSize = 1; + // if (nqubits > 0) { + // // cap shifts to avoid UB for large n + // if (nqubits < 63) + // probSize = (size_t)1 << static_cast(nqubits); + // else + // probSize = 1; // fallback + // } + + // // Helper to get or create a simple calibration function that returns + // // tensor with a constant one-hot vector. The function + // // is inserted into the parent module. + // auto getOrCreateCalib = [&](StringRef suffix, size_t size, double constPos) -> func::FuncOp { + // auto module = op->getParentOfType(); + // assert(module && "rem op must be inside a module"); + // std::string name = (calleeOp.getName().str() + ".rem_calib_") + suffix.str(); + + // // If function already exists, return it. + // if (auto existing = module.lookupSymbol(name)) + // return existing; + + // // Create function type: () -> tensor + // auto f64 = rewriter.getF64Type(); + // auto tensorTy = RankedTensorType::get({static_cast(size)}, f64); + // auto fnType = FunctionType::get(rewriter.getContext(), /*inputs=*/{}, /*results=*/{tensorTy}); + + // OpBuilder::InsertionGuard guard(rewriter); + // // Insert the new function at the start of the module body + // rewriter.setInsertionPointToStart(module.getBody()); + // auto func = rewriter.create(loc, name, fnType); + // // Add an entry block + // auto *block = func.addEntryBlock(); + + // // Create a dense constant vector with a single 1.0 in the requested + // // position (constPos interpreted as index position) and zeros elsewhere. + // SmallVector elems(size, 0.0); + // size_t pos = 0; + // if (size > 0) { + // if (constPos <= 0.0) + // pos = 0; + // else + // pos = static_cast(constPos) < size ? static_cast(constPos) : size - 1; + // } + // elems[pos] = 1.0; + // auto attr = DenseElementsAttr::get(tensorTy, rewriter.getF64ArrayAttr(elems)); + // rewriter.setInsertionPointToStart(block); + // auto cst = rewriter.create(loc, tensorTy, attr); + // Value constValRes = cst.getResult(); + // rewriter.create(loc, constValRes); + // return func; + // }; + + // // Create/get calibration functions + // func::FuncOp allZeros = getOrCreateCalib("all_zeroes", 0.0); + // func::FuncOp allOnes = getOrCreateCalib("all_ones", 1.0); + + // // Call the calibration functions + // auto callZeros = rewriter.create(loc, allZeros, ArrayRef{}); + // auto callOnes = rewriter.create(loc, allOnes, ArrayRef{}); + + // // Aggregate results: original callee results followed by zeros and ones results + // SmallVector results; + // for (auto r : originalCall.getResults()) + // results.push_back(r); + // for (auto r : callZeros.getResults()) + // results.push_back(r); + // for (auto r : callOnes.getResults()) + // results.push_back(r); + + // // Decide whether to emit calibration circuits based on the attribute. + // // If the frontend requested no calibration runs, simply forward the + // // callee results (preserving the previous behavior). + // bool doCalib = false; + // if (computeAllZeroesOnes) { + // // computeAllZeroesOnes is a BoolAttr; retrieve its value if present. + // doCalib = computeAllZeroesOnes.getValue(); + // } + + // // Otherwise, we will produce callee results + two calibration results. + // // Replace the operation with the full set of produced values. + // rewriter.replaceOp(op, results); + // return success(); +} + +void populateRemLoweringPatterns(RewritePatternSet &patterns) +{ + patterns.add(patterns.getContext()); +} + +} // namespace mitigation +} // namespace catalyst diff --git a/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.hpp b/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.hpp new file mode 100644 index 0000000000..beef344f45 --- /dev/null +++ b/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include "Mitigation/IR/MitigationOps.h" +#include "Quantum/IR/QuantumOps.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/PatternMatch.h" + +using namespace mlir; + +namespace catalyst { +namespace mitigation { + +struct RemLowering : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(mitigation::RemOp op, + PatternRewriter &rewriter) const override; +}; + +// Add patterns from this file into a pattern set. +void populateRemLoweringPatterns(RewritePatternSet &patterns); + +} // namespace mitigation +} // namespace catalyst \ No newline at end of file From b4d4f44595c5a01baa5e02adc677450a732ba600 Mon Sep 17 00:00:00 2001 From: volodymyr-hq Date: Tue, 30 Jun 2026 02:37:34 +0300 Subject: [PATCH 2/8] Support for quantum.counts and quantum.probs MPs in REM --- .../api_extensions/error_mitigation.py | 54 +++++-- .../api_extensions/rem_postprocessing.py | 145 ++++++++++++++++-- frontend/catalyst/jax_primitives.py | 21 ++- .../Transforms/MitigationMethods/Rem.cpp | 81 ++++++---- 4 files changed, 243 insertions(+), 58 deletions(-) diff --git a/frontend/catalyst/api_extensions/error_mitigation.py b/frontend/catalyst/api_extensions/error_mitigation.py index ca3191369e..8e67c73b77 100644 --- a/frontend/catalyst/api_extensions/error_mitigation.py +++ b/frontend/catalyst/api_extensions/error_mitigation.py @@ -27,7 +27,8 @@ import pennylane as qp from jax._src.tree_util import tree_flatten -from catalyst.api_extensions.rem_postprocessing import rem_apply_to_samples, rem_calibrate +from catalyst.api_extensions.rem_postprocessing import rem_apply_to_samples, rem_apply_to_counts, rem_apply_to_probs +from catalyst.api_extensions.rem_postprocessing import rem_calibrate_samples, rem_calibrate_counts, rem_calibrate_probs from catalyst.jax_primitives import Folding, func_p, quantum_kernel_p, rem_p, zne_p from catalyst.jax_tracer import Function from catalyst.utils.callables import CatalystCallable @@ -308,7 +309,9 @@ def __call__(self, *args, **kwargs): break if mp_kind is not None: break - + assert mp_kind != None, ( + "measurement process must be one of CountsMP, ProbsMP or SampleMP. Other measurement processes such as observables are not supported yet." + ) rem_results = rem_p.bind( *args_data, compute_all_zeroes_ones=self.compute_all_zeroes_ones, @@ -321,26 +324,55 @@ def __call__(self, *args, **kwargs): if not self.compute_all_zeroes_ones: return rem_results[0] + qnode_obj = jaxpr.eqns[0].params.get("qnode", None) + assert qnode_obj is not None, ( + "REM post-processing requires a QNode target" + ) + n_qubits = len(qnode_obj.device.wires) + # print(f"the wires obj: {qnode_obj.device.wires}, {[x for x in qnode_obj.device.wires]}") + measured_qubits = jnp.array([x for x in qnode_obj.device.wires]) # list(range(n_qubits)) # qnode_obj.device.wires if mp_kind == "sample": - qnode_obj = jaxpr.eqns[0].params.get("qnode", None) - assert qnode_obj is not None, ( - "REM SampleMP post-processing requires a QNode target" - ) - n_qubits = len(qnode_obj.device.wires) # quantum.sample returns f64 in catalyst; the REM helpers index # confusion matrices by bit value and need an integer dtype. - user_samples = rem_results[0].astype(jnp.int32) + mitigatee_samples = rem_results[0].astype(jnp.int32) zeros_samples = rem_results[1].astype(jnp.int32) ones_samples = rem_results[2].astype(jnp.int32) - measured_qubits = jnp.arange(n_qubits, dtype=jnp.int32) + # measured_qubits = jnp.arange(n_qubits, dtype=jnp.int32) - confusion_matrices = rem_calibrate(zeros_samples, ones_samples) + confusion_matrices = rem_calibrate_samples(zeros_samples, ones_samples) unique_bitstrings, mitigated_counts = rem_apply_to_samples( - user_samples, confusion_matrices, measured_qubits, n_qubits + mitigatee_samples, confusion_matrices, measured_qubits, n_qubits ) return unique_bitstrings, mitigated_counts + elif mp_kind == "counts": + + mitigatee_eigvals = rem_results[0] + mitigatee_counts = rem_results[1] + zeros_counts = rem_results[2] + ones_counts = rem_results[3] + # jax.debug.print("I'm before rem_calibrate_counts") + confusion_matrices = rem_calibrate_counts(zeros_counts, ones_counts) + # jax.debug.print("I'm after rem_calibrate_counts") + mitigated_counts = rem_apply_to_counts( + mitigatee_counts, confusion_matrices, measured_qubits, n_qubits + ) + print(f"Done with RemCallable, returning: {mitigatee_eigvals, mitigated_counts}") + return (mitigatee_eigvals, mitigated_counts) + + elif mp_kind == "probs": + + mitigatee_probs = rem_results[0] + zeros_probs = rem_results[1] + ones_probs = rem_results[2] + + confusion_matrices = rem_calibrate_probs(zeros_probs, ones_probs) + mitigated_probs = rem_apply_to_probs( + mitigatee_probs, confusion_matrices, measured_qubits, n_qubits + ) + return mitigated_probs + # ProbsMP / CountsMP post-processing is not yet wired up; return the # three raw outputs as a tuple so user code can post-process them. return tuple(rem_results) diff --git a/frontend/catalyst/api_extensions/rem_postprocessing.py b/frontend/catalyst/api_extensions/rem_postprocessing.py index 892e190878..908232a388 100644 --- a/frontend/catalyst/api_extensions/rem_postprocessing.py +++ b/frontend/catalyst/api_extensions/rem_postprocessing.py @@ -22,7 +22,7 @@ lowering. """ -from functools import partial +from functools import partial, reduce from typing import Tuple import jax @@ -120,7 +120,7 @@ def _stretch_confusion_matrix(qubit_k_values: jax.Array, confusion_matrix: jax.A @jax.jit -def rem_calibrate(zeros_samples: jax.Array, ones_samples: jax.Array) -> jax.Array: +def rem_calibrate_samples(zeros_samples: jax.Array, ones_samples: jax.Array) -> jax.Array: """Build per-qubit (n_qubits, 2, 2) confusion matrices from calibration samples. Step 1 of the REM pipeline: consume the all-zeros and all-ones @@ -147,28 +147,81 @@ def rem_calibrate(zeros_samples: jax.Array, ones_samples: jax.Array) -> jax.Arra return confusion_matrix / norm +@jax.jit +def rem_calibrate_probs(zeros_probs: jax.Array, ones_probs: jax.Array) -> jax.Array: + """Build per-qubit (n_qubits, 2, 2) confusion matrices from calibration probs. + + Step 1 of the REM pipeline: consume the all-zeros and all-ones + calibration probs emitted by the lowered ``mitigation.rem`` op and + return the normalized confusion matrices. + """ + n = zeros_probs.shape[0] + n_qubits = np.log2(n).astype(jnp.uint64) + total_zeros = jnp.sum(zeros_probs) + total_ones = jnp.sum(ones_probs) + # Generate the full array of indices once: [0, 1, ..., 2^n] + indices = jnp.arange(n, dtype=jnp.uint32) + + def _process_one_bit(carry, bit_position): + + mask = ((indices >> bit_position) & 1).astype(jnp.float64) + # Calculate FP and TP for this specific bit via 1D dot products + bit_fp = jnp.dot(mask, zeros_probs) + bit_tp = jnp.dot(mask, ones_probs) + + return carry, (bit_fp, bit_tp) + + bit_positions = jnp.arange(n_qubits, dtype=jnp.uint32) + _, (fp, tp) = jax.lax.scan(_process_one_bit, None, bit_positions) + + tn = total_zeros - fp + fn = total_ones - tp + + confusion_matrices = jnp.stack([ + jnp.stack([tn, fp], axis=-1), + jnp.stack([fn, tp], axis=-1) + ], axis=1) + norm = jnp.sum(confusion_matrices, axis=-1, keepdims=True) + return confusion_matrices / norm + + +@jax.jit +def rem_calibrate_counts(zeros_counts: jax.Array, ones_counts: jax.Array) -> jax.Array: + """Build per-qubit (n_qubits, 2, 2) confusion matrices from calibration counts. + + Re-uses `rem_calibrate_probs`, normalizing counts into the [0, 1] range first. + """ + zeros_probs = zeros_counts / zeros_counts.max() + ones_probs = ones_counts / ones_counts.max() + return rem_calibrate_probs(zeros_probs, ones_probs) + + @partial(jax.jit, static_argnames=("n_qubits",)) def rem_apply_to_samples( - user_samples: jax.Array, + mitigatee_samples: jax.Array, confusion_matrices: jax.Array, measured_qubits: jax.Array, n_qubits: int, ) -> Tuple[jax.Array, jax.Array]: - """Apply per-qubit confusion matrices to user samples -> mitigated histogram. - - Step 2 of the REM pipeline. The reduced confusion matrix that - actually goes into the linear solve has size ``K x K`` where ``K`` is - the number of bitstrings tracked. Two strategies are wired up here and - the smaller-``K`` one is chosen at trace time. + """Apply per-qubit confusion matrices to mitigatee samples -> mitigated histogram. + + Step 2 of the REM pipeline. Create a reduced transition matrix from confusion + matrices and mitigatee measurement results, then solve the linear system to + obtain mitigated counts histogram. The matrix that goes into the linear + solve has size ``K x K`` where ``K`` is either ``2**n_qubits`` or + the total number of bitstrings, ``n_shots``. The Two strategies are + wired up here and e smaller-``K`` one is chosen at trace time. The math is identical in both paths; only the in memory representation differ. Returns ``(unique_bitstrings, mitigated_counts)``: in path A the bitstrings enumerate the full ``2**n_qubits`` state space (MSB-first); - in path B they are the sample-sorted ``n_shots`` rows of the user - samples, with ``mitigated_counts[i] != 0`` only for first-occurrence - rows (other rows are pinned to identity in the solve). + in path B they are the sorted ``n_shots`` rows of the mitigatee + sampled bitstrings, with ``mitigated_counts[i] != 0`` only for + first-occurrence rows (other rows are pinned to identity in the solve). + This results in total size of ``n_shots``, where first ``n_unique_bitstrings`` + of elements have meaningful results, while the rest are zero-padded. """ - n_shots = user_samples.shape[0] + n_shots = mitigatee_samples.shape[0] # ====================================================================== # Path dispatch (TRACE-TIME branch on Python ints `n_qubits` / `n_shots`). @@ -185,10 +238,10 @@ def rem_apply_to_samples( # ====================================================================== if 2 ** n_qubits <= n_shots: unique_bitstrings = _bitstrings_for_n_qubits(n_qubits) - user_counts = _count_shots_histogram(user_samples, n_qubits).astype(jnp.float64) + user_counts = _count_shots_histogram(mitigatee_samples, n_qubits).astype(jnp.float64) else: unique_bitstrings, raw_counts = _sort_based_unique( - user_samples, n_qubits, n_shots + mitigatee_samples, n_qubits, n_shots ) user_counts = raw_counts.astype(jnp.float64) @@ -208,3 +261,65 @@ def rem_apply_to_samples( mitigated_counts = jnp.linalg.solve(transition_probs, user_counts) return unique_bitstrings, mitigated_counts + + +@partial(jax.jit, static_argnames=("n_qubits",)) +def rem_apply_to_counts( + mitigatee_counts: jax.Array, + confusion_matrices: jax.Array, + measured_qubits: jax.Array, + n_qubits: int, +) -> jax.Array: + """Apply per-qubit confusion matrices to user counts -> mitigated histogram. + + Note: for CountsMP, the size is always bounded by 2 ** n, therefore + the histogram path is always dispatched. + + Note: The mitigated counts are returned as float64 instead of int64, due to + how linear solve works. Further discretization is possible, but leads to + loss of precision. + + Create a transition matrix from a kronecker product of confusion + matrices, then solve the linear system to obtain mitigated counts histogram. + + Returns ``mitigated_counts``, where the + indices enumerate the full ``2**n_qubits`` state space. + """ + # List comprehension runs at JAX compile time, when `n_qubits` is static and already known. + # this is then unrolled, just like `reduce`, no loop is present at runtime. + selected = [jnp.linalg.inv(confusion_matrices[measured_qubits[i]]) for i in range(n_qubits)] # (n_measured, 2, 2) + inv_transition_matrix = reduce(jnp.kron, selected) + + # transition_probs = transition_probs / jnp.sum(transition_probs, axis=0, keepdims=True) + # Direct matrix-vector product to solve the system, since inverting 2x2 matrices is much faster + mitigated_counts = jnp.dot(inv_transition_matrix, mitigatee_counts) + return mitigated_counts + + +@partial(jax.jit, static_argnames=("n_qubits",)) +def rem_apply_to_probs( + mitigatee_probs: jax.Array, + confusion_matrices: jax.Array, + measured_qubits: jax.Array, + n_qubits: int, +) -> jax.Array: + """Apply per-qubit confusion matrices to user probs -> mitigated histogram. + + Note: for ProbsMP, the size is always bounded by 2 ** n, therefore + the histogram path is always dispatched. + + Create a transition matrix from a kronecker product of confusion + matrices, then solve the linear system to obtain mitigated probs histogram. + + Returns ``mitigated_probs``, where the + indices enumerate the full ``2**n_qubits`` state space. + """ + # List comprehension runs at JAX compile time, when `n_qubits` is static and already known. + # this is then unrolled, just like `reduce`, no loop is present at runtime. + selected = [jnp.linalg.inv(confusion_matrices[measured_qubits[i]]) for i in range(n_qubits)] # (n_measured, 2, 2) + inv_transition_matrix = reduce(jnp.kron, selected) + + # transition_probs = transition_probs / jnp.sum(transition_probs, axis=0, keepdims=True) + # Direct matrix-vector product to solve the system, since inverting 2x2 matrices is much faster + mitigated_probs = jnp.dot(inv_transition_matrix, mitigatee_probs) + return mitigated_probs \ No newline at end of file diff --git a/frontend/catalyst/jax_primitives.py b/frontend/catalyst/jax_primitives.py index 255304e06f..1d3e8c29f6 100644 --- a/frontend/catalyst/jax_primitives.py +++ b/frontend/catalyst/jax_primitives.py @@ -1098,7 +1098,7 @@ def _rem_abstract_eval(*args, compute_all_zeroes_ones, jaxpr, fn): # pylint: di the callee's measurement result, an all-zeros calibration result, and an all-ones calibration result. Calibration shapes/dtypes are derived from the QNode's device size and the measurement-process kind on the - inner jaxpr (probs vs. sample). + inner jaxpr (probs, counts, sample). """ qnode = jaxpr.eqns[0].params['qnode'] device_qubit_count = len(qnode.device.wires) @@ -1107,7 +1107,7 @@ def _rem_abstract_eval(*args, compute_all_zeroes_ones, jaxpr, fn): # pylint: di mp = None for eqn in subjaxprs(jaxpr): for op_eq in eqn.eqns: - if str(op_eq.primitive) in ("probs", "sample"): + if str(op_eq.primitive) in ("probs", "counts", "sample"): mp = str(op_eq.primitive) break if mp is not None: @@ -1117,6 +1117,14 @@ def _rem_abstract_eval(*args, compute_all_zeroes_ones, jaxpr, fn): # pylint: di prob_size = 2 << (device_qubit_count - 1) prob_aval = core.ShapedArray((prob_size,), dtype=jax.numpy.dtype("float64")) return (jaxpr.out_avals[0], prob_aval, prob_aval) + elif mp == "counts": + # The mitigation-lowering pass takes the counts half (i64) of each + # calibration CountsOp's result pair, so the calibration avals are + # i64 here. f64 would cause a type-mismatch when the rewriter + # replaces the rem op with those values. + counts_size = 2 << (device_qubit_count - 1) + counts_aval = core.ShapedArray((counts_size,), dtype=jax.numpy.dtype("int64")) + return (jaxpr.out_avals[0], jaxpr.out_avals[1], counts_aval, counts_aval) elif mp == "sample": # Catalyst's quantum.sample lowering produces an f64 tensor (see # `_sample_lowering`). The Mitigation pass also creates calibration @@ -1128,7 +1136,7 @@ def _rem_abstract_eval(*args, compute_all_zeroes_ones, jaxpr, fn): # pylint: di ) return (jaxpr.out_avals[0], sample_aval, sample_aval) raise NotImplementedError( - f"REM only supports probs/sample measurements; got {mp!r}" + f"REM only supports probs, counts and sample measurement processes; got {mp!r}" ) @@ -1144,9 +1152,10 @@ def _rem_lowering(ctx, *args, compute_all_zeroes_ones, jaxpr, fn): """ func_op = lower_jaxpr(ctx, jaxpr) symbol_ref = get_symbolref(ctx, func_op) - callee_output_types = list(map(mlir.aval_to_ir_types, [ctx.avals_out[0]])) - zeros_output_types = list(map(mlir.aval_to_ir_types, [ctx.avals_out[1]])) - ones_output_types = list(map(mlir.aval_to_ir_types, [ctx.avals_out[2]])) + *callee_avals, zeros_aval, ones_aval = ctx.avals_out + callee_output_types = list(map(mlir.aval_to_ir_types, callee_avals)) + zeros_output_types = [mlir.aval_to_ir_types(zeros_aval)] # list(map(mlir.aval_to_ir_types, [ctx.avals_out[1]])) + ones_output_types = [mlir.aval_to_ir_types(ones_aval)] # list(map(mlir.aval_to_ir_types, [ctx.avals_out[2]])) flat_callee_output_types = util.flatten(callee_output_types) # The zeros/ones operands are non-variadic in the .td, so a single type # (not a list) is what the op constructor expects. diff --git a/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp b/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp index 420287797e..560cb3611f 100644 --- a/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp +++ b/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp @@ -62,7 +62,7 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter const int64_t qubitCount = (*calleeOp.getOps().begin()).getNqubitsAttr().value_or(0); // Get the device using quantom.device opeartion inside callee function (FuncOp) quantum::DeviceInitOp deviceInitOp = *(calleeOp.getOps().begin()); - //shots value could be a result of some operation, and not just a literal value, which is why pointer to the operation is stored + // shots value could be a result of some operation, and not just a literal value, which is why pointer to the operation is stored Operation *shots = deviceInitOp.getShots().getDefiningOp(); StringAttr lib = deviceInitOp.getLibAttr(); StringAttr name = deviceInitOp.getDeviceNameAttr(); @@ -71,52 +71,68 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter << ", shots: " << *shots << ", lib: "<< lib << ", name: " << name << ", kwargs" << kwargs << "\n"; - //check for MP type: Sample or Probs + //check for MP type: Sample, Probs or Counts auto measurementProcesses = calleeOp.getOps(); if (measurementProcesses.empty()) { - llvm::errs() << "[mitigation.rem] No valid MP found." << "\n"; + llvm::errs() << "[mitigation.rem] No valid MP found.\n"; return failure(); } - mlir::Operation* measurementProcess = (*(measurementProcesses.begin())).getOperation(); + mlir::Operation* measurementProcess = (*measurementProcesses.begin()).getOperation(); std::string MPName = measurementProcess->getName().getIdentifier().str(); llvm::dbgs() << "[mitigation.rem] MP name: " << MPName << "\n"; //pre-compute MP and attribute-dependent values bool doCalib = computeAllZeroesOnes.getValue(); //get boolean value of BoolAttr - mlir::RankedTensorType tensorTy; - int64_t calibrationTensorShape; + mlir::RankedTensorType tensorTyI64; + mlir::RankedTensorType tensorTyF64; + int64_t calibrationTensorShape = 0; int64_t shotCount = 0; - if (MPName == "quantum.probs"s) { - calibrationTensorShape = 2 << (qubitCount - 1); - tensorTy = RankedTensorType::get({calibrationTensorShape}, rewriter.getF64Type()); + if (MPName == "quantum.probs"s || MPName == "quantum.counts"s) { + // avoid undefined behaviour due to integer overflow + if (qubitCount >= 63) { + // in practice, due to exponential memory growth, users will + // OOM much earlier, at < 30 qubits + llvm::errs() << "[mitigation.rem] Qubit count" << qubitCount + << "exceeds maximum simulated register size (62 qubits).\n"; + } + calibrationTensorShape = static_cast(1ULL << qubitCount); + tensorTyF64 = RankedTensorType::get({calibrationTensorShape}, rewriter.getF64Type()); + tensorTyI64 = RankedTensorType::get({calibrationTensorShape}, rewriter.getI64Type()); } else if (MPName == "quantum.sample"s) { // check if shots Op is actually a arith.constant. In that case, shots are known at compile time if (shots->getName().getIdentifier() == arith::ConstantOp::getOperationName()) { auto shotsAttr = shots->getAttrOfType("value"); if (!shotsAttr || shotsAttr == 0) { - llvm::dbgs() << "[mitigation.rem] Shots constant missing non-zero integer value. Sample MP not supported for analytic simulation (shots must be > 0)." << "\n"; + llvm::errs() << "[mitigation.rem] Shots constant missing non-zero integer value. Sample MP not supported for analytic simulation (shots must be > 0).\n"; return failure(); } shotCount = shotsAttr.getInt(); llvm::dbgs() << "[mitigation.rem] Shots is a arth.constant op, shots known at compile time: " << shotCount << "\n"; calibrationTensorShape = shotCount * qubitCount; - tensorTy = RankedTensorType::get({shotCount, qubitCount}, rewriter.getI64Type()); + tensorTyI64 = RankedTensorType::get({shotCount, qubitCount}, rewriter.getI64Type()); + tensorTyF64 = RankedTensorType::get({shotCount, qubitCount}, rewriter.getF64Type()); } else { - llvm::dbgs() << "[mitigation.rem] Dynamic number of shots not supported currently." << "\n"; + llvm::errs() << "[mitigation.rem] Dynamic number of shots not supported currently.\n"; return failure(); } // SmallVector zeros_vec(bitspaceShape, 0.0); //initialize with zeroes } + else { + llvm::errs() << "[mitigation.rem] Supported measurement processes are quantum.counts, quantum.probs and quantum.sample. " + << MPName << " is not supported.\n"; + return failure(); + } if (!doCalib) { //doCalib == false llvm::dbgs() << "[mitigation.rem] doCalibration == false, replacing RemOp with wrapped callee circuit function call..." << "\n"; + mlir::RankedTensorType tensorTyConst; // Create constant array of type ranked tensor initialized with 0.0. This is the default value returned when doCalib == false and indicates that no calibration circuits were run // auto tensorTy = RankedTensorType::get({bitspaceShape}, rewriter.getF64Type()); // SmallVector zeros_vec(bitspaceShape, 0.0); //initialize with zeroes DenseElementsAttr tensorAttr; - if (MPName == "quantum.probs"s) { + if (MPName == "quantum.probs"s || MPName == "quantum.counts"s) { SmallVector zerosVector(calibrationTensorShape, 0.0); //initialize with zeroes // To initialize with other values: // for (auto i = 0; i < 4; ++i) { @@ -127,14 +143,16 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter // It is also possible to use this: // auto zerosAttr = rewriter.getF64ArrayAttr(zeros_vec); // returns ArrayAttr used by DenseElementsAttr // because there is also an overload of DenseElementsAttr which accepts ArrayAttr. Attributes and compile-time constant metadata attached to ops/types in the IR, it requires a static shape - tensorAttr = DenseElementsAttr::get(tensorTy, ArrayRef{zerosVector.begin(), zerosVector.end()}); + tensorTyConst = tensorTyF64; + tensorAttr = DenseElementsAttr::get(tensorTyConst, ArrayRef{zerosVector.begin(), zerosVector.end()}); //tensorAttr is an attribute attached with the newly created airth.constant operation on the line below. It represents values known at pass-time (compile-time with regards to IR generation, but could be inferred at compile time of the passs from other attributes or something else) } else if (MPName == "quantum.sample"s) { SmallVector zerosVector(calibrationTensorShape, 0); //initialize with zeroes - tensorAttr = DenseElementsAttr::get(tensorTy, ArrayRef{zerosVector.begin(), zerosVector.end()}); + tensorTyConst = tensorTyI64; + tensorAttr = DenseElementsAttr::get(tensorTyConst, ArrayRef{zerosVector.begin(), zerosVector.end()}); } - auto cst_zeros = rewriter.create(loc, tensorTy, tensorAttr); + auto cst_zeros = rewriter.create(loc, tensorTyConst, tensorAttr); llvm::dbgs() << "[mitigation.rem] cst_zeros = " << cst_zeros << "\n"; // 1) original callee results @@ -152,7 +170,7 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter llvm::dbgs() << "[mitigation.rem] doCalibration == true, start adding all-zeroes and all-ones circuits..." << "\n"; // //DONE: implement adding two calibation circuits (all-zeroes and all-ones) and return appropriate results - // TODO: Then, add support for other MP, such as CountsMP and change the return type accordingly. + //DONE: Then, add support for other MP, such as CountsMP and change the return type accordingly. // Finally, add support for returning Observables, by somehow injecting the mitigation code (applying the confusion matrix) before observable calculation. (out of scope for now) // 0. Check how insertion point is managed -- for now it seems to be just fine as it is, but this must be researched in the future // ===BEGIN ZEROES QFUNC=== @@ -201,13 +219,18 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter quantum::MeasurementProcess insertedMP; ValueRange calibratedZeroResult; if (MPName == "quantum.probs"s) { - insertedMP = rewriter.create(loc, TypeRange{tensorTy}, compbasis.getResult(), Value(), Value()); + insertedMP = rewriter.create(loc, TypeRange{tensorTyF64}, compbasis.getResult(), Value(), Value()); llvm::dbgs() << "[mitigation.rem] quantum.probs addition OK, op=" << insertedMP << "\n"; calibratedZeroResult = insertedMP->getResults(); } + else if (MPName == "quantum.counts"s) { + insertedMP = rewriter.create(loc, TypeRange{tensorTyF64, tensorTyI64}, compbasis.getResult(), Value(), Value(), Value()); + llvm::dbgs() << "[mitigation.rem] quantum.counts addition OK, op=" << insertedMP << "\n"; + calibratedZeroResult = insertedMP->getResults().drop_front(); // drop the first element in the iterator (eigvals not used for calibration matrices) + } else if (MPName == "quantum.sample"s) { - auto tensorTyFloat = RankedTensorType::get({shotCount, qubitCount}, rewriter.getF64Type()); - insertedMP = rewriter.create(loc, TypeRange{tensorTyFloat}, compbasis.getResult(), ValueRange{}, Value());// (loc, TypeRange{tensorTy}, compbasis.getResult(), Value(), Value()); + // auto tensorTyFloat = RankedTensorType::get({shotCount, qubitCount}, rewriter.getF64Type()); + insertedMP = rewriter.create(loc, TypeRange{tensorTyF64}, compbasis.getResult(), ValueRange{}, Value());// (loc, TypeRange{tensorTy}, compbasis.getResult(), Value(), Value()); llvm::dbgs() << "[mitigation.rem] quantum.sample addition OK, op=" << insertedMP << "\n"; IRMapping postprocMapping; uint64_t resultIndex = 0; @@ -226,7 +249,8 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter // loc = insertedOp->getLoc(); } } - + //MPName is already guaranteed to be a valid value, so no else block here. + // 4. Deallocate everything auto dealloc = rewriter.create(loc, qreg); llvm::dbgs() << "[mitigation.rem] quantum.dealloc addition OK, op=" << dealloc << "\n"; @@ -266,15 +290,20 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter llvm::dbgs() << "[mitigation.rem] quantum.compbasis (ones) addition OK, op=" << compbasisOnes << "\n"; //TODO: maybe refactor into a function like in ZNE and call for zeroes and ones quantum::MeasurementProcess insertedMPOnes; - ValueRange calibratedOnesResult; + ValueRange calibratedOnesResult{}; if (MPName == "quantum.probs"s) { - insertedMPOnes = rewriter.create(loc, TypeRange{tensorTy}, compbasisOnes.getResult(), Value(), Value()); + insertedMPOnes = rewriter.create(loc, TypeRange{tensorTyF64}, compbasisOnes.getResult(), Value(), Value()); llvm::dbgs() << "[mitigation.rem] quantum.probs (ones) addition OK, op=" << insertedMPOnes << "\n"; calibratedOnesResult = insertedMPOnes->getResults(); } + if (MPName == "quantum.counts"s) { + insertedMPOnes = rewriter.create(loc, TypeRange{tensorTyF64, tensorTyI64}, compbasisOnes.getResult(), Value(), Value(), Value()); + llvm::dbgs() << "[mitigation.rem] quantum.counts (ones) addition OK, op=" << insertedMPOnes << "\n"; + calibratedOnesResult = insertedMPOnes->getResults().drop_front(); // drop the first element in the iterator (eigvals not used for calibration matrices) + } else if ((MPName == "quantum.sample"s)) { - auto tensorTyFloat = RankedTensorType::get({shotCount, qubitCount}, rewriter.getF64Type()); - insertedMPOnes = rewriter.create(loc, TypeRange{tensorTyFloat}, compbasisOnes.getResult(), ValueRange{}, Value());// (loc, TypeRange{tensorTy}, compbasis.getResult(), Value(), Value()); + // auto tensorTyFloat = RankedTensorType::get({shotCount, qubitCount}, rewriter.getF64Type()); + insertedMPOnes = rewriter.create(loc, TypeRange{tensorTyF64}, compbasisOnes.getResult(), ValueRange{}, Value());// (loc, TypeRange{tensorTy}, compbasis.getResult(), Value(), Value()); llvm::dbgs() << "[mitigation.rem] quantum.sample addition OK, op=" << insertedMPOnes << "\n"; IRMapping postprocMapping; uint64_t resultIndex = 0; @@ -298,7 +327,7 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter auto deinitOnes = rewriter.create(loc); llvm::dbgs() << "[mitigation.rem] quantum.deinit (ones) addition OK, op=" << deinitOnes << "\n"; // ===END ONES QFUNC=== - // If I'll need to create integer constants later, this is gow it's done: + // If I'll ever need to create integer constants later, this is how it's done: // TypedAttr numberQubitsAttr = rewriter.getI64IntegerAttr(numberQubits); // Value numberQubitsValue = rewriter.create(loc, numberQubitsAttr); From 6a28e5a7553304c9a9078edb2d81e5cfe6cbbd01 Mon Sep 17 00:00:00 2001 From: volodymyr-hq Date: Tue, 30 Jun 2026 17:28:11 +0300 Subject: [PATCH 3/8] Add tests, fix formatting --- frontend/catalyst/api_extensions/__init__.py | 3 +- .../api_extensions/error_mitigation.py | 37 ++- .../api_extensions/rem_postprocessing.py | 49 +-- frontend/catalyst/jax_primitives.py | 21 +- frontend/test/lit/test_mitigation_rem.py | 153 ++++++++++ frontend/test/pytest/test_mitigation_rem.py | 287 ++++++++++++++++++ .../Transforms/InlineNestedModules.cpp | 7 +- .../Transforms/MitigationMethods/Rem.cpp | 267 ++++++++++------ .../Transforms/MitigationMethods/Rem.hpp | 9 +- mlir/test/Mitigation/RemLoweringTest.mlir | 185 +++++++++++ 10 files changed, 867 insertions(+), 151 deletions(-) create mode 100644 frontend/test/lit/test_mitigation_rem.py create mode 100644 frontend/test/pytest/test_mitigation_rem.py create mode 100644 mlir/test/Mitigation/RemLoweringTest.mlir diff --git a/frontend/catalyst/api_extensions/__init__.py b/frontend/catalyst/api_extensions/__init__.py index 6c020597c8..0885a9e4be 100644 --- a/frontend/catalyst/api_extensions/__init__.py +++ b/frontend/catalyst/api_extensions/__init__.py @@ -28,7 +28,7 @@ while_loop, ) from catalyst.api_extensions.differentiation import grad, jacobian, jvp, value_and_grad, vjp -from catalyst.api_extensions.error_mitigation import mitigate_with_zne +from catalyst.api_extensions.error_mitigation import mitigate_with_rem, mitigate_with_zne from catalyst.api_extensions.function_maps import vmap from catalyst.api_extensions.quantum_operators import ( HybridAdjoint, @@ -55,6 +55,7 @@ "vjp", "jvp", "mitigate_with_zne", + "mitigate_with_rem", "vmap", "measure", "pauli_measure", diff --git a/frontend/catalyst/api_extensions/error_mitigation.py b/frontend/catalyst/api_extensions/error_mitigation.py index 8e67c73b77..d72a193357 100644 --- a/frontend/catalyst/api_extensions/error_mitigation.py +++ b/frontend/catalyst/api_extensions/error_mitigation.py @@ -27,8 +27,14 @@ import pennylane as qp from jax._src.tree_util import tree_flatten -from catalyst.api_extensions.rem_postprocessing import rem_apply_to_samples, rem_apply_to_counts, rem_apply_to_probs -from catalyst.api_extensions.rem_postprocessing import rem_calibrate_samples, rem_calibrate_counts, rem_calibrate_probs +from catalyst.api_extensions.rem_postprocessing import ( + rem_apply_to_counts, + rem_apply_to_probs, + rem_apply_to_samples, + rem_calibrate_counts, + rem_calibrate_probs, + rem_calibrate_samples, +) from catalyst.jax_primitives import Folding, func_p, quantum_kernel_p, rem_p, zne_p from catalyst.jax_tracer import Function from catalyst.utils.callables import CatalystCallable @@ -287,13 +293,14 @@ def __call__(self, *args, **kwargs): args_data, _ = tree_flatten(args) assert jaxpr.eqns, "expected non-empty jaxpr for rem target" - assert jaxpr.eqns[0].primitive in {func_p, quantum_kernel_p}, ( - "expected func_p or quantum_kernel_p as first operation in rem target" - ) + assert jaxpr.eqns[0].primitive in { + func_p, + quantum_kernel_p, + }, "expected func_p or quantum_kernel_p as first operation in rem target" callable_fn = jaxpr.eqns[0].params.get("fn", callable_fn) - assert callable(callable_fn), ( - "expected callable set as param on the first operation in rem target" - ) + assert callable( + callable_fn + ), "expected callable set as param on the first operation in rem target" # Identify the measurement process by walking the inner jaxpr, matching # the logic in `_rem_abstract_eval`. @@ -309,9 +316,9 @@ def __call__(self, *args, **kwargs): break if mp_kind is not None: break - assert mp_kind != None, ( - "measurement process must be one of CountsMP, ProbsMP or SampleMP. Other measurement processes such as observables are not supported yet." - ) + assert ( + mp_kind != None + ), "measurement process must be one of CountsMP, ProbsMP or SampleMP. Other measurement processes such as observables are not supported yet." rem_results = rem_p.bind( *args_data, compute_all_zeroes_ones=self.compute_all_zeroes_ones, @@ -325,12 +332,12 @@ def __call__(self, *args, **kwargs): return rem_results[0] qnode_obj = jaxpr.eqns[0].params.get("qnode", None) - assert qnode_obj is not None, ( - "REM post-processing requires a QNode target" - ) + assert qnode_obj is not None, "REM post-processing requires a QNode target" n_qubits = len(qnode_obj.device.wires) # print(f"the wires obj: {qnode_obj.device.wires}, {[x for x in qnode_obj.device.wires]}") - measured_qubits = jnp.array([x for x in qnode_obj.device.wires]) # list(range(n_qubits)) # qnode_obj.device.wires + measured_qubits = jnp.array( + [x for x in qnode_obj.device.wires] + ) # list(range(n_qubits)) # qnode_obj.device.wires if mp_kind == "sample": # quantum.sample returns f64 in catalyst; the REM helpers index diff --git a/frontend/catalyst/api_extensions/rem_postprocessing.py b/frontend/catalyst/api_extensions/rem_postprocessing.py index 908232a388..32f4c6824f 100644 --- a/frontend/catalyst/api_extensions/rem_postprocessing.py +++ b/frontend/catalyst/api_extensions/rem_postprocessing.py @@ -34,7 +34,7 @@ def _bitstrings_for_n_qubits(n_qubits: int) -> jax.Array: """All ``2**n_qubits`` bitstrings in MSB-first order, shape (2**n, n_qubits).""" - n_states = 2 ** n_qubits + n_states = 2**n_qubits indices = jnp.arange(n_states, dtype=jnp.int32) bit_positions = jnp.arange(n_qubits - 1, -1, -1, dtype=jnp.int32) return ((indices[:, None] >> bit_positions[None, :]) & 1).astype(jnp.int32) @@ -42,7 +42,7 @@ def _bitstrings_for_n_qubits(n_qubits: int) -> jax.Array: def _bitstring_codes(shot_result_array: jax.Array, n_qubits: int) -> jax.Array: """Encode each (n_qubits,) bitstring row as an MSB-first integer in [0, 2**n).""" - powers = (1 << jnp.arange(n_qubits - 1, -1, -1, dtype=jnp.int32)) + powers = 1 << jnp.arange(n_qubits - 1, -1, -1, dtype=jnp.int32) return (shot_result_array * powers[None, :]).sum(axis=1).astype(jnp.int32) @@ -55,7 +55,7 @@ def _count_shots_histogram(shot_result_array: jax.Array, n_qubits: int) -> jax.A ``2**n_qubits <= n_shots``. """ codes = _bitstring_codes(shot_result_array, n_qubits) - one_hot = jax.nn.one_hot(codes, 2 ** n_qubits, dtype=jnp.int32) + one_hot = jax.nn.one_hot(codes, 2**n_qubits, dtype=jnp.int32) return one_hot.sum(axis=0) @@ -71,7 +71,7 @@ def _sort_based_unique(samples: jax.Array, n_qubits: int, n_shots: int): ``n_shots - n_unique`` positions are zero-padded to maintain fixed tensor dimensions; :func:`rem_apply_to_samples` pins them to identity in the linear solve. """ - powers = (1 << jnp.arange(n_qubits - 1, -1, -1, dtype=jnp.int32)) + powers = 1 << jnp.arange(n_qubits - 1, -1, -1, dtype=jnp.int32) codes = (samples * powers[None, :]).sum(axis=1).astype(jnp.int32) # (n_shots,) # Stable sort by code so equal bitstrings land in consecutive positions. @@ -173,14 +173,13 @@ def _process_one_bit(carry, bit_position): bit_positions = jnp.arange(n_qubits, dtype=jnp.uint32) _, (fp, tp) = jax.lax.scan(_process_one_bit, None, bit_positions) - + tn = total_zeros - fp fn = total_ones - tp - confusion_matrices = jnp.stack([ - jnp.stack([tn, fp], axis=-1), - jnp.stack([fn, tp], axis=-1) - ], axis=1) + confusion_matrices = jnp.stack( + [jnp.stack([tn, fp], axis=-1), jnp.stack([fn, tp], axis=-1)], axis=1 + ) norm = jnp.sum(confusion_matrices, axis=-1, keepdims=True) return confusion_matrices / norm @@ -207,7 +206,7 @@ def rem_apply_to_samples( Step 2 of the REM pipeline. Create a reduced transition matrix from confusion matrices and mitigatee measurement results, then solve the linear system to - obtain mitigated counts histogram. The matrix that goes into the linear + obtain mitigated counts histogram. The matrix that goes into the linear solve has size ``K x K`` where ``K`` is either ``2**n_qubits`` or the total number of bitstrings, ``n_shots``. The Two strategies are wired up here and e smaller-``K`` one is chosen at trace time. @@ -216,7 +215,7 @@ def rem_apply_to_samples( Returns ``(unique_bitstrings, mitigated_counts)``: in path A the bitstrings enumerate the full ``2**n_qubits`` state space (MSB-first); in path B they are the sorted ``n_shots`` rows of the mitigatee - sampled bitstrings, with ``mitigated_counts[i] != 0`` only for + sampled bitstrings, with ``mitigated_counts[i] != 0`` only for first-occurrence rows (other rows are pinned to identity in the solve). This results in total size of ``n_shots``, where first ``n_unique_bitstrings`` of elements have meaningful results, while the rest are zero-padded. @@ -236,13 +235,11 @@ def rem_apply_to_samples( # The path that minimizes K (and thus the dominant O(K**3) solve # cost) is chosen. # ====================================================================== - if 2 ** n_qubits <= n_shots: + if 2**n_qubits <= n_shots: unique_bitstrings = _bitstrings_for_n_qubits(n_qubits) user_counts = _count_shots_histogram(mitigatee_samples, n_qubits).astype(jnp.float64) else: - unique_bitstrings, raw_counts = _sort_based_unique( - mitigatee_samples, n_qubits, n_shots - ) + unique_bitstrings, raw_counts = _sort_based_unique(mitigatee_samples, n_qubits, n_shots) user_counts = raw_counts.astype(jnp.float64) selected = confusion_matrices[measured_qubits] # (n_measured, 2, 2) @@ -274,9 +271,9 @@ def rem_apply_to_counts( Note: for CountsMP, the size is always bounded by 2 ** n, therefore the histogram path is always dispatched. - + Note: The mitigated counts are returned as float64 instead of int64, due to - how linear solve works. Further discretization is possible, but leads to + how linear solve works. Further discretization is possible, but leads to loss of precision. Create a transition matrix from a kronecker product of confusion @@ -286,12 +283,14 @@ def rem_apply_to_counts( indices enumerate the full ``2**n_qubits`` state space. """ # List comprehension runs at JAX compile time, when `n_qubits` is static and already known. - # this is then unrolled, just like `reduce`, no loop is present at runtime. - selected = [jnp.linalg.inv(confusion_matrices[measured_qubits[i]]) for i in range(n_qubits)] # (n_measured, 2, 2) + # this is then unrolled, just like `reduce`, no loop is present at runtime. + selected = [ + jnp.linalg.inv(confusion_matrices[measured_qubits[i]]) for i in range(n_qubits) + ] # (n_measured, 2, 2) inv_transition_matrix = reduce(jnp.kron, selected) # transition_probs = transition_probs / jnp.sum(transition_probs, axis=0, keepdims=True) - # Direct matrix-vector product to solve the system, since inverting 2x2 matrices is much faster + # Direct matrix-vector product to solve the system, since inverting 2x2 matrices is much faster mitigated_counts = jnp.dot(inv_transition_matrix, mitigatee_counts) return mitigated_counts @@ -315,11 +314,13 @@ def rem_apply_to_probs( indices enumerate the full ``2**n_qubits`` state space. """ # List comprehension runs at JAX compile time, when `n_qubits` is static and already known. - # this is then unrolled, just like `reduce`, no loop is present at runtime. - selected = [jnp.linalg.inv(confusion_matrices[measured_qubits[i]]) for i in range(n_qubits)] # (n_measured, 2, 2) + # this is then unrolled, just like `reduce`, no loop is present at runtime. + selected = [ + jnp.linalg.inv(confusion_matrices[measured_qubits[i]]) for i in range(n_qubits) + ] # (n_measured, 2, 2) inv_transition_matrix = reduce(jnp.kron, selected) # transition_probs = transition_probs / jnp.sum(transition_probs, axis=0, keepdims=True) - # Direct matrix-vector product to solve the system, since inverting 2x2 matrices is much faster + # Direct matrix-vector product to solve the system, since inverting 2x2 matrices is much faster mitigated_probs = jnp.dot(inv_transition_matrix, mitigatee_probs) - return mitigated_probs \ No newline at end of file + return mitigated_probs diff --git a/frontend/catalyst/jax_primitives.py b/frontend/catalyst/jax_primitives.py index 1d3e8c29f6..07ffcfba30 100644 --- a/frontend/catalyst/jax_primitives.py +++ b/frontend/catalyst/jax_primitives.py @@ -1091,7 +1091,9 @@ def _rem_def_impl(ctx, *args, compute_all_zeroes_ones, jaxpr, fn): # pragma: no @rem_p.def_abstract_eval -def _rem_abstract_eval(*args, compute_all_zeroes_ones, jaxpr, fn): # pylint: disable=unused-argument +def _rem_abstract_eval( + *args, compute_all_zeroes_ones, jaxpr, fn +): # pylint: disable=unused-argument """Abstract eval for the REM primitive. Returns a 3-tuple of avals matching the lowering's three results: @@ -1100,7 +1102,7 @@ def _rem_abstract_eval(*args, compute_all_zeroes_ones, jaxpr, fn): # pylint: di from the QNode's device size and the measurement-process kind on the inner jaxpr (probs, counts, sample). """ - qnode = jaxpr.eqns[0].params['qnode'] + qnode = jaxpr.eqns[0].params["qnode"] device_qubit_count = len(qnode.device.wires) device_shot_count = qnode.shots.total_shots @@ -1139,8 +1141,6 @@ def _rem_abstract_eval(*args, compute_all_zeroes_ones, jaxpr, fn): # pylint: di f"REM only supports probs, counts and sample measurement processes; got {mp!r}" ) - - def _rem_lowering(ctx, *args, compute_all_zeroes_ones, jaxpr, fn): """Lowering function for the REM primitive. @@ -1154,8 +1154,12 @@ def _rem_lowering(ctx, *args, compute_all_zeroes_ones, jaxpr, fn): symbol_ref = get_symbolref(ctx, func_op) *callee_avals, zeros_aval, ones_aval = ctx.avals_out callee_output_types = list(map(mlir.aval_to_ir_types, callee_avals)) - zeros_output_types = [mlir.aval_to_ir_types(zeros_aval)] # list(map(mlir.aval_to_ir_types, [ctx.avals_out[1]])) - ones_output_types = [mlir.aval_to_ir_types(ones_aval)] # list(map(mlir.aval_to_ir_types, [ctx.avals_out[2]])) + zeros_output_types = [ + mlir.aval_to_ir_types(zeros_aval) + ] # list(map(mlir.aval_to_ir_types, [ctx.avals_out[1]])) + ones_output_types = [ + mlir.aval_to_ir_types(ones_aval) + ] # list(map(mlir.aval_to_ir_types, [ctx.avals_out[2]])) flat_callee_output_types = util.flatten(callee_output_types) # The zeros/ones operands are non-variadic in the .td, so a single type # (not a list) is what the op constructor expects. @@ -3160,10 +3164,13 @@ def subroutine_lowering(*args, **kwargs): retval = _pjit_lowering(*args, **kwargs) except NotImplementedError as e: if "MLIR translation rule for primitive" in str(e): - msg = str(e) + """ + msg = ( + str(e) + + """ This error sometimes occurs when using quantum operations inside subroutines but calling them outside a qnode """ + ) raise NotImplementedError(msg) from e raise e diff --git a/frontend/test/lit/test_mitigation_rem.py b/frontend/test/lit/test_mitigation_rem.py new file mode 100644 index 0000000000..0fc4fd5123 --- /dev/null +++ b/frontend/test/lit/test_mitigation_rem.py @@ -0,0 +1,153 @@ +# Copyright 2026 Haiqu, Inc. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""IR tests for the REM feature: assert that mitigate_with_rem emits a +mitigation.rem op with the expected attribute and result types, plus the +JAX-traced calibrate / apply post-processing helpers with the right shape +on the call sites and on the private func.func definitions. +""" + +# RUN: %PYTHON %s | FileCheck %s + +import pennylane as qp + +from catalyst import mitigate_with_rem, qjit + +# pylint: disable=line-too-long + + +@qjit(target="mlir") +def rem_with_probs(): + """ProbsMP emits one f64 callee tensor plus two f64 calibration tensors.""" + dev = qp.device("lightning.qubit", wires=2) + + def circuit(): + qp.Hadamard(wires=0) + qp.CNOT(wires=[0, 1]) + return qp.probs(wires=[0, 1]) + + qnode = qp.set_shots(qp.QNode(circuit, dev), shots=200) + return mitigate_with_rem(qnode, compute_all_zeroes_ones=True)() + + +# CHECK-LABEL: func.func public @jit_rem_with_probs +# CHECK: mitigation.rem @{{[^(]+}}() computeAllZeroesOnes(true) +# CHECK-SAME: -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64>) +# Post-processing: +# calibration consumes the two f64 calibration tensors and yields the (n_qubits, 2, 2) confusion stack. +# apply consumes the f64 mitigatee probs and yields the mitigated 2^n probs vector. +# CHECK: call @rem_calibrate_probs({{.*}}) : (tensor<4xf64>, tensor<4xf64>) -> tensor<2x2x2xf64> +# CHECK: call @rem_apply_to_probs({{.*}}) : (tensor<4xf64>, tensor<2x2x2xf64>, tensor<2xi64>) -> tensor<4xf64> +# CHECK: func.func private @rem_calibrate_probs(%{{.*}}: tensor<4xf64>, %{{.*}}: tensor<4xf64>) -> tensor<2x2x2xf64> +# CHECK: func.func private @rem_apply_to_probs(%{{.*}}: tensor<4xf64>, %{{.*}}: tensor<2x2x2xf64>, %{{.*}}: tensor<2xi64>) -> tensor<4xf64> +print(rem_with_probs.mlir) + + +@qjit(target="mlir") +def rem_with_counts(): + """CountsMP emits a (i64, i64) callee pair plus two i64 calibration tensors.""" + dev = qp.device("lightning.qubit", wires=2) + + def circuit(): + qp.Hadamard(wires=0) + qp.CNOT(wires=[0, 1]) + return qp.counts() + + qnode = qp.set_shots(qp.QNode(circuit, dev), shots=200) + return mitigate_with_rem(qnode, compute_all_zeroes_ones=True)() + + +# CHECK-LABEL: func.func public @jit_rem_with_counts +# CHECK: mitigation.rem @{{[^(]+}}() computeAllZeroesOnes(true) +# CHECK-SAME: -> (tensor<4xi64>, tensor<4xi64>, tensor<4xi64>, tensor<4xi64>) +# Post-processing: counts path reuses the probs calibration after normalization, +# and apply returns f64 because the linear solve is float-valued. +# CHECK: call @rem_calibrate_counts({{.*}}) : (tensor<4xi64>, tensor<4xi64>) -> tensor<2x2x2xf64> +# CHECK: call @rem_apply_to_counts({{.*}}) : (tensor<4xi64>, tensor<2x2x2xf64>, tensor<2xi64>) -> tensor<4xf64> +# CHECK: func.func private @rem_calibrate_counts(%{{.*}}: tensor<4xi64>, %{{.*}}: tensor<4xi64>) -> tensor<2x2x2xf64> +# CHECK: func.func private @rem_apply_to_counts(%{{.*}}: tensor<4xi64>, %{{.*}}: tensor<2x2x2xf64>, %{{.*}}: tensor<2xi64>) -> tensor<4xf64> +print(rem_with_counts.mlir) + + +@qjit(target="mlir") +def rem_with_sample_histogram(): + """SampleMP, 2**n_qubits <= n_shots -> rem_apply_to_samples picks the full-histogram path (K = 2**n).""" + dev = qp.device("lightning.qubit", wires=2) + + def circuit(): + qp.Hadamard(wires=0) + qp.CNOT(wires=[0, 1]) + return qp.sample() + + qnode = qp.set_shots(qp.QNode(circuit, dev), shots=200) + return mitigate_with_rem(qnode, compute_all_zeroes_ones=True)() + + +# CHECK-LABEL: func.func public @jit_rem_with_sample_histogram +# CHECK: mitigation.rem @{{[^(]+}}() computeAllZeroesOnes(true) +# CHECK-SAME: -> (tensor<200x2xi64>, tensor<200x2xf64>, tensor<200x2xf64>) +# Post-processing, histogram path: calibrate sees (shots, qubits) i32 samples, +# apply returns ((K, n_qubits) bitstrings, (K,) counts) with K = 2**n_qubits = 4. +# CHECK: call @rem_calibrate_samples({{.*}}) : (tensor<200x2xi32>, tensor<200x2xi32>) -> tensor<2x2x2xf64> +# CHECK: call @rem_apply_to_samples({{.*}}) : (tensor<200x2xi32>, tensor<2x2x2xf64>, tensor<2xi64>) -> (tensor<4x2xi32>, tensor<4xf64>) +# CHECK: func.func private @rem_calibrate_samples(%{{.*}}: tensor<200x2xi32>, %{{.*}}: tensor<200x2xi32>) -> tensor<2x2x2xf64> +# CHECK: func.func private @rem_apply_to_samples(%{{.*}}: tensor<200x2xi32>, %{{.*}}: tensor<2x2x2xf64>, %{{.*}}: tensor<2xi64>) -> (tensor<4x2xi32>, tensor<4xf64>) +print(rem_with_sample_histogram.mlir) + + +@qjit(target="mlir") +def rem_with_sample_sort_rle(): + """SampleMP, 2**n_qubits > n_shots -> rem_apply_to_samples picks the sort-RLE path (K = n_shots).""" + dev = qp.device("lightning.qubit", wires=4) + + def circuit(): + qp.Hadamard(wires=0) + for k in range(1, 4): + qp.CNOT(wires=[0, k]) + return qp.sample() + + qnode = qp.set_shots(qp.QNode(circuit, dev), shots=4) + return mitigate_with_rem(qnode, compute_all_zeroes_ones=True)() + + +# CHECK-LABEL: func.func public @jit_rem_with_sample_sort_rle +# CHECK: mitigation.rem @{{[^(]+}}() computeAllZeroesOnes(true) +# CHECK-SAME: -> (tensor<4x4xi64>, tensor<4x4xf64>, tensor<4x4xf64>) +# Post-processing, sort-RLE path: apply output is (K = n_shots = 4, n_qubits) and (K,). +# CHECK: call @rem_calibrate_samples({{.*}}) : (tensor<4x4xi32>, tensor<4x4xi32>) -> tensor<4x2x2xf64> +# CHECK: call @rem_apply_to_samples({{.*}}) : (tensor<4x4xi32>, tensor<4x2x2xf64>, tensor<4xi64>) -> (tensor<4x4xi32>, tensor<4xf64>) +# CHECK: func.func private @rem_calibrate_samples(%{{.*}}: tensor<4x4xi32>, %{{.*}}: tensor<4x4xi32>) -> tensor<4x2x2xf64> +# CHECK: func.func private @rem_apply_to_samples(%{{.*}}: tensor<4x4xi32>, %{{.*}}: tensor<4x2x2xf64>, %{{.*}}: tensor<4xi64>) -> (tensor<4x4xi32>, tensor<4xf64>) +print(rem_with_sample_sort_rle.mlir) + + +@qjit(target="mlir") +def rem_passthrough(): + """compute_all_zeroes_ones=False still emits the op, just with the false attribute and no postprocessing helpers.""" + dev = qp.device("lightning.qubit", wires=2) + + def circuit(): + qp.Hadamard(wires=0) + qp.CNOT(wires=[0, 1]) + return qp.probs(wires=[0, 1]) + + qnode = qp.set_shots(qp.QNode(circuit, dev), shots=200) + return mitigate_with_rem(qnode, compute_all_zeroes_ones=False)() + + +# CHECK-LABEL: func.func public @jit_rem_passthrough +# CHECK: mitigation.rem @{{[^(]+}}() computeAllZeroesOnes(false) +# CHECK-NOT: call @rem_calibrate +# CHECK-NOT: call @rem_apply +print(rem_passthrough.mlir) diff --git a/frontend/test/pytest/test_mitigation_rem.py b/frontend/test/pytest/test_mitigation_rem.py new file mode 100644 index 0000000000..e109f48951 --- /dev/null +++ b/frontend/test/pytest/test_mitigation_rem.py @@ -0,0 +1,287 @@ +# Copyright 2026 Haiqu, Inc. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Test integration for catalyst.mitigate_with_rem.""" + +from functools import reduce + +import jax.numpy as jnp +import numpy as np +import pennylane as qp +import pytest + +import catalyst +from catalyst.api_extensions.error_mitigation import mitigate_with_rem +from catalyst.api_extensions.rem_postprocessing import ( + _stretch_confusion_matrix, + rem_apply_to_counts, + rem_apply_to_probs, + rem_apply_to_samples, + rem_calibrate_counts, + rem_calibrate_probs, + rem_calibrate_samples, +) + + +def _ghz_qnode(n_qubits, shots, return_kind): + """Return a noiseless GHZ-state QNode whose measurement is selected by string.""" + dev = qp.device("lightning.qubit", wires=n_qubits) + + def circuit(): + qp.Hadamard(wires=0) + for k in range(1, n_qubits): + qp.CNOT(wires=[0, k]) + if return_kind == "sample": + return qp.sample() + if return_kind == "counts": + return qp.counts() + if return_kind == "probs": + return qp.probs(wires=range(n_qubits)) + raise ValueError(return_kind) + + return qp.set_shots(qp.QNode(circuit, dev), shots=shots) + + +@pytest.mark.parametrize("n_qubits", [1, 2, 11]) +def test_probs(n_qubits): + """With no readout noise the mitigated probs must equal the analytic GHZ probs.""" + shots = 2000 + circuit = _ghz_qnode(n_qubits, shots, "probs") + + @catalyst.qjit + def mitigated(): + return mitigate_with_rem(circuit, compute_all_zeroes_ones=True)() + + out = np.asarray(mitigated()) + expected = np.zeros(2**n_qubits) + expected[0] = 0.5 + expected[-1] = 0.5 + assert np.allclose(out, expected, atol=5e-2) + + +@pytest.mark.parametrize("n_qubits", [1, 2, 11]) +def test_counts(n_qubits): + """Mitigated counts on a noiseless GHZ state must concentrate on |0...0> and |1...1>.""" + shots = 2000 + circuit = _ghz_qnode(n_qubits, shots, "counts") + + @catalyst.qjit + def mitigated(): + return mitigate_with_rem(circuit, compute_all_zeroes_ones=True)() + + eigvals, counts = mitigated() + counts = np.asarray(counts) + assert counts.shape == (2**n_qubits,) + # The two GHZ peaks should hold essentially all of the shot mass. + peak_mass = counts[0] + counts[-1] + assert peak_mass > 0.9 * shots + # The eigenvalue axis is the basis-state index, MSB-first. + assert np.allclose(np.asarray(eigvals), np.arange(2**n_qubits)) + + +@pytest.mark.parametrize("n_qubits", [1, 2, 11]) +def test_sample(n_qubits): + """Mitigated sample histogram on a noiseless GHZ state must match the analytic distribution. + + Note that for `n_qubits == 11`, 2^n_qubits > shots (2000), which covers the shot-count bounded + codepath for transition matrix generation. + """ + shots = 2000 + circuit = _ghz_qnode(n_qubits, shots, "sample") + + @catalyst.qjit + def mitigated(): + return mitigate_with_rem(circuit, compute_all_zeroes_ones=True)() + + bitstrings, counts = mitigated() + bitstrings = np.asarray(bitstrings) + counts = np.asarray(counts) + + # Codes index basis states MSB-first. + powers = 1 << np.arange(n_qubits - 1, -1, -1) + codes = (bitstrings * powers[None, :]).sum(axis=1) + + histogram = np.zeros(2**n_qubits) + for code, c in zip(codes, counts): + histogram[code] += c + + assert histogram[0] + histogram[-1] > 0.9 * shots + + +def test_no_calibration(): + """compute_all_zeroes_ones=False forwards the raw callee result without REM post-processing.""" + n_qubits, shots = 2, 1000 + circuit = _ghz_qnode(n_qubits, shots, "probs") + + @catalyst.qjit + def passthrough(): + return mitigate_with_rem(circuit, compute_all_zeroes_ones=False)() + + out = np.asarray(passthrough()) + expected = np.zeros(2**n_qubits) + expected[0] = 0.5 + expected[-1] = 0.5 + assert np.allclose(out, expected, atol=5e-2) + + +def _per_qubit_symmetric_confusion(n_qubits, p): + """Stack of ``(n_qubits, 2, 2)`` symmetric bit-flip channels with flip probability ``p``.""" + one = np.array([[1.0 - p, p], [p, 1.0 - p]]) + return jnp.asarray(np.broadcast_to(one, (n_qubits, 2, 2))) + + +@pytest.mark.parametrize( + "n_qubits,n_shots,expected_k", + [ + # 2**n <= n_shots: histogram path, output K = 2**n_qubits. + (2, 8, 4), + (3, 16, 8), + # 2**n > n_shots: sort-RLE path, output K = n_shots. + (4, 8, 8), + (5, 16, 16), + ], +) +def test_sample_path_dispatch(n_qubits, n_shots, expected_k): + """Cover both sample post-processing paths: histogram (K=2**n) and sort-RLE (K=n_shots). + + With an identity noise channel, the linear solve reduces to the user + histogram itself, so total mass must equal n_shots in either path and + the output K must follow the trace-time branch in rem_apply_to_samples. + """ + cm = _per_qubit_symmetric_confusion(n_qubits, 0.0) + + rng = np.random.default_rng(0) + samples = rng.integers(0, 2, size=(n_shots, n_qubits)).astype(np.int32) + + bitstrings, counts = rem_apply_to_samples( + jnp.asarray(samples), + cm, + jnp.arange(n_qubits), + n_qubits, + ) + + assert counts.shape == (expected_k,) + assert bitstrings.shape == (expected_k, n_qubits) + assert np.isclose(float(jnp.sum(counts)), float(n_shots)) + + +def test_sample_sort_rle_path(): + """Sort-RLE path (2**n > n_shots) must preserve total shot mass through the linear solve.""" + n_qubits, n_shots, p = 4, 8, 0.10 + cm = _per_qubit_symmetric_confusion(n_qubits, p) + + rng = np.random.default_rng(7) + codes = rng.choice(2**n_qubits, size=n_shots, replace=False) + bits_msb_first = ((codes[:, None] >> np.arange(n_qubits - 1, -1, -1)) & 1).astype(np.int32) + + _, counts = rem_apply_to_samples( + jnp.asarray(bits_msb_first), + cm, + jnp.arange(n_qubits), + n_qubits, + ) + assert counts.shape == (n_shots,) + assert np.isclose(float(jnp.sum(counts)), float(n_shots), atol=1e-6) + + +def test_inverse_recovers_probs(): + """Round-trip: apply known confusion to a clean prob vector, invert with REM, recover input.""" + n_qubits, p = 2, 0.10 + clean = jnp.asarray([0.5, 0.0, 0.0, 0.5]) # 2-qubit GHZ analytic probs + cm = _per_qubit_symmetric_confusion(n_qubits, p) + + flip = np.array([[1.0 - p, p], [p, 1.0 - p]]) + forward = reduce(np.kron, [flip] * n_qubits) + noisy = jnp.asarray(forward @ np.asarray(clean)) + + recovered = rem_apply_to_probs(noisy, cm, jnp.arange(n_qubits), n_qubits) + assert np.allclose(np.asarray(recovered), np.asarray(clean), atol=1e-8) + + +def test_inverse_recovers_counts(): + """Same round-trip for counts: float linear solve must invert the analytic channel exactly.""" + n_qubits, p = 2, 0.15 + clean = jnp.asarray([500.0, 0.0, 0.0, 500.0]) + cm = _per_qubit_symmetric_confusion(n_qubits, p) + + flip = np.array([[1.0 - p, p], [p, 1.0 - p]]) + forward = reduce(np.kron, [flip] * n_qubits) + noisy = jnp.asarray(forward @ np.asarray(clean)) + + recovered = rem_apply_to_counts(noisy, cm, jnp.arange(n_qubits), n_qubits) + assert np.allclose(np.asarray(recovered), np.asarray(clean), atol=1e-6) + + +def test_calibrate_paths_agree(): + """rem_calibrate_{samples,counts,probs} on the same logical channel must yield matching matrices.""" + n_qubits, n_shots = 2, 4000 + p = 0.10 + rng = np.random.default_rng(0) + + flips_zero = rng.random((n_shots, n_qubits)) < p + zeros_samples = flips_zero.astype(np.int32) + flips_one = rng.random((n_shots, n_qubits)) < p + ones_samples = (1 - flips_one).astype(np.int32) + + cm_from_samples = rem_calibrate_samples(jnp.asarray(zeros_samples), jnp.asarray(ones_samples)) + + # Build counts/probs from the same samples so the three paths consume the same channel. + powers = (1 << np.arange(n_qubits - 1, -1, -1)).astype(np.int32) + z_codes = (zeros_samples[:, ::-1] * powers[None, :]).sum(axis=1) + o_codes = (ones_samples[:, ::-1] * powers[None, :]).sum(axis=1) + zeros_counts = np.bincount(z_codes, minlength=2**n_qubits).astype(np.int64) + ones_counts = np.bincount(o_codes, minlength=2**n_qubits).astype(np.int64) + + cm_from_counts = rem_calibrate_counts(jnp.asarray(zeros_counts), jnp.asarray(ones_counts)) + cm_from_probs = rem_calibrate_probs( + jnp.asarray(zeros_counts / zeros_counts.sum()), + jnp.asarray(ones_counts / ones_counts.sum()), + ) + + expected = np.array([[1.0 - p, p], [p, 1.0 - p]]) + for k in range(n_qubits): + for cm in (cm_from_samples, cm_from_counts, cm_from_probs): + assert np.allclose(np.asarray(cm[k]), expected, atol=3e-2) + + +@pytest.mark.parametrize("n_qubits", [1, 2, 3]) +def test_stretch_confusion_matrix(n_qubits): + """_stretch_confusion_matrix on a bitstring row must agree with an explicit lookup.""" + np_rng = np.random.default_rng(42) + cm = np_rng.random((2, 2)) + bitstring = np_rng.integers(0, 2, size=n_qubits).astype(np.int32) + + expected = np.empty((n_qubits, n_qubits)) + for i in range(n_qubits): + for j in range(n_qubits): + expected[i, j] = cm[bitstring[i], bitstring[j]] + + got = _stretch_confusion_matrix(jnp.asarray(bitstring), jnp.asarray(cm)) + assert np.allclose(np.asarray(got), expected) + + +def test_unsupported_measurement(): + """REM rejects non-sample/counts/probs measurements with a clear trace-time error.""" + dev = qp.device("lightning.qubit", wires=2) + + def circuit(): + qp.Hadamard(wires=0) + return qp.expval(qp.PauliZ(0)) + + qnode = qp.set_shots(qp.QNode(circuit, dev), shots=200) + + def mitigated(): + return mitigate_with_rem(qnode, compute_all_zeroes_ones=True)() + + with pytest.raises(AssertionError, match="measurement process must be one of"): + catalyst.qjit(mitigated) diff --git a/mlir/lib/Catalyst/Transforms/InlineNestedModules.cpp b/mlir/lib/Catalyst/Transforms/InlineNestedModules.cpp index 91b37243c9..f8c015174c 100644 --- a/mlir/lib/Catalyst/Transforms/InlineNestedModules.cpp +++ b/mlir/lib/Catalyst/Transforms/InlineNestedModules.cpp @@ -72,7 +72,6 @@ #include "mlir/Transforms/GreedyPatternRewriteDriver.h" #include "Catalyst/IR/CatalystOps.h" -#include "Catalyst/Transforms/Patterns.h" #include "Gradient/IR/GradientInterfaces.h" #include "Mitigation/IR/MitigationOps.h" @@ -553,9 +552,9 @@ struct InlineNestedSymbolTablePass : PassWrapper( - context, &old_to_new); + RewritePatternSet nestedToFlat(context); + nestedToFlat.add(context, &old_to_new); run = _stopAfterStep >= 4 || _stopAfterStep == 0; if (run && failed(applyPatternsGreedily(symbolTable, std::move(nestedToFlat), config))) { signalPassFailure(); diff --git a/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp b/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp index 560cb3611f..1f7cfe7cec 100644 --- a/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp +++ b/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp @@ -2,19 +2,19 @@ #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/IRMapping.h" #include "mlir/IR/SymbolTable.h" #include "mlir/IR/Types.h" + #include "Catalyst/Utils/CallGraph.h" -#include "mlir/IR/BuiltinAttributes.h" // Quantum dialect types and ops -#include "Quantum/IR/QuantumOps.h" - #include "llvm/Support/Debug.h" #include "llvm/Support/raw_ostream.h" +#include "Quantum/IR/QuantumOps.h" + #define DEBUG_TYPE "mitigation-rem" using namespace mlir; @@ -33,56 +33,64 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter // functions that return a 1-element tensor. These will be replaced // with proper calibration circuits in a subsequent step. Location loc = op.getLoc(); - + auto computeAllZeroesOnes = op.getComputeAllZeroesOnesAttr(); llvm::dbgs() << "[mitigation.rem] computeAllZeroesOnes: " << computeAllZeroesOnes << "\n"; - for (auto v : op.getResultTypes()) - { + for (auto v : op.getResultTypes()) { llvm::dbgs() << "[mitigation.rem] Result type: " << v << "\n"; - } + } // Resolve callee - func::FuncOp calleeOp = SymbolTable::lookupNearestSymbolFrom(op, op.getCalleeAttr()); + func::FuncOp calleeOp = + SymbolTable::lookupNearestSymbolFrom(op, op.getCalleeAttr()); if (!calleeOp) { return rewriter.notifyMatchFailure(op, "cannot resolve callee"); } - + // Call the original callee SmallVector callArgs(op.getArgs().begin(), op.getArgs().end()); auto callOp = rewriter.create(loc, calleeOp, callArgs); - - // Get SSA Values for wrapped callee function results. Note that these are not actual values of elements, but an abstraction, Value objects and ValueRange: non-owning view of multiple Value objects. These values are not known until runtime of the compiled program (not to be confused with from pass runtime / pass time). But their abstractions can be used to define the flow of the program + + // Get SSA Values for wrapped callee function results. Note that these are not actual values of + // elements, but an abstraction, Value objects and ValueRange: non-owning view of multiple Value + // objects. These values are not known until runtime of the compiled program (not to be confused + // with from pass runtime / pass time). But their abstractions can be used to define the flow of + // the program auto results = callOp.getResults(); - // Declare vector of SSA Values representing the total result (callee results + two calibration circuit results, or zero-tensors if no calibration is present) + // Declare vector of SSA Values representing the total result (callee results + two calibration + // circuit results, or zero-tensors if no calibration is present) SmallVector completeResults; - // The type (and shape is part of it) of the return value is known, and we know that we have 3 times as many Value objects (because of zeros and ones calibration circuits) + // The type (and shape is part of it) of the return value is known, and we know that we have 3 + // times as many Value objects (because of zeros and ones calibration circuits) completeResults.reserve(results.size() * 3); - + // Determine quantum device information, specifically number of qubits and shots. - // Get the number of qubits from quantum.alloc operation inside callee function (FuncOp). Note the due to iterator (pointer) - const int64_t qubitCount = (*calleeOp.getOps().begin()).getNqubitsAttr().value_or(0); - // Get the device using quantom.device opeartion inside callee function (FuncOp) + // Get the number of qubits from quantum.alloc operation inside callee function (FuncOp). Note + // the due to iterator (pointer) + const int64_t qubitCount = + (*calleeOp.getOps().begin()).getNqubitsAttr().value_or(0); + // Get the device using quantom.device opeartion inside callee function (FuncOp) quantum::DeviceInitOp deviceInitOp = *(calleeOp.getOps().begin()); - // shots value could be a result of some operation, and not just a literal value, which is why pointer to the operation is stored + // shots value could be a result of some operation, and not just a literal value, which is why + // pointer to the operation is stored Operation *shots = deviceInitOp.getShots().getDefiningOp(); StringAttr lib = deviceInitOp.getLibAttr(); StringAttr name = deviceInitOp.getDeviceNameAttr(); StringAttr kwargs = deviceInitOp.getKwargsAttr(); - llvm::dbgs() << "[mitigation.rem] qubitCount: " << qubitCount - << ", shots: " << *shots << ", lib: "<< lib - << ", name: " << name << ", kwargs" << kwargs << "\n"; - - //check for MP type: Sample, Probs or Counts + llvm::dbgs() << "[mitigation.rem] qubitCount: " << qubitCount << ", shots: " << *shots + << ", lib: " << lib << ", name: " << name << ", kwargs" << kwargs << "\n"; + + // check for MP type: Sample, Probs or Counts auto measurementProcesses = calleeOp.getOps(); if (measurementProcesses.empty()) { llvm::errs() << "[mitigation.rem] No valid MP found.\n"; return failure(); } - mlir::Operation* measurementProcess = (*measurementProcesses.begin()).getOperation(); + mlir::Operation *measurementProcess = (*measurementProcesses.begin()).getOperation(); std::string MPName = measurementProcess->getName().getIdentifier().str(); llvm::dbgs() << "[mitigation.rem] MP name: " << MPName << "\n"; - - //pre-compute MP and attribute-dependent values - bool doCalib = computeAllZeroesOnes.getValue(); //get boolean value of BoolAttr + + // pre-compute MP and attribute-dependent values + bool doCalib = computeAllZeroesOnes.getValue(); // get boolean value of BoolAttr mlir::RankedTensorType tensorTyI64; mlir::RankedTensorType tensorTyF64; int64_t calibrationTensorShape = 0; @@ -100,15 +108,20 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter tensorTyI64 = RankedTensorType::get({calibrationTensorShape}, rewriter.getI64Type()); } else if (MPName == "quantum.sample"s) { - // check if shots Op is actually a arith.constant. In that case, shots are known at compile time + // check if shots Op is actually a arith.constant. In that case, shots are known at compile + // time if (shots->getName().getIdentifier() == arith::ConstantOp::getOperationName()) { auto shotsAttr = shots->getAttrOfType("value"); if (!shotsAttr || shotsAttr == 0) { - llvm::errs() << "[mitigation.rem] Shots constant missing non-zero integer value. Sample MP not supported for analytic simulation (shots must be > 0).\n"; + llvm::errs() + << "[mitigation.rem] Shots constant missing non-zero integer value. Sample MP " + "not supported for analytic simulation (shots must be > 0).\n"; return failure(); } shotCount = shotsAttr.getInt(); - llvm::dbgs() << "[mitigation.rem] Shots is a arth.constant op, shots known at compile time: " << shotCount << "\n"; + llvm::dbgs() + << "[mitigation.rem] Shots is a arth.constant op, shots known at compile time: " + << shotCount << "\n"; calibrationTensorShape = shotCount * qubitCount; tensorTyI64 = RankedTensorType::get({shotCount, qubitCount}, rewriter.getI64Type()); tensorTyF64 = RankedTensorType::get({shotCount, qubitCount}, rewriter.getF64Type()); @@ -117,41 +130,52 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter llvm::errs() << "[mitigation.rem] Dynamic number of shots not supported currently.\n"; return failure(); } - + // SmallVector zeros_vec(bitspaceShape, 0.0); //initialize with zeroes } else { - llvm::errs() << "[mitigation.rem] Supported measurement processes are quantum.counts, quantum.probs and quantum.sample. " + llvm::errs() << "[mitigation.rem] Supported measurement processes are quantum.counts, " + "quantum.probs and quantum.sample. " << MPName << " is not supported.\n"; return failure(); } - if (!doCalib) { //doCalib == false - llvm::dbgs() << "[mitigation.rem] doCalibration == false, replacing RemOp with wrapped callee circuit function call..." << "\n"; + if (!doCalib) { // doCalib == false + llvm::dbgs() << "[mitigation.rem] doCalibration == false, replacing RemOp with wrapped " + "callee circuit function call..." + << "\n"; mlir::RankedTensorType tensorTyConst; - // Create constant array of type ranked tensor initialized with 0.0. This is the default value returned when doCalib == false and indicates that no calibration circuits were run - // auto tensorTy = RankedTensorType::get({bitspaceShape}, rewriter.getF64Type()); + // Create constant array of type ranked tensor initialized with 0.0. This is the + // default value returned when doCalib == false and indicates that no calibration circuits + // were run auto tensorTy = RankedTensorType::get({bitspaceShape}, rewriter.getF64Type()); // SmallVector zeros_vec(bitspaceShape, 0.0); //initialize with zeroes DenseElementsAttr tensorAttr; if (MPName == "quantum.probs"s || MPName == "quantum.counts"s) { - SmallVector zerosVector(calibrationTensorShape, 0.0); //initialize with zeroes + SmallVector zerosVector(calibrationTensorShape, 0.0); // initialize with zeroes // To initialize with other values: // for (auto i = 0; i < 4; ++i) { // zeros_vec[i] = static_cast(i); // } // ----------------------------------------------------------------- - // ArrayRef is a non-owning view and needs to be created because it is often required by MLIR APIs to avoid copies. - // It is also possible to use this: - // auto zerosAttr = rewriter.getF64ArrayAttr(zeros_vec); // returns ArrayAttr used by DenseElementsAttr - // because there is also an overload of DenseElementsAttr which accepts ArrayAttr. Attributes and compile-time constant metadata attached to ops/types in the IR, it requires a static shape + // ArrayRef is a non-owning view and needs to be created because it is often required by + // MLIR APIs to avoid copies. It is also possible to use this: auto zerosAttr = + // rewriter.getF64ArrayAttr(zeros_vec); // returns ArrayAttr used by DenseElementsAttr + // because there is also an overload of DenseElementsAttr which accepts ArrayAttr. + // Attributes and compile-time constant metadata attached to ops/types in the IR, it + // requires a static shape tensorTyConst = tensorTyF64; - tensorAttr = DenseElementsAttr::get(tensorTyConst, ArrayRef{zerosVector.begin(), zerosVector.end()}); - //tensorAttr is an attribute attached with the newly created airth.constant operation on the line below. It represents values known at pass-time (compile-time with regards to IR generation, but could be inferred at compile time of the passs from other attributes or something else) + tensorAttr = DenseElementsAttr::get( + tensorTyConst, ArrayRef{zerosVector.begin(), zerosVector.end()}); + // tensorAttr is an attribute attached with the newly created airth.constant operation + // on the line below. It represents values known at pass-time (compile-time with regards + // to IR generation, but could be inferred at compile time of the passs from other + // attributes or something else) } else if (MPName == "quantum.sample"s) { - SmallVector zerosVector(calibrationTensorShape, 0); //initialize with zeroes + SmallVector zerosVector(calibrationTensorShape, 0); // initialize with zeroes tensorTyConst = tensorTyI64; - tensorAttr = DenseElementsAttr::get(tensorTyConst, ArrayRef{zerosVector.begin(), zerosVector.end()}); - } + tensorAttr = DenseElementsAttr::get( + tensorTyConst, ArrayRef{zerosVector.begin(), zerosVector.end()}); + } auto cst_zeros = rewriter.create(loc, tensorTyConst, tensorAttr); llvm::dbgs() << "[mitigation.rem] cst_zeros = " << cst_zeros << "\n"; @@ -160,27 +184,38 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter completeResults.push_back(v); // 2) zeros calibration placeholders (reuse same tensors for now) completeResults.push_back(cst_zeros.getResult()); - // 3) ones calibration placeholders + // 3) ones calibration placeholders completeResults.push_back(cst_zeros.getResult()); - // Replace the results of the original operation (mitigation.rem) with this vector of Values (could be any object castable to ValueRange) + // Replace the results of the original operation (mitigation.rem) with this vector of Values + // (could be any object castable to ValueRange) rewriter.replaceOp(op, completeResults); return success(); } // else, doCalib == true - llvm::dbgs() << "[mitigation.rem] doCalibration == true, start adding all-zeroes and all-ones circuits..." << "\n"; + llvm::dbgs() << "[mitigation.rem] doCalibration == true, start adding all-zeroes and all-ones " + "circuits..." + << "\n"; // - //DONE: implement adding two calibation circuits (all-zeroes and all-ones) and return appropriate results - //DONE: Then, add support for other MP, such as CountsMP and change the return type accordingly. - // Finally, add support for returning Observables, by somehow injecting the mitigation code (applying the confusion matrix) before observable calculation. (out of scope for now) - // 0. Check how insertion point is managed -- for now it seems to be just fine as it is, but this must be researched in the future + // DONE: implement adding two calibation circuits (all-zeroes and all-ones) and return + // appropriate results DONE: Then, add support for other MP, such as CountsMP and change the + // return type accordingly. + // Finally, add support for returning Observables, by somehow injecting the mitigation code + // (applying the confusion matrix) before observable calculation. (out of scope for now) 0. + // Check how insertion point is managed -- for now it seems to be just fine as it is, but this + // must be researched in the future // ===BEGIN ZEROES QFUNC=== - // 1. Add device creation and allocation (quantum.device and quantum.alloc with previously determined attributes) - llvm::dbgs() << "[mitigation.rem] doCalibration == true, start adding all-zeroes and all-ones circuits..." << "\n"; + // 1. Add device creation and allocation (quantum.device and quantum.alloc with previously + // determined attributes) + llvm::dbgs() << "[mitigation.rem] doCalibration == true, start adding all-zeroes and all-ones " + "circuits..." + << "\n"; Operation *shotsLocal = shots->clone(); // not sure why this is needed - rewriter.insert(shotsLocal); // not sure why this is needed - auto devInit = rewriter.create(loc, shotsLocal->getResult(0), lib, name, kwargs); + rewriter.insert(shotsLocal); // not sure why this is needed + auto devInit = + rewriter.create(loc, shotsLocal->getResult(0), lib, name, kwargs); Type qregType = quantum::QuregType::get(rewriter.getContext()); - // not sure why empty IntegerAttr is needed, probably need to research more about how quantum.alloc works + // not sure why empty IntegerAttr is needed, probably need to research more about how + // quantum.alloc works IntegerAttr qubitCountAttr = rewriter.getI64IntegerAttr(qubitCount); Value numberQubitsValue = rewriter.create(loc, qubitCountAttr); // IntegerAttr intAttr{}; @@ -202,7 +237,8 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter // treats the first argument as the result type (observable) rather than // mistakenly interpreting operands. This avoids accidental operand ordering // that can produce explicit operandSegmentSizes or wrong printed form. - auto compbasis = rewriter.create(loc, TypeRange{obsType}, ValueRange{}, qregVal); + auto compbasis = rewriter.create(loc, TypeRange{obsType}, + ValueRange{}, qregVal); llvm::dbgs() << "[mitigation.rem] quantum.compbasis addition OK, op=" << compbasis << "\n"; // 3. Add quantum.probs op from result of compbasis and get Value of tensor // Probs returns a tensor of size 2^nqubits; build the appropriate RankedTensorType @@ -219,18 +255,26 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter quantum::MeasurementProcess insertedMP; ValueRange calibratedZeroResult; if (MPName == "quantum.probs"s) { - insertedMP = rewriter.create(loc, TypeRange{tensorTyF64}, compbasis.getResult(), Value(), Value()); + insertedMP = rewriter.create(loc, TypeRange{tensorTyF64}, + compbasis.getResult(), Value(), Value()); llvm::dbgs() << "[mitigation.rem] quantum.probs addition OK, op=" << insertedMP << "\n"; calibratedZeroResult = insertedMP->getResults(); } else if (MPName == "quantum.counts"s) { - insertedMP = rewriter.create(loc, TypeRange{tensorTyF64, tensorTyI64}, compbasis.getResult(), Value(), Value(), Value()); + insertedMP = + rewriter.create(loc, TypeRange{tensorTyF64, tensorTyI64}, + compbasis.getResult(), Value(), Value(), Value()); llvm::dbgs() << "[mitigation.rem] quantum.counts addition OK, op=" << insertedMP << "\n"; - calibratedZeroResult = insertedMP->getResults().drop_front(); // drop the first element in the iterator (eigvals not used for calibration matrices) + calibratedZeroResult = + insertedMP->getResults().drop_front(); // drop the first element in the iterator + // (eigvals not used for calibration matrices) } else if (MPName == "quantum.sample"s) { - // auto tensorTyFloat = RankedTensorType::get({shotCount, qubitCount}, rewriter.getF64Type()); - insertedMP = rewriter.create(loc, TypeRange{tensorTyF64}, compbasis.getResult(), ValueRange{}, Value());// (loc, TypeRange{tensorTy}, compbasis.getResult(), Value(), Value()); + // auto tensorTyFloat = RankedTensorType::get({shotCount, qubitCount}, + // rewriter.getF64Type()); + insertedMP = rewriter.create( + loc, TypeRange{tensorTyF64}, compbasis.getResult(), ValueRange{}, + Value()); // (loc, TypeRange{tensorTy}, compbasis.getResult(), Value(), Value()); llvm::dbgs() << "[mitigation.rem] quantum.sample addition OK, op=" << insertedMP << "\n"; IRMapping postprocMapping; uint64_t resultIndex = 0; @@ -239,17 +283,22 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter // OpBuilder::InsertionGuard guard(rewriter); rewriter.setInsertionPointAfter(insertedMP.getOperation()); calibratedZeroResult = insertedMP->getResults(); - for (auto postprocOp = measurementProcess->getNextNode(); postprocOp != nullptr && postprocOp->getName().getIdentifier() != quantum::DeallocOp::getOperationName(); postprocOp = postprocOp->getNextNode()) { - llvm::dbgs() << "[mitigation.rem] looking for next op: " << postprocOp->getName().getIdentifier() << "\n"; + for (auto postprocOp = measurementProcess->getNextNode(); + postprocOp != nullptr && + postprocOp->getName().getIdentifier() != quantum::DeallocOp::getOperationName(); + postprocOp = postprocOp->getNextNode()) { + llvm::dbgs() << "[mitigation.rem] looking for next op: " + << postprocOp->getName().getIdentifier() << "\n"; Operation *insertedOp = rewriter.clone(*postprocOp, postprocMapping); if (!insertedOp->getResults().empty()) calibratedZeroResult = insertedOp->getResults(); rewriter.setInsertionPointAfter(insertedOp); - llvm::dbgs() << "[mitigation.rem] insertedOp: " << insertedOp->getName().getIdentifier() << "\n"; + llvm::dbgs() << "[mitigation.rem] insertedOp: " << insertedOp->getName().getIdentifier() + << "\n"; // loc = insertedOp->getLoc(); } } - //MPName is already guaranteed to be a valid value, so no else block here. + // MPName is already guaranteed to be a valid value, so no else block here. // 4. Deallocate everything auto dealloc = rewriter.create(loc, qreg); @@ -259,12 +308,16 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter // ===END ZEROES QFUNC=== // 5. Add quantum.custom 'X' gates to implement ones circuit // ===BEGIN ONES QFUNC=== - // Reuse numberQubitsValue and qubitCountAttr created above. Clone shots and create a new device/alloc sequence for the ones circuit. + // Reuse numberQubitsValue and qubitCountAttr created above. Clone shots and create a new + // device/alloc sequence for the ones circuit. Operation *shotsLocalOnes = shots->clone(); rewriter.insert(shotsLocalOnes); - auto devInitOnes = rewriter.create(loc, shotsLocalOnes->getResult(0), lib, name, kwargs); - llvm::dbgs() << "[mitigation.rem] quantum.device initialization (ones) OK, op=" << devInitOnes << "\n"; - auto qregOnes = rewriter.create(loc, qregType, numberQubitsValue, qubitCountAttr); + auto devInitOnes = rewriter.create(loc, shotsLocalOnes->getResult(0), + lib, name, kwargs); + llvm::dbgs() << "[mitigation.rem] quantum.device initialization (ones) OK, op=" << devInitOnes + << "\n"; + auto qregOnes = + rewriter.create(loc, qregType, numberQubitsValue, qubitCountAttr); llvm::dbgs() << "[mitigation.rem] quantum.alloc (ones) addition OK, op=" << qregOnes << "\n"; // For each qubit in the register: extract -> apply PauliX -> insert back @@ -276,35 +329,51 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter for (int64_t i = 0; i < qubitCount; ++i) { auto idxAttr = rewriter.getI64IntegerAttr(i); // Extract qubit - auto extracted = rewriter.create(loc, rewriter.getType(), currentQreg, nullptr, idxAttr); + auto extracted = rewriter.create( + loc, rewriter.getType(), currentQreg, nullptr, idxAttr); // Apply PauliX (1-qubit gate) - auto xgate = rewriter.create(loc, /*gate_name=*/"PauliX", mlir::ValueRange({extracted.getResult()})); + auto xgate = rewriter.create(loc, /*gate_name=*/"PauliX", + mlir::ValueRange({extracted.getResult()})); // Insert back into the register, producing a new register Value - auto inserted = rewriter.create(loc, qregOnes.getType(), currentQreg, nullptr, idxAttr, xgate.getResult(0)); + auto inserted = rewriter.create(loc, qregOnes.getType(), currentQreg, + nullptr, idxAttr, xgate.getResult(0)); currentQreg = inserted.getResult(); } // Measure and get probs for ones circuit Value qregValOnes = currentQreg; - auto compbasisOnes = rewriter.create(loc, TypeRange{obsType}, ValueRange{}, qregValOnes); - llvm::dbgs() << "[mitigation.rem] quantum.compbasis (ones) addition OK, op=" << compbasisOnes << "\n"; - //TODO: maybe refactor into a function like in ZNE and call for zeroes and ones + auto compbasisOnes = rewriter.create(loc, TypeRange{obsType}, + ValueRange{}, qregValOnes); + llvm::dbgs() << "[mitigation.rem] quantum.compbasis (ones) addition OK, op=" << compbasisOnes + << "\n"; + // TODO: maybe refactor into a function like in ZNE and call for zeroes and ones quantum::MeasurementProcess insertedMPOnes; ValueRange calibratedOnesResult{}; if (MPName == "quantum.probs"s) { - insertedMPOnes = rewriter.create(loc, TypeRange{tensorTyF64}, compbasisOnes.getResult(), Value(), Value()); - llvm::dbgs() << "[mitigation.rem] quantum.probs (ones) addition OK, op=" << insertedMPOnes << "\n"; + insertedMPOnes = rewriter.create( + loc, TypeRange{tensorTyF64}, compbasisOnes.getResult(), Value(), Value()); + llvm::dbgs() << "[mitigation.rem] quantum.probs (ones) addition OK, op=" << insertedMPOnes + << "\n"; calibratedOnesResult = insertedMPOnes->getResults(); } if (MPName == "quantum.counts"s) { - insertedMPOnes = rewriter.create(loc, TypeRange{tensorTyF64, tensorTyI64}, compbasisOnes.getResult(), Value(), Value(), Value()); - llvm::dbgs() << "[mitigation.rem] quantum.counts (ones) addition OK, op=" << insertedMPOnes << "\n"; - calibratedOnesResult = insertedMPOnes->getResults().drop_front(); // drop the first element in the iterator (eigvals not used for calibration matrices) + insertedMPOnes = rewriter.create( + loc, TypeRange{tensorTyF64, tensorTyI64}, compbasisOnes.getResult(), Value(), Value(), + Value()); + llvm::dbgs() << "[mitigation.rem] quantum.counts (ones) addition OK, op=" << insertedMPOnes + << "\n"; + calibratedOnesResult = insertedMPOnes->getResults() + .drop_front(); // drop the first element in the iterator (eigvals + // not used for calibration matrices) } else if ((MPName == "quantum.sample"s)) { - // auto tensorTyFloat = RankedTensorType::get({shotCount, qubitCount}, rewriter.getF64Type()); - insertedMPOnes = rewriter.create(loc, TypeRange{tensorTyF64}, compbasisOnes.getResult(), ValueRange{}, Value());// (loc, TypeRange{tensorTy}, compbasis.getResult(), Value(), Value()); - llvm::dbgs() << "[mitigation.rem] quantum.sample addition OK, op=" << insertedMPOnes << "\n"; + // auto tensorTyFloat = RankedTensorType::get({shotCount, qubitCount}, + // rewriter.getF64Type()); + insertedMPOnes = rewriter.create( + loc, TypeRange{tensorTyF64}, compbasisOnes.getResult(), ValueRange{}, + Value()); // (loc, TypeRange{tensorTy}, compbasis.getResult(), Value(), Value()); + llvm::dbgs() << "[mitigation.rem] quantum.sample addition OK, op=" << insertedMPOnes + << "\n"; IRMapping postprocMapping; uint64_t resultIndex = 0; for (auto res : measurementProcess->getResults()) @@ -312,18 +381,24 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter // OpBuilder::InsertionGuard guard(rewriter); rewriter.setInsertionPointAfter(insertedMPOnes.getOperation()); calibratedOnesResult = insertedMPOnes->getResults(); - for (auto postprocOp = measurementProcess->getNextNode(); postprocOp != nullptr && postprocOp->getName().getIdentifier() != quantum::DeallocOp::getOperationName(); postprocOp = postprocOp->getNextNode()) { - llvm::dbgs() << "[mitigation.rem] looking for next op: " << postprocOp->getName().getIdentifier() << "\n"; + for (auto postprocOp = measurementProcess->getNextNode(); + postprocOp != nullptr && + postprocOp->getName().getIdentifier() != quantum::DeallocOp::getOperationName(); + postprocOp = postprocOp->getNextNode()) { + llvm::dbgs() << "[mitigation.rem] looking for next op: " + << postprocOp->getName().getIdentifier() << "\n"; Operation *insertedOp = rewriter.clone(*postprocOp, postprocMapping); if (!insertedOp->getResults().empty()) calibratedOnesResult = insertedOp->getResults(); rewriter.setInsertionPointAfter(insertedOp); // loc = insertedOp->getLoc(); - llvm::dbgs() << "[mitigation.rem] insertedOp: " << insertedOp->getName().getIdentifier() << "\n"; - } + llvm::dbgs() << "[mitigation.rem] insertedOp: " << insertedOp->getName().getIdentifier() + << "\n"; + } } auto deallocOnes = rewriter.create(loc, qregOnes); - llvm::dbgs() << "[mitigation.rem] quantum.dealloc (ones) addition OK, op=" << deallocOnes << "\n"; + llvm::dbgs() << "[mitigation.rem] quantum.dealloc (ones) addition OK, op=" << deallocOnes + << "\n"; auto deinitOnes = rewriter.create(loc); llvm::dbgs() << "[mitigation.rem] quantum.deinit (ones) addition OK, op=" << deinitOnes << "\n"; // ===END ONES QFUNC=== @@ -343,7 +418,7 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter completeResults.push_back(*calibratedZeroResult.begin()); // 3) ones calibration result completeResults.push_back(*calibratedOnesResult.begin()); - + rewriter.replaceOp(op, completeResults); return success(); @@ -379,7 +454,8 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter // // Create function type: () -> tensor // auto f64 = rewriter.getF64Type(); // auto tensorTy = RankedTensorType::get({static_cast(size)}, f64); - // auto fnType = FunctionType::get(rewriter.getContext(), /*inputs=*/{}, /*results=*/{tensorTy}); + // auto fnType = FunctionType::get(rewriter.getContext(), /*inputs=*/{}, + // /*results=*/{tensorTy}); // OpBuilder::InsertionGuard guard(rewriter); // // Insert the new function at the start of the module body @@ -396,7 +472,8 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter // if (constPos <= 0.0) // pos = 0; // else - // pos = static_cast(constPos) < size ? static_cast(constPos) : size - 1; + // pos = static_cast(constPos) < size ? static_cast(constPos) : size + // - 1; // } // elems[pos] = 1.0; // auto attr = DenseElementsAttr::get(tensorTy, rewriter.getF64ArrayAttr(elems)); diff --git a/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.hpp b/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.hpp index beef344f45..341fb75aba 100644 --- a/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.hpp +++ b/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.hpp @@ -1,20 +1,19 @@ #pragma once -#include "Mitigation/IR/MitigationOps.h" -#include "Quantum/IR/QuantumOps.h" #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/IR/PatternMatch.h" +#include "Mitigation/IR/MitigationOps.h" + using namespace mlir; namespace catalyst { namespace mitigation { struct RemLowering : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; + using OpRewritePattern::OpRewritePattern; - LogicalResult matchAndRewrite(mitigation::RemOp op, - PatternRewriter &rewriter) const override; + LogicalResult matchAndRewrite(mitigation::RemOp op, PatternRewriter &rewriter) const override; }; // Add patterns from this file into a pattern set. diff --git a/mlir/test/Mitigation/RemLoweringTest.mlir b/mlir/test/Mitigation/RemLoweringTest.mlir new file mode 100644 index 0000000000..b357dcae9b --- /dev/null +++ b/mlir/test/Mitigation/RemLoweringTest.mlir @@ -0,0 +1,185 @@ +// Copyright 2026 Haiqu, Inc. + +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at + +// http://www.apache.org/licenses/LICENSE-2.0 + +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// RUN: quantum-opt %s --lower-mitigation --split-input-file --verify-diagnostics | FileCheck %s + +// =================================================================== +// computeAllZeroesOnes(true) on a ProbsMP callee: the pass must keep the +// original func.call, then emit two cloned quantum.device / quantum.alloc / +// quantum.compbasis / quantum.probs blocks. The all-ones clone is +// distinguished from the all-zeroes one by per-qubit PauliX gates inserted +// before the compbasis observation. +// =================================================================== + +func.func @probsCircuit() -> tensor<4xf64> attributes {qnode} { + %shots = arith.constant 0 : i64 + quantum.device shots(%shots) ["rtd_lightning.so", "LightningQubit", "{}"] + %r = quantum.alloc(2) : !quantum.reg + %idx0 = arith.constant 0 : i64 + %idx1 = arith.constant 1 : i64 + %q_0 = quantum.extract %r[%idx0] : !quantum.reg -> !quantum.bit + %q_1 = quantum.extract %r[%idx1] : !quantum.reg -> !quantum.bit + %h_0 = quantum.custom "h"() %q_0 : !quantum.bit + %r_0 = quantum.insert %r[ 0], %h_0 : !quantum.reg, !quantum.bit + %r_1 = quantum.insert %r_0[ 1], %q_1 : !quantum.reg, !quantum.bit + %obs = quantum.compbasis qubits %h_0, %q_1 : !quantum.obs + %probs = quantum.probs %obs : tensor<4xf64> + quantum.dealloc %r_1 : !quantum.reg + quantum.device_release + func.return %probs : tensor<4xf64> +} + +func.func @remProbsCalibrate() -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64>) { + %out:3 = mitigation.rem @probsCircuit() computeAllZeroesOnes(true) : () -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64>) + func.return %out#0, %out#1, %out#2 : tensor<4xf64>, tensor<4xf64>, tensor<4xf64> +} + +// CHECK-LABEL: func.func @remProbsCalibrate +// CHECK: call @probsCircuit() : () -> tensor<4xf64> +// CHECK: quantum.device shots({{.*}}) ["rtd_lightning.so", "LightningQubit", "{}"] +// CHECK: quantum.alloc(%{{.*}}) : !quantum.reg +// CHECK: quantum.compbasis +// CHECK: quantum.probs {{.*}} : tensor<4xf64> +// CHECK: quantum.device_release +// CHECK: quantum.device shots({{.*}}) ["rtd_lightning.so", "LightningQubit", "{}"] +// CHECK: quantum.alloc(%{{.*}}) : !quantum.reg +// CHECK: quantum.custom "PauliX" +// CHECK: quantum.custom "PauliX" +// CHECK: quantum.compbasis +// CHECK: quantum.probs {{.*}} : tensor<4xf64> +// CHECK: quantum.device_release + +// ----- + +// =================================================================== +// computeAllZeroesOnes(true) on a CountsMP callee: callee result is the +// (eigvals, counts) pair, and the rem op surfaces two i64 calibration +// tensors -- one per all-zeroes / all-ones run. +// =================================================================== + +func.func @countsCircuit() -> (tensor<4xf64>, tensor<4xi64>) attributes {qnode} { + %shots = arith.constant 100 : i64 + quantum.device shots(%shots) ["rtd_lightning.so", "LightningQubit", "{}"] + %r = quantum.alloc(2) : !quantum.reg + %idx0 = arith.constant 0 : i64 + %idx1 = arith.constant 1 : i64 + %q_0 = quantum.extract %r[%idx0] : !quantum.reg -> !quantum.bit + %q_1 = quantum.extract %r[%idx1] : !quantum.reg -> !quantum.bit + %h_0 = quantum.custom "h"() %q_0 : !quantum.bit + %r_0 = quantum.insert %r[ 0], %h_0 : !quantum.reg, !quantum.bit + %r_1 = quantum.insert %r_0[ 1], %q_1 : !quantum.reg, !quantum.bit + %obs = quantum.compbasis qubits %h_0, %q_1 : !quantum.obs + %eig, %cnt = quantum.counts %obs : tensor<4xf64>, tensor<4xi64> + quantum.dealloc %r_1 : !quantum.reg + quantum.device_release + func.return %eig, %cnt : tensor<4xf64>, tensor<4xi64> +} + +func.func @remCountsCalibrate() -> (tensor<4xf64>, tensor<4xi64>, tensor<4xi64>, tensor<4xi64>) { + %out:4 = mitigation.rem @countsCircuit() computeAllZeroesOnes(true) : () -> (tensor<4xf64>, tensor<4xi64>, tensor<4xi64>, tensor<4xi64>) + func.return %out#0, %out#1, %out#2, %out#3 : tensor<4xf64>, tensor<4xi64>, tensor<4xi64>, tensor<4xi64> +} + +// CHECK-LABEL: func.func @remCountsCalibrate +// CHECK: call @countsCircuit() : () -> (tensor<4xf64>, tensor<4xi64>) +// CHECK: quantum.counts {{.*}} : tensor<4xf64>, tensor<4xi64> +// CHECK: quantum.custom "PauliX" +// CHECK: quantum.counts {{.*}} : tensor<4xf64>, tensor<4xi64> + +// ----- + +// =================================================================== +// computeAllZeroesOnes(true) on a SampleMP callee with shots = 200: the +// calibration tensors take the (shots, qubits) shape inferred from the +// callee's quantum.device + quantum.alloc. +// =================================================================== + +func.func @sampleCircuit() -> tensor<200x2xf64> attributes {qnode} { + %shots = arith.constant 200 : i64 + quantum.device shots(%shots) ["rtd_lightning.so", "LightningQubit", "{}"] + %r = quantum.alloc(2) : !quantum.reg + %idx0 = arith.constant 0 : i64 + %idx1 = arith.constant 1 : i64 + %q_0 = quantum.extract %r[%idx0] : !quantum.reg -> !quantum.bit + %q_1 = quantum.extract %r[%idx1] : !quantum.reg -> !quantum.bit + %h_0 = quantum.custom "h"() %q_0 : !quantum.bit + %r_0 = quantum.insert %r[ 0], %h_0 : !quantum.reg, !quantum.bit + %r_1 = quantum.insert %r_0[ 1], %q_1 : !quantum.reg, !quantum.bit + %obs = quantum.compbasis qubits %h_0, %q_1 : !quantum.obs + %samples = quantum.sample %obs : tensor<200x2xf64> + quantum.dealloc %r_1 : !quantum.reg + quantum.device_release + func.return %samples : tensor<200x2xf64> +} + +func.func @remSampleCalibrate() -> (tensor<200x2xf64>, tensor<200x2xf64>, tensor<200x2xf64>) { + %out:3 = mitigation.rem @sampleCircuit() computeAllZeroesOnes(true) : () -> (tensor<200x2xf64>, tensor<200x2xf64>, tensor<200x2xf64>) + func.return %out#0, %out#1, %out#2 : tensor<200x2xf64>, tensor<200x2xf64>, tensor<200x2xf64> +} + +// CHECK-LABEL: func.func @remSampleCalibrate +// CHECK: call @sampleCircuit() : () -> tensor<200x2xf64> +// CHECK: quantum.sample {{.*}} : tensor<200x2xf64> +// CHECK: quantum.custom "PauliX" +// CHECK: quantum.sample {{.*}} : tensor<200x2xf64> + +// ----- + +// =================================================================== +// computeAllZeroesOnes(false): the pass forwards the callee unchanged and +// stamps two zero-filled placeholder tensors for the calibration slots -- +// no cloned circuits and no extra quantum.device calls. +// =================================================================== + +func.func @passthroughCircuit() -> tensor<4xf64> attributes {qnode} { + %shots = arith.constant 0 : i64 + quantum.device shots(%shots) ["rtd_lightning.so", "LightningQubit", "{}"] + %r = quantum.alloc(2) : !quantum.reg + %idx0 = arith.constant 0 : i64 + %q_0 = quantum.extract %r[%idx0] : !quantum.reg -> !quantum.bit + %h_0 = quantum.custom "h"() %q_0 : !quantum.bit + %r_0 = quantum.insert %r[ 0], %h_0 : !quantum.reg, !quantum.bit + %obs = quantum.compbasis qubits %h_0 : !quantum.obs + %probs = quantum.probs %obs : tensor<4xf64> + quantum.dealloc %r_0 : !quantum.reg + quantum.device_release + func.return %probs : tensor<4xf64> +} + +func.func @remProbsPassthrough() -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64>) { + %out:3 = mitigation.rem @passthroughCircuit() computeAllZeroesOnes(false) : () -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64>) + func.return %out#0, %out#1, %out#2 : tensor<4xf64>, tensor<4xf64>, tensor<4xf64> +} + +// CHECK-LABEL: func.func @remProbsPassthrough +// CHECK: arith.constant dense<0.000000e+00> : tensor<4xf64> +// CHECK: call @passthroughCircuit() : () -> tensor<4xf64> +// CHECK-NOT: quantum.compbasis +// CHECK-NOT: quantum.custom "PauliX" + +// ----- + +// =================================================================== +// A missing callee leaves the mitigation.rem op untouched: the lowering +// pattern bails out via notifyMatchFailure, which keeps the IR in a state +// that downstream passes can either re-try or surface as a hard error. +// =================================================================== + +func.func @remMissingCallee() -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64>) { + %out:3 = mitigation.rem @doesNotExist() computeAllZeroesOnes(true) : () -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64>) + func.return %out#0, %out#1, %out#2 : tensor<4xf64>, tensor<4xf64>, tensor<4xf64> +} + +// CHECK-LABEL: func.func @remMissingCallee +// CHECK: mitigation.rem @doesNotExist() computeAllZeroesOnes(true) From 4744f034260ed2d2dce9bb49045d9e7901388b30 Mon Sep 17 00:00:00 2001 From: volodymyr-hq Date: Tue, 30 Jun 2026 18:02:10 +0300 Subject: [PATCH 4/8] Remove commented-out code --- .../api_extensions/error_mitigation.py | 9 +- frontend/catalyst/jax_primitives.py | 8 +- frontend/test/pytest/test_mitigation_rem.py | 2 +- .../Transforms/MitigationMethods/Rem.cpp | 123 +----------------- 4 files changed, 6 insertions(+), 136 deletions(-) diff --git a/frontend/catalyst/api_extensions/error_mitigation.py b/frontend/catalyst/api_extensions/error_mitigation.py index d72a193357..ce80386cd9 100644 --- a/frontend/catalyst/api_extensions/error_mitigation.py +++ b/frontend/catalyst/api_extensions/error_mitigation.py @@ -334,10 +334,7 @@ def __call__(self, *args, **kwargs): qnode_obj = jaxpr.eqns[0].params.get("qnode", None) assert qnode_obj is not None, "REM post-processing requires a QNode target" n_qubits = len(qnode_obj.device.wires) - # print(f"the wires obj: {qnode_obj.device.wires}, {[x for x in qnode_obj.device.wires]}") - measured_qubits = jnp.array( - [x for x in qnode_obj.device.wires] - ) # list(range(n_qubits)) # qnode_obj.device.wires + measured_qubits = jnp.array([x for x in qnode_obj.device.wires]) if mp_kind == "sample": # quantum.sample returns f64 in catalyst; the REM helpers index @@ -345,7 +342,6 @@ def __call__(self, *args, **kwargs): mitigatee_samples = rem_results[0].astype(jnp.int32) zeros_samples = rem_results[1].astype(jnp.int32) ones_samples = rem_results[2].astype(jnp.int32) - # measured_qubits = jnp.arange(n_qubits, dtype=jnp.int32) confusion_matrices = rem_calibrate_samples(zeros_samples, ones_samples) unique_bitstrings, mitigated_counts = rem_apply_to_samples( @@ -359,13 +355,10 @@ def __call__(self, *args, **kwargs): mitigatee_counts = rem_results[1] zeros_counts = rem_results[2] ones_counts = rem_results[3] - # jax.debug.print("I'm before rem_calibrate_counts") confusion_matrices = rem_calibrate_counts(zeros_counts, ones_counts) - # jax.debug.print("I'm after rem_calibrate_counts") mitigated_counts = rem_apply_to_counts( mitigatee_counts, confusion_matrices, measured_qubits, n_qubits ) - print(f"Done with RemCallable, returning: {mitigatee_eigvals, mitigated_counts}") return (mitigatee_eigvals, mitigated_counts) elif mp_kind == "probs": diff --git a/frontend/catalyst/jax_primitives.py b/frontend/catalyst/jax_primitives.py index 07ffcfba30..b6001ff4ac 100644 --- a/frontend/catalyst/jax_primitives.py +++ b/frontend/catalyst/jax_primitives.py @@ -1154,12 +1154,8 @@ def _rem_lowering(ctx, *args, compute_all_zeroes_ones, jaxpr, fn): symbol_ref = get_symbolref(ctx, func_op) *callee_avals, zeros_aval, ones_aval = ctx.avals_out callee_output_types = list(map(mlir.aval_to_ir_types, callee_avals)) - zeros_output_types = [ - mlir.aval_to_ir_types(zeros_aval) - ] # list(map(mlir.aval_to_ir_types, [ctx.avals_out[1]])) - ones_output_types = [ - mlir.aval_to_ir_types(ones_aval) - ] # list(map(mlir.aval_to_ir_types, [ctx.avals_out[2]])) + zeros_output_types = [mlir.aval_to_ir_types(zeros_aval)] + ones_output_types = [mlir.aval_to_ir_types(ones_aval)] flat_callee_output_types = util.flatten(callee_output_types) # The zeros/ones operands are non-variadic in the .td, so a single type # (not a list) is what the op constructor expects. diff --git a/frontend/test/pytest/test_mitigation_rem.py b/frontend/test/pytest/test_mitigation_rem.py index e109f48951..e6ae46a137 100644 --- a/frontend/test/pytest/test_mitigation_rem.py +++ b/frontend/test/pytest/test_mitigation_rem.py @@ -92,7 +92,7 @@ def mitigated(): @pytest.mark.parametrize("n_qubits", [1, 2, 11]) def test_sample(n_qubits): """Mitigated sample histogram on a noiseless GHZ state must match the analytic distribution. - + Note that for `n_qubits == 11`, 2^n_qubits > shots (2000), which covers the shot-count bounded codepath for transition matrix generation. """ diff --git a/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp b/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp index 1f7cfe7cec..51c938cbfa 100644 --- a/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp +++ b/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp @@ -130,8 +130,6 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter llvm::errs() << "[mitigation.rem] Dynamic number of shots not supported currently.\n"; return failure(); } - - // SmallVector zeros_vec(bitspaceShape, 0.0); //initialize with zeroes } else { llvm::errs() << "[mitigation.rem] Supported measurement processes are quantum.counts, " @@ -206,9 +204,6 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter // ===BEGIN ZEROES QFUNC=== // 1. Add device creation and allocation (quantum.device and quantum.alloc with previously // determined attributes) - llvm::dbgs() << "[mitigation.rem] doCalibration == true, start adding all-zeroes and all-ones " - "circuits..." - << "\n"; Operation *shotsLocal = shots->clone(); // not sure why this is needed rewriter.insert(shotsLocal); // not sure why this is needed auto devInit = @@ -218,7 +213,6 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter // quantum.alloc works IntegerAttr qubitCountAttr = rewriter.getI64IntegerAttr(qubitCount); Value numberQubitsValue = rewriter.create(loc, qubitCountAttr); - // IntegerAttr intAttr{}; llvm::dbgs() << "[mitigation.rem] quantum.device initialization OK, op=" << devInit << "\n"; auto qreg = rewriter.create(loc, qregType, numberQubitsValue, qubitCountAttr); llvm::dbgs() << "[mitigation.rem] quantum.alloc addition OK, op=" << qreg << "\n"; @@ -242,8 +236,6 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter llvm::dbgs() << "[mitigation.rem] quantum.compbasis addition OK, op=" << compbasis << "\n"; // 3. Add quantum.probs op from result of compbasis and get Value of tensor // Probs returns a tensor of size 2^nqubits; build the appropriate RankedTensorType - // const int64_t bitspaceShape = 2 << (qubitCount - 1); - // auto probsTensorTy = RankedTensorType::get({bitspaceShape}, rewriter.getF64Type()); // Be explicit about the builder overload to avoid overload-resolution // ambiguity that can lead to incorrect operandSegmentSizes being set. // Use the TypeRange overload so the result type is unambiguous and the @@ -270,17 +262,13 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter // (eigvals not used for calibration matrices) } else if (MPName == "quantum.sample"s) { - // auto tensorTyFloat = RankedTensorType::get({shotCount, qubitCount}, - // rewriter.getF64Type()); insertedMP = rewriter.create( - loc, TypeRange{tensorTyF64}, compbasis.getResult(), ValueRange{}, - Value()); // (loc, TypeRange{tensorTy}, compbasis.getResult(), Value(), Value()); + loc, TypeRange{tensorTyF64}, compbasis.getResult(), ValueRange{}, Value()); llvm::dbgs() << "[mitigation.rem] quantum.sample addition OK, op=" << insertedMP << "\n"; IRMapping postprocMapping; uint64_t resultIndex = 0; for (auto res : measurementProcess->getResults()) postprocMapping.map(res, insertedMP->getResult(resultIndex++)); - // OpBuilder::InsertionGuard guard(rewriter); rewriter.setInsertionPointAfter(insertedMP.getOperation()); calibratedZeroResult = insertedMP->getResults(); for (auto postprocOp = measurementProcess->getNextNode(); @@ -295,7 +283,6 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter rewriter.setInsertionPointAfter(insertedOp); llvm::dbgs() << "[mitigation.rem] insertedOp: " << insertedOp->getName().getIdentifier() << "\n"; - // loc = insertedOp->getLoc(); } } // MPName is already guaranteed to be a valid value, so no else block here. @@ -367,18 +354,14 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter // not used for calibration matrices) } else if ((MPName == "quantum.sample"s)) { - // auto tensorTyFloat = RankedTensorType::get({shotCount, qubitCount}, - // rewriter.getF64Type()); insertedMPOnes = rewriter.create( - loc, TypeRange{tensorTyF64}, compbasisOnes.getResult(), ValueRange{}, - Value()); // (loc, TypeRange{tensorTy}, compbasis.getResult(), Value(), Value()); + loc, TypeRange{tensorTyF64}, compbasisOnes.getResult(), ValueRange{}, Value()); llvm::dbgs() << "[mitigation.rem] quantum.sample addition OK, op=" << insertedMPOnes << "\n"; IRMapping postprocMapping; uint64_t resultIndex = 0; for (auto res : measurementProcess->getResults()) postprocMapping.map(res, insertedMPOnes->getResult(resultIndex++)); - // OpBuilder::InsertionGuard guard(rewriter); rewriter.setInsertionPointAfter(insertedMPOnes.getOperation()); calibratedOnesResult = insertedMPOnes->getResults(); for (auto postprocOp = measurementProcess->getNextNode(); @@ -391,7 +374,6 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter if (!insertedOp->getResults().empty()) calibratedOnesResult = insertedOp->getResults(); rewriter.setInsertionPointAfter(insertedOp); - // loc = insertedOp->getLoc(); llvm::dbgs() << "[mitigation.rem] insertedOp: " << insertedOp->getName().getIdentifier() << "\n"; } @@ -402,14 +384,6 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter auto deinitOnes = rewriter.create(loc); llvm::dbgs() << "[mitigation.rem] quantum.deinit (ones) addition OK, op=" << deinitOnes << "\n"; // ===END ONES QFUNC=== - // If I'll ever need to create integer constants later, this is how it's done: - // TypedAttr numberQubitsAttr = rewriter.getI64IntegerAttr(numberQubits); - // Value numberQubitsValue = rewriter.create(loc, numberQubitsAttr); - - // placeholder fallback: return 3 groups (callee, zeros, ones) by - // replicating the call results. Each returned Value is an SSA use of - // the same tensor produced by the callee; this yields three separate - // results in the caller, each referencing the same underlying tensor. // 1) original callee results for (Value v : results) @@ -421,99 +395,6 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter rewriter.replaceOp(op, completeResults); return success(); - - // // Determine device qubit count if present on the op. This controls the - // // size of the probability vector returned by the calibration circuits - // // (2^n_qubits). If absent, fallback to a single-element tensor used - // // during iterative development. - // int64_t nqubits = 0; - // if (auto attr = op->getAttrOfType("deviceNumQubits")) { - // nqubits = attr.getInt(); - // } - // size_t probSize = 1; - // if (nqubits > 0) { - // // cap shifts to avoid UB for large n - // if (nqubits < 63) - // probSize = (size_t)1 << static_cast(nqubits); - // else - // probSize = 1; // fallback - // } - - // // Helper to get or create a simple calibration function that returns - // // tensor with a constant one-hot vector. The function - // // is inserted into the parent module. - // auto getOrCreateCalib = [&](StringRef suffix, size_t size, double constPos) -> func::FuncOp { - // auto module = op->getParentOfType(); - // assert(module && "rem op must be inside a module"); - // std::string name = (calleeOp.getName().str() + ".rem_calib_") + suffix.str(); - - // // If function already exists, return it. - // if (auto existing = module.lookupSymbol(name)) - // return existing; - - // // Create function type: () -> tensor - // auto f64 = rewriter.getF64Type(); - // auto tensorTy = RankedTensorType::get({static_cast(size)}, f64); - // auto fnType = FunctionType::get(rewriter.getContext(), /*inputs=*/{}, - // /*results=*/{tensorTy}); - - // OpBuilder::InsertionGuard guard(rewriter); - // // Insert the new function at the start of the module body - // rewriter.setInsertionPointToStart(module.getBody()); - // auto func = rewriter.create(loc, name, fnType); - // // Add an entry block - // auto *block = func.addEntryBlock(); - - // // Create a dense constant vector with a single 1.0 in the requested - // // position (constPos interpreted as index position) and zeros elsewhere. - // SmallVector elems(size, 0.0); - // size_t pos = 0; - // if (size > 0) { - // if (constPos <= 0.0) - // pos = 0; - // else - // pos = static_cast(constPos) < size ? static_cast(constPos) : size - // - 1; - // } - // elems[pos] = 1.0; - // auto attr = DenseElementsAttr::get(tensorTy, rewriter.getF64ArrayAttr(elems)); - // rewriter.setInsertionPointToStart(block); - // auto cst = rewriter.create(loc, tensorTy, attr); - // Value constValRes = cst.getResult(); - // rewriter.create(loc, constValRes); - // return func; - // }; - - // // Create/get calibration functions - // func::FuncOp allZeros = getOrCreateCalib("all_zeroes", 0.0); - // func::FuncOp allOnes = getOrCreateCalib("all_ones", 1.0); - - // // Call the calibration functions - // auto callZeros = rewriter.create(loc, allZeros, ArrayRef{}); - // auto callOnes = rewriter.create(loc, allOnes, ArrayRef{}); - - // // Aggregate results: original callee results followed by zeros and ones results - // SmallVector results; - // for (auto r : originalCall.getResults()) - // results.push_back(r); - // for (auto r : callZeros.getResults()) - // results.push_back(r); - // for (auto r : callOnes.getResults()) - // results.push_back(r); - - // // Decide whether to emit calibration circuits based on the attribute. - // // If the frontend requested no calibration runs, simply forward the - // // callee results (preserving the previous behavior). - // bool doCalib = false; - // if (computeAllZeroesOnes) { - // // computeAllZeroesOnes is a BoolAttr; retrieve its value if present. - // doCalib = computeAllZeroesOnes.getValue(); - // } - - // // Otherwise, we will produce callee results + two calibration results. - // // Replace the operation with the full set of produced values. - // rewriter.replaceOp(op, results); - // return success(); } void populateRemLoweringPatterns(RewritePatternSet &patterns) From 0bef86f70102b4c667fe50658c6615fa5ed5d849 Mon Sep 17 00:00:00 2001 From: volodymyr-hq Date: Wed, 1 Jul 2026 20:59:58 +0300 Subject: [PATCH 5/8] Fix confusion matrix caching via runCalibration attribute --- .../api_extensions/error_mitigation.py | 98 ++++++++++------ frontend/catalyst/jax_primitives.py | 12 +- frontend/test/lit/test_mitigation_rem.py | 36 +++--- frontend/test/pytest/test_mitigation_rem.py | 111 ++++++++++++++++-- mlir/include/Mitigation/IR/MitigationOps.td | 11 +- .../Transforms/MitigationMethods/Rem.cpp | 6 +- mlir/test/Mitigation/RemLoweringTest.mlir | 20 ++-- 7 files changed, 205 insertions(+), 89 deletions(-) diff --git a/frontend/catalyst/api_extensions/error_mitigation.py b/frontend/catalyst/api_extensions/error_mitigation.py index ce80386cd9..a3d12d8dd9 100644 --- a/frontend/catalyst/api_extensions/error_mitigation.py +++ b/frontend/catalyst/api_extensions/error_mitigation.py @@ -168,14 +168,18 @@ def workflow(weights, s): return ZNECallable(fn, scale_factors, extrapolate, folding) -def mitigate_with_rem(fn=None, *, compute_all_zeroes_ones: bool = False): +def mitigate_with_rem( + fn=None, *, calibration_matrices=None, return_calibration_matrices: bool = False +): """A qjit-compatible frontend that emits a `mitigation.rem` operation. - Takes a boolean ``compute_all_zeroes_ones`` flag that is forwarded to the - lowering as an attribute on the `mitigation.rem` op. When ``True``, the - mitigation-lowering pass also emits the all-zeros and all-ones calibration - circuits and this callable runs the classical REM post-processing on the - three resulting sample arrays. + When ``calibration_matrices`` is ``None`` the lowering pass emits the + all-zeros and all-ones calibration circuits and this callable computes + the confusion matrices itself. When a precomputed ``(n_qubits, 2, 2)`` + matrix stack is supplied, ``rem_calibrate_*`` is skipped and the matrices + are passed straight into ``rem_apply_to_*``. Pass + ``return_calibration_matrices=True`` to get the matrices back alongside + the mitigated result so they can be cached for subsequent calls. """ kwargs = copy.copy(locals()) @@ -184,7 +188,11 @@ def mitigate_with_rem(fn=None, *, compute_all_zeroes_ones: bool = False): if fn is None: return functools.partial(mitigate_with_rem, **kwargs) - return RemCallable(fn, compute_all_zeroes_ones=compute_all_zeroes_ones) + return RemCallable( + fn, + calibration_matrices=calibration_matrices, + return_calibration_matrices=return_calibration_matrices, + ) ## IMPL ## @@ -267,23 +275,32 @@ class RemCallable(CatalystCallable): Args: fn (Callable): the circuit to be mitigated with REM. - compute_all_zeroes_ones (bool): forwarded to the lowering as an - attribute on the ``mitigation.rem`` op. When ``True``, the pass - emits the all-zeros and all-ones calibration circuits next to the - user circuit and this callable applies classical post-processing - to the three sample arrays. SampleMP is currently the only - measurement type whose post-processing is wired up; for ProbsMP - and CountsMP the three raw outputs are returned as a tuple. + calibration_matrices: optional precomputed per-qubit confusion stack + of shape ``(n_qubits, 2, 2)``. When ``None`` the lowering pass + emits the all-zeros and all-ones calibration circuits and this + callable computes the matrices from their outputs via + ``rem_calibrate_*``. When provided, ``rem_calibrate_*`` is not + traced and the supplied matrices are passed straight into + ``rem_apply_to_*``. + return_calibration_matrices (bool): when ``True`` the callable + returns ``(mitigated_result, calibration_matrices)`` instead of + just the mitigated result. Raises: TypeError: Non-QNode object was passed as ``fn``. """ - def __init__(self, fn: Callable, compute_all_zeroes_ones: bool = False): + def __init__( + self, + fn: Callable, + calibration_matrices=None, + return_calibration_matrices: bool = False, + ): functools.update_wrapper(self, fn) self.fn = fn self.__name__ = f"rem.{getattr(fn, '__name__', 'unknown')}" - self.compute_all_zeroes_ones = bool(compute_all_zeroes_ones) + self.calibration_matrices = calibration_matrices + self.return_calibration_matrices = bool(return_calibration_matrices) super().__init__("fn") def __call__(self, *args, **kwargs): @@ -319,18 +336,14 @@ def __call__(self, *args, **kwargs): assert ( mp_kind != None ), "measurement process must be one of CountsMP, ProbsMP or SampleMP. Other measurement processes such as observables are not supported yet." + run_calibration = self.calibration_matrices is None rem_results = rem_p.bind( *args_data, - compute_all_zeroes_ones=self.compute_all_zeroes_ones, + run_calibration=run_calibration, jaxpr=jaxpr, fn=callable_fn, ) - # With calibration disabled the lowering forwards only the callee - # result plus two placeholder zero tensors, so surface just the result. - if not self.compute_all_zeroes_ones: - return rem_results[0] - qnode_obj = jaxpr.eqns[0].params.get("qnode", None) assert qnode_obj is not None, "REM post-processing requires a QNode target" n_qubits = len(qnode_obj.device.wires) @@ -340,42 +353,49 @@ def __call__(self, *args, **kwargs): # quantum.sample returns f64 in catalyst; the REM helpers index # confusion matrices by bit value and need an integer dtype. mitigatee_samples = rem_results[0].astype(jnp.int32) - zeros_samples = rem_results[1].astype(jnp.int32) - ones_samples = rem_results[2].astype(jnp.int32) - - confusion_matrices = rem_calibrate_samples(zeros_samples, ones_samples) + if run_calibration: + zeros_samples = rem_results[1].astype(jnp.int32) + ones_samples = rem_results[2].astype(jnp.int32) + confusion_matrices = rem_calibrate_samples(zeros_samples, ones_samples) + else: + confusion_matrices = jnp.asarray(self.calibration_matrices) unique_bitstrings, mitigated_counts = rem_apply_to_samples( mitigatee_samples, confusion_matrices, measured_qubits, n_qubits ) - return unique_bitstrings, mitigated_counts + mitigated_result = (unique_bitstrings, mitigated_counts) elif mp_kind == "counts": mitigatee_eigvals = rem_results[0] mitigatee_counts = rem_results[1] - zeros_counts = rem_results[2] - ones_counts = rem_results[3] - confusion_matrices = rem_calibrate_counts(zeros_counts, ones_counts) + if run_calibration: + zeros_counts = rem_results[2] + ones_counts = rem_results[3] + confusion_matrices = rem_calibrate_counts(zeros_counts, ones_counts) + else: + confusion_matrices = jnp.asarray(self.calibration_matrices) mitigated_counts = rem_apply_to_counts( mitigatee_counts, confusion_matrices, measured_qubits, n_qubits ) - return (mitigatee_eigvals, mitigated_counts) + mitigated_result = (mitigatee_eigvals, mitigated_counts) elif mp_kind == "probs": mitigatee_probs = rem_results[0] - zeros_probs = rem_results[1] - ones_probs = rem_results[2] - - confusion_matrices = rem_calibrate_probs(zeros_probs, ones_probs) + if run_calibration: + zeros_probs = rem_results[1] + ones_probs = rem_results[2] + confusion_matrices = rem_calibrate_probs(zeros_probs, ones_probs) + else: + confusion_matrices = jnp.asarray(self.calibration_matrices) mitigated_probs = rem_apply_to_probs( mitigatee_probs, confusion_matrices, measured_qubits, n_qubits ) - return mitigated_probs + mitigated_result = mitigated_probs - # ProbsMP / CountsMP post-processing is not yet wired up; return the - # three raw outputs as a tuple so user code can post-process them. - return tuple(rem_results) + if self.return_calibration_matrices: + return mitigated_result, confusion_matrices + return mitigated_result def polynomial_extrapolation(degree): diff --git a/frontend/catalyst/jax_primitives.py b/frontend/catalyst/jax_primitives.py index b6001ff4ac..6d47fbfc06 100644 --- a/frontend/catalyst/jax_primitives.py +++ b/frontend/catalyst/jax_primitives.py @@ -1086,13 +1086,13 @@ def _zne_lowering(ctx, *args, folding, jaxpr, fn): @rem_p.def_impl -def _rem_def_impl(ctx, *args, compute_all_zeroes_ones, jaxpr, fn): # pragma: no cover +def _rem_def_impl(ctx, *args, run_calibration, jaxpr, fn): # pragma: no cover raise NotImplementedError() @rem_p.def_abstract_eval def _rem_abstract_eval( - *args, compute_all_zeroes_ones, jaxpr, fn + *args, run_calibration, jaxpr, fn ): # pylint: disable=unused-argument """Abstract eval for the REM primitive. @@ -1142,11 +1142,11 @@ def _rem_abstract_eval( ) -def _rem_lowering(ctx, *args, compute_all_zeroes_ones, jaxpr, fn): +def _rem_lowering(ctx, *args, run_calibration, jaxpr, fn): """Lowering function for the REM primitive. Emits a ``mitigation.rem`` op with the callee symbol-ref + the boolean - ``computeAllZeroesOnes`` attribute. The op has one variadic callee result + ``runCalibration`` attribute. The op has one variadic callee result plus two fixed-shape calibration results (zeros / ones), matching the three avals produced by :func:`_rem_abstract_eval`. """ @@ -1173,7 +1173,7 @@ def _rem_lowering(ctx, *args, compute_all_zeroes_ones, jaxpr, fn): constants.append(constantVals) args_and_consts = constants + list(args) - bool_attr = ir.BoolAttr.get(bool(compute_all_zeroes_ones)) + bool_attr = ir.BoolAttr.get(bool(run_calibration)) rem_op = RemOp( flat_callee_output_types, @@ -1181,7 +1181,7 @@ def _rem_lowering(ctx, *args, compute_all_zeroes_ones, jaxpr, fn): flat_ones_output_types, symbol_ref, mlir.flatten_ir_values(args_and_consts), - computeAllZeroesOnes=bool_attr, + runCalibration=bool_attr, ) return rem_op.results diff --git a/frontend/test/lit/test_mitigation_rem.py b/frontend/test/lit/test_mitigation_rem.py index 0fc4fd5123..7a204b0ddb 100644 --- a/frontend/test/lit/test_mitigation_rem.py +++ b/frontend/test/lit/test_mitigation_rem.py @@ -1,4 +1,4 @@ -# Copyright 2026 Haiqu, Inc. +# Copyright 2025 Haiqu, Inc. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -20,6 +20,7 @@ # RUN: %PYTHON %s | FileCheck %s +import jax.numpy as jnp import pennylane as qp from catalyst import mitigate_with_rem, qjit @@ -38,11 +39,11 @@ def circuit(): return qp.probs(wires=[0, 1]) qnode = qp.set_shots(qp.QNode(circuit, dev), shots=200) - return mitigate_with_rem(qnode, compute_all_zeroes_ones=True)() + return mitigate_with_rem(qnode)() # CHECK-LABEL: func.func public @jit_rem_with_probs -# CHECK: mitigation.rem @{{[^(]+}}() computeAllZeroesOnes(true) +# CHECK: mitigation.rem @{{[^(]+}}() runCalibration(true) # CHECK-SAME: -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64>) # Post-processing: # calibration consumes the two f64 calibration tensors and yields the (n_qubits, 2, 2) confusion stack. @@ -65,11 +66,11 @@ def circuit(): return qp.counts() qnode = qp.set_shots(qp.QNode(circuit, dev), shots=200) - return mitigate_with_rem(qnode, compute_all_zeroes_ones=True)() + return mitigate_with_rem(qnode)() # CHECK-LABEL: func.func public @jit_rem_with_counts -# CHECK: mitigation.rem @{{[^(]+}}() computeAllZeroesOnes(true) +# CHECK: mitigation.rem @{{[^(]+}}() runCalibration(true) # CHECK-SAME: -> (tensor<4xi64>, tensor<4xi64>, tensor<4xi64>, tensor<4xi64>) # Post-processing: counts path reuses the probs calibration after normalization, # and apply returns f64 because the linear solve is float-valued. @@ -91,11 +92,11 @@ def circuit(): return qp.sample() qnode = qp.set_shots(qp.QNode(circuit, dev), shots=200) - return mitigate_with_rem(qnode, compute_all_zeroes_ones=True)() + return mitigate_with_rem(qnode)() # CHECK-LABEL: func.func public @jit_rem_with_sample_histogram -# CHECK: mitigation.rem @{{[^(]+}}() computeAllZeroesOnes(true) +# CHECK: mitigation.rem @{{[^(]+}}() runCalibration(true) # CHECK-SAME: -> (tensor<200x2xi64>, tensor<200x2xf64>, tensor<200x2xf64>) # Post-processing, histogram path: calibrate sees (shots, qubits) i32 samples, # apply returns ((K, n_qubits) bitstrings, (K,) counts) with K = 2**n_qubits = 4. @@ -118,11 +119,11 @@ def circuit(): return qp.sample() qnode = qp.set_shots(qp.QNode(circuit, dev), shots=4) - return mitigate_with_rem(qnode, compute_all_zeroes_ones=True)() + return mitigate_with_rem(qnode)() # CHECK-LABEL: func.func public @jit_rem_with_sample_sort_rle -# CHECK: mitigation.rem @{{[^(]+}}() computeAllZeroesOnes(true) +# CHECK: mitigation.rem @{{[^(]+}}() runCalibration(true) # CHECK-SAME: -> (tensor<4x4xi64>, tensor<4x4xf64>, tensor<4x4xf64>) # Post-processing, sort-RLE path: apply output is (K = n_shots = 4, n_qubits) and (K,). # CHECK: call @rem_calibrate_samples({{.*}}) : (tensor<4x4xi32>, tensor<4x4xi32>) -> tensor<4x2x2xf64> @@ -133,8 +134,9 @@ def circuit(): @qjit(target="mlir") -def rem_passthrough(): - """compute_all_zeroes_ones=False still emits the op, just with the false attribute and no postprocessing helpers.""" +def rem_cached(): + """Passing precomputed calibration matrices runs the op with runCalibration(false) + and emits rem_apply_to_probs but not rem_calibrate_probs.""" dev = qp.device("lightning.qubit", wires=2) def circuit(): @@ -143,11 +145,13 @@ def circuit(): return qp.probs(wires=[0, 1]) qnode = qp.set_shots(qp.QNode(circuit, dev), shots=200) - return mitigate_with_rem(qnode, compute_all_zeroes_ones=False)() + cm = jnp.broadcast_to(jnp.eye(2, dtype=jnp.float64), (2, 2, 2)) + return mitigate_with_rem(qnode, calibration_matrices=cm)() -# CHECK-LABEL: func.func public @jit_rem_passthrough -# CHECK: mitigation.rem @{{[^(]+}}() computeAllZeroesOnes(false) +# CHECK-LABEL: func.func public @jit_rem_cached +# CHECK: mitigation.rem @{{[^(]+}}() runCalibration(false) +# CHECK: call @rem_apply_to_probs({{.*}}) : (tensor<4xf64>, tensor<2x2x2xf64>, tensor<2xi64>) -> tensor<4xf64> # CHECK-NOT: call @rem_calibrate -# CHECK-NOT: call @rem_apply -print(rem_passthrough.mlir) +# CHECK-NOT: func.func private @rem_calibrate +print(rem_cached.mlir) diff --git a/frontend/test/pytest/test_mitigation_rem.py b/frontend/test/pytest/test_mitigation_rem.py index e6ae46a137..9dcb5a41b0 100644 --- a/frontend/test/pytest/test_mitigation_rem.py +++ b/frontend/test/pytest/test_mitigation_rem.py @@ -60,7 +60,7 @@ def test_probs(n_qubits): @catalyst.qjit def mitigated(): - return mitigate_with_rem(circuit, compute_all_zeroes_ones=True)() + return mitigate_with_rem(circuit)() out = np.asarray(mitigated()) expected = np.zeros(2**n_qubits) @@ -77,7 +77,7 @@ def test_counts(n_qubits): @catalyst.qjit def mitigated(): - return mitigate_with_rem(circuit, compute_all_zeroes_ones=True)() + return mitigate_with_rem(circuit)() eigvals, counts = mitigated() counts = np.asarray(counts) @@ -101,7 +101,7 @@ def test_sample(n_qubits): @catalyst.qjit def mitigated(): - return mitigate_with_rem(circuit, compute_all_zeroes_ones=True)() + return mitigate_with_rem(circuit)() bitstrings, counts = mitigated() bitstrings = np.asarray(bitstrings) @@ -118,22 +118,113 @@ def mitigated(): assert histogram[0] + histogram[-1] > 0.9 * shots -def test_no_calibration(): - """compute_all_zeroes_ones=False forwards the raw callee result without REM post-processing.""" - n_qubits, shots = 2, 1000 +def test_return_calibration_matrices(): + """return_calibration_matrices=True wraps the mitigated probs in a (result, cm) tuple.""" + n_qubits, shots = 2, 2000 circuit = _ghz_qnode(n_qubits, shots, "probs") @catalyst.qjit - def passthrough(): - return mitigate_with_rem(circuit, compute_all_zeroes_ones=False)() + def mitigated(): + return mitigate_with_rem(circuit, return_calibration_matrices=True)() + + result, cm = mitigated() + cm = np.asarray(cm) + assert cm.shape == (n_qubits, 2, 2) + # Each 2x2 confusion matrix is column-stochastic. + assert np.allclose(cm.sum(axis=-2), 1.0, atol=1e-6) + + expected = np.zeros(2**n_qubits) + expected[0] = 0.5 + expected[-1] = 0.5 + assert np.allclose(np.asarray(result), expected, atol=5e-2) + + +def test_caching_probs_roundtrip(): + """Compute calibration once, reuse it on a second qjit call; result still matches the GHZ probs.""" + n_qubits, shots = 2, 2000 + circuit = _ghz_qnode(n_qubits, shots, "probs") + + @catalyst.qjit + def fresh(): + return mitigate_with_rem(circuit, return_calibration_matrices=True)() + + _, cm = fresh() - out = np.asarray(passthrough()) + @catalyst.qjit + def cached(cm_in): + return mitigate_with_rem(circuit, calibration_matrices=cm_in)() + + out = np.asarray(cached(cm)) expected = np.zeros(2**n_qubits) expected[0] = 0.5 expected[-1] = 0.5 assert np.allclose(out, expected, atol=5e-2) +def test_caching_counts_roundtrip(): + """Counts MP: cached calibration matrices yield concentrated GHZ peaks on a noiseless device.""" + n_qubits, shots = 2, 2000 + circuit = _ghz_qnode(n_qubits, shots, "counts") + + @catalyst.qjit + def fresh(): + return mitigate_with_rem(circuit, return_calibration_matrices=True)() + + _, cm = fresh() + + @catalyst.qjit + def cached(cm_in): + return mitigate_with_rem(circuit, calibration_matrices=cm_in)() + + eigvals, counts = cached(cm) + counts = np.asarray(counts) + assert counts.shape == (2**n_qubits,) + assert counts[0] + counts[-1] > 0.9 * shots + assert np.allclose(np.asarray(eigvals), np.arange(2**n_qubits)) + + +def test_caching_sample_roundtrip(): + """Sample MP histogram-path cache reuse on a noiseless GHZ device.""" + n_qubits, shots = 2, 2000 + circuit = _ghz_qnode(n_qubits, shots, "sample") + + @catalyst.qjit + def fresh(): + return mitigate_with_rem(circuit, return_calibration_matrices=True)() + + _, cm = fresh() + + @catalyst.qjit + def cached(cm_in): + return mitigate_with_rem(circuit, calibration_matrices=cm_in)() + + bitstrings, counts = cached(cm) + bitstrings = np.asarray(bitstrings) + counts = np.asarray(counts) + powers = 1 << np.arange(n_qubits - 1, -1, -1) + codes = (bitstrings * powers[None, :]).sum(axis=1) + histogram = np.zeros(2**n_qubits) + for code, c in zip(codes, counts): + histogram[code] += c + assert histogram[0] + histogram[-1] > 0.9 * shots + + +def test_caching_skips_rem_calibrate_in_ir(): + """Cached path must not emit rem_calibrate_* helpers in the traced MLIR.""" + n_qubits, shots = 2, 200 + circuit = _ghz_qnode(n_qubits, shots, "probs") + cm = jnp.broadcast_to(jnp.eye(2, dtype=jnp.float64), (n_qubits, 2, 2)) + + @catalyst.qjit(target="mlir") + def cached(): + return mitigate_with_rem(circuit, calibration_matrices=cm)() + + ir = cached.mlir + assert "runCalibration(false)" in ir + assert "rem_apply_to_probs" in ir + assert "rem_calibrate_probs" not in ir + + def _per_qubit_symmetric_confusion(n_qubits, p): """Stack of ``(n_qubits, 2, 2)`` symmetric bit-flip channels with flip probability ``p``.""" one = np.array([[1.0 - p, p], [p, 1.0 - p]]) @@ -281,7 +372,7 @@ def circuit(): qnode = qp.set_shots(qp.QNode(circuit, dev), shots=200) def mitigated(): - return mitigate_with_rem(qnode, compute_all_zeroes_ones=True)() + return mitigate_with_rem(qnode)() with pytest.raises(AssertionError, match="measurement process must be one of"): catalyst.qjit(mitigated) diff --git a/mlir/include/Mitigation/IR/MitigationOps.td b/mlir/include/Mitigation/IR/MitigationOps.td index e8d130d92a..ad225b9fd5 100644 --- a/mlir/include/Mitigation/IR/MitigationOps.td +++ b/mlir/include/Mitigation/IR/MitigationOps.td @@ -65,19 +65,20 @@ def RemOp : Mitigation_Op<"rem", [DeclareOpInterfaceMethods, let description = [{ The `mitigation.rem` operation wraps a quantum function call so the mitigation-lowering pass can insert calibration circuits next to it. - When `computeAllZeroesOnes` is true the pass emits an all-zeros and an + When `runCalibration` is true the pass emits an all-zeros and an all-ones calibration circuit alongside the original callee; their sample/probs results are returned as the `zeros` and `ones` SSA values and consumed by the classical REM post-processing on the JAX side. When false the calibration outputs are zero-filled placeholder tensors - of the same type, which leaves a hook for a future cached-confusion- - matrices path without changing the op's result signature. + of the same type, which lets the caller supply cached confusion + matrices and skip the calibration circuits without changing the op's + result signature. }]; let arguments = (ins SymbolRefAttr:$callee, Variadic:$args, - BoolAttr:$computeAllZeroesOnes, + BoolAttr:$runCalibration, OptionalAttr:$arg_attrs, OptionalAttr:$res_attrs ); @@ -87,7 +88,7 @@ def RemOp : Mitigation_Op<"rem", [DeclareOpInterfaceMethods, RankedTensorOf<[AnyFloat, AnyInteger]>:$ones ); let assemblyFormat = [{ - $callee `(` $args `)` `computeAllZeroesOnes` `(` $computeAllZeroesOnes `)` attr-dict `:` functional-type($args, results) + $callee `(` $args `)` `runCalibration` `(` $runCalibration `)` attr-dict `:` functional-type($args, results) }]; } diff --git a/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp b/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp index 51c938cbfa..8ea239baf6 100644 --- a/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp +++ b/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp @@ -34,8 +34,8 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter // with proper calibration circuits in a subsequent step. Location loc = op.getLoc(); - auto computeAllZeroesOnes = op.getComputeAllZeroesOnesAttr(); - llvm::dbgs() << "[mitigation.rem] computeAllZeroesOnes: " << computeAllZeroesOnes << "\n"; + auto runCalibration = op.getRunCalibrationAttr(); + llvm::dbgs() << "[mitigation.rem] runCalibration: " << runCalibration << "\n"; for (auto v : op.getResultTypes()) { llvm::dbgs() << "[mitigation.rem] Result type: " << v << "\n"; } @@ -90,7 +90,7 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter llvm::dbgs() << "[mitigation.rem] MP name: " << MPName << "\n"; // pre-compute MP and attribute-dependent values - bool doCalib = computeAllZeroesOnes.getValue(); // get boolean value of BoolAttr + bool doCalib = runCalibration.getValue(); // get boolean value of BoolAttr mlir::RankedTensorType tensorTyI64; mlir::RankedTensorType tensorTyF64; int64_t calibrationTensorShape = 0; diff --git a/mlir/test/Mitigation/RemLoweringTest.mlir b/mlir/test/Mitigation/RemLoweringTest.mlir index b357dcae9b..741a81cdda 100644 --- a/mlir/test/Mitigation/RemLoweringTest.mlir +++ b/mlir/test/Mitigation/RemLoweringTest.mlir @@ -15,7 +15,7 @@ // RUN: quantum-opt %s --lower-mitigation --split-input-file --verify-diagnostics | FileCheck %s // =================================================================== -// computeAllZeroesOnes(true) on a ProbsMP callee: the pass must keep the +// runCalibration(true) on a ProbsMP callee: the pass must keep the // original func.call, then emit two cloned quantum.device / quantum.alloc / // quantum.compbasis / quantum.probs blocks. The all-ones clone is // distinguished from the all-zeroes one by per-qubit PauliX gates inserted @@ -41,7 +41,7 @@ func.func @probsCircuit() -> tensor<4xf64> attributes {qnode} { } func.func @remProbsCalibrate() -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64>) { - %out:3 = mitigation.rem @probsCircuit() computeAllZeroesOnes(true) : () -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64>) + %out:3 = mitigation.rem @probsCircuit() runCalibration(true) : () -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64>) func.return %out#0, %out#1, %out#2 : tensor<4xf64>, tensor<4xf64>, tensor<4xf64> } @@ -63,7 +63,7 @@ func.func @remProbsCalibrate() -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64>) // ----- // =================================================================== -// computeAllZeroesOnes(true) on a CountsMP callee: callee result is the +// runCalibration(true) on a CountsMP callee: callee result is the // (eigvals, counts) pair, and the rem op surfaces two i64 calibration // tensors -- one per all-zeroes / all-ones run. // =================================================================== @@ -87,7 +87,7 @@ func.func @countsCircuit() -> (tensor<4xf64>, tensor<4xi64>) attributes {qnode} } func.func @remCountsCalibrate() -> (tensor<4xf64>, tensor<4xi64>, tensor<4xi64>, tensor<4xi64>) { - %out:4 = mitigation.rem @countsCircuit() computeAllZeroesOnes(true) : () -> (tensor<4xf64>, tensor<4xi64>, tensor<4xi64>, tensor<4xi64>) + %out:4 = mitigation.rem @countsCircuit() runCalibration(true) : () -> (tensor<4xf64>, tensor<4xi64>, tensor<4xi64>, tensor<4xi64>) func.return %out#0, %out#1, %out#2, %out#3 : tensor<4xf64>, tensor<4xi64>, tensor<4xi64>, tensor<4xi64> } @@ -100,7 +100,7 @@ func.func @remCountsCalibrate() -> (tensor<4xf64>, tensor<4xi64>, tensor<4xi64>, // ----- // =================================================================== -// computeAllZeroesOnes(true) on a SampleMP callee with shots = 200: the +// runCalibration(true) on a SampleMP callee with shots = 200: the // calibration tensors take the (shots, qubits) shape inferred from the // callee's quantum.device + quantum.alloc. // =================================================================== @@ -124,7 +124,7 @@ func.func @sampleCircuit() -> tensor<200x2xf64> attributes {qnode} { } func.func @remSampleCalibrate() -> (tensor<200x2xf64>, tensor<200x2xf64>, tensor<200x2xf64>) { - %out:3 = mitigation.rem @sampleCircuit() computeAllZeroesOnes(true) : () -> (tensor<200x2xf64>, tensor<200x2xf64>, tensor<200x2xf64>) + %out:3 = mitigation.rem @sampleCircuit() runCalibration(true) : () -> (tensor<200x2xf64>, tensor<200x2xf64>, tensor<200x2xf64>) func.return %out#0, %out#1, %out#2 : tensor<200x2xf64>, tensor<200x2xf64>, tensor<200x2xf64> } @@ -137,7 +137,7 @@ func.func @remSampleCalibrate() -> (tensor<200x2xf64>, tensor<200x2xf64>, tensor // ----- // =================================================================== -// computeAllZeroesOnes(false): the pass forwards the callee unchanged and +// runCalibration(false): the pass forwards the callee unchanged and // stamps two zero-filled placeholder tensors for the calibration slots -- // no cloned circuits and no extra quantum.device calls. // =================================================================== @@ -158,7 +158,7 @@ func.func @passthroughCircuit() -> tensor<4xf64> attributes {qnode} { } func.func @remProbsPassthrough() -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64>) { - %out:3 = mitigation.rem @passthroughCircuit() computeAllZeroesOnes(false) : () -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64>) + %out:3 = mitigation.rem @passthroughCircuit() runCalibration(false) : () -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64>) func.return %out#0, %out#1, %out#2 : tensor<4xf64>, tensor<4xf64>, tensor<4xf64> } @@ -177,9 +177,9 @@ func.func @remProbsPassthrough() -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64> // =================================================================== func.func @remMissingCallee() -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64>) { - %out:3 = mitigation.rem @doesNotExist() computeAllZeroesOnes(true) : () -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64>) + %out:3 = mitigation.rem @doesNotExist() runCalibration(true) : () -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64>) func.return %out#0, %out#1, %out#2 : tensor<4xf64>, tensor<4xf64>, tensor<4xf64> } // CHECK-LABEL: func.func @remMissingCallee -// CHECK: mitigation.rem @doesNotExist() computeAllZeroesOnes(true) +// CHECK: mitigation.rem @doesNotExist() runCalibration(true) From 61a9ca74ebf19a18b359e74e2c068de942bfb37a Mon Sep 17 00:00:00 2001 From: volodymyr-hq Date: Wed, 1 Jul 2026 21:33:55 +0300 Subject: [PATCH 6/8] Change output shape for RemOp to variadic, do not return confusion matrices when calibration is omitted --- frontend/catalyst/jax_primitives.py | 38 +++++++++------- frontend/test/lit/test_mitigation_rem.py | 2 + mlir/include/Mitigation/IR/MitigationOps.td | 7 +-- .../Transforms/MitigationMethods/Rem.cpp | 45 +------------------ mlir/test/Mitigation/RemLoweringTest.mlir | 22 ++++----- 5 files changed, 42 insertions(+), 72 deletions(-) diff --git a/frontend/catalyst/jax_primitives.py b/frontend/catalyst/jax_primitives.py index 6d47fbfc06..aa64df0ac2 100644 --- a/frontend/catalyst/jax_primitives.py +++ b/frontend/catalyst/jax_primitives.py @@ -1118,7 +1118,9 @@ def _rem_abstract_eval( if mp == "probs": prob_size = 2 << (device_qubit_count - 1) prob_aval = core.ShapedArray((prob_size,), dtype=jax.numpy.dtype("float64")) - return (jaxpr.out_avals[0], prob_aval, prob_aval) + if run_calibration: + return (jaxpr.out_avals[0], prob_aval, prob_aval) + return (jaxpr.out_avals[0],) elif mp == "counts": # The mitigation-lowering pass takes the counts half (i64) of each # calibration CountsOp's result pair, so the calibration avals are @@ -1126,7 +1128,9 @@ def _rem_abstract_eval( # replaces the rem op with those values. counts_size = 2 << (device_qubit_count - 1) counts_aval = core.ShapedArray((counts_size,), dtype=jax.numpy.dtype("int64")) - return (jaxpr.out_avals[0], jaxpr.out_avals[1], counts_aval, counts_aval) + if run_calibration: + return (jaxpr.out_avals[0], jaxpr.out_avals[1], counts_aval, counts_aval) + return (jaxpr.out_avals[0], jaxpr.out_avals[1]) elif mp == "sample": # Catalyst's quantum.sample lowering produces an f64 tensor (see # `_sample_lowering`). The Mitigation pass also creates calibration @@ -1136,7 +1140,9 @@ def _rem_abstract_eval( sample_aval = core.ShapedArray( (device_shot_count, device_qubit_count), dtype=jax.numpy.dtype("float64") ) - return (jaxpr.out_avals[0], sample_aval, sample_aval) + if run_calibration: + return (jaxpr.out_avals[0], sample_aval, sample_aval) + return (jaxpr.out_avals[0],) raise NotImplementedError( f"REM only supports probs, counts and sample measurement processes; got {mp!r}" ) @@ -1146,21 +1152,23 @@ def _rem_lowering(ctx, *args, run_calibration, jaxpr, fn): """Lowering function for the REM primitive. Emits a ``mitigation.rem`` op with the callee symbol-ref + the boolean - ``runCalibration`` attribute. The op has one variadic callee result - plus two fixed-shape calibration results (zeros / ones), matching the - three avals produced by :func:`_rem_abstract_eval`. + ``runCalibration`` attribute. All three result groups (callee_result, + zeros, ones) are variadic; the zeros/ones groups are empty when + ``run_calibration`` is False so the op leaves no calibration SSA values in + the lowered IR. """ func_op = lower_jaxpr(ctx, jaxpr) symbol_ref = get_symbolref(ctx, func_op) - *callee_avals, zeros_aval, ones_aval = ctx.avals_out + if run_calibration: + *callee_avals, zeros_aval, ones_aval = ctx.avals_out + zeros_output_types = util.flatten([mlir.aval_to_ir_types(zeros_aval)]) + ones_output_types = util.flatten([mlir.aval_to_ir_types(ones_aval)]) + else: + callee_avals = list(ctx.avals_out) + zeros_output_types = [] + ones_output_types = [] callee_output_types = list(map(mlir.aval_to_ir_types, callee_avals)) - zeros_output_types = [mlir.aval_to_ir_types(zeros_aval)] - ones_output_types = [mlir.aval_to_ir_types(ones_aval)] flat_callee_output_types = util.flatten(callee_output_types) - # The zeros/ones operands are non-variadic in the .td, so a single type - # (not a list) is what the op constructor expects. - flat_zeros_output_types = util.flatten(zeros_output_types)[0] - flat_ones_output_types = util.flatten(ones_output_types)[0] constants = [] for const in jaxpr.consts: @@ -1177,8 +1185,8 @@ def _rem_lowering(ctx, *args, run_calibration, jaxpr, fn): rem_op = RemOp( flat_callee_output_types, - flat_zeros_output_types, - flat_ones_output_types, + zeros_output_types, + ones_output_types, symbol_ref, mlir.flatten_ir_values(args_and_consts), runCalibration=bool_attr, diff --git a/frontend/test/lit/test_mitigation_rem.py b/frontend/test/lit/test_mitigation_rem.py index 7a204b0ddb..fff46e3c78 100644 --- a/frontend/test/lit/test_mitigation_rem.py +++ b/frontend/test/lit/test_mitigation_rem.py @@ -150,7 +150,9 @@ def circuit(): # CHECK-LABEL: func.func public @jit_rem_cached +# The op's result tuple must shrink to only the callee's return, no zeros/ones slots. # CHECK: mitigation.rem @{{[^(]+}}() runCalibration(false) +# CHECK-SAME: -> tensor<4xf64> # CHECK: call @rem_apply_to_probs({{.*}}) : (tensor<4xf64>, tensor<2x2x2xf64>, tensor<2xi64>) -> tensor<4xf64> # CHECK-NOT: call @rem_calibrate # CHECK-NOT: func.func private @rem_calibrate diff --git a/mlir/include/Mitigation/IR/MitigationOps.td b/mlir/include/Mitigation/IR/MitigationOps.td index ad225b9fd5..1e22b3fbba 100644 --- a/mlir/include/Mitigation/IR/MitigationOps.td +++ b/mlir/include/Mitigation/IR/MitigationOps.td @@ -60,7 +60,8 @@ def ZneOp : Mitigation_Op<"zne", [DeclareOpInterfaceMethods, } def RemOp : Mitigation_Op<"rem", [DeclareOpInterfaceMethods, - DeclareOpInterfaceMethods]> { + DeclareOpInterfaceMethods, + AttrSizedResultSegments]> { let summary = "Compute a quantum function with Readout Error Mitigation (REM)."; let description = [{ The `mitigation.rem` operation wraps a quantum function call so the @@ -84,8 +85,8 @@ def RemOp : Mitigation_Op<"rem", [DeclareOpInterfaceMethods, ); let results = (outs Variadic]>>:$callee_result, - RankedTensorOf<[AnyFloat, AnyInteger]>:$zeros, - RankedTensorOf<[AnyFloat, AnyInteger]>:$ones + Variadic>:$zeros, + Variadic>:$ones ); let assemblyFormat = [{ $callee `(` $args `)` `runCalibration` `(` $runCalibration `)` attr-dict `:` functional-type($args, results) diff --git a/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp b/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp index 8ea239baf6..d9ad48c58f 100644 --- a/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp +++ b/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp @@ -141,51 +141,10 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter llvm::dbgs() << "[mitigation.rem] doCalibration == false, replacing RemOp with wrapped " "callee circuit function call..." << "\n"; - mlir::RankedTensorType tensorTyConst; - // Create constant array of type ranked tensor initialized with 0.0. This is the - // default value returned when doCalib == false and indicates that no calibration circuits - // were run auto tensorTy = RankedTensorType::get({bitspaceShape}, rewriter.getF64Type()); - // SmallVector zeros_vec(bitspaceShape, 0.0); //initialize with zeroes - DenseElementsAttr tensorAttr; - if (MPName == "quantum.probs"s || MPName == "quantum.counts"s) { - SmallVector zerosVector(calibrationTensorShape, 0.0); // initialize with zeroes - // To initialize with other values: - // for (auto i = 0; i < 4; ++i) { - // zeros_vec[i] = static_cast(i); - // } - // ----------------------------------------------------------------- - // ArrayRef is a non-owning view and needs to be created because it is often required by - // MLIR APIs to avoid copies. It is also possible to use this: auto zerosAttr = - // rewriter.getF64ArrayAttr(zeros_vec); // returns ArrayAttr used by DenseElementsAttr - // because there is also an overload of DenseElementsAttr which accepts ArrayAttr. - // Attributes and compile-time constant metadata attached to ops/types in the IR, it - // requires a static shape - tensorTyConst = tensorTyF64; - tensorAttr = DenseElementsAttr::get( - tensorTyConst, ArrayRef{zerosVector.begin(), zerosVector.end()}); - // tensorAttr is an attribute attached with the newly created airth.constant operation - // on the line below. It represents values known at pass-time (compile-time with regards - // to IR generation, but could be inferred at compile time of the passs from other - // attributes or something else) - } - else if (MPName == "quantum.sample"s) { - SmallVector zerosVector(calibrationTensorShape, 0); // initialize with zeroes - tensorTyConst = tensorTyI64; - tensorAttr = DenseElementsAttr::get( - tensorTyConst, ArrayRef{zerosVector.begin(), zerosVector.end()}); - } - auto cst_zeros = rewriter.create(loc, tensorTyConst, tensorAttr); - llvm::dbgs() << "[mitigation.rem] cst_zeros = " << cst_zeros << "\n"; - - // 1) original callee results + // The zeros/ones result groups are empty variadics when doCalib is false, + // so the op is replaced by only the callee's SSA values. for (Value v : results) completeResults.push_back(v); - // 2) zeros calibration placeholders (reuse same tensors for now) - completeResults.push_back(cst_zeros.getResult()); - // 3) ones calibration placeholders - completeResults.push_back(cst_zeros.getResult()); - // Replace the results of the original operation (mitigation.rem) with this vector of Values - // (could be any object castable to ValueRange) rewriter.replaceOp(op, completeResults); return success(); } diff --git a/mlir/test/Mitigation/RemLoweringTest.mlir b/mlir/test/Mitigation/RemLoweringTest.mlir index 741a81cdda..ea19cfb0ef 100644 --- a/mlir/test/Mitigation/RemLoweringTest.mlir +++ b/mlir/test/Mitigation/RemLoweringTest.mlir @@ -41,7 +41,7 @@ func.func @probsCircuit() -> tensor<4xf64> attributes {qnode} { } func.func @remProbsCalibrate() -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64>) { - %out:3 = mitigation.rem @probsCircuit() runCalibration(true) : () -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64>) + %out:3 = mitigation.rem @probsCircuit() runCalibration(true) {resultSegmentSizes = array} : () -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64>) func.return %out#0, %out#1, %out#2 : tensor<4xf64>, tensor<4xf64>, tensor<4xf64> } @@ -87,7 +87,7 @@ func.func @countsCircuit() -> (tensor<4xf64>, tensor<4xi64>) attributes {qnode} } func.func @remCountsCalibrate() -> (tensor<4xf64>, tensor<4xi64>, tensor<4xi64>, tensor<4xi64>) { - %out:4 = mitigation.rem @countsCircuit() runCalibration(true) : () -> (tensor<4xf64>, tensor<4xi64>, tensor<4xi64>, tensor<4xi64>) + %out:4 = mitigation.rem @countsCircuit() runCalibration(true) {resultSegmentSizes = array} : () -> (tensor<4xf64>, tensor<4xi64>, tensor<4xi64>, tensor<4xi64>) func.return %out#0, %out#1, %out#2, %out#3 : tensor<4xf64>, tensor<4xi64>, tensor<4xi64>, tensor<4xi64> } @@ -124,7 +124,7 @@ func.func @sampleCircuit() -> tensor<200x2xf64> attributes {qnode} { } func.func @remSampleCalibrate() -> (tensor<200x2xf64>, tensor<200x2xf64>, tensor<200x2xf64>) { - %out:3 = mitigation.rem @sampleCircuit() runCalibration(true) : () -> (tensor<200x2xf64>, tensor<200x2xf64>, tensor<200x2xf64>) + %out:3 = mitigation.rem @sampleCircuit() runCalibration(true) {resultSegmentSizes = array} : () -> (tensor<200x2xf64>, tensor<200x2xf64>, tensor<200x2xf64>) func.return %out#0, %out#1, %out#2 : tensor<200x2xf64>, tensor<200x2xf64>, tensor<200x2xf64> } @@ -137,9 +137,9 @@ func.func @remSampleCalibrate() -> (tensor<200x2xf64>, tensor<200x2xf64>, tensor // ----- // =================================================================== -// runCalibration(false): the pass forwards the callee unchanged and -// stamps two zero-filled placeholder tensors for the calibration slots -- -// no cloned circuits and no extra quantum.device calls. +// runCalibration(false): the pass forwards the callee unchanged. With the +// zeros/ones result groups empty, the op replacement is just the callee's +// SSA values -- no placeholder constants and no calibration circuits. // =================================================================== func.func @passthroughCircuit() -> tensor<4xf64> attributes {qnode} { @@ -157,14 +157,14 @@ func.func @passthroughCircuit() -> tensor<4xf64> attributes {qnode} { func.return %probs : tensor<4xf64> } -func.func @remProbsPassthrough() -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64>) { - %out:3 = mitigation.rem @passthroughCircuit() runCalibration(false) : () -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64>) - func.return %out#0, %out#1, %out#2 : tensor<4xf64>, tensor<4xf64>, tensor<4xf64> +func.func @remProbsPassthrough() -> tensor<4xf64> { + %out = mitigation.rem @passthroughCircuit() runCalibration(false) {resultSegmentSizes = array} : () -> tensor<4xf64> + func.return %out : tensor<4xf64> } // CHECK-LABEL: func.func @remProbsPassthrough -// CHECK: arith.constant dense<0.000000e+00> : tensor<4xf64> // CHECK: call @passthroughCircuit() : () -> tensor<4xf64> +// CHECK-NOT: arith.constant dense<0.000000e+00> // CHECK-NOT: quantum.compbasis // CHECK-NOT: quantum.custom "PauliX" @@ -177,7 +177,7 @@ func.func @remProbsPassthrough() -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64> // =================================================================== func.func @remMissingCallee() -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64>) { - %out:3 = mitigation.rem @doesNotExist() runCalibration(true) : () -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64>) + %out:3 = mitigation.rem @doesNotExist() runCalibration(true) {resultSegmentSizes = array} : () -> (tensor<4xf64>, tensor<4xf64>, tensor<4xf64>) func.return %out#0, %out#1, %out#2 : tensor<4xf64>, tensor<4xf64>, tensor<4xf64> } From 20df5a3e5230c86c223d7319f24202fa9485f35b Mon Sep 17 00:00:00 2001 From: volodymyr-hq Date: Thu, 2 Jul 2026 13:01:24 +0300 Subject: [PATCH 7/8] Refactor to satisfy CodeFactor complaints --- .../api_extensions/error_mitigation.py | 136 +++--- .../api_extensions/rem_postprocessing.py | 2 - frontend/catalyst/jax_primitives.py | 4 +- frontend/test/pytest/test_mitigation_rem.py | 4 +- .../Transforms/MitigationMethods/Rem.cpp | 401 +++++++----------- 5 files changed, 235 insertions(+), 312 deletions(-) diff --git a/frontend/catalyst/api_extensions/error_mitigation.py b/frontend/catalyst/api_extensions/error_mitigation.py index a3d12d8dd9..9d6763fd51 100644 --- a/frontend/catalyst/api_extensions/error_mitigation.py +++ b/frontend/catalyst/api_extensions/error_mitigation.py @@ -319,23 +319,12 @@ def __call__(self, *args, **kwargs): callable_fn ), "expected callable set as param on the first operation in rem target" - # Identify the measurement process by walking the inner jaxpr, matching - # the logic in `_rem_abstract_eval`. - mp_kind = None - for eqn in jaxpr.eqns: - inner = eqn.params.get("call_jaxpr", None) - if inner is None: - continue - for op_eq in inner.eqns: - pname = str(op_eq.primitive) - if pname in ("probs", "sample", "counts"): - mp_kind = pname - break - if mp_kind is not None: - break - assert ( - mp_kind != None - ), "measurement process must be one of CountsMP, ProbsMP or SampleMP. Other measurement processes such as observables are not supported yet." + mp_kind = _detect_measurement_process(jaxpr) + assert mp_kind is not None, ( + "measurement process must be one of CountsMP, ProbsMP or SampleMP. " + "Other measurement processes such as observables are not supported yet." + ) + run_calibration = self.calibration_matrices is None rem_results = rem_p.bind( *args_data, @@ -347,56 +336,63 @@ def __call__(self, *args, **kwargs): qnode_obj = jaxpr.eqns[0].params.get("qnode", None) assert qnode_obj is not None, "REM post-processing requires a QNode target" n_qubits = len(qnode_obj.device.wires) - measured_qubits = jnp.array([x for x in qnode_obj.device.wires]) - if mp_kind == "sample": - - # quantum.sample returns f64 in catalyst; the REM helpers index - # confusion matrices by bit value and need an integer dtype. - mitigatee_samples = rem_results[0].astype(jnp.int32) - if run_calibration: - zeros_samples = rem_results[1].astype(jnp.int32) - ones_samples = rem_results[2].astype(jnp.int32) - confusion_matrices = rem_calibrate_samples(zeros_samples, ones_samples) - else: - confusion_matrices = jnp.asarray(self.calibration_matrices) - unique_bitstrings, mitigated_counts = rem_apply_to_samples( - mitigatee_samples, confusion_matrices, measured_qubits, n_qubits - ) - mitigated_result = (unique_bitstrings, mitigated_counts) - - elif mp_kind == "counts": - - mitigatee_eigvals = rem_results[0] - mitigatee_counts = rem_results[1] - if run_calibration: - zeros_counts = rem_results[2] - ones_counts = rem_results[3] - confusion_matrices = rem_calibrate_counts(zeros_counts, ones_counts) - else: - confusion_matrices = jnp.asarray(self.calibration_matrices) - mitigated_counts = rem_apply_to_counts( - mitigatee_counts, confusion_matrices, measured_qubits, n_qubits - ) - mitigated_result = (mitigatee_eigvals, mitigated_counts) - - elif mp_kind == "probs": - - mitigatee_probs = rem_results[0] - if run_calibration: - zeros_probs = rem_results[1] - ones_probs = rem_results[2] - confusion_matrices = rem_calibrate_probs(zeros_probs, ones_probs) - else: - confusion_matrices = jnp.asarray(self.calibration_matrices) - mitigated_probs = rem_apply_to_probs( - mitigatee_probs, confusion_matrices, measured_qubits, n_qubits - ) - mitigated_result = mitigated_probs + measured_qubits = jnp.array(list(qnode_obj.device.wires)) + + handler = { + "sample": self._apply_sample, + "counts": self._apply_counts, + "probs": self._apply_probs, + }[mp_kind] + mitigated_result, confusion_matrices = handler( + rem_results, run_calibration, measured_qubits, n_qubits + ) if self.return_calibration_matrices: return mitigated_result, confusion_matrices return mitigated_result + def _apply_sample(self, rem_results, run_calibration, measured_qubits, n_qubits): + # quantum.sample returns f64 in catalyst; the REM helpers index + # confusion matrices by bit value and need an integer dtype. + mitigatee_samples = rem_results[0].astype(jnp.int32) + if run_calibration: + zeros_samples = rem_results[1].astype(jnp.int32) + ones_samples = rem_results[2].astype(jnp.int32) + confusion_matrices = rem_calibrate_samples(zeros_samples, ones_samples) + else: + confusion_matrices = jnp.asarray(self.calibration_matrices) + unique_bitstrings, mitigated_counts = rem_apply_to_samples( + mitigatee_samples, confusion_matrices, measured_qubits, n_qubits + ) + return (unique_bitstrings, mitigated_counts), confusion_matrices + + def _apply_counts(self, rem_results, run_calibration, measured_qubits, n_qubits): + mitigatee_eigvals = rem_results[0] + mitigatee_counts = rem_results[1] + if run_calibration: + zeros_counts = rem_results[2] + ones_counts = rem_results[3] + confusion_matrices = rem_calibrate_counts(zeros_counts, ones_counts) + else: + confusion_matrices = jnp.asarray(self.calibration_matrices) + mitigated_counts = rem_apply_to_counts( + mitigatee_counts, confusion_matrices, measured_qubits, n_qubits + ) + return (mitigatee_eigvals, mitigated_counts), confusion_matrices + + def _apply_probs(self, rem_results, run_calibration, measured_qubits, n_qubits): + mitigatee_probs = rem_results[0] + if run_calibration: + zeros_probs = rem_results[1] + ones_probs = rem_results[2] + confusion_matrices = rem_calibrate_probs(zeros_probs, ones_probs) + else: + confusion_matrices = jnp.asarray(self.calibration_matrices) + mitigated_probs = rem_apply_to_probs( + mitigatee_probs, confusion_matrices, measured_qubits, n_qubits + ) + return mitigated_probs, confusion_matrices + def polynomial_extrapolation(degree): """utility to generate polynomial fitting functions of arbitrary degree""" @@ -410,3 +406,19 @@ def _wrap_callable(fn): elif isinstance(fn, Callable): # Keep at the bottom return Function(fn) raise TypeError(f"Target must be callable, got: {type(fn)}") + + +def _detect_measurement_process(jaxpr): + """Walk the inner jaxpr for a probs/sample/counts primitive; return its name or None. + + Mirrors the detection logic in :func:`_rem_abstract_eval`. + """ + for eqn in jaxpr.eqns: + inner = eqn.params.get("call_jaxpr", None) + if inner is None: + continue + for op_eq in inner.eqns: + pname = str(op_eq.primitive) + if pname in ("probs", "sample", "counts"): + return pname + return None diff --git a/frontend/catalyst/api_extensions/rem_postprocessing.py b/frontend/catalyst/api_extensions/rem_postprocessing.py index 32f4c6824f..c0cdeb0382 100644 --- a/frontend/catalyst/api_extensions/rem_postprocessing.py +++ b/frontend/catalyst/api_extensions/rem_postprocessing.py @@ -28,8 +28,6 @@ import jax import jax.numpy as jnp import numpy as np -from jax import Array -from jax.typing import ArrayLike def _bitstrings_for_n_qubits(n_qubits: int) -> jax.Array: diff --git a/frontend/catalyst/jax_primitives.py b/frontend/catalyst/jax_primitives.py index aa64df0ac2..518330217a 100644 --- a/frontend/catalyst/jax_primitives.py +++ b/frontend/catalyst/jax_primitives.py @@ -1091,9 +1091,7 @@ def _rem_def_impl(ctx, *args, run_calibration, jaxpr, fn): # pragma: no cover @rem_p.def_abstract_eval -def _rem_abstract_eval( - *args, run_calibration, jaxpr, fn -): # pylint: disable=unused-argument +def _rem_abstract_eval(*args, run_calibration, jaxpr, fn): # pylint: disable=unused-argument """Abstract eval for the REM primitive. Returns a 3-tuple of avals matching the lowering's three results: diff --git a/frontend/test/pytest/test_mitigation_rem.py b/frontend/test/pytest/test_mitigation_rem.py index 9dcb5a41b0..c7d5123679 100644 --- a/frontend/test/pytest/test_mitigation_rem.py +++ b/frontend/test/pytest/test_mitigation_rem.py @@ -140,7 +140,7 @@ def mitigated(): def test_caching_probs_roundtrip(): - """Compute calibration once, reuse it on a second qjit call; result still matches the GHZ probs.""" + """Compute calibration once, reuse it on a second qjit call, get matching GHZ probs.""" n_qubits, shots = 2, 2000 circuit = _ghz_qnode(n_qubits, shots, "probs") @@ -314,7 +314,7 @@ def test_inverse_recovers_counts(): def test_calibrate_paths_agree(): - """rem_calibrate_{samples,counts,probs} on the same logical channel must yield matching matrices.""" + """rem_calibrate_{samples,counts,probs} on the same channel must produce matching matrices.""" n_qubits, n_shots = 2, 4000 p = 0.10 rng = np.random.default_rng(0) diff --git a/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp b/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp index d9ad48c58f..17abe544a8 100644 --- a/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp +++ b/mlir/lib/Mitigation/Transforms/MitigationMethods/Rem.cpp @@ -23,15 +23,126 @@ using namespace std::string_literals; namespace catalyst { namespace mitigation { +namespace { + +// Emit one calibration circuit next to the user callee. It reuses the callee's +// `quantum.device` kwargs verbatim (so it inherits the same noise channel) and, +// when `applyPauliX` is true, applies a PauliX gate on every qubit before +// compbasis to produce the all-ones basis state. Returns the calibration result +// values (drops the eigvals half of CountsOp). +ValueRange emitCalibrationCircuit(PatternRewriter &rewriter, Location loc, func::FuncOp calleeOp, + RankedTensorType tensorTyF64, RankedTensorType tensorTyI64, + bool applyPauliX) +{ + const char *tag = applyPauliX ? "(ones)" : "(zeros)"; + + const int64_t qubitCount = + (*calleeOp.getOps().begin()).getNqubitsAttr().value_or(0); + quantum::DeviceInitOp deviceInitOp = *(calleeOp.getOps().begin()); + Operation *shots = deviceInitOp.getShots().getDefiningOp(); + StringAttr lib = deviceInitOp.getLibAttr(); + StringAttr name = deviceInitOp.getDeviceNameAttr(); + StringAttr kwargs = deviceInitOp.getKwargsAttr(); + + Operation *measurementProcess = + (*calleeOp.getOps().begin()).getOperation(); + const std::string MPName = measurementProcess->getName().getIdentifier().str(); + + Operation *shotsLocal = shots->clone(); + rewriter.insert(shotsLocal); + auto devInit = + rewriter.create(loc, shotsLocal->getResult(0), lib, name, kwargs); + llvm::dbgs() << "[mitigation.rem] quantum.device initialization " << tag + << " OK, op=" << devInit << "\n"; + + IntegerAttr qubitCountAttr = rewriter.getI64IntegerAttr(qubitCount); + Value numberQubitsValue = rewriter.create(loc, qubitCountAttr); + Type qregType = quantum::QuregType::get(rewriter.getContext()); + auto qreg = rewriter.create(loc, qregType, numberQubitsValue, qubitCountAttr); + llvm::dbgs() << "[mitigation.rem] quantum.alloc " << tag << " addition OK, op=" << qreg << "\n"; + + Value currentQreg = qreg.getResult(); + if (applyPauliX) { + // For each qubit: extract -> PauliX -> insert. Each insert produces a + // fresh register SSA Value so the chain threads correctly. + for (int64_t i = 0; i < qubitCount; ++i) { + auto idxAttr = rewriter.getI64IntegerAttr(i); + auto extracted = rewriter.create( + loc, rewriter.getType(), currentQreg, nullptr, idxAttr); + auto xgate = rewriter.create( + loc, /*gate_name=*/"PauliX", mlir::ValueRange({extracted.getResult()})); + auto inserted = rewriter.create( + loc, qreg.getType(), currentQreg, nullptr, idxAttr, xgate.getResult(0)); + currentQreg = inserted.getResult(); + } + } + + Type obsType = ::catalyst::quantum::ObservableType::get(rewriter.getContext()); + auto compbasis = rewriter.create(loc, TypeRange{obsType}, + ValueRange{}, currentQreg); + llvm::dbgs() << "[mitigation.rem] quantum.compbasis " << tag << " addition OK, op=" << compbasis + << "\n"; + + ValueRange calibratedResult; + if (MPName == "quantum.probs"s) { + auto insertedMP = rewriter.create( + loc, TypeRange{tensorTyF64}, compbasis.getResult(), Value(), Value()); + llvm::dbgs() << "[mitigation.rem] quantum.probs " << tag + << " addition OK, op=" << insertedMP << "\n"; + calibratedResult = insertedMP->getResults(); + } + else if (MPName == "quantum.counts"s) { + auto insertedMP = + rewriter.create(loc, TypeRange{tensorTyF64, tensorTyI64}, + compbasis.getResult(), Value(), Value(), Value()); + llvm::dbgs() << "[mitigation.rem] quantum.counts " << tag + << " addition OK, op=" << insertedMP << "\n"; + // Drop the eigvals half; only the counts tensor feeds calibration. + calibratedResult = insertedMP->getResults().drop_front(); + } + else if (MPName == "quantum.sample"s) { + auto insertedMP = rewriter.create( + loc, TypeRange{tensorTyF64}, compbasis.getResult(), ValueRange{}, Value()); + llvm::dbgs() << "[mitigation.rem] quantum.sample " << tag + << " addition OK, op=" << insertedMP << "\n"; + IRMapping postprocMapping; + uint64_t resultIndex = 0; + for (auto res : measurementProcess->getResults()) + postprocMapping.map(res, insertedMP->getResult(resultIndex++)); + rewriter.setInsertionPointAfter(insertedMP.getOperation()); + calibratedResult = insertedMP->getResults(); + for (auto postprocOp = measurementProcess->getNextNode(); + postprocOp != nullptr && + postprocOp->getName().getIdentifier() != quantum::DeallocOp::getOperationName(); + postprocOp = postprocOp->getNextNode()) { + llvm::dbgs() << "[mitigation.rem] looking for next op: " + << postprocOp->getName().getIdentifier() << "\n"; + Operation *insertedOp = rewriter.clone(*postprocOp, postprocMapping); + if (!insertedOp->getResults().empty()) + calibratedResult = insertedOp->getResults(); + rewriter.setInsertionPointAfter(insertedOp); + llvm::dbgs() << "[mitigation.rem] insertedOp: " << insertedOp->getName().getIdentifier() + << "\n"; + } + } + + auto dealloc = rewriter.create(loc, qreg); + llvm::dbgs() << "[mitigation.rem] quantum.dealloc " << tag << " addition OK, op=" << dealloc + << "\n"; + auto deinit = rewriter.create(loc); + llvm::dbgs() << "[mitigation.rem] quantum.deinit " << tag << " addition OK, op=" << deinit + << "\n"; + return calibratedResult; +} + +} // namespace + LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter &rewriter) const { // The Rem lowering will produce three groups of results: // 1) the original callee results // 2) probabilities from an all-zeroes calibration circuit // 3) probabilities from an all-ones calibration circuit - // For the initial implementation we create simple placeholder calibration - // functions that return a 1-element tensor. These will be replaced - // with proper calibration circuits in a subsequent step. Location loc = op.getLoc(); auto runCalibration = op.getRunCalibrationAttr(); @@ -50,11 +161,6 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter SmallVector callArgs(op.getArgs().begin(), op.getArgs().end()); auto callOp = rewriter.create(loc, calleeOp, callArgs); - // Get SSA Values for wrapped callee function results. Note that these are not actual values of - // elements, but an abstraction, Value objects and ValueRange: non-owning view of multiple Value - // objects. These values are not known until runtime of the compiled program (not to be confused - // with from pass runtime / pass time). But their abstractions can be used to define the flow of - // the program auto results = callOp.getResults(); // Declare vector of SSA Values representing the total result (callee results + two calibration // circuit results, or zero-tensors if no calibration is present) @@ -63,9 +169,21 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter // times as many Value objects (because of zeros and ones calibration circuits) completeResults.reserve(results.size() * 3); + const bool doCalib = runCalibration.getValue(); + if (!doCalib) { + llvm::dbgs() << "[mitigation.rem] doCalibration == false, replacing RemOp with wrapped " + "callee circuit function call..." + << "\n"; + // The zeros/ones result groups are empty variadics when doCalib is false, + // so the op is replaced by only the callee's SSA values. + for (Value v : results) + completeResults.push_back(v); + rewriter.replaceOp(op, completeResults); + return success(); + } + // Determine quantum device information, specifically number of qubits and shots. - // Get the number of qubits from quantum.alloc operation inside callee function (FuncOp). Note - // the due to iterator (pointer) + // Get the number of qubits from quantum.alloc operation inside callee function (FuncOp). const int64_t qubitCount = (*calleeOp.getOps().begin()).getNqubitsAttr().value_or(0); // Get the device using quantom.device opeartion inside callee function (FuncOp) @@ -73,11 +191,10 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter // shots value could be a result of some operation, and not just a literal value, which is why // pointer to the operation is stored Operation *shots = deviceInitOp.getShots().getDefiningOp(); - StringAttr lib = deviceInitOp.getLibAttr(); - StringAttr name = deviceInitOp.getDeviceNameAttr(); - StringAttr kwargs = deviceInitOp.getKwargsAttr(); llvm::dbgs() << "[mitigation.rem] qubitCount: " << qubitCount << ", shots: " << *shots - << ", lib: " << lib << ", name: " << name << ", kwargs" << kwargs << "\n"; + << ", lib: " << deviceInitOp.getLibAttr() + << ", name: " << deviceInitOp.getDeviceNameAttr() << ", kwargs" + << deviceInitOp.getKwargsAttr() << "\n"; // check for MP type: Sample, Probs or Counts auto measurementProcesses = calleeOp.getOps(); @@ -85,16 +202,13 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter llvm::errs() << "[mitigation.rem] No valid MP found.\n"; return failure(); } - mlir::Operation *measurementProcess = (*measurementProcesses.begin()).getOperation(); - std::string MPName = measurementProcess->getName().getIdentifier().str(); + const std::string MPName = + (*measurementProcesses.begin()).getOperation()->getName().getIdentifier().str(); llvm::dbgs() << "[mitigation.rem] MP name: " << MPName << "\n"; // pre-compute MP and attribute-dependent values - bool doCalib = runCalibration.getValue(); // get boolean value of BoolAttr - mlir::RankedTensorType tensorTyI64; - mlir::RankedTensorType tensorTyF64; - int64_t calibrationTensorShape = 0; - int64_t shotCount = 0; + RankedTensorType tensorTyI64; + RankedTensorType tensorTyF64; if (MPName == "quantum.probs"s || MPName == "quantum.counts"s) { // avoid undefined behaviour due to integer overflow if (qubitCount >= 63) { @@ -103,33 +217,29 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter llvm::errs() << "[mitigation.rem] Qubit count" << qubitCount << "exceeds maximum simulated register size (62 qubits).\n"; } - calibrationTensorShape = static_cast(1ULL << qubitCount); + const int64_t calibrationTensorShape = static_cast(1ULL << qubitCount); tensorTyF64 = RankedTensorType::get({calibrationTensorShape}, rewriter.getF64Type()); tensorTyI64 = RankedTensorType::get({calibrationTensorShape}, rewriter.getI64Type()); } else if (MPName == "quantum.sample"s) { - // check if shots Op is actually a arith.constant. In that case, shots are known at compile - // time - if (shots->getName().getIdentifier() == arith::ConstantOp::getOperationName()) { - auto shotsAttr = shots->getAttrOfType("value"); - if (!shotsAttr || shotsAttr == 0) { - llvm::errs() - << "[mitigation.rem] Shots constant missing non-zero integer value. Sample MP " - "not supported for analytic simulation (shots must be > 0).\n"; - return failure(); - } - shotCount = shotsAttr.getInt(); - llvm::dbgs() - << "[mitigation.rem] Shots is a arth.constant op, shots known at compile time: " - << shotCount << "\n"; - calibrationTensorShape = shotCount * qubitCount; - tensorTyI64 = RankedTensorType::get({shotCount, qubitCount}, rewriter.getI64Type()); - tensorTyF64 = RankedTensorType::get({shotCount, qubitCount}, rewriter.getF64Type()); - } - else { + // shots must be a compile-time arith.constant with a non-zero value. + if (shots->getName().getIdentifier() != arith::ConstantOp::getOperationName()) { llvm::errs() << "[mitigation.rem] Dynamic number of shots not supported currently.\n"; return failure(); } + auto shotsAttr = shots->getAttrOfType("value"); + if (!shotsAttr || shotsAttr == 0) { + llvm::errs() + << "[mitigation.rem] Shots constant missing non-zero integer value. " + "Sample MP not supported for analytic simulation (shots must be > 0).\n"; + return failure(); + } + const int64_t shotCount = shotsAttr.getInt(); + llvm::dbgs() + << "[mitigation.rem] Shots is a arth.constant op, shots known at compile time: " + << shotCount << "\n"; + tensorTyI64 = RankedTensorType::get({shotCount, qubitCount}, rewriter.getI64Type()); + tensorTyF64 = RankedTensorType::get({shotCount, qubitCount}, rewriter.getF64Type()); } else { llvm::errs() << "[mitigation.rem] Supported measurement processes are quantum.counts, " @@ -137,212 +247,17 @@ LogicalResult RemLowering::matchAndRewrite(mitigation::RemOp op, PatternRewriter << MPName << " is not supported.\n"; return failure(); } - if (!doCalib) { // doCalib == false - llvm::dbgs() << "[mitigation.rem] doCalibration == false, replacing RemOp with wrapped " - "callee circuit function call..." - << "\n"; - // The zeros/ones result groups are empty variadics when doCalib is false, - // so the op is replaced by only the callee's SSA values. - for (Value v : results) - completeResults.push_back(v); - rewriter.replaceOp(op, completeResults); - return success(); - } - // else, doCalib == true + llvm::dbgs() << "[mitigation.rem] doCalibration == true, start adding all-zeroes and all-ones " "circuits..." << "\n"; - // - // DONE: implement adding two calibation circuits (all-zeroes and all-ones) and return - // appropriate results DONE: Then, add support for other MP, such as CountsMP and change the - // return type accordingly. - // Finally, add support for returning Observables, by somehow injecting the mitigation code - // (applying the confusion matrix) before observable calculation. (out of scope for now) 0. - // Check how insertion point is managed -- for now it seems to be just fine as it is, but this - // must be researched in the future - // ===BEGIN ZEROES QFUNC=== - // 1. Add device creation and allocation (quantum.device and quantum.alloc with previously - // determined attributes) - Operation *shotsLocal = shots->clone(); // not sure why this is needed - rewriter.insert(shotsLocal); // not sure why this is needed - auto devInit = - rewriter.create(loc, shotsLocal->getResult(0), lib, name, kwargs); - Type qregType = quantum::QuregType::get(rewriter.getContext()); - // not sure why empty IntegerAttr is needed, probably need to research more about how - // quantum.alloc works - IntegerAttr qubitCountAttr = rewriter.getI64IntegerAttr(qubitCount); - Value numberQubitsValue = rewriter.create(loc, qubitCountAttr); - llvm::dbgs() << "[mitigation.rem] quantum.device initialization OK, op=" << devInit << "\n"; - auto qreg = rewriter.create(loc, qregType, numberQubitsValue, qubitCountAttr); - llvm::dbgs() << "[mitigation.rem] quantum.alloc addition OK, op=" << qreg << "\n"; - // 2. Add quantum.compbasis op with the quantum observation result type and - // the allocated quantum register Value as its operand. Use the dialect's - // observation type and pass the AllocOp result Value (qreg.getResult()) - // as the operand; passing the wrong type or an incorrect operand can - // cause an `operandSegmentSizes` attribute to appear in the printed IR. - Type obsType = ::catalyst::quantum::ObservableType::get(rewriter.getContext()); - Value qregVal = qreg.getResult(); - // ComputationalBasisOp has the signature build(OpBuilder, OperationState, Type obs, - // ValueRange qubits, optional Value qreg). To create the "qreg" form (measure - // the entire register), pass an empty qubits ValueRange and the qreg Value as - // the optional argument. - // Use the TypeRange overload to ensure the builder picks the overload that - // treats the first argument as the result type (observable) rather than - // mistakenly interpreting operands. This avoids accidental operand ordering - // that can produce explicit operandSegmentSizes or wrong printed form. - auto compbasis = rewriter.create(loc, TypeRange{obsType}, - ValueRange{}, qregVal); - llvm::dbgs() << "[mitigation.rem] quantum.compbasis addition OK, op=" << compbasis << "\n"; - // 3. Add quantum.probs op from result of compbasis and get Value of tensor - // Probs returns a tensor of size 2^nqubits; build the appropriate RankedTensorType - // Be explicit about the builder overload to avoid overload-resolution - // ambiguity that can lead to incorrect operandSegmentSizes being set. - // Use the TypeRange overload so the result type is unambiguous and the - // `obs` operand is passed in the correct position. - // Pass explicit placeholders for the optional `dynamic_shape` and `state_in` - // operands (as null Values) so that the ODS builder records the correct - // operandSegmentSizes = {1, 0, 0} and the assembly printer can elide that - // property producing a clean printed form. - quantum::MeasurementProcess insertedMP; - ValueRange calibratedZeroResult; - if (MPName == "quantum.probs"s) { - insertedMP = rewriter.create(loc, TypeRange{tensorTyF64}, - compbasis.getResult(), Value(), Value()); - llvm::dbgs() << "[mitigation.rem] quantum.probs addition OK, op=" << insertedMP << "\n"; - calibratedZeroResult = insertedMP->getResults(); - } - else if (MPName == "quantum.counts"s) { - insertedMP = - rewriter.create(loc, TypeRange{tensorTyF64, tensorTyI64}, - compbasis.getResult(), Value(), Value(), Value()); - llvm::dbgs() << "[mitigation.rem] quantum.counts addition OK, op=" << insertedMP << "\n"; - calibratedZeroResult = - insertedMP->getResults().drop_front(); // drop the first element in the iterator - // (eigvals not used for calibration matrices) - } - else if (MPName == "quantum.sample"s) { - insertedMP = rewriter.create( - loc, TypeRange{tensorTyF64}, compbasis.getResult(), ValueRange{}, Value()); - llvm::dbgs() << "[mitigation.rem] quantum.sample addition OK, op=" << insertedMP << "\n"; - IRMapping postprocMapping; - uint64_t resultIndex = 0; - for (auto res : measurementProcess->getResults()) - postprocMapping.map(res, insertedMP->getResult(resultIndex++)); - rewriter.setInsertionPointAfter(insertedMP.getOperation()); - calibratedZeroResult = insertedMP->getResults(); - for (auto postprocOp = measurementProcess->getNextNode(); - postprocOp != nullptr && - postprocOp->getName().getIdentifier() != quantum::DeallocOp::getOperationName(); - postprocOp = postprocOp->getNextNode()) { - llvm::dbgs() << "[mitigation.rem] looking for next op: " - << postprocOp->getName().getIdentifier() << "\n"; - Operation *insertedOp = rewriter.clone(*postprocOp, postprocMapping); - if (!insertedOp->getResults().empty()) - calibratedZeroResult = insertedOp->getResults(); - rewriter.setInsertionPointAfter(insertedOp); - llvm::dbgs() << "[mitigation.rem] insertedOp: " << insertedOp->getName().getIdentifier() - << "\n"; - } - } - // MPName is already guaranteed to be a valid value, so no else block here. - - // 4. Deallocate everything - auto dealloc = rewriter.create(loc, qreg); - llvm::dbgs() << "[mitigation.rem] quantum.dealloc addition OK, op=" << dealloc << "\n"; - auto deinit = rewriter.create(loc); - llvm::dbgs() << "[mitigation.rem] quantum.deinit addition OK, op=" << deinit << "\n"; - // ===END ZEROES QFUNC=== - // 5. Add quantum.custom 'X' gates to implement ones circuit - // ===BEGIN ONES QFUNC=== - // Reuse numberQubitsValue and qubitCountAttr created above. Clone shots and create a new - // device/alloc sequence for the ones circuit. - Operation *shotsLocalOnes = shots->clone(); - rewriter.insert(shotsLocalOnes); - auto devInitOnes = rewriter.create(loc, shotsLocalOnes->getResult(0), - lib, name, kwargs); - llvm::dbgs() << "[mitigation.rem] quantum.device initialization (ones) OK, op=" << devInitOnes - << "\n"; - auto qregOnes = - rewriter.create(loc, qregType, numberQubitsValue, qubitCountAttr); - llvm::dbgs() << "[mitigation.rem] quantum.alloc (ones) addition OK, op=" << qregOnes << "\n"; - // For each qubit in the register: extract -> apply PauliX -> insert back - // Use a Value to track the current register SSA value so each insert - // produces a new SSA Value derived from the previous one. Reusing the - // previous InsertOp object (or its Op pointer) risks confusing the - // builder/operands and producing identical SSA uses. - Value currentQreg = qregOnes.getResult(); - for (int64_t i = 0; i < qubitCount; ++i) { - auto idxAttr = rewriter.getI64IntegerAttr(i); - // Extract qubit - auto extracted = rewriter.create( - loc, rewriter.getType(), currentQreg, nullptr, idxAttr); - // Apply PauliX (1-qubit gate) - auto xgate = rewriter.create(loc, /*gate_name=*/"PauliX", - mlir::ValueRange({extracted.getResult()})); - // Insert back into the register, producing a new register Value - auto inserted = rewriter.create(loc, qregOnes.getType(), currentQreg, - nullptr, idxAttr, xgate.getResult(0)); - currentQreg = inserted.getResult(); - } - - // Measure and get probs for ones circuit - Value qregValOnes = currentQreg; - auto compbasisOnes = rewriter.create(loc, TypeRange{obsType}, - ValueRange{}, qregValOnes); - llvm::dbgs() << "[mitigation.rem] quantum.compbasis (ones) addition OK, op=" << compbasisOnes - << "\n"; - // TODO: maybe refactor into a function like in ZNE and call for zeroes and ones - quantum::MeasurementProcess insertedMPOnes; - ValueRange calibratedOnesResult{}; - if (MPName == "quantum.probs"s) { - insertedMPOnes = rewriter.create( - loc, TypeRange{tensorTyF64}, compbasisOnes.getResult(), Value(), Value()); - llvm::dbgs() << "[mitigation.rem] quantum.probs (ones) addition OK, op=" << insertedMPOnes - << "\n"; - calibratedOnesResult = insertedMPOnes->getResults(); - } - if (MPName == "quantum.counts"s) { - insertedMPOnes = rewriter.create( - loc, TypeRange{tensorTyF64, tensorTyI64}, compbasisOnes.getResult(), Value(), Value(), - Value()); - llvm::dbgs() << "[mitigation.rem] quantum.counts (ones) addition OK, op=" << insertedMPOnes - << "\n"; - calibratedOnesResult = insertedMPOnes->getResults() - .drop_front(); // drop the first element in the iterator (eigvals - // not used for calibration matrices) - } - else if ((MPName == "quantum.sample"s)) { - insertedMPOnes = rewriter.create( - loc, TypeRange{tensorTyF64}, compbasisOnes.getResult(), ValueRange{}, Value()); - llvm::dbgs() << "[mitigation.rem] quantum.sample addition OK, op=" << insertedMPOnes - << "\n"; - IRMapping postprocMapping; - uint64_t resultIndex = 0; - for (auto res : measurementProcess->getResults()) - postprocMapping.map(res, insertedMPOnes->getResult(resultIndex++)); - rewriter.setInsertionPointAfter(insertedMPOnes.getOperation()); - calibratedOnesResult = insertedMPOnes->getResults(); - for (auto postprocOp = measurementProcess->getNextNode(); - postprocOp != nullptr && - postprocOp->getName().getIdentifier() != quantum::DeallocOp::getOperationName(); - postprocOp = postprocOp->getNextNode()) { - llvm::dbgs() << "[mitigation.rem] looking for next op: " - << postprocOp->getName().getIdentifier() << "\n"; - Operation *insertedOp = rewriter.clone(*postprocOp, postprocMapping); - if (!insertedOp->getResults().empty()) - calibratedOnesResult = insertedOp->getResults(); - rewriter.setInsertionPointAfter(insertedOp); - llvm::dbgs() << "[mitigation.rem] insertedOp: " << insertedOp->getName().getIdentifier() - << "\n"; - } - } - auto deallocOnes = rewriter.create(loc, qregOnes); - llvm::dbgs() << "[mitigation.rem] quantum.dealloc (ones) addition OK, op=" << deallocOnes - << "\n"; - auto deinitOnes = rewriter.create(loc); - llvm::dbgs() << "[mitigation.rem] quantum.deinit (ones) addition OK, op=" << deinitOnes << "\n"; - // ===END ONES QFUNC=== + ValueRange calibratedZeroResult = + emitCalibrationCircuit(rewriter, loc, calleeOp, tensorTyF64, tensorTyI64, + /*applyPauliX=*/false); + ValueRange calibratedOnesResult = + emitCalibrationCircuit(rewriter, loc, calleeOp, tensorTyF64, tensorTyI64, + /*applyPauliX=*/true); // 1) original callee results for (Value v : results) From 361185becfb4cab3369a2f4668af1c3e083d11ef Mon Sep 17 00:00:00 2001 From: volodymyr-hq Date: Thu, 2 Jul 2026 17:12:35 +0300 Subject: [PATCH 8/8] Update outdated documentation --- frontend/catalyst/api_extensions/rem_postprocessing.py | 2 +- mlir/include/Mitigation/IR/MitigationOps.td | 7 +++---- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/frontend/catalyst/api_extensions/rem_postprocessing.py b/frontend/catalyst/api_extensions/rem_postprocessing.py index c0cdeb0382..a77d0f864d 100644 --- a/frontend/catalyst/api_extensions/rem_postprocessing.py +++ b/frontend/catalyst/api_extensions/rem_postprocessing.py @@ -16,7 +16,7 @@ JAX-compilable post-processing helpers for Readout Error Mitigation (REM). These helpers integrate with the catalyst frontend so ``mitigate_with_rem`` -can apply classical SampleMP post-processing - per-qubit confusion-matrix +can apply classical post-processing - per-qubit confusion-matrix calibration, reduced-confusion-matrix construction, and a linear solve - to the three sample arrays emitted by the ``mitigation.rem`` op after lowering. diff --git a/mlir/include/Mitigation/IR/MitigationOps.td b/mlir/include/Mitigation/IR/MitigationOps.td index 1e22b3fbba..b496fdd7d6 100644 --- a/mlir/include/Mitigation/IR/MitigationOps.td +++ b/mlir/include/Mitigation/IR/MitigationOps.td @@ -70,10 +70,9 @@ def RemOp : Mitigation_Op<"rem", [DeclareOpInterfaceMethods, all-ones calibration circuit alongside the original callee; their sample/probs results are returned as the `zeros` and `ones` SSA values and consumed by the classical REM post-processing on the JAX side. - When false the calibration outputs are zero-filled placeholder tensors - of the same type, which lets the caller supply cached confusion - matrices and skip the calibration circuits without changing the op's - result signature. + When false the calibration is ommited, which lets the caller supply + cached confusion matrices and skip the calibration circuits without + changing the op's result signature. }]; let arguments = (ins