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>

Author: vincentmr

.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)
+
 * `lightning.qubit` supports mid-circuit measurements.
   [(#650)](https://github.com/PennyLaneAI/pennylane-lightning/pull/650)
 

.github/workflows/tests_gpu_cuda.yml

@@ -316,7 +316,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

.github/workflows/tests_gpu_kokkos.yml

@@ -325,7 +325,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.kokkos python -m pytest tests/ $COVERAGE_FLAGS

.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 }}
@@ -393,8 +394,8 @@ jobs:
         run: |
           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 }}
@@ -592,7 +593,8 @@ jobs:
         run: |
           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 }}
@@ -611,11 +613,11 @@ jobs:
         if: ${{ matrix.pl_backend == 'all' }}
         run: |
           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 }}

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

@@ -767,48 +767,42 @@ class StateVectorKokkos final
      * @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});
+        KokkosVector matrix("gate_matrix", 4);
         Kokkos::parallel_for(
-            policy_2d,
-            collapseFunctor<fp_t>(*data_, num_qubits, stride, negbranch));
-
+            matrix.size(), KOKKOS_LAMBDA(const 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() {
-        KokkosVector sv_view =
-            getView(); // circumvent error capturing this with KOKKOS_LAMBDA
+        auto sv_view = getView();
 
         // 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));
+            KOKKOS_LAMBDA(const 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");
 
-        std::complex<PrecisionT> inv_norm = 1. / std::sqrt(squaredNorm);
+        const std::complex<PrecisionT> inv_norm =
+            1. / Kokkos::sqrt(squaredNorm);
         Kokkos::parallel_for(
             sv_view.size(),
-            KOKKOS_LAMBDA(const size_t i) { sv_view(i) *= inv_norm; });
+            KOKKOS_LAMBDA(const std::size_t i) { sv_view(i) *= inv_norm; });
     }
 
     /**

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,
+             "Normalizes the statevector to norm 1.");
 }
 
 /**

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

@@ -973,34 +973,4 @@ template <class PrecisionT, bool adj = false> struct generatorMultiRZFunctor {
     }
 };
 
-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/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),

pennylane_lightning/lightning_kokkos/lightning_kokkos.py

@@ -21,6 +21,8 @@
 from warnings import warn
 
 import numpy as np
+from pennylane.measurements import MidMeasureMP
+from pennylane.ops import Conditional
 
 from pennylane_lightning.core.lightning_base import (
     LightningBase,
@@ -136,6 +138,8 @@ def _kokkos_configuration():
         "QFT",
         "ECR",
         "BlockEncode",
+        "MidMeasureMP",
+        "Conditional",
     }
 
     allowed_observables = {
@@ -210,6 +214,15 @@ def __init__(
             if not LightningKokkos.kokkos_config:
                 LightningKokkos.kokkos_config = _kokkos_configuration()
 
+        # pylint: disable=missing-function-docstring
+        @classmethod
+        def capabilities(cls):
+            capabilities = super().capabilities().copy()
+            capabilities.update(
+                supports_mid_measure=True,
+            )
+            return capabilities
+
         @staticmethod
         def _asarray(arr, dtype=None):
             arr = np.asarray(arr)  # arr is not copied
@@ -370,7 +383,25 @@ def _apply_basis_state(self, state, wires):
             num = self._get_basis_state_index(state, wires)
             self._create_basis_state(num)
 
-        def apply_lightning(self, operations):
+        def _apply_lightning_midmeasure(self, operation: MidMeasureMP, mid_measurements: dict):
+            """Execute a MidMeasureMP operation and return the sample in mid_measurements.
+            Args:
+                operation (~pennylane.operation.Operation): mid-circuit measurement
+            Returns:
+                None
+            """
+            wires = self.wires.indices(operation.wires)
+            wire = list(wires)[0]
+            sample = qml.math.reshape(self.generate_samples(shots=1), (-1,))[wire]
+            if operation.postselect is not None and sample != operation.postselect:
+                mid_measurements[operation] = -1
+                return
+            mid_measurements[operation] = sample
+            getattr(self.state_vector, "collapse")(wire, bool(sample))
+            if operation.reset and bool(sample):
+                self.apply([qml.PauliX(operation.wires)], mid_measurements=mid_measurements)
+
+        def apply_lightning(self, operations, mid_measurements=None):
             """Apply a list of operations to the state tensor.
 
             Args:
@@ -392,12 +423,17 @@ def apply_lightning(self, operations):
                 else:
                     name = ops.name
                     invert_param = False
-                if name == "Identity":
+                if isinstance(ops, qml.Identity):
                     continue
                 method = getattr(state, name, None)
                 wires = self.wires.indices(ops.wires)
 
-                if ops.name == "C(GlobalPhase)":
+                if isinstance(ops, Conditional):
+                    if ops.meas_val.concretize(mid_measurements):
+                        self.apply_lightning([ops.then_op])
+                elif isinstance(ops, MidMeasureMP):
+                    self._apply_lightning_midmeasure(ops, mid_measurements)
+                elif ops.name == "C(GlobalPhase)":
                     controls = ops.control_wires
                     control_values = ops.control_values
                     param = ops.base.parameters[0]
@@ -425,7 +461,7 @@ def apply_lightning(self, operations):
                     method(wires, invert_param, param)
 
         # pylint: disable=unused-argument
-        def apply(self, operations, rotations=None, **kwargs):
+        def apply(self, operations, rotations=None, mid_measurements=None, **kwargs):
             """Applies a list of operations to the state tensor."""
             # State preparation is currently done in Python
             if operations:  # make sure operations[0] exists
@@ -445,7 +481,9 @@ def apply(self, operations, rotations=None, **kwargs):
                         + f"Operations have already been applied on a {self.short_name} device."
                     )
 
-            self.apply_lightning(operations)
+            self.apply_lightning(operations, mid_measurements=mid_measurements)
+            if mid_measurements is not None and any(v == -1 for v in mid_measurements.values()):
+                self._apply_basis_state(np.zeros(self.num_wires), wires=self.wires)
 
         # pylint: disable=protected-access
         def expval(self, observable, shot_range=None, bin_size=None):
@@ -575,19 +613,20 @@ def var(self, observable, shot_range=None, bin_size=None):
 
             return measure.var(observable.name, observable_wires)
 
-        def generate_samples(self):
+        def generate_samples(self, shots=None):
             """Generate samples
 
             Returns:
                 array[int]: array of samples in binary representation with shape
                 ``(dev.shots, dev.num_wires)``
             """
+            shots = self.shots if shots is None else shots
             measure = (
                 MeasurementsC64(self._kokkos_state)
                 if self.use_csingle
                 else MeasurementsC128(self._kokkos_state)
             )
-            return measure.generate_samples(len(self.wires), self.shots).astype(int, copy=False)
+            return measure.generate_samples(len(self.wires), shots).astype(int, copy=False)
 
         def probability_lightning(self, wires):
             """Return the probability of each computational basis state.

tests/test_native_mcm.py

@@ -22,8 +22,8 @@
 from flaky import flaky
 from pennylane._device import DeviceError
 
-if not LightningDevice._new_API:
-    pytest.skip("Exclusive tests for new device API. Skipping.", allow_module_level=True)
+if device_name not in ("lightning.qubit", "lightning.kokkos"):
+    pytest.skip("Native MCM not supported. Skipping.", allow_module_level=True)
 
 if not LightningDevice._CPP_BINARY_AVAILABLE:  # pylint: disable=protected-access
     pytest.skip("No binary module found. Skipping.", allow_module_level=True)
@@ -79,11 +79,18 @@ def func(x, y):
         qml.cond(m0, qml.RY)(y, wires=1)
         return qml.classical_shadow(wires=0)
 
-    with pytest.raises(
-        DeviceError,
-        match=f"not accepted with finite shots on lightning.qubit",
-    ):
-        func(*params)
+    if device_name == "lightning.qubit":
+        with pytest.raises(
+            DeviceError,
+            match=f"not accepted with finite shots on lightning.qubit",
+        ):
+            func(*params)
+    else:
+        with pytest.raises(
+            TypeError,
+            match=f"Native mid-circuit measurement mode does not support ClassicalShadowMP measurements.",
+        ):
+            func(*params)
 
 
 @flaky(max_runs=5)