Add MidMeasureMP Kokkos

pennylane_lightning/core/_version.py

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

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

@@ -757,6 +757,60 @@ class StateVectorKokkos final
         return -static_cast<fp_t>(0.5);
     }
 
+    /**
+     * @brief Collapse the state vector as after having measured one of the
+     * qubits.
+     *
+     * The branch parameter imposes the measurement result on the given wire.
+     *
+     * @param wire Wire to collapse.
+     * @param branch Branch 0 or 1.
+     */
+    void collapse(const std::size_t wire, const bool branch) {
+        auto &&num_qubits = this->getNumQubits();
+
+        const size_t stride = pow(2, num_qubits_ - (1 + wire));
+        const size_t vec_size = pow(2, num_qubits_);
+        const auto section_size = vec_size / stride;
+        const auto half_section_size = section_size / 2;
+
+        const size_t negbranch = branch ? 0 : 1;
+
+        Kokkos::MDRangePolicy<DoubleLoopRank> policy_2d(
+            {0, 0}, {half_section_size, stride});
+        Kokkos::parallel_for(
+            policy_2d,
+            collapseFunctor<fp_t>(*data_, num_qubits, stride, negbranch));
+
+        normalize();
+    }
+
+    /**
+     * @brief Normalize vector (to have norm 1).
+     */
+    void normalize() {
+        KokkosVector sv_view =
+            getView(); // circumvent error capturing this with KOKKOS_LAMBDA
+
+        // TODO: @tomlqc what about squaredNorm()
+        PrecisionT squaredNorm = 0.0;
+        Kokkos::parallel_reduce(
+            sv_view.size(),
+            KOKKOS_LAMBDA(const size_t i, PrecisionT &sum) {
+                sum += std::norm<PrecisionT>(sv_view(i));
+            },
+            squaredNorm);
+
+        PL_ABORT_IF(squaredNorm <
+                        std::numeric_limits<PrecisionT>::epsilon() * 1e2,
+                    "vector has norm close to zero and can't be normalized");
+
+        std::complex<PrecisionT> inv_norm = 1. / std::sqrt(squaredNorm);
+        Kokkos::parallel_for(
+            sv_view.size(),
+            KOKKOS_LAMBDA(const size_t i) { sv_view(i) *= inv_norm; });
+    }
+
     /**
      * @brief Update data of the class
      *

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

@@ -972,4 +972,35 @@ template <class PrecisionT, bool adj = false> struct generatorMultiRZFunctor {
             1 - 2 * int(Kokkos::Impl::bit_count(k & wires_parity) % 2));
     }
 };
+
+template <class PrecisionT> struct collapseFunctor {
+    using ComplexT = Kokkos::complex<PrecisionT>;
+    using KokkosComplexVector = Kokkos::View<ComplexT *>;
+
+    KokkosComplexVector arr;
+    std::size_t num_qubits;
+    std::size_t stride;
+    std::size_t negbranch;
+
+    collapseFunctor(KokkosComplexVector &arr_, std::size_t num_qubits_,
+                    std::size_t stride_, std::size_t negbranch_) {
+        arr = arr_;
+        num_qubits = num_qubits_;
+        stride = stride_;
+        negbranch = negbranch_;
+    }
+
+    // zero half the entries
+    // the "half" entries depend on the stride
+    // *_*_*_*_ for stride 1
+    // **__**__ for stride 2
+    // ****____ for stride 4
+
+    KOKKOS_INLINE_FUNCTION
+    void operator()(const std::size_t left, const std::size_t right) const {
+        const size_t offset = stride * (negbranch + 2 * left);
+        arr[offset + right] = ComplexT{0., 0.};
+    }
+};
+
 } // namespace Pennylane::LightningKokkos::Functors

pennylane_lightning/core/src/simulators/lightning_kokkos/tests/Test_StateVectorLKokkos.cpp

@@ -308,4 +308,72 @@ TEMPLATE_TEST_CASE("StateVectorKokkos::StateVectorKokkos",
         REQUIRE(sv.getDataVector() == data_);
         // REQUIRE(sv.getDataVector() == approx(st_data));
     }
-}
+}
+
+TEMPLATE_TEST_CASE("StateVectorKokkos::collapse", "[StateVectorKokkos]", float,
+                   double) {
+    using PrecisionT = TestType;
+    using ComplexT = typename StateVectorKokkos<PrecisionT>::ComplexT;
+    using TestVectorT = TestVector<ComplexT>;
+
+    const std::size_t num_qubits = 3;
+
+    // TODO @tomlqc use same template for testing all Lightning flavours?
+
+    SECTION("Collapse the state vector as after having measured one of the "
+            "qubits.") {
+        TestVectorT init_state = createPlusState_<ComplexT>(num_qubits);
+
+        const ComplexT coef{0.5, PrecisionT{0.0}};
+        const ComplexT zero{PrecisionT{0.0}, PrecisionT{0.0}};
+
+        std::vector<std::vector<std::vector<ComplexT>>> expected_state = {
+            {{coef, coef, coef, coef, zero, zero, zero, zero},
+             {coef, coef, zero, zero, coef, coef, zero, zero},
+             {coef, zero, coef, zero, coef, zero, coef, zero}},
+            {{zero, zero, zero, zero, coef, coef, coef, coef},
+             {zero, zero, coef, coef, zero, zero, coef, coef},
+             {zero, coef, zero, coef, zero, coef, zero, coef}},
+        };
+
+        std::size_t wire = GENERATE(0, 1, 2);
+        std::size_t branch = GENERATE(0, 1);
+        StateVectorKokkos<PrecisionT> sv(
+            reinterpret_cast<ComplexT *>(init_state.data()), init_state.size());
+        sv.collapse(wire, branch);
+
+        PrecisionT eps = std::numeric_limits<PrecisionT>::epsilon() * 10e3;
+        REQUIRE(isApproxEqual(sv.getData(), sv.getDataVector().size(),
+                              expected_state[branch][wire].data(),
+                              expected_state[branch][wire].size(), eps));
+    }
+}
+
+TEMPLATE_TEST_CASE("StateVectorKokkos::normalize", "[StateVectorKokkos]", float,
+                   double) {
+    using PrecisionT = TestType;
+    using ComplexT = typename StateVectorKokkos<PrecisionT>::ComplexT;
+
+    // TODO @tomlqc use same template for testing all Lightning flavours?
+
+    SECTION("Normalize state vector.") {
+        const ComplexT init{1.0, PrecisionT{0.0}};
+        const ComplexT half{0.5, PrecisionT{0.0}};
+        const ComplexT zero{PrecisionT{0.0}, PrecisionT{0.0}};
+
+        std::vector<ComplexT> init_state = {init, zero, init, init,
+                                            zero, zero, init, zero};
+
+        std::vector<ComplexT> expected_state = {half, zero, half, half,
+                                                zero, zero, half, zero};
+
+        StateVectorKokkos<PrecisionT> sv(
+            reinterpret_cast<ComplexT *>(init_state.data()), init_state.size());
+        sv.normalize();
+
+        PrecisionT eps = std::numeric_limits<PrecisionT>::epsilon() * 1e3;
+        REQUIRE(isApproxEqual(sv.getData(), sv.getDataVector().size(),
+                              expected_state.data(), expected_state.size(),
+                              eps));
+    }
+}

pennylane_lightning/core/src/simulators/lightning_qubit/tests/Test_StateVectorLQubit.cpp

@@ -245,9 +245,35 @@ TEMPLATE_PRODUCT_TEST_CASE("StateVectorLQubit::collapse", "[StateVectorLQubit]",
 
         std::size_t wire = GENERATE(0, 1, 2);
         std::size_t branch = GENERATE(0, 1);
-        StateVectorLQubitManaged<PrecisionT> sv(init_state);
+        StateVectorT sv(init_state.data(), init_state.size());
         sv.collapse(wire, branch);
 
         REQUIRE(sv.getDataVector() == approx(expected_state[branch][wire]));
     }
 }
+
+TEMPLATE_PRODUCT_TEST_CASE("StateVectorLQubit::normalize",
+                           "[StateVectorLQubit]",
+                           (StateVectorLQubitManaged, StateVectorLQubitRaw),
+                           (float, double)) {
+    using StateVectorT = TestType;
+    using PrecisionT = typename StateVectorT::PrecisionT;
+    using ComplexT = typename StateVectorT::ComplexT;
+
+    SECTION("Normalize state vector.") {
+        const ComplexT init{1.0, PrecisionT{0.0}};
+        const ComplexT half{0.5, PrecisionT{0.0}};
+        const ComplexT zero{PrecisionT{0.0}, PrecisionT{0.0}};
+
+        std::vector<ComplexT> init_state = {init, zero, init, init,
+                                            zero, zero, init, zero};
+
+        std::vector<ComplexT> expected_state = {half, zero, half, half,
+                                                zero, zero, half, zero};
+
+        StateVectorT sv(init_state.data(), init_state.size());
+        sv.normalize();
+
+        REQUIRE(sv.getDataVector() == approx(expected_state));
+    }
+}

pennylane_lightning/core/src/utils/TestHelpers.hpp

@@ -288,6 +288,19 @@ auto createZeroState(size_t num_qubits) -> TestVector<ComplexT> {
     return res;
 }
 
+/**
+ * @brief create |+>^N
+ */
+template <typename ComplexT>
+auto createPlusState_(size_t num_qubits) -> TestVector<ComplexT> {
+    TestVector<ComplexT> res(size_t{1U} << num_qubits, 1.0,
+                             getBestAllocator<ComplexT>());
+    for (auto &elem : res) {
+        elem /= std::sqrt(1U << num_qubits);
+    }
+    return res;
+}
+
 /**
  * @brief create |+>^N
  */