Skip to content

Commit

Permalink
fix statevector test
Browse files Browse the repository at this point in the history
  • Loading branch information
@gandalfr-KY
gandalfr-KY committed Jan 24, 2025
1 parent 3642eb2 commit ac2197c
Show file tree
Hide file tree
Showing 6 changed files with 241 additions and 339 deletions.
20 changes: 16 additions & 4 deletions include/scaluq/util/utility.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -135,12 +135,24 @@ inline internal::ComplexMatrix<Fp> get_expanded_matrix(
}

// Host std::vector を Device Kokkos::View に変換する関数
template <typename T>
Kokkos::View<T*> convert_host_vector_to_device_view(const std::vector<T>& vec);
template <typename T, ExecutionSpace Sp>
Kokkos::View<T*, Sp> convert_vector_to_view(const std::vector<T>& vec) {
Kokkos::View<const T*, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> host_view(
vec.data(), vec.size());
Kokkos::View<T*, Sp> device_view("device_view", vec.size());
Kokkos::deep_copy(device_view, host_view);
return device_view;
}

// Device Kokkos::View を Host std::vector に変換する関数
template <typename T>
std::vector<T> convert_device_view_to_host_vector(const Kokkos::View<T*>& device_view);
template <typename T, ExecutionSpace Sp>
std::vector<T> convert_view_to_vector(const Kokkos::View<T*, Sp>& device_view) {
std::vector<T> host_vector(device_view.extent(0));
Kokkos::View<T*, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> host_view(
host_vector.data(), host_vector.size());
Kokkos::deep_copy(host_view, device_view);
return host_vector;
}

// Device Kokkos::View を Host std::vector に変換する関数
template <typename T, typename Layout>
Expand Down
8 changes: 3 additions & 5 deletions src/operator/operator.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -183,11 +183,9 @@ Complex<Fp> Operator<Fp, Sp>::get_transition_amplitude(
_terms.begin(), _terms.end(), coefs_vector.begin(), [](const PauliOperator<Fp, Sp>& pauli) {
return pauli._ptr->_coef;
});
Kokkos::View<std::uint64_t*> bmasks =
internal::convert_host_vector_to_device_view(bmasks_vector);
Kokkos::View<std::uint64_t*> pmasks =
internal::convert_host_vector_to_device_view(pmasks_vector);
Kokkos::View<Complex<Fp>*> coefs = internal::convert_host_vector_to_device_view(coefs_vector);
auto bmasks = internal::convert_vector_to_view<std::uint64_t, Sp>(bmasks_vector);
auto pmasks = internal::convert_vector_to_view<std::uint64_t, Sp>(pmasks_vector);
auto coefs = internal::convert_vector_to_view<Complex<Fp>, Sp>(coefs_vector);
std::uint64_t dim = state_vector_bra.dim();
Complex<Fp> res;
Kokkos::parallel_reduce(
Expand Down
26 changes: 11 additions & 15 deletions src/state/state_vector.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ void StateVector<Fp, Sp>::set_Haar_random_state(std::uint64_t n_qubits, std::uin
FLOAT_AND_SPACE(Fp, Sp)
[[nodiscard]] std::vector<Complex<Fp>> StateVector<Fp, Sp>::get_amplitudes() const {
if constexpr (std::is_same_v<Sp, DefaultSpace>) {
return internal::convert_device_view_to_host_vector<ComplexType>(_raw);
return internal::convert_view_to_vector<ComplexType>(_raw);
} else {
return std::vector<Complex<Fp>>(_raw.data(), _raw.data() + _raw.size());
}
Expand Down Expand Up @@ -132,8 +132,8 @@ Fp StateVector<Fp, Sp>::get_marginal_probability(
}

Fp sum = 0.;
auto d_target_index = internal::convert_host_vector_to_device_view(target_index);
auto d_target_value = internal::convert_host_vector_to_device_view(target_value);
auto d_target_index = internal::convert_vector_to_view<std::uint64_t, Sp>(target_index);
auto d_target_value = internal::convert_vector_to_view<std::uint64_t, Sp>(target_value);

Kokkos::parallel_reduce(
"marginal_prob",
Expand Down Expand Up @@ -169,33 +169,33 @@ Fp StateVector<Fp, Sp>::get_entropy() const {
FLOAT_AND_SPACE(Fp, Sp)
void StateVector<Fp, Sp>::add_state_vector_with_coef(ComplexType coef, const StateVector& state) {
Kokkos::parallel_for(
Kokkos::RangePolicy<Sp>(0, this->_dim >> 1),
Kokkos::RangePolicy<Sp>(0, this->_dim),
KOKKOS_CLASS_LAMBDA(std::uint64_t i) { this->_raw[i] += coef * state._raw[i]; });
Kokkos::fence();
}
FLOAT_AND_SPACE(Fp, Sp)
void StateVector<Fp, Sp>::multiply_coef(ComplexType coef) {
Kokkos::parallel_for(
Kokkos::RangePolicy<Sp>(0, this->_dim >> 1),
Kokkos::RangePolicy<Sp>(0, this->_dim),
KOKKOS_CLASS_LAMBDA(std::uint64_t i) { this->_raw[i] *= coef; });
Kokkos::fence();
}
FLOAT_AND_SPACE(Fp, Sp)
std::vector<std::uint64_t> StateVector<Fp, Sp>::sampling(std::uint64_t sampling_count,
std::uint64_t seed) const {
Kokkos::View<Fp*> stacked_prob("prob", _dim + 1);
Kokkos::View<Fp*, Sp> stacked_prob("prob", _dim + 1);
Kokkos::parallel_scan(
"compute_stacked_prob",
Kokkos::RangePolicy<Sp>(0, this->_dim >> 1),
Kokkos::RangePolicy<Sp>(0, this->_dim),
KOKKOS_CLASS_LAMBDA(std::uint64_t i, Fp & update, const bool final) {
update += internal::squared_norm(this->_raw[i]);
if (final) {
stacked_prob[i + 1] = update;
}
});

Kokkos::View<std::uint64_t*> result(Kokkos::ViewAllocateWithoutInitializing("result"),
sampling_count);
Kokkos::View<std::uint64_t*, Sp> result(Kokkos::ViewAllocateWithoutInitializing("result"),
sampling_count);
Kokkos::Random_XorShift64_Pool<Sp> rand_pool(seed);
Kokkos::parallel_for(
Kokkos::RangePolicy<Sp>(0, sampling_count), KOKKOS_LAMBDA(std::uint64_t i) {
Expand All @@ -214,7 +214,7 @@ std::vector<std::uint64_t> StateVector<Fp, Sp>::sampling(std::uint64_t sampling_
rand_pool.free_state(rand_gen);
});
Kokkos::fence();
return internal::convert_device_view_to_host_vector<std::uint64_t>(result);
return internal::convert_view_to_vector<std::uint64_t, Sp>(result);
}
FLOAT_AND_SPACE(Fp, Sp)
void StateVector<Fp, Sp>::load(const std::vector<Complex<Fp>>& other) {
Expand All @@ -223,11 +223,7 @@ void StateVector<Fp, Sp>::load(const std::vector<Complex<Fp>>& other) {
"Error: StateVector::load(const vector<ComplexType>&): invalid "
"length of state");
}
if constexpr (std::is_same_v<Sp, DefaultSpace>) {
_raw = internal::convert_host_vector_to_device_view(other);
} else {
std::copy(other.begin(), other.end(), _raw.data());
}
_raw = internal::convert_vector_to_view<Complex<Fp>, Sp>(other);
}
FLOAT_AND_SPACE(Fp, Sp)
StateVector<Fp, Sp> StateVector<Fp, Sp>::copy() const {
Expand Down
12 changes: 6 additions & 6 deletions src/state/state_vector_batched.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -184,7 +184,7 @@ std::vector<Fp> StateVectorBatched<Fp, Sp>::get_squared_norm() const {
Kokkos::single(Kokkos::PerTeam(team), [&] { norms[batch_id] = nrm; });
});
Kokkos::fence();
return internal::convert_device_view_to_host_vector<Fp>(norms);
return internal::convert_view_to_vector<Fp, Sp>(norms);
}

FLOAT_AND_SPACE(Fp, Sp)
Expand Down Expand Up @@ -234,7 +234,7 @@ std::vector<Fp> StateVectorBatched<Fp, Sp>::get_zero_probability(
Kokkos::single(Kokkos::PerTeam(team), [&] { probs[batch_id] = sum; });
});
Kokkos::fence();
return internal::convert_device_view_to_host_vector<Fp>(probs);
return internal::convert_view_to_vector<Fp, Sp>(probs);
}

FLOAT_AND_SPACE(Fp, Sp)
Expand Down Expand Up @@ -262,8 +262,8 @@ std::vector<Fp> StateVectorBatched<Fp, Sp>::get_marginal_probability(
}
}

auto target_index_d = internal::convert_host_vector_to_device_view(target_index);
auto target_value_d = internal::convert_host_vector_to_device_view(target_value);
auto target_index_d = internal::convert_vector_to_view<std::uint64_t, Sp>(target_index);
auto target_value_d = internal::convert_vector_to_view<std::uint64_t, Sp>(target_value);
Kokkos::View<Fp*> probs("probs", _batch_size);
Kokkos::parallel_for(
Kokkos::TeamPolicy<Sp>(Sp(), _batch_size, Kokkos::AUTO),
Expand All @@ -286,7 +286,7 @@ std::vector<Fp> StateVectorBatched<Fp, Sp>::get_marginal_probability(
Kokkos::single(Kokkos::PerTeam(team), [&] { probs[batch_id] = sum; });
});
Kokkos::fence();
return internal::convert_device_view_to_host_vector<Fp>(probs);
return internal::convert_view_to_vector<Fp, Sp>(probs);
}

FLOAT_AND_SPACE(Fp, Sp)
Expand All @@ -310,7 +310,7 @@ std::vector<Fp> StateVectorBatched<Fp, Sp>::get_entropy() const {
Kokkos::single(Kokkos::PerTeam(team), [&] { ents[batch_id] = sum; });
});
Kokkos::fence();
return internal::convert_device_view_to_host_vector(ents);
return internal::convert_view_to_vector<Fp, Sp>(ents);
}

FLOAT_AND_SPACE(Fp, Sp)
Expand Down
31 changes: 0 additions & 31 deletions src/util/utility.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,37 +4,6 @@

namespace scaluq {
namespace internal {
// Host std::vector を Device Kokkos::View に変換する関数
template <typename T>
Kokkos::View<T*> convert_host_vector_to_device_view(const std::vector<T>& vec) {
Kokkos::View<const T*, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> host_view(
vec.data(), vec.size());
Kokkos::View<T*> device_view("device_view", vec.size());
Kokkos::deep_copy(device_view, host_view);
return device_view;
}
#define FUNC_MACRO(T) \
template Kokkos::View<T*> convert_host_vector_to_device_view(const std::vector<T>&);
CALL_MACRO_FOR_FLOAT(FUNC_MACRO)
CALL_MACRO_FOR_COMPLEX(FUNC_MACRO)
CALL_MACRO_FOR_UINT(FUNC_MACRO)
#undef FUNC_MACRO

// Device Kokkos::View を Host std::vector に変換する関数
template <typename T>
std::vector<T> convert_device_view_to_host_vector(const Kokkos::View<T*>& device_view) {
std::vector<T> host_vector(device_view.extent(0));
Kokkos::View<T*, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> host_view(
host_vector.data(), host_vector.size());
Kokkos::deep_copy(host_view, device_view);
return host_vector;
}
#define FUNC_MACRO(T) \
template std::vector<T> convert_device_view_to_host_vector(const Kokkos::View<T*>&);
CALL_MACRO_FOR_FLOAT(FUNC_MACRO)
CALL_MACRO_FOR_COMPLEX(FUNC_MACRO)
CALL_MACRO_FOR_UINT(FUNC_MACRO)
#undef FUNC_MACRO

// Device Kokkos::View を Host std::vector に変換する関数
template <typename T, typename Layout>
Expand Down
Loading

0 comments on commit ac2197c

Please sign in to comment.