clean up indexer size calls

parthenon-hpc-lab · Oct 31, 2024 · 5a01dbe · 5a01dbe
1 parent ed4911f
commit 5a01dbe
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 27 deletions.
diff --git a/src/bvals/comms/boundary_communication.cpp b/src/bvals/comms/boundary_communication.cpp
@@ -274,20 +274,9 @@ TaskStatus SetBounds(std::shared_ptr<MeshData<Real>> &md) {
   size_t Nel_max = 0;
   for (int b = 0; b < nbound; b++) {
     for (int it = 0; it < bnd_info_h(b).ntopological_elements; ++it) {
-      const int Nt = bnd_info_h(b).idxer[it].template EndIdx<0>() -
-                     bnd_info_h(b).idxer[it].template StartIdx<0>() + 1;
-      const int Nu = bnd_info_h(b).idxer[it].template EndIdx<1>() -
-                     bnd_info_h(b).idxer[it].template StartIdx<1>() + 1;
-      const int Nv = bnd_info_h(b).idxer[it].template EndIdx<2>() -
-                     bnd_info_h(b).idxer[it].template StartIdx<2>() + 1;
-      // const int Nk = bnd_info_h(b).idxer[it].template EndIdx<3>() -
-      //                bnd_info_h(b).idxer[it].template StartIdx<3>() + 1;
-      // const int Nj = bnd_info_h(b).idxer[it].template EndIdx<4>() -
-      //                bnd_info_h(b).idxer[it].template StartIdx<4>() + 1;
-      // const int Ni = bnd_info_h(b).idxer[it].template EndIdx<5>() -
-      //                bnd_info_h(b).idxer[it].template StartIdx<5>() + 1;
-      if (Nt * Nu * Nv > Nel_max) {
-        Nel_max = Nt * Nu * Nv;
+      const int tensor_size = bnd_info_h(b).idxer[it].sizeAtAndBelow(2);
+      if (tensor_size > Nel_max) {
+        Nel_max = tensor_size;
       }
     }
   }
@@ -310,20 +299,10 @@ TaskStatus SetBounds(std::shared_ptr<MeshData<Real>> &md) {
           Real fac = ftemp; // Can't capture structured bindings
           const int iel = static_cast<int>(tel) % 3;
           // Element t, u, v in variable
-          const int Nel = (bnd_info(b).idxer[it].template EndIdx<0>() -
-                           bnd_info(b).idxer[it].template StartIdx<0>() + 1) *
-                          (bnd_info(b).idxer[it].template EndIdx<1>() -
-                           bnd_info(b).idxer[it].template StartIdx<1>() + 1) *
-                          (bnd_info(b).idxer[it].template EndIdx<2>() -
-                           bnd_info(b).idxer[it].template StartIdx<2>() + 1);
-          const int Nidx = (bnd_info(b).idxer[it].template EndIdx<3>() -
-                            bnd_info(b).idxer[it].template StartIdx<3>() + 1) *
-                           (bnd_info(b).idxer[it].template EndIdx<4>() -
-                            bnd_info(b).idxer[it].template StartIdx<4>() + 1) *
-                           (bnd_info(b).idxer[it].template EndIdx<5>() -
-                            bnd_info(b).idxer[it].template StartIdx<5>() + 1);
+          const int Nel = idxer.sizeAtAndBelow(2);
+          const int Nidx = idxer.sizeAtAndAbove(3);
           if (el >= Nel) return;
-          const int Ni = idxer.template EndIdx<5>() - idxer.template StartIdx<5>() + 1;
+          const int Ni = idxer.size(5);
           if (bnd_info(b).buf_allocated && bnd_info(b).allocated) {
             Kokkos::parallel_for(
                 Kokkos::TeamThreadRange<>(team_member, idxer.size() / Nel / Ni),

diff --git a/src/utils/indexer.hpp b/src/utils/indexer.hpp
@@ -90,6 +90,25 @@ struct Indexer {
   KOKKOS_FORCEINLINE_FUNCTION
   std::size_t size() const { return _size; }
 
+  KOKKOS_FORCEINLINE_FUNCTION
+  std::size_t size(int dim) const { return end[dim] - start[dim] + 1; }
+
+  KOKKOS_FORCEINLINE_FUNCTION
+  std::size_t sizeAtAndAbove(int dim) const {
+    std::size_t out = dim < rank;
+    for (int i = dim; i < rank; ++i)
+      out *= size(i);
+    return out;
+  }
+
+  KOKKOS_FORCEINLINE_FUNCTION
+  std::size_t sizeAtAndBelow(int dim) const {
+    std::size_t out = dim >= 0;
+    for (int i = 0; i <= dim; ++i)
+      out *= size(i);
+    return out;
+  }
+
   KOKKOS_FORCEINLINE_FUNCTION
   std::tuple<Ts...> operator()(int idx) const {
     return GetIndicesImpl(idx, std::make_index_sequence<sizeof...(Ts)>());