Use native implementation for adjoints in (control) operations

.github/CHANGELOG.md

@@ -6,6 +6,9 @@
 
 ### Improvements
 
+* Use native C++ kernels for controlled-adjoint and adjoint-controlled of supported operations.
+  [(#1063)](https://github.com/PennyLaneAI/pennylane-lightning/pull/1063)
+
 * In Lightning-Tensor, allow `qml.MPSPrep` to accept an MPS with `len(MPS) = n_wires-1`.
   [(#1064)](https://github.com/PennyLaneAI/pennylane-lightning/pull/1064)
 

pennylane_lightning/core/_serialize.py

@@ -443,8 +443,18 @@ def serialize_ops(self, tape: QuantumTape, wires_map: dict = None) -> Tuple[
         uses_stateprep = False
 
         def get_wires(operation, single_op):
-            if isinstance(operation, qml.ops.op_math.Controlled) and not isinstance(
-                operation,
+            # Serialize adjoint(op) and adjoint(ctrl(op))
+            if isinstance(operation, qml.ops.op_math.Adjoint):
+                inverse = True
+                op_base = operation.base
+                single_op_base = single_op.base
+            else:
+                inverse = False
+                op_base = operation
+                single_op_base = single_op
+
+            if isinstance(op_base, qml.ops.op_math.Controlled) and not isinstance(
+                op_base,
                 (
                     qml.CNOT,
                     qml.CY,
@@ -457,19 +467,41 @@ def get_wires(operation, single_op):
                     qml.CSWAP,
                 ),
             ):
-                name = operation.base.name
-                wires_list = list(operation.target_wires)
-                controlled_wires_list = list(operation.control_wires)
-                control_values_list = operation.control_values
+                wires_list = list(op_base.target_wires)
+                controlled_wires_list = list(op_base.control_wires)
+                control_values_list = op_base.control_values
+                # Serialize ctrl(adjoint(op))
+                if isinstance(op_base.base, qml.ops.op_math.Adjoint):
+                    ctrl_adjoint = True
+                    name = op_base.base.base.name
+                else:
+                    ctrl_adjoint = False
+                    name = op_base.base.name
+
+                # Inside the controlled operation, if the base operation (of the adjoint)
+                # is supported natively, we apply the the base operation and invert the
+                # inverse flag; otherwise we apply the QubitUnitary of a matrix which
+                # contains the inverse and leave the inverse flag as is.
                 if not hasattr(self.sv_type, name):
-                    single_op = QubitUnitary(matrix(single_op.base), single_op.base.wires)
-                    name = single_op.name
+                    single_op_base = QubitUnitary(
+                        matrix(single_op_base.base), single_op_base.base.wires
+                    )
+                    name = single_op_base.name
+                else:
+                    inverse ^= ctrl_adjoint
             else:
-                name = single_op.name
-                wires_list = single_op.wires.tolist()
+                name = single_op_base.name
+                wires_list = single_op_base.wires.tolist()
                 controlled_wires_list = []
                 control_values_list = []
-            return single_op, name, list(wires_list), controlled_wires_list, control_values_list
+            return (
+                single_op_base,
+                name,
+                inverse,
+                list(wires_list),
+                controlled_wires_list,
+                control_values_list,
+            )
 
         for operation in tape.operations:
             if isinstance(operation, (BasisState, StatePrep)):
@@ -480,30 +512,29 @@ def get_wires(operation, single_op):
             else:
                 op_list = [operation]
 
-            inverse = isinstance(operation, qml.ops.op_math.Adjoint)
-
             for single_op in op_list:
                 (
-                    single_op,
+                    single_op_base,
                     name,
+                    inverse,
                     wires_list,
                     controlled_wires_list,
                     controlled_values_list,
                 ) = get_wires(operation, single_op)
                 inverses.append(inverse)
-                names.append(single_op.base.name if inverse else name)
+                names.append(name)
                 # QubitUnitary is a special case, it has a parameter which is not differentiable.
                 # We thus pass a dummy 0.0 parameter which will not be referenced
-                if isinstance(single_op, qml.QubitUnitary):
+                if isinstance(single_op_base, qml.QubitUnitary):
                     params.append([0.0])
-                    mats.append(matrix(single_op))
+                    mats.append(matrix(single_op_base))
                 else:
-                    if hasattr(self.sv_type, single_op.base.name if inverse else name):
-                        params.append(single_op.parameters)
+                    if hasattr(self.sv_type, name):
+                        params.append(single_op_base.parameters)
                         mats.append([])
                     else:
                         params.append([])
-                        mats.append(matrix(single_op))
+                        mats.append(matrix(single_op_base))
 
                 controlled_values.append(controlled_values_list)
                 controlled_wires.append(

pennylane_lightning/core/_state_vector_base.py

@@ -21,6 +21,7 @@
 import numpy as np
 from pennylane import BasisState, StatePrep
 from pennylane.measurements import MidMeasureMP
+from pennylane.ops import Controlled
 from pennylane.tape import QuantumScript
 from pennylane.wires import Wires
 
@@ -131,11 +132,12 @@ def _apply_basis_state(self, state, wires):
         self._qubit_state.setBasisState(list(state), list(wires))
 
     @abstractmethod
-    def _apply_lightning_controlled(self, operation):
+    def _apply_lightning_controlled(self, operation: Controlled, adjoint: bool):
         """Apply an arbitrary controlled operation to the state tensor.
 
         Args:
             operation (~pennylane.operation.Operation): controlled operation to apply
+            adjoint (bool): Apply the adjoint of the operation if True
 
         Returns:
             None

pennylane_lightning/core/_version.py

@@ -16,4 +16,4 @@
 Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.41.0-dev24"
+__version__ = "0.41.0-dev25"

pennylane_lightning/lightning_gpu/_state_vector.py

@@ -229,34 +229,39 @@
         # set the state vector on GPU with provided state and their corresponding wires
         self._qubit_state.setStateVector(state, list(device_wires), use_async)
 
-    def _apply_lightning_controlled(self, operation):
+    def _apply_lightning_controlled(self, operation, adjoint):
         """Apply an arbitrary controlled operation to the state tensor.
 
         Args:
             operation (~pennylane.operation.Operation): controlled operation to apply
+            adjoint (bool): Apply the adjoint of the operation if True
 
         Returns:
             None
         """
         state = self.state_vector
 
-        basename = operation.base.name
-        method = getattr(state, f"{basename}", None)
+        if isinstance(operation.base, Adjoint):
+            base_operation = operation.base.base
+            adjoint = not adjoint
+        else:
+            base_operation = operation.base
+
+        method = getattr(state, f"{base_operation.name}", None)
         control_wires = list(operation.control_wires)
         control_values = operation.control_values
         target_wires = list(operation.target_wires)
         if method:  # apply n-controlled specialized gate
-            inv = False
             param = operation.parameters
-            method(control_wires, control_values, target_wires, inv, param)
+            method(control_wires, control_values, target_wires, adjoint, param)
         else:  # apply gate as an n-controlled matrix
             method = getattr(state, "applyControlledMatrix")
             method(
-                qml.matrix(operation.base),
+                qml.matrix(base_operation),
                 control_wires,
                 control_values,
                 target_wires,
-                False,
+                adjoint,
             )
 
     def _apply_lightning_midmeasure(
@@ -300,6 +305,7 @@
                 postselection. Use ``"hw-like"`` to discard invalid shots and ``"fill-shots"`` to
                 keep the same number of shots. Default is ``None``.
 
+
         Returns:
             None
         """
@@ -311,11 +317,12 @@
             if isinstance(operation, qml.Identity):
                 continue
             if isinstance(operation, Adjoint):
-                name = operation.base.name
+                op_adjoint_base = operation.base
                 invert_param = True
             else:
-                name = operation.name
+                op_adjoint_base = operation
                 invert_param = False
+            name = op_adjoint_base.name
             method = getattr(state, name, None)
             wires = list(operation.wires)
 
@@ -330,13 +337,13 @@
                 param = operation.parameters
                 method(wires, invert_param, param)
             elif (
-                isinstance(operation, qml.ops.Controlled) and not self._mpi_handler.use_mpi
+                isinstance(op_adjoint_base, qml.ops.Controlled) and not self._mpi_handler.use_mpi
             ):  # MPI backend does not have native controlled gates support
-                self._apply_lightning_controlled(operation)
+                self._apply_lightning_controlled(op_adjoint_base, invert_param)
             elif (
                 self._mpi_handler.use_mpi
-                and isinstance(operation, qml.ops.Controlled)
-                and isinstance(operation.base, qml.GlobalPhase)
+                and isinstance(op_adjoint_base, qml.ops.Controlled)
+                and isinstance(op_adjoint_base.base, qml.GlobalPhase)
             ):
                 # TODO: To move this line to the _apply_lightning_controlled method once the MPI backend supports controlled gates natively
                 raise DeviceError(
@@ -348,7 +355,6 @@
                 except AttributeError:  # pragma: no cover
                     # To support older versions of PL
                     mat = operation.matrix
-
                 r_dtype = np.float32 if self.dtype == np.complex64 else np.float64
                 param = (
                     [[r_dtype(operation.hash)]]

pennylane_lightning/lightning_kokkos/_state_vector.py

@@ -181,34 +181,39 @@ def _apply_state_vector(self, state, device_wires: Wires):
         # This operate on device
         self._qubit_state.setStateVector(state, list(device_wires))
 
-    def _apply_lightning_controlled(self, operation):
+    def _apply_lightning_controlled(self, operation, adjoint):
         """Apply an arbitrary controlled operation to the state tensor.
 
         Args:
             operation (~pennylane.operation.Operation): controlled operation to apply
+            adjoint (bool): Apply the adjoint of the operation if True
 
         Returns:
             None
         """
         state = self.state_vector
 
-        basename = operation.base.name
-        method = getattr(state, f"{basename}", None)
+        if isinstance(operation.base, Adjoint):
+            base_operation = operation.base.base
+            adjoint = not adjoint
+        else:
+            base_operation = operation.base
+
+        method = getattr(state, f"{base_operation.name}", None)
         control_wires = list(operation.control_wires)
         control_values = operation.control_values
         target_wires = list(operation.target_wires)
-        inv = False  # TODO: update to use recursive _apply_lightning to handle nested adjoint/ctrl
         if method is not None:  # apply n-controlled specialized gate
             param = operation.parameters
-            method(control_wires, control_values, target_wires, inv, param)
+            method(control_wires, control_values, target_wires, adjoint, param)
         else:  # apply gate as an n-controlled matrix
             method = getattr(state, "applyControlledMatrix")
             method(
-                qml.matrix(operation.base),
+                qml.matrix(base_operation),
                 control_wires,
                 control_values,
                 target_wires,
-                inv,
+                adjoint,
             )
 
     def _apply_lightning_midmeasure(
@@ -262,11 +267,12 @@ def _apply_lightning(
             if isinstance(operation, qml.Identity):
                 continue
             if isinstance(operation, Adjoint):
-                name = operation.base.name
+                op_adjoint_base = operation.base
                 invert_param = True
             else:
-                name = operation.name
+                op_adjoint_base = operation
                 invert_param = False
+            name = op_adjoint_base.name
             method = getattr(state, name, None)
             wires = list(operation.wires)
 
@@ -279,18 +285,17 @@ def _apply_lightning(
                 )
             elif isinstance(operation, qml.PauliRot):
                 method = getattr(state, "applyPauliRot")
-                # pylint: disable=protected-access
-                paulis = operation._hyperparameters[
+                paulis = operation._hyperparameters[  # pylint: disable=protected-access
                     "pauli_word"
-                ]  # pylint: disable=protected-access
+                ]
                 wires = [i for i, w in zip(wires, paulis) if w != "I"]
                 word = "".join(p for p in paulis if p != "I")
                 method(wires, invert_param, operation.parameters, word)
             elif method is not None:  # apply specialized gate
                 param = operation.parameters
                 method(wires, invert_param, param)
-            elif isinstance(operation, qml.ops.Controlled):  # apply n-controlled gate
-                self._apply_lightning_controlled(operation)
+            elif isinstance(op_adjoint_base, qml.ops.Controlled):  # apply n-controlled gate
+                self._apply_lightning_controlled(op_adjoint_base, invert_param)
             else:  # apply gate as a matrix
                 # Inverse can be set to False since qml.matrix(operation) is already in
                 # inverted form