Fix hipcc + lightning_kokkos issue (#1018)

### 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`](../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`. - [X] Add a new entry to the `.github/CHANGELOG.md` file, summarizing the change, and including a link back to the PR. - [X] 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:** This PR fixes an issue with the hipcc (llvm) compiler, which does not allow `[[maybe_unused]]` attributes in kokkos lambdas, and fails to correctly resolve `pennylane::utils::exp2` function. E.g. in nightly rocm docker build https://github.com/PennyLaneAI/pennylane-lightning/actions/runs/12173013064 **Description of the Change:** In lightning kokkos, `exp2` are explicitly replaced by `pennylane::utils::exp2`, and `[[maybe_unused]]` attributes are now moved from the lambda arguments to the lambda body. This has been tested to work with the rocm docker build. (https://github.com/PennyLaneAI/pennylane-lightning/actions/runs/12188525751/job/34001703721) **Benefits:** Fixes rocm build failures **Possible Drawbacks:** **Related GitHub Issues:** [sc-79824] --------- Co-authored-by: ringo-but-quantum <github-ringo-but-quantum@xanadu.ai>
PennyLaneAI · Dec 10, 2024 · e6827b2 · e6827b2
1 parent 9eb6123
commit e6827b2
Show file tree

Hide file tree

Showing 5 changed files with 59 additions and 43 deletions.
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
@@ -83,6 +83,9 @@
 
 ### Bug fixes
 
+* Fix Lightning Kokkos `[[maybe_unused]]` and `exp2` errors with hipcc.
+  [(#1018)](https://github.com/PennyLaneAI/pennylane-lightning/pull/1018)
+
 * Pin `jax[cpu]==0.4.28` for compatibility with PennyLane and Catalyst.
   [(#1019)](https://github.com/PennyLaneAI/pennylane-lightning/pull/1019)
 

diff --git a/pennylane_lightning/core/_version.py b/pennylane_lightning/core/_version.py
@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.40.0-dev32"
+__version__ = "0.40.0-dev33"
diff --git a/pennylane_lightning/core/src/simulators/lightning_kokkos/gates/BasicGateFunctors.hpp b/pennylane_lightning/core/src/simulators/lightning_kokkos/gates/BasicGateFunctors.hpp
@@ -58,9 +58,9 @@ template <class PrecisionT, class FuncT> class applyNCNFunctor {
                     const std::vector<bool> &controlled_values,
                     const std::vector<std::size_t> &wires, FuncT core_function_)
         : arr(arr_), core_function(core_function_) {
-        const std::size_t two2N =
-            exp2(num_qubits - wires.size() - controlled_wires.size());
-        dim = exp2(wires.size());
+        const std::size_t two2N = Pennylane::Util::exp2(
+            num_qubits - wires.size() - controlled_wires.size());
+        dim = Pennylane::Util::exp2(wires.size());
         const auto &[parity_, rev_wires_] =
             reverseWires(num_qubits, wires, controlled_wires);
         parity = parity_;
@@ -115,7 +115,8 @@ class applyNC1Functor<PrecisionT, FuncT, true> {
         indices = vector2view(indices_);
         Kokkos::parallel_for(
             Kokkos::RangePolicy<ExecutionSpace>(
-                0, exp2(num_qubits - controlled_wires.size() - wires.size())),
+                0, Pennylane::Util::exp2(num_qubits - controlled_wires.size() -
+                                         wires.size())),
             *this);
     }
     KOKKOS_FUNCTION void operator()(std::size_t k) const {
@@ -147,9 +148,9 @@ class applyNC1Functor<PrecisionT, FuncT, false> {
           rev_wire_shift((static_cast<std::size_t>(1U) << rev_wire)),
           wire_parity(fillTrailingOnes(rev_wire)),
           wire_parity_inv(fillLeadingOnes(rev_wire + 1)) {
-        Kokkos::parallel_for(
-            Kokkos::RangePolicy<ExecutionSpace>(0, exp2(num_qubits - 1)),
-            *this);
+        Kokkos::parallel_for(Kokkos::RangePolicy<ExecutionSpace>(
+                                 0, Pennylane::Util::exp2(num_qubits - 1)),
+                             *this);
     }
     KOKKOS_FUNCTION void operator()(std::size_t k) const {
         std::size_t i0 = ((k << 1U) & wire_parity_inv) | (wire_parity & k);
@@ -644,7 +645,8 @@ class applyNC2Functor<PrecisionT, FuncT, true> {
         indices = vector2view(indices_);
         Kokkos::parallel_for(
             Kokkos::RangePolicy<ExecutionSpace>(
-                0, exp2(num_qubits - controlled_wires.size() - wires.size())),
+                0, Pennylane::Util::exp2(num_qubits - controlled_wires.size() -
+                                         wires.size())),
             *this);
     }
     KOKKOS_FUNCTION void operator()(std::size_t k) const {
@@ -692,9 +694,9 @@ class applyNC2Functor<PrecisionT, FuncT, false> {
           parity_high(fillLeadingOnes(rev_wire_max + 1)),
           parity_middle(fillLeadingOnes(rev_wire_min + 1) &
                         fillTrailingOnes(rev_wire_max)) {
-        Kokkos::parallel_for(
-            Kokkos::RangePolicy<ExecutionSpace>(0, exp2(num_qubits - 2)),
-            *this);
+        Kokkos::parallel_for(Kokkos::RangePolicy<ExecutionSpace>(
+                                 0, Pennylane::Util::exp2(num_qubits - 2)),
+                             *this);
     }
     KOKKOS_FUNCTION void operator()(std::size_t k) const {
         std::size_t i00 = ((k << 2U) & parity_high) |
@@ -1267,9 +1269,9 @@ template <class PrecisionT, class FuncT> class applyNC3Functor {
             fillLeadingOnes(rev_wire_min + 1) & fillTrailingOnes(rev_wire_mid);
         parity_hmiddle =
             fillLeadingOnes(rev_wire_mid + 1) & fillTrailingOnes(rev_wire_max);
-        Kokkos::parallel_for(
-            Kokkos::RangePolicy<ExecutionSpace>(0, exp2(num_qubits - 3)),
-            *this);
+        Kokkos::parallel_for(Kokkos::RangePolicy<ExecutionSpace>(
+                                 0, Pennylane::Util::exp2(num_qubits - 3)),
+                             *this);
     }
     KOKKOS_FUNCTION void operator()(std::size_t k) const {
         std::size_t i000 = ((k << 3U) & parity_high) |
@@ -1364,7 +1366,8 @@ class applyNC4Functor<PrecisionT, FuncT, true> {
         indices = vector2view(indices_);
         Kokkos::parallel_for(
             Kokkos::RangePolicy<ExecutionSpace>(
-                0, exp2(num_qubits - controlled_wires.size() - wires.size())),
+                0, Pennylane::Util::exp2(num_qubits - controlled_wires.size() -
+                                         wires.size())),
             *this);
     }
     KOKKOS_FUNCTION void operator()(std::size_t k) const {
@@ -1473,9 +1476,9 @@ class applyNC4Functor<PrecisionT, FuncT, false> {
                          fillTrailingOnes(rev_wire_max);
         parity_middle = fillLeadingOnes(rev_wire_min_mid + 1) &
                         fillTrailingOnes(rev_wire_max_mid);
-        Kokkos::parallel_for(
-            Kokkos::RangePolicy<ExecutionSpace>(0, exp2(num_qubits - 4)),
-            *this);
+        Kokkos::parallel_for(Kokkos::RangePolicy<ExecutionSpace>(
+                                 0, Pennylane::Util::exp2(num_qubits - 4)),
+                             *this);
     }
     KOKKOS_FUNCTION void operator()(std::size_t k) const {
         std::size_t i0000 = ((k << 4U) & parity_high) |
@@ -1699,7 +1702,8 @@ void applyMultiRZ(Kokkos::View<Kokkos::complex<PrecisionT> *> arr_,
             (static_cast<std::size_t>(1U) << (num_qubits - wire - 1));
     }
     Kokkos::parallel_for(
-        Kokkos::RangePolicy<ExecutionSpace>(0, exp2(num_qubits)),
+        Kokkos::RangePolicy<ExecutionSpace>(0,
+                                            Pennylane::Util::exp2(num_qubits)),
         KOKKOS_LAMBDA(std::size_t k) {
             arr_(k) *= (Kokkos::Impl::bit_count(k & wires_parity) % 2 == 0)
                            ? shift_0
@@ -1777,7 +1781,8 @@ void applyPauliRot(Kokkos::View<Kokkos::complex<PrecisionT> *> arr_,
         get_mask([&word](const int a) { return word[a] == 'Z'; });
     const auto count_mask_y = std::popcount(mask_y);
     Kokkos::parallel_for(
-        Kokkos::RangePolicy<ExecutionSpace>(0, exp2(num_qubits)),
+        Kokkos::RangePolicy<ExecutionSpace>(0,
+                                            Pennylane::Util::exp2(num_qubits)),
         KOKKOS_LAMBDA(std::size_t i0) {
             std::size_t i1 = i0 ^ mask_xy;
             if (i0 <= i1) {

diff --git a/pennylane_lightning/core/src/simulators/lightning_kokkos/gates/BasicGeneratorFunctors.hpp b/pennylane_lightning/core/src/simulators/lightning_kokkos/gates/BasicGeneratorFunctors.hpp
@@ -346,10 +346,11 @@ void applyGenMultiRZ(Kokkos::View<Kokkos::complex<PrecisionT> *> arr_,
                      [[maybe_unused]] bool inverse = false) {
     std::size_t wires_parity = static_cast<std::size_t>(0U);
     for (std::size_t wire : wires) {
-        wires_parity |= exp2(num_qubits - wire - 1);
+        wires_parity |= Pennylane::Util::exp2(num_qubits - wire - 1);
     }
     Kokkos::parallel_for(
-        Kokkos::RangePolicy<ExecutionSpace>(0, exp2(num_qubits)),
+        Kokkos::RangePolicy<ExecutionSpace>(0,
+                                            Pennylane::Util::exp2(num_qubits)),
         KOKKOS_LAMBDA(std::size_t k) {
             arr_(k) *= static_cast<PrecisionT>(
                 1 - 2 * int(Kokkos::Impl::bit_count(k & wires_parity) % 2));
@@ -411,7 +412,8 @@ template <class PrecisionT, class FuncT> class applyNCGenerator1Functor {
         indices = vector2view(indices_);
         Kokkos::parallel_for(
             Kokkos::RangePolicy<ExecutionSpace>(
-                0, exp2(num_qubits - controlled_wires.size() - wires.size())),
+                0, Pennylane::Util::exp2(num_qubits - controlled_wires.size() -
+                                         wires.size())),
             *this);
     }
     KOKKOS_FUNCTION void operator()(std::size_t k) const {
@@ -506,9 +508,10 @@ void applyNCGenGlobalPhase(
     const std::vector<bool> &controlled_values,
     [[maybe_unused]] const std::vector<std::size_t> &wires,
     [[maybe_unused]] bool inverse = false) {
-    auto core_function = KOKKOS_LAMBDA(
-        [[maybe_unused]] Kokkos::View<Kokkos::complex<PrecisionT> *> arr,
-        std::size_t i0, std::size_t i1) {
+    auto core_function =
+        KOKKOS_LAMBDA(Kokkos::View<Kokkos::complex<PrecisionT> *> arr,
+                      std::size_t i0, std::size_t i1) {
+        [[maybe_unused]] const auto &arr_ = arr;
         [[maybe_unused]] const auto i0_ = i0;
         [[maybe_unused]] const auto i1_ = i1;
     };
@@ -586,7 +589,8 @@ template <class PrecisionT, class FuncT> class applyNCGenerator2Functor {
         indices = vector2view(indices_);
         Kokkos::parallel_for(
             Kokkos::RangePolicy<ExecutionSpace>(
-                0, exp2(num_qubits - controlled_wires.size() - wires.size())),
+                0, Pennylane::Util::exp2(num_qubits - controlled_wires.size() -
+                                         wires.size())),
             *this);
     }
     KOKKOS_FUNCTION void operator()(std::size_t k) const {
@@ -800,7 +804,8 @@ template <class PrecisionT, class FuncT> class applyNCGenerator4Functor {
         indices = vector2view(indices_);
         Kokkos::parallel_for(
             Kokkos::RangePolicy<ExecutionSpace>(
-                0, exp2(num_qubits - controlled_wires.size() - wires.size())),
+                0, Pennylane::Util::exp2(num_qubits - controlled_wires.size() -
+                                         wires.size())),
             *this);
     }
     KOKKOS_FUNCTION void operator()(std::size_t k) const {
@@ -847,8 +852,9 @@ void applyNCGenDoubleExcitationMinus(
     [[maybe_unused]] bool inverse = false) {
     auto core_function = KOKKOS_LAMBDA(
         Kokkos::View<Kokkos::complex<PrecisionT> *> arr, std::size_t i0011,
-        std::size_t i1100, [[maybe_unused]] const KokkosIntVector &indices,
-        [[maybe_unused]] std::size_t offset) {
+        std::size_t i1100, const KokkosIntVector &indices, std::size_t offset) {
+        [[maybe_unused]] const auto &indices_ = indices;
+        [[maybe_unused]] const auto offset_ = offset;
         const auto v0011 = arr(i0011);
         const auto v1100 = arr(i1100);
         arr(i0011) = Kokkos::complex<PrecisionT>{imag(v1100), -real(v1100)};
@@ -868,8 +874,9 @@ void applyNCGenDoubleExcitationPlus(
     [[maybe_unused]] bool inverse = false) {
     auto core_function = KOKKOS_LAMBDA(
         Kokkos::View<Kokkos::complex<PrecisionT> *> arr, std::size_t i0011,
-        std::size_t i1100, [[maybe_unused]] const KokkosIntVector &indices,
-        [[maybe_unused]] std::size_t offset) {
+        std::size_t i1100, const KokkosIntVector &indices, std::size_t offset) {
+        [[maybe_unused]] const auto &indices_ = indices;
+        [[maybe_unused]] const auto offset_ = offset;
         const auto v0011 = arr(i0011);
         const auto v1100 = arr(i1100);
         arr(i0011) = Kokkos::complex<PrecisionT>{-imag(v1100), real(v1100)};
@@ -895,16 +902,17 @@ void applyNCGenMultiRZ(Kokkos::View<Kokkos::complex<PrecisionT> *> arr_,
     std::size_t ctrls_parity = std::accumulate(
         controlled_wires.begin(), controlled_wires.end(), std::size_t{0},
         [num_qubits](std::size_t acc, std::size_t wire) {
-            return acc | exp2(num_qubits - wire - 1);
+            return acc | Pennylane::Util::exp2(num_qubits - wire - 1);
+        });
+    std::size_t wires_parity = std::accumulate(
+        wires.begin(), wires.end(), std::size_t{0},
+        [num_qubits](std::size_t acc, std::size_t wire) {
+            return acc | Pennylane::Util::exp2(num_qubits - wire - 1);
         });
-    std::size_t wires_parity =
-        std::accumulate(wires.begin(), wires.end(), std::size_t{0},
-                        [num_qubits](std::size_t acc, std::size_t wire) {
-                            return acc | exp2(num_qubits - wire - 1);
-                        });
 
     Kokkos::parallel_for(
-        Kokkos::RangePolicy<ExecutionSpace>(0, exp2(num_qubits)),
+        Kokkos::RangePolicy<ExecutionSpace>(0,
+                                            Pennylane::Util::exp2(num_qubits)),
         KOKKOS_LAMBDA(std::size_t k) {
             if (ctrls_mask == (ctrls_parity & k)) {
                 arr_(k) *= static_cast<PrecisionT>(

diff --git a/pennylane_lightning/core/src/simulators/lightning_kokkos/gates/MatrixGateFunctors.hpp b/pennylane_lightning/core/src/simulators/lightning_kokkos/gates/MatrixGateFunctors.hpp
@@ -60,7 +60,7 @@ template <class Precision> struct multiQubitOpFunctor {
                         const KokkosComplexVector &matrix_,
                         const std::vector<std::size_t> &wires_) {
         wires = vector2view(wires_);
-        dim = exp2(wires_.size());
+        dim = Pennylane::Util::exp2(wires_.size());
         num_qubits = num_qubits_;
         arr = arr_;
         matrix = matrix_;
@@ -137,7 +137,7 @@ template <class Precision> struct NCMultiQubitOpFunctor {
                           const std::vector<std::size_t> &controlled_wires_,
                           const std::vector<bool> &controlled_values_,
                           const std::vector<std::size_t> &wires_) {
-        dim = exp2(wires_.size());
+        dim = Pennylane::Util::exp2(wires_.size());
         arr = arr_;
         matrix = matrix_;
         num_qubits = num_qubits_;
@@ -202,7 +202,7 @@ template <class PrecisionT> struct apply1QubitOpFunctor {
         num_qubits = num_qubits_;
 
         rev_wire = num_qubits - wires_[0] - 1;
-        rev_wire_shift = exp2(rev_wire);
+        rev_wire_shift = Pennylane::Util::exp2(rev_wire);
         wire_parity = fillTrailingOnes(rev_wire);
         wire_parity_inv = fillLeadingOnes(rev_wire + 1);
     }