Program Listing for File MeasurementsKokkos.hpp

↰ Return to documentation for file (pennylane_lightning/core/src/simulators/lightning_kokkos/measurements/MeasurementsKokkos.hpp)

// Copyright 2018-2023 Xanadu Quantum Technologies Inc.

// Licensed under the Apache License, Version 2.0 (the License);
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

// http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an AS IS BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <chrono>

#include <Kokkos_Core.hpp>
#include <Kokkos_Random.hpp>

#include "ExpValFunctors.hpp"
#include "LinearAlgebraKokkos.hpp" // getRealOfComplexInnerProduct
#include "MeasurementsBase.hpp"
#include "MeasuresFunctors.hpp"
#include "Observables.hpp"
#include "ObservablesKokkos.hpp"
#include "StateVectorKokkos.hpp"
#include "Util.hpp"

namespace {
using namespace Pennylane::Measures;
using namespace Pennylane::Observables;
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 = 1,
    Identity = 1,
    PauliX,
    PauliY,
    PauliZ,
    Hadamard,
    END
};
} // namespace

namespace Pennylane::LightningKokkos::Measures {
template <class StateVectorT>
class Measurements final
    : public MeasurementsBase<StateVectorT, Measurements<StateVectorT>> {
  private:
    using PrecisionT = typename StateVectorT::PrecisionT;
    using ComplexT = typename StateVectorT::ComplexT;
    using BaseType = MeasurementsBase<StateVectorT, Measurements<StateVectorT>>;

    using KokkosExecSpace = typename StateVectorT::KokkosExecSpace;
    using KokkosVector = typename StateVectorT::KokkosVector;
    using KokkosSizeTVector = typename StateVectorT::KokkosSizeTVector;
    using UnmanagedSizeTHostView =
        typename StateVectorT::UnmanagedSizeTHostView;
    using UnmanagedConstComplexHostView =
        typename StateVectorT::UnmanagedConstComplexHostView;
    using UnmanagedConstSizeTHostView =
        typename StateVectorT::UnmanagedConstSizeTHostView;
    using UnmanagedPrecisionHostView =
        typename StateVectorT::UnmanagedPrecisionHostView;
    using ScratchViewComplex = typename StateVectorT::ScratchViewComplex;
    using TeamPolicy = typename StateVectorT::TeamPolicy;

  public:
    explicit Measurements(const StateVectorT &statevector)
        : BaseType{statevector} {
        init_expval_funcs_();
    };

    template <template <class> class functor_t, int num_wires>
    PrecisionT applyExpValNamedFunctor(const std::vector<std::size_t> &wires) {
        if constexpr (num_wires > 0)
            PL_ASSERT(wires.size() == num_wires);

        const std::size_t num_qubits = this->_statevector.getNumQubits();
        const Kokkos::View<ComplexT *> arr_data = this->_statevector.getView();
        PrecisionT expval = 0.0;
        Kokkos::parallel_reduce(exp2(num_qubits - num_wires),
                                functor_t(arr_data, num_qubits, wires), expval);
        return expval;
    }

    template <template <class> class functor_t, int num_wires>
    PrecisionT applyExpValFunctor(const KokkosVector &matrix,
                                  const std::vector<std::size_t> &wires) {
        PL_ASSERT(wires.size() == num_wires);
        const std::size_t num_qubits = this->_statevector.getNumQubits();
        Kokkos::View<ComplexT *> arr_data = this->_statevector.getView();
        PrecisionT expval = 0.0;
        Kokkos::parallel_reduce(
            exp2(num_qubits - num_wires),
            functor_t<PrecisionT>(arr_data, num_qubits, matrix, wires), expval);
        return expval;
    }

    auto getExpValMatrix(const KokkosVector &matrix,
                         const std::vector<std::size_t> &wires) -> PrecisionT {
        std::size_t num_qubits = this->_statevector.getNumQubits();
        std::size_t two2N = std::exp2(num_qubits - wires.size());
        std::size_t dim = std::exp2(wires.size());
        const KokkosVector arr_data = this->_statevector.getView();

        PrecisionT expval = 0.0;
        switch (wires.size()) {
        case 1:
            Kokkos::parallel_reduce(two2N,
                                    getExpVal1QubitOpFunctor<PrecisionT>(
                                        arr_data, num_qubits, matrix, wires),
                                    expval);
            break;
        case 2:
            Kokkos::parallel_reduce(two2N,
                                    getExpVal2QubitOpFunctor<PrecisionT>(
                                        arr_data, num_qubits, matrix, wires),
                                    expval);
            break;
        case 3:
            Kokkos::parallel_reduce(two2N,
                                    getExpVal3QubitOpFunctor<PrecisionT>(
                                        arr_data, num_qubits, matrix, wires),
                                    expval);
            break;
        case 4:
            Kokkos::parallel_reduce(two2N,
                                    getExpVal4QubitOpFunctor<PrecisionT>(
                                        arr_data, num_qubits, matrix, wires),
                                    expval);
            break;
        default:
            std::size_t scratch_size = ScratchViewComplex::shmem_size(dim);
            Kokkos::parallel_reduce(
                "getExpValMultiQubitOpFunctor",
                TeamPolicy(two2N, Kokkos::AUTO, dim)
                    .set_scratch_size(0, Kokkos::PerTeam(scratch_size)),
                getExpValMultiQubitOpFunctor<PrecisionT>(arr_data, num_qubits,
                                                         matrix, wires),
                expval);
            break;
        }
        return expval;
    }

    PrecisionT expval(const Observable<StateVectorT> &ob) {
        StateVectorT ob_sv{this->_statevector};
        ob.applyInPlace(ob_sv);
        return getRealOfComplexInnerProduct(this->_statevector.getView(),
                                            ob_sv.getView());
    }

    PrecisionT
    expval(const Pennylane::LightningKokkos::Observables::HermitianObs<
           StateVectorT> &ob) {
        return expval(ob.getMatrix(), ob.getWires());
    }

    PrecisionT expval(const std::vector<ComplexT> &matrix_,
                      const std::vector<std::size_t> &wires) {
        PL_ABORT_IF(matrix_.size() != exp2(2 * wires.size()),
                    "The size of matrix does not match with the given "
                    "number of wires");
        KokkosVector matrix("matrix_", matrix_.size());
        Kokkos::deep_copy(matrix, UnmanagedConstComplexHostView(
                                      matrix_.data(), matrix_.size()));
        return getExpValMatrix(matrix, wires);
    };

    PrecisionT expval(const std::string &operation,
                      const std::vector<std::size_t> &wires) {
        switch (expval_funcs_[operation]) {
        case ExpValFunc::Identity:
            return applyExpValNamedFunctor<getExpectationValueIdentityFunctor,
                                           0>(wires);
        case ExpValFunc::PauliX:
            return applyExpValNamedFunctor<getExpectationValuePauliXFunctor, 1>(
                wires);
        case ExpValFunc::PauliY:
            return applyExpValNamedFunctor<getExpectationValuePauliYFunctor, 1>(
                wires);
        case ExpValFunc::PauliZ:
            return applyExpValNamedFunctor<getExpectationValuePauliZFunctor, 1>(
                wires);
        case ExpValFunc::Hadamard:
            return applyExpValNamedFunctor<getExpectationValueHadamardFunctor,
                                           1>(wires);
        default:
            PL_ABORT(
                std::string("Expval does not exist for named observable ") +
                operation);
        }
    };

    template <typename op_type>
    std::vector<PrecisionT>
    expval(const std::vector<op_type> &operations_list,
           const std::vector<std::vector<std::size_t>> &wires_list) {
        PL_ABORT_IF(
            (operations_list.size() != wires_list.size()),
            "The lengths of the list of operations and wires do not match.");
        std::vector<PrecisionT> expected_value_list;

        for (size_t index = 0; index < operations_list.size(); index++) {
            expected_value_list.emplace_back(
                expval(operations_list[index], wires_list[index]));
        }

        return expected_value_list;
    }

    auto expval(const Observable<StateVectorT> &obs,
                const std::size_t &num_shots,
                const std::vector<std::size_t> &shot_range) -> PrecisionT {
        return BaseType::expval(obs, num_shots, shot_range);
    }

    template <class index_type>
    PrecisionT expval(const index_type *row_map_ptr,
                      const index_type row_map_size,
                      const index_type *entries_ptr, const ComplexT *values_ptr,
                      const index_type numNNZ) {
        const Kokkos::View<ComplexT *> arr_data = this->_statevector.getView();
        PrecisionT expval = 0.0;
        KokkosSizeTVector kok_row_map("row_map", row_map_size);
        KokkosSizeTVector kok_indices("indices", numNNZ);
        KokkosVector kok_data("data", numNNZ);

        Kokkos::deep_copy(kok_data,
                          UnmanagedConstComplexHostView(values_ptr, numNNZ));
        Kokkos::deep_copy(kok_indices,
                          UnmanagedConstSizeTHostView(entries_ptr, numNNZ));
        Kokkos::deep_copy(kok_row_map, UnmanagedConstSizeTHostView(
                                           row_map_ptr, row_map_size));

        Kokkos::parallel_reduce(
            row_map_size - 1,
            getExpectationValueSparseFunctor<PrecisionT>(
                arr_data, kok_data, kok_indices, kok_row_map),
            expval);
        return expval;
    };

    auto var(const Observable<StateVectorT> &ob) -> PrecisionT {
        StateVectorT ob_sv{this->_statevector};
        ob.applyInPlace(ob_sv);

        const PrecisionT mean_square =
            getRealOfComplexInnerProduct(ob_sv.getView(), ob_sv.getView());
        const PrecisionT squared_mean = static_cast<PrecisionT>(
            std::pow(getRealOfComplexInnerProduct(this->_statevector.getView(),
                                                  ob_sv.getView()),
                     2));
        return (mean_square - squared_mean);
    }

    PrecisionT var(const std::string &operation,
                   const std::vector<std::size_t> &wires) {
        StateVectorT ob_sv{this->_statevector};
        ob_sv.applyOperation(operation, wires);

        const PrecisionT mean_square =
            getRealOfComplexInnerProduct(ob_sv.getView(), ob_sv.getView());
        const PrecisionT squared_mean = static_cast<PrecisionT>(
            std::pow(getRealOfComplexInnerProduct(this->_statevector.getView(),
                                                  ob_sv.getView()),
                     2));
        return (mean_square - squared_mean);
    };

    PrecisionT var(const std::vector<ComplexT> &matrix,
                   const std::vector<std::size_t> &wires) {
        StateVectorT ob_sv{this->_statevector};
        ob_sv.applyMatrix(matrix, wires);

        const PrecisionT mean_square =
            getRealOfComplexInnerProduct(ob_sv.getView(), ob_sv.getView());
        const PrecisionT squared_mean = static_cast<PrecisionT>(
            std::pow(getRealOfComplexInnerProduct(this->_statevector.getView(),
                                                  ob_sv.getView()),
                     2));
        return (mean_square - squared_mean);
    };

    template <typename op_type>
    std::vector<PrecisionT>
    var(const std::vector<op_type> &operations_list,
        const std::vector<std::vector<std::size_t>> &wires_list) {
        PL_ABORT_IF(
            (operations_list.size() != wires_list.size()),
            "The lengths of the list of operations and wires do not match.");

        std::vector<PrecisionT> expected_value_list;

        for (size_t index = 0; index < operations_list.size(); index++) {
            expected_value_list.emplace_back(
                var(operations_list[index], wires_list[index]));
        }

        return expected_value_list;
    };

    template <class index_type>
    PrecisionT var(const index_type *row_map_ptr, const index_type row_map_size,
                   const index_type *entries_ptr, const ComplexT *values_ptr,
                   const index_type numNNZ) {
        PL_ABORT_IF(
            (this->_statevector.getLength() != (size_t(row_map_size) - 1)),
            "Statevector and Hamiltonian have incompatible sizes.");

        StateVectorT ob_sv{this->_statevector};

        SparseMV_Kokkos<PrecisionT>(this->_statevector.getView(),
                                    ob_sv.getView(), row_map_ptr, row_map_size,
                                    entries_ptr, values_ptr, numNNZ);

        const PrecisionT mean_square =
            getRealOfComplexInnerProduct(ob_sv.getView(), ob_sv.getView());
        const PrecisionT squared_mean = static_cast<PrecisionT>(
            std::pow(getRealOfComplexInnerProduct(this->_statevector.getView(),
                                                  ob_sv.getView()),
                     2));
        return (mean_square - squared_mean);
    };

    auto var(const Observable<StateVectorT> &obs, const std::size_t &num_shots)
        -> PrecisionT {
        return BaseType::var(obs, num_shots);
    }

    auto probs() -> std::vector<PrecisionT> {
        const std::size_t N = this->_statevector.getLength();

        Kokkos::View<ComplexT *> arr_data = this->_statevector.getView();
        Kokkos::View<PrecisionT *> d_probability("d_probability", N);

        // Compute probability distribution from StateVector
        Kokkos::parallel_for(
            Kokkos::RangePolicy<KokkosExecSpace>(0, N),
            getProbFunctor<PrecisionT>(arr_data, d_probability));

        std::vector<PrecisionT> probabilities(N, 0);

        Kokkos::deep_copy(UnmanagedPrecisionHostView(probabilities.data(),
                                                     probabilities.size()),
                          d_probability);
        return probabilities;
    }

    std::vector<PrecisionT> probs(const std::vector<std::size_t> &wires) {
        PL_ABORT_IF_NOT(
            std::is_sorted(wires.cbegin(), wires.cend()),
            "LightningKokkos does not currently support out-of-order wire "
            "indices with probability calculations");
        using MDPolicyType_2D =
            Kokkos::MDRangePolicy<Kokkos::Rank<2, Kokkos::Iterate::Left>>;

        //  Determining probabilities for the sorted wires.
        const Kokkos::View<ComplexT *> arr_data = this->_statevector.getView();
        const std::size_t num_qubits = this->_statevector.getNumQubits();

        std::vector<std::size_t> sorted_ind_wires(wires);
        const bool is_sorted_wires =
            std::is_sorted(sorted_ind_wires.begin(), sorted_ind_wires.end());
        std::vector<std::size_t> sorted_wires(wires);

        if (!is_sorted_wires) {
            sorted_ind_wires = Pennylane::Util::sorting_indices(wires);
            for (size_t pos = 0; pos < wires.size(); pos++)
                sorted_wires[pos] = wires[sorted_ind_wires[pos]];
        }

        std::vector<std::size_t> all_indices =
            Pennylane::Util::generateBitsPatterns(sorted_wires, num_qubits);

        std::vector<std::size_t> all_offsets =
            Pennylane::Util::generateBitsPatterns(
                Pennylane::Util::getIndicesAfterExclusion(sorted_wires,
                                                          num_qubits),
                num_qubits);

        Kokkos::View<PrecisionT *> d_probabilities("d_probabilities",
                                                   all_indices.size());

        Kokkos::View<size_t *> d_sorted_ind_wires("d_sorted_ind_wires",
                                                  sorted_ind_wires.size());
        Kokkos::View<size_t *> d_all_indices("d_all_indices",
                                             all_indices.size());
        Kokkos::View<size_t *> d_all_offsets("d_all_offsets",
                                             all_offsets.size());

        Kokkos::deep_copy(
            d_all_indices,
            UnmanagedSizeTHostView(all_indices.data(), all_indices.size()));
        Kokkos::deep_copy(
            d_all_offsets,
            UnmanagedSizeTHostView(all_offsets.data(), all_offsets.size()));
        Kokkos::deep_copy(d_sorted_ind_wires,
                          UnmanagedSizeTHostView(sorted_ind_wires.data(),
                                                 sorted_ind_wires.size()));

        const int num_all_indices =
            all_indices.size(); // int is required by Kokkos::MDRangePolicy
        const int num_all_offsets = all_offsets.size();

        MDPolicyType_2D mdpolicy_2d0({{0, 0}},
                                     {{num_all_indices, num_all_offsets}});

        Kokkos::parallel_for(
            "Set_Prob", mdpolicy_2d0,
            KOKKOS_LAMBDA(const std::size_t i, const std::size_t j) {
                const std::size_t index = d_all_indices(i) + d_all_offsets(j);
                const PrecisionT REAL = arr_data(index).real();
                const PrecisionT IMAG = arr_data(index).imag();
                const PrecisionT value = REAL * REAL + IMAG * IMAG;
                Kokkos::atomic_add(&d_probabilities(i), value);
            });

        std::vector<PrecisionT> probabilities(all_indices.size(), 0);

        if (is_sorted_wires) {
            Kokkos::deep_copy(UnmanagedPrecisionHostView(probabilities.data(),
                                                         probabilities.size()),
                              d_probabilities);
            return probabilities;
        } else {
            Kokkos::View<PrecisionT *> transposed_tensor("transposed_tensor",
                                                         all_indices.size());

            Kokkos::View<size_t *> d_trans_index("d_trans_index",
                                                 all_indices.size());

            const int num_trans_tensor = transposed_tensor.size();
            const int num_sorted_ind_wires = sorted_ind_wires.size();

            MDPolicyType_2D mdpolicy_2d1(
                {{0, 0}}, {{num_trans_tensor, num_sorted_ind_wires}});

            Kokkos::parallel_for(
                "TransIndex", mdpolicy_2d1,
                getTransposedIndexFunctor(d_sorted_ind_wires, d_trans_index,
                                          num_sorted_ind_wires));

            Kokkos::parallel_for(
                "Transpose",
                Kokkos::RangePolicy<KokkosExecSpace>(0, num_trans_tensor),
                getTransposedFunctor<PrecisionT>(
                    transposed_tensor, d_probabilities, d_trans_index));

            Kokkos::deep_copy(UnmanagedPrecisionHostView(probabilities.data(),
                                                         probabilities.size()),
                              transposed_tensor);

            return probabilities;
        }
    }

    std::vector<PrecisionT> probs(const Observable<StateVectorT> &obs,
                                  std::size_t num_shots = 0) {
        return BaseType::probs(obs, num_shots);
    }

    std::vector<PrecisionT> probs(size_t num_shots) {
        return BaseType::probs(num_shots);
    }

    std::vector<PrecisionT> probs(const std::vector<std::size_t> &wires,
                                  std::size_t num_shots) {
        PL_ABORT_IF_NOT(
            std::is_sorted(wires.cbegin(), wires.cend()),
            "LightningKokkos does not currently support out-of-order wire "
            "indices with probability calculations");

        return BaseType::probs(wires, num_shots);
    }

    auto generate_samples(size_t num_samples) -> std::vector<std::size_t> {
        const std::size_t num_qubits = this->_statevector.getNumQubits();
        const std::size_t N = this->_statevector.getLength();

        Kokkos::View<ComplexT *> arr_data = this->_statevector.getView();
        Kokkos::View<PrecisionT *> probability("probability", N);
        Kokkos::View<size_t *> samples("num_samples", num_samples * num_qubits);

        // Compute probability distribution from StateVector
        Kokkos::parallel_for(Kokkos::RangePolicy<KokkosExecSpace>(0, N),
                             getProbFunctor<PrecisionT>(arr_data, probability));

        // Convert probability distribution to cumulative distribution
        Kokkos::parallel_scan(
            Kokkos::RangePolicy<KokkosExecSpace>(0, N),
            KOKKOS_LAMBDA(const std::size_t k, PrecisionT &update_value,
                          const bool is_final) {
                const PrecisionT val_k = probability(k);
                if (is_final)
                    probability(k) = update_value;
                update_value += val_k;
            });

        // Sampling using Random_XorShift64_Pool
        Kokkos::Random_XorShift64_Pool<> rand_pool(
            std::chrono::high_resolution_clock::now()
                .time_since_epoch()
                .count());

        Kokkos::parallel_for(
            Kokkos::RangePolicy<KokkosExecSpace>(0, num_samples),
            Sampler<PrecisionT, Kokkos::Random_XorShift64_Pool>(
                samples, probability, rand_pool, num_qubits, N));

        std::vector<std::size_t> samples_h(num_samples * num_qubits);

        using UnmanagedSize_tHostView =
            Kokkos::View<size_t *, Kokkos::HostSpace,
                         Kokkos::MemoryTraits<Kokkos::Unmanaged>>;

        Kokkos::deep_copy(
            UnmanagedSize_tHostView(samples_h.data(), samples_h.size()),
            samples);

        return samples_h;
    }

  private:
    std::unordered_map<std::string, ExpValFunc> expval_funcs_;

    // clang-format off
    void init_expval_funcs_() {
        expval_funcs_["Identity"] = ExpValFunc::Identity;
        expval_funcs_["PauliX"]   = ExpValFunc::PauliX;
        expval_funcs_["PauliY"]   = ExpValFunc::PauliY;
        expval_funcs_["PauliZ"]   = ExpValFunc::PauliZ;
        expval_funcs_["Hadamard"] = ExpValFunc::Hadamard;
    }
    // clang-format on
};

} // namespace Pennylane::LightningKokkos::Measures

api/program_listing_file_pennylane_lightning_core_src_simulators_lightning_kokkos_measurements_MeasurementsKokkos.hpp

 

Download Python script

 

Download Notebook

 

View on GitHub