Accel/expval

[vincentmr]

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:
In the LKokkos backend, there are two ways to implement expectation values:

1. Copy the statevector, apply the observable to the copy, compute the expectation value with a BLAS-like inner product.
2. Accumulate the expectation value on-the-fly, applying the observable to a portion of the statevector.

The first method is currently in use, but is wasteful in a couple ways:

• The statevector copy can require a lot of memory, which means running out with one or two fewer qubits depending on the system.
• Computational efficiency is asymptotically memory-bound and reduce operations can achieve roughly twice the bandwidth of for loop operations on Nvidia devices, for example, which favors on-the-fly implementations.

A simple benchmarking script like the one below yields, for 29 qubits (on Perlmutter's A100)

1. 0.330 sec.
2. 0.012 sec.

i.e. a 27.5x speed-up.

import pennylane as qml
import time

n_wires = 29
n_repeat = 100

dev = qml.device("lightning.kokkos", wires=n_wires)

@qml.qnode(dev)
def circuit():
    return qml.expval(qml.PauliZ(n_wires//2))

t0 = time.time()
for _ in range(n_repeat):
    circuit()
dt = (time.time() - t0)/n_repeat
print(f"{dt}")

Description of the Change:

Benefits:

Possible Drawbacks:

Related GitHub Issues:

[codecov]

Codecov Report

❗ No coverage uploaded for pull request base (bugfix/cuda12@6dc7883). Click here to learn what that means.
The diff coverage is n/a.

@@               Coverage Diff                @@
##             bugfix/cuda12     #481   +/-   ##
================================================
  Coverage                 ?   96.84%           
================================================
  Files                    ?      142           
  Lines                    ?    16275           
  Branches                 ?        0           
================================================
  Hits                     ?    15761           
  Misses                   ?      514           
  Partials                 ?        0           

[AmintorDusko]

Amazing work. Do we know if these changes will allow us to reach a larger number of qubits?
If yes, what are the numbers we have?

[vincentmr]

Amazing work. Do we know if these changes will allow us to reach a larger number of qubits? If yes, what are the numbers we have?

From what I could see on Perlmutter, we can reach 30 instead of 29 on a A100 card. The main thing is that 1- and 2-qubit expvals are much faster.

[vincentmr]

In support of this PR, I made a gist that benchmarks expval on a CPU (OMP_NUM_THREADS=64) and a GPU (A100). The results are generally faster, especially the 1- and 2-qubit operators above 20 qubits. In the following figures, inner and team stand for the inner-product-based and TeamPolicy (this PR) algorithms respectively; first stands for targeting low-index qubits (only first shown since timings are similar across first, mid and last targets); 1, 2, 3 stand for the number of wires targeted by the Hermitian unitary. The timings for 1-qubit Hermitian unitaries are similar to that of named gates (e.g. PauliZ) with an equal number of wires.

[vincentmr]

Bonus with LKokkos benchmarks run on LUMI's AMD cards

[AmintorDusko]

Nothing more to ask. You did a great job here!

[mlxd]

💯

.

[mlxd]

Happy with the revision. Thanks @vincentmr