Add sparseH for LGPU (#526)

* Init commit * Fix std::endl; * Use more generic indices in base std::size_t. * add pybind layer * add python layer * Quick and dirty spham bindings. * Add sparse_ham serialization. * Add sparse_ham tests in tests/test_adjoint_jacobian.py' * Bug fix sparse product. * add sparseH * Trigger CI * Fix python bindings LGPU idxT * Fix serial tests and update changelog. * add more unit tests for sparseH base class * Fix tidy & sparse adjoint test device name. * Fix tidy warning for sparse_ham. * Send backend-specific ops in respective modules. * Fix sparse_hamiltonianmpi_c and add getWires test. * Add sparseH diff capability in LQ. * Add sparse Hamiltonian support for Lightning-Kokkos (#527) * Use more generic indices in base std::size_t. * Quick and dirty spham bindings. * Add sparse_ham serialization. * Add sparse_ham tests in tests/test_adjoint_jacobian.py' * Bug fix sparse product. * Fix python bindings LGPU idxT * Fix serial tests and update changelog. * Fix tidy & sparse adjoint test device name. * Fix tidy warning for sparse_ham. * Send backend-specific ops in respective modules. * Fix sparse_hamiltonianmpi_c and add getWires test. * Fix clang tidy * Comment workflows but tidy. * Fix tidy warn * Add override to sp::getWires * Restore triggers * Update tests_linux_x86_mpi.yml * Add constructibility tests. * Move L-Kokkos-CUDA tests to workflow call, called from tests_gpu_cu11.yml. * Remove GPU deadlock. * Bug fix Python MPI. * Upload both outputs. * Update gcc version in format.yml. * Update .github/CHANGELOG.md [skip ci] Co-authored-by: Amintor Dusko <[email protected]> * Update .github/workflows/tests_gpu_kokkos.yml [skip ci] Co-authored-by: Amintor Dusko <[email protected]> * rename argn [skip ci] * Remove unused lines [skip ci] * Fix SparseHamiltonianBase::isEqual. [skip ci] * Trigger CI * Auto update version * Trigger CI * resolve comments * rename dev_kokkos to dev * Fix tidy. --------- Co-authored-by: Vincent Michaud-Rioux <[email protected]> Co-authored-by: Vincent Michaud-Rioux <[email protected]> Co-authored-by: Amintor Dusko <[email protected]> Co-authored-by: Dev version update bot <github-actions[bot]@users.noreply.github.com>
PennyLaneAI · Oct 24, 2023 · b1b698d · b1b698d
1 parent ccea774
commit b1b698d
Show file tree

Hide file tree

Showing 39 changed files with 1,245 additions and 151 deletions.
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
@@ -2,6 +2,12 @@
 
 ### New features since last release
 
+* Add `SparseHamiltonian` support for Lightning-GPU.
+  [(#526)] (https://github.com/PennyLaneAI/pennylane-lightning/pull/526)
+
+* Add `SparseHamiltonian` support for Lightning-Kokkos.
+  [(#527)] (https://github.com/PennyLaneAI/pennylane-lightning/pull/527)
+
 * Integrate python/pybind layer of distributed Lightning-GPU into the Lightning monorepo with python unit tests.
   [(#518)] (https://github.com/PennyLaneAI/pennylane-lightning/pull/518)
 

diff --git a/.github/workflows/format.yml b/.github/workflows/format.yml
@@ -84,7 +84,7 @@ jobs:
           cp -rf ${{ github.workspace}}/Kokkos_install/${{ matrix.exec_model }}/* Kokkos/
 
       - name: Install dependencies
-        run: sudo apt update && sudo apt -y install clang-tidy-14 cmake g++-10 ninja-build libomp-14-dev
+        run: sudo apt update && sudo apt -y install clang-tidy-14 cmake gcc-11 g++-11 ninja-build libomp-14-dev
         env:
           DEBIAN_FRONTEND: noninteractive
 
@@ -96,5 +96,6 @@ jobs:
             -DBUILD_TESTS=ON \
             -DENABLE_WARNINGS=ON \
             -DPL_BACKEND=${{ matrix.pl_backend }} \
-            -DCMAKE_CXX_COMPILER="$(which g++-10)"
+            -DCMAKE_CXX_COMPILER="$(which g++-11)" \
+            -DCMAKE_C_COMPILER="$(which gcc-11)"
           cmake --build ./Build
diff --git a/.github/workflows/tests_gpu.yml → .github/workflows/tests_gpu_kokkos.yml b/.github/workflows/tests_gpu.yml → .github/workflows/tests_gpu_kokkos.yml
@@ -1,4 +1,4 @@
-name: Testing (GPU)
+name: Testing::LKokkos::GPU
 on:
   pull_request:
   push:

diff --git a/.github/workflows/tests_linux_x86_mpi_gpu.yml b/.github/workflows/tests_linux_x86_mpi_gpu.yml
@@ -246,15 +246,16 @@ jobs:
         run: |
           source /etc/profile.d/modules.sh && module use /opt/modules/ && module load ${{ matrix.mpilib }}
           PL_DEVICE=lightning.gpu python -m pytest ./tests/ $COVERAGE_FLAGS
+          mv coverage.xml coverage-${{ github.job }}-lightning_gpu_${{ matrix.mpilib }}-0.xml
           PL_DEVICE=lightning.gpu /opt/mpi/${{ matrix.mpilib }}/bin/mpirun -np 2 python -m pytest ./mpitests $COVERAGE_FLAGS
+          mv coverage.xml coverage-${{ github.job }}-lightning_gpu_${{ matrix.mpilib }}-1.xml
           # PL_DEVICE=lightning.gpu /opt/mpi/${{ matrix.mpilib }}/bin/mpirun --oversubscribe -n 4 pytest -s -x mpitests/test_device.py -k test_create_device $COVERAGE_FLAGS
-          mv coverage.xml coverage-${{ github.job }}-lightning_gpu_${{ matrix.mpilib }}.xml
 
       - name: Upload code coverage results
         uses: actions/upload-artifact@v3
         with:
           name: ubuntu-codecov-results-python
-          path: coverage-${{ github.job }}-lightning_gpu_${{ matrix.mpilib }}.xml
+          path: coverage-${{ github.job }}-lightning_gpu_${{ matrix.mpilib }}-*.xml
           if-no-files-found: error
 
       - name: Cleanup

diff --git a/pennylane_lightning/core/_serialize.py b/pennylane_lightning/core/_serialize.py
@@ -25,6 +25,8 @@
     Identity,
     StatePrep,
     Rot,
+    Hamiltonian,
+    SparseHamiltonian,
 )
 from pennylane.operation import Tensor
 from pennylane.tape import QuantumTape
@@ -52,6 +54,7 @@ class QuantumScriptSerializer:
     # pylint: disable=import-outside-toplevel, too-many-instance-attributes, c-extension-no-member
     def __init__(self, device_name, use_csingle: bool = False, use_mpi: bool = False):
         self.use_csingle = use_csingle
+        self.device_name = device_name
         if device_name == "lightning.qubit":
             try:
                 import pennylane_lightning.lightning_qubit_ops as lightning_ops
@@ -75,6 +78,7 @@ def __init__(self, device_name, use_csingle: bool = False, use_mpi: bool = False
                 ) from exception
         else:
             raise DeviceError(f'The device name "{device_name}" is not a valid option.')
+        self.statevector_c64 = lightning_ops.StateVectorC64
         self.statevector_c128 = lightning_ops.StateVectorC128
         self.named_obs_c64 = lightning_ops.observables.NamedObsC64
         self.named_obs_c128 = lightning_ops.observables.NamedObsC128
@@ -84,20 +88,25 @@ def __init__(self, device_name, use_csingle: bool = False, use_mpi: bool = False
         self.tensor_prod_obs_c128 = lightning_ops.observables.TensorProdObsC128
         self.hamiltonian_c64 = lightning_ops.observables.HamiltonianC64
         self.hamiltonian_c128 = lightning_ops.observables.HamiltonianC128
+        self.sparse_hamiltonian_c64 = lightning_ops.observables.SparseHamiltonianC64
+        self.sparse_hamiltonian_c128 = lightning_ops.observables.SparseHamiltonianC128
 
         self._use_mpi = False
 
         if use_mpi:
             self._use_mpi = use_mpi
-            self.statevectormpi_c128 = lightning_ops.StateVectorMPIC128
-            self.named_obsmpi_c64 = lightning_ops.observablesMPI.NamedObsMPIC64
-            self.named_obsmpi_c128 = lightning_ops.observablesMPI.NamedObsMPIC128
-            self.hermitian_obsmpi_c64 = lightning_ops.observablesMPI.HermitianObsMPIC64
-            self.hermitian_obsmpi_c128 = lightning_ops.observablesMPI.HermitianObsMPIC128
-            self.tensor_prod_obsmpi_c64 = lightning_ops.observablesMPI.TensorProdObsMPIC64
-            self.tensor_prod_obsmpi_c128 = lightning_ops.observablesMPI.TensorProdObsMPIC128
-            self.hamiltonianmpi_c64 = lightning_ops.observablesMPI.HamiltonianMPIC64
-            self.hamiltonianmpi_c128 = lightning_ops.observablesMPI.HamiltonianMPIC128
+            self.statevector_mpi_c64 = lightning_ops.StateVectorMPIC64
+            self.statevector_mpi_c128 = lightning_ops.StateVectorMPIC128
+            self.named_obs_mpi_c64 = lightning_ops.observablesMPI.NamedObsMPIC64
+            self.named_obs_mpi_c128 = lightning_ops.observablesMPI.NamedObsMPIC128
+            self.hermitian_obs_mpi_c64 = lightning_ops.observablesMPI.HermitianObsMPIC64
+            self.hermitian_obs_mpi_c128 = lightning_ops.observablesMPI.HermitianObsMPIC128
+            self.tensor_prod_obs_mpi_c64 = lightning_ops.observablesMPI.TensorProdObsMPIC64
+            self.tensor_prod_obs_mpi_c128 = lightning_ops.observablesMPI.TensorProdObsMPIC128
+            self.hamiltonian_mpi_c64 = lightning_ops.observablesMPI.HamiltonianMPIC64
+            self.hamiltonian_mpi_c128 = lightning_ops.observablesMPI.HamiltonianMPIC128
+
+            self._mpi_manager = lightning_ops.MPIManager
 
     @property
     def ctype(self):
@@ -113,37 +122,44 @@ def rtype(self):
     def sv_type(self):
         """State vector matching ``use_csingle`` precision (and MPI if it is supported)."""
         if self._use_mpi:
-            return self.statevectormpi_c128
-        return self.statevector_c128
+            return self.statevector_mpi_c64 if self.use_csingle else self.statevector_mpi_c128
+        return self.statevector_c64 if self.use_csingle else self.statevector_c128
 
     @property
     def named_obs(self):
         """Named observable matching ``use_csingle`` precision."""
         if self._use_mpi:
-            return self.named_obsmpi_c64 if self.use_csingle else self.named_obsmpi_c128
+            return self.named_obs_mpi_c64 if self.use_csingle else self.named_obs_mpi_c128
         return self.named_obs_c64 if self.use_csingle else self.named_obs_c128
 
     @property
     def hermitian_obs(self):
         """Hermitian observable matching ``use_csingle`` precision."""
         if self._use_mpi:
-            return self.hermitian_obsmpi_c64 if self.use_csingle else self.hermitian_obsmpi_c128
+            return self.hermitian_obs_mpi_c64 if self.use_csingle else self.hermitian_obs_mpi_c128
         return self.hermitian_obs_c64 if self.use_csingle else self.hermitian_obs_c128
 
     @property
     def tensor_obs(self):
         """Tensor product observable matching ``use_csingle`` precision."""
         if self._use_mpi:
-            return self.tensor_prod_obsmpi_c64 if self.use_csingle else self.tensor_prod_obsmpi_c128
+            return (
+                self.tensor_prod_obs_mpi_c64 if self.use_csingle else self.tensor_prod_obs_mpi_c128
+            )
         return self.tensor_prod_obs_c64 if self.use_csingle else self.tensor_prod_obs_c128
 
     @property
     def hamiltonian_obs(self):
         """Hamiltonian observable matching ``use_csingle`` precision."""
         if self._use_mpi:
-            return self.hamiltonianmpi_c64 if self.use_csingle else self.hamiltonianmpi_c128
+            return self.hamiltonian_mpi_c64 if self.use_csingle else self.hamiltonian_mpi_c128
         return self.hamiltonian_c64 if self.use_csingle else self.hamiltonian_c128
 
+    @property
+    def sparse_hamiltonian_obs(self):
+        """SparseHamiltonian observable matching ``use_csingle`` precision."""
+        return self.sparse_hamiltonian_c64 if self.use_csingle else self.sparse_hamiltonian_c128
+
     def _named_obs(self, observable, wires_map: dict):
         """Serializes a Named observable"""
         wires = [wires_map[w] for w in observable.wires]
@@ -168,6 +184,37 @@ def _hamiltonian(self, observable, wires_map: dict):
         terms = [self._ob(t, wires_map) for t in observable.ops]
         return self.hamiltonian_obs(coeffs, terms)
 
+    def _sparse_hamiltonian(self, observable, wires_map: dict):
+        """Serialize an observable (Sparse Hamiltonian)
+
+        Args:
+            observable (Observable): the input observable (Sparse Hamiltonian)
+            wire_map (dict): a dictionary mapping input wires to the device's backend wires
+
+        Returns:
+            sparse_hamiltonian_obs (SparseHamiltonianC64 or SparseHamiltonianC128): A Sparse Hamiltonian observable object compatible with the C++ backend
+        """
+
+        if self._use_mpi:
+            Hmat = Hamiltonian([1.0], [Identity(0)]).sparse_matrix()
+            H_sparse = SparseHamiltonian(Hmat, wires=range(1))
+            spm = H_sparse.sparse_matrix()
+            # Only root 0 needs the overall sparsematrix data
+            if self._mpi_manager().getRank() == 0:
+                spm = observable.sparse_matrix()
+            self._mpi_manager().Barrier()
+        else:
+            spm = observable.sparse_matrix()
+        data = np.array(spm.data).astype(self.ctype)
+        indices = np.array(spm.indices).astype(np.int64)
+        offsets = np.array(spm.indptr).astype(np.int64)
+
+        wires = []
+        wires_list = observable.wires.tolist()
+        wires.extend([wires_map[w] for w in wires_list])
+
+        return self.sparse_hamiltonian_obs(data, indices, offsets, wires)
+
     def _pauli_word(self, observable, wires_map: dict):
         """Serialize a :class:`pennylane.pauli.PauliWord` into a Named or Tensor observable."""
         if len(observable) == 1:
@@ -195,6 +242,8 @@ def _ob(self, observable, wires_map):
             return self._tensor_ob(observable, wires_map)
         if observable.name == "Hamiltonian":
             return self._hamiltonian(observable, wires_map)
+        if observable.name == "SparseHamiltonian":
+            return self._sparse_hamiltonian(observable, wires_map)
         if isinstance(observable, (PauliX, PauliY, PauliZ, Identity, Hadamard)):
             return self._named_obs(observable, wires_map)
         if observable._pauli_rep is not None:

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.33.0-dev23"
+__version__ = "0.33.0-dev24"
diff --git a/pennylane_lightning/core/src/bindings/Bindings.hpp b/pennylane_lightning/core/src/bindings/Bindings.hpp
@@ -287,7 +287,8 @@ void registerInfo(py::module_ &m) {
  * @tparam StateVectorT
  * @param m Pybind module
  */
-template <class StateVectorT> void registerObservables(py::module_ &m) {
+template <class StateVectorT>
+void registerBackendAgnosticObservables(py::module_ &m) {
     using PrecisionT =
         typename StateVectorT::PrecisionT; // Statevector's precision.
     using ComplexT =
@@ -627,7 +628,8 @@ template <class StateVectorT> void lightningClassBindings(py::module_ &m) {
     /* Observables submodule */
     py::module_ obs_submodule =
         m.def_submodule("observables", "Submodule for observables classes.");
-    registerObservables<StateVectorT>(obs_submodule);
+    registerBackendAgnosticObservables<StateVectorT>(obs_submodule);
+    registerBackendSpecificObservables<StateVectorT>(obs_submodule);
 
     //***********************************************************************//
     //                             Measurements

diff --git a/pennylane_lightning/core/src/bindings/BindingsMPI.hpp b/pennylane_lightning/core/src/bindings/BindingsMPI.hpp
@@ -82,6 +82,10 @@ template <class StateVectorT> void registerObservablesMPI(py::module_ &m) {
 
     using np_arr_c = py::array_t<std::complex<ParamT>, py::array::c_style>;
     using np_arr_r = py::array_t<ParamT, py::array::c_style>;
+    using np_arr_sparse_ind = typename std::conditional<
+        std::is_same<ParamT, float>::value,
+        py::array_t<int32_t, py::array::c_style | py::array::forcecast>,
+        py::array_t<int64_t, py::array::c_style | py::array::forcecast>>::type;
 
     std::string class_name;
 
@@ -191,6 +195,48 @@ template <class StateVectorT> void registerObservablesMPI(py::module_ &m) {
                 return self == other_cast;
             },
             "Compare two observables");
+#ifdef _ENABLE_PLGPU
+    class_name = "SparseHamiltonianMPIC" + bitsize;
+    using SpIDX = typename SparseHamiltonianMPI<StateVectorT>::IdxT;
+    py::class_<SparseHamiltonianMPI<StateVectorT>,
+               std::shared_ptr<SparseHamiltonianMPI<StateVectorT>>,
+               Observable<StateVectorT>>(m, class_name.c_str(),
+                                         py::module_local())
+        .def(py::init([](const np_arr_c &data, const np_arr_sparse_ind &indices,
+                         const np_arr_sparse_ind &offsets,
+                         const std::vector<std::size_t> &wires) {
+            const py::buffer_info buffer_data = data.request();
+            const auto *data_ptr = static_cast<ComplexT *>(buffer_data.ptr);
+
+            const py::buffer_info buffer_indices = indices.request();
+            const auto *indices_ptr = static_cast<SpIDX *>(buffer_indices.ptr);
+
+            const py::buffer_info buffer_offsets = offsets.request();
+            const auto *offsets_ptr = static_cast<SpIDX *>(buffer_offsets.ptr);
+
+            return SparseHamiltonianMPI<StateVectorT>{
+                std::vector<ComplexT>({data_ptr, data_ptr + data.size()}),
+                std::vector<SpIDX>({indices_ptr, indices_ptr + indices.size()}),
+                std::vector<SpIDX>({offsets_ptr, offsets_ptr + offsets.size()}),
+                wires};
+        }))
+        .def("__repr__", &SparseHamiltonianMPI<StateVectorT>::getObsName)
+        .def("get_wires", &SparseHamiltonianMPI<StateVectorT>::getWires,
+             "Get wires of observables")
+        .def(
+            "__eq__",
+            [](const SparseHamiltonianMPI<StateVectorT> &self,
+               py::handle other) -> bool {
+                if (!py::isinstance<SparseHamiltonianMPI<StateVectorT>>(
+                        other)) {
+                    return false;
+                }
+                auto other_cast =
+                    other.cast<SparseHamiltonianMPI<StateVectorT>>();
+                return self == other_cast;
+            },
+            "Compare two observables");
+#endif
 }
 
 /**