Accel/expval

.github/CHANGELOG.md

@@ -20,6 +20,10 @@
   [(#477)](https://github.com/PennyLaneAI/pennylane-lightning/pull/477)
 
 ### 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)
 

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>>;

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,17 @@ 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) {
+                auto matrix_data =
+                    static_cast<ComplexT *>(matrix.request().ptr);
+                std::vector<ComplexT> matrix_v{
+                    matrix_data, matrix_data + exp2(2 * wires.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,

pennylane_lightning/core/src/simulators/lightning_kokkos/measurements/ExpValFunctors.hpp

@@ -254,6 +254,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 *>;