Add `sample(wires)` support in LightningQubit #809

vincentmr · 2024-07-19T19:49:05Z

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:
sample call generate_samples which computes the full probabilities and uses the alias method to generate samples for all wires. This is wasteful whenever samples are required only for a subset of all wires.

Description of the Change:
Move alias method logic to the discrete_random_variable class.

Benefits:
Compute minimal probs and samples.
We benchmark the current changes against master, which already benefits from some good speed-ups introduce in #795 and #800 .

We use ISAIC's AMD EPYC-Milan Processor using a single core/thread. The times are obtained using at least 5 experiments and running for at least 250 milliseconds. We begin comparing master's generate_samples(num_samples) with ours generate_samples({0}, num_samples). For 4-12 qubits, overheads dominate the calculation (the absolute times range from 6 microseconds to 18 milliseconds, which is not a lot. Already at 12 qubits however, a trend appears where our implementation is significantly faster. This is to be expected for two reason: probs(wires) itself is faster than probs() (for enough qubits) and sample(wires) starts also requiring significantly less work than sample().

Next we turn to comparing master's generate_samples(num_samples) with ours generate_samples({0..num_qubits/2}, num_samples). The situation there is similar, with speed-ups close to 1 for the smaller qubit counts and (sometimes) beyond 20x for qubit counts above 20.

Finally we compare master's generate_samples(num_samples) with ours generate_samples({0..num_qubits-1}, num_samples) (i.e. computing samples on all wires). We expect similar performance since the main difference comes from the caching mechanism in master's discrete random variable generator. The data suggests caching samples is counter-productive compared with calculating the sample values on the fly.

Turning OMP ON, using 16 threads and comparing master's generate_samples(num_samples) with ours generate_samples({0}, num_samples) we get good speed-ups above 12 qubits. Below that the overhead of spawning threads isn't repaid, but absolute times remain low.

Possible Drawbacks:

Related GitHub Issues:
[sc-65127]

…wires, num_shots).

codecov · 2024-07-19T19:52:01Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 98.12%. Comparing base (9a83ff4) to head (3a2a341).
Report is 92 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #809      +/-   ##
==========================================
- Coverage   98.65%   98.12%   -0.53%     
==========================================
  Files         114       73      -41     
  Lines       17779    11172    -6607     
==========================================
- Hits        17540    10963    -6577     
+ Misses        239      209      -30

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

maliasadi

Nice work @vincentmr! 🥳 Just a few suggestions to consider before merging the PR.

pennylane_lightning/core/src/simulators/lightning_qubit/measurements/MeasurementKernels.hpp