(LK-C-4) Add controlled QubitUnitary support to Lightning Kokkos #955

josephleekl · 2024-10-21T18:21:47Z

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:

Description of the Change:
This PR adds support for controlled QubitUnitary for Lightning Kokkos. There are specialized implementations for 1/2/3-qubits (applyNC1/2/3QubitOpFunctor) and a general case implementation (NCMultiQubitOpFunctor). These functors are defined in MatrixGateFunctors.hpp. These are called by applyControlledMatrix or applyOperation/applyNCMultiQubitOp in StateVectorKokkos.hpp

Benefits:
Performance benchmarks for gates are shown here: https://www.notion.so/xanaduai/Lightning-Kokkos-Native-Controlled-Operation-Gate-Benchmarks-12ebc6bd17648017a2dcd237748b24fe

Possible Drawbacks:

Related GitHub Issues:

[sc-76776]

github-actions · 2024-10-21T18:22: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 · 2024-10-21T21:45:17Z

Codecov Report

Attention: Patch coverage is 99.32584% with 3 lines in your changes missing coverage. Please review.

Please upload report for BASE (lk-control-base@7ea9296). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
...nylane_lightning/lightning_kokkos/_state_vector.py	0.00%	3 Missing ⚠️

Additional details and impacted files

@@                Coverage Diff                 @@
##             lk-control-base     #955   +/-   ##
==================================================
  Coverage                   ?   97.05%           
==================================================
  Files                      ?      221           
  Lines                      ?    34895           
  Branches                   ?        0           
==================================================
  Hits                       ?    33866           
  Misses                     ?     1029           
  Partials                   ?        0

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

josephleekl · 2024-11-12T20:38:14Z

pennylane_lightning/core/src/simulators/lightning_kokkos/gates/MatrixGateFunctors.hpp

@@ -104,6 +111,71 @@ template <class Precision> struct multiQubitOpFunctor {
    }
 };

+template <class Precision> struct NCMultiQubitOpFunctor {


This construction is similar to multiQubitOpFunctor

josephleekl · 2024-11-12T20:38:48Z

pennylane_lightning/core/src/simulators/lightning_kokkos/gates/MatrixGateFunctors.hpp

@@ -142,6 +214,51 @@ template <class PrecisionT> struct apply1QubitOpFunctor {
        arr(i1) = matrix(0B10) * v0 + matrix(0B11) * v1;
    }
 };
+
+template <class PrecisionT> struct applyNC1QubitOpFunctor {


and this with apply1QubitOpFunctor etc.

multiphaseCFD

Thanks @josephleekl ! A couple of Q.

multiphaseCFD · 2024-11-14T21:19:31Z

pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp

+
+        auto &&num_qubits = this->getNumQubits();
+        std::size_t two2N =
+            std::exp2(num_qubits - wires.size() - controlled_wires.size());


What about using Pennylane::Util::exp2?

can we mark two2N as a const?

I believe that's the idea behind Util::exp2. Could you ensure that we use this exp2 instead of std::exp2. I think you used it in a few other places in your previous PRs.

Thanks, I'll update to Util::exp2

multiphaseCFD · 2024-11-14T21:25:59Z

pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp

+        KokkosVector matrix_trans("matrix_trans", matrix.size());
+
+        if (inverse) {
+            Kokkos::MDRangePolicy<DoubleLoopRank> policy_2d({0, 0}, {dim, dim});


do we want to use shmem to optimize the transpose here?

multiphaseCFD · 2024-11-14T21:26:48Z

pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp

+            break;
+        default:
+            std::size_t scratch_size = ScratchViewComplex::shmem_size(dim) +
+                                       ScratchViewSizeT::shmem_size(dim);


any concern on the shmem usage here?

maliasadi

Nice work @josephleekl! 🤩

pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp

maliasadi · 2024-11-15T00:02:21Z

pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp

+
+        auto &&num_qubits = this->getNumQubits();
+        std::size_t two2N =
+            std::exp2(num_qubits - wires.size() - controlled_wires.size());


I believe that's the idea behind Util::exp2. Could you ensure that we use this exp2 instead of std::exp2. I think you used it in a few other places in your previous PRs.

maliasadi · 2024-11-15T00:07:17Z

pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp

+            Kokkos::MDRangePolicy<DoubleLoopRank> policy_2d({0, 0}, {dim, dim});
+            Kokkos::parallel_for(
+                policy_2d,
+                KOKKOS_LAMBDA(const std::size_t i, const std::size_t j) {


We don't need to specify const for size_t! Could you ensure that this is fixes everywhere?

Why don't we need to specify const for size_t?

I think our cpp-dev-guide doesn't specify using const for primitive data type for parameter passed by values, but some guides like google C++ guide suggests not? I am happy to take it all out to keep it consistent?

maliasadi · 2024-11-15T00:09:36Z

pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp

+        const std::vector<std::size_t> &wires, bool inverse = false) {
+        PL_ABORT_IF(wires.empty(), "Number of wires must be larger than 0");
+        std::size_t n = static_cast<std::size_t>(1U) << wires.size();
+        KokkosVector matrix_(matrix, n * n);


Suggested change

KokkosVector matrix_(matrix, n * n);

KokkosVector matrix_(matrix, exp2(n, 2));

maliasadi · 2024-11-15T00:10:20Z

pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp

+                          bool inverse = false) {
+        PL_ABORT_IF(wires.empty(), "Number of wires must be larger than 0");
+        std::size_t n = static_cast<std::size_t>(1U) << wires.size();
+        std::size_t n2 = n * n;


Suggested change

std::size_t n2 = n * n;

std::size_t n2 {exp2(n, 2)};

maliasadi · 2024-11-15T00:23:46Z

pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp

+            PL_ABORT_IF(gate_matrix.empty(),
+                        std::string("Operation does not exist for ") + opName +
+                            std::string(" and no matrix provided."));


Where do you test this?

maliasadi · 2024-11-15T00:24:32Z

pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp

+        ComplexT *matrix, const std::vector<std::size_t> &controlled_wires,
+        const std::vector<bool> &controlled_values,
+        const std::vector<std::size_t> &wires, bool inverse = false) {
+        PL_ABORT_IF(wires.empty(), "Number of wires must be larger than 0");


Where have you tested this abort?

maliasadi · 2024-11-15T00:24:47Z

pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp

+        PL_ABORT_IF(wires.empty(), "Number of wires must be larger than 0");
+        PL_ABORT_IF(matrix.size() != exp2(2 * wires.size()),
+                    "The size of matrix does not match with the given "
+                    "number of wires");


Where have you tested these aborts?

maliasadi · 2024-11-15T00:28:14Z

pennylane_lightning/lightning_kokkos/_state_vector.py

+                control_wires,
+                control_values,
+                target_wires,
+                False,


Could you also add a code comment about this inv == False for the future reference? (this is a missing feature that we'll tackle later)

Suggested change

False,

inv,

maliasadi · 2024-11-15T00:28:38Z

pennylane_lightning/lightning_kokkos/lightning_kokkos.py

@@ -129,6 +129,7 @@
        "C(DoubleExcitationMinus)",
        "C(DoubleExcitationPlus)",
        "C(MultiRZ)",
+        "C(QubitUnitary)",


What about ControlledQubitUnitary?

It's already in L72

AmintorDusko

Some comments, and questions. Most of it I did some time ago and forgot to submit 😓.
I couldn't track down all comments so, please disregard any outdated comment. I may delete some after submitting.

AmintorDusko · 2024-11-06T19:06:42Z

pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp

+        case 1:
+            Kokkos::parallel_for(two2N, applyNC1QubitOpFunctor<fp_t>(
+                                            *data_, num_qubits, matrix_trans,
+                                            controlled_wires, controlled_values,
+                                            wires));
+            break;
+        case 2:
+            Kokkos::parallel_for(two2N, applyNC2QubitOpFunctor<fp_t>(
+                                            *data_, num_qubits, matrix_trans,
+                                            controlled_wires, controlled_values,
+                                            wires));
+            break;
+        case 3:
+            Kokkos::parallel_for(two2N, applyNC3QubitOpFunctor<fp_t>(
+                                            *data_, num_qubits, matrix_trans,
+                                            controlled_wires, controlled_values,
+                                            wires));
+            break;


Is there something we can do to reduce the code here? Maybe determine the apply function in a case-like structure and reuse the Kokkos::parallel_for call. I mean, because all functions have the same signature, right? No need to lose sleep about this though. 😅

AmintorDusko · 2024-11-06T19:10:09Z

pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp

+    inline void applyControlledMatrix(
+        ComplexT *matrix, const std::vector<std::size_t> &controlled_wires,
+        const std::vector<bool> &controlled_values,
+        const std::vector<std::size_t> &wires, bool inverse = false) {
+        PL_ABORT_IF(wires.empty(), "Number of wires must be larger than 0");
+        std::size_t n = static_cast<std::size_t>(1U) << wires.size();
+        KokkosVector matrix_(matrix, n * n);
+        applyNCMultiQubitOp(matrix_, controlled_wires, controlled_values, wires,
+                            inverse);
+    }
+
+    inline void
+    applyControlledMatrix(const ComplexT *matrix,
+                          const std::vector<std::size_t> &controlled_wires,
+                          const std::vector<bool> &controlled_values,
+                          const std::vector<std::size_t> &wires,
+                          bool inverse = false) {
+        PL_ABORT_IF(wires.empty(), "Number of wires must be larger than 0");
+        std::size_t n = static_cast<std::size_t>(1U) << wires.size();
+        std::size_t n2 = n * n;
+        KokkosVector matrix_("matrix_", n2);
+        Kokkos::deep_copy(matrix_, UnmanagedConstComplexHostView(matrix, n2));
+        applyNCMultiQubitOp(matrix_, controlled_wires, controlled_values, wires,
+                            inverse);
+    }


In the constant matrix overload, you can probably copy the constant matrix to a matrix_ and call the non-constant matrix overload. Reusing all logic. What are your thoughts?

pennylane_lightning/core/src/simulators/lightning_kokkos/gates/MatrixGateFunctors.hpp

AmintorDusko · 2024-11-06T19:26:40Z

pennylane_lightning/core/src/simulators/lightning_kokkos/gates/MatrixGateFunctors.hpp

+            for (std::size_t i = 0; i < dim; i++) {
+                coeffs_in(i) = arr(indices(i) + offset);
+            }


Suggested change

for (std::size_t i = 0; i < dim; i++) {

coeffs_in(i) = arr(indices(i) + offset);

}

std::transform(indices.begin(), indices.end(), coeffs_in.begin(), [&](std::size_t index) { return arr(index + offset)};

right? (please check)

I don't remember if we can do this in a KOKKOS_INLINE_FUNCTION.

pennylane_lightning/core/src/simulators/lightning_kokkos/gates/MatrixGateFunctors.hpp

AmintorDusko · 2024-11-15T13:04:47Z

pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp

+            Kokkos::MDRangePolicy<DoubleLoopRank> policy_2d({0, 0}, {dim, dim});
+            Kokkos::parallel_for(
+                policy_2d,
+                KOKKOS_LAMBDA(const std::size_t i, const std::size_t j) {


Why don't we need to specify const for size_t?

josephleekl added 12 commits October 18, 2024 19:34

initial commit

4b9062e

initial commit

ba70971

support controlled 1 qubit gate

6ac216a

add controlled globalphase test

bea8bf7

fix controlled globalphase return

38c48b9

format

968b86e

format

8c97d82

format

09b07d5

add globalphase python test

a7106b7

apply controlled only for 1 qubit gate

34e6400

format

b1263fa

trigger CI

1245270

josephleekl changed the title ~~Lk control matrix nq~~ Add controlled N-qubit matrix support to Lightning Kokkos Oct 21, 2024

fix parity_2_offset to kokkos_inline_function

7ebee0f

josephleekl force-pushed the lk-control-gate-NQ-multiRZ branch 2 times, most recently from 42f9cd3 to 23f398c Compare October 21, 2024 21:01

josephleekl force-pushed the lk-control-matrix-NQ branch from 8807e39 to 07c38a1 Compare October 21, 2024 21:02

josephleekl added ci:build_wheels Activate wheel building. ci:use-gpu-runner Enable usage of GPU runner for this Pull Request urgent Mark a pull request as high priority labels Oct 21, 2024

josephleekl marked this pull request as ready for review October 21, 2024 21:40

josephleekl added 7 commits October 22, 2024 11:41

separate generate/control bit patterns

cec599b

format

a01b38a

initial commit

8dec286

support 2/3/4 qubit control gates

17ef49b

update statevectorkokkos error

e3911c3

update statevectorkokkos error

a136669

test control=wire exception

bead29b

josephleekl commented Nov 12, 2024

View reviewed changes

update toml

b68dd52

josephleekl requested review from maliasadi, multiphaseCFD, AmintorDusko and LuisAlfredoNu and removed request for maliasadi and multiphaseCFD November 12, 2024 22:30

josephleekl added 12 commits November 13, 2024 16:57

update test name

21c4a8c

add contiguous control test

0a62ea2

add test

63114fd

format test

e93965b

remove scratch for NCMultiRZ

3cef4e8

remove scratch using

e8ba058

restore measure test

8e4ff26

format

bd8b70a

Merge branch 'lk-control-gate-NQ-multiRZ' into lk-control-matrix-NQ

224cfa1

update scratch indices in NCmultiqubitop

e37bdee

add TODO for using shmem for indices in multirz

cf1d66e

Merge branch 'lk-control-gate-NQ-multiRZ' into lk-control-matrix-NQ

2a10925

Base automatically changed from lk-control-gate-NQ-multiRZ to lk-control-base November 13, 2024 21:36

josephleekl added 3 commits November 13, 2024 21:45

Merge branch 'lk-control-base' into lk-control-matrix-NQ

e8653e5

add tests

42725dc

Merge branch 'lk-control-base' into lk-control-matrix-NQ

a073877

multiphaseCFD reviewed Nov 14, 2024

View reviewed changes

maliasadi requested changes Nov 15, 2024

View reviewed changes

AmintorDusko reviewed Nov 15, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

(LK-C-4) Add controlled QubitUnitary support to Lightning Kokkos #955

(LK-C-4) Add controlled QubitUnitary support to Lightning Kokkos #955

josephleekl commented Oct 21, 2024 •

edited

Loading

github-actions bot commented Oct 21, 2024

codecov bot commented Oct 21, 2024 •

edited

Loading

josephleekl Nov 12, 2024

josephleekl Nov 12, 2024

multiphaseCFD left a comment

multiphaseCFD Nov 14, 2024

multiphaseCFD Nov 14, 2024

maliasadi Nov 15, 2024

josephleekl Nov 15, 2024

multiphaseCFD Nov 14, 2024

multiphaseCFD Nov 14, 2024

maliasadi left a comment

maliasadi Nov 15, 2024

maliasadi Nov 15, 2024

AmintorDusko Nov 15, 2024

josephleekl Nov 15, 2024

maliasadi Nov 15, 2024

maliasadi Nov 15, 2024

maliasadi Nov 15, 2024

maliasadi Nov 15, 2024

maliasadi Nov 15, 2024

maliasadi Nov 15, 2024

maliasadi Nov 15, 2024

josephleekl Nov 15, 2024

AmintorDusko left a comment

AmintorDusko Nov 6, 2024

AmintorDusko Nov 6, 2024

AmintorDusko Nov 6, 2024

AmintorDusko Nov 6, 2024

AmintorDusko Nov 15, 2024

	KokkosVector matrix_(matrix, n * n);
	KokkosVector matrix_(matrix, exp2(n, 2));

(LK-C-4) Add controlled QubitUnitary support to Lightning Kokkos #955

Are you sure you want to change the base?

(LK-C-4) Add controlled QubitUnitary support to Lightning Kokkos #955

Conversation

josephleekl commented Oct 21, 2024 • edited Loading

Before submitting

github-actions bot commented Oct 21, 2024

codecov bot commented Oct 21, 2024 • edited Loading

Codecov Report

Choose a reason for hiding this comment

Choose a reason for hiding this comment

multiphaseCFD 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

maliasadi 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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

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

josephleekl commented Oct 21, 2024 •

edited

Loading

codecov bot commented Oct 21, 2024 •

edited

Loading