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

josephleekl · 2025-01-21T20:55:43Z

Before submitting

Please complete the following checklist when submitting a PR:

All new features must include a unit test.
If you've fixed a bug or added code that should be tested, add a test to the
tests directory!
All new functions and code must be clearly commented and documented.
If you do make documentation changes, make sure that the docs build and
render correctly by running make docs.
Ensure that the test suite passes, by running make test.
Add a new entry to the .github/CHANGELOG.md file, summarizing the
change, and including a link back to the PR.
Ensure that code is properly formatted by running make format.

When all the above are checked, delete everything above the dashed
line and fill in the pull request template.

Context:
This PR enables lightning devices to dynamically allocate memory for statevectors just before execution. The usual qml.device("lightning.x", wires=y) is still supported, but user can also not provide wires, i.e. `qml.device("lightning.x"), and the qubits required for the statevector is provided by the circuit

Description of the Change:

Benefits:

dev = qml.device("lightning.qubit")
@qml.qnode(dev)
def circuit():
    qml.Identity(wires=3)
    qml.Identity(wires=1)
    qml.Hadamard(wires=0)
    return qml.state()


@qml.qnode(dev)
def circuit1():
    qml.Identity(wires=1)
    qml.Hadamard(wires=0)
    return qml.state()

circuit()
> array([0.70710678+0.j, 0.70710678+0.j, 0.        +0.j, 0.        +0.j,
       0.        +0.j, 0.        +0.j, 0.        +0.j, 0.        +0.j])
 
circuit1()
> array([0.70710678+0.j, 0.        +0.j, 0.70710678+0.j, 0.        +0.j])

Possible Drawbacks:

Related GitHub Issues:

[sc-82702]

github-actions · 2025-01-21T20:56:02Z

Hello. You may have forgotten to update the changelog!
Please edit .github/CHANGELOG.md with:

A one-to-two sentence description of the change. You may include a small working example for new features.
A link back to this PR.
Your name (or GitHub username) in the contributors section.

codecov · 2025-01-22T14:12:45Z

Codecov Report

Attention: Patch coverage is 44.61538% with 36 lines in your changes missing coverage. Please review.

Project coverage is 91.79%. Comparing base (24e429e) to head (544da03).

Files with missing lines	Patch %	Lines
pennylane_lightning/core/lightning_newAPI_base.py	8.82%	31 Missing ⚠️
...ylane_lightning/lightning_qubit/lightning_qubit.py	0.00%	3 Missing ⚠️
...nnylane_lightning/lightning_qubit/_state_vector.py	33.33%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1043      +/-   ##
==========================================
- Coverage   91.89%   91.79%   -0.10%     
==========================================
  Files         178      178              
  Lines       27465    27523      +58     
==========================================
+ Hits        25238    25266      +28     
- Misses       2227     2257      +30

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

AmintorDusko

Nice job! A few comments and suggestions.

pennylane_lightning/core/lightning_newAPI_base.py

AmintorDusko · 2025-01-24T14:23:09Z

pennylane_lightning/core/lightning_newAPI_base.py

        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
+        results = []
+        for circuit in circuits:
+            if self.wires is None:  # Dynamic wires allocation
+                if self._statevector is None:
+                    self._statevector = self.LightningStateVector(
+                        num_wires=circuit.num_wires, dtype=self._c_dtype
+                    )
+                else:
+                    if self._statevector.num_wires != circuit.num_wires:
+                        self._statevector.update_num_qubits(circuit.num_wires)
+                circuit = circuit.map_to_standard_wires()
+
+            results.append(
+                self.simulate_and_jacobian(
+                    circuit, self._statevector, batch_obs=batch_obs, wire_map=self._wire_map
+                )
            )
-            for c in circuits
-        )
        return tuple(zip(*results))


And you know what? I'm not saying you need to do that, but the whole logic inside compute_derivative and execute_and_compute derivatives is essentially the same, right? We could have a single internal method that would depend on an external (to it) defined function (self.jacobian and self.simulate_and_jacobian in the case) and we call it in both cases.

Indeed it's a bit confusing to have the combinations like [execute_and_, nothing] x [compute_derivative, vjp, ...]. I'll look at if this could be simplified

pennylane_lightning/core/lightning_newAPI_base.py

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

AmintorDusko · 2025-01-24T15:23:23Z

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

+    void updateNumQubits(std::size_t num_qubits) {
+        BaseType::num_qubits_ = num_qubits;
+        BaseType::setKernels(num_qubits, BaseType::threading_,
+                             BaseType::memory_model_);
+        data_.resize(exp2(num_qubits));
+    }


Is changing the number of qubits in place in the current scope?
If yes we may need to implement how we treat the data already there when we change the statevector's number of qubits. @maliasadi has experience with it, he might have a good intuition about this decision.

The reason I believe we need to update num qubits in place is that we want the same device qml.device('lightning.qubit') to be able to execute circuit with different sizes one after another, so resizing is needed.

We need to determine what is the expected behaviour for the data already at the statevector when resizing.

pennylane_lightning/lightning_qubit/_state_vector.py

tests/lightning_qubit/test_state_vector_class.py

tests/test_apply.py

AmintorDusko · 2025-01-24T15:41:39Z

tests/test_apply.py

+    def test_dynamic_allocate_qubit(self, qubit_device, tol):
+        """Test the dynamic allocation of qubits in Lightning devices"""
+
+        dev = qubit_device(None)
+        dev_default = qml.device("default.qubit")
+
+        def circuit0():
+            qml.Hadamard(wires=0)
+            qml.Identity(wires=1)
+            return qml.state()
+
+        def circuit1():
+            qml.Identity(wires=2)
+            qml.Identity(wires=1)
+            qml.Hadamard(wires=0)
+            return qml.state()
+
+        def circuit2():
+            qml.Identity(wires=3)
+            qml.Identity(wires=1)
+            qml.Hadamard(wires=0)
+            return qml.state()
+
+        def circuit3():
+            qml.Identity(wires=3)
+            qml.Hadamard(wires=0)
+            return qml.state()

-        with pytest.raises(TypeError, match="missing 1 required positional argument: 'wires'"):
-            qml.device(device_name)
+        for circuit in [circuit0, circuit1, circuit2, circuit3]:
+            results = qml.qnode(dev)(circuit)()
+            expected = qml.qnode(dev_default)(circuit)()
+            assert np.allclose(results, expected, atol=tol, rtol=0)


What do you think about making the circuit cases a pytest parameter and providing the expected result explicitly?

I am happy to use the explicit results later, just before we merge the PR; for now it's convenient to make sure the results are consistent with default.qubit (the results can be counter-intuitive depending on the wires!)

OK. So you are going to make the change when the PR is already fully reviewed, right?

josephleekl · 2025-01-24T19:34:40Z

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

+    void updateNumQubits(std::size_t num_qubits) {
+        BaseType::num_qubits_ = num_qubits;
+        length_ = exp2(num_qubits);
+        BaseType::setKernels(num_qubits, BaseType::threading_,


Need to double check whether I need to manually dealloc and alloc, or will this setKernels call handle the memory correctly

setKernel calls the method of a new instance created by getInstance, and all the data stored in an OperationKernelMap instance is stored in STL classes with their memory management. Right?
Please double-check, though.

pennylane_lightning/core/lightning_newAPI_base.py

lightning qubit can change numqubits

69e6269

ringo-but-quantum and others added 2 commits January 21, 2025 20:56

Auto update version from '0.41.0-dev5' to '0.41.0-dev6'

5d6ddd7

LQ support updating numqubit

fc6a9e6

josephleekl and others added 12 commits January 22, 2025 21:51

LQ dyn

3cb66a3

remove print statements

ee62471

add test and format

4835b1f

style fix

f34bf83

update docstring

2f1d0ce

update

550cbaf

restore lightning_base

d0d5662

update

7c32b82

Auto update version from '0.41.0-dev6' to '0.41.0-dev7'

3863001

revert arg

3e7864d

Merge branch 'master' into dyn-alloc

bc85851

update test

55fa41e

josephleekl marked this pull request as ready for review January 23, 2025 20:50

josephleekl requested review from AmintorDusko, multiphaseCFD, maliasadi and LuisAlfredoNu and removed request for AmintorDusko and multiphaseCFD January 23, 2025 20:50

josephleekl added 3 commits January 23, 2025 21:12

num_qubit property

04be8ed

property

3f853de

update doc

9e5225e

AmintorDusko reviewed Jan 24, 2025

View reviewed changes

josephleekl and others added 8 commits January 24, 2025 15:46

add tests

b454039

amintor comments

6f9e173

refactor

43b3c0f

format

a8e75d7

update changelog

05fbfa9

update doc

900c6e5

Merge branch 'master' into dyn-alloc

97e40fc

Auto update version from '0.41.0-dev7' to '0.41.0-dev8'

544da03

josephleekl changed the title ~~Lightning dynamically allocate qubit memory~~ [dyn-alloc-base] Lightning dynamically allocate qubit memory Jan 24, 2025

josephleekl commented Jan 24, 2025

View reviewed changes

pennylane_lightning/core/lightning_newAPI_base.py Show resolved Hide resolved

josephleekl added 3 commits January 24, 2025 21:45

make abstract

142e6ec

make abstract

3f3e532

update test

bbf2cae

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

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

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

josephleekl commented Jan 21, 2025 •

edited

Loading

github-actions bot commented Jan 21, 2025

codecov bot commented Jan 22, 2025 •

edited

Loading

AmintorDusko left a comment

AmintorDusko Jan 24, 2025

josephleekl Jan 24, 2025

AmintorDusko Jan 24, 2025

josephleekl Jan 24, 2025

AmintorDusko Jan 24, 2025

AmintorDusko Jan 24, 2025

josephleekl Jan 24, 2025

AmintorDusko Jan 24, 2025

josephleekl Jan 24, 2025

AmintorDusko Jan 24, 2025 •

edited

Loading

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

Are you sure you want to change the base?

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

Conversation

josephleekl commented Jan 21, 2025 • edited Loading

Before submitting

github-actions bot commented Jan 21, 2025

codecov bot commented Jan 22, 2025 • edited Loading

Codecov Report

AmintorDusko left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

AmintorDusko Jan 24, 2025 • edited Loading

Choose a reason for hiding this comment

josephleekl commented Jan 21, 2025 •

edited

Loading

codecov bot commented Jan 22, 2025 •

edited

Loading

AmintorDusko Jan 24, 2025 •

edited

Loading