MPI support for Sparse Hamiltonian expval (#127)

* add cpp layer

* add docstring

* add python layer

* quick update

* make format

* Add changelog

* make format

* Trigger MPI CI

* update base on comments

* update based on comments

* make format & add to-do work

* make format

* sparse matrix struct to class

* update unit test cases

* update docstring

* add more docstring

* make format

* bug fix

* quick update

* add more docstring

.github/CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ### New features since last release
 
+ * Add Sparse Hamiltonian support for expectation value calculation.
+ [(#127)] (https://github.com/PennyLaneAI/pennylane-lightning-gpu/pull/127)
+
+
 ### Breaking changes
 
 ### Improvements
@@ -14,6 +18,8 @@
 
 This release contains contributions from (in alphabetical order):
 
+Shuli Shu
+
 ---
 
 # Release 0.31.0

mpitests/test_adjoint_jacobian.py

@@ -1044,6 +1044,7 @@ def circuit(params):
         j_gpu = qml.jacobian(qnode_gpu)(params)
 
 
+'''
 @pytest.mark.parametrize(
     "returns",
     [
@@ -1088,6 +1089,8 @@ def circuit(params):
     ):
         j_gpu = qml.jacobian(qnode_gpu)(params)
 
+'''
+
 
 @pytest.mark.parametrize(
     "obs,obs_type_c64,obs_type_c128",

mpitests/test_apply.py

@@ -553,6 +553,52 @@ def circuit():
         assert np.allclose(local_state_vector, local_expected_output_cpu, atol=tol, rtol=0)
 
 
+class TestSparseHamExpval:
+    """Tests sparse hamiltonian expectation values."""
+
+    def test_sparse_hamiltonian_expectation(self, tol):
+        comm = MPI.COMM_WORLD
+        commSize = comm.Get_size()
+        num_global_wires = commSize.bit_length() - 1
+        num_local_wires = 3 - num_global_wires
+
+        obs = qml.Identity(0) @ qml.PauliX(1) @ qml.PauliY(2)
+        obs1 = qml.Identity(1)
+        Hmat = qml.Hamiltonian([1.0, 1.0], [obs1, obs]).sparse_matrix()
+
+        state_vector = np.array(
+            [
+                0.0 + 0.0j,
+                0.0 + 0.1j,
+                0.1 + 0.1j,
+                0.1 + 0.2j,
+                0.2 + 0.2j,
+                0.2 + 0.3j,
+                0.3 + 0.3j,
+                0.3 + 0.5j,
+            ],
+            dtype=np.complex128,
+        )
+
+        local_state_vector = np.zeros(1 << num_local_wires).astype(np.complex128)
+        comm.Scatter(state_vector, local_state_vector, root=0)
+
+        dev_gpumpi = qml.device("lightning.gpu", wires=3, mpi=True, c_dtype=np.complex128)
+        dev_gpu = qml.device("lightning.gpu", wires=3, mpi=False, c_dtype=np.complex128)
+
+        dev_gpumpi.syncH2D(local_state_vector)
+        dev_gpu.syncH2D(state_vector)
+
+        H_sparse = qml.SparseHamiltonian(Hmat, wires=range(3))
+
+        comm.Barrier()
+
+        res = dev_gpumpi.expval(H_sparse)
+        expected = dev_gpu.expval(H_sparse)
+
+        assert np.allclose(res, expected)
+
+
 class TestExpval:
     """Tests that expectation values are properly calculated or that the proper errors are raised."""
 
@@ -573,6 +619,7 @@ def test_expval_single_wire_no_parameters(self, tol, operation, wires):
         obs = operation(wires)
         expval_single_wire_no_param(tol, obs)
 
+    @pytest.fixture(params=[np.complex64, np.complex128])
     @pytest.mark.parametrize(
         "obs",
         [
@@ -584,13 +631,13 @@ def test_expval_single_wire_no_parameters(self, tol, operation, wires):
             qml.PauliZ(numQubits - 2) @ qml.PauliZ(numQubits - 1),
         ],
     )
-    def test_expval_multiple_obs(self, obs, tol):
+    def test_expval_multiple_obs(self, obs, request, tol):
         """Test expval with Hamiltonian"""
         num_wires = numQubits
         comm = MPI.COMM_WORLD
 
-        dev_cpu = qml.device("default.qubit", wires=num_wires, c_dtype=np.complex128)
-        dev_gpumpi = qml.device("lightning.gpu", wires=num_wires, mpi=True, c_dtype=np.complex128)
+        dev_cpu = qml.device("default.qubit", wires=num_wires, c_dtype=request.params)
+        dev_gpumpi = qml.device("lightning.gpu", wires=num_wires, mpi=True, c_dtype=request.params)
 
         def circuit():
             qml.RX(0.4, wires=[0])

pennylane_lightning_gpu/lightning_gpu.py

@@ -832,8 +832,19 @@ def expval(self, observable, shot_range=None, bin_size=None):
                         CSR_SparseHamiltonian.data,
                     )
                 else:
-                    raise RuntimeError(
-                        "LightningGPU-MPI does not currently support SparseHamiltonian."
+                    # Identity for CSR_SparseHamiltonian to pass to processes with rank != 0 to reduce
+                    # host(cpu) memory requirements
+                    obs = qml.Identity(0)
+                    Hmat = qml.Hamiltonian([1.0], [obs]).sparse_matrix()
+                    H_sparse = qml.SparseHamiltonian(Hmat, wires=range(1))
+                    CSR_SparseHamiltonian = H_sparse.sparse_matrix().tocsr()
+                    # CSR_SparseHamiltonian for rank == 0
+                    if self._mpi_manager.getRank() == 0:
+                        CSR_SparseHamiltonian = observable.sparse_matrix().tocsr()
+                    return self._gpu_state.ExpectationValue(
+                        CSR_SparseHamiltonian.indptr,
+                        CSR_SparseHamiltonian.indices,
+                        CSR_SparseHamiltonian.data,
                     )
 
             if observable.name in ["Hamiltonian"]:

pennylane_lightning_gpu/src/bindings/Bindings.cpp

@@ -1352,6 +1352,26 @@ void StateVectorCudaMPI_class_bindings(py::module &m) {
             },
             "Calculate the expectation value of the Hamiltonian observable "
             "with custatevecComputeExpectation.")
+
+        .def(
+            "ExpectationValue",
+            [](StateVectorCudaMPI<PrecisionT> &sv,
+               const np_arr_sparse_ind &csrOffsets,
+               const np_arr_sparse_ind &columns, const np_arr_c values) {
+                using index_type = typename std::conditional<
+                    std::is_same<ParamT, float>::value, int32_t, int64_t>::type;
+                return sv.template getExpectationValueOnSparseSpMV<index_type>(
+                    static_cast<index_type *>(csrOffsets.request().ptr),
+                    static_cast<index_type>(
+                        csrOffsets.request()
+                            .size), // num_rows + 1 or csrOffsets
+                    static_cast<index_type *>(columns.request().ptr), // columns
+                    static_cast<std::complex<PrecisionT> *>(
+                        values.request().ptr),
+                    static_cast<index_type>(values.request().size)); // nnz
+            },
+            "Calculate the expectation value of a sparse Hamiltonian.")
+
         .def(
             "ExpectationValue",
             [](StateVectorCudaMPI<PrecisionT> &sv,