Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Allow multiple receivers in an element. #425

Merged
merged 5 commits into from
Jan 27, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 16 additions & 24 deletions include/compute/receivers/receivers.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -329,9 +329,10 @@ struct receivers : public impl::StationIterator,
* @return Kokkos::View<int *, Kokkos::DefaultExecutionSpace> View of the
* elements indices associated with the receivers
*/
Kokkos::View<int *, Kokkos::DefaultExecutionSpace>
get_elements_on_device(const specfem::element::medium_tag medium,
const specfem::element::property_tag property) const;
std::tuple<Kokkos::View<int *, Kokkos::DefaultExecutionSpace>,
Kokkos::View<int *, Kokkos::DefaultExecutionSpace> >
get_indices_on_device(const specfem::element::medium_tag medium,
const specfem::element::property_tag property) const;

/**
* @brief Get the spectral element indices in which the receivers are located
Expand All @@ -344,9 +345,10 @@ struct receivers : public impl::StationIterator,
* @return Kokkos::View<int *, Kokkos::DefaultExecutionSpace> View of the
* elements indices associated with the receivers
*/
Kokkos::View<int *, Kokkos::DefaultHostExecutionSpace>
get_elements_on_host(const specfem::element::medium_tag medium,
const specfem::element::property_tag property) const;
std::tuple<Kokkos::View<int *, Kokkos::DefaultHostExecutionSpace>,
Kokkos::View<int *, Kokkos::DefaultHostExecutionSpace> >
get_indices_on_host(const specfem::element::medium_tag medium,
const specfem::element::property_tag property) const;

/**
* @brief Get the seismogram types
Expand All @@ -364,8 +366,6 @@ struct receivers : public impl::StationIterator,
IndexViewType::HostMirror h_elements; ///< Host view to store the
///< elements associated with the
///< receivers
IndexViewType receiver_domain_index_mapping;
IndexViewType::HostMirror h_receiver_domain_index_mapping;
LagrangeInterpolantType lagrange_interpolant; ///< Lagrange interpolant for
///< every receiver
LagrangeInterpolantType::HostMirror
Expand All @@ -380,6 +380,12 @@ struct receivers : public impl::StationIterator,
GET_NAME(PROPERTY_TAG)); \
IndexViewType::HostMirror CREATE_VARIABLE_NAME( \
h_elements, GET_NAME(DIMENSION_TAG), GET_NAME(MEDIUM_TAG), \
GET_NAME(PROPERTY_TAG)); \
IndexViewType CREATE_VARIABLE_NAME( \
receiver_indices, GET_NAME(DIMENSION_TAG), GET_NAME(MEDIUM_TAG), \
GET_NAME(PROPERTY_TAG)); \
IndexViewType::HostMirror CREATE_VARIABLE_NAME( \
h_receiver_indices, GET_NAME(DIMENSION_TAG), GET_NAME(MEDIUM_TAG), \
GET_NAME(PROPERTY_TAG));

CALL_MACRO_FOR_ALL_MATERIAL_SYSTEMS(
Expand Down Expand Up @@ -439,16 +445,9 @@ load_on_device(const MemberType &team_member, const IteratorType &iterator,
Kokkos::abort(message.c_str());
}

if (receivers.receiver_domain_index_mapping(index.ispec) == -1) {
std::string message = "Invalid element detected in kernel at " +
std::string(__FILE__) + ":" +
std::to_string(__LINE__);
Kokkos::abort(message.c_str());
}

#endif

const int irec = receivers.receiver_domain_index_mapping(index.ispec);
const int irec = iterator_index.imap;

lagrange_interpolant(iterator_index.ielement, index.iz, index.ix, 0) =
receivers.lagrange_interpolant(irec, index.iz, index.ix, 0);
Expand Down Expand Up @@ -500,16 +499,9 @@ store_on_device(const MemberType &team_member, const IteratorType &iterator,
Kokkos::abort(message.c_str());
}

if (receivers.receiver_domain_index_mapping(index.ispec) == -1) {
std::string message = "Invalid element detected in kernel at " +
std::string(__FILE__) +
std::to_string(__LINE__);
Kokkos::abort(message.c_str());
}

#endif

const int irec = receivers.receiver_domain_index_mapping(index.ispec);
const int irec = iterator_index.imap;

receivers.seismogram_components(isig_step, iseis, irec, 0) =
seismogram_components(iterator_index.ielement, 0);
Expand Down
16 changes: 8 additions & 8 deletions include/kokkos_kernels/impl/compute_seismogram.tpp
Original file line number Diff line number Diff line change
Expand Up @@ -27,12 +27,12 @@ void specfem::kokkos_kernels::impl::compute_seismograms(
constexpr int ngll = NGLL;
constexpr auto dimension = DimensionType;

const auto elements =
assembly.receivers.get_elements_on_device(MediumTag, PropertyTag);
const auto [elements, receiver_indices] =
assembly.receivers.get_indices_on_device(MediumTag, PropertyTag);

const int nelements = elements.extent(0);
const int nreceivers = receiver_indices.extent(0);

if (nelements == 0)
if (nreceivers == 0)
return;

auto &receivers = assembly.receivers;
Expand Down Expand Up @@ -61,7 +61,7 @@ void specfem::kokkos_kernels::impl::compute_seismograms(
lane_size, no_simd,
Kokkos::DefaultExecutionSpace>;

using ChunkPolicy = specfem::policy::element_chunk<ParallelConfig>;
using ChunkPolicy = specfem::policy::mapped_element_chunk<ParallelConfig>;
using ChunkElementFieldType = specfem::chunk_element::field<
ParallelConfig::chunk_size, ngll, DimensionType, MediumTag,
specfem::kokkos::DevScratchSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>,
Expand All @@ -84,7 +84,7 @@ void specfem::kokkos_kernels::impl::compute_seismograms(

receivers.set_seismogram_step(isig_step);

ChunkPolicy policy(elements, ngll, ngll);
ChunkPolicy policy(elements, receiver_indices, ngll, ngll);
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I ended up updating the naming of the policy, index, iterator etc. for more clarity:

policy -> mapped_policy
iterator -> mapped_iterator
index -> element_index

But it is fine the way it is here


for (int iseis = 0; iseis < nseismograms; ++iseis) {

Expand Down Expand Up @@ -112,12 +112,12 @@ void specfem::kokkos_kernels::impl::compute_seismograms(
team_member.league_rank() * ChunkPolicy::tile_size * simd_size +
tile;

if (starting_element_index >= nelements) {
if (starting_element_index >= nreceivers) {
break;
}

const auto iterator =
policy.league_iterator(starting_element_index);
policy.mapped_league_iterator(starting_element_index);

specfem::compute::load_on_device(team_member, iterator, field,
element_field);
Expand Down
91 changes: 87 additions & 4 deletions include/policies/chunk.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,26 @@ struct chunk_index_type<false, DimensionType> {
const specfem::point::index<dimension> index)
: ielement(ielement), index(index){};
};

/**
* @brief Struct to store the index of a quadrature point generated by chunk
* policy.
*
* @tparam UseSIMD Indicates whether SIMD is used or not.
* @tparam DimensionType Dimension type of the elements within this iterator.
*/
template <bool UseSIMD, specfem::dimension::type DimensionType>
struct mapped_chunk_index_type
: public chunk_index_type<UseSIMD, DimensionType> {
using Base = chunk_index_type<UseSIMD, DimensionType>;
int imap; ///< Index of the mapped element

mapped_chunk_index_type(const int ielement,
const specfem::point::index<DimensionType> index,
const int imap)
: Base(ielement, index), imap(imap) {}
};

} // namespace impl

/**
Expand Down Expand Up @@ -93,7 +113,7 @@ class chunk<ViewType, specfem::dimension::type::dim2, SIMD> {
///< type
///@}

private:
protected:
constexpr static bool using_simd = simd::using_simd;
constexpr static int simd_size = simd::size();

Expand Down Expand Up @@ -223,6 +243,35 @@ class chunk<ViewType, specfem::dimension::type::dim2, SIMD> {
return Kokkos::make_pair(indices(0), indices(num_elements - 1) + 1);
}
};

template <typename ViewType, specfem::dimension::type DimensionType,
typename SIMD>
class mapped_chunk;

template <typename ViewType, typename SIMD>
class mapped_chunk<ViewType, specfem::dimension::type::dim2, SIMD>
: public chunk<ViewType, specfem::dimension::type::dim2, SIMD> {
using Base = chunk<ViewType, specfem::dimension::type::dim2, SIMD>;
using mapped_index_type =
typename impl::mapped_chunk_index_type<Base::using_simd,
Base::dimension>; ///< Index

public:
mapped_chunk(const ViewType &indices, const ViewType &mapping,
const int ngllz, const int ngllx)
: Base(indices, ngllz, ngllx), mapping(mapping) {}

KOKKOS_INLINE_FUNCTION
mapped_index_type operator()(const int i) const {
const auto base_index = Base::operator()(i);
return mapped_index_type(base_index.ielement, base_index.index,
mapping(base_index.ielement));
}

private:
ViewType mapping;
};

} // namespace iterator

namespace policy {
Expand All @@ -237,7 +286,7 @@ template <typename ParallelConfig>
struct element_chunk
: public Kokkos::TeamPolicy<typename ParallelConfig::execution_space> {

private:
protected:
using IndexViewType = Kokkos::View<
int *,
typename ParallelConfig::execution_space::memory_space>; ///< View
Expand Down Expand Up @@ -289,7 +338,7 @@ struct element_chunk
true; ///< Indicates that this is a Kokkos team policy
///@}

private:
protected:
constexpr static int simd_size = simd::size();
constexpr static bool using_simd = simd::using_simd;

Expand Down Expand Up @@ -345,10 +394,44 @@ struct element_chunk
return iterator_type(my_indices, ngllz, ngllx);
}

private:
protected:
IndexViewType elements; ///< View of elements
int ngllz; ///< Number of GLL points in the z-direction
int ngllx; ///< Number of GLL points in the x-direction
};

template <typename ParallelConfig>
struct mapped_element_chunk : public element_chunk<ParallelConfig> {
using simd = typename ParallelConfig::simd;
using Base = element_chunk<ParallelConfig>;
using IndexViewType = typename Base::IndexViewType;

using mapped_iterator_type =
specfem::iterator::mapped_chunk<IndexViewType, ParallelConfig::dimension,
simd>; ///< Iterator

mapped_element_chunk(const IndexViewType &view, const IndexViewType &mapping,
int ngllz, int ngllx)
: Base(view, ngllz, ngllx), mapping(mapping) {}

KOKKOS_INLINE_FUNCTION
mapped_iterator_type mapped_league_iterator(const int start_index) const {
const int start = start_index;
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  1. Can this just be called league_iterator?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I got the error mapped_element_chunk and element_chunk are not compatible when trying to override league_iterator.

const int end =
(start + Base::chunk_size * Base::simd_size > Base::elements.extent(0))
? Base::elements.extent(0)
: start + Base::chunk_size * Base::simd_size;
const auto elem_indices =
Kokkos::subview(Base::elements, Kokkos::make_pair(start, end));
const auto map_indices =
Kokkos::subview(mapping, Kokkos::make_pair(start, end));
return mapped_iterator_type(elem_indices, map_indices, Base::ngllz,
Base::ngllx);
}

protected:
IndexViewType mapping; ///< View of elements
};

} // namespace policy
} // namespace specfem
56 changes: 32 additions & 24 deletions src/compute/compute_receivers.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -24,19 +24,11 @@ specfem::compute::receivers::receivers(
h_lagrange_interpolant(Kokkos::create_mirror_view(lagrange_interpolant)),
elements("specfem::compute::receivers::elements", receivers.size()),
h_elements(Kokkos::create_mirror_view(elements)),
receiver_domain_index_mapping(
"specfem::compute::receivers::receiver_domain_index_mapping", nspec),
h_receiver_domain_index_mapping(
Kokkos::create_mirror_view(receiver_domain_index_mapping)),
element_types(element_types), impl::StationIterator(receivers.size(),
stypes.size()),
impl::SeismogramIterator(receivers.size(), stypes.size(), max_sig_step,
dt, t0, nsteps_between_samples) {

for (int ispec = 0; ispec < nspec; ++ispec) {
h_receiver_domain_index_mapping(ispec) = -1;
}

for (int isies = 0; isies < stypes.size(); ++isies) {
auto seis_type = stypes[isies];
switch (seis_type) {
Expand Down Expand Up @@ -75,12 +67,6 @@ specfem::compute::receivers::receivers(
};
const auto lcoord = specfem::algorithms::locate_point(gcoord, mesh);

if (h_receiver_domain_index_mapping(lcoord.ispec) != -1) {
throw std::runtime_error(
"Multiple receivers are detected in the same element");
}

h_receiver_domain_index_mapping(lcoord.ispec) = ireceiver;
h_elements(ireceiver) = lcoord.ispec;

const auto xi = mesh.quadratures.gll.h_xi;
Expand Down Expand Up @@ -143,6 +129,17 @@ specfem::compute::receivers::receivers(
GET_NAME(MEDIUM_TAG), GET_NAME(PROPERTY_TAG)) = \
Kokkos::create_mirror_view( \
CREATE_VARIABLE_NAME(elements, GET_NAME(DIMENTION_TAG), \
GET_NAME(MEDIUM_TAG), GET_NAME(PROPERTY_TAG))); \
CREATE_VARIABLE_NAME(receiver_indices, GET_NAME(DIMENTION_TAG), \
GET_NAME(MEDIUM_TAG), GET_NAME(PROPERTY_TAG)) = \
IndexViewType("specfem::compute::receivers::elements", \
CREATE_VARIABLE_NAME(count, GET_NAME(DIMENTION_TAG), \
GET_NAME(MEDIUM_TAG), \
GET_NAME(PROPERTY_TAG))); \
CREATE_VARIABLE_NAME(h_receiver_indices, GET_NAME(DIMENTION_TAG), \
GET_NAME(MEDIUM_TAG), GET_NAME(PROPERTY_TAG)) = \
Kokkos::create_mirror_view( \
CREATE_VARIABLE_NAME(receiver_indices, GET_NAME(DIMENTION_TAG), \
GET_NAME(MEDIUM_TAG), GET_NAME(PROPERTY_TAG)));

CALL_MACRO_FOR_ALL_MATERIAL_SYSTEMS(
Expand All @@ -165,6 +162,11 @@ specfem::compute::receivers::receivers(
(CREATE_VARIABLE_NAME(index, GET_NAME(DIMENTION_TAG), \
GET_NAME(MEDIUM_TAG), GET_NAME(PROPERTY_TAG))) = \
ispec; \
CREATE_VARIABLE_NAME(h_receiver_indices, GET_NAME(DIMENTION_TAG), \
GET_NAME(MEDIUM_TAG), GET_NAME(PROPERTY_TAG)) \
(CREATE_VARIABLE_NAME(index, GET_NAME(DIMENTION_TAG), \
GET_NAME(MEDIUM_TAG), GET_NAME(PROPERTY_TAG))) = \
ireceiver; \
CREATE_VARIABLE_NAME(index, GET_NAME(DIMENTION_TAG), \
GET_NAME(MEDIUM_TAG), GET_NAME(PROPERTY_TAG)) \
++; \
Expand All @@ -185,22 +187,24 @@ specfem::compute::receivers::receivers(

Kokkos::deep_copy(lagrange_interpolant, h_lagrange_interpolant);
Kokkos::deep_copy(elements, h_elements);
Kokkos::deep_copy(receiver_domain_index_mapping,
h_receiver_domain_index_mapping);

return;
}

Kokkos::View<int *, Kokkos::DefaultHostExecutionSpace>
specfem::compute::receivers::get_elements_on_host(
std::tuple<Kokkos::View<int *, Kokkos::DefaultHostExecutionSpace>,
Kokkos::View<int *, Kokkos::DefaultHostExecutionSpace> >
specfem::compute::receivers::get_indices_on_host(
const specfem::element::medium_tag medium_tag,
const specfem::element::property_tag property_tag) const {

#define RETURN_VALUE(DIMENTION_TAG, MEDIUM_TAG, PROPERTY_TAG) \
if (medium_tag == GET_TAG(MEDIUM_TAG) && \
property_tag == GET_TAG(PROPERTY_TAG)) { \
return CREATE_VARIABLE_NAME(h_elements, GET_NAME(DIMENTION_TAG), \
GET_NAME(MEDIUM_TAG), GET_NAME(PROPERTY_TAG)); \
return std::make_tuple( \
CREATE_VARIABLE_NAME(h_elements, GET_NAME(DIMENTION_TAG), \
GET_NAME(MEDIUM_TAG), GET_NAME(PROPERTY_TAG)), \
CREATE_VARIABLE_NAME(h_receiver_indices, GET_NAME(DIMENTION_TAG), \
GET_NAME(MEDIUM_TAG), GET_NAME(PROPERTY_TAG))); \
}

CALL_MACRO_FOR_ALL_MATERIAL_SYSTEMS(
Expand All @@ -211,16 +215,20 @@ specfem::compute::receivers::get_elements_on_host(
#undef RETURN_VALUE
}

Kokkos::View<int *, Kokkos::DefaultExecutionSpace>
specfem::compute::receivers::get_elements_on_device(
std::tuple<Kokkos::View<int *, Kokkos::DefaultExecutionSpace>,
Kokkos::View<int *, Kokkos::DefaultExecutionSpace> >
specfem::compute::receivers::get_indices_on_device(
const specfem::element::medium_tag medium_tag,
const specfem::element::property_tag property_tag) const {

#define RETURN_VALUE(DIMENTION_TAG, MEDIUM_TAG, PROPERTY_TAG) \
if (medium_tag == GET_TAG(MEDIUM_TAG) && \
property_tag == GET_TAG(PROPERTY_TAG)) { \
return CREATE_VARIABLE_NAME(elements, GET_NAME(DIMENTION_TAG), \
GET_NAME(MEDIUM_TAG), GET_NAME(PROPERTY_TAG)); \
return std::make_tuple( \
CREATE_VARIABLE_NAME(elements, GET_NAME(DIMENTION_TAG), \
GET_NAME(MEDIUM_TAG), GET_NAME(PROPERTY_TAG)), \
CREATE_VARIABLE_NAME(receiver_indices, GET_NAME(DIMENTION_TAG), \
GET_NAME(MEDIUM_TAG), GET_NAME(PROPERTY_TAG))); \
}

CALL_MACRO_FOR_ALL_MATERIAL_SYSTEMS(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ S0002 AA 2250.0000000 3000.0000000 0.0 0.0
S0003 AA 2500.0000000 3000.0000000 0.0 0.0
S0004 AA 2750.0000000 3000.0000000 0.0 0.0
S0005 AA 3000.0000000 3000.0000000 0.0 0.0
S0006 AA 2500.0000000 2500.0000000 0.0 0.0
S0007 AA 2500.0000000 2300.0000000 0.0 0.0
S0008 AA 2500.0000000 2100.0000000 0.0 0.0
S0009 AA 2500.0000000 1900.0000000 0.0 0.0
Expand Down
Loading