(LK-C-3) Add controlled MultiRZ support to Lightning Kokkos

pennylane_lightning/core/src/simulators/lightning_kokkos/gates/BasicGateFunctors.hpp

@@ -36,6 +36,64 @@ using Pennylane::LightningKokkos::Util::vector2view;
 /// @endcond
 
 namespace Pennylane::LightningKokkos::Functors {
+template <class PrecisionT, class FuncT> class applyNCNFunctor {
+    using KokkosComplexVector = Kokkos::View<Kokkos::complex<PrecisionT> *>;
+    using KokkosIntVector = Kokkos::View<std::size_t *>;
+    using ScratchViewComplex =
+        Kokkos::View<Kokkos::complex<PrecisionT> *,
+                     Kokkos::DefaultExecutionSpace::scratch_memory_space,
+                     Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
+    using ScratchViewSizeT =
+        Kokkos::View<std::size_t *,
+                     Kokkos::DefaultExecutionSpace::scratch_memory_space,
+                     Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
+    using MemberType = Kokkos::TeamPolicy<>::member_type;
+
+    Kokkos::View<Kokkos::complex<PrecisionT> *> arr;
+    const FuncT core_function;
+    KokkosIntVector indices;
+    KokkosIntVector parity;
+    KokkosIntVector rev_wires;
+    KokkosIntVector rev_wire_shifts;
+    std::size_t dim;
+
+  public:
+    template <class ExecutionSpace>
+    applyNCNFunctor([[maybe_unused]] ExecutionSpace exec,
+                    Kokkos::View<Kokkos::complex<PrecisionT> *> arr_,
+                    std::size_t num_qubits,
+                    const std::vector<std::size_t> &controlled_wires,
+                    const std::vector<bool> &controlled_values,
+                    const std::vector<std::size_t> &wires, FuncT core_function_)
+        : arr(arr_), core_function(core_function_) {
+
+        std::size_t two2N =
+            std::exp2(num_qubits - wires.size() - controlled_wires.size());
+        dim = std::exp2(wires.size());
+        const auto &[parity_, rev_wires_] =
+            reverseWires(num_qubits, wires, controlled_wires);
+        parity = parity_;
+        std::vector<std::size_t> indices_ =
+            generateBitPatterns(wires, num_qubits);
+        ControlBitPatterns(indices_, num_qubits, controlled_wires,
+                           controlled_values);
+        indices = vector2view(indices_);
+        std::size_t scratch_size = ScratchViewComplex::shmem_size(dim) +
+                                   ScratchViewSizeT::shmem_size(dim);
+        Kokkos::parallel_for(
+            Kokkos::TeamPolicy(two2N, Kokkos::AUTO, dim)
+                .set_scratch_size(0, Kokkos::PerTeam(scratch_size)),
+            *this);
+    }
+    KOKKOS_FUNCTION void operator()(const MemberType &teamMember) const {
+        const std::size_t k = teamMember.league_rank();
+        const std::size_t offset = Util::parity_2_offset(parity, k);
+        Kokkos::parallel_for(Kokkos::TeamThreadRange(teamMember, dim),
+                             [&](const std::size_t i) {
+                                 core_function(arr, i, indices, offset);
+                             });
+    }
+};
 
 template <class PrecisionT, class FuncT, bool has_controls>
 class applyNC1Functor {};
@@ -1714,6 +1772,40 @@ void applyMultiRZ(Kokkos::View<Kokkos::complex<PrecisionT> *> arr_,
         });
 }
 
+template <class ExecutionSpace, class PrecisionT>
+void applyNCMultiRZ(Kokkos::View<Kokkos::complex<PrecisionT> *> arr_,
+                    const std::size_t num_qubits,
+                    const std::vector<std::size_t> &controlled_wires,
+                    const std::vector<bool> &controlled_values,
+                    const std::vector<std::size_t> &wires,
+                    const bool inverse = false,
+                    const std::vector<PrecisionT> &params = {}) {
+    const PrecisionT &angle = params[0];
+    const Kokkos::complex<PrecisionT> shift_0 = Kokkos::complex<PrecisionT>{
+        std::cos(angle / 2),
+        (inverse) ? std::sin(angle / 2) : -std::sin(angle / 2)};
+    const Kokkos::complex<PrecisionT> shift_1 = Kokkos::conj(shift_0);
+    std::size_t wires_parity = 0U;
+    wires_parity =
+        std::accumulate(wires.begin(), wires.end(), std::size_t{0},
+                        [num_qubits](std::size_t acc, std::size_t wire) {
+                            return acc | (static_cast<std::size_t>(1U)
+                                          << (num_qubits - wire - 1));
+                        });
+    auto core_function = KOKKOS_LAMBDA(
+        Kokkos::View<Kokkos::complex<PrecisionT> *> arr, const std::size_t i,
+        Kokkos::View<std::size_t *> indices, std::size_t offset) {
+        const std::size_t index = indices(i);
+        arr(index + offset) *=
+            (Kokkos::Impl::bit_count((index + offset) & wires_parity) % 2 == 0)
+                ? shift_0
+                : shift_1;
+    };
+
+    applyNCNFunctor(ExecutionSpace{}, arr_, num_qubits, controlled_wires,
+                    controlled_values, wires, core_function);
+}
+
 template <class ExecutionSpace, class PrecisionT>
 void applyPauliRot(Kokkos::View<Kokkos::complex<PrecisionT> *> arr_,
                    const std::size_t num_qubits,
@@ -2009,6 +2101,11 @@ void applyNCNamedOperation(const ControlledGateOperation gateop,
                                            controlled_values, wires, inverse,
                                            params);
         return;
+    case ControlledGateOperation::MultiRZ:
+        applyNCMultiRZ<ExecutionSpace>(arr_, num_qubits, controlled_wires,
+                                       controlled_values, wires, inverse,
+                                       params);
+        return;
     default:
         PL_ABORT("Controlled gate operation does not exist.");
     }

pennylane_lightning/core/src/simulators/lightning_kokkos/gates/tests/Test_StateVectorKokkos_Param.cpp

@@ -958,6 +958,41 @@ TEMPLATE_TEST_CASE("StateVectorKokkos::applyOperation param "
             CHECK(real(sv_gate_host[j]) == Approx(real(expected_result[j])));
         }
     }
+
+    SECTION("2-controlled SingleExcitationMinus") {
+        Kokkos::deep_copy(sv_gate.getView(), ini_sv);
+
+        const std::vector<std::size_t> control_wires = {0, 2};
+        const std::vector<bool> control_values = {true, false};
+        const std::vector<std::size_t> target_wire = {1, 3};
+        const TestType param = 0.234;
+        sv_gate.applyOperation("MultiRZ", control_wires, control_values,
+                               target_wire, inverse, {param});
+        auto sv_gate_host = Kokkos::create_mirror_view_and_copy(
+            Kokkos::HostSpace{}, sv_gate.getView());
+
+        std::vector<ComplexT> expected_result{// Generated using Pennylane
+                                              ComplexT{0.25, 0.0},
+                                              ComplexT{0.25, 0.0},
+                                              ComplexT{0.25, 0.0},
+                                              ComplexT{0.25, 0.0},
+                                              ComplexT{0.25, 0.0},
+                                              ComplexT{0.25, 0.0},
+                                              ComplexT{0.25, 0.0},
+                                              ComplexT{0.25, 0.0},
+                                              ComplexT{0.24829083, -0.02918331},
+                                              ComplexT{0.24829083, +0.02918331},
+                                              ComplexT{0.25, 0.0},
+                                              ComplexT{0.25, 0.0},
+                                              ComplexT{0.24829083, +0.02918331},
+                                              ComplexT{0.24829083, -0.02918331},
+                                              ComplexT{0.25, 0.0},
+                                              ComplexT{0.25, 0.0}};
+        for (std::size_t j = 0; j < exp2(num_qubits); j++) {
+            CHECK(imag(sv_gate_host[j]) == Approx(imag(expected_result[j])));
+            CHECK(real(sv_gate_host[j]) == Approx(real(expected_result[j])));
+        }
+    }
 }
 
 TEMPLATE_TEST_CASE(

pennylane_lightning/lightning_kokkos/_state_vector.py

@@ -285,7 +285,7 @@ def _apply_lightning(
                 param = operation.parameters
                 method(wires, invert_param, param)
             elif isinstance(operation, qml.ops.Controlled) and not isinstance(
-                operation.base, (qml.QubitUnitary, qml.BlockEncode, qml.MultiRZ)
+                operation.base, (qml.QubitUnitary, qml.BlockEncode)
             ):  # apply n-controlled gate
                 # Kokkos does not support controlled gates except for GlobalPhase and single-qubit
                 self._apply_lightning_controlled(operation)

pennylane_lightning/lightning_kokkos/lightning_kokkos.py

@@ -128,6 +128,7 @@
         "C(DoubleExcitation)",
         "C(DoubleExcitationMinus)",
         "C(DoubleExcitationPlus)",
+        "C(MultiRZ)",
         "C(GlobalPhase)",
         "CRot",
         "IsingXX",

pennylane_lightning/lightning_kokkos/lightning_kokkos.toml

@@ -23,7 +23,7 @@ IsingXX                = { properties = [ "invertible", "controllable", "differe
 IsingXY                = { properties = [ "invertible", "controllable", "differentiable" ] }
 IsingYY                = { properties = [ "invertible", "controllable", "differentiable" ] }
 IsingZZ                = { properties = [ "invertible", "controllable", "differentiable" ] }
-MultiRZ                = { properties = [ "invertible",                 "differentiable" ] }
+MultiRZ                = { properties = [ "invertible", "controllable", "differentiable" ] }
 PauliX                 = { properties = [ "invertible", "controllable", "differentiable" ] }
 PauliY                 = { properties = [ "invertible", "controllable", "differentiable" ] }
 PauliZ                 = { properties = [ "invertible", "controllable", "differentiable" ] }

tests/test_gates.py

@@ -487,8 +487,6 @@ def test_controlled_qubit_gates(operation, n_qubits, control_value, tol):
     dev = qml.device(device_name, wires=n_qubits)
     threshold = 5 if device_name == "lightning.tensor" else 250
     num_wires = max(operation.num_wires, 1)
-    if device_name == "lightning.kokkos" and isinstance(operation, qml.MultiRZ):
-        pytest.skip("lightning.kokkos does not support controlled-multiRZ")
 
     for n_wires in range(num_wires + 1, num_wires + 4):
         wire_lists = list(itertools.permutations(range(0, n_qubits), n_wires))