Merge branch 'lk-control-generator-NQ' into lk-control-adjoint

PennyLaneAI · Nov 13, 2024 · 8bef76f · 8bef76f
2 parents 694cfef + 001d353
commit 8bef76f
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 21 deletions.
diff --git a/pennylane_lightning/core/src/simulators/lightning_kokkos/gates/BasicGateFunctors.hpp b/pennylane_lightning/core/src/simulators/lightning_kokkos/gates/BasicGateFunctors.hpp
@@ -39,14 +39,6 @@ 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;
@@ -78,12 +70,8 @@ template <class PrecisionT, class FuncT> class applyNCNFunctor {
         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::parallel_for(Kokkos::TeamPolicy(two2N, Kokkos::AUTO, dim),
+                             *this);
     }
     KOKKOS_FUNCTION void operator()(const MemberType &teamMember) const {
         const std::size_t k = teamMember.league_rank();

diff --git a/pennylane_lightning/core/src/simulators/lightning_kokkos/gates/MatrixGateFunctors.hpp b/pennylane_lightning/core/src/simulators/lightning_kokkos/gates/MatrixGateFunctors.hpp
@@ -156,16 +156,22 @@ template <class Precision> struct NCMultiQubitOpFunctor {
     void operator()(const MemberType &teamMember) const {
         const std::size_t k = teamMember.league_rank();
         ScratchViewComplex coeffs_in(teamMember.team_scratch(0), dim);
+        ScratchViewSizeT indices_scratch(teamMember.team_scratch(0), dim);
         const std::size_t offset = parity_2_offset(parity, k);
+
         if (teamMember.team_rank() == 0) {
-            for (std::size_t i = 0; i < dim; i++) {
-                coeffs_in(i) = arr(indices(i) + offset);
-            }
+            Kokkos::parallel_for(Kokkos::ThreadVectorRange(teamMember, dim),
+                                 [&](const std::size_t inner_idx) {
+                                     coeffs_in(inner_idx) =
+                                         arr(indices(inner_idx) + offset);
+                                     indices_scratch(inner_idx) =
+                                         indices(inner_idx);
+                                 });
         }
-
+        teamMember.team_barrier();
         Kokkos::parallel_for(
             Kokkos::TeamThreadRange(teamMember, dim), [&](const std::size_t i) {
-                const auto idx = indices(i) + offset;
+                const auto idx = indices_scratch(i) + offset;
                 arr(idx) = 0.0;
                 const std::size_t base_idx = i * dim;
 

diff --git a/...ing/core/src/simulators/lightning_kokkos/gates/tests/Test_StateVectorKokkos_Generator.cpp b/...ing/core/src/simulators/lightning_kokkos/gates/tests/Test_StateVectorKokkos_Generator.cpp
@@ -642,4 +642,4 @@ TEMPLATE_TEST_CASE("StateVectorKokkos::applyControlledGenerator CRX/Y/Z",
             }
         }
     }
-}
+}
diff --git a/...ghtning/core/src/simulators/lightning_kokkos/gates/tests/Test_StateVectorKokkos_Param.cpp b/...ghtning/core/src/simulators/lightning_kokkos/gates/tests/Test_StateVectorKokkos_Param.cpp
@@ -1022,7 +1022,7 @@ TEMPLATE_TEST_CASE("StateVectorKokkos::applyOperation param "
         }
     }
 
-    SECTION("2-controlled SingleExcitationMinus") {
+    SECTION("2-controlled MultiRZ - c{0,2}") {
         Kokkos::deep_copy(sv_gate.getView(), ini_sv);
 
         const std::vector<std::size_t> control_wires = {0, 2};
@@ -1056,6 +1056,41 @@ TEMPLATE_TEST_CASE("StateVectorKokkos::applyOperation param "
             CHECK(real(sv_gate_host[j]) == Approx(real(expected_result[j])));
         }
     }
+
+    SECTION("3-controlled MultiRZ - c{0,1,2}") {
+        Kokkos::deep_copy(sv_gate.getView(), ini_sv);
+
+        const std::vector<std::size_t> control_wires = {0, 1, 2};
+        const std::vector<bool> control_values = {false, true, false};
+        const std::vector<std::size_t> target_wire = {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.24829083, -0.02918331},
+                                              ComplexT{0.24829083, +0.02918331},
+                                              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.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(