From 629095304b36f7017b7ddfc8e38d4ffa6a3fb427 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?= Date: Wed, 11 Sep 2024 10:09:42 -0400 Subject: [PATCH] apply format --- .../lightning_gpu/_measurements.py | 10 +-- .../lightning_gpu/_mpi_handler.py | 72 ++++++++++--------- .../lightning_gpu/_state_vector.py | 65 ++++++++--------- .../lightning_gpu/lightning_gpu.py | 39 +++++----- .../lightning_kokkos/_measurements.py | 6 +- .../lightning_qubit/_measurements.py | 6 +- tests/conftest.py | 6 +- .../test_state_vector_class.py | 8 ++- 8 files changed, 107 insertions(+), 105 deletions(-) diff --git a/pennylane_lightning/lightning_gpu/_measurements.py b/pennylane_lightning/lightning_gpu/_measurements.py index a2bbfcfc9d..82ccb9e30f 100644 --- a/pennylane_lightning/lightning_gpu/_measurements.py +++ b/pennylane_lightning/lightning_gpu/_measurements.py @@ -17,9 +17,9 @@ try: from pennylane_lightning.lightning_gpu_ops import MeasurementsC64, MeasurementsC128 - + try: - from pennylane_lightning.lightning_gpu_ops import MeasurementsMPIC64,MeasurementsMPIC128 + from pennylane_lightning.lightning_gpu_ops import MeasurementsMPIC64, MeasurementsMPIC128 MPI_SUPPORT = True except ImportError: @@ -32,7 +32,7 @@ import numpy as np import pennylane as qml -from pennylane.measurements import CountsMP, SampleMeasurement, Shots, MeasurementProcess +from pennylane.measurements import CountsMP, MeasurementProcess, SampleMeasurement, Shots from pennylane.typing import TensorLike from pennylane_lightning.core._measurements_base import LightningBaseMeasurements @@ -136,12 +136,12 @@ def probs(self, measurementprocess: MeasurementProcess): 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 diff --git a/pennylane_lightning/lightning_gpu/_mpi_handler.py b/pennylane_lightning/lightning_gpu/_mpi_handler.py index 11271d139d..46c7e81cc8 100644 --- a/pennylane_lightning/lightning_gpu/_mpi_handler.py +++ b/pennylane_lightning/lightning_gpu/_mpi_handler.py @@ -17,10 +17,8 @@ try: # pylint: disable=no-name-in-module - from pennylane_lightning.lightning_gpu_ops import ( - DevTag, - MPIManager, - ) + from pennylane_lightning.lightning_gpu_ops import DevTag, MPIManager + MPI_SUPPORT = True except ImportError: MPI_SUPPORT = False @@ -29,46 +27,50 @@ import numpy as np + # MPI options -class LightningGPU_MPIHandler(): - """MPI handler for PennyLane Lightning GPU device - +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. - + + 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 wit.h + num_wires (int): the number of wires to initialize the device wit.h 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: - + 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 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.") - + 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 parameter are ok self.mpi_manager, self.devtag = self._mpi_init_helper(num_wires) @@ -76,44 +78,44 @@ def __init__(self, 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: + 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) diff --git a/pennylane_lightning/lightning_gpu/_state_vector.py b/pennylane_lightning/lightning_gpu/_state_vector.py index 20fab56840..126f0d973d 100644 --- a/pennylane_lightning/lightning_gpu/_state_vector.py +++ b/pennylane_lightning/lightning_gpu/_state_vector.py @@ -15,22 +15,16 @@ Class implementation for lightning_gpu state-vector manipulation. """ try: - from pennylane_lightning.lightning_gpu_ops import ( - StateVectorC64, - StateVectorC128, - ) - - try: # Try to import the MPI modules + from pennylane_lightning.lightning_gpu_ops import StateVectorC64, StateVectorC128 + + try: # Try to import the MPI modules # pylint: disable=no-name-in-module - from pennylane_lightning.lightning_gpu_ops import ( - StateVectorMPIC64, - StateVectorMPIC128, - ) + from pennylane_lightning.lightning_gpu_ops import StateVectorMPIC64, StateVectorMPIC128 MPI_SUPPORT = True except ImportError: MPI_SUPPORT = False - + except ImportError: pass @@ -39,11 +33,10 @@ import numpy as np import pennylane as qml from pennylane import DeviceError -from pennylane.ops.op_math import Adjoint -from pennylane.wires import Wires from pennylane.measurements import MidMeasureMP from pennylane.ops import Conditional -from pennylane import DeviceError +from pennylane.ops.op_math import Adjoint +from pennylane.wires import Wires from pennylane_lightning.core._serialize import global_phase_diagonal from pennylane_lightning.core._state_vector_base import LightningBaseStateVector @@ -60,6 +53,7 @@ qml.QubitUnitary, ) + class LightningGPUStateVector(LightningBaseStateVector): """Lightning GPU state-vector class. @@ -72,18 +66,19 @@ class LightningGPUStateVector(LightningBaseStateVector): device_name(string): state vector device name. Options: ["lightning.gpu"] """ - def __init__(self, - num_wires, - dtype=np.complex128, - device_name="lightning.gpu", - mpi_handler = None, - sync=True, - ): + def __init__( + self, + num_wires, + dtype=np.complex128, + device_name="lightning.gpu", + mpi_handler=None, + sync=True, + ): super().__init__(num_wires, dtype) self._device_name = device_name - + if mpi_handler is None: mpi_handler = LightningGPU_MPIHandler(False, 0, None, num_wires, dtype) @@ -105,7 +100,7 @@ def __init__(self, if not self._mpi_handler.use_mpi: self._qubit_state = self._state_dtype()(self.num_wires) - + self._create_basis_state(0) def _state_dtype(self): @@ -139,7 +134,6 @@ def syncD2H(self, state_vector, use_async=False): [0.+0.j 1.+0.j] """ self._qubit_state.DeviceToHost(state_vector.ravel(order="C"), use_async) - @property def state(self): @@ -158,7 +152,6 @@ def state(self): self.syncD2H(state) return state - def syncH2D(self, state_vector, use_async=False): """Copy the state vector data on host provided by the user to the state vector on the device Args: @@ -179,7 +172,7 @@ def syncH2D(self, state_vector, use_async=False): 1.0 """ self._qubit_state.HostToDevice(state_vector.ravel(order="C"), use_async) - + @staticmethod def _asarray(arr, dtype=None): arr = np.asarray(arr) # arr is not copied @@ -211,10 +204,10 @@ def _apply_state_vector(self, state, device_wires, use_async=False): use_async(bool): indicates whether to use asynchronous memory copy from host to device or not. Note: This function only supports synchronized memory copy from host to device. """ - + if isinstance(state, self._qubit_state.__class__): raise DeviceError("LightningGPU does not support allocate external state_vector.") - + # TODO # state_data = allocate_aligned_array(state.size, np.dtype(self.dtype), True) # state.getState(state_data) @@ -294,7 +287,7 @@ def _apply_lightning_controlled(self, operation): wires = self.wires.indices(operation.wires) matrix = global_phase_diagonal(param, self.wires, control_wires, control_values) state.apply(name, wires, inv, [[param]], matrix) - + def _apply_lightning_midmeasure( self, operation: MidMeasureMP, mid_measurements: dict, postselect_mode: str ): @@ -311,7 +304,7 @@ def _apply_lightning_midmeasure( None """ raise DeviceError("LightningGPU does not support Mid-circuit measurements.") - + def _apply_lightning( self, operations, mid_measurements: dict = None, postselect_mode: str = None ): @@ -342,7 +335,7 @@ def _apply_lightning( invert_param = False method = getattr(state, name, None) wires = list(operation.wires) - + if isinstance(operation, Conditional): if operation.meas_val.concretize(mid_measurements): self._apply_lightning([operation.base]) @@ -364,10 +357,13 @@ def _apply_lightning( except AttributeError: # pragma: no cover # To support older versions of PL mat = operation.matrix - - + r_dtype = np.float32 if self.dtype == np.complex64 else np.float64 - param = [[r_dtype(operation.hash)]] if isinstance(operation, gate_cache_needs_hash) else [] + param = ( + [[r_dtype(operation.hash)]] + if isinstance(operation, gate_cache_needs_hash) + else [] + ) if len(mat) == 0: raise ValueError("Unsupported operation") @@ -378,4 +374,3 @@ def _apply_lightning( param, mat.ravel(order="C"), # inv = False: Matrix already in correct form; ) # Parameters can be ignored for explicit matrices; F-order for cuQuantum - diff --git a/pennylane_lightning/lightning_gpu/lightning_gpu.py b/pennylane_lightning/lightning_gpu/lightning_gpu.py index f9ef0048ac..ad20da9e78 100644 --- a/pennylane_lightning/lightning_gpu/lightning_gpu.py +++ b/pennylane_lightning/lightning_gpu/lightning_gpu.py @@ -20,7 +20,7 @@ from ctypes.util import find_library from importlib import util as imp_util from pathlib import Path -from typing import Optional, Tuple Callable, Union +from typing import Callable, Optional, Tuple, Union import numpy as np import pennylane as qml @@ -37,9 +37,9 @@ Result_or_ResultBatch, ) -from ._mpi_handler import LightningGPU_MPIHandler from ._adjoint_jacobian import LightningGPUAdjointJacobian from ._measurements import LightningGPUMeasurements +from ._mpi_handler import LightningGPU_MPIHandler from ._state_vector import LightningGPUStateVector try: @@ -54,11 +54,10 @@ try: # pylint: disable=no-name-in-module - from pennylane_lightning.lightning_gpu_ops import ( - DevTag, - MPIManager, - ) - from ._mpi_handler import LightningGPU_MPIHandler + from pennylane_lightning.lightning_gpu_ops import DevTag, MPIManager + + from ._mpi_handler import LightningGPU_MPIHandler + MPI_SUPPORT = True except ImportError: MPI_SUPPORT = False @@ -150,6 +149,7 @@ } ) + def stopping_condition(op: Operator) -> bool: """A function that determines whether or not an operation is supported by ``lightning.gpu``.""" # To avoid building matrices beyond the given thresholds. @@ -203,15 +203,15 @@ def _add_adjoint_transforms(program: TransformProgram) -> None: name = "adjoint + lightning.gpu" return 0 + def check_gpu_resources() -> None: - """ Check the available resources of each Nvidia GPU """ - if (find_library("custatevec") is None - and not imp_util.find_spec("cuquantum")): - + """Check the available resources of each Nvidia GPU""" + if find_library("custatevec") is None and not imp_util.find_spec("cuquantum"): + raise ImportError( "custatevec libraries not found. Please pip install the appropriate custatevec library in a virtual environment." ) - + if not DevPool.getTotalDevices(): raise ValueError("No supported CUDA-capable device found") @@ -280,7 +280,7 @@ def __init__( # pylint: disable=too-many-arguments "To manually compile from source, follow the instructions at " "https://docs.pennylane.ai/projects/lightning/en/stable/dev/installation.html." ) - + check_gpu_resources() super().__init__( @@ -298,13 +298,16 @@ def __init__( # pylint: disable=too-many-arguments self._sync = sync # Creating the state vector - - self._mpi_handler = LightningGPU_MPIHandler(mpi, mpi_buf_size, self._dp, self.num_wires, c_dtype) - - self._num_local_wires = self._mpi_handler.num_local_wires - self._statevector = self.LightningStateVector(self.num_wires, dtype=c_dtype, mpi_handler=self._mpi_handler, sync=self._sync) + self._mpi_handler = LightningGPU_MPIHandler( + mpi, mpi_buf_size, self._dp, self.num_wires, c_dtype + ) + self._num_local_wires = self._mpi_handler.num_local_wires + + self._statevector = self.LightningStateVector( + self.num_wires, dtype=c_dtype, mpi_handler=self._mpi_handler, sync=self._sync + ) @property def name(self): diff --git a/pennylane_lightning/lightning_kokkos/_measurements.py b/pennylane_lightning/lightning_kokkos/_measurements.py index 8070c26f80..5e43a9c8cb 100644 --- a/pennylane_lightning/lightning_kokkos/_measurements.py +++ b/pennylane_lightning/lightning_kokkos/_measurements.py @@ -25,7 +25,7 @@ import numpy as np import pennylane as qml -from pennylane.measurements import CountsMP, SampleMeasurement, Shots, MeasurementProcess +from pennylane.measurements import CountsMP, MeasurementProcess, SampleMeasurement, Shots from pennylane.typing import TensorLike from pennylane_lightning.core._measurements_base import LightningBaseMeasurements @@ -131,14 +131,14 @@ def probs(self, measurementprocess: MeasurementProcess): if diagonalizing_gates: self._qubit_state.apply_operations(diagonalizing_gates) results = self._measurement_lightning.probs(measurementprocess.wires.tolist()) - + # print("probs: result:",results) # print('*'*100) 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 diff --git a/pennylane_lightning/lightning_qubit/_measurements.py b/pennylane_lightning/lightning_qubit/_measurements.py index 214e57f351..6958f5b1e5 100644 --- a/pennylane_lightning/lightning_qubit/_measurements.py +++ b/pennylane_lightning/lightning_qubit/_measurements.py @@ -26,7 +26,7 @@ import numpy as np import pennylane as qml -from pennylane.measurements import CountsMP, SampleMeasurement, Shots, MeasurementProcess +from pennylane.measurements import CountsMP, MeasurementProcess, SampleMeasurement, Shots from pennylane.typing import TensorLike from pennylane_lightning.core._measurements_base import LightningBaseMeasurements @@ -157,10 +157,10 @@ def probs(self, measurementprocess: MeasurementProcess): 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 diff --git a/tests/conftest.py b/tests/conftest.py index 4cc0e4c5c0..ace8debfd9 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -142,12 +142,12 @@ def get_device(): from pennylane_lightning.lightning_kokkos_ops import LightningException elif device_name == "lightning.gpu": from pennylane_lightning.lightning_gpu import LightningGPU as LightningDevice - from pennylane_lightning.lightning_gpu._state_vector import ( - LightningGPUStateVector as LightningStateVector, - ) from pennylane_lightning.lightning_gpu._measurements import ( LightningGPUMeasurements as LightningMeasurements, ) + from pennylane_lightning.lightning_gpu._state_vector import ( + LightningGPUStateVector as LightningStateVector, + ) LightningAdjointJacobian = None diff --git a/tests/lightning_qubit/test_state_vector_class.py b/tests/lightning_qubit/test_state_vector_class.py index 6e562d80cc..b1bcdf1de1 100644 --- a/tests/lightning_qubit/test_state_vector_class.py +++ b/tests/lightning_qubit/test_state_vector_class.py @@ -29,7 +29,7 @@ from pennylane_lightning.lightning_kokkos_ops import InitializationSettings except ImportError: pass - + if device_name == "lightning.gpu": from pennylane_lightning.lightning_gpu._mpi_handler import LightningGPU_MPIHandler @@ -90,8 +90,10 @@ def test_apply_state_vector_with_lightning_handle(tol): state_vector_1 = LightningStateVector(2) state_vector_1.apply_operations([qml.BasisState(np.array([0, 1]), wires=[0, 1])]) - if device_name == 'lightning.gpu': - with pytest.raises(qml.DeviceError, match="LightningGPU does not support allocate external state_vector."): + if device_name == "lightning.gpu": + with pytest.raises( + qml.DeviceError, match="LightningGPU does not support allocate external state_vector." + ): state_vector_2 = LightningStateVector(2) state_vector_2._apply_state_vector(state_vector_1.state_vector, Wires([0, 1]))