Add operations for setting state and basis state and integrate this with `skip_initial_state_prep` #955

rauletorresc · 2024-07-23T23:49:44Z

Context: It has been determined that qml.StatePrep takes too long to compile. To facilitate this, it was determined that the optimization of setting the initial state to the operand of qml.StatePrep could help. This is already implemented in lightning.qubit and lightning.kokkos when using them without @qjit.

Description of the Change: This change adds two new JAXPR and MLIR operations. One for qml.StatePrep and another for qml.BasisState. This change includes:

Changes in the runtime to add support for this new operations, including changes to the QuantumDevice interface
Changes in the TOML specification
Changes in the frontend
A runtime implementation of qml.BasisState which provides validation at runtime. This is needed because when using jax.numpy.arrays all values are staged.

It was necessary to add new operations instead of reusing quantum.custom op because of the signature of these operations. I.e., quantum.custom expects floating point parameters, but qml.StatePrep expects a vector of complex numbers and qml.BasisState expects an integer index.

Benefits: Possibly faster compile times (has not been verified).

Possible Drawbacks: The introduction of new operations and extension of the runtime for operations that not all devices will need can be seen as a code smell. Kokkos does not implement this since it does not implement the new device API.

Related GitHub Issues: #939

close #939

[sc-69069]

frontend/catalyst/api_extensions/quantum_operators.py

frontend/catalyst/device/qjit_device.py

erick-xanadu · 2024-07-30T16:07:18Z

@dime10 any comments about the name of the flag in the TOML files being: skip_initial_state_prep?

paul0403 · 2024-07-30T16:15:54Z

@dime10 any comments about the name of the flag in the TOML files being: skip_initial_state_prep?

Maybe do_not_decompose_initial_state_prep? I can see people possibly interpreting skip_initial_state_prep as "ignore this state prep operation altogether"

dime10 · 2024-07-30T16:17:20Z

@dime10 any comments about the name of the flag in the TOML files being: skip_initial_state_prep?

Yes, skip makes no sense to me in the context of the TOML file. Looking at the existing flags that indicate "support" for something:

        # If the device supports mid circuit measurements natively
        mid_circuit_measurement = true
        # This field is currently unchecked but it is reserved for the purpose of
        # determining if the device supports dynamic qubit allocation/deallocation.
        dynamic_qubit_management = false

I would say just initial_state_prep is fine.

**Context:** We would like Catalyst to optimize the first instance of `qml.StatePrep` by just setting the state vector in PennyLane lightning. **Description of the Change:** To achieve this, we set a flag in the TOML files where it denotes whether the value for the flag `skip_initial_state_prep` used during decomposition. This flag is already in PennyLane and will skip the initial state preparation. There are other ways to achieve this but this is the least invasive way to do so. Some other alternatives that were discussed: 1. Adding StatePrep to the list of supported operations: this is not a good idea as it would lead StatePrep to not be decomposed ever. 2. Creating a new operation and map StatePrep directly into this new operation only in the context of Catalyst. This would lead to setting this new `qml` Operation in the list of supported operations on the device, even though it would never be seen except when using with Catalyst. The way this was achieved was not to create a new `qml` Operation, only an JAXPR/MLIR operation and bind `StatePrep` to this JAXPR operation and lower it through MLIR. So, no new `qml` operations. **Benefits:** Optimization **Possible Drawbacks:** More flags. No one likes flags. **Related GitHub Issues:** Will be followed by this PennyLaneAI/catalyst#955 pull request in Catalyst. [sc-69069] --------- Co-authored-by: ringo-but-quantum <[email protected]> Co-authored-by: Ali Asadi <[email protected]>

codecov · 2024-08-02T20:34:37Z

Codecov Report

Attention: Patch coverage is 95.65217% with 6 lines in your changes missing coverage. Please review.

Project coverage is 97.89%. Comparing base (4f18276) to head (507302c).
Report is 212 commits behind head on main.

Files with missing lines	Patch %	Lines
runtime/include/QuantumDevice.hpp	0.00%	4 Missing ⚠️
frontend/catalyst/jax_tracer.py	93.93%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #955      +/-   ##
==========================================
- Coverage   97.92%   97.89%   -0.04%     
==========================================
  Files          75       75              
  Lines       10561    10684     +123     
  Branches     1218     1226       +8     
==========================================
+ Hits        10342    10459     +117     
- Misses        172      177       +5     
- Partials       47       48       +1

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

erick-xanadu · 2024-08-02T21:24:24Z

@dime10 you mentioned that we shouldn't support contextual decomposition. (I.e., if qml.StatePrep is inside grad, it is okay to emit an error). There is a demo with this. I added a small hack.

dime10 · 2024-08-02T22:18:42Z

@dime10 you mentioned that we shouldn't support contextual decomposition. (I.e., if qml.StatePrep is inside grad, it is okay to emit an error). There is a demo with this. I added a small hack.

I just mentioned that we don't do this for any other operations yet, so it's not a requirement in order to add the support for this op. If you implemented it anyways that's fine 👍
Which demo?

(currently reviewing)

erick-xanadu · 2024-08-02T22:19:59Z

@dime10 adaptive circuits demo. I just added qml.RX(0, wires=[0]) before BasisState. Not a great hack. 😅

paul0403 · 2024-08-03T14:44:40Z

Reminder to update the same changes on lightning kokkos to the lightning repo (e.g. PennyLaneAI/pennylane-lightning#819)

**Context:** We would like Catalyst to optimize the first instance of `qml.StatePrep` by just setting the state vector in PennyLane lightning. **Description of the Change:** To achieve this, we set a flag in the TOML files where it denotes whether the value for the flag `skip_initial_state_prep` used during decomposition. This flag is already in PennyLane and will skip the initial state preparation. There are other ways to achieve this but this is the least invasive way to do so. Some other alternatives that were discussed: 1. Adding StatePrep to the list of supported operations: this is not a good idea as it would lead StatePrep to not be decomposed ever. 2. Creating a new operation and map StatePrep directly into this new operation only in the context of Catalyst. This would lead to setting this new `qml` Operation in the list of supported operations on the device, even though it would never be seen except when using with Catalyst. The way this was achieved was not to create a new `qml` Operation, only an JAXPR/MLIR operation and bind `StatePrep` to this JAXPR operation and lower it through MLIR. So, no new `qml` operations. **Benefits:** Optimization **Possible Drawbacks:** More flags. No one likes flags. **Related GitHub Issues:** Will be followed by this PennyLaneAI/catalyst#955 pull request in Catalyst. [sc-69069] --------- Co-authored-by: ringo-but-quantum <[email protected]> Co-authored-by: Ali Asadi <[email protected]>

erick-xanadu · 2024-08-05T14:39:32Z

@paul0403 changes will be on this PR: PennyLaneAI/pennylane-lightning#838 Thanks!

Co-authored-by: David Ittah <[email protected]>

doc/changelog.md

erick-xanadu · 2024-08-13T19:45:41Z

I added building wheels as I want to test macOS specifically for the use of int8_t for the basis state.

Co-authored-by: David Ittah <[email protected]>

dime10 · 2024-08-13T19:58:37Z

I added building wheels as I want to test macOS specifically for the use of int8_t for the basis state.

Sounds good, I'm trying to test locally right now but the lightning wheels weren't available, I have notified the perf team though.

dime10 · 2024-08-13T20:54:05Z

I added building wheels as I want to test macOS specifically for the use of int8_t for the basis state.

Sounds good, I'm trying to test locally right now but the lightning wheels weren't available, I have notified the perf team though.

Alright they should all be available now 👍

demos/adaptive_circuits_demo.ipynb

dime10

🚀

erick-xanadu · 2024-08-14T13:41:34Z

Do note that coverage is wrong here, coverage is not calculated with OpenQASM3 runtime.

**Context:** #955 could not test lightning.kokkos due to the changes in the QuantumDevice interface along with the separation of repositories. It can only be tested once lightning.kokkos has been built and released with the changes in the QuantumDevice interface. **Description of the Change:** Restore kokkos tests EDIT: Do not merge until there's a new dev release from lightning.

rauletorresc requested a review from a team July 23, 2024 23:49

rauletorresc self-assigned this Jul 23, 2024

erick-xanadu reviewed Jul 24, 2024

View reviewed changes

frontend/catalyst/api_extensions/quantum_operators.py Outdated Show resolved Hide resolved

dime10 reviewed Jul 24, 2024

View reviewed changes

frontend/catalyst/device/qjit_device.py Outdated Show resolved Hide resolved

rauletorresc force-pushed the raultorres/stateprep_lightning branch from da5eff0 to 3e36012 Compare July 24, 2024 21:46

paul0403 mentioned this pull request Jul 25, 2024

[Bug] qml.StatePrep decomposes unnecessarily for Lightning, and the decomposition is very slow #939

Closed

erick-xanadu mentioned this pull request Jul 30, 2024

TOML files update PennyLaneAI/pennylane-lightning#826

Merged

erick-xanadu force-pushed the raultorres/stateprep_lightning branch from 49b70cf to 47ed503 Compare July 31, 2024 17:57

erick-xanadu changed the title ~~StatePrep does not decompose for Lightning~~ [frontend/middle-end/runtime] Add operations for setting state and basis state and integrate this with skip_initial_state_prep Jul 31, 2024

erick-xanadu changed the title ~~[frontend/middle-end/runtime] Add operations for setting state and basis state and integrate this with skip_initial_state_prep~~ Add operations for setting state and basis state and integrate this with skip_initial_state_prep Jul 31, 2024

erick-xanadu force-pushed the raultorres/stateprep_lightning branch from 5b30977 to 93e195c Compare July 31, 2024 23:44

erick-xanadu marked this pull request as ready for review July 31, 2024 23:46

erick-xanadu force-pushed the raultorres/stateprep_lightning branch 2 times, most recently from 1b8c01f to 449dc68 Compare August 2, 2024 12:43

erick-xanadu requested review from rmoyard and mehrdad2m August 2, 2024 16:43

erick-xanadu force-pushed the raultorres/stateprep_lightning branch from 398d9e0 to e4928a7 Compare August 2, 2024 20:19

erick-xanadu mentioned this pull request Aug 5, 2024

Add SetBasisState and SetState to lightning.kokkos PennyLaneAI/pennylane-lightning#838

Closed

erick-xanadu added the author:build-wheels Run the wheel building workflows on this Pull Request label Aug 5, 2024

erick-xanadu and others added 3 commits August 13, 2024 15:40

Update doc/changelog.md

c76b8fd

Co-authored-by: David Ittah <[email protected]>

Update doc/changelog.md

0451092

Co-authored-by: David Ittah <[email protected]>

comment

3e7625f

erick-xanadu requested a review from dime10 August 13, 2024 19:43

erick-xanadu added the author:build-wheels Run the wheel building workflows on this Pull Request label Aug 13, 2024

dime10 reviewed Aug 13, 2024

View reviewed changes

doc/changelog.md Outdated Show resolved Hide resolved

Update doc/changelog.md

0923e2b

Co-authored-by: David Ittah <[email protected]>

erick-xanadu and others added 3 commits August 13, 2024 16:08

Merge branch 'main' into raultorres/stateprep_lightning

041990b

RTD

3586386

Documentation

ee9c008

test

8baa687

erick-xanadu reviewed Aug 13, 2024

View reviewed changes

demos/adaptive_circuits_demo.ipynb Outdated Show resolved Hide resolved

Update demos/adaptive_circuits_demo.ipynb

5a80cae

erick-xanadu reviewed Aug 13, 2024

View reviewed changes

demos/adaptive_circuits_demo.ipynb Outdated Show resolved Hide resolved

erick-xanadu reviewed Aug 13, 2024

View reviewed changes

demos/adaptive_circuits_demo.ipynb Outdated Show resolved Hide resolved

erick-xanadu added 2 commits August 13, 2024 17:10

Apply suggestions from code review

8de8d9b

Merge branch 'main' into raultorres/stateprep_lightning

9fce036

dime10 approved these changes Aug 13, 2024

View reviewed changes

fix

803822f

erick-xanadu removed the author:build-wheels Run the wheel building workflows on this Pull Request label Aug 13, 2024

erick-xanadu added 2 commits August 13, 2024 18:38

Install correct kokkos

9c76b5f

Disable kokkos, use correct kokkos version on wheels

507302c

erick-xanadu merged commit dcbc683 into main Aug 14, 2024
43 of 45 checks passed

erick-xanadu deleted the raultorres/stateprep_lightning branch August 14, 2024 15:50

erick-xanadu mentioned this pull request Aug 14, 2024

[Build System] Restore kokkos tests #1021

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add operations for setting state and basis state and integrate this with `skip_initial_state_prep` #955

Add operations for setting state and basis state and integrate this with `skip_initial_state_prep` #955

rauletorresc commented Jul 23, 2024 •

edited by erick-xanadu

Loading

erick-xanadu commented Jul 30, 2024

paul0403 commented Jul 30, 2024

dime10 commented Jul 30, 2024

codecov bot commented Aug 2, 2024 •

edited

Loading

erick-xanadu commented Aug 2, 2024

dime10 commented Aug 2, 2024 •

edited

Loading

erick-xanadu commented Aug 2, 2024

paul0403 commented Aug 3, 2024

erick-xanadu commented Aug 5, 2024

erick-xanadu commented Aug 13, 2024

dime10 commented Aug 13, 2024

dime10 commented Aug 13, 2024

dime10 left a comment

erick-xanadu commented Aug 14, 2024

Add operations for setting state and basis state and integrate this with skip_initial_state_prep #955

Add operations for setting state and basis state and integrate this with skip_initial_state_prep #955

Conversation

rauletorresc commented Jul 23, 2024 • edited by erick-xanadu Loading

erick-xanadu commented Jul 30, 2024

paul0403 commented Jul 30, 2024

dime10 commented Jul 30, 2024

codecov bot commented Aug 2, 2024 • edited Loading

Codecov Report

erick-xanadu commented Aug 2, 2024

dime10 commented Aug 2, 2024 • edited Loading

erick-xanadu commented Aug 2, 2024

paul0403 commented Aug 3, 2024

erick-xanadu commented Aug 5, 2024

erick-xanadu commented Aug 13, 2024

dime10 commented Aug 13, 2024

dime10 commented Aug 13, 2024

dime10 left a comment

Choose a reason for hiding this comment

erick-xanadu commented Aug 14, 2024

Add operations for setting state and basis state and integrate this with `skip_initial_state_prep` #955

Add operations for setting state and basis state and integrate this with `skip_initial_state_prep` #955

rauletorresc commented Jul 23, 2024 •

edited by erick-xanadu

Loading

codecov bot commented Aug 2, 2024 •

edited

Loading

dime10 commented Aug 2, 2024 •

edited

Loading