[FTQC] Add support for parametric mid-circuit measurements

.dep-versions

@@ -8,7 +8,7 @@ enzyme=v0.0.149
 
 # For a custom PL version, update the package version here and at
 # 'doc/requirements.txt
-pennylane=0.42.0-dev15
+pennylane=0.42.0-dev19
 
 # For a custom LQ/LK version, update the package version here and at
 # 'doc/requirements.txt'

doc/conf.py

@@ -96,6 +96,7 @@ def __getattr__(cls, name):
     "mlir_quantum.dialects.quantum",
     "mlir_quantum.dialects.gradient",
     "mlir_quantum.dialects.catalyst",
+    "mlir_quantum.dialects.mbqc",
     "mlir_quantum.dialects.mitigation",
     "mlir_quantum.dialects._transform_ops_gen",
     "pybind11",

doc/releases/changelog-dev.md

@@ -100,6 +100,11 @@
   This runtime stub is currently for mock execution only and should be treated as a placeholder
   operation. Internally, it functions just as a computational-basis measurement instruction.
 
+* PennyLane's arbitrary-basis measurement operations, such as [`qml.ftqc.measure_arbitrary_basis()`
+  ](https://docs.pennylane.ai/en/stable/code/api/pennylane.ftqc.measure_arbitrary_basis.html), are
+  now QJIT-compatible with program capture enabled.
+  [(#1645)](https://github.com/PennyLaneAI/catalyst/pull/1645)
+
 * The utility function `EnsureFunctionDeclaration` is refactored into the `Utils` of the `Catalyst`
   dialect, instead of being duplicated in each individual dialect.
   [(#1683)](https://github.com/PennyLaneAI/catalyst/pull/1683)

doc/requirements.txt

@@ -32,4 +32,4 @@ lxml_html_clean
 --extra-index-url https://test.pypi.org/simple/
 pennylane-lightning-kokkos==0.42.0-dev11
 pennylane-lightning==0.42.0-dev11
-pennylane==0.42.0-dev15
+pennylane==0.42.0-dev19

frontend/catalyst/from_plxpr.py

@@ -29,6 +29,7 @@
 from pennylane.capture.primitives import cond_prim as plxpr_cond_prim
 from pennylane.capture.primitives import for_loop_prim as plxpr_for_loop_prim
 from pennylane.capture.primitives import while_loop_prim as plxpr_while_loop_prim
+from pennylane.ftqc.primitives import measure_in_basis_prim as plxpr_measure_in_basis_prim
 from pennylane.ops.functions.map_wires import _map_wires_transform as pl_map_wires
 from pennylane.transforms import cancel_inverses as pl_cancel_inverses
 from pennylane.transforms import commute_controlled as pl_commute_controlled
@@ -47,12 +48,14 @@
 from catalyst.jax_primitives import (
     AbstractQbit,
     AbstractQreg,
+    MeasurementPlane,
     compbasis_p,
     cond_p,
     counts_p,
     expval_p,
     for_p,
     gphase_p,
+    measure_in_basis_p,
     namedobs_p,
     probs_p,
     qalloc_p,
@@ -640,6 +643,27 @@ def handle_while_loop(
     return outvals
 
 
+# pylint: disable=unused-argument, too-many-positional-arguments
+@QFuncPlxprInterpreter.register_primitive(plxpr_measure_in_basis_prim)
+def handle_measure_in_basis(self, angle, wire, plane, reset, postselect):
+    """Handle the conversion from plxpr to Catalyst jaxpr for the measure_in_basis primitive"""
+    _angle = jax.lax.convert_element_type(angle, jnp.dtype(jnp.float64))
+
+    try:
+        _plane = MeasurementPlane(plane)
+    except ValueError as e:
+        raise ValueError(
+            f"Measurement plane must be one of {[plane.value for plane in MeasurementPlane]}"
+        ) from e
+
+    in_wire = self.get_wire(wire)
+    result, out_wire = measure_in_basis_p.bind(_angle, in_wire, plane=_plane, postselect=postselect)
+
+    self.wire_map[wire] = out_wire
+
+    return result
+
+
 # Derived interpreters must be declared after the primitive registrations of their
 # parents or be placed in a separate file, in order to access those registrations.
 # This is due to the registrations being done outside the parent class definition.

frontend/catalyst/jax_primitives.py

@@ -61,6 +61,7 @@
     ValueAndGradOp,
     VJPOp,
 )
+from mlir_quantum.dialects.mbqc import MeasureInBasisOp
 from mlir_quantum.dialects.mitigation import ZneOp
 from mlir_quantum.dialects.quantum import (
     AdjointOp,
@@ -230,6 +231,16 @@ class Folding(Enum):
     ALL = "local-all"
 
 
+class MeasurementPlane(Enum):
+    """
+    Measurement planes for arbitrary-basis measurements in MBQC
+    """
+
+    XY = "XY"
+    YZ = "YZ"
+    ZX = "ZX"
+
+
 ##############
 # Primitives #
 ##############
@@ -292,6 +303,8 @@ class Folding(Enum):
 set_basis_state_p.multiple_results = True
 quantum_kernel_p = core.CallPrimitive("quantum_kernel")
 quantum_kernel_p.multiple_results = True
+measure_in_basis_p = core.Primitive("measure_in_basis")
+measure_in_basis_p.multiple_results = True
 
 
 def _assert_jaxpr_without_constants(jaxpr: ClosedJaxpr):
@@ -1214,6 +1227,79 @@ def _qmeasure_lowering(jax_ctx: mlir.LoweringRuleContext, qubit: ir.Value, posts
     )
 
 
+#
+# arbitrary-basis measurements
+#
+@measure_in_basis_p.def_abstract_eval
+def _measure_in_basis_abstract_eval(
+    angle: float, qubit: AbstractQbit, plane: MeasurementPlane, postselect: int = None
+):
+    assert isinstance(qubit, AbstractQbit)
+    return core.ShapedArray((), bool), qubit
+
+
+@measure_in_basis_p.def_impl
+def _measure_in_basis_def_impl(
+    ctx, angle: float, qubit: AbstractQbit, plane: MeasurementPlane, postselect: int = None
+):  # pragma: no cover
+    raise NotImplementedError()
+
+
+def _measurement_plane_attribute(ctx, plane: MeasurementPlane):
+    return ir.OpaqueAttr.get(
+        "mbqc",
+        ("measurement_plane " + MeasurementPlane(plane).name).encode("utf-8"),
+        ir.NoneType.get(ctx),
+        ctx,
+    )
+
+
+def _measure_in_basis_lowering(
+    jax_ctx: mlir.LoweringRuleContext,
+    angle: float,
+    qubit: ir.Value,
+    plane: MeasurementPlane,
+    postselect: int = None,
+):
+    ctx = jax_ctx.module_context.context
+    ctx.allow_unregistered_dialects = True
+
+    assert ir.OpaqueType.isinstance(qubit.type)
+    assert ir.OpaqueType(qubit.type).dialect_namespace == "quantum"
+    assert ir.OpaqueType(qubit.type).data == "bit"
+
+    angle = safe_cast_to_f64(angle, "angle")
+    angle = extract_scalar(angle, "angle")
+
+    assert ir.F64Type.isinstance(
+        angle.type
+    ), "Only scalar double parameters are allowed for quantum gates!"
+
+    # Prepare postselect attribute
+    if postselect is not None:
+        i32_type = ir.IntegerType.get_signless(32, ctx)
+        postselect = ir.IntegerAttr.get(i32_type, postselect)
+
+    result_type = ir.IntegerType.get_signless(1)
+
+    result, new_qubit = MeasureInBasisOp(
+        result_type,
+        qubit.type,
+        qubit,
+        plane=_measurement_plane_attribute(ctx, plane),
+        angle=angle,
+        postselect=postselect,
+    ).results
+
+    result_from_elements_op = ir.RankedTensorType.get((), result.type)
+    from_elements_op = FromElementsOp(result_from_elements_op, result)
+
+    return (
+        from_elements_op.results[0],
+        new_qubit,
+    )
+
+
 #
 # compbasis observable
 #
@@ -2250,6 +2336,7 @@ def _cos_lowering2(ctx, x):
     (sin_p, _sin_lowering2),
     (cos_p, _cos_lowering2),
     (quantum_kernel_p, _quantum_kernel_lowering),
+    (measure_in_basis_p, _measure_in_basis_lowering),
 )