diff --git a/include/compute/sources/source_medium.hpp b/include/compute/sources/source_medium.hpp index 3f56d910..e5d0feba 100644 --- a/include/compute/sources/source_medium.hpp +++ b/include/compute/sources/source_medium.hpp @@ -95,36 +95,42 @@ struct source_medium { SourceArrayView source_array; ///< Lagrange interpolants for every source SourceArrayView::HostMirror h_source_array; ///< Host mirror of source_array - template + template KOKKOS_INLINE_FUNCTION void - load_on_device(const int timestep, const IndexType index, + load_on_device(const int timestep, const IteratorIndexType &iterator_index, PointSourceType &point_source) const { - // For the source it is important to remember that we are using the - // mapped index to access the element and source indices - // that means that index actually is a mapped_chunk_index - // and we need to use index.ispec to access the element index - // and index.imap to access the source index + /* For the source it is important to remember that we are using the + * mapped index to access the element and source indices + * that means that index actually is a mapped_chunk_index + * and we need to use index.ispec to access the element index + * and index.imap to access the source index + */ + const auto index = iterator_index.index; + const auto isource = iterator_index.imap; for (int component = 0; component < components; component++) { point_source.stf(component) = - source_time_function(timestep, index.imap, component); + source_time_function(timestep, isource, component); point_source.lagrange_interpolant(component) = - source_array(index.imap, component, index.iz, index.ix); + source_array(isource, component, index.iz, index.ix); } } - template + template KOKKOS_INLINE_FUNCTION void - // For the source it is important to remember that we are using the - // mapped index to access the element and source indices - // that means that index actually is a mapped_chunk_index - // and we need to use index.ispec to access the element index - // and index.imap to access the source index - store_on_device(const int timestep, const IndexType index, + store_on_device(const int timestep, const IteratorIndexType &iterator_index, const PointSourceType &point_source) const { + /* For the source it is important to remember that we are using the + * mapped index to access the element and source indices + * that means that index actually is a mapped_chunk_index + * and we need to use index.ispec to access the element index + * and index.imap to access the source index + */ + const auto index = iterator_index.index; + const auto isource = iterator_index.imap; for (int component = 0; component < components; component++) { - source_time_function(timestep, index.imap, component) = + source_time_function(timestep, isource, component) = point_source.stf(component); - source_array(index.imap, component, index.iz, index.ix) = + source_array(isource, component, index.iz, index.ix) = point_source.lagrange_interpolant(component); } } diff --git a/include/compute/sources/sources.hpp b/include/compute/sources/sources.hpp index b8adb6eb..de9d4417 100644 --- a/include/compute/sources/sources.hpp +++ b/include/compute/sources/sources.hpp @@ -250,12 +250,15 @@ struct sources { * @param sources Source information for the domain * @param point_source Point source object to load source information into */ -template +template KOKKOS_INLINE_FUNCTION void -load_on_device(const IndexType index, const specfem::compute::sources &sources, +load_on_device(const IteratorIndexType iterator_index, + const specfem::compute::sources &sources, PointSourceType &point_source) { - static_assert(IndexType::using_simd == false, + const auto index = iterator_index.index; + + static_assert(index.using_simd == false, "IndexType must not use SIMD when loading sources"); static_assert( @@ -265,7 +268,7 @@ load_on_device(const IndexType index, const specfem::compute::sources &sources, static_assert(PointSourceType::dimension == specfem::dimension::type::dim2, "PointSourceType must be a 2D point source type"); - static_assert(IndexType::dimension == specfem::dimension::type::dim2, + static_assert(index.dimension == specfem::dimension::type::dim2, "IndexType must be a 2D index type"); static_assert( @@ -295,16 +298,13 @@ load_on_device(const IndexType index, const specfem::compute::sources &sources, } #endif - IndexType lcoord = index; - // lcoord.ispec = index.imap; - #define SOURCE_MEDIUM_LOAD_ON_DEVICE(DIMENSION_TAG, MEDIUM_TAG) \ if constexpr (GET_TAG(DIMENSION_TAG) == specfem::dimension::type::dim2) { \ if constexpr (GET_TAG(MEDIUM_TAG) == PointSourceType::medium_tag) { \ sources \ .CREATE_VARIABLE_NAME(source, GET_NAME(DIMENSION_TAG), \ GET_NAME(MEDIUM_TAG)) \ - .load_on_device(sources.timestep, lcoord, point_source); \ + .load_on_device(sources.timestep, iterator_index, point_source); \ } \ } @@ -412,12 +412,15 @@ void load_on_host(const IndexType index, * @param point_source Point source object to load source information into * @param sources Source information for the domain */ -template +template KOKKOS_INLINE_FUNCTION void -store_on_device(const IndexType index, const PointSourceType &point_source, +store_on_device(const IteratorIndexType iterator_index, + const PointSourceType &point_source, const specfem::compute::sources &sources) { - static_assert(IndexType::using_simd == false, + const auto index = iterator_index.index; + + static_assert(index.using_simd == false, "IndexType must not use SIMD when storing sources"); static_assert( @@ -427,7 +430,7 @@ store_on_device(const IndexType index, const PointSourceType &point_source, static_assert(PointSourceType::dimension == specfem::dimension::type::dim2, "PointSourceType must be a 2D point source type"); - static_assert(IndexType::dimension == specfem::dimension::type::dim2, + static_assert(index.dimension == specfem::dimension::type::dim2, "IndexType must be a 2D index type"); static_assert( @@ -454,16 +457,13 @@ store_on_device(const IndexType index, const PointSourceType &point_source, } #endif - IndexType lcoord = index; - lcoord.ispec = sources.source_domain_index_mapping(index.ispec); - #define SOURCE_MEDIUM_STORE_ON_DEVICE(DIMENSION_TAG, MEDIUM_TAG) \ if constexpr (GET_TAG(DIMENSION_TAG) == specfem::dimension::type::dim2) { \ if constexpr (GET_TAG(MEDIUM_TAG) == PointSourceType::medium_tag) { \ sources \ .CREATE_VARIABLE_NAME(source, GET_NAME(DIMENSION_TAG), \ GET_NAME(MEDIUM_TAG)) \ - .store_on_device(sources.timestep, lcoord, point_source); \ + .store_on_device(sources.timestep, iterator_index, point_source); \ } \ } diff --git a/include/kokkos_kernels/impl/compute_source_interaction.tpp b/include/kokkos_kernels/impl/compute_source_interaction.tpp index 23b24666..26707150 100644 --- a/include/kokkos_kernels/impl/compute_source_interaction.tpp +++ b/include/kokkos_kernels/impl/compute_source_interaction.tpp @@ -90,18 +90,23 @@ Kokkos::parallel_for( break; } - const auto iterator = + // This is a mapped_chunk iterator + const auto mapped_chunk_iterator = policy.mapped_league_iterator(starting_element_index); Kokkos::parallel_for( - Kokkos::TeamThreadRange(team, iterator.chunk_size()), + Kokkos::TeamThreadRange(team, mapped_chunk_iterator.chunk_size()), [&](const int i) { - const auto mapped_iterator_index = iterator(i); - const auto mapped_index = mapped_iterator_index.index; + // mapped_chunk_index_type + const auto mapped_chunked_index = mapped_chunk_iterator(i); + // mapped_index is specfem::point::index + const auto mapped_index = mapped_chunked_index.index; + // need mapped_chunk_index here to get the imap=isource PointSourcesType point_source; - specfem::compute::load_on_device(mapped_iterator_index, sources, point_source); + specfem::compute::load_on_device(mapped_chunked_index, sources, + point_source); PointPropertiesType point_property; specfem::compute::load_on_device(mapped_index, properties, diff --git a/include/policies/chunk.hpp b/include/policies/chunk.hpp index fbfdb7ee..9699d0df 100644 --- a/include/policies/chunk.hpp +++ b/include/policies/chunk.hpp @@ -52,6 +52,26 @@ struct chunk_index_type { const specfem::point::index 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 +struct mapped_chunk_index_type + : public chunk_index_type { + using Base = chunk_index_type; + int imap; ///< Index of the mapped element + + mapped_chunk_index_type(const int ielement, + const specfem::point::index index, + const int imap) + : Base(ielement, index), imap(imap) {} +}; + } // namespace impl /** @@ -232,6 +252,9 @@ template class mapped_chunk : public chunk { using Base = chunk; + using mapped_index_type = + typename impl::mapped_chunk_index_type; ///< Index public: mapped_chunk(const ViewType &indices, const ViewType &mapping, @@ -239,16 +262,10 @@ class mapped_chunk : Base(indices, ngllz, ngllx), mapping(mapping) {} KOKKOS_INLINE_FUNCTION - const int imap(const int i) const { -#ifdef KOKKOS_ENABLE_CUDA - const int ielement = i % num_elements; - return mapping(ielement); -#else - const int ix = i % Base::ngllx; - const int iz = (i / Base::ngllx) % Base::ngllz; - const int ielement = i / (Base::ngllz * Base::ngllx); - return mapping(ielement); -#endif + 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: diff --git a/tests/unit-tests/assembly/sources/sources.cpp b/tests/unit-tests/assembly/sources/sources.cpp index 52aeb49a..82f7b2da 100644 --- a/tests/unit-tests/assembly/sources/sources.cpp +++ b/tests/unit-tests/assembly/sources/sources.cpp @@ -5,6 +5,7 @@ #include "enumerations/medium.hpp" #include "enumerations/wavefield.hpp" #include "point/sources.hpp" +#include "policies/chunk.hpp" #include "gtest/gtest.h" #include @@ -43,17 +44,23 @@ void check_store(specfem::compute::assembly &assembly) { Kokkos::deep_copy(values_to_store, h_values_to_store); - using PointType = specfem::point::source; - + using PointSourceType = + specfem::point::source; + using mapped_chunk_index_type = + specfem::iterator::impl::mapped_chunk_index_type< + false, specfem::dimension::type::dim2>; Kokkos::parallel_for( "check_store_on_device", Kokkos::MDRangePolicy >( { 0, 0, 0 }, { nelements, ngllz, ngllx }), KOKKOS_LAMBDA(const int &i, const int &iz, const int &ix) { const int ielement = element_indices(i); + const int isource = source_indices(i); const auto index = specfem::point::index(ielement, iz, ix); + const auto mapped_iterator_index = + mapped_chunk_index_type(ielement, index, isource); specfem::datatype::ScalarPointViewType stf; specfem::datatype::ScalarPointViewType @@ -62,8 +69,9 @@ void check_store(specfem::compute::assembly &assembly) { stf(ic) = values_to_store(i); lagrange_interpolant(ic) = values_to_store(i); } - PointType point(stf, lagrange_interpolant); - specfem::compute::store_on_device(index, point, sources); + PointSourceType point(stf, lagrange_interpolant); + specfem::compute::store_on_device(mapped_iterator_index, point, + sources); }); Kokkos::fence();