Accel/expval (#481)

* Introduce std::unordered_map<std::string, ExpValFunc> expval_funcs_. * Introduce applyExpectationValueFunctor. * Add binding to LKokkos expval(matrix, wires). Combine expval functor calls into two templated methods. Call specialized expval methods when possible. Remove obsolete 'Apply directly' tests. * Update changelog. * Add test for arbitrary expval(Hermitian). * Add getExpectationValueMultiQubitOpFunctor. * Add typename hint for macos. * Add typename macos. * Use Kokkos::ThreadVectorRange policy for innerloop in getExpectationValueMultiQubitOpFunctor. * Couple fix for HIP. * Use inner product scheme instead of getExpectationValueMultiQubitOpFunctor to compute multi-qubit expval.
PennyLaneAI · Aug 25, 2023 · fb82dea · fb82dea
1 parent 6dc7883
commit fb82dea
Show file tree

Hide file tree

Showing 9 changed files with 210 additions and 348 deletions.
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
@@ -21,6 +21,9 @@
 
 ### Improvements
 
+* Refactor LKokkos `Measurements` class to use (fast) specialized functors whenever possible.
+  [(#481)] (https://github.com/PennyLaneAI/pennylane-lightning/pull/481)
+
 * Merge Lightning Qubit and Lightning Kokkos backends in the new repository.
   [(#472)] (https://github.com/PennyLaneAI/pennylane-lightning/pull/472)
 

diff --git a/pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp b/pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkos.hpp
@@ -65,7 +65,6 @@ class StateVectorKokkos final
     using HostExecSpace = Kokkos::DefaultHostExecutionSpace;
     using KokkosVector = Kokkos::View<ComplexT *>;
     using KokkosSizeTVector = Kokkos::View<size_t *>;
-    using KokkosRangePolicy = Kokkos::RangePolicy<KokkosExecSpace>;
     using UnmanagedComplexHostView =
         Kokkos::View<ComplexT *, Kokkos::HostSpace,
                      Kokkos::MemoryTraits<Kokkos::Unmanaged>>;

diff --git a/pennylane_lightning/core/src/simulators/lightning_kokkos/bindings/LKokkosBindings.hpp b/pennylane_lightning/core/src/simulators/lightning_kokkos/bindings/LKokkosBindings.hpp
@@ -26,6 +26,7 @@
 #include "MeasurementsKokkos.hpp"
 #include "StateVectorKokkos.hpp"
 #include "TypeList.hpp"
+#include "Util.hpp" // exp2
 
 /// @cond DEV
 namespace {
@@ -34,6 +35,7 @@ using namespace Pennylane::LightningKokkos::Algorithms;
 using namespace Pennylane::LightningKokkos::Measures;
 using Kokkos::InitializationSettings;
 using Pennylane::LightningKokkos::StateVectorKokkos;
+using Pennylane::Util::exp2;
 } // namespace
 /// @endcond
 
@@ -164,6 +166,18 @@ void registerBackendSpecificMeasurements(PyClass &pyclass) {
                  const std::string &, const std::vector<size_t> &)>(
                  &Measurements<StateVectorT>::expval),
              "Expected value of an operation by name.")
+        .def(
+            "expval",
+            [](Measurements<StateVectorT> &M, const np_arr_c &matrix,
+               const std::vector<size_t> &wires) {
+                const std::size_t matrix_size = exp2(2 * wires.size());
+                auto matrix_data =
+                    static_cast<ComplexT *>(matrix.request().ptr);
+                std::vector<ComplexT> matrix_v{matrix_data,
+                                               matrix_data + matrix_size};
+                return M.expval(matrix_v, wires);
+            },
+            "Expected value of a Hermitian observable.")
         .def(
             "expval",
             [](Measurements<StateVectorT> &M, const np_arr_sparse_ind &row_map,

diff --git a/pennylane_lightning/core/src/simulators/lightning_kokkos/gates/GateFunctorsNonparam.hpp b/pennylane_lightning/core/src/simulators/lightning_kokkos/gates/GateFunctorsNonparam.hpp
@@ -173,7 +173,8 @@ template <class PrecisionT, bool inverse = false> struct sFunctor {
         rev_wire_shift = (static_cast<size_t>(1U) << rev_wire);
         wire_parity = fillTrailingOnes(rev_wire);
         wire_parity_inv = fillLeadingOnes(rev_wire + 1);
-        shift = (inverse) ? -Kokkos::complex(0, 1) : Kokkos::complex(0, 1);
+        shift =
+            (inverse) ? -Kokkos::complex{0.0, 1.0} : Kokkos::complex{0.0, 1.0};
     }
 
     KOKKOS_INLINE_FUNCTION

diff --git a/pennylane_lightning/core/src/simulators/lightning_kokkos/measurements/ExpValFunctors.hpp b/pennylane_lightning/core/src/simulators/lightning_kokkos/measurements/ExpValFunctors.hpp
@@ -256,6 +256,76 @@ template <class PrecisionT> struct getExpectationValueTwoQubitOpFunctor {
     }
 };
 
+template <class PrecisionT> struct getExpectationValueMultiQubitOpFunctor {
+    using ComplexT = Kokkos::complex<PrecisionT>;
+    using KokkosComplexVector = Kokkos::View<ComplexT *>;
+    using KokkosIntVector = Kokkos::View<std::size_t *>;
+    using ScratchViewComplex =
+        Kokkos::View<ComplexT *,
+                     Kokkos::DefaultExecutionSpace::scratch_memory_space,
+                     Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
+    using MemberType = Kokkos::TeamPolicy<>::member_type;
+
+    KokkosComplexVector arr;
+    KokkosComplexVector matrix;
+    KokkosIntVector wires;
+    std::size_t dim;
+    std::size_t num_qubits;
+
+    getExpectationValueMultiQubitOpFunctor(const KokkosComplexVector &arr_,
+                                           std::size_t num_qubits_,
+                                           const KokkosComplexVector &matrix_,
+                                           KokkosIntVector &wires_) {
+        dim = 1U << wires_.size();
+        num_qubits = num_qubits_;
+        wires = wires_;
+        arr = arr_;
+        matrix = matrix_;
+    }
+
+    KOKKOS_INLINE_FUNCTION
+    void operator()(const MemberType &teamMember, PrecisionT &expval) const {
+        const std::size_t k = teamMember.league_rank() * dim;
+        PrecisionT tempExpVal = 0.0;
+        ScratchViewComplex coeffs_in(teamMember.team_scratch(0), dim);
+        if (teamMember.team_rank() == 0) {
+            Kokkos::parallel_for(
+                Kokkos::ThreadVectorRange(teamMember, dim),
+                [&](const std::size_t inner_idx) {
+                    std::size_t idx = k | inner_idx;
+                    const std::size_t n_wires = wires.size();
+                    for (std::size_t pos = 0; pos < n_wires; pos++) {
+                        std::size_t x =
+                            ((idx >> (n_wires - pos - 1)) ^
+                             (idx >> (num_qubits - wires(pos) - 1))) &
+                            1U;
+                        idx = idx ^ ((x << (n_wires - pos - 1)) |
+                                     (x << (num_qubits - wires(pos) - 1)));
+                    }
+                    coeffs_in(inner_idx) = arr(idx);
+                });
+        }
+        teamMember.team_barrier();
+        Kokkos::parallel_reduce(
+            Kokkos::TeamThreadRange(teamMember, dim),
+            [&](const std::size_t i, PrecisionT &innerExpVal) {
+                const std::size_t base_idx = i * dim;
+                ComplexT tmp{0.0};
+                Kokkos::parallel_reduce(
+                    Kokkos::ThreadVectorRange(teamMember, dim),
+                    [&](const std::size_t j, ComplexT &isum) {
+                        isum = isum + matrix(base_idx + j) * coeffs_in(j);
+                    },
+                    tmp);
+                innerExpVal += real(conj(coeffs_in(i)) * tmp);
+            },
+            tempExpVal);
+        if (teamMember.team_rank() == 0) {
+            expval += tempExpVal;
+        }
+    }
+};
+
 template <class PrecisionT> struct getExpectationValueSparseFunctor {
     using KokkosComplexVector = Kokkos::View<Kokkos::complex<PrecisionT> *>;
     using KokkosSizeTVector = Kokkos::View<std::size_t *>;