Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Optimize LM controlled kernels [sc-73461] #882

Merged
merged 73 commits into from
Sep 11, 2024
Merged

Optimize LM controlled kernels [sc-73461] #882

merged 73 commits into from
Sep 11, 2024

Conversation

vincentmr
Copy link
Contributor

@vincentmr vincentmr commented Aug 30, 2024

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:
As a first step toward adding controls in Lightning Kokkos kernels, the controlled Lightning Qubit kernels are simplified.

Description of the Change:

  • Introduce controlBitPatterns, the controlled version of generateBitPatterns, and remove obsolete parity2indices implementations.
  • Avoid temporary arrays/vectors as much as possible.
  • Change core_function signature from coefficients & indices to indices & offset.

Benefits:

  • Combine and eliminate a few branches and loops.
  • indices are now precomputed (thereby saving time) and only offset needs be updated on the fly.
  • All omp parallel for loops are now free of private arguments.

We illustrate the performance improvement running the XAS workflow from the benchmark suite. This workflow has a circuit with the following specs

{'resources': Resources(num_wires=11, num_gates=17382, gate_types=defaultdict(<class 'int'>, {'StatePrep': 1, 'Hadamard': 1, 'PhaseShift': 4380, 'SingleExcitation': 4000, 'C(MultiRZ)': 9000}), gate_sizes=defaultdict(<class 'int'>, {10: 1, 1: 4381, 2: 4000, 3: 9000}), depth=9117, shots=Shots(total_shots=None, shot_vector=())), 'errors': {}, 'num_observables': 2, 'num_diagonalizing_gates': 4, 'num_trainable_params': 17381, 'num_device_wires': 11, 'num_tape_wires': 11, 'device_name': 'lightning.qubit', 'level': 'device', 'gradient_options': {}, 'interface': 'auto', 'diff_method': 'best', 'gradient_fn': 'adjoint'}

and a bottleneck is applying the 9000 C(MultiRZ) gates. With v0.38.0 we get (zooming on simulate_and_jacobian with SnakeViz)
Screenshot from 2024-09-11 16-08-09
and for the current PR
Screenshot from 2024-09-11 16-08-18
We get a 6.6x speed-up on the C(MultiRZ) gates.

Possible Drawbacks:

Related GitHub Issues:

@vincentmr vincentmr changed the title Lq new control Optimize LM controlled kernels Aug 30, 2024
Copy link
Contributor

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.

Base automatically changed from lk_pauli_rot to master September 10, 2024 19:26
@vincentmr vincentmr changed the title Optimize LM controlled kernels Optimize LM controlled kernels [sc-73461] Sep 10, 2024
@vincentmr vincentmr marked this pull request as ready for review September 10, 2024 19:44
@vincentmr vincentmr added the ci:build_wheels Activate wheel building. label Sep 10, 2024
Copy link

codecov bot commented Sep 10, 2024

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 94.52%. Comparing base (bbb3eb4) to head (3751426).
Report is 1 commits behind head on master.

Additional details and impacted files
@@            Coverage Diff             @@
##           master     #882      +/-   ##
==========================================
- Coverage   96.23%   94.52%   -1.72%     
==========================================
  Files         212      168      -44     
  Lines       28318    21134    -7184     
==========================================
- Hits        27252    19977    -7275     
- Misses       1066     1157      +91     

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

@vincentmr vincentmr requested a review from a team September 10, 2024 19:59
@maliasadi maliasadi added the urgent Mark a pull request as high priority label Sep 10, 2024
Copy link
Member

@maliasadi maliasadi left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @vincentmr! Just a few minor comments but happy to approve after resolving them.

Copy link
Contributor

@LuisAlfredoNu LuisAlfredoNu left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice work @vincentmr. I have a few comments and questions. Do you have some quick data about the performance of this implementation? 🚀

Copy link
Member

@multiphaseCFD multiphaseCFD left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @vincentmr for the nice work! LGTM!

Copy link
Member

@multiphaseCFD multiphaseCFD left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM! Thanks for the nice work, @vincentmr !

Copy link
Contributor

@LuisAlfredoNu LuisAlfredoNu left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice optimization and good performance improvement. Thank you @vincentmr

@vincentmr vincentmr merged commit ef3a8cc into master Sep 11, 2024
95 of 96 checks passed
@vincentmr vincentmr deleted the lq_new_control branch September 11, 2024 21:03
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
ci:build_wheels Activate wheel building. urgent Mark a pull request as high priority
Projects
None yet
Development

Successfully merging this pull request may close these issues.

5 participants