From fb82dea189ab1afc1209d0ed7e2147ae8ee1d713 Mon Sep 17 00:00:00 2001 From: Vincent Michaud-Rioux Date: Fri, 25 Aug 2023 10:48:19 -0400 Subject: [PATCH] Accel/expval (#481) * Introduce std::unordered_map 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. --- .github/CHANGELOG.md | 3 + .../lightning_kokkos/StateVectorKokkos.hpp | 1 - .../bindings/LKokkosBindings.hpp | 14 + .../gates/GateFunctorsNonparam.hpp | 3 +- .../measurements/ExpValFunctors.hpp | 70 ++++ .../measurements/MeasurementsKokkos.hpp | 316 ++++++------------ .../tests/Test_StateVectorKokkos_Expval.cpp | 125 ------- .../lightning_kokkos/lightning_kokkos.py | 6 + tests/test_expval.py | 20 ++ 9 files changed, 210 insertions(+), 348 deletions(-) diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md index 3d2c249ced..710fd2b1eb 100644 --- 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 index fdcc1147b7..1683f7f770 100644 --- 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; using KokkosSizeTVector = Kokkos::View; - using KokkosRangePolicy = Kokkos::RangePolicy; using UnmanagedComplexHostView = Kokkos::View>; diff --git a/pennylane_lightning/core/src/simulators/lightning_kokkos/bindings/LKokkosBindings.hpp b/pennylane_lightning/core/src/simulators/lightning_kokkos/bindings/LKokkosBindings.hpp index 431952580d..4f54ca0a0f 100644 --- 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 &)>( &Measurements::expval), "Expected value of an operation by name.") + .def( + "expval", + [](Measurements &M, const np_arr_c &matrix, + const std::vector &wires) { + const std::size_t matrix_size = exp2(2 * wires.size()); + auto matrix_data = + static_cast(matrix.request().ptr); + std::vector matrix_v{matrix_data, + matrix_data + matrix_size}; + return M.expval(matrix_v, wires); + }, + "Expected value of a Hermitian observable.") .def( "expval", [](Measurements &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 index 0dd5a386a6..a66ce4e72b 100644 --- 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 struct sFunctor { rev_wire_shift = (static_cast(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 index 0526014f33..23de0e44a8 100644 --- 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 struct getExpectationValueTwoQubitOpFunctor { } }; +template struct getExpectationValueMultiQubitOpFunctor { + using ComplexT = Kokkos::complex; + using KokkosComplexVector = Kokkos::View; + using KokkosIntVector = Kokkos::View; + using ScratchViewComplex = + Kokkos::View>; + 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 struct getExpectationValueSparseFunctor { using KokkosComplexVector = Kokkos::View *>; using KokkosSizeTVector = Kokkos::View; diff --git a/pennylane_lightning/core/src/simulators/lightning_kokkos/measurements/MeasurementsKokkos.hpp b/pennylane_lightning/core/src/simulators/lightning_kokkos/measurements/MeasurementsKokkos.hpp index 6183c3fa94..6b80fd0058 100644 --- a/pennylane_lightning/core/src/simulators/lightning_kokkos/measurements/MeasurementsKokkos.hpp +++ b/pennylane_lightning/core/src/simulators/lightning_kokkos/measurements/MeasurementsKokkos.hpp @@ -32,6 +32,15 @@ using Pennylane::LightningKokkos::StateVectorKokkos; using Pennylane::LightningKokkos::Util::getRealOfComplexInnerProduct; using Pennylane::LightningKokkos::Util::SparseMV_Kokkos; using Pennylane::Util::exp2; +enum class ExpValFunc : uint32_t { + BEGIN = 0, + Identity, + PauliX, + PauliY, + PauliZ, + Hadamard, + END +}; } // namespace /// @endcond @@ -55,231 +64,55 @@ class Measurements final typename StateVectorT::UnmanagedConstSizeTHostView; using UnmanagedPrecisionHostView = typename StateVectorT::UnmanagedPrecisionHostView; - - using ExpValFunc = std::function &, const std::vector &)>; - using ExpValMap = std::unordered_map; - - ExpValMap expval_funcs; + using ScratchViewComplex = typename StateVectorT::ScratchViewComplex; + using TeamPolicy = typename StateVectorT::TeamPolicy; public: explicit Measurements(const StateVectorT &statevector) - : BaseType{statevector}, - expval_funcs{{"Identity", - [&](auto &&wires, auto &¶ms) { - return getExpectationValueIdentity( - std::forward(wires), - std::forward(params)); - }}, - {"PauliX", - [&](auto &&wires, auto &¶ms) { - return getExpectationValuePauliX( - std::forward(wires), - std::forward(params)); - }}, - {"PauliY", - [&](auto &&wires, auto &¶ms) { - return getExpectationValuePauliY( - std::forward(wires), - std::forward(params)); - }}, - {"PauliZ", - [&](auto &&wires, auto &¶ms) { - return getExpectationValuePauliZ( - std::forward(wires), - std::forward(params)); - }}, - {"Hadamard", [&](auto &&wires, auto &¶ms) { - return getExpectationValueHadamard( - std::forward(wires), - std::forward(params)); - }}} {}; - - /** - * @brief Calculate the expectation value of a named observable. - * - * @param obsName observable name - * @param wires wires the observable acts on - * @param params parameters for the observable - * @param gate_matrix optional matrix - */ - PrecisionT getExpectationValue( - const std::string &obsName, const std::vector &wires, - [[maybe_unused]] const std::vector ¶ms = {0.0}, - const std::vector &gate_matrix = {}) { - auto &&par = (params.empty()) ? std::vector{0.0} : params; - auto &&local_wires = - (gate_matrix.empty()) - ? wires - : std::vector{ - wires.rbegin(), - wires.rend()}; // ensure wire indexing correctly preserved - // for tensor-observables - - if (expval_funcs.find(obsName) != expval_funcs.end()) { - return expval_funcs.at(obsName)(local_wires, par); - } - - KokkosVector matrix("gate_matrix", gate_matrix.size()); - Kokkos::deep_copy(matrix, UnmanagedConstComplexHostView( - gate_matrix.data(), gate_matrix.size())); - return getExpectationValueMultiQubitOp(matrix, wires, par); - } - - /** - * @brief Calculate expectation value with respect to identity observable on - * specified wire. For normalised states this function will always return 1. - * - * @param wires Wire to apply observable to. - * @param params Not used. - * @return Squared norm of state. - */ - auto getExpectationValueIdentity( - const std::vector &wires, - [[maybe_unused]] const std::vector ¶ms = {0.0}) { - const size_t num_qubits = this->_statevector.getNumQubits(); - const Kokkos::View arr_data = this->_statevector.getView(); - PrecisionT expval = 0; - Kokkos::parallel_reduce( - exp2(num_qubits), - getExpectationValueIdentityFunctor(arr_data, num_qubits, wires), - expval); - return expval; - } - - /** - * @brief Calculate expectation value with respect to Pauli X observable on - * specified wire. - * - * @param wires Wire to apply observable to. - * @param params Not used. - * @return Expectation value with respect to Pauli X applied to specified - * wire. - */ - auto getExpectationValuePauliX( - const std::vector &wires, - [[maybe_unused]] const std::vector ¶ms = {0.0}) { - const size_t num_qubits = this->_statevector.getNumQubits(); - const Kokkos::View arr_data = this->_statevector.getView(); - PrecisionT expval = 0; - Kokkos::parallel_reduce( - exp2(num_qubits - 1), - getExpectationValuePauliXFunctor(arr_data, num_qubits, wires), - expval); - return expval; - } - - /** - * @brief Calculate expectation value with respect to Pauli Y observable on - * specified wire. - * - * @param wires Wire to apply observable to. - * @param params Not used. - * @return Expectation value with respect to Pauli Y applied to specified - * wire. - */ - auto getExpectationValuePauliY( - const std::vector &wires, - [[maybe_unused]] const std::vector ¶ms = {0.0}) { - const size_t num_qubits = this->_statevector.getNumQubits(); - const Kokkos::View arr_data = this->_statevector.getView(); - PrecisionT expval = 0; - Kokkos::parallel_reduce( - exp2(num_qubits - 1), - getExpectationValuePauliYFunctor(arr_data, num_qubits, wires), - expval); - return expval; - } + : BaseType{statevector} { + init_expval_funcs_(); + }; /** - * @brief Calculate expectation value with respect to Pauli Z observable on - * specified wire. + * @brief Templated method that returns the expectation value of named + * observables. * - * @param wires Wire to apply observable to. - * @param params Not used. - * @return Expectation value with respect to Pauli Z applied to specified - * wire. + * @tparam functor_t Expectation value functor class for Kokkos dispatcher. + * @tparam nqubits Number of wires. + * @param wires Wires to apply the observable to. */ - auto getExpectationValuePauliZ( - const std::vector &wires, - [[maybe_unused]] const std::vector ¶ms = {0.0}) { - const size_t num_qubits = this->_statevector.getNumQubits(); - const Kokkos::View arr_data = this->_statevector.getView(); - PrecisionT expval = 0; - Kokkos::parallel_reduce( - exp2(num_qubits - 1), - getExpectationValuePauliZFunctor(arr_data, num_qubits, wires), - expval); - return expval; - } + template