Add support for lowering adjoint and control with Operator2

doc/releases/changelog-dev.md

@@ -9,7 +9,8 @@
   without non-lowerable arguments. `Operator2` classes are now lowered to specialized operations
   where applicable, unlocking compilation and execution for these cases.
   [(#2979)](https://github.com/PennyLaneAI/catalyst/pull/2979)
-  [(#2969)](https://github.com/PennyLaneAI/catalyst/pull/2969/)
+  [(#2969)](https://github.com/PennyLaneAI/catalyst/pull/2969)
+  [(#2980)](https://github.com/PennyLaneAI/catalyst/pull/2980)
   [(#2990)](https://github.com/PennyLaneAI/catalyst/pull/2990)
 
 * The `ResourceAnalysis` pass now reports each loop body and each subroutine as its own entry

frontend/catalyst/from_plxpr/qfunc_interpreter.py

@@ -28,9 +28,9 @@
 from pennylane.capture.primitives import cond_prim as pl_cond_prim
 from pennylane.capture.primitives import ctrl_transform_prim as plxpr_ctrl_transform_prim
 from pennylane.capture.primitives import measure_prim as plxpr_measure_prim
+from pennylane.capture.primitives import operator_p
 from pennylane.capture.primitives import pauli_measure_prim as plxpr_pauli_measure_prim
 from pennylane.capture.primitives import quantum_subroutine_prim, transform_prim
-from pennylane.core.operator.operator2 import operator_p
 from pennylane.ftqc.primitives import measure_in_basis_prim as plxpr_measure_in_basis_prim
 from pennylane.measurements import CountsMP
 from pennylane.wires import AbstractQubit, is_abstract_qubit
@@ -300,14 +300,27 @@ def _handle_operator(self, *args, op_cls, hybrid_lens, hybrid_trees, adjoint, n_
         # only support compilable_argnames for the moment
         raise NotImplementedError
 
-    wire_inputs = args[len(op_cls.dynamic_argnames) :]
+    if n_ctrls:
+        wire_inputs = args[len(op_cls.dynamic_argnames) : -2 * n_ctrls]
+        control_wire_inputs = args[-2 * n_ctrls : -n_ctrls]
+        control_values = args[-n_ctrls:]
+    else:
+        wire_inputs = args[len(op_cls.dynamic_argnames) :]
+        control_wire_inputs = control_values = ()
+
     new_wires = [
         w if is_abstract_qubit(w) else qref_get_p.bind(self.init_qreg, w) for w in wire_inputs
     ]
+    new_control_wires = [
+        w if is_abstract_qubit(w) else qref_get_p.bind(self.init_qreg, w)
+        for w in control_wire_inputs
+    ]
 
     qref_operator_p.bind(
         *args[: len(op_cls.dynamic_argnames)],
         *new_wires,
+        *new_control_wires,
+        *control_values,
         op_cls=op_cls,
         hybrid_lens=hybrid_lens,
         hybrid_trees=hybrid_trees,

frontend/catalyst/from_plxpr/qref_operator2_primitives.py

@@ -104,19 +104,40 @@ def _qref_operator_p_lowering(
     op_cls,
     **kwargs,
 ):
-    _general_validation(*args, op_cls=op_cls, **kwargs)
     ctx = jax_ctx.module_context.context
     ctx.allow_unregistered_dialects = True
-    if op_cls.__name__ in _SPECIAL_LOWERINGS:
-        return _SPECIAL_LOWERINGS[op_cls.__name__](jax_ctx, *args, op_cls=op_cls, **kwargs)
-    # will be used in future improvements
+    _general_validation(*args, op_cls=op_cls, **kwargs)
+
     hybrid_lens = kwargs.pop("hybrid_lens")  # pylint: disable=unused-variable
     hybrid_trees = kwargs.pop("hybrid_trees")  # pylint: disable=unused-variable
-    adjoint = kwargs.pop("adjoint")  # pylint: disable=unused-variable
-    n_ctrls = kwargs.pop("n_ctrls")  # pylint: disable=unused-variable
+    adjoint = kwargs.pop("adjoint")
+    n_ctrls = kwargs.pop("n_ctrls")
     wire_lens = kwargs.pop("wire_lens")
+
+    if n_ctrls:
+        ctrl_qubits = args[-2 * n_ctrls : -n_ctrls]
+        ctrl_values = [
+            TensorExtractOp(ir.IntegerType.get_signless(1), val, []).result
+            for val in args[-n_ctrls:]
+        ]
+        qubits_slice = slice(len(op_cls.dynamic_argnames), -2 * n_ctrls)
+    else:
+        ctrl_qubits = ctrl_values = ()
+        qubits_slice = slice(len(op_cls.dynamic_argnames), None)
+
     params = args[: len(op_cls.dynamic_argnames)]
-    qubits = args[len(op_cls.dynamic_argnames) :]
+    qubits = args[qubits_slice]
+
+    if op_cls.__name__ in _SPECIAL_LOWERINGS:
+        return _SPECIAL_LOWERINGS[op_cls.__name__](
+            *params,
+            *qubits,
+            ctrl_qubits=ctrl_qubits,
+            ctrl_values=ctrl_values,
+            adjoint=adjoint,
+            **kwargs,
+        )
+    # will be used in future improvements
 
     name_attr = get_mlir_attribute_from_pyval(op_cls.__name__)
 
@@ -136,10 +157,6 @@ def _qref_operator_p_lowering(
 
     processed_qubit_map = get_mlir_attribute_from_pyval(qubit_map)
 
-    ctrl_qubits = []
-    ctrl_values = []
-    adjoint = False
-
     if _is_custom_op(op_cls, jax_ctx.avals_in[: len(op_cls.dynamic_argnames)]):
         params = [extract_scalar(safe_cast_to_f64(p, op_cls), op_cls) for p in params]
         CustomOp(
@@ -157,9 +174,9 @@ def _qref_operator_p_lowering(
             qubits=qubits,
             qreg=None,
             forward_args=[],
-            ctrl_qubits=[],
-            ctrl_values=[],
-            adjoint=False,
+            ctrl_qubits=ctrl_qubits,
+            ctrl_values=ctrl_values,
+            adjoint=adjoint,
             UID=None,
             arr_qubit_indices=[],
             param_map=processed_param_map,
@@ -170,58 +187,43 @@ def _qref_operator_p_lowering(
 
 
 @_register_special_lowering("MultiRZ")
-def _multirz_lowering(jax_ctx: mlir.LoweringRuleContext, *args, **_):
-    theta = extract_scalar(safe_cast_to_f64(args[0], "MultiRZ"), "MultiRZ")
-    qubits = args[1:]
+def _multirz_lowering(theta, *qubits, ctrl_qubits, ctrl_values, adjoint):
     MultiRZOp(
-        theta=theta,
+        theta=extract_scalar(safe_cast_to_f64(theta, "MultiRZ"), "MultiRZ"),
         qubits=qubits,
-        ctrl_qubits=[],
-        ctrl_values=[],
-        adjoint=False,
+        ctrl_qubits=ctrl_qubits,
+        ctrl_values=ctrl_values,
+        adjoint=adjoint,
     )
     return []
 
 
 @_register_special_lowering("PCPhase")
-def _pcphase_lowering(
-    jax_ctx: mlir.LoweringRuleContext,
-    *args,
-    **_,
-):
-    qubits = args[2:]
+def _pcphase_lowering(theta, dim, *qubits, ctrl_qubits, ctrl_values, adjoint):
     PCPhaseOp(
-        theta=extract_scalar(safe_cast_to_f64(args[0], "PCPhase"), "PCPhase"),
-        dim=extract_scalar(safe_cast_to_f64(args[1], "PCPhase"), "PCPhase"),
+        theta=extract_scalar(safe_cast_to_f64(theta, "PCPhase"), "PCPhase"),
+        dim=extract_scalar(safe_cast_to_f64(dim, "PCPhase"), "PCPhase"),
         qubits=qubits,
-        ctrl_qubits=[],
-        ctrl_values=[],
-        adjoint=False,
+        ctrl_qubits=ctrl_qubits,
+        ctrl_values=ctrl_values,
+        adjoint=adjoint,
     )
     return ()
 
 
 @_register_special_lowering("GlobalPhase")
-def _special_gphase_lowering(
-    jax_ctx: mlir.LoweringRuleContext,
-    *args,
-    op_cls,
-    **_,
-):
+def _special_gphase_lowering(angle, *_, ctrl_qubits, ctrl_values, adjoint):
     GlobalPhaseOp(
-        angle=extract_scalar(safe_cast_to_f64(args[0], "GlobalPhase"), "GlobalPhase"),
-        ctrl_qubits=[],
-        ctrl_values=[],
-        adjoint=False,
+        angle=extract_scalar(safe_cast_to_f64(angle, "GlobalPhase"), "GlobalPhase"),
+        ctrl_qubits=ctrl_qubits,
+        ctrl_values=ctrl_values,
+        adjoint=adjoint,
     )
     return ()
 
 
 @_register_special_lowering("QubitUnitary")
-def _special_unitary_lowering(jax_ctx: mlir.LoweringRuleContext, matrix, *qubits, **_):
-    ctrl_qubits = []
-    ctrl_values = []
-
+def _special_unitary_lowering(matrix, *qubits, ctrl_qubits, ctrl_values, adjoint):
     matrix_type = matrix.type
     is_tensor = ir.RankedTensorType.isinstance(matrix_type)
     shape = ir.RankedTensorType(matrix_type).shape if is_tensor else None
@@ -245,50 +247,29 @@ def _special_unitary_lowering(jax_ctx: mlir.LoweringRuleContext, matrix, *qubits
         tensor_complex_f64_type = ir.RankedTensorType.get(shape, complex_f64_type)
         matrix = StableHLOConvertOp(tensor_complex_f64_type, matrix).result
 
-    ctrl_values_i1 = [
-        TensorExtractOp(ir.IntegerType.get_signless(1), v, []).result for v in ctrl_values
-    ]
-
     QubitUnitaryOp(
         matrix=matrix,
         qubits=qubits,
         ctrl_qubits=ctrl_qubits,
-        ctrl_values=ctrl_values_i1,
-        adjoint=False,
+        ctrl_values=ctrl_values,
+        adjoint=adjoint,
     )
 
     return ()
 
 
 @_register_special_lowering("PauliRot")
-def _special_paulirot_lowering(
-    jax_ctx: mlir.LoweringRuleContext,
-    angle,
-    *qubits,
-    pauli_word,
-    **_,
-):
+def _special_paulirot_lowering(angle, *qubits, ctrl_qubits, ctrl_values, adjoint, pauli_word):
     pauli_word = unflatten(*pauli_word)
-    ctrl_qubits = []
-    ctrl_values = []
-
-    angle = safe_cast_to_f64(angle, "PauliRot")
-    angle = extract_scalar(angle, "PauliRot")
-    assert ir.F64Type.isinstance(angle.type)
-
     pauli_word = ir.ArrayAttr.get([ir.StringAttr.get(p) for p in pauli_word])
 
-    ctrl_values_i1 = [
-        TensorExtractOp(ir.IntegerType.get_signless(1), v, []).result for v in ctrl_values
-    ]
-
     PauliRotOp(
-        angle=angle,
+        angle=extract_scalar(safe_cast_to_f64(angle, "PauliRot"), "PauliRot"),
         pauli_product=pauli_word,
         qubits=qubits,
         ctrl_qubits=ctrl_qubits,
-        ctrl_values=ctrl_values_i1,
-        adjoint=False,
+        ctrl_values=ctrl_values,
+        adjoint=adjoint,
     )
 
     return ()