Program Listing for File StateVectorKokkos.hpp¶
↰ Return to documentation for file (pennylane_lightning/core/simulators/lightning_kokkos/StateVectorKokkos.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 <complex>
#include <cstddef>
#include <cstdlib>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>
#include <Kokkos_Core.hpp>
#include <Kokkos_Random.hpp>
#include "BitUtil.hpp" // isPerfectPowerOf2
#include "Constant.hpp"
#include "ConstantUtil.hpp"
#include "Error.hpp"
#include "GateFunctors.hpp"
#include "GateOperation.hpp"
#include "StateVectorBase.hpp"
#include "Util.hpp"
#include "UtilKokkos.hpp"
#include "CPUMemoryModel.hpp"
namespace {
using namespace Pennylane::Gates::Constant;
using namespace Pennylane::LightningKokkos::Functors;
using namespace Pennylane::LightningKokkos::Util;
using Pennylane::Gates::GateOperation;
using Pennylane::Gates::GeneratorOperation;
using Pennylane::Util::array_contains;
using Pennylane::Util::exp2;
using Pennylane::Util::isPerfectPowerOf2;
using Pennylane::Util::log2;
using Pennylane::Util::reverse_lookup;
using std::size_t;
} // namespace
namespace Pennylane::LightningKokkos {
template <class fp_t = double>
class StateVectorKokkos final
: public StateVectorBase<fp_t, StateVectorKokkos<fp_t>> {
private:
using BaseType = StateVectorBase<fp_t, StateVectorKokkos<fp_t>>;
public:
using PrecisionT = fp_t;
using ComplexT = Kokkos::complex<fp_t>;
using CFP_t = ComplexT;
using DoubleLoopRank = Kokkos::Rank<2>;
using HostExecSpace = Kokkos::DefaultHostExecutionSpace;
using KokkosExecSpace = Kokkos::DefaultExecutionSpace;
using KokkosVector = Kokkos::View<ComplexT *>;
using KokkosSizeTVector = Kokkos::View<std::size_t *>;
using UnmanagedComplexHostView =
Kokkos::View<ComplexT *, Kokkos::HostSpace,
Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
using UnmanagedSizeTHostView =
Kokkos::View<std::size_t *, Kokkos::HostSpace,
Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
using UnmanagedConstComplexHostView =
Kokkos::View<const ComplexT *, Kokkos::HostSpace,
Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
using UnmanagedConstSizeTHostView =
Kokkos::View<const std::size_t *, Kokkos::HostSpace,
Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
using UnmanagedPrecisionHostView =
Kokkos::View<PrecisionT *, Kokkos::HostSpace,
Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
using ScratchViewComplex =
Kokkos::View<ComplexT *, KokkosExecSpace::scratch_memory_space,
Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
using ScratchViewSizeT =
Kokkos::View<std::size_t *, KokkosExecSpace::scratch_memory_space,
Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
using TeamPolicy = Kokkos::TeamPolicy<>;
using MemoryStorageT = Pennylane::Util::MemoryStorageLocation::Undefined;
StateVectorKokkos() = delete;
StateVectorKokkos(std::size_t num_qubits,
const Kokkos::InitializationSettings &kokkos_args = {})
: BaseType{num_qubits} {
num_qubits_ = num_qubits;
{
const std::lock_guard<std::mutex> lock(init_mutex_);
if (!Kokkos::is_initialized()) {
Kokkos::initialize(kokkos_args);
}
}
#ifdef _WIN32
PL_ABORT_IF_NOT(num_qubits,
"LightningKokkos zero-qubit device initialization is "
"not supported on Windows.");
#endif
data_ = std::make_unique<KokkosVector>("data_", exp2(num_qubits));
setBasisState(0U);
};
void initZeros() { Kokkos::deep_copy(getView(), ComplexT{0.0, 0.0}); }
void setBasisState(std::size_t index) {
KokkosVector sv_view =
getView(); // circumvent error capturing this with KOKKOS_LAMBDA
Kokkos::parallel_for(
sv_view.size(), KOKKOS_LAMBDA(std::size_t i) {
sv_view(i) =
(i == index) ? ComplexT{1.0, 0.0} : ComplexT{0.0, 0.0};
});
}
void setBasisState(const std::vector<std::size_t> &state,
const std::vector<std::size_t> &wires) {
PL_ABORT_IF_NOT(state.size() == wires.size(),
"state and wires must have equal dimensions.");
const auto num_qubits = this->getNumQubits();
PL_ABORT_IF_NOT(
std::find_if(wires.begin(), wires.end(),
[&num_qubits](const auto i) {
return i >= num_qubits;
}) == wires.end(),
"wires must take values lower than the number of qubits.");
const auto n_wires = wires.size();
std::size_t index{0U};
for (std::size_t k = 0; k < n_wires; k++) {
const auto bit = static_cast<std::size_t>(state[k]);
index |= bit << (num_qubits - 1 - wires[k]);
}
setBasisState(index);
}
void resetStateVector() {
if (this->getLength() > 0) {
setBasisState(0U);
}
}
void setStateVector(const std::vector<std::size_t> &indices,
const std::vector<ComplexT> &values) {
PL_ABORT_IF_NOT(indices.size() == values.size(),
"Inconsistent indices and values dimensions.")
initZeros();
auto d_indices = vector2view(indices);
auto d_values = vector2view(values);
KokkosVector sv_view =
getView(); // circumvent error capturing this with KOKKOS_LAMBDA
Kokkos::parallel_for(
indices.size(), KOKKOS_LAMBDA(std::size_t i) {
sv_view(d_indices[i]) = d_values[i];
});
}
void setStateVector(const std::vector<ComplexT> &state,
const std::vector<std::size_t> &wires) {
PL_ABORT_IF_NOT(state.size() == exp2(wires.size()),
"Inconsistent state and wires dimensions.");
setStateVector(state.data(), wires);
}
void setStateVector(const ComplexT *state,
const std::vector<std::size_t> &wires) {
constexpr std::size_t one{1U};
const auto num_qubits = this->getNumQubits();
PL_ABORT_IF_NOT(
std::find_if(wires.begin(), wires.end(),
[&num_qubits](const auto i) {
return i >= num_qubits;
}) == wires.end(),
"wires must take values lower than the number of qubits.");
const auto num_state = exp2(wires.size());
auto d_sv = getView();
auto d_state = pointer2view(state, num_state);
auto d_wires = vector2view(wires);
initZeros();
Kokkos::parallel_for(
num_state, KOKKOS_LAMBDA(std::size_t i) {
std::size_t index{0U};
for (std::size_t w = 0; w < d_wires.size(); w++) {
const std::size_t bit = (i & (one << w)) >> w;
index |= bit << (num_qubits - 1 -
d_wires(d_wires.size() - 1 - w));
}
d_sv(index) = d_state(i);
});
}
StateVectorKokkos(ComplexT *hostdata_, std::size_t length,
const Kokkos::InitializationSettings &kokkos_args = {})
: StateVectorKokkos(log2(length), kokkos_args) {
PL_ABORT_IF_NOT(isPerfectPowerOf2(length),
"The size of provided data must be a power of 2.");
HostToDevice(hostdata_, length);
}
StateVectorKokkos(std::complex<PrecisionT> *hostdata_, std::size_t length,
const Kokkos::InitializationSettings &kokkos_args = {})
: StateVectorKokkos(log2(length), kokkos_args) {
PL_ABORT_IF_NOT(isPerfectPowerOf2(length),
"The size of provided data must be a power of 2.");
HostToDevice(reinterpret_cast<ComplexT *>(hostdata_), length);
}
StateVectorKokkos(const ComplexT *hostdata_, std::size_t length,
const Kokkos::InitializationSettings &kokkos_args = {})
: StateVectorKokkos(log2(length), kokkos_args) {
PL_ABORT_IF_NOT(isPerfectPowerOf2(length),
"The size of provided data must be a power of 2.");
std::vector<ComplexT> hostdata_copy(hostdata_, hostdata_ + length);
HostToDevice(hostdata_copy.data(), length);
}
StateVectorKokkos(std::vector<ComplexT> hostdata_,
const Kokkos::InitializationSettings &kokkos_args = {})
: StateVectorKokkos(hostdata_.data(), hostdata_.size(), kokkos_args) {}
template <class complex>
StateVectorKokkos(std::vector<complex> hostdata_,
const Kokkos::InitializationSettings &kokkos_args = {})
: StateVectorKokkos(reinterpret_cast<ComplexT *>(hostdata_.data()),
hostdata_.size(), kokkos_args) {}
StateVectorKokkos(const StateVectorKokkos &other,
const Kokkos::InitializationSettings &kokkos_args = {})
: StateVectorKokkos(other.getNumQubits(), kokkos_args) {
this->DeviceToDevice(other.getView());
}
~StateVectorKokkos() {
data_.reset();
{
const std::lock_guard<std::mutex> lock(init_mutex_);
if (!is_exit_reg_) {
is_exit_reg_ = true;
std::atexit([]() {
if (!Kokkos::is_finalized()) {
Kokkos::finalize();
}
});
}
}
}
void applyOperation(const std::string &opName,
const std::vector<std::size_t> &wires,
bool inverse = false,
const std::vector<fp_t> ¶ms = {},
const std::vector<ComplexT> &gate_matrix = {}) {
if (opName == "Identity") {
// No op
} else if (array_contains(gate_names, std::string_view{opName})) {
const std::size_t num_qubits = this->getNumQubits();
const GateOperation gateop =
reverse_lookup(gate_names, std::string_view{opName});
applyNamedOperation<KokkosExecSpace>(gateop, *data_, num_qubits,
wires, inverse, params);
} else {
PL_ABORT_IF(gate_matrix.empty(),
std::string("Operation does not exist for ") + opName +
std::string(" and no matrix provided."));
return applyMultiQubitOp(vector2view(gate_matrix), wires, inverse);
}
}
void applyPauliRot(const std::vector<std::size_t> &wires, bool inverse,
const std::vector<PrecisionT> ¶ms,
const std::string &word) {
PL_ABORT_IF_NOT(wires.size() == word.size(),
"wires and word have incompatible dimensions.");
Pennylane::LightningKokkos::Functors::applyPauliRot<KokkosExecSpace,
PrecisionT>(
getView(), this->getNumQubits(), wires, inverse, params[0], word);
}
void applyMultiQubitOp(const KokkosVector matrix,
const std::vector<std::size_t> &wires,
bool inverse = false) {
const std::size_t num_qubits = this->getNumQubits();
const std::size_t two2N = exp2(num_qubits - wires.size());
const std::size_t dim = exp2(wires.size());
KokkosVector matrix_trans("matrix_trans", matrix.size());
if (inverse) {
Kokkos::MDRangePolicy<DoubleLoopRank> policy_2d({0, 0}, {dim, dim});
Kokkos::parallel_for(
policy_2d, KOKKOS_LAMBDA(std::size_t i, std::size_t j) {
matrix_trans(i + j * dim) =
Kokkos::conj(matrix(i * dim + j));
});
} else {
matrix_trans = matrix;
}
switch (wires.size()) {
case 1:
Kokkos::parallel_for(
two2N, apply1QubitOpFunctor<fp_t>(*data_, num_qubits,
matrix_trans, wires));
break;
case 2:
Kokkos::parallel_for(
two2N, apply2QubitOpFunctor<fp_t>(*data_, num_qubits,
matrix_trans, wires));
break;
case 3:
Kokkos::parallel_for(
two2N, apply3QubitOpFunctor<fp_t>(*data_, num_qubits,
matrix_trans, wires));
break;
case 4:
Kokkos::parallel_for(
two2N, apply4QubitOpFunctor<fp_t>(*data_, num_qubits,
matrix_trans, wires));
break;
default:
// TODO: explore runtime determine L0 or L1 scratch level (for GPU
// shared memory), or SIMD
std::size_t scratch_size = ScratchViewComplex::shmem_size(dim) +
ScratchViewSizeT::shmem_size(dim);
Kokkos::parallel_for(
"multiQubitOpFunctor",
TeamPolicy(two2N, Kokkos::AUTO, dim)
.set_scratch_size(1, Kokkos::PerTeam(scratch_size)),
multiQubitOpFunctor<PrecisionT>(*data_, num_qubits,
matrix_trans, wires));
break;
}
}
void applyOperation(const std::string &opName,
const std::vector<std::size_t> &controlled_wires,
const std::vector<bool> &controlled_values,
const std::vector<std::size_t> &wires,
bool inverse = false,
const std::vector<fp_t> ¶ms = {},
const std::vector<ComplexT> &gate_matrix = {}) {
PL_ABORT_IF_NOT(
areVecsDisjoint<std::size_t>(controlled_wires, wires),
"`controlled_wires` and target wires must be disjoint.");
PL_ABORT_IF_NOT(controlled_wires.size() == controlled_values.size(),
"`controlled_wires` must have the same size as "
"`controlled_values`.");
if (controlled_wires.empty()) {
return applyOperation(opName, wires, inverse, params, gate_matrix);
}
if (array_contains(controlled_gate_names, std::string_view{opName})) {
const std::size_t num_qubits = this->getNumQubits();
const ControlledGateOperation gateop =
reverse_lookup(controlled_gate_names, std::string_view{opName});
applyNCNamedOperation<KokkosExecSpace>(
gateop, *data_, num_qubits, controlled_wires, controlled_values,
wires, inverse, params);
} else {
PL_ABORT_IF(gate_matrix.empty(),
std::string("Operation does not exist for ") + opName +
std::string(" and no matrix provided."));
return applyNCMultiQubitOp(vector2view(gate_matrix),
controlled_wires, controlled_values,
wires, inverse);
}
}
void applyNCMultiQubitOp(const KokkosVector matrix,
const std::vector<std::size_t> &controlled_wires,
const std::vector<bool> &controlled_values,
const std::vector<std::size_t> &wires,
bool inverse = false) {
const std::size_t num_qubits = this->getNumQubits();
PL_ASSERT(num_qubits >= wires.size() + controlled_wires.size());
const std::size_t two2N =
exp2(num_qubits - wires.size() - controlled_wires.size());
const std::size_t dim = exp2(wires.size());
KokkosVector matrix_trans("matrix_trans", matrix.size());
if (inverse) {
Kokkos::MDRangePolicy<DoubleLoopRank> policy_2d({0, 0}, {dim, dim});
Kokkos::parallel_for(
policy_2d, KOKKOS_LAMBDA(std::size_t i, std::size_t j) {
matrix_trans(i + j * dim) =
Kokkos::conj(matrix(i * dim + j));
});
} else {
matrix_trans = matrix;
}
switch (wires.size()) {
case 1:
Kokkos::parallel_for(two2N, applyNC1QubitOpFunctor<fp_t>(
*data_, num_qubits, matrix_trans,
controlled_wires, controlled_values,
wires));
break;
case 2:
Kokkos::parallel_for(two2N, applyNC2QubitOpFunctor<fp_t>(
*data_, num_qubits, matrix_trans,
controlled_wires, controlled_values,
wires));
break;
case 3:
Kokkos::parallel_for(two2N, applyNC3QubitOpFunctor<fp_t>(
*data_, num_qubits, matrix_trans,
controlled_wires, controlled_values,
wires));
break;
default:
// TODO: explore runtime determine L0 or L1 scratch level (for GPU
// shared memory), or SIMD
const std::size_t scratch_size =
ScratchViewComplex::shmem_size(dim) +
ScratchViewSizeT::shmem_size(dim);
Kokkos::parallel_for(
"multiNCQubitOpFunctor",
TeamPolicy(two2N, Kokkos::AUTO, dim)
.set_scratch_size(1, Kokkos::PerTeam(scratch_size)),
NCMultiQubitOpFunctor<PrecisionT>(
*data_, num_qubits, matrix_trans, controlled_wires,
controlled_values, wires));
break;
}
}
inline void applyMatrix(ComplexT *matrix,
const std::vector<std::size_t> &wires,
bool inverse = false) {
PL_ABORT_IF(wires.empty(), "Number of wires must be larger than 0");
const std::size_t n2 = exp2(wires.size() * 2);
KokkosVector matrix_("matrix_", n2);
// Note that when copying data between different memory spaces (host !=
// device), Kokkos::View<Kokkos::complex*> cannot perform a deep copy of
// unmanaged complex numbers during initialization via its constructor.
// Thus, we need to explicitly deep-copy the matrix data using
// Kokkos::deep_copy().
Kokkos::deep_copy(matrix_, UnmanagedComplexHostView(matrix, n2));
applyMultiQubitOp(matrix_, wires, inverse);
}
inline void applyMatrix(const ComplexT *matrix,
const std::vector<std::size_t> &wires,
bool inverse = false) {
PL_ABORT_IF(wires.empty(), "Number of wires must be larger than 0");
const std::size_t n2 = exp2(wires.size() * 2);
KokkosVector matrix_("matrix_", n2);
Kokkos::deep_copy(matrix_, UnmanagedConstComplexHostView(matrix, n2));
applyMultiQubitOp(matrix_, wires, inverse);
}
inline void applyMatrix(const std::vector<ComplexT> &matrix,
const std::vector<std::size_t> &wires,
bool inverse = false) {
PL_ABORT_IF(wires.empty(), "Number of wires must be larger than 0");
PL_ABORT_IF(matrix.size() != exp2(2 * wires.size()),
"The size of matrix does not match with the given "
"number of wires");
applyMatrix(matrix.data(), wires, inverse);
}
inline void applyControlledMatrix(
ComplexT *matrix, const std::vector<std::size_t> &controlled_wires,
const std::vector<bool> &controlled_values,
const std::vector<std::size_t> &wires, bool inverse = false) {
PL_ABORT_IF(wires.empty(), "Number of wires must be larger than 0");
const std::size_t n2 = exp2(wires.size() * 2);
KokkosVector matrix_("matrix_", n2);
Kokkos::deep_copy(matrix_, UnmanagedComplexHostView(matrix, n2));
applyNCMultiQubitOp(matrix_, controlled_wires, controlled_values, wires,
inverse);
}
inline void
applyControlledMatrix(const ComplexT *matrix,
const std::vector<std::size_t> &controlled_wires,
const std::vector<bool> &controlled_values,
const std::vector<std::size_t> &wires,
bool inverse = false) {
PL_ABORT_IF(wires.empty(), "Number of wires must be larger than 0");
const std::size_t n2 = exp2(wires.size() * 2);
KokkosVector matrix_("matrix_", n2);
Kokkos::deep_copy(matrix_, UnmanagedConstComplexHostView(matrix, n2));
applyNCMultiQubitOp(matrix_, controlled_wires, controlled_values, wires,
inverse);
}
inline void
applyControlledMatrix(const std::vector<ComplexT> &matrix,
const std::vector<std::size_t> &controlled_wires,
const std::vector<bool> &controlled_values,
const std::vector<std::size_t> &wires,
bool inverse = false) {
PL_ABORT_IF(wires.empty(), "Number of wires must be larger than 0");
PL_ABORT_IF(matrix.size() != exp2(2 * wires.size()),
"The size of matrix does not match with the given "
"number of wires");
applyControlledMatrix(matrix.data(), controlled_wires,
controlled_values, wires, inverse);
}
auto applyGenerator(const std::string &opName,
const std::vector<std::size_t> &wires,
bool inverse = false) -> PrecisionT {
const std::size_t num_qubits = this->getNumQubits();
const GeneratorOperation generator_op =
reverse_lookup(generator_names, std::string_view{opName});
return applyNamedGenerator<KokkosExecSpace>(generator_op, *data_,
num_qubits, wires, inverse);
}
auto
applyControlledGenerator(const std::string &opName,
const std::vector<std::size_t> &controlled_wires,
const std::vector<bool> &controlled_values,
const std::vector<std::size_t> &wires,
bool inverse = false) -> PrecisionT {
const std::size_t num_qubits = this->getNumQubits();
const ControlledGeneratorOperation generator_op = reverse_lookup(
controlled_generator_names, std::string_view{opName});
return applyNCNamedGenerator<KokkosExecSpace>(
generator_op, *data_, num_qubits, controlled_wires,
controlled_values, wires, inverse);
}
auto applyGenerator(const std::string &opName,
const std::vector<std::size_t> &controlled_wires,
const std::vector<bool> &controlled_values,
const std::vector<std::size_t> &wires,
bool inverse = false) -> PrecisionT {
PL_ABORT_IF_NOT(
areVecsDisjoint<std::size_t>(controlled_wires, wires),
"`controlled_wires` and `target wires` must be disjoint.");
PL_ABORT_IF_NOT(controlled_wires.size() == controlled_values.size(),
"`controlled_wires` must have the same size as "
"`controlled_values`.");
if (controlled_wires.empty()) {
return applyGenerator(opName, wires, inverse);
}
return applyControlledGenerator(opName, controlled_wires,
controlled_values, wires, inverse);
}
void collapse(std::size_t wire, bool branch) {
KokkosVector matrix("gate_matrix", 4);
Kokkos::parallel_for(
matrix.size(), KOKKOS_LAMBDA(std::size_t k) {
matrix(k) = ((k == 0 && branch == 0) || (k == 3 && branch == 1))
? ComplexT{1.0, 0.0}
: ComplexT{0.0, 0.0};
});
applyMultiQubitOp(matrix, {wire}, false);
normalize();
}
void normalize() {
auto sv_view = getView();
PrecisionT squaredNorm = 0.0;
Kokkos::parallel_reduce(
sv_view.size(),
KOKKOS_LAMBDA(std::size_t i, PrecisionT &sum) {
const PrecisionT norm = Kokkos::abs(sv_view(i));
sum += norm * norm;
},
squaredNorm);
PL_ABORT_IF(squaredNorm <
std::numeric_limits<PrecisionT>::epsilon() * 1e2,
"Vector has norm close to zero and cannot be normalized");
const std::complex<PrecisionT> inv_norm =
1. / Kokkos::sqrt(squaredNorm);
Kokkos::parallel_for(
sv_view.size(),
KOKKOS_LAMBDA(std::size_t i) { sv_view(i) *= inv_norm; });
}
void updateData(const KokkosVector other) {
Kokkos::deep_copy(*data_, other);
}
void updateData(const StateVectorKokkos<fp_t> &other) {
updateData(other.getView());
}
void updateData(ComplexT *new_data, std::size_t new_size) {
updateData(KokkosVector(new_data, new_size));
}
void updateData(std::vector<ComplexT> &other) {
updateData(other.data(), other.size());
}
[[nodiscard]] auto getData() -> ComplexT * { return getView().data(); }
[[nodiscard]] auto getData() const -> const ComplexT * {
return getView().data();
}
[[nodiscard]] auto getView() const -> KokkosVector & { return *data_; }
[[nodiscard]] auto getView() -> KokkosVector & { return *data_; }
[[nodiscard]] auto getDataVector() -> std::vector<ComplexT> {
return view2vector(getView());
}
[[nodiscard]] auto getDataVector() const -> const std::vector<ComplexT> {
return view2vector(getView());
}
inline void HostToDevice(ComplexT *sv, std::size_t length) {
Kokkos::deep_copy(*data_, UnmanagedComplexHostView(sv, length));
}
inline void DeviceToHost(ComplexT *sv, std::size_t length) const {
Kokkos::deep_copy(UnmanagedComplexHostView(sv, length), *data_);
}
inline void DeviceToDevice(KokkosVector vector_to_copy) {
Kokkos::deep_copy(*data_, vector_to_copy);
}
private:
std::size_t num_qubits_;
std::mutex init_mutex_;
std::unique_ptr<KokkosVector> data_;
inline static bool is_exit_reg_ = false;
};
}; // namespace Pennylane::LightningKokkos
api/program_listing_file_pennylane_lightning_core_simulators_lightning_kokkos_StateVectorKokkos.hpp
Download Python script
Download Notebook
View on GitHub