Updated the get mask utility for ~10% performance gain on apple silicon

core/specfem/assembly/boundaries/dim2/impl/acoustic_free_surface.hpp

@@ -244,7 +244,7 @@ struct acoustic_free_surface<specfem::element::dimension_tag::dim2> {
     using mask_type = typename simd::mask_type;
     using tag_type = typename simd::tag_type;
 
-    mask_type mask([&](std::size_t lane) { return index.mask(lane); });
+    const auto mask = index.template get_mask<simd>();
 
     for (int lane = 0; lane < mask_type::size(); ++lane) {
       if (index.mask(lane)) {
@@ -284,7 +284,7 @@ struct acoustic_free_surface<specfem::element::dimension_tag::dim2> {
     using mask_type = typename simd::mask_type;
     using tag_type = typename simd::tag_type;
 
-    mask_type mask([&](std::size_t lane) { return index.mask(lane); });
+    const auto mask = index.template get_mask<simd>();
 
     for (int lane = 0; lane < mask_type::size(); ++lane) {
       if (index.mask(lane)) {

core/specfem/assembly/boundaries/dim2/impl/stacey.hpp

@@ -327,7 +327,7 @@ template <> struct stacey<specfem::element::dimension_tag::dim2> {
     using mask_type = typename simd::mask_type;
     using tag_type = typename simd::tag_type;
 
-    mask_type mask([&](std::size_t lane) { return index.mask(lane); });
+    const auto mask = index.template get_mask<simd>();
 
     for (int lane = 0; lane < mask_type::size(); ++lane) {
       if (index.mask(lane)) {
@@ -375,7 +375,7 @@ template <> struct stacey<specfem::element::dimension_tag::dim2> {
     using mask_type = typename simd::mask_type;
     using tag_type = typename simd::tag_type;
 
-    mask_type mask([&](std::size_t lane) { return index.mask(lane); });
+    const auto mask = index.template get_mask<simd>();
 
     for (int lane = 0; lane < mask_type::size(); ++lane) {
       if (index.mask(lane)) {
@@ -496,7 +496,7 @@ template <> struct stacey<specfem::element::dimension_tag::dim2> {
     using mask_type = typename simd::mask_type;
     using tag_type = typename simd::tag_type;
 
-    mask_type mask([&](std::size_t lane) { return index.mask(lane); });
+    const auto mask = index.template get_mask<simd>();
 
     for (int lane = 0; lane < mask_type::size(); ++lane) {
       if (index.mask(lane)) {
@@ -546,7 +546,7 @@ template <> struct stacey<specfem::element::dimension_tag::dim2> {
     using mask_type = typename simd::mask_type;
     using tag_type = typename simd::tag_type;
 
-    mask_type mask([&](std::size_t lane) { return index.mask(lane); });
+    const auto mask = index.template get_mask<simd>();
 
     for (int lane = 0; lane < mask_type::size(); ++lane) {
       if (index.mask(lane)) {

core/specfem/assembly/boundary_values/dim2/impl/boundary_medium_container.hpp

@@ -116,7 +116,7 @@ class boundary_medium_container<specfem::element::dimension_tag::dim2,
     using mask_type = typename simd::mask_type;
     using tag_type = typename simd::tag_type;
 
-    mask_type mask([&](std::size_t lane) { return index.mask(lane); });
+    const auto mask = index.template get_mask<simd>();
 
     for (int icomp = 0; icomp < components; ++icomp)
       Kokkos::Experimental::where(mask, acceleration(icomp))
@@ -144,7 +144,7 @@ class boundary_medium_container<specfem::element::dimension_tag::dim2,
     using mask_type = typename simd::mask_type;
     using tag_type = typename simd::tag_type;
 
-    mask_type mask([&](std::size_t lane) { return index.mask(lane); });
+    const auto mask = index.template get_mask<simd>();
 
     for (int icomp = 0; icomp < components; ++icomp)
       Kokkos::Experimental::where(mask, acceleration(icomp))

core/specfem/assembly/boundary_values/dim3/impl/boundary_medium_container.hpp

@@ -119,7 +119,7 @@ class boundary_medium_container<specfem::element::dimension_tag::dim3,
     using mask_type = typename simd::mask_type;
     using tag_type = typename simd::tag_type;
 
-    mask_type mask([&](std::size_t lane) { return index.mask(lane); });
+    const auto mask = index.template get_mask<simd>();
 
     for (int icomp = 0; icomp < components; ++icomp)
       Kokkos::Experimental::where(mask, acceleration(icomp))
@@ -148,7 +148,7 @@ class boundary_medium_container<specfem::element::dimension_tag::dim3,
     using mask_type = typename simd::mask_type;
     using tag_type = typename simd::tag_type;
 
-    mask_type mask([&](std::size_t lane) { return index.mask(lane); });
+    const auto mask = index.template get_mask<simd>();
 
     for (int icomp = 0; icomp < components; ++icomp)
       Kokkos::Experimental::where(mask, acceleration(icomp))

core/specfem/assembly/fields/impl/add_access_functions.hpp

@@ -75,9 +75,10 @@ KOKKOS_FORCEINLINE_FUNCTION void add_after_simd_dispatch(
 
   check_accessor_compatibility<AccessorTypes...>();
 
-  using mask_type = typename std::tuple_element_t<
-      0, std::tuple<AccessorTypes...> >::simd::mask_type;
-  mask_type mask([&](std::size_t lane) { return index.mask(lane); });
+  using simd =
+      typename std::tuple_element_t<0, std::tuple<AccessorTypes...> >::simd;
+  using mask_type = typename simd::mask_type;
+  const auto mask = index.template get_mask<simd>();
 
   const int iglob = index.iglob;
 

core/specfem/assembly/fields/impl/atomic_add_access_functions.hpp

@@ -89,9 +89,10 @@ KOKKOS_FORCEINLINE_FUNCTION void atomic_add_after_simd_dispatch(
 
   check_accessor_compatibility<AccessorTypes...>();
 
-  using mask_type = typename std::tuple_element_t<
-      0, std::tuple<AccessorTypes...> >::simd::mask_type;
-  mask_type mask([&](std::size_t lane) { return index.mask(lane); });
+  using simd =
+      typename std::tuple_element_t<0, std::tuple<AccessorTypes...> >::simd;
+  using mask_type = typename simd::mask_type;
+  const auto mask = index.template get_mask<simd>();
 
   const int iglob = index.iglob;
 

core/specfem/assembly/fields/impl/load_access_functions.hpp

@@ -89,9 +89,10 @@ KOKKOS_FORCEINLINE_FUNCTION void load_after_simd_dispatch(
 
   check_accessor_compatibility<AccessorTypes...>();
 
-  using mask_type = typename std::tuple_element_t<
-      0, std::tuple<AccessorTypes...> >::simd::mask_type;
-  mask_type mask([&](std::size_t lane) { return index.mask(lane); });
+  using simd =
+      typename std::tuple_element_t<0, std::tuple<AccessorTypes...> >::simd;
+  using mask_type = typename simd::mask_type;
+  const auto mask = index.template get_mask<simd>();
 
   const int iglob = index.iglob;
 

core/specfem/assembly/fields/impl/store_access_functions.hpp

@@ -71,9 +71,10 @@ KOKKOS_FORCEINLINE_FUNCTION void store_after_simd_dispatch(
 
   check_accessor_compatibility<AccessorTypes...>();
 
-  using mask_type = typename std::tuple_element_t<
-      0, std::tuple<AccessorTypes...> >::simd::mask_type;
-  mask_type mask([&](std::size_t lane) { return index.mask(lane); });
+  using simd =
+      typename std::tuple_element_t<0, std::tuple<AccessorTypes...> >::simd;
+  using mask_type = typename simd::mask_type;
+  const auto mask = index.template get_mask<simd>();
 
   const int iglob = index.iglob;
 

core/specfem/assembly/impl/domain_accessor.hpp

@@ -54,7 +54,7 @@ class DomainAccessor {
     using mask_type = typename simd::mask_type;
     using tag_type = typename simd::tag_type;
 
-    mask_type mask([&](std::size_t lane) { return index.mask(lane); });
+    const auto mask = index.template get_mask<simd>();
     static_cast<const DataContainer *>(this)->for_each_on_device(
         index, [&](const type_real &value, const std::size_t i) mutable {
           Kokkos::Experimental::where(mask, values[i])
@@ -70,7 +70,7 @@ class DomainAccessor {
     using mask_type = typename simd::mask_type;
     using tag_type = typename simd::tag_type;
 
-    mask_type mask([&](std::size_t lane) { return index.mask(lane); });
+    const auto mask = index.template get_mask<simd>();
     static_cast<const DataContainer *>(this)->for_each_on_host(
         index, [&](const type_real &value, const std::size_t i) mutable {
           Kokkos::Experimental::where(mask, values[i])
@@ -104,7 +104,7 @@ class DomainAccessor {
     using mask_type = typename simd::mask_type;
     using tag_type = typename simd::tag_type;
 
-    mask_type mask([&](std::size_t lane) { return index.mask(lane); });
+    const auto mask = index.template get_mask<simd>();
     static_cast<const DataContainer *>(this)->for_each_on_device(
         index, [&](type_real &value, const std::size_t i) {
           Kokkos::Experimental::where(mask, values[i])
@@ -120,7 +120,7 @@ class DomainAccessor {
     using mask_type = typename simd::mask_type;
     using tag_type = typename simd::tag_type;
 
-    mask_type mask([&](std::size_t lane) { return index.mask(lane); });
+    const auto mask = index.template get_mask<simd>();
     static_cast<const DataContainer *>(this)->for_each_on_host(
         index, [&](type_real &value, const std::size_t i) {
           Kokkos::Experimental::where(mask, values[i])

core/specfem/assembly/jacobian_matrix/dim2/impl_load.hpp

@@ -33,7 +33,7 @@ KOKKOS_FORCEINLINE_FUNCTION void impl_load(const IndexType &index,
   const auto &mapping = container.xix.get_mapping();
   const std::size_t _index = mapping(ispec, iz, ix);
 
-  mask_type mask([&](std::size_t lane) { return index.mask(lane); });
+  const auto mask = index.template get_mask<simd>();
 
   if constexpr (on_device) {
     Kokkos::Experimental::where(mask, point.xix)

core/specfem/assembly/jacobian_matrix/dim2/impl_store.hpp

@@ -30,7 +30,7 @@ inline void impl_store(const IndexType &index, const ContainerType &derivatives,
   using mask_type = typename simd::mask_type;
   using tag_type = typename simd::tag_type;
 
-  mask_type mask([&](std::size_t lane) { return index.mask(lane); });
+  const auto mask = index.template get_mask<simd>();
 
   const auto &mapping = derivatives.xix.get_mapping();
   const std::size_t _index = mapping(ispec, iz, ix);

core/specfem/point/assembly_index.hpp

@@ -107,6 +107,26 @@ struct assembly_index<true>
   KOKKOS_FUNCTION
   bool mask(const std::size_t &lane) const { return int(lane) < number_points; }
 
+  /**
+   * @brief Returns a SIMD mask for valid lanes.
+   *
+   * For full chunks (all lanes valid) returns an all-true mask using the cheap
+   * broadcast constructor, avoiding the per-lane generator. For the rare
+   * partial last chunk falls back to per-lane evaluation.
+   *
+   * @tparam simd_type SIMD wrapper type (e.g. specfem::datatype::simd<float,
+   * true>)
+   * @return SIMD mask with true for valid lanes
+   */
+  template <typename simd_type>
+  KOKKOS_INLINE_FUNCTION typename simd_type::mask_type get_mask() const {
+    if (number_points >= simd_type::size()) {
+      return typename simd_type::mask_type(true);
+    }
+    return typename simd_type::mask_type(
+        [&](std::size_t lane) { return int(lane) < number_points; });
+  }
+
   /**
    * @name Constructors
    * @brief Constructors for initializing the SIMD assembly index.

core/specfem/point/index.hpp

@@ -164,6 +164,26 @@ struct index<specfem::element::dimension_tag::dim2, true>
   bool mask(const std::size_t &lane) const {
     return int(lane) < number_elements;
   }
+
+  /**
+   * @brief Returns a SIMD mask for valid lanes.
+   *
+   * For full chunks (all lanes valid) returns an all-true mask using the cheap
+   * broadcast constructor, avoiding the per-lane generator. For the rare
+   * partial last chunk falls back to per-lane evaluation.
+   *
+   * @tparam simd_type SIMD wrapper type (e.g. specfem::datatype::simd<float,
+   * true>)
+   * @return SIMD mask with true for valid lanes
+   */
+  template <typename simd_type>
+  KOKKOS_INLINE_FUNCTION typename simd_type::mask_type get_mask() const {
+    if (number_elements >= simd_type::size()) {
+      return typename simd_type::mask_type(true);
+    }
+    return typename simd_type::mask_type(
+        [&](std::size_t lane) { return int(lane) < number_elements; });
+  }
 };
 
 //-------------------------- 3D Specializations ------------------------------//
@@ -334,6 +354,26 @@ struct index<specfem::element::dimension_tag::dim3, true>
   bool mask(const std::size_t &lane) const {
     return int(lane) < number_elements;
   }
+
+  /**
+   * @brief Returns a SIMD mask for valid lanes.
+   *
+   * For full chunks (all lanes valid) returns an all-true mask using the cheap
+   * broadcast constructor, avoiding the per-lane generator. For the rare
+   * partial last chunk falls back to per-lane evaluation.
+   *
+   * @tparam simd_type SIMD wrapper type (e.g. specfem::datatype::simd<float,
+   * true>)
+   * @return SIMD mask with true for valid lanes
+   */
+  template <typename simd_type>
+  KOKKOS_INLINE_FUNCTION typename simd_type::mask_type get_mask() const {
+    if (number_elements >= simd_type::size()) {
+      return typename simd_type::mask_type(true);
+    }
+    return typename simd_type::mask_type(
+        [&](std::size_t lane) { return int(lane) < number_elements; });
+  }
 };
 
 } // namespace point