Add MidMeasureMP Kokkos (#658)

* collapse() and normalized() with Kokkos * Auto update version * CodeFactor * normalize() abort if norm is zero * unittest for normalize() * Auto update version * unittest for normalize() in LighningQubit * Auto update version * Add MCM bindings and tests for L-Kokkos. (#672) * Add MCM bindings and tests for L-Kokkos. * Auto update version * Refactor collapse. * Try parallelizing tests. * Only run test_native_mcm once amongst all workflows since it takes quite a bit of time. * trigger ci --------- Co-authored-by: Dev version update bot <github-actions[bot]@users.noreply.github.com> * Do not run test_native_mcm with LK-GPU. * Auto update version * trigger ci * Update pennylane_lightning/core/src/simulators/lightning_kokkos/bindings/LKokkosBindings.hpp Co-authored-by: Ali Asadi <10773383+maliasadi@users.noreply.github.com> * Auto update version * Update pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp Co-authored-by: Ali Asadi <10773383+maliasadi@users.noreply.github.com> * Update pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp Co-authored-by: Ali Asadi <10773383+maliasadi@users.noreply.github.com> * Update pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp Co-authored-by: Ali Asadi <10773383+maliasadi@users.noreply.github.com> * Update pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp Co-authored-by: Ali Asadi <10773383+maliasadi@users.noreply.github.com> * Increase shotsfor flaky test_composite_mcm_single_measure_obs * Create separate device_allowed_operations list for ops supported by the device but not Catalyst. * Use isinstance instead of name comparison. * Check projector. * Update operations at device init. * Modify stopping condition in LK. * Cannot use instance with Hamiltonian check. * Update pennylane_lightning/core/src/simulators/lightning_kokkos/tests/Test_StateVectorLKokkos.cpp Co-authored-by: Amintor Dusko <87949283+AmintorDusko@users.noreply.github.com> * Update pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp Co-authored-by: Amintor Dusko <87949283+AmintorDusko@users.noreply.github.com> * Remove obsolete comment. * Auto update version * trigger ci --------- Co-authored-by: Dev version update bot <github-actions[bot]@users.noreply.github.com> Co-authored-by: Vincent Michaud-Rioux <vincent.michaud-rioux@xanadu.ai> Co-authored-by: Vincent Michaud-Rioux <vincentm@nanoacademic.com> Co-authored-by: Ali Asadi <10773383+maliasadi@users.noreply.github.com> Co-authored-by: Amintor Dusko <87949283+AmintorDusko@users.noreply.github.com>
PennyLaneAI · Apr 12, 2024 · d3b8bb0 · d3b8bb0
1 parent b76db37
commit d3b8bb0
Show file tree

Hide file tree

Showing 14 changed files with 262 additions and 38 deletions.
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
@@ -2,6 +2,9 @@
 
 ### New features since last release
 
+* `lightning.kokkos` supports mid-circuit measurements.
+  [(#672)](https://github.com/PennyLaneAI/pennylane-lightning/pull/672)
+
 * Add dynamic linking to LAPACK/OpenBlas shared objects in scipy.libs for both C++ and Python layer.
   [(#653)](https://github.com/PennyLaneAI/pennylane-lightning/pull/653)
 

diff --git a/.github/workflows/tests_gpu_cuda.yml b/.github/workflows/tests_gpu_cuda.yml
@@ -311,7 +311,7 @@ jobs:
           OMP_PROC_BIND: false
         run: |
           cd main/
-          PL_DEVICE=lightning.qubit python -m pytest tests/ $COVERAGE_FLAGS
+          PL_DEVICE=lightning.qubit python -m pytest tests/ -k "not test_native_mcm" $COVERAGE_FLAGS
           pl-device-test --device lightning.qubit --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
           pl-device-test --device lightning.qubit --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
           PL_DEVICE=lightning.gpu python -m pytest tests/ $COVERAGE_FLAGS

diff --git a/.github/workflows/tests_gpu_kokkos.yml b/.github/workflows/tests_gpu_kokkos.yml
@@ -303,9 +303,9 @@ jobs:
         run: |
           cd main/
           DEVICENAME=`echo ${{ matrix.pl_backend }} | sed "s/_/./g"`
+          PL_DEVICE=${DEVICENAME} python -m pytest tests/ -k "not test_native_mcm" $COVERAGE_FLAGS 2> /dev/null || echo Something went wrong with Pytest
           pl-device-test --device ${DEVICENAME} --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append 2> /dev/null || echo Something went wrong with pl-device-test shot=20000
           pl-device-test --device ${DEVICENAME} --shots=None --skip-ops $COVERAGE_FLAGS --cov-append 2> /dev/null || echo Something went wrong with pl-device-test shot=None
-          PL_DEVICE=${DEVICENAME} python -m pytest tests/ $COVERAGE_FLAGS 2> /dev/null || echo Something went wrong with Pytest
           mv coverage.xml coverage-${{ github.job }}-${{ matrix.pl_backend }}.xml
 
       - name: Install all backend devices
@@ -325,10 +325,10 @@ jobs:
           OMP_PROC_BIND: false
         run: |
           cd main/
-          PL_DEVICE=lightning.qubit python -m pytest tests/ $COVERAGE_FLAGS
+          PL_DEVICE=lightning.qubit python -m pytest tests/ -k "not test_native_mcm" $COVERAGE_FLAGS
           pl-device-test --device lightning.qubit --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
           pl-device-test --device lightning.qubit --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
-          PL_DEVICE=lightning.kokkos python -m pytest tests/ $COVERAGE_FLAGS
+          PL_DEVICE=lightning.kokkos python -m pytest tests/ -k "not test_native_mcm" $COVERAGE_FLAGS --cov-append
           pl-device-test --device lightning.kokkos --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
           pl-device-test --device lightning.kokkos --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
           mv coverage.xml coverage-${{ github.job }}-${{ matrix.pl_backend }}.xml

diff --git a/.github/workflows/tests_linux.yml b/.github/workflows/tests_linux.yml
@@ -21,6 +21,7 @@ env:
   COVERAGE_FLAGS: "--cov=pennylane_lightning --cov-report=term-missing --no-flaky-report -p no:warnings --tb=native"
   GCC_VERSION: 11
   OMP_NUM_THREADS: "2"
+  OMP_PROC_BIND: "false"
 
 concurrency:
   group: tests_linux-${{ github.ref }}-${{ inputs.lightning-version }}-${{ inputs.pennylane-version }}
@@ -353,8 +354,8 @@ jobs:
           python -m pip install pytest-xdist
           cd main/
           DEVICENAME=`echo ${{ matrix.pl_backend }} | sed "s/_/./g"`
-          OMP_NUM_THREADS=1 PL_DEVICE=${DEVICENAME} python -m pytest -n auto tests/ -k "not unitary_correct" $COVERAGE_FLAGS
-          PL_DEVICE=${DEVICENAME} python -m pytest tests/ -k "unitary_correct" $COVERAGE_FLAGS --cov-append
+          OMP_NUM_THREADS=1 PL_DEVICE=${DEVICENAME} python -m pytest -n auto tests/ -k "not unitary_correct and not test_native_mcm" $COVERAGE_FLAGS
+          PL_DEVICE=${DEVICENAME} python -m pytest tests/ -k "unitary_correct and not test_native_mcm" $COVERAGE_FLAGS --cov-append
           pl-device-test --device ${DEVICENAME} --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
           pl-device-test --device ${DEVICENAME} --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
           mv .coverage .coverage-${{ github.job }}-${{ matrix.pl_backend }}
@@ -558,7 +559,8 @@ jobs:
           python -m pip install pytest-xdist
           cd main/
           DEVICENAME=`echo ${{ matrix.pl_backend }} | sed "s/_/./g"`
-          PL_DEVICE=${DEVICENAME} python -m pytest tests/ $COVERAGE_FLAGS
+          PL_DEVICE=${DEVICENAME} python -m pytest tests/ -k "not test_native_mcm" $COVERAGE_FLAGS
+          OMP_NUM_THREADS=1 PL_DEVICE=${DEVICENAME} python -m pytest -n auto tests/ -k "test_native_mcm" $COVERAGE_FLAGS --cov-append
           pl-device-test --device ${DEVICENAME} --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
           pl-device-test --device ${DEVICENAME} --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
           mv .coverage .coverage-${{ github.job }}-${{ matrix.pl_backend }}
@@ -579,11 +581,11 @@ jobs:
           # TODO: Remove installing pytest-xdist with release v0.36.0
           python -m pip install pytest-xdist
           cd main/
-          OMP_NUM_THREADS=1 PL_DEVICE=lightning.qubit python -m pytest -n auto tests/ -k "not unitary_correct" $COVERAGE_FLAGS
-          PL_DEVICE=lightning.qubit python -m pytest tests/ -k "unitary_correct" $COVERAGE_FLAGS --cov-append
+          OMP_NUM_THREADS=1 PL_DEVICE=lightning.qubit python -m pytest -n auto tests/ -k "not unitary_correct and not test_native_mcm" $COVERAGE_FLAGS
+          PL_DEVICE=lightning.qubit python -m pytest tests/ -k "unitary_correct and not test_native_mcm" $COVERAGE_FLAGS --cov-append
           pl-device-test --device lightning.qubit --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
           pl-device-test --device lightning.qubit --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
-          PL_DEVICE=lightning.kokkos python -m pytest tests/ $COVERAGE_FLAGS --cov-append
+          PL_DEVICE=lightning.kokkos python -m pytest tests/ -k "not test_native_mcm" $COVERAGE_FLAGS --cov-append
           pl-device-test --device lightning.kokkos --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
           pl-device-test --device lightning.kokkos --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
           mv .coverage .coverage-${{ github.job }}-${{ matrix.pl_backend }}

diff --git a/pennylane_lightning/core/lightning_base.py b/pennylane_lightning/core/lightning_base.py
@@ -94,9 +94,9 @@ def stopping_condition(self):
         and observable) and returns ``True`` if supported by the device."""
 
         def accepts_obj(obj):
-            if obj.name == "QFT":
+            if isinstance(obj, qml.QFT):
                 return len(obj.wires) < 10
-            if obj.name == "GroverOperator":
+            if isinstance(obj, qml.GroverOperator):
                 return len(obj.wires) < 13
             is_not_tape = not isinstance(obj, qml.tape.QuantumTape)
             is_supported = getattr(self, "supports_operation", lambda name: False)(obj.name)

diff --git a/pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp b/pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp
@@ -757,6 +757,53 @@ class StateVectorKokkos final
         return -static_cast<fp_t>(0.5);
     }
 
+    /**
+     * @brief Collapse the state vector 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(std::size_t wire, bool branch) {
+        KokkosVector matrix("gate_matrix", 4);
+        Kokkos::parallel_for(
+            matrix.size(), KOKKOS_LAMBDA(std::size_t k) {
+                matrix(k) = ((k == 0 && branch == 0) || (k == 3 && branch == 1))
+                                ? ComplexT{1.0, 0.0}
+                                : ComplexT{0.0, 0.0};
+            });
+        applyMultiQubitOp(matrix, {wire}, false);
+        normalize();
+    }
+
+    /**
+     * @brief Normalize vector (to have norm 1).
+     */
+    void normalize() {
+        auto sv_view = getView();
+
+        PrecisionT squaredNorm = 0.0;
+        Kokkos::parallel_reduce(
+            sv_view.size(),
+            KOKKOS_LAMBDA(std::size_t i, PrecisionT & sum) {
+                const PrecisionT norm = Kokkos::abs(sv_view(i));
+                sum += norm * norm;
+            },
+            squaredNorm);
+
+        PL_ABORT_IF(squaredNorm <
+                        std::numeric_limits<PrecisionT>::epsilon() * 1e2,
+                    "vector has norm close to zero and can't be normalized");
+
+        const std::complex<PrecisionT> inv_norm =
+            1. / Kokkos::sqrt(squaredNorm);
+        Kokkos::parallel_for(
+            sv_view.size(),
+            KOKKOS_LAMBDA(std::size_t i) { sv_view(i) *= inv_norm; });
+    }
+
     /**
      * @brief Update data of the class
      *

diff --git a/pennylane_lightning/core/src/simulators/lightning_kokkos/bindings/LKokkosBindings.hpp b/pennylane_lightning/core/src/simulators/lightning_kokkos/bindings/LKokkosBindings.hpp
@@ -137,7 +137,11 @@ void registerBackendClassSpecificBindings(PyClass &pyclass) {
                 sv.applyOperation(str, wires, inv, std::vector<ParamT>{},
                                   conv_matrix);
             },
-            "Apply operation via the gate matrix");
+            "Apply operation via the gate matrix")
+        .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.");
 }
 
 /**

diff --git a/pennylane_lightning/core/src/simulators/lightning_kokkos/gates/GateFunctorsGenerator.hpp b/pennylane_lightning/core/src/simulators/lightning_kokkos/gates/GateFunctorsGenerator.hpp
@@ -972,4 +972,5 @@ template <class PrecisionT, bool adj = false> struct generatorMultiRZFunctor {
             1 - 2 * int(Kokkos::Impl::bit_count(k & wires_parity) % 2));
     }
 };
+
 } // namespace Pennylane::LightningKokkos::Functors
diff --git a/pennylane_lightning/core/src/simulators/lightning_kokkos/measurements/MeasurementsKokkos.hpp b/pennylane_lightning/core/src/simulators/lightning_kokkos/measurements/MeasurementsKokkos.hpp
@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 #pragma once
+#include <chrono>
 
 #include <Kokkos_Core.hpp>
 #include <Kokkos_Random.hpp>
@@ -682,7 +683,10 @@ class Measurements final
             });
 
         // Sampling using Random_XorShift64_Pool
-        Kokkos::Random_XorShift64_Pool<> rand_pool(5374857);
+        Kokkos::Random_XorShift64_Pool<> rand_pool(
+            std::chrono::high_resolution_clock::now()
+                .time_since_epoch()
+                .count());
 
         Kokkos::parallel_for(
             Kokkos::RangePolicy<KokkosExecSpace>(0, num_samples),

diff --git a/pennylane_lightning/core/src/simulators/lightning_kokkos/tests/Test_StateVectorLKokkos.cpp b/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 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));
+    }
+}
diff --git a/pennylane_lightning/core/src/simulators/lightning_qubit/tests/Test_StateVectorLQubit.cpp b/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));
+    }
+}
diff --git a/pennylane_lightning/core/src/utils/TestHelpers.hpp b/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
  */
-Original file line number
+Diff line change
@@ Expand Up @@
 - 2 * int(Kokkos::Impl::bit_count(k & wires_parity) % 2));
         }
     };
     } // namespace Pennylane::LightningKokkos::Functors