[dyn-alloc-base] Lightning dynamically allocate qubit memory

.github/CHANGELOG.md

@@ -6,6 +6,9 @@
 
 ### Improvements
 
+* Lightning devices support dynamically allocated wires (e.g. `qml.device("lightning.qubit")`)
+  [(#1043)](https://github.com/PennyLaneAI/pennylane-lightning/pull/1043)
+
 * Remove the old device API references in the Lightning repo and test suite.
   [(#1057)](https://github.com/PennyLaneAI/pennylane-lightning/pull/1057)
 

mpitests/test_apply.py

@@ -119,7 +119,7 @@ def circuit():
     assert np.allclose(local_state_vector, local_expected_output_cpu, atol=tol, rtol=0)
 
 
-class TestApply:  # pylint: disable=missing-function-docstring,too-many-arguments
+class TestApply:  # pylint: disable=missing-function-docstring,too-many-arguments,too-many-positional-arguments
     """Tests whether the device can apply supported quantum gates."""
 
     @pytest.fixture(params=fixture_params)
@@ -371,8 +371,6 @@ def test_dev_reset(self, tol, dev_mpi):
         comm.Scatter(state_vector, local_state_vector, root=0)
         dev_cpu = qml.device("lightning.qubit", wires=num_wires, c_dtype=c_dtype)
 
-        dev_cpu._statevector.reset_state()
-
         def circuit():
             qml.PauliX(wires=[0])
             qml.PauliX(wires=[0])

mpitests/test_device.py

@@ -39,6 +39,62 @@ def test_create_device():
         dev = qml.device(device_name, mpi=True, wires=4)
 
 
+def test_unsupported_dynamic_wires():
+    with pytest.raises(
+        DeviceError,
+        match="does not support dynamic wires allocation.",
+    ):
+        dev = qml.device(device_name, mpi=True)
+
+
+@pytest.mark.parametrize(
+    "circuit_in, n_wires, wires_list",
+    [
+        (
+            QuantumScript(
+                [
+                    qml.RX(0.1, 0),
+                    qml.CNOT([1, 0]),
+                    qml.RZ(0.1, 1),
+                    qml.CNOT([2, 1]),
+                ],
+                [qml.expval(qml.Z(0))],
+            ),
+            3,
+            [0, 1, 2],
+        ),
+        (
+            QuantumScript(
+                [
+                    qml.RX(0.1, 0),
+                    qml.CNOT([1, 4]),
+                    qml.RZ(0.1, 4),
+                    qml.CNOT([2, 1]),
+                ],
+                [qml.expval(qml.Z(6))],
+            ),
+            7,
+            [0, 1, 4, 2, 6],
+        ),
+    ],
+)
+def test_dynamic_wires_from_circuit_fixed_wires(circuit_in, n_wires, wires_list):
+    """Test that dynamic_wires_from_circuit creates correct statevector and circuit."""
+    dev = qml.device(device_name, mpi=True, wires=n_wires)
+    circuit_out = dev.dynamic_wires_from_circuit(circuit_in)
+
+    assert circuit_out.num_wires == circuit_in.num_wires
+    assert circuit_out.wires == qml.wires.Wires(wires_list)
+    assert circuit_out.operations == circuit_in.operations
+    assert circuit_out.measurements == circuit_in.measurements
+
+    assert dev._statevector._mpi_handler.use_mpi
+    assert (
+        dev._statevector._mpi_handler.num_local_wires
+        + dev._statevector._mpi_handler.num_global_wires
+    ) == n_wires
+
+
 def test_unsupported_mpi_buf_size():
     with pytest.raises(ValueError, match="Unsupported mpi_buf_size value"):
         dev = qml.device(device_name, mpi=True, wires=4, mpi_buf_size=-1)

pennylane_lightning/core/_version.py

@@ -16,4 +16,4 @@
 Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.41.0-dev19"
+__version__ = "0.41.0-dev20"

pennylane_lightning/core/lightning_base.py

@@ -43,7 +43,7 @@ class LightningBase(Device):
     A class that serves as a base class for Lightning simulators.
 
     Args:
-        wires (int or list): number or list of wires to initialize the device with
+        wires (Optional[int, list]): number or list of wires to initialize the device with. Defaults to ``None`` if not specified, and the device will allocate the number of wires depending on the circuit to execute.
         sync (bool): immediately sync with host after applying operations on the device
         c_dtype: Datatypes for statevector representation. Must be one of
             ``np.complex64`` or ``np.complex128``.
@@ -59,7 +59,7 @@ class LightningBase(Device):
 
     def __init__(  # pylint: disable=too-many-arguments
         self,
-        wires: Union[int, List],
+        wires: Union[int, List] = None,
         *,
         c_dtype: Union[np.complex64, np.complex128],
         shots: Union[int, List],
@@ -70,7 +70,10 @@ def __init__(  # pylint: disable=too-many-arguments
         self._c_dtype = c_dtype
         self._batch_obs = batch_obs
 
-        if isinstance(wires, int):
+        # State-vector is dynamically allocated just before execution
+        self._statevector = None
+
+        if isinstance(wires, int) or wires is None:
             self._wire_map = None  # should just use wires as is
         else:
             self._wire_map = {w: i for i, w in enumerate(self.wires)}
@@ -102,6 +105,17 @@ def _setup_execution_config(self, config: ExecutionConfig):
 
         """
 
+    @abstractmethod
+    def dynamic_wires_from_circuit(self, circuit):
+        """Allocate the underlying quantum state from the pre-defined wires or a given circuit if applicable. Circuit wires will be mapped to Pennylane ``default.qubit`` standard wire order.
+
+        Args:
+            circuit (QuantumTape): The circuit to execute.
+
+        Returns:
+            QuantumTape: The updated circuit with the wires mapped to the standard wire order.
+        """
+
     @abstractmethod
     def preprocess(self, execution_config: ExecutionConfig = DefaultExecutionConfig):
         """This function defines the device transform program to be applied and an updated device configuration.
@@ -316,7 +330,12 @@ def compute_derivatives(
         batch_obs = execution_config.device_options.get("batch_obs", self._batch_obs)
 
         return tuple(
-            self.jacobian(circuit, self._statevector, batch_obs=batch_obs, wire_map=self._wire_map)
+            self.jacobian(
+                self.dynamic_wires_from_circuit(circuit),
+                self._statevector,
+                batch_obs=batch_obs,
+                wire_map=self._wire_map,
+            )
             for circuit in circuits
         )
 
@@ -337,9 +356,12 @@ def execute_and_compute_derivatives(
         batch_obs = execution_config.device_options.get("batch_obs", self._batch_obs)
         results = tuple(
             self.simulate_and_jacobian(
-                c, self._statevector, batch_obs=batch_obs, wire_map=self._wire_map
+                self.dynamic_wires_from_circuit(circuit),
+                self._statevector,
+                batch_obs=batch_obs,
+                wire_map=self._wire_map,
             )
-            for c in circuits
+            for circuit in circuits
         )
         return tuple(zip(*results))
 
@@ -392,7 +414,13 @@ def compute_vjp(
         """
         batch_obs = execution_config.device_options.get("batch_obs", self._batch_obs)
         return tuple(
-            self.vjp(circuit, cots, self._statevector, batch_obs=batch_obs, wire_map=self._wire_map)
+            self.vjp(
+                self.dynamic_wires_from_circuit(circuit),
+                cots,
+                self._statevector,
+                batch_obs=batch_obs,
+                wire_map=self._wire_map,
+            )
             for circuit, cots in zip(circuits, cotangents)
         )
 
@@ -415,7 +443,11 @@ def execute_and_compute_vjp(
         batch_obs = execution_config.device_options.get("batch_obs", self._batch_obs)
         results = tuple(
             self.simulate_and_vjp(
-                circuit, cots, self._statevector, batch_obs=batch_obs, wire_map=self._wire_map
+                self.dynamic_wires_from_circuit(circuit),
+                cots,
+                self._statevector,
+                batch_obs=batch_obs,
+                wire_map=self._wire_map,
             )
             for circuit, cots in zip(circuits, cotangents)
         )
@@ -460,6 +492,13 @@ def loop(i, y):
         # has jax dependency, so can't import up top
         from .lightning_interpreter import LightningInterpreter
 
+        if self.wires is None:
+            raise NotImplementedError("Wires must be specified for integration with plxpr capture.")
+
+        self._statevector = self.LightningStateVector(
+            num_wires=len(self.wires), dtype=self._c_dtype
+        )
+
         interpreter = LightningInterpreter(
             self._statevector, self.LightningMeasurements, shots=self.shots
         )

pennylane_lightning/core/src/simulators/base/StateVectorBase.hpp

@@ -204,4 +204,4 @@ template <class PrecisionT, class Derived> class StateVectorBase {
     }
 };
 
-} // namespace Pennylane
+} // namespace Pennylane

pennylane_lightning/core/src/simulators/lightning_qubit/StateVectorLQubitRaw.hpp

@@ -138,4 +138,4 @@ class StateVectorLQubitRaw final
         updateData(new_data.data(), new_data.size());
     }
 };
-} // namespace Pennylane::LightningQubit
+} // namespace Pennylane::LightningQubit

pennylane_lightning/lightning_gpu/lightning_gpu.py

@@ -203,7 +203,7 @@
     :doc:`/lightning_gpu/installation` guide for more details.
 
     Args:
-        wires (int): the number of wires to initialize the device with
+        wires (Optional[int, list]): the number of wires to initialize the device with. Defaults to ``None`` if not specified, and the device will allocate the number of wires depending on the circuit to execute.
         c_dtype: Datatypes for statevector representation. Must be one of
             ``np.complex64`` or ``np.complex128``.
         shots (int): How many times the circuit should be evaluated (or sampled) to estimate
@@ -235,7 +235,7 @@
 
     def __init__(  # pylint: disable=too-many-arguments
         self,
-        wires: Union[int, List],
+        wires: Union[int, List] = None,
         *,
         c_dtype: Union[np.complex128, np.complex64] = np.complex128,
         shots: Union[int, List] = None,
@@ -268,15 +268,22 @@
         self._dp = DevPool()
         self._use_async = use_async
 
-        # Creating the state vector
-        self._mpi_handler = MPIHandler(mpi, mpi_buf_size, len(self.wires), c_dtype)
-
-        self._statevector = self.LightningStateVector(
-            num_wires=len(self.wires),
-            dtype=c_dtype,
-            mpi_handler=self._mpi_handler,
-            use_async=self._use_async,
-        )
+        # Create the state vector only for MPI, otherwise created dynamically before execution
+        if mpi:
+            if wires is None:
+                raise qml.DeviceError(
+                    "Lightning-GPU-MPI does not support dynamic wires allocation."
+                )
+            self._mpi_handler = MPIHandler(mpi, mpi_buf_size, len(self.wires), c_dtype)
+            self._statevector = self.LightningStateVector(
+                num_wires=len(self.wires),
+                dtype=c_dtype,
+                mpi_handler=self._mpi_handler,
+                use_async=self._use_async,
+            )
+        else:
+            self._statevector = None
+            self._mpi_handler = None
 
     @property
     def name(self):
@@ -316,6 +323,35 @@
 
         return replace(config, **updated_values, device_options=new_device_options)
 
+    def dynamic_wires_from_circuit(self, circuit):
+        """Allocate a state-vector from the pre-defined wires or a given circuit if applicable. Circuit wires will be mapped to Pennylane ``default.qubit`` standard wire order.
+
+        Args:
+            circuit (QuantumTape): The circuit to execute.
+
+        Returns:
+            QuantumTape: The updated circuit with the wires mapped to the standard wire order.
+        """
+        if self._mpi_handler and self._mpi_handler.use_mpi:
+            return circuit
+
+        if self.wires is None:
+            num_wires = circuit.num_wires
+            # Map to follow default.qubit wire order for dynamic wires
+            circuit = circuit.map_to_standard_wires()
+        else:
+            num_wires = len(self.wires)
+
+        if (self._statevector is None) or (self._statevector.num_wires != num_wires):
+            self._statevector = self.LightningStateVector(
+                num_wires=num_wires,
+                dtype=self._c_dtype,
+                mpi_handler=None,
+                use_async=self._use_async,
+            )
+
+        return circuit
+
     def preprocess(self, execution_config: ExecutionConfig = DefaultExecutionConfig):
         """This function defines the device transform program to be applied and an updated device configuration.
 
@@ -379,7 +415,7 @@
                 [circuit], _ = qml.map_wires(circuit, self._wire_map)
             results.append(
                 self.simulate(
-                    circuit,
+                    self.dynamic_wires_from_circuit(circuit),
                     self._statevector,
                     postselect_mode=execution_config.mcm_config.postselect_mode,
                 )
@@ -433,7 +469,7 @@
         Note that this function can return measurements for non-commuting observables simultaneously.
         """
         if circuit.shots and (any(isinstance(op, MidMeasureMP) for op in circuit.operations)):
-            if self._mpi_handler.use_mpi:
+            if self._mpi_handler and self._mpi_handler.use_mpi:
                 raise qml.DeviceError(
                     "Lightning-GPU-MPI does not support Mid-circuit measurements."
                 )