Update lightning qubit memory management (#601)

* update dev version * update clean * expand .gitignore * expand lightning qubit class-specific Python bindings * add some statevector manipulation methods to the statevector managed simulator * update the lightning qubit Python class * add c++ unit tests * add comments to test * update lightning_qubit.py * remove obsolete tests * rename statevector instance * return some important tests * implement PR review suggestions * add review suggestion * Add semi-colon. * Fix Projector obs in L-Qubit and add Proj support in L-Kokkos. * Implement previous commit fix for None shots only. * Add tests for Proj expval/var * Auto update version * Trigger CI * remove comment * add projector support to LGPU * Format * revert changes for LGPU * skip tests for Projector observable not supported * expand tests for lightning.qubit * update changelog * Auto update version * Trigger CI * add some review suggestions * remove identities * update LKokkos and LQubit _apply_state_vector * Update tests/test_apply.py --------- Co-authored-by: Vincent Michaud-Rioux <[email protected]>
PennyLaneAI · Jan 31, 2024 · cfe2455 · cfe2455
1 parent ebc4ed6
commit cfe2455
Show file tree

Hide file tree

Showing 13 changed files with 311 additions and 398 deletions.
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
@@ -9,6 +9,12 @@
 
 ### Improvements
 
+* Decouple LightningQubit memory ownership from numpy and migrate it to Lightning-Qubit managed state-vector class.
+  [(#601)](https://github.com/PennyLaneAI/pennylane-lightning/pull/601)
+
+* Expand support for Projector observables on Lightning-Kokkos.
+  [(#601)](https://github.com/PennyLaneAI/pennylane-lightning/pull/601)
+
 * Split Docker build cron job into two jobs: master and latest. This is mainly for reporting in the `plugin-test-matrix` repo.
   [(#600)](https://github.com/PennyLaneAI/pennylane-lightning/pull/600)
 

diff --git a/.gitignore b/.gitignore
@@ -6,6 +6,7 @@ doc/code/api/
 PennyLane_Lightning.egg-info/
 PennyLane_Lightning_Kokkos.egg-info/
 build/
+build_lightning_*/
 Build/
 BuildCov/
 BuildGBench/

diff --git a/Makefile b/Makefile
@@ -59,6 +59,7 @@ help:
 clean:
 	find . -type d -name '__pycache__' -exec rm -r {} \+
 	rm -rf build Build BuildTests BuildTidy BuildGBench
+	rm -rf build_*
 	rm -rf .coverage coverage_html_report/
 	rm -rf pennylane_lightning/*_ops*
 

diff --git a/pennylane_lightning/core/_version.py b/pennylane_lightning/core/_version.py
@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.35.0-dev9"
+__version__ = "0.35.0-dev10"
diff --git a/pennylane_lightning/core/src/simulators/lightning_qubit/StateVectorLQubitManaged.hpp b/pennylane_lightning/core/src/simulators/lightning_qubit/StateVectorLQubitManaged.hpp
@@ -19,6 +19,7 @@
 
 #pragma once
 
+#include <algorithm> // fill
 #include <complex>
 #include <vector>
 
@@ -71,12 +72,12 @@ class StateVectorLQubitManaged final
      * @param memory_model Memory model the statevector will use
      */
     explicit StateVectorLQubitManaged(
-        size_t num_qubits, Threading threading = Threading::SingleThread,
+        std::size_t num_qubits, Threading threading = Threading::SingleThread,
         CPUMemoryModel memory_model = bestCPUMemoryModel())
         : BaseType{num_qubits, threading, memory_model},
           data_{exp2(num_qubits), ComplexT{0.0, 0.0},
                 getAllocator<ComplexT>(this->memory_model_)} {
-        data_[0] = {1, 0};
+        setBasisState(0U);
     }
 
     /**
@@ -102,7 +103,7 @@ class StateVectorLQubitManaged final
      * @param threading Threading option the statevector to use
      * @param memory_model Memory model the statevector will use
      */
-    StateVectorLQubitManaged(const ComplexT *other_data, size_t other_size,
+    StateVectorLQubitManaged(const ComplexT *other_data, std::size_t other_size,
                              Threading threading = Threading::SingleThread,
                              CPUMemoryModel memory_model = bestCPUMemoryModel())
         : BaseType(log2PerfectPower(other_size), threading, memory_model),
@@ -139,6 +140,39 @@ class StateVectorLQubitManaged final
 
     ~StateVectorLQubitManaged() = default;
 
+    /**
+     * @brief Prepares a single computational basis state.
+     *
+     * @param index Index of the target element.
+     */
+    void setBasisState(const std::size_t index) {
+        std::fill(data_.begin(), data_.end(), 0);
+        data_[index] = {1, 0};
+    }
+
+    /**
+     * @brief Set values for a batch of elements of the state-vector.
+     *
+     * @param values Values to be set for the target elements.
+     * @param indices Indices of the target elements.
+     */
+    void setStateVector(const std::vector<std::size_t> &indices,
+                        const std::vector<ComplexT> &values) {
+        for (std::size_t n = 0; n < indices.size(); n++) {
+            data_[indices[n]] = values[n];
+        }
+    }
+
+    /**
+     * @brief Reset the data back to the \f$\ket{0}\f$ state.
+     *
+     */
+    void resetStateVector() {
+        if (this->getLength() > 0) {
+            setBasisState(0U);
+        }
+    }
+
     [[nodiscard]] auto getData() -> ComplexT * { return data_.data(); }
 
     [[nodiscard]] auto getData() const -> const ComplexT * {
@@ -164,7 +198,7 @@ class StateVectorLQubitManaged final
      * @param new_data data pointer to new data.
      * @param new_size size of underlying data storage.
      */
-    void updateData(const ComplexT *new_data, size_t new_size) {
+    void updateData(const ComplexT *new_data, std::size_t new_size) {
         PL_ASSERT(data_.size() == new_size);
         std::copy(new_data, new_data + new_size, data_.data());
     }

diff --git a/pennylane_lightning/core/src/simulators/lightning_qubit/bindings/LQubitBindings.hpp b/pennylane_lightning/core/src/simulators/lightning_qubit/bindings/LQubitBindings.hpp
@@ -26,7 +26,7 @@
 #include "GateOperation.hpp"
 #include "MeasurementsLQubit.hpp"
 #include "ObservablesLQubit.hpp"
-#include "StateVectorLQubitRaw.hpp"
+#include "StateVectorLQubitManaged.hpp"
 #include "TypeList.hpp"
 #include "VectorJacobianProduct.hpp"
 
@@ -36,16 +36,16 @@ using namespace Pennylane::Bindings;
 using namespace Pennylane::LightningQubit::Algorithms;
 using namespace Pennylane::LightningQubit::Measures;
 using namespace Pennylane::LightningQubit::Observables;
-using Pennylane::LightningQubit::StateVectorLQubitRaw;
+using Pennylane::LightningQubit::StateVectorLQubitManaged;
 } // namespace
 /// @endcond
 
 namespace py = pybind11;
 
 namespace Pennylane::LightningQubit {
 using StateVectorBackends =
-    Pennylane::Util::TypeList<StateVectorLQubitRaw<float>,
-                              StateVectorLQubitRaw<double>, void>;
+    Pennylane::Util::TypeList<StateVectorLQubitManaged<float>,
+                              StateVectorLQubitManaged<double>, void>;
 
 /**
  * @brief Get a gate kernel map for a statevector.
@@ -162,12 +162,68 @@ void registerControlledGate(PyClass &pyclass) {
  */
 template <class StateVectorT, class PyClass>
 void registerBackendClassSpecificBindings(PyClass &pyclass) {
+    using PrecisionT =
+        typename StateVectorT::PrecisionT; // Statevector's precision
+    using ComplexT = typename StateVectorT::ComplexT;
+    using ParamT = PrecisionT; // Parameter's data precision
+    using np_arr_c = py::array_t<std::complex<ParamT>,
+                                 py::array::c_style | py::array::forcecast>;
+
     registerGatesForStateVector<StateVectorT>(pyclass);
     registerControlledGate<StateVectorT>(pyclass);
-    pyclass.def("applyControlledMatrix", &applyControlledMatrix<StateVectorT>,
-                "Apply controlled operation");
-    pyclass.def("kernel_map", &svKernelMap<StateVectorT>,
-                "Get internal kernels for operations");
+
+    pyclass
+        .def(py::init([](std::size_t num_qubits) {
+            return new StateVectorT(num_qubits);
+        }))
+        .def("resetStateVector", &StateVectorT::resetStateVector)
+        .def(
+            "setBasisState",
+            [](StateVectorT &sv, const size_t index) {
+                sv.setBasisState(index);
+            },
+            "Create Basis State.")
+        .def(
+            "setStateVector",
+            [](StateVectorT &sv, const std::vector<std::size_t> &indices,
+               const np_arr_c &state) {
+                const auto buffer = state.request();
+                std::vector<ComplexT> state_in;
+                if (buffer.size) {
+                    const auto ptr = static_cast<const ComplexT *>(buffer.ptr);
+                    state_in = std::vector<ComplexT>{ptr, ptr + buffer.size};
+                }
+                sv.setStateVector(indices, state_in);
+            },
+            "Set State Vector with values and their corresponding indices")
+        .def(
+            "getState",
+            [](const StateVectorT &sv, np_arr_c &state) {
+                py::buffer_info numpyArrayInfo = state.request();
+                auto *data_ptr =
+                    static_cast<std::complex<PrecisionT> *>(numpyArrayInfo.ptr);
+                if (state.size()) {
+                    std::copy(sv.getData(), sv.getData() + sv.getLength(),
+                              data_ptr);
+                }
+            },
+            "Copy StateVector data into a Numpy array.")
+        .def(
+            "UpdateData",
+            [](StateVectorT &device_sv, const np_arr_c &state) {
+                const py::buffer_info numpyArrayInfo = state.request();
+                auto *data_ptr = static_cast<ComplexT *>(numpyArrayInfo.ptr);
+                const auto length =
+                    static_cast<size_t>(numpyArrayInfo.shape[0]);
+                if (length) {
+                    device_sv.updateData(data_ptr, length);
+                }
+            },
+            "Copy StateVector data into a Numpy array.")
+        .def("applyControlledMatrix", &applyControlledMatrix<StateVectorT>,
+             "Apply controlled operation")
+        .def("kernel_map", &svKernelMap<StateVectorT>,
+             "Get internal kernels for operations");
 }
 
 /**

diff --git a/...ane_lightning/core/src/simulators/lightning_qubit/tests/Test_StateVectorLQubitManaged.cpp b/...ane_lightning/core/src/simulators/lightning_qubit/tests/Test_StateVectorLQubitManaged.cpp
@@ -21,7 +21,8 @@
 
 #include <catch2/catch.hpp>
 
-#include "LinearAlgebra.hpp" //randomUnitary
+#include "CPUMemoryModel.hpp" // getBestAllocator
+#include "LinearAlgebra.hpp"  //randomUnitary
 #include "StateVectorLQubitManaged.hpp"
 #include "StateVectorLQubitRaw.hpp"
 #include "TestHelpers.hpp" // createRandomStateVectorData, TestVector
@@ -107,4 +108,58 @@ TEMPLATE_TEST_CASE("StateVectorLQubitManaged::StateVectorLQubitManaged",
 
         REQUIRE(sv.getDataVector() == approx(st_data));
     }
-}
+}
+
+TEMPLATE_TEST_CASE("StateVectorLQubitManaged::setBasisState",
+                   "[StateVectorLQubitManaged]", float, double) {
+    using PrecisionT = TestType;
+    using ComplexT = std::complex<PrecisionT>;
+    using TestVectorT = TestVector<ComplexT>;
+
+    const std::size_t num_qubits = 3;
+
+    SECTION("Prepares a single computational basis state.") {
+        TestVectorT init_state =
+            createRandomStateVectorData<PrecisionT>(re, num_qubits);
+
+        TestVectorT expected_state(size_t{1U} << num_qubits, 0.0,
+                                   getBestAllocator<ComplexT>());
+        std::size_t index = GENERATE(0, 1, 2, 3, 4, 5, 6, 7);
+        expected_state[index] = {1.0, 0.0};
+        StateVectorLQubitManaged<PrecisionT> sv(init_state);
+        sv.setBasisState(index);
+
+        REQUIRE(sv.getDataVector() == approx(expected_state));
+    }
+}
+
+TEMPLATE_TEST_CASE("StateVectorLQubitManaged::SetStateVector",
+                   "[StateVectorLQubitManaged]", float, double) {
+    using PrecisionT = TestType;
+    using ComplexT = std::complex<PrecisionT>;
+    using TestVectorT = TestVector<ComplexT>;
+
+    const std::size_t num_qubits = 3;
+
+    SECTION("Set state vector with values and indices") {
+        TestVectorT init_state =
+            createRandomStateVectorData<PrecisionT>(re, num_qubits);
+
+        auto expected_state = init_state;
+
+        for (size_t i = 0; i < Pennylane::Util::exp2(num_qubits - 1); i++) {
+            std::swap(expected_state[i * 2], expected_state[i * 2 + 1]);
+        }
+
+        std::vector<std::size_t> indices = {0, 2, 4, 6, 1, 3, 5, 7};
+
+        std::vector<ComplexT> values = {
+            init_state[1], init_state[3], init_state[5], init_state[7],
+            init_state[0], init_state[2], init_state[4], init_state[6]};
+
+        StateVectorLQubitManaged<PrecisionT> sv{num_qubits};
+        sv.setStateVector(indices, values);
+
+        REQUIRE(sv.getDataVector() == approx(expected_state));
+    }
+}
diff --git a/pennylane_lightning/lightning_kokkos/lightning_kokkos.py b/pennylane_lightning/lightning_kokkos/lightning_kokkos.py
@@ -150,6 +150,7 @@ def _kokkos_configuration():
         "Hadamard",
         "Hermitian",
         "Identity",
+        "Projector",
         "SparseHamiltonian",
         "Hamiltonian",
         "Sum",
@@ -196,7 +197,7 @@ def __init__(
             shots=None,
             batch_obs=False,
             kokkos_args=None,
-        ):  # pylint: disable=unused-argument
+        ):  # pylint: disable=unused-argument, too-many-arguments
             super().__init__(wires, shots=shots, c_dtype=c_dtype)
 
             if kokkos_args is None:
@@ -213,10 +214,6 @@ def __init__(
             if not LightningKokkos.kokkos_config:
                 LightningKokkos.kokkos_config = _kokkos_configuration()
 
-            # Create the initial state. Internally, we store the
-            # state as an array of dimension [2]*wires.
-            self._pre_rotated_state = _kokkos_dtype(c_dtype)(self.num_wires)
-
         @staticmethod
         def _asarray(arr, dtype=None):
             arr = np.asarray(arr)  # arr is not copied
@@ -347,9 +344,8 @@ def _apply_state_vector(self, state, device_wires):
             """
 
             if isinstance(state, self._kokkos_state.__class__):
-                state_data = np.zeros(state.size, dtype=self.C_DTYPE)
-                state_data = self._asarray(state_data, dtype=self.C_DTYPE)
-                state.DeviceToHost(state_data.ravel(order="C"))
+                state_data = allocate_aligned_array(state.size, np.dtype(self.C_DTYPE), True)
+                state.DeviceToHost(state_data)
                 state = state_data
 
             ravelled_indices, state = self._preprocess_state_vector(state, device_wires)
@@ -468,6 +464,9 @@ def expval(self, observable, shot_range=None, bin_size=None):
             if observable.name in [
                 "Projector",
             ]:
+                if self.shots is None:
+                    qs = qml.tape.QuantumScript([], [qml.expval(observable)])
+                    self.apply(self._get_diagonalizing_gates(qs))
                 return super().expval(observable, shot_range=shot_range, bin_size=bin_size)
 
             if self.shots is not None:
@@ -528,6 +527,9 @@ def var(self, observable, shot_range=None, bin_size=None):
             if observable.name in [
                 "Projector",
             ]:
+                if self.shots is None:
+                    qs = qml.tape.QuantumScript([], [qml.var(observable)])
+                    self.apply(self._get_diagonalizing_gates(qs))
                 return super().var(observable, shot_range=shot_range, bin_size=bin_size)
 
             if self.shots is not None: