Add the state_vector, measurement class and simulate method for the LightningGPU with the new device API

pennylane_lightning/core/_measurements_base.py

@@ -130,6 +130,7 @@ def expval(self, measurementprocess: MeasurementProcess):
             measurementprocess.obs.name, measurementprocess.obs.wires
         )
 
+    @abstractmethod
     def probs(self, measurementprocess: MeasurementProcess):
         """Probabilities of the supplied observable or wires contained in the MeasurementProcess.
 
@@ -139,15 +140,6 @@ def probs(self, measurementprocess: MeasurementProcess):
         Returns:
             Probabilities of the supplied observable or wires
         """
-        diagonalizing_gates = measurementprocess.diagonalizing_gates()
-        if diagonalizing_gates:
-            self._qubit_state.apply_operations(diagonalizing_gates)
-        results = self._measurement_lightning.probs(measurementprocess.wires.tolist())
-        if diagonalizing_gates:
-            self._qubit_state.apply_operations(
-                [qml.adjoint(g, lazy=False) for g in reversed(diagonalizing_gates)]
-            )
-        return results
 
     def var(self, measurementprocess: MeasurementProcess):
         """Variance of the supplied observable contained in the MeasurementProcess.

pennylane_lightning/core/src/simulators/lightning_gpu/measurements/MeasurementsGPU.hpp

@@ -273,7 +273,7 @@ class Measurements final
         PL_CUSTATEVEC_IS_SUCCESS(custatevecSamplerSample(
             this->_statevector.getCusvHandle(), sampler, bitStrings.data(),
             bitOrdering.data(), bitStringLen, rand_nums.data(), num_samples,
-            CUSTATEVEC_SAMPLER_OUTPUT_ASCENDING_ORDER));
+            CUSTATEVEC_SAMPLER_OUTPUT_RANDNUM_ORDER));
         PL_CUDA_IS_SUCCESS(cudaStreamSynchronize(
             this->_statevector.getDataBuffer().getDevTag().getStreamID()));
 

pennylane_lightning/lightning_gpu/_measurements.py

@@ -15,8 +15,30 @@
 Class implementation for state vector measurements.
 """
 
+from warnings import warn
+
+try:
+    from pennylane_lightning.lightning_gpu_ops import MeasurementsC64, MeasurementsC128
+
+    try:
+        from pennylane_lightning.lightning_gpu_ops import MeasurementsMPIC64, MeasurementsMPIC128
+
+        MPI_SUPPORT = True
+    except ImportError as ex:
+        warn(str(ex), UserWarning)
+
+        MPI_SUPPORT = False
+
+except ImportError as ex:
+    warn(str(ex), UserWarning)
+
+    pass
+
+from typing import List
+
 import numpy as np
 import pennylane as qml
+from pennylane.measurements import CountsMP, MeasurementProcess, SampleMeasurement, Shots
 from pennylane.typing import TensorLike
 
 from pennylane_lightning.core._measurements_base import LightningBaseMeasurements
@@ -37,3 +59,98 @@
     ) -> TensorLike:
 
         super().__init__(lgpu_state)
+
+        self._measurement_lightning = self._measurement_dtype()(lgpu_state.state_vector)
+
+    def _measurement_dtype(self):
+        """Binding to Lightning GPU Measurements C++ class.
+
+        Returns: the Measurements class
+        """
+        return MeasurementsC64 if self.dtype == np.complex64 else MeasurementsC128
+
+    def _measure_with_samples_diagonalizing_gates(
+        self,
+        mps: List[SampleMeasurement],
+        shots: Shots,
+    ) -> TensorLike:
+        """
+        Returns the samples of the measurement process performed on the given state,
+        by rotating the state into the measurement basis using the diagonalizing gates
+        given by the measurement process.
+
+        Args:
+            mps (~.measurements.SampleMeasurement): The sample measurements to perform
+            shots (~.measurements.Shots): The number of samples to take
+
+        Returns:
+            TensorLike[Any]: Sample measurement results
+        """
+        # apply diagonalizing gates
+        self._apply_diagonalizing_gates(mps)
+
+        # Specific for LGPU:
+        total_indices = self._qubit_state.num_wires
+        wires = qml.wires.Wires(range(total_indices))
+
+        def _process_single_shot(samples):
+            processed = []
+            for mp in mps:
+                res = mp.process_samples(samples, wires)
+                if not isinstance(mp, CountsMP):
+                    res = qml.math.squeeze(res)
+
+                processed.append(res)
+
+            return tuple(processed)
+
+        try:
+            samples = self._measurement_lightning.generate_samples(
+                len(wires), shots.total_shots
+            ).astype(int, copy=False)
+
+        except ValueError as e:
+            if str(e) != "probabilities contain NaN":
+                raise e
+            samples = qml.math.full((shots.total_shots, len(wires)), 0)
+
+        self._apply_diagonalizing_gates(mps, adjoint=True)
+
+        # if there is a shot vector, use the shots.bins generator to
+        # split samples w.r.t. the shots
+        processed_samples = []
+        for lower, upper in shots.bins():
+            result = _process_single_shot(samples[..., lower:upper, :])
+            processed_samples.append(result)
+
+        return (
+            tuple(zip(*processed_samples)) if shots.has_partitioned_shots else processed_samples[0]
+        )
+
+    def probs(self, measurementprocess: MeasurementProcess):
+        """Probabilities of the supplied observable or wires contained in the MeasurementProcess.
+
+        Args:
+            measurementprocess (StateMeasurement): measurement to apply to the state
+
+        Returns:
+            Probabilities of the supplied observable or wires
+        """
+        diagonalizing_gates = measurementprocess.diagonalizing_gates()
+
+        if diagonalizing_gates:
+            self._qubit_state.apply_operations(diagonalizing_gates)
+
+        results = self._measurement_lightning.probs(measurementprocess.wires.tolist())
+
+        if diagonalizing_gates:
+            self._qubit_state.apply_operations(
+                [qml.adjoint(g, lazy=False) for g in reversed(diagonalizing_gates)]
+            )
+
+        # Device returns as col-major orderings, so perform transpose on data for bit-index shuffle for now.
+        if len(results) > 0:
+            num_local_wires = len(results).bit_length() - 1 if len(results) > 0 else 0
+            return results.reshape([2] * num_local_wires).transpose().reshape(-1)
+
+        return results

pennylane_lightning/lightning_gpu/_mpi_handler.py

@@ -0,0 +1,120 @@
+# Copyright 2022-2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+This module contains the :class:`~.LightningGPU_MPIHandler` class, a MPI handler to use LightningGPU device with multi-GPU on multi-node system.
+"""
+
+try:
+    # pylint: disable=no-name-in-module
+    from pennylane_lightning.lightning_gpu_ops import DevTag, MPIManager
+
+    MPI_SUPPORT = True
+except ImportError:
+    MPI_SUPPORT = False
+
+from typing import Callable, Union
+
+import numpy as np
+
+
+# MPI options
+class LightningGPU_MPIHandler:
+    """MPI handler for PennyLane Lightning GPU device
+
+    MPI handler to use a GPU-backed Lightning device using NVIDIA cuQuantum SDK with parallel capabilities.
+
+    Use the MPI library is necessary to initialize different variables and methods to handle the data across nodes and perform checks for memory allocation on each device.
+
+    Args:
+        mpi (bool): declare if the device will use the MPI support.
+        mpi_buf_size (int): size of GPU memory (in MiB) set for MPI operation and its default value is 64 MiB.
+        dev_pool (Callable): Method to handle the GPU devices available.
+        num_wires (int): the number of wires to initialize the device with.
+        c_dtype (np.complex64, np.complex128): Datatypes for statevector representation
+    """
+
+    def __init__(
+        self,
+        mpi: bool,
+        mpi_buf_size: int,
+        dev_pool: Callable,
+        num_wires: int,
+        c_dtype: Union[np.complex64, np.complex128],
+    ) -> None:
+
+        self.use_mpi = mpi
+        self.mpi_but_size = mpi_buf_size
+        self._dp = dev_pool
+
+        if self.use_mpi:
+
+            if not MPI_SUPPORT:
+                raise ImportError("MPI related APIs are not found.")
+
+            if mpi_buf_size < 0:
+                raise TypeError(f"Unsupported mpi_buf_size value: {mpi_buf_size}, should be >= 0")
+
+            if mpi_buf_size > 0 and (mpi_buf_size & (mpi_buf_size - 1)):
+                raise ValueError(
+                    f"Unsupported mpi_buf_size value: {mpi_buf_size}. mpi_buf_size should be power of 2."
+                )
+
+            # After check if all MPI parameters are ok
+            self.mpi_manager, self.devtag = self._mpi_init_helper(num_wires)
+
+            # set the number of global and local wires
+            commSize = self._mpi_manager.getSize()
+            self.num_global_wires = commSize.bit_length() - 1
+            self.num_local_wires = num_wires - self._num_global_wires
+
+            # Memory size in bytes
+            sv_memsize = np.dtype(c_dtype).itemsize * (1 << self.num_local_wires)
+            if self._mebibytesToBytes(mpi_buf_size) > sv_memsize:
+                raise ValueError("The MPI buffer size is larger than the local state vector size.")
+
+        if not self.use_mpi:
+            self.num_local_wires = num_wires
+            self.num_global_wires = num_wires
+
+    def _mebibytesToBytes(mebibytes):
+        return mebibytes * 1024 * 1024
+
+    def _mpi_init_helper(self, num_wires):
+        """Set up MPI checks and initializations."""
+
+        # initialize MPIManager and config check in the MPIManager ctor
+        mpi_manager = MPIManager()
+
+        # check if number of GPUs per node is larger than number of processes per node
+        numDevices = self._dp.getTotalDevices()
+        numProcsNode = mpi_manager.getSizeNode()
+
+        if numDevices < numProcsNode:
+            raise ValueError(
+                "Number of devices should be larger than or equal to the number of processes on each node."
+            )
+
+        # check if the process number is larger than number of statevector elements
+        if mpi_manager.getSize() > (1 << (num_wires - 1)):
+            raise ValueError(
+                "Number of processes should be smaller than the number of statevector elements."
+            )
+
+        # set GPU device
+        rank = self._mpi_manager.getRank()
+        deviceid = rank % numProcsNode
+        self._dp.setDeviceID(deviceid)
+        devtag = DevTag(deviceid)
+
+        return (mpi_manager, devtag)