Update lightning qubit memory management

.github/CHANGELOG.md

@@ -9,6 +9,12 @@
 
 ### Improvements
 
+* Decouple LightningQubit memory ownership from numpy and migrate it to LightningQubit managed state-vector class.
+  [(#601)](https://github.com/PennyLaneAI/pennylane-lightning/pull/601)
+
+* Expand support for Projector observables on LightningKokkos.
+  [(#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)
 

.gitignore

@@ -6,9 +6,7 @@ doc/code/api/
 PennyLane_Lightning.egg-info/
 PennyLane_Lightning_Kokkos.egg-info/
 build/
-build_lightning_qubit/
-build_lightning_kokkos/
-build_lightning_gpu/
+build_lightning_*/
 Build/
 BuildCov/
 BuildGBench/

Makefile

@@ -59,9 +59,7 @@ help:
 clean:
 	find . -type d -name '__pycache__' -exec rm -r {} \+
 	rm -rf build Build BuildTests BuildTidy BuildGBench
-	rm -rf build_lightning_qubit
-	rm -rf build_lightning_kokkos
-	rm -rf build_lightning_gpu
+	rm -rf build_*
 	rm -rf .coverage coverage_html_report/
 	rm -rf pennylane_lightning/*_ops*
 

pennylane_lightning/core/_version.py

@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.35.0-dev8"
+__version__ = "0.35.0-dev10"

pennylane_lightning/core/src/simulators/lightning_qubit/StateVectorLQubitManaged.hpp

@@ -72,7 +72,7 @@ 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},
@@ -103,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),
@@ -145,7 +145,7 @@ class StateVectorLQubitManaged final
      *
      * @param index Index of the target element.
      */
-    void setBasisState(const size_t index) {
+    void setBasisState(const std::size_t index) {
         std::fill(data_.begin(), data_.end(), 0);
         data_[index] = {1, 0};
     }
@@ -158,7 +158,7 @@ class StateVectorLQubitManaged final
      */
     void setStateVector(const std::vector<std::size_t> &indices,
                         const std::vector<ComplexT> &values) {
-        for (size_t n = 0; n < indices.size(); n++) {
+        for (std::size_t n = 0; n < indices.size(); n++) {
             data_[indices[n]] = values[n];
         }
     }
@@ -198,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());
     }

pennylane_lightning/core/src/simulators/lightning_qubit/bindings/LQubitBindings.hpp

@@ -207,7 +207,7 @@ void registerBackendClassSpecificBindings(PyClass &pyclass) {
                               data_ptr);
                 }
             },
-            "Synchronize data from the GPU device to host.")
+            "Copy StateVector data into a Numpy array.")
         .def(
             "UpdateData",
             [](StateVectorT &device_sv, const np_arr_c &state) {
@@ -219,7 +219,7 @@ void registerBackendClassSpecificBindings(PyClass &pyclass) {
                     device_sv.updateData(data_ptr, length);
                 }
             },
-            "Synchronize data from the host device to GPU.")
+            "Copy StateVector data into a Numpy array.")
         .def("applyControlledMatrix", &applyControlledMatrix<StateVectorT>,
              "Apply controlled operation")
         .def("kernel_map", &svKernelMap<StateVectorT>,

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

@@ -124,11 +124,9 @@ TEMPLATE_TEST_CASE("StateVectorLQubitManaged::setBasisState",
 
         TestVectorT expected_state(size_t{1U} << num_qubits, 0.0,
                                    getBestAllocator<ComplexT>());
-        expected_state[3] = {1.0, 0.0};
-
+        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);
-
-        size_t index = 3;
         sv.setBasisState(index);
 
         REQUIRE(sv.getDataVector() == approx(expected_state));

pennylane_lightning/lightning_kokkos/lightning_kokkos.py

@@ -150,6 +150,7 @@
         "Hadamard",
         "Hermitian",
         "Identity",
+        "Projector",
         "SparseHamiltonian",
         "Hamiltonian",
         "Sum",
@@ -196,7 +197,7 @@
             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 @@
             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
@@ -466,8 +463,12 @@
                 Expectation value of the observable
             """
             if observable.name in [
+                "Identity",
                 "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:
@@ -526,8 +527,12 @@
                 Variance of the observable
             """
             if observable.name in [
+                "Identity",
                 "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:

pennylane_lightning/lightning_qubit/lightning_qubit.py

@@ -78,7 +78,7 @@
     def _state_dtype(dtype):
         if dtype not in [np.complex128, np.complex64]:  # pragma: no cover
             raise ValueError(f"Data type is not supported for state-vector computation: {dtype}")
         return StateVectorC128 if dtype == np.complex128 else StateVectorC64
 
     allowed_operations = {
         "Identity",
@@ -231,9 +231,7 @@
 
             # Create the initial state. Internally, we store the
             # state as an array of dimension [2]*wires.
-            self._state = _state_dtype(c_dtype)(self.num_wires)
-            self._pre_rotated_state = _state_dtype(c_dtype)(self.num_wires)
-
+            self._qubit_state = _state_dtype(c_dtype)(self.num_wires)
             self._batch_obs = batch_obs
             self._mcmc = mcmc
             if self._mcmc:
@@ -280,14 +278,14 @@
             Args:
                 index (int): integer representing the computational basis state.
             """
-            self._state.setBasisState(index)
+            self._qubit_state.setBasisState(index)
 
         def reset(self):
             """Reset the device"""
             super().reset()
 
             # init the state vector to |00..0>
-            self._state.resetStateVector()
+            self._qubit_state.resetStateVector()
 
         @property
         def create_ops_list(self):
@@ -314,15 +312,15 @@
             >>> print(dev.state)
             [0.+0.j 1.+0.j]
             """
             state = np.zeros(2**self.num_wires, dtype=self.C_DTYPE)
             state = self._asarray(state, dtype=self.C_DTYPE)
-            self._state.getState(state.ravel(order="C"))
+            self._qubit_state.getState(state.ravel(order="C"))
             return state
 
         @property
         def state_vector(self):
             """Returns a handle to the statevector."""
-            return self._state
+            return self._qubit_state
 
         def _apply_state_vector(self, state, device_wires):
             """Initialize the internal state vector in a specified state.
@@ -332,11 +330,11 @@
                 device_wires (Wires): wires that get initialized in the state
             """
 
-            if isinstance(state, self._state.__class__):
+            if isinstance(state, self._qubit_state.__class__):
                 state_data = np.zeros(state.size, dtype=self.C_DTYPE)
                 state_data = self._asarray(state_data, dtype=self.C_DTYPE)
-                self._state.getState(state_data.ravel(order="C"))
+                self._qubit_state.getState(state_data.ravel(order="C"))
                 state = state_data
 
             ravelled_indices, state = self._preprocess_state_vector(state, device_wires)
 
@@ -346,11 +344,11 @@
 
             if len(device_wires) == self.num_wires and Wires(sorted(device_wires)) == device_wires:
                 # Initialize the entire device state with the input state
                 state = self._reshape(state, output_shape).ravel(order="C")
-                self._state.UpdateData(state)
+                self._qubit_state.UpdateData(state)
                 return
 
-            self._state.setStateVector(ravelled_indices, state)  # this operation on device
+            self._qubit_state.setStateVector(ravelled_indices, state)  # this operation on device
 
         def _apply_basis_state(self, state, wires):
             """Initialize the state vector in a specified computational basis state.
@@ -364,7 +362,7 @@
             Note: This function does not support broadcasted inputs yet.
             """
             num = self._get_basis_state_index(state, wires)
             self._create_basis_state(num)
 
         def _apply_lightning_controlled(self, operation):
             """Apply an arbitrary controlled operation to the state tensor.
@@ -375,12 +373,12 @@
             Returns:
                 array[complex]: the output state tensor
             """
             state = self.state_vector
 
             basename = "PauliX" if operation.name == "MultiControlledX" else operation.base.name
             if basename == "Identity":
                 return
             method = getattr(state, f"{basename}", None)
             control_wires = self.wires.indices(operation.control_wires)
             control_values = (
                 [bool(int(i)) for i in operation.hyperparameters["control_values"]]
@@ -396,7 +394,7 @@
                 param = operation.parameters
                 method(control_wires, control_values, target_wires, inv, param)
             else:  # apply gate as an n-controlled matrix
                 method = getattr(state, "applyControlledMatrix")
                 target_wires = self.wires.indices(operation.target_wires)
                 try:
                     method(
@@ -421,15 +419,15 @@
             Returns:
                 array[complex]: the output state tensor
             """
             state = self.state_vector
 
             # Skip over identity operations instead of performing
             # matrix multiplication with it.
             for operation in operations:
                 name = operation.name
                 if name == "Identity":
                     continue
                 method = getattr(state, name, None)
                 wires = self.wires.indices(operation.wires)
 
                 if method is not None:  # apply specialized gate
@@ -441,11 +439,11 @@
                     or name == "ControlledQubitUnitary"
                     or name == "MultiControlledX"
                 ):  # apply n-controlled gate
                     self._apply_lightning_controlled(operation)
                 else:  # apply gate as a matrix
                     # Inverse can be set to False since qml.matrix(operation) is already in
                     # inverted form
                     method = getattr(state, "applyMatrix")
                     try:
                         method(qml.matrix(operation), wires, False)
                     except AttributeError:  # pragma: no cover
@@ -473,7 +471,7 @@
                         f"Operations have already been applied on a {self.short_name} device."
                     )
 
             self.apply_lightning(operations)
 
         # pylint: disable=protected-access
         def expval(self, observable, shot_range=None, bin_size=None):
@@ -494,6 +492,9 @@
                 "Identity",
                 "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:
@@ -548,6 +549,9 @@
                 "Identity",
                 "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:
@@ -624,10 +628,10 @@
         # pylint: disable=attribute-defined-outside-init
         def sample(self, observable, shot_range=None, bin_size=None, counts=False):
             """Return samples of an observable."""
             if observable.name != "PauliZ":
                 self.apply_lightning(observable.diagonalizing_gates())
                 self._samples = self.generate_samples()
             return super().sample(
                 observable, shot_range=shot_range, bin_size=bin_size, counts=counts
             )
 
@@ -694,11 +698,11 @@
                         "The number of qubits of starting_state must be the same as "
                         "that of the device."
                     )
                 self._apply_state_vector(starting_state, self.wires)
             elif not use_device_state:
                 self.reset()
                 self.apply(tape.operations)
             return self.state_vector
 
         def adjoint_jacobian(self, tape, starting_state=None, use_device_state=False):
             """Computes and returns the Jacobian with the adjoint method."""

tests/test_adjoint_jacobian.py

@@ -583,8 +583,8 @@ def test_provide_starting_state(self, tol, dev):
 
         dM1 = dev.adjoint_jacobian(tape)
 
-        if device_name == "lightning.kokkos":
-            dev._pre_rotated_state = dev.state_vector  # necessary for lightning.kokkos
+        if device_name in ["lightning.kokkos", "lightning.qubit"]:
+            dev._pre_rotated_state = dev.state_vector
 
             qml.execute([tape], dev, None)
             dM2 = dev.adjoint_jacobian(tape, starting_state=dev._pre_rotated_state)