(LK-C-4) Add controlled QubitUnitary support to Lightning Kokkos (#955)

### Before submitting

Please complete the following checklist when submitting a PR:

- [ ] All new features must include a unit test.
If you've fixed a bug or added code that should be tested, add a test to
the
      [`tests`](../tests) directory!

- [ ] All new functions and code must be clearly commented and
documented.
If you do make documentation changes, make sure that the docs build and
      render correctly by running `make docs`.

- [ ] Ensure that the test suite passes, by running `make test`.

- [ ] Add a new entry to the `.github/CHANGELOG.md` file, summarizing
the
      change, and including a link back to the PR.

- [ ] Ensure that code is properly formatted by running `make format`. 

When all the above are checked, delete everything above the dashed
line and fill in the pull request template.


------------------------------------------------------------------------------------------------------------

**Context:**

**Description of the Change:**
This PR adds support for controlled QubitUnitary for Lightning Kokkos.
There are specialized implementations for 1/2/3-qubits
(`applyNC1/2/3QubitOpFunctor`) and a general case implementation
(`NCMultiQubitOpFunctor`). These functors are defined in
`MatrixGateFunctors.hpp`. These are called by `applyControlledMatrix` or
`applyOperation`/`applyNCMultiQubitOp` in `StateVectorKokkos.hpp`

**Benefits:**
Performance benchmarks for gates are shown here:
https://www.notion.so/xanaduai/Lightning-Kokkos-Native-Controlled-Operation-Gate-Benchmarks-12ebc6bd17648017a2dcd237748b24fe

**Possible Drawbacks:**

**Related GitHub Issues:**

[sc-76776]

pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp

@@ -134,11 +134,11 @@ class StateVectorKokkos final
      *
      * @param index Index of the target element.
      */
-    void setBasisState(const std::size_t index) {
+    void setBasisState(std::size_t index) {
         KokkosVector sv_view =
             getView(); // circumvent error capturing this with KOKKOS_LAMBDA
         Kokkos::parallel_for(
-            sv_view.size(), KOKKOS_LAMBDA(const std::size_t i) {
+            sv_view.size(), KOKKOS_LAMBDA(std::size_t i) {
                 sv_view(i) =
                     (i == index) ? ComplexT{1.0, 0.0} : ComplexT{0.0, 0.0};
             });
@@ -197,7 +197,7 @@ class StateVectorKokkos final
         KokkosVector sv_view =
             getView(); // circumvent error capturing this with KOKKOS_LAMBDA
         Kokkos::parallel_for(
-            indices.size(), KOKKOS_LAMBDA(const std::size_t i) {
+            indices.size(), KOKKOS_LAMBDA(std::size_t i) {
                 sv_view(d_indices[i]) = d_values[i];
             });
     }
@@ -237,7 +237,7 @@ class StateVectorKokkos final
         auto d_wires = vector2view(wires);
         initZeros();
         Kokkos::parallel_for(
-            num_state, KOKKOS_LAMBDA(const std::size_t i) {
+            num_state, KOKKOS_LAMBDA(std::size_t i) {
                 std::size_t index{0U};
                 for (std::size_t w = 0; w < d_wires.size(); w++) {
                     const std::size_t bit = (i & (one << w)) >> w;
@@ -367,8 +367,7 @@ class StateVectorKokkos final
      * @param params Rotation angle.
      * @param word A Pauli word (e.g. "XYYX").
      */
-    void applyPauliRot(const std::vector<std::size_t> &wires,
-                       const bool inverse,
+    void applyPauliRot(const std::vector<std::size_t> &wires, bool inverse,
                        const std::vector<PrecisionT> &params,
                        const std::string &word) {
         PL_ABORT_IF_NOT(wires.size() == word.size(),
@@ -384,8 +383,9 @@ class StateVectorKokkos final
         auto two2N = BaseType::getLength();
         auto dataview = getView();
         Kokkos::parallel_for(
-            two2N, KOKKOS_LAMBDA(const std::size_t i) {
-                dataview(i) *= (inverse) ? conj(diagonal_(i)) : diagonal_(i);
+            two2N, KOKKOS_LAMBDA(std::size_t i) {
+                dataview(i) *=
+                    (inverse) ? Kokkos::conj(diagonal_(i)) : diagonal_(i);
             });
     }
 
@@ -399,17 +399,17 @@ class StateVectorKokkos final
     void applyMultiQubitOp(const KokkosVector matrix,
                            const std::vector<std::size_t> &wires,
                            bool inverse = false) {
-        auto &&num_qubits = this->getNumQubits();
-        std::size_t two2N = std::exp2(num_qubits - wires.size());
-        std::size_t dim = std::exp2(wires.size());
+        const std::size_t num_qubits = this->getNumQubits();
+        const std::size_t two2N = exp2(num_qubits - wires.size());
+        const std::size_t dim = exp2(wires.size());
         KokkosVector matrix_trans("matrix_trans", matrix.size());
 
         if (inverse) {
             Kokkos::MDRangePolicy<DoubleLoopRank> policy_2d({0, 0}, {dim, dim});
             Kokkos::parallel_for(
-                policy_2d,
-                KOKKOS_LAMBDA(const std::size_t i, const std::size_t j) {
-                    matrix_trans(i + j * dim) = conj(matrix(i * dim + j));
+                policy_2d, KOKKOS_LAMBDA(std::size_t i, std::size_t j) {
+                    matrix_trans(i + j * dim) =
+                        Kokkos::conj(matrix(i * dim + j));
                 });
         } else {
             matrix_trans = matrix;
@@ -436,12 +436,13 @@ class StateVectorKokkos final
                                                   matrix_trans, wires));
             break;
         default:
+            // TODO: explore runtime determine scratch space level (L0 vs L1)
             std::size_t scratch_size = ScratchViewComplex::shmem_size(dim) +
                                        ScratchViewSizeT::shmem_size(dim);
             Kokkos::parallel_for(
                 "multiQubitOpFunctor",
                 TeamPolicy(two2N, Kokkos::AUTO, dim)
-                    .set_scratch_size(0, Kokkos::PerTeam(scratch_size)),
+                    .set_scratch_size(1, Kokkos::PerTeam(scratch_size)),
                 multiQubitOpFunctor<PrecisionT>(*data_, num_qubits,
                                                 matrix_trans, wires));
             break;
@@ -472,20 +473,93 @@ class StateVectorKokkos final
         PL_ABORT_IF_NOT(controlled_wires.size() == controlled_values.size(),
                         "`controlled_wires` must have the same size as "
                         "`controlled_values`.");
-        PL_ABORT_IF_NOT(
-            array_contains(controlled_gate_names, std::string_view{opName}),
-            "Controlled matrix operation not yet supported.");
 
         if (controlled_wires.empty()) {
             return applyOperation(opName, wires, inverse, params, gate_matrix);
         }
+        if (array_contains(controlled_gate_names, std::string_view{opName})) {
+            const std::size_t num_qubits = this->getNumQubits();
+            const ControlledGateOperation gateop =
+                reverse_lookup(controlled_gate_names, std::string_view{opName});
+            applyNCNamedOperation<KokkosExecSpace>(
+                gateop, *data_, num_qubits, controlled_wires, controlled_values,
+                wires, inverse, params);
+        } else {
+            PL_ABORT_IF(gate_matrix.empty(),
+                        std::string("Operation does not exist for ") + opName +
+                            std::string(" and no matrix provided."));
+            return applyNCMultiQubitOp(vector2view(gate_matrix),
+                                       controlled_wires, controlled_values,
+                                       wires, inverse);
+        }
+    }
+
+    /**
+     * @brief Apply a controlled-multi qubit operator to the state vector using
+     * a matrix
+     *
+     * @param matrix Kokkos gate matrix in the device space.
+     * @param controlled_wires Control wires.
+     * @param controlled_values Control values (true or false).
+     * @param wires Wires to apply gate to.
+     * @param inverse Indicates whether to use adjoint of gate. (Default to
+     * false)
+     */
+    void applyNCMultiQubitOp(const KokkosVector matrix,
+                             const std::vector<std::size_t> &controlled_wires,
+                             const std::vector<bool> &controlled_values,
+                             const std::vector<std::size_t> &wires,
+                             bool inverse = false) {
 
         const std::size_t num_qubits = this->getNumQubits();
-        const ControlledGateOperation gateop =
-            reverse_lookup(controlled_gate_names, std::string_view{opName});
-        applyNCNamedOperation<KokkosExecSpace>(
-            gateop, *data_, num_qubits, controlled_wires, controlled_values,
-            wires, inverse, params);
+        const std::size_t two2N =
+            exp2(num_qubits - wires.size() - controlled_wires.size());
+        const std::size_t dim = exp2(wires.size());
+        KokkosVector matrix_trans("matrix_trans", matrix.size());
+
+        if (inverse) {
+            Kokkos::MDRangePolicy<DoubleLoopRank> policy_2d({0, 0}, {dim, dim});
+            Kokkos::parallel_for(
+                policy_2d, KOKKOS_LAMBDA(std::size_t i, std::size_t j) {
+                    matrix_trans(i + j * dim) =
+                        Kokkos::conj(matrix(i * dim + j));
+                });
+        } else {
+            matrix_trans = matrix;
+        }
+
+        switch (wires.size()) {
+        case 1:
+            Kokkos::parallel_for(two2N, applyNC1QubitOpFunctor<fp_t>(
+                                            *data_, num_qubits, matrix_trans,
+                                            controlled_wires, controlled_values,
+                                            wires));
+            break;
+        case 2:
+            Kokkos::parallel_for(two2N, applyNC2QubitOpFunctor<fp_t>(
+                                            *data_, num_qubits, matrix_trans,
+                                            controlled_wires, controlled_values,
+                                            wires));
+            break;
+        case 3:
+            Kokkos::parallel_for(two2N, applyNC3QubitOpFunctor<fp_t>(
+                                            *data_, num_qubits, matrix_trans,
+                                            controlled_wires, controlled_values,
+                                            wires));
+            break;
+        default:
+            // TODO: explore runtime determine scratch space level (L0 vs L1)
+            std::size_t scratch_size = ScratchViewComplex::shmem_size(dim) +
+                                       ScratchViewSizeT::shmem_size(dim);
+            Kokkos::parallel_for(
+                "multiNCQubitOpFunctor",
+                TeamPolicy(two2N, Kokkos::AUTO, dim)
+                    .set_scratch_size(1, Kokkos::PerTeam(scratch_size)),
+                NCMultiQubitOpFunctor<PrecisionT>(
+                    *data_, num_qubits, matrix_trans, controlled_wires,
+                    controlled_values, wires));
+            break;
+        }
     }
 
     /**
@@ -500,8 +574,8 @@ class StateVectorKokkos final
                             const std::vector<std::size_t> &wires,
                             bool inverse = false) {
         PL_ABORT_IF(wires.empty(), "Number of wires must be larger than 0");
-        std::size_t n = static_cast<std::size_t>(1U) << wires.size();
-        KokkosVector matrix_(matrix, n * n);
+        const std::size_t n2 = exp2(wires.size() * 2);
+        KokkosVector matrix_(matrix, n2);
         applyMultiQubitOp(matrix_, wires, inverse);
     }
 
@@ -517,8 +591,7 @@ class StateVectorKokkos final
                             const std::vector<std::size_t> &wires,
                             bool inverse = false) {
         PL_ABORT_IF(wires.empty(), "Number of wires must be larger than 0");
-        std::size_t n = static_cast<std::size_t>(1U) << wires.size();
-        std::size_t n2 = n * n;
+        const std::size_t n2 = exp2(wires.size() * 2);
         KokkosVector matrix_("matrix_", n2);
         Kokkos::deep_copy(matrix_, UnmanagedConstComplexHostView(matrix, n2));
         applyMultiQubitOp(matrix_, wires, inverse);
@@ -541,6 +614,75 @@ class StateVectorKokkos final
         applyMatrix(matrix.data(), wires, inverse);
     }
 
+    /**
+     * @brief Apply a given matrix for controlled operations directly to the
+     * statevector using a raw matrix pointer vector.
+     *
+     * @param matrix Pointer to the array data (in row-major format).
+     * @param controlled_wires Controlled wires
+     * @param controlled_values Controlled values (true or false)
+     * @param wires Wires to apply gate to.
+     * @param inverse Indicate whether inverse should be taken.
+     */
+    inline void applyControlledMatrix(
+        ComplexT *matrix, const std::vector<std::size_t> &controlled_wires,
+        const std::vector<bool> &controlled_values,
+        const std::vector<std::size_t> &wires, bool inverse = false) {
+        PL_ABORT_IF(wires.empty(), "Number of wires must be larger than 0");
+        const std::size_t n2 = exp2(wires.size() * 2);
+        KokkosVector matrix_(matrix, n2);
+        applyNCMultiQubitOp(matrix_, controlled_wires, controlled_values, wires,
+                            inverse);
+    }
+
+    /**
+     * @brief Apply a given matrix directly to the statevector using a
+     * raw matrix pointer vector.
+     *
+     * @param matrix Pointer to the array data (in row-major format).
+     * @param controlled_wires Controlled wires
+     * @param controlled_values Controlled values (true or false)
+     * @param wires Wires to apply gate to.
+     * @param inverse Indicate whether inverse should be taken.
+     */
+    inline void
+    applyControlledMatrix(const ComplexT *matrix,
+                          const std::vector<std::size_t> &controlled_wires,
+                          const std::vector<bool> &controlled_values,
+                          const std::vector<std::size_t> &wires,
+                          bool inverse = false) {
+        PL_ABORT_IF(wires.empty(), "Number of wires must be larger than 0");
+        const std::size_t n2 = exp2(wires.size() * 2);
+        KokkosVector matrix_("matrix_", n2);
+        Kokkos::deep_copy(matrix_, UnmanagedConstComplexHostView(matrix, n2));
+        applyNCMultiQubitOp(matrix_, controlled_wires, controlled_values, wires,
+                            inverse);
+    }
+
+    /**
+     * @brief Apply a given controlled-matrix directly to the statevector.
+     *
+     * @param matrix Vector containing the statevector data (in row-major
+     * format).
+     * @param controlled_wires Control wires.
+     * @param controlled_values Control values (false or true).
+     * @param wires Wires to apply gate to.
+     * @param inverse Indicate whether inverse should be taken.
+     */
+    inline void
+    applyControlledMatrix(const std::vector<ComplexT> &matrix,
+                          const std::vector<std::size_t> &controlled_wires,
+                          const std::vector<bool> &controlled_values,
+                          const std::vector<std::size_t> &wires,
+                          bool inverse = false) {
+        PL_ABORT_IF(wires.empty(), "Number of wires must be larger than 0");
+        PL_ABORT_IF(matrix.size() != exp2(2 * wires.size()),
+                    "The size of matrix does not match with the given "
+                    "number of wires");
+        applyControlledMatrix(matrix.data(), controlled_wires,
+                              controlled_values, wires, inverse);
+    }
+
     /**
      * @brief Apply a single generator to the state vector using the given
      * kernel.

pennylane_lightning/core/src/simulators/lightning_kokkos/bindings/LKokkosBindings.hpp

@@ -47,7 +47,26 @@ namespace Pennylane::LightningKokkos {
 using StateVectorBackends =
     Pennylane::Util::TypeList<StateVectorKokkos<float>,
                               StateVectorKokkos<double>, void>;
+/**
+ * @brief Register controlled matrix kernel.
+ */
+template <class StateVectorT>
+void applyControlledMatrix(
+    StateVectorT &st,
+    const py::array_t<std::complex<typename StateVectorT::PrecisionT>,
+                      py::array::c_style | py::array::forcecast> &matrix,
+    const std::vector<std::size_t> &controlled_wires,
+    const std::vector<bool> &controlled_values,
+    const std::vector<std::size_t> &wires, bool inverse = false) {
+    using ComplexT = typename StateVectorT::ComplexT;
+    st.applyControlledMatrix(
+        static_cast<const ComplexT *>(matrix.request().ptr), controlled_wires,
+        controlled_values, wires, inverse);
+}
 
+/**
+ * @brief Register controlled gates.
+ */
 template <class StateVectorT, class PyClass>
 void registerControlledGate(PyClass &pyclass) {
     using ParamT = typename StateVectorT::PrecisionT;
@@ -170,7 +189,9 @@ void registerBackendClassSpecificBindings(PyClass &pyclass) {
         .def("collapse", &StateVectorT::collapse,
              "Collapse the statevector onto the 0 or 1 branch of a given wire.")
         .def("normalize", &StateVectorT::normalize,
-             "Normalize the statevector to norm 1.");
+             "Normalize the statevector to norm 1.")
+        .def("applyControlledMatrix", &applyControlledMatrix<StateVectorT>,
+             "Apply controlled operation");
 }
 
 /**