Merge pull request #1351 from PrincetonUniversity/issue-1350

Fix bug in coupled interfaces

Author: Rohit-Kakodkar

core/specfem/assembly/edge_types/dim2/edge_types.cpp

@@ -72,6 +72,7 @@ specfem::assembly::edge_types<specfem::dimension::type::dim2>::edge_types(
           std::vector<specfem::mesh_entity::edge<dimension_tag> >
               coupled_collect;
 
+          int edge_index = 0;
           for (const auto &edge :
                boost::make_iterator_range(boost::edges(nc_graph))) {
             const int ispec1 = boost::source(edge, nc_graph);
@@ -95,8 +96,11 @@ specfem::assembly::edge_types<specfem::dimension::type::dim2>::edge_types(
               count++;
               // we do not need orientation flipping -- that's handled by
               // the transfer function
-              self_collect.push_back({ ispec1, self_orientation, false });
-              coupled_collect.push_back({ ispec2, coupled_orientation, false });
+              self_collect.push_back(
+                  { ispec1, edge_index, self_orientation, false });
+              coupled_collect.push_back(
+                  { ispec2, edge_index, coupled_orientation, false });
+              edge_index++;
             }
           }
 
@@ -137,6 +141,7 @@ specfem::assembly::edge_types<specfem::dimension::type::dim2>::edge_types(
               coupled_interfaces.template get<self_medium, coupled_medium>();
           const int nedges =
               interface_container.num_interfaces; // number of edges
+          int edge_index = 0;
           for (int iedge = 0; iedge < nedges; ++iedge) {
             const int ispec1_mesh =
                 interface_container.medium1_index_mapping(iedge);
@@ -152,13 +157,14 @@ specfem::assembly::edge_types<specfem::dimension::type::dim2>::edge_types(
                   connection_mapping.flip_orientation(edge1, edge2);
               _h_self_edges_(index) =
                   specfem::mesh_entity::edge<specfem::dimension::type::dim2>{
-                    ispec1, edge1, false
+                    ispec1, edge_index, edge1, false
                   };
               _h_coupled_edges_(index) =
                   specfem::mesh_entity::edge<specfem::dimension::type::dim2>{
-                    ispec2, edge2, flip
+                    ispec2, edge_index, edge2, flip
                   };
               index++;
+              edge_index++;
             }
           }
           Kokkos::deep_copy(_self_edges_, _h_self_edges_);

core/specfem/timescheme/newmark.tpp

@@ -41,7 +41,8 @@ int corrector_phase_impl(
 
   specfem::execution::for_all(
       "specfem::TimeScheme::Newmark::corrector_phase_impl", range,
-      KOKKOS_LAMBDA(const IndexType &index) {
+      KOKKOS_LAMBDA(const typename decltype(range)::base_index_type &iterator_index) {
+        const auto index = iterator_index.get_index();
         PointAccelerationType acceleration;
         PointVelocityType velocity;
 
@@ -99,7 +100,8 @@ int predictor_phase_impl(
 
   specfem::execution::for_all(
       "specfem::TimeScheme::Newmark::corrector_phase_impl", range,
-      KOKKOS_LAMBDA(const IndexType &index) {
+      KOKKOS_LAMBDA(const typename decltype(range)::base_index_type &iterator_index) {
+        const auto index = iterator_index.get_index();
         PointDisplacementType displacement;
         PointVelocityType velocity;
         PointAccelerationType acceleration;

include/enumerations/dim2/mesh_entities.hpp

@@ -142,12 +142,14 @@ template <typename T> bool contains(const T &list, const dim2::type &value) {
 template <> struct edge<specfem::dimension::type::dim2> {
   specfem::mesh_entity::dim2::type edge_type;
   int ispec;
+  int iedge;
   bool reverse_orientation;
 
   KOKKOS_INLINE_FUNCTION
-  edge(const int ispec, const specfem::mesh_entity::dim2::type edge_type,
+  edge(const int ispec, const int iedge,
+       const specfem::mesh_entity::dim2::type edge_type,
        const bool reverse_orientation = false)
-      : edge_type(edge_type), ispec(ispec),
+      : edge_type(edge_type), ispec(ispec), iedge(iedge),
         reverse_orientation(reverse_orientation) {}
 
   KOKKOS_INLINE_FUNCTION

include/execution/chunked_domain_iterator.hpp

@@ -501,15 +501,14 @@ class ChunkedDomainIterator : public TeamPolicy<ParallelConfig> {
                           ///< specfem::execution::for_each_level
   using execution_space =
       typename base_type::execution_space; ///< Execution space type.
-  using base_index_type = specfem::point::index<
-      ParallelConfig::dimension,
-      ParallelConfig::simd::using_simd>; ///< Index type
-                                         ///< to be used
-                                         ///< when calling
-                                         ///< @ref
-                                         ///< specfem::execution::for_all
-                                         ///< with this
-                                         ///< iterator.
+  using base_index_type = PointIndex<
+      dimension_tag, typename ViewType::value_type,
+      ParallelConfig::simd::using_simd,
+      typename base_type::execution_space>; ///< Index type to be
+                                            ///< used when calling
+                                            ///< @ref
+                                            ///< specfem::execution::for_all
+                                            ///< with this iterator.
 
   /**
    * @brief Construct a new Chunked Domain Iterator object

include/execution/chunked_edge_iterator.hpp

@@ -33,7 +33,8 @@
  * // Create chunked iterator and process edges
  * specfem::execution::ChunkedEdgeIterator iterator(ParallelConfig(), edges,
  * num_points); specfem::execution::for_all("process_edges", iterator,
- *     KOKKOS_LAMBDA(const auto& index) {
+ *     KOKKOS_LAMBDA(const auto& iterator_index) {
+ *       const auto index = iterator_index.get_index();
  *         // Access edge point data
  *         int ispec = index.ispec;     // Element index
  *         int ipoint = index.ipoint;   // Point index along edge
@@ -80,7 +81,8 @@ namespace specfem::execution {
  * @code{.cpp}
  * // Typically used within chunked edge iterator lambdas
  * specfem::execution::for_all("process_edges", iterator,
- *     KOKKOS_LAMBDA(const EdgePointIndex<dim2, int, ExecutionSpace>& index) {
+ *     KOKKOS_LAMBDA(const auto & iterator_index) {
+*        const auto index = iterator_index.get_index();
  *         int element_id = index.ispec;    // Element containing this edge point
  *         int z_coord = index.iz;          // Local z-coordinate within element
  *         int x_coord = index.ix;          // Local x-coordinate within element
@@ -121,6 +123,11 @@ class EdgePointIndex {
   KOKKOS_INLINE_FUNCTION
   constexpr const index_type &get_index() const { return this->index; }
 
+  KOKKOS_INLINE_FUNCTION
+  constexpr const index_type &get_local_index() const {
+    return this->local_index; ///< Returns the local point index
+  }
+
   /**
    * @brief Constructor for EdgePointIndex
    *
@@ -130,8 +137,10 @@ class EdgePointIndex {
    */
   KOKKOS_INLINE_FUNCTION
   EdgePointIndex(const specfem::point::index<DimensionTag, false> &index,
-                 const int ipoint, const KokkosIndexType &kokkos_index)
-      : index(index.ispec, kokkos_index, ipoint, index.iz, index.ix),
+                 const int gedge, const int iedge, const int ipoint,
+                 const KokkosIndexType &kokkos_index)
+      : index(index.ispec, gedge, ipoint, index.iz, index.ix),
+        local_index(index.ispec, iedge, ipoint, index.iz, index.ix),
         kokkos_index(kokkos_index) {}
 
   /**
@@ -144,7 +153,8 @@ class EdgePointIndex {
   constexpr const iterator_type get_iterator() const { return iterator_type{}; }
 
 private:
-  index_type index; ///< Local element coordinates of the edge point
+  index_type index;       ///< Local element coordinates of the edge point
+  index_type local_index; ///< Local element coordinates relative to chunk
 
   KokkosIndexType kokkos_index; ///< Kokkos index type
 };
@@ -201,6 +211,7 @@ class ChunkEdgeIterator : public TeamThreadRangePolicy<TeamMemberType, int> {
     const int iedge = i % nedges;
     const int ipoint = i / nedges;
     const int ispec = edges(iedge).ispec;
+    const int gedge = edges(iedge).iedge;
     const specfem::mesh_entity::dim2::type edge_type = edges(iedge).edge_type;
     const auto index =
         (edges(iedge).reverse_orientation)
@@ -209,7 +220,7 @@ class ChunkEdgeIterator : public TeamThreadRangePolicy<TeamMemberType, int> {
             : ChunkEdgeIterator::compute_index(ispec, ipoint, num_points,
                                                edge_type);
 
-    return index_type{ index, ipoint, iedge };
+    return index_type{ index, gedge, iedge, ipoint, i };
   }
 
   /**
@@ -463,10 +474,15 @@ class ChunkedEdgeIterator : public TeamPolicy<ParallelConfig> {
                                          ///< specfem::execution::for_each_level
   using execution_space =
       typename base_type::execution_space; ///< Execution space type.
-  using base_index_type = specfem::point::edge_index<
-      ParallelConfig::dimension>; ///< Index type to be used when calling
-                                  ///< @ref specfem::execution::for_all
-                                  ///< with this iterator.
+  using base_index_type = EdgePointIndex<
+      ParallelConfig::dimension, int,
+      typename base_type::execution_space>; ///< Index type
+                                            ///< to be used
+                                            ///< when calling
+                                            ///< @ref
+                                            ///< specfem::execution::for_all
+                                            ///< with this
+                                            ///< iterator.
 
   /**
    * @brief Constructor with explicit edge view and point count

include/execution/chunked_intersection_iterator.hpp

@@ -129,6 +129,11 @@ class InterfacePointIndex {
                         ///< the interface
   }
 
+  KOKKOS_INLINE_FUNCTION
+  constexpr const index_type &get_local_index() const {
+    return this->local_index; ///< Returns the local point index
+  }
+
   /**
    * @brief Constructor for InterfacePointIndex
    *
@@ -158,6 +163,7 @@ class InterfacePointIndex {
 
 private:
   index_type index;             ///< Index of the GLL point on the interface
+  index_type local_index;       ///< Index of the GLL point relative to chunk
   KokkosIndexType kokkos_index; ///< Kokkos index type
 };
 
@@ -552,10 +558,15 @@ class ChunkedIntersectionIterator : public TeamPolicy<ParallelConfig> {
 
   using execution_space =
       typename base_type::execution_space; ///< Execution space type.
-  using base_index_type = specfem::point::interface_index<
-      ParallelConfig::dimension>; ///< Index type to be used when calling @c
-                                  ///< specfem::execution::for_all with this
-                                  ///< iterator.
+  using base_index_type = InterfacePointIndex<
+      ParallelConfig::dimension, int,
+      typename base_type::execution_space>; ///< Index type
+                                            ///< to be used
+                                            ///< when calling
+                                            ///< @ref
+                                            ///< specfem::execution::for_all
+                                            ///< with this
+                                            ///< iterator.
 
   /**
    * @brief Constructor with explicit edge views and point count

include/execution/for_all.hpp

@@ -17,8 +17,7 @@ std::enable_if_t<
                      Kokkos::DefaultHostExecutionSpace>)),
     void>
 impl_for_all(const IndexType &index, const ClosureType &closure) {
-  const auto i = index.get_index();
-  closure(i);
+  closure(index);
 }
 
 #if defined(KOKKOS_ENABLE_CUDA) || defined(KOKKOS_ENABLE_HIP)
@@ -32,8 +31,7 @@ KOKKOS_FORCEINLINE_FUNCTION std::enable_if_t<
                      Kokkos::DefaultExecutionSpace>)),
     void>
 impl_for_all(const IndexType &index, const ClosureType &closure) {
-  const auto i = index.get_index();
-  closure(i);
+  closure(index);
 }
 
 #endif

include/execution/mapped_chunked_domain_iterator.hpp

@@ -169,6 +169,10 @@ class MappedChunkedDomainIterator
   using execution_space =
       typename base_type::execution_space; ///< Execution space type
 
+  using base_index_type =
+      MappedPointIndex<DimensionTag, typename ViewType::value_type,
+                       ParallelConfig::simd::using_simd, execution_space>;
+
   MappedChunkedDomainIterator(
       const ViewType indices, const ViewType mapping,
       const specfem::mesh_entity::element_grid<DimensionTag> &element_grid)

include/execution/range_iterator.hpp

@@ -130,6 +130,14 @@ class RangeIterator : public RangePolicy<ParallelConfig> {
                                        ///< calling @ref
                                        ///< specfem::execution::for_each_level
 
+  using base_index_type = RangeIndex<
+      typename base_type::index_type, ParallelConfig::simd::using_simd,
+      typename base_type::execution_space>; ///< Index type to be
+                                            ///< used when calling
+                                            ///< @ref
+                                            ///< specfem::execution::for_all
+                                            ///< with this iterator.
+
   using execution_space =
       typename base_type::execution_space; ///< Execution space type.
 

include/kokkos_kernels/impl/compute_coupling.tpp

@@ -63,34 +63,32 @@ void specfem::kokkos_kernels::impl::compute_coupling(
 
   specfem::execution::for_all(
       "specfem::kokkos_kernels::impl::compute_coupling", chunk,
-      KOKKOS_LAMBDA(const typename decltype(chunk)::base_index_type &index) {
-        const auto self_index = index.self_index;
-        const auto coupled_index = index.coupled_index;
+      KOKKOS_LAMBDA(const typename decltype(chunk)::base_index_type &iterator_index) {
+        const auto index = iterator_index.get_index();
 
         specfem::point::coupled_interface<dimension_tag, connection_tag,
                                           interface_tag, boundary_tag>
             point_interface_data;
-        specfem::assembly::load_on_device(self_index, coupled_interfaces,
+        specfem::assembly::load_on_device(index.self_index, coupled_interfaces,
                                           point_interface_data);
 
         CoupledFieldType coupled_field;
-        specfem::assembly::load_on_device(coupled_index, field, coupled_field);
-
+        specfem::assembly::load_on_device(index.coupled_index, field, coupled_field);
         SelfFieldType self_field;
 
         specfem::medium::compute_coupling(point_interface_data, coupled_field,
                                           self_field);
 
         PointBoundaryType point_boundary;
-        specfem::assembly::load_on_device(self_index, boundaries,
+        specfem::assembly::load_on_device(index.self_index, boundaries,
                                           point_boundary);
         if constexpr (BoundaryTag ==
                       specfem::element::boundary_tag::acoustic_free_surface) {
           specfem::boundary_conditions::apply_boundary_conditions(
               point_boundary, self_field);
         }
 
-        specfem::assembly::atomic_add_on_device(self_index, field, self_field);
+        specfem::assembly::atomic_add_on_device(index.self_index, field, self_field);
       });
 
   return;