MPI Sparse Hamiltonian for Adjoint Jacobian

.github/CHANGELOG.md

@@ -2,6 +2,22 @@
 
 ### New features since last release
 
+ * Add sparse Hamiltonian support to multi-node/multi-GPU adjoint methods. 
+ [(#128)] (https://github.com/PennyLaneAI/pennylane-lightning-gpu/pull/128)
+
+ Note each MPI process will return the overall result of the adjoint method. To ensure wires index arranged in a way aligned with cuQuantum backend, `wires` should be reordered in a descending manner in a python script as follows:
+ ```python
+    qml.SparseHamiltonian(
+        qml.Hamiltonian(
+            [0.1], [qml.PauliX(0) @ qml.PauliY(1)]
+        ).sparse_matrix(wires.[::-1]),
+    )
+ ```
+
+ * Add Sparse Hamiltonian support for expectation value calculation.
+ [(#127)] (https://github.com/PennyLaneAI/pennylane-lightning-gpu/pull/127)
+
+
 ### Breaking changes
 
 ### Improvements
@@ -14,6 +30,8 @@
 
 This release contains contributions from (in alphabetical order):
 
+Shuli Shu
+
 ---
 
 # Release 0.31.0

mpitests/test_adjoint_jacobian.py

@@ -31,6 +31,8 @@
     TensorProdObsGPUMPI_C128,
     HamiltonianGPUMPI_C64,
     HamiltonianGPUMPI_C128,
+    SparseHamiltonianGPUMPI_C64,
+    SparseHamiltonianGPUMPI_C128,
 )
 from pennylane_lightning_gpu._serialize import _serialize_ob
 
@@ -1044,49 +1046,142 @@ def circuit(params):
         j_gpu = qml.jacobian(qnode_gpu)(params)
 
 
+custom_wires0 = ["alice", 3.14, -1, 0, "bob", "luc"]
+""" Important Note: Ensure wires index arranged from high to low
+    qml.SparseHamiltonian(
+        qml.Hamiltonian(
+            [0.1], [qml.PauliX(wires=custom_wires0[1]) @ qml.PauliY(wires=custom_wires0[0])]
+        ).sparse_matrix(custom_wires0.[::-1]),
+    ) 
+"""
+
+
 @pytest.mark.parametrize(
     "returns",
     [
         qml.SparseHamiltonian(
             qml.Hamiltonian(
-                [0.1], [qml.PauliX(wires=custom_wires[0]) @ qml.PauliZ(wires=custom_wires[1])]
-            ).sparse_matrix(custom_wires),
-            wires=custom_wires,
+                [0.1], [qml.PauliX(wires=custom_wires0[0]) @ qml.PauliY(wires=custom_wires0[1])]
+            ).sparse_matrix(custom_wires0[::-1]),
+            wires=custom_wires0,
         ),
         qml.SparseHamiltonian(
             qml.Hamiltonian(
-                [2.0], [qml.PauliX(wires=custom_wires[2]) @ qml.PauliZ(wires=custom_wires[0])]
-            ).sparse_matrix(custom_wires),
-            wires=custom_wires,
+                [2.0], [qml.PauliX(wires=custom_wires0[2]) @ qml.PauliZ(wires=custom_wires0[0])]
+            ).sparse_matrix(custom_wires0[::-1]),
+            wires=custom_wires0,
         ),
         qml.SparseHamiltonian(
             qml.Hamiltonian(
-                [1.1], [qml.PauliX(wires=custom_wires[0]) @ qml.PauliZ(wires=custom_wires[2])]
-            ).sparse_matrix(custom_wires),
-            wires=custom_wires,
+                [1.1], [qml.PauliX(wires=custom_wires0[0]) @ qml.PauliZ(wires=custom_wires0[2])]
+            ).sparse_matrix(custom_wires0[::-1]),
+            wires=custom_wires0,
         ),
     ],
 )
-def test_failed_adjoint_SparseHamiltonian(returns):
+def test_adjoint_SparseHamiltonian_custom_wires(returns):
     """Integration tests that compare to default.qubit for a large circuit containing parametrized
     operations and when using custom wire labels"""
 
-    dev_gpu = qml.device("lightning.gpu", wires=custom_wires, mpi=True)
+    comm = MPI.COMM_WORLD
+    dev_gpu = qml.device("lightning.gpu", wires=custom_wires0, mpi=True)
+    dev_cpu = qml.device("default.qubit", wires=custom_wires0)
 
     def circuit(params):
-        circuit_ansatz(params, wires=custom_wires)
+        circuit_ansatz(params, wires=custom_wires0)
         return qml.expval(returns)
 
-    n_params = 30
-    np.random.seed(1337)
-    params = np.random.rand(n_params)
+    if comm.Get_rank() == 0:
+        n_params = 30
+        np.random.seed(1337)
+        params = np.random.rand(n_params)
+    else:
+        params = None
+
+    params = comm.bcast(params, root=0)
 
     qnode_gpu = qml.QNode(circuit, dev_gpu, diff_method="adjoint")
+    qnode_cpu = qml.QNode(circuit, dev_cpu, diff_method="parameter-shift")
 
-    with pytest.raises(
-        RuntimeError, match="LightningGPU-MPI does not currently support SparseHamiltonian."
-    ):
-        j_gpu = qml.jacobian(qnode_gpu)(params)
+    j_gpu = qml.jacobian(qnode_gpu)(params)
+    j_cpu = qml.jacobian(qnode_cpu)(params)
+
+    assert np.allclose(j_cpu, j_gpu)
+
+
+""" Important Note: Ensure wires index arranged from high to low
+    qml.SparseHamiltonian(
+    qml.Hamiltonian(
+                [0.1],
+                [qml.PauliX(1) @ qml.PauliZ(0)],
+            ).sparse_matrix(range(5, -1, -1)),
+            wires=range(6),
+    )
+"""
+
+
+@pytest.mark.parametrize(
+    "returns",
+    [
+        qml.SparseHamiltonian(
+            qml.Hamiltonian(
+                [0.1],
+                [qml.PauliZ(1) @ qml.PauliX(0) @ qml.Identity(2) @ qml.PauliX(4) @ qml.Identity(5)],
+            ).sparse_matrix(range(5, -1, -1)),
+            wires=range(6),
+        ),
+        qml.SparseHamiltonian(
+            qml.Hamiltonian(
+                [0.1],
+                [qml.PauliX(1) @ qml.PauliZ(0)],
+            ).sparse_matrix(range(5, -1, -1)),
+            wires=range(6),
+        ),
+        qml.SparseHamiltonian(
+            qml.Hamiltonian(
+                [0.1],
+                [qml.PauliX(0) @ qml.PauliZ(1)],
+            ).sparse_matrix(range(5, -1, -1)),
+            wires=range(6),
+        ),
+        qml.SparseHamiltonian(
+            qml.Hamiltonian([2.0], [qml.PauliX(1) @ qml.PauliZ(2)]).sparse_matrix(range(5, -1, -1)),
+            wires=range(6),
+        ),
+        qml.SparseHamiltonian(
+            qml.Hamiltonian([1.1], [qml.PauliX(2) @ qml.PauliZ(0)]).sparse_matrix(range(5, -1, -1)),
+            wires=range(6),
+        ),
+    ],
+)
+def test_adjoint_SparseHamiltonian(returns):
+    """Integration tests that compare to default.qubit for a large circuit containing parametrized
+    operations and when using custom wire labels"""
+
+    comm = MPI.COMM_WORLD
+    dev_gpu = qml.device("lightning.gpu", wires=6, mpi=True)
+    dev_cpu = qml.device("default.qubit", wires=6)
+
+    def circuit(params):
+        circuit_ansatz(params, wires=range(6))
+        return qml.expval(returns)
+
+    if comm.Get_rank() == 0:
+        n_params = 30
+        np.random.seed(1337)
+        params = np.random.rand(n_params)
+    else:
+        params = None
+
+    params = comm.bcast(params, root=0)
+
+    qnode_gpu = qml.QNode(circuit, dev_gpu, diff_method="adjoint")
+    qnode_cpu = qml.QNode(circuit, dev_cpu, diff_method="parameter-shift")
+
+    j_gpu = qml.jacobian(qnode_gpu)(params)
+    j_cpu = qml.jacobian(qnode_cpu)(params)
+
+    assert np.allclose(j_cpu, j_gpu)
 
 
 @pytest.mark.parametrize(
@@ -1101,6 +1196,11 @@ def circuit(params):
             HamiltonianGPUMPI_C64,
             HamiltonianGPUMPI_C128,
         ),
+        (
+            qml.SparseHamiltonian(qml.Hamiltonian([1], [qml.PauliZ(0)]).sparse_matrix(), wires=[0]),
+            SparseHamiltonianGPUMPI_C64,
+            SparseHamiltonianGPUMPI_C128,
+        ),
     ],
 )
 @pytest.mark.parametrize("use_csingle", [True, False])

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/_serialize.py

@@ -64,6 +64,8 @@
             TensorProdObsGPUMPI_C128,
             HamiltonianGPUMPI_C64,
             HamiltonianGPUMPI_C128,
+            SparseHamiltonianGPUMPI_C64,
+            SparseHamiltonianGPUMPI_C128,
             HermitianObsGPUMPI_C64,
             HermitianObsGPUMPI_C128,
         )
@@ -117,6 +119,20 @@ def _hamiltonian_ob_dtype(use_csingle, use_mpi: bool):
     )
 
 
+def _sparsehamiltonian_ob_dtype(use_csingle, use_mpi: bool):
+    if not use_mpi:
+        return (
+            [SparseHamiltonianGPU_C64, np.float32, np.int32]
+            if use_csingle
+            else [SparseHamiltonianGPU_C128, np.float64, np.int64]
+        )
+    return (
+        [SparseHamiltonianGPUMPI_C64, np.float32, np.int32]
+        if use_csingle
+        else [SparseHamiltonianGPUMPI_C128, np.float64, np.int64]
+    )
+
+
 def _sv_py_dtype(use_csingle, use_mpi: bool):
     if not use_mpi:
         return LightningGPU_C64 if use_csingle else LightningGPU_C128
@@ -151,24 +167,42 @@ def _serialize_hamiltonian(
     return hamiltonian_obs(coeffs, terms)
 
 
-def _serialize_sparsehamiltonian(ob, wires_map: dict, use_csingle: bool):
-    if use_csingle:
-        ctype = np.complex64
-        rtype = np.int32
-        sparsehamiltonian_obs = SparseHamiltonianGPU_C64
+def _serialize_sparsehamiltonian(ob, wires_map: dict, use_csingle: bool, use_mpi: bool):
+    sparsehamiltonian_obs, ctype, rtype = _sparsehamiltonian_ob_dtype(use_csingle, use_mpi)
+
+    if use_mpi:
+        mpi_manager_local = MPIManager()
+        # Only root 0 needs the overall sparsematrix data
+        if mpi_manager_local.getRank() == 0:
+            spm = ob.sparse_matrix()
+            data = np.array(spm.data).astype(ctype)
+            indices = np.array(spm.indices).astype(rtype)
+            offsets = np.array(spm.indptr).astype(rtype)
+            wires = []
+            wires_list = ob.wires.tolist()
+            wires.extend([wires_map[w] for w in wires_list])
+        else:
+            # Other root only needs non-null some sparsematrix data to pass
+            obs = qml.Identity(0)
+            Hmat = qml.Hamiltonian([1.0], [obs]).sparse_matrix()
+            H_sparse = qml.SparseHamiltonian(Hmat, wires=range(1))
+            spm = H_sparse.sparse_matrix()
+
+            data = np.array(spm.data).astype(ctype)
+            indices = np.array(spm.indices).astype(rtype)
+            offsets = np.array(spm.indptr).astype(rtype)
+
+            wires = []
+            wires_list = ob.wires.tolist()
+            wires.extend([wires_map[w] for w in wires_list])
     else:
-        ctype = np.complex128
-        rtype = np.int64
-        sparsehamiltonian_obs = SparseHamiltonianGPU_C128
-
-    spm = ob.sparse_matrix()
-    data = np.array(spm.data).astype(ctype)
-    indices = np.array(spm.indices).astype(rtype)
-    offsets = np.array(spm.indptr).astype(rtype)
-
-    wires = []
-    wires_list = ob.wires.tolist()
-    wires.extend([wires_map[w] for w in wires_list])
+        spm = ob.sparse_matrix()
+        data = np.array(spm.data).astype(ctype)
+        indices = np.array(spm.indices).astype(rtype)
+        offsets = np.array(spm.indptr).astype(rtype)
+        wires = []
+        wires_list = ob.wires.tolist()
+        wires.extend([wires_map[w] for w in wires_list])
 
     return sparsehamiltonian_obs(data, indices, offsets, wires)
 
@@ -214,9 +248,7 @@ def _serialize_ob(ob, wires_map, use_csingle, use_mpi: bool = False, use_splitti
     elif ob.name == "Hamiltonian":
         return _serialize_hamiltonian(ob, wires_map, use_csingle, use_mpi, use_splitting)
     elif ob.name == "SparseHamiltonian":
-        if use_mpi:
-            raise TypeError("SparseHamiltonian is not supported for MPI backend.")
-        return _serialize_sparsehamiltonian(ob, wires_map, use_csingle)
+        return _serialize_sparsehamiltonian(ob, wires_map, use_csingle, use_mpi)
     elif isinstance(ob, (PauliX, PauliY, PauliZ, Identity, Hadamard)):
         return _serialize_named_ob(ob, wires_map, use_csingle, use_mpi)
     elif ob._pauli_rep is not None: