Support arbitrary controlled operations in lightning.qubit

.github/CHANGELOG.md

@@ -7,11 +7,14 @@
 * Add shot-noise support for probs in the C++ layer. Probabilities are calculated from generated samples. All Lightning backends support this feature. Please note that target wires should be sorted in ascending manner.
 [(#568)](https://github.com/PennyLaneAI/pennylane-lightning/pull/568)
 
+* Add `LM` kernels to apply arbitrary controlled operations efficiently.
+  [(#516)](https://github.com/PennyLaneAI/pennylane-lightning/pull/516)
+
 * Add shots support for variance value, probs, sample, counts calculation for given observables (`NamedObs`, `TensorProd` and `Hamiltonian`) based on Pauli words, `Identity` and `Hadamard` in the C++ layer. All Lightning backends support this support feature.
-[(#561)](https://github.com/PennyLaneAI/pennylane-lightning/pull/561)
+  [(#561)](https://github.com/PennyLaneAI/pennylane-lightning/pull/561)
 
 * Add shots support for expectation value calculation for given observables (`NamedObs`, `TensorProd` and `Hamiltonian`) based on Pauli words, `Identity` and `Hadamard` in the C++ layer by adding `measure_with_samples` to the measurement interface. All Lightning backends support this support feature.
-[(#556)](https://github.com/PennyLaneAI/pennylane-lightning/pull/556)
+  [(#556)](https://github.com/PennyLaneAI/pennylane-lightning/pull/556)
 
 * `qml.QubitUnitary` operators can be included in a circuit differentiated with the adjoint method. Lightning handles circuits with arbitrary non-differentiable `qml.QubitUnitary` operators. 1,2-qubit `qml.QubitUnitary` operators with differentiable parameters can be differentiated using decomposition.
   [(#540)] (https://github.com/PennyLaneAI/pennylane-lightning/pull/540)

pennylane_lightning/core/_serialize.py

@@ -27,6 +27,7 @@
     Rot,
     Hamiltonian,
     SparseHamiltonian,
+    QubitUnitary,
 )
 from pennylane.operation import Tensor
 from pennylane.tape import QuantumTape
@@ -311,11 +312,32 @@ def serialize_ops(
         """
         names = []
         params = []
+        controlled_wires = []
         wires = []
         mats = []
 
         uses_stateprep = False
 
+        def get_wires(operation, single_op):
+            if operation.name[0:2] == "C(" or (
+                operation.name == "MultiControlledX"
+                and all(char == "1" for char in operation.hyperparameters["control_values"])
+            ):
+                name = "PauliX" if operation.name == "MultiControlledX" else operation.base.name
+                controlled_wires_list = operation.control_wires
+                if operation.name == "MultiControlledX":
+                    wires_list = list(set(operation.wires) - set(controlled_wires_list))
+                else:
+                    wires_list = operation.target_wires
+                if not hasattr(self.sv_type, name):
+                    single_op = QubitUnitary(matrix(single_op.base), single_op.base.wires)
+                    name = single_op.name
+            else:
+                name = single_op.name
+                wires_list = single_op.wires.tolist()
+                controlled_wires_list = []
+            return single_op, name, wires_list, controlled_wires_list
+
         for operation in tape.operations:
             if isinstance(operation, (BasisState, StatePrep)):
                 uses_stateprep = True
@@ -326,7 +348,7 @@ def serialize_ops(
                 op_list = [operation]
 
             for single_op in op_list:
-                name = single_op.name
+                single_op, name, wires_list, controlled_wires_list = get_wires(operation, single_op)
                 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
@@ -340,8 +362,8 @@ def serialize_ops(
                     params.append(single_op.parameters)
                     mats.append([])
 
-                wires_list = single_op.wires.tolist()
+                controlled_wires.append([wires_map[w] for w in controlled_wires_list])
                 wires.append([wires_map[w] for w in wires_list])
 
         inverses = [False] * len(names)
-        return (names, params, wires, inverses, mats), uses_stateprep
+        return (names, params, wires, inverses, mats, controlled_wires), uses_stateprep

pennylane_lightning/core/_version.py

@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.34.0-dev14"
+__version__ = "0.34.0-dev15"

pennylane_lightning/core/src/algorithms/AdjointJacobianBase.hpp

@@ -17,7 +17,6 @@
  * method.
  */
 #pragma once
-
 #include <span>
 
 #include "JacobianData.hpp"
@@ -56,11 +55,20 @@ template <class StateVectorT, class Derived> class AdjointJacobianBase {
                                 bool adj = false) {
         for (size_t op_idx = 0; op_idx < operations.getOpsName().size();
              op_idx++) {
-            state.applyOperation(operations.getOpsName()[op_idx],
-                                 operations.getOpsWires()[op_idx],
-                                 operations.getOpsInverses()[op_idx] ^ adj,
-                                 operations.getOpsParams()[op_idx],
-                                 operations.getOpsMatrices()[op_idx]);
+            if (operations.getOpsControlledWires()[op_idx].empty()) {
+                state.applyOperation(operations.getOpsName()[op_idx],
+                                     operations.getOpsWires()[op_idx],
+                                     operations.getOpsInverses()[op_idx] ^ adj,
+                                     operations.getOpsParams()[op_idx],
+                                     operations.getOpsMatrices()[op_idx]);
+            } else {
+                state.applyOperation(operations.getOpsName()[op_idx],
+                                     operations.getOpsControlledWires()[op_idx],
+                                     operations.getOpsWires()[op_idx],
+                                     operations.getOpsInverses()[op_idx] ^ adj,
+                                     operations.getOpsParams()[op_idx],
+                                     operations.getOpsMatrices()[op_idx]);
+            }
         }
     }
 
@@ -77,11 +85,20 @@ template <class StateVectorT, class Derived> class AdjointJacobianBase {
     inline void applyOperationAdj(UpdatedStateVectorT &state,
                                   const OpsData<StateVectorT> &operations,
                                   size_t op_idx) {
-        state.applyOperation(operations.getOpsName()[op_idx],
-                             operations.getOpsWires()[op_idx],
-                             !operations.getOpsInverses()[op_idx],
-                             operations.getOpsParams()[op_idx],
-                             operations.getOpsMatrices()[op_idx]);
+        if (operations.getOpsControlledWires()[op_idx].empty()) {
+            state.applyOperation(operations.getOpsName()[op_idx],
+                                 operations.getOpsWires()[op_idx],
+                                 !operations.getOpsInverses()[op_idx],
+                                 operations.getOpsParams()[op_idx],
+                                 operations.getOpsMatrices()[op_idx]);
+        } else {
+            state.applyOperation(operations.getOpsName()[op_idx],
+                                 operations.getOpsControlledWires()[op_idx],
+                                 operations.getOpsWires()[op_idx],
+                                 !operations.getOpsInverses()[op_idx],
+                                 operations.getOpsParams()[op_idx],
+                                 operations.getOpsMatrices()[op_idx]);
+        }
     }
 
     /**
@@ -117,6 +134,24 @@ template <class StateVectorT, class Derived> class AdjointJacobianBase {
         return sv.applyGenerator(op_name, wires, adj);
     }
 
+    /**
+     * @brief Applies the gate generator for a given parametric gate. Returns
+     * the associated scaling coefficient.
+     *
+     * @param sv Statevector data to operate upon.
+     * @param op_name Name of parametric gate.
+     * @param controlled_wires Control wires.
+     * @param wires Wires to operate upon.
+     * @param adj Indicate whether to take the adjoint of the operation.
+     * @return PrecisionT Generator scaling coefficient.
+     */
+    inline auto applyGenerator(StateVectorT &sv, const std::string &op_name,
+                               const std::vector<size_t> &controlled_wires,
+                               const std::vector<size_t> &wires, const bool adj)
+        -> PrecisionT {
+        return sv.applyGenerator(op_name, controlled_wires, wires, adj);
+    }
+
     /**
      * @brief Apply a given `%Observable<StateVectorT>` object to
      * `%StateVectorT`.

pennylane_lightning/core/src/algorithms/JacobianData.hpp

@@ -49,8 +49,39 @@ template <class StateVectorT> class OpsData {
     const std::vector<std::vector<size_t>> ops_wires_;
     const std::vector<bool> ops_inverses_;
     const std::vector<std::vector<ComplexT>> ops_matrices_;
+    const std::vector<std::vector<size_t>> ops_controlled_wires_;
 
   public:
+    /**
+     * @brief Construct an OpsData object, representing the serialized
+     * operations to apply upon the `%StateVector`.
+     *
+     * @param ops_name Name of each operation to apply.
+     * @param ops_params Parameters for a given operation ({} if optional).
+     * @param ops_wires Wires upon which to apply operation.
+     * @param ops_inverses Value to represent whether given operation is
+     * adjoint.
+     * @param ops_matrices Numerical representation of given matrix if not
+     * supported.
+     * @param ops_controlled_wires Control wires
+     */
+    OpsData(std::vector<std::string> ops_name,
+            const std::vector<std::vector<PrecisionT>> &ops_params,
+            std::vector<std::vector<size_t>> ops_wires,
+            std::vector<bool> ops_inverses,
+            std::vector<std::vector<ComplexT>> ops_matrices,
+            std::vector<std::vector<size_t>> ops_controlled_wires)
+        : num_par_ops_{0}, ops_name_{std::move(ops_name)},
+          ops_params_{ops_params}, ops_wires_{std::move(ops_wires)},
+          ops_inverses_{std::move(ops_inverses)},
+          ops_matrices_{std::move(ops_matrices)},
+          ops_controlled_wires_{std::move(ops_controlled_wires)} {
+        for (const auto &p : ops_params) {
+            num_par_ops_ += static_cast<size_t>(!p.empty());
+        }
+        num_nonpar_ops_ = ops_params.size() - num_par_ops_;
+    };
+
     /**
      * @brief Construct an OpsData object, representing the serialized
      * operations to apply upon the `%StateVector`.
@@ -71,11 +102,10 @@ template <class StateVectorT> class OpsData {
         : num_par_ops_{0}, ops_name_{std::move(ops_name)},
           ops_params_{ops_params}, ops_wires_{std::move(ops_wires)},
           ops_inverses_{std::move(ops_inverses)},
-          ops_matrices_{std::move(ops_matrices)} {
+          ops_matrices_{std::move(ops_matrices)},
+          ops_controlled_wires_(ops_name.size()) {
         for (const auto &p : ops_params) {
-            if (!p.empty()) {
-                num_par_ops_++;
-            }
+            num_par_ops_ += static_cast<size_t>(!p.empty());
         }
         num_nonpar_ops_ = ops_params.size() - num_par_ops_;
     };
@@ -97,11 +127,10 @@ template <class StateVectorT> class OpsData {
         : num_par_ops_{0}, ops_name_{ops_name}, ops_params_{ops_params},
           ops_wires_{std::move(ops_wires)},
           ops_inverses_{std::move(ops_inverses)},
-          ops_matrices_(ops_name.size()) {
+          ops_matrices_(ops_name.size()),
+          ops_controlled_wires_(ops_name.size()) {
         for (const auto &p : ops_params) {
-            if (p.size() > 0) {
-                num_par_ops_++;
-            }
+            num_par_ops_ += static_cast<size_t>(!p.empty());
         }
         num_nonpar_ops_ = ops_params.size() - num_par_ops_;
     };
@@ -140,6 +169,15 @@ template <class StateVectorT> class OpsData {
         -> const std::vector<std::vector<size_t>> & {
         return ops_wires_;
     }
+    /**
+     * @brief Get the controlled wires for each operation.
+     *
+     * @return const std::vector<std::vector<size_t>>&
+     */
+    [[nodiscard]] auto getOpsControlledWires() const
+        -> const std::vector<std::vector<size_t>> & {
+        return ops_controlled_wires_;
+    }
     /**
      * @brief Get the adjoint flag for each operation.
      *

pennylane_lightning/core/src/algorithms/tests/Test_AdjointJacobian.cpp

@@ -25,6 +25,7 @@ using namespace Pennylane::Util;
 
 #ifdef _ENABLE_PLQUBIT
 constexpr bool BACKEND_FOUND = true;
+constexpr bool SUPPORTS_CTRL = true;
 
 #include "AdjointJacobianLQubit.hpp"
 #include "ObservablesLQubit.hpp"
@@ -40,6 +41,7 @@ using namespace Pennylane::LightningQubit::Observables;
 
 #elif _ENABLE_PLKOKKOS == 1
 constexpr bool BACKEND_FOUND = true;
+constexpr bool SUPPORTS_CTRL = false;
 
 #include "AdjointJacobianKokkos.hpp"
 #include "ObservablesKokkos.hpp"
@@ -55,6 +57,7 @@ using namespace Pennylane::LightningKokkos::Observables;
 
 #elif _ENABLE_PLGPU == 1
 constexpr bool BACKEND_FOUND = true;
+constexpr bool SUPPORTS_CTRL = false;
 #include "AdjointJacobianGPU.hpp"
 #include "ObservablesGPU.hpp"
 #include "TestHelpersStateVectors.hpp"
@@ -69,6 +72,7 @@ using namespace Pennylane::LightningGPU::Observables;
 
 #else
 constexpr bool BACKEND_FOUND = false;
+constexpr bool SUPPORTS_CTRL = false;
 using TestStateVectorBackends = Pennylane::Util::TypeList<void>;
 
 template <class StateVector> struct StateVectorToName {};
@@ -134,6 +138,44 @@ template <typename TypeList> void testAdjointJacobian() {
             }
         }
 
+        DYNAMIC_SECTION("Op=PhaseShift, Obs=Y - "
+                        << StateVectorToName<StateVectorT>::name) {
+            if (SUPPORTS_CTRL) {
+                const std::vector<size_t> tp{0};
+                const size_t num_qubits = GENERATE(2, 3, 4);
+
+                const size_t num_params = 3;
+                const size_t num_obs = 1;
+                const auto obs = std::make_shared<NamedObs<StateVectorT>>(
+                    "PauliY", std::vector<size_t>{num_qubits - 1});
+                std::vector<PrecisionT> jacobian(num_obs * tp.size(), 0);
+
+                for (const auto &p : param) {
+                    std::vector<std::vector<size_t>> controls{
+                        std::vector<size_t>(num_qubits - 1)};
+                    std::iota(controls[0].begin(), controls[0].end(), 0);
+                    auto ops = OpsData<StateVectorT>({"PhaseShift"}, {{p}},
+                                                     {{num_qubits - 1}},
+                                                     {false}, {{}}, controls);
+
+                    std::vector<ComplexT> cdata(1U << num_qubits);
+                    cdata[cdata.size() - 2] =
+                        Pennylane::Util::INVSQRT2<PrecisionT>();
+                    cdata[cdata.size() - 1] =
+                        Pennylane::Util::INVSQRT2<PrecisionT>();
+
+                    StateVectorT psi(cdata.data(), cdata.size());
+                    JacobianData<StateVectorT> tape{
+                        num_params, psi.getLength(), psi.getData(), {obs}, ops,
+                        tp};
+                    adj.adjointJacobian(std::span{jacobian}, tape, psi, true);
+
+                    CAPTURE(jacobian);
+                    CHECK(cos(p) == Approx(jacobian[0]));
+                }
+            }
+        }
+
         DYNAMIC_SECTION("Op=RX, Obs=Z - "
                         << StateVectorToName<StateVectorT>::name) {
             const std::vector<size_t> tp{0};
@@ -189,6 +231,7 @@ template <typename TypeList> void testAdjointJacobian() {
                 CHECK(cos(p) == Approx(jacobian[0]).margin(1e-7));
             }
         }
+
         DYNAMIC_SECTION("Op=RX, Obs=[Z,Z] - "
                         << StateVectorToName<StateVectorT>::name) {
             std::vector<size_t> tp{0};
@@ -366,7 +409,9 @@ template <typename TypeList> void testAdjointJacobian() {
                 {{0}, {0}, {0}, {0, 1}, {1, 2}, {1}, {1}, {1}},
                 {false, false, false, false, false, false, false, false},
                 std::vector<std::vector<ComplexT>>{
-                    {}, {}, {}, cnot, {}, {}, {}, {}});
+                    {}, {}, {}, cnot, {}, {}, {}, {}},
+                std::vector<std::vector<size_t>>{
+                    {}, {}, {}, {}, {}, {}, {}, {}});
 
             JacobianData<StateVectorT> tape{
                 num_params, psi.getLength(), psi.getData(), {obs}, ops, tp};
@@ -436,7 +481,7 @@ template <typename TypeList> void testAdjointJacobian() {
             }
         }
 
-        DYNAMIC_SECTION("Mixed Ops, Obs and TParams- "
+        DYNAMIC_SECTION("Mixed Ops, Obs and TParams - "
                         << StateVectorToName<StateVectorT>::name) {
             std::vector<PrecisionT> param{-M_PI / 7, M_PI / 5, 2 * M_PI / 3};
             const std::vector<size_t> t_params{1, 2, 3};

pennylane_lightning/core/src/algorithms/tests/mpi/Test_AdjointJacobianMPI.cpp

@@ -264,7 +264,9 @@ template <typename TypeList> void testAdjointJacobian() {
                 {{0}, {0}, {0}, {0, 1}, {1, 2}, {1}, {1}, {1}},
                 {false, false, false, false, false, false, false, false},
                 std::vector<std::vector<ComplexT>>{
-                    {}, {}, {}, cnot, {}, {}, {}, {}});
+                    {}, {}, {}, cnot, {}, {}, {}, {}},
+                std::vector<std::vector<size_t>>{
+                    {}, {}, {}, {}, {}, {}, {}, {}});
 
             JacobianDataMPI<StateVectorT> tape{num_params, psi, {obs}, ops, tp};