Exawind · psakievich · Apr 18, 2024 · Apr 18, 2024
diff --git a/include/matrix_free/ValidSimdLength.h b/include/matrix_free/ValidSimdLength.h
@@ -20,7 +20,7 @@ namespace sierra {
 namespace nalu {
 namespace matrix_free {
 
-static constexpr int invalid_offset = -1;
+inline constexpr int invalid_offset = -1;
 
 constexpr stk::mesh::FastMeshIndex invalid_mesh_index =
   stk::mesh::FastMeshIndex{

diff --git a/unit_tests/matrix_free/SetOffsetsAndCoordinates.h b/unit_tests/matrix_free/SetOffsetsAndCoordinates.h
@@ -0,0 +1,57 @@
+// Copyright 2017 National Technology & Engineering Solutions of Sandia, LLC
+// (NTESS), National Renewable Energy Laboratory, University of Texas Austin,
+// Northwest Research Associates. Under the terms of Contract DE-NA0003525
+// with NTESS, the U.S. Government retains certain rights in this software.
+//
+// This software is released under the BSD 3-clause license. See LICENSE file
+// for more details.
+//
+
+#include <Kokkos_Core.hpp>
+
+#include "matrix_free/KokkosViewTypes.h"
+#include "matrix_free/ValidSimdLength.h"
+#include "matrix_free/LobattoQuadratureRule.h"
+
+#include "gtest/gtest.h"
+
+namespace sierra {
+namespace nalu {
+namespace matrix_free {
+
+template <int p>
+void
+set_offsets_and_coordinates(
+  elem_offset_view<p> offsets, vector_view<p> coordinates, int nrepeated)
+{
+  constexpr int nodes_per_elem = (p + 1) * (p + 1) * (p + 1);
+  constexpr auto nodes = GLL<p>::nodes;
+  Kokkos::parallel_for(
+    nrepeated, KOKKOS_LAMBDA(int index) {
+      const int elem_offset = index * nodes_per_elem;
+
+      for (int k = 0; k < p + 1; ++k) {
+        const auto cz = nodes[k];
+        for (int j = 0; j < p + 1; ++j) {
+          const auto cy = nodes[j];
+          for (int i = 0; i < p + 1; ++i) {
+            const auto cx = nodes[i];
+            coordinates(index, k, j, i, 0) = cx;
+            coordinates(index, k, j, i, 1) = cy;
+            coordinates(index, k, j, i, 2) = cz;
+
+            const int simd_lane_0 = 0;
+            offsets(index, k, j, i, simd_lane_0) =
+              elem_offset + k * (p + 1) * (p + 1) + j * (p + 1) + i;
+            for (int n = 1; n < simd_len; ++n) {
+              offsets(index, k, j, i, n) = invalid_offset;
+            }
+          }
+        }
+      }
+    });
+}
+
+} // namespace matrix_free
+} // namespace nalu
+} // namespace sierra
diff --git a/unit_tests/matrix_free/UnitTestConductionDiagonal.C b/unit_tests/matrix_free/UnitTestConductionDiagonal.C
@@ -26,6 +26,7 @@
 #include "matrix_free/LinearDiffusionMetric.h"
 #include "matrix_free/LinearVolume.h"
 #include "matrix_free/KokkosViewTypes.h"
+#include "SetOffsetsAndCoordinates.h"
 
 #include "gtest/gtest.h"
 #include "mpi.h"
@@ -47,39 +48,14 @@ make_map()
     Teuchos::make_rcp<Teuchos::MpiComm<int>>(MPI_COMM_WORLD));
 }
 
-template <int p>
-void
-set_aux_fields(elem_offset_view<p> offsets, vector_view<p> coordinates)
-{
-  constexpr auto nodes = GLL<p>::nodes;
-  Kokkos::parallel_for(
-    num_elems, KOKKOS_LAMBDA(int index) {
-      for (int k = 0; k < p + 1; ++k) {
-        const auto cz = nodes[k];
-        for (int j = 0; j < p + 1; ++j) {
-          const auto cy = nodes[j];
-          for (int i = 0; i < p + 1; ++i) {
-            const auto cx = nodes[i];
-            coordinates(index, k, j, i, 0) = cx;
-            coordinates(index, k, j, i, 1) = cy;
-            coordinates(index, k, j, i, 2) = cz;
-            offsets(index, 0, k, j, i) = 1 + index * simd_len * nodes_per_elem +
-                                         k * (order + 1) * (order + 1) +
-                                         j * (order + 1) + i;
-          }
-        }
-      }
-    });
-}
-
 } // namespace test_diagonal
 class DiagonalFixture : public ::testing::Test
 {
 public:
   DiagonalFixture()
   {
     vector_view<order> coordinates{"coordinates", num_elems};
-    test_diagonal::set_aux_fields<order>(offsets, coordinates);
+    set_offsets_and_coordinates<order>(offsets, coordinates, num_elems);
 
     scalar_view<order> alpha{"alpha", num_elems};
     Kokkos::deep_copy(alpha, 1.0);
@@ -101,10 +77,6 @@ public:
 
 TEST_F(DiagonalFixture, diagonal_executes)
 {
-  node_offset_view dirichlet_offsets("empty_dirichlet", 1);
-  decltype(rhs.getLocalViewDevice(Tpetra::Access::ReadWrite)) shared_rhs(
-    "empty_rhs", 1, 1);
-
   rhs.putScalar(0.);
   conduction_diagonal<order>(
     1, offsets, volume_metric, diffusion_metric,

diff --git a/unit_tests/matrix_free/UnitTestConductionInterior.C b/unit_tests/matrix_free/UnitTestConductionInterior.C
@@ -27,6 +27,8 @@
 #include "matrix_free/LinearVolume.h"
 #include "matrix_free/KokkosViewTypes.h"
 
+#include "SetOffsetsAndCoordinates.h"
+
 #include "gtest/gtest.h"
 #include "mpi.h"
 
@@ -48,32 +50,8 @@ make_map()
     Teuchos::make_rcp<Teuchos::MpiComm<int>>(MPI_COMM_WORLD));
 }
 
-template <int p>
-void
-set_aux_fields(elem_offset_view<p> offsets, vector_view<p> coordinates)
-{
-  constexpr auto nodes = GLL<p>::nodes;
-  Kokkos::parallel_for(
-    num_elems, KOKKOS_LAMBDA(int index) {
-      for (int k = 0; k < p + 1; ++k) {
-        const auto cz = nodes[k];
-        for (int j = 0; j < p + 1; ++j) {
-          const auto cy = nodes[j];
-          for (int i = 0; i < p + 1; ++i) {
-            const auto cx = nodes[i];
-            coordinates(index, k, j, i, 0) = cx;
-            coordinates(index, k, j, i, 1) = cy;
-            coordinates(index, k, j, i, 2) = cz;
-            offsets(index, 0, k, j, i) = 1 + index * simd_len * nodes_per_elem +
-                                         k * (order + 1) * (order + 1) +
-                                         j * (order + 1) + i;
-          }
-        }
-      }
-    });
-}
-
 } // namespace test_conduction
+
 class ConductionResidualFixture : public ::testing::Test
 {
 public:
@@ -84,7 +62,7 @@ public:
     Kokkos::deep_copy(qm1, 1.0);
 
     vector_view<order> coordinates{"coordinates", num_elems};
-    test_conduction::set_aux_fields<order>(offsets, coordinates);
+    set_offsets_and_coordinates<order>(offsets, coordinates, num_elems);
 
     scalar_view<order> alpha{"alpha", num_elems};
     Kokkos::deep_copy(alpha, 1.0);

diff --git a/unit_tests/matrix_free/UnitTestContinuityInterior.C b/unit_tests/matrix_free/UnitTestContinuityInterior.C
@@ -27,6 +27,7 @@
 #include "matrix_free/LinearDiffusionMetric.h"
 #include "matrix_free/LinearAdvectionMetric.h"
 #include "matrix_free/KokkosViewTypes.h"
+#include "SetOffsetsAndCoordinates.h"
 
 #include "gtest/gtest.h"
 #include "mpi.h"
@@ -49,32 +50,8 @@ make_map()
     Teuchos::make_rcp<Teuchos::MpiComm<int>>(MPI_COMM_WORLD));
 }
 
-template <int p>
-void
-set_aux_fields(elem_offset_view<p> offsets, vector_view<p> coordinates)
-{
-  constexpr auto nodes = GLL<p>::nodes;
-  Kokkos::parallel_for(
-    num_elems, KOKKOS_LAMBDA(int index) {
-      for (int k = 0; k < p + 1; ++k) {
-        const auto cz = nodes[k];
-        for (int j = 0; j < p + 1; ++j) {
-          const auto cy = nodes[j];
-          for (int i = 0; i < p + 1; ++i) {
-            const auto cx = nodes[i];
-            coordinates(index, k, j, i, 0) = cx;
-            coordinates(index, k, j, i, 1) = cy;
-            coordinates(index, k, j, i, 2) = cz;
-            offsets(index, 0, k, j, i) = 1 + index * simd_len * nodes_per_elem +
-                                         k * (order + 1) * (order + 1) +
-                                         j * (order + 1) + i;
-          }
-        }
-      }
-    });
-}
-
 } // namespace test_continuity
+
 class ContinuityResidualFixture : public ::testing::Test
 {
 public:
@@ -84,7 +61,7 @@ public:
     Kokkos::deep_copy(pressure, 0);
 
     vector_view<order> coordinates{"coordinates", num_elems};
-    test_continuity::set_aux_fields<order>(offsets, coordinates);
+    set_offsets_and_coordinates<order>(offsets, coordinates, num_elems);
 
     scalar_view<order> density{"density", num_elems};
     Kokkos::deep_copy(density, 1.0);

diff --git a/unit_tests/matrix_free/UnitTestFilterDiagonal.C b/unit_tests/matrix_free/UnitTestFilterDiagonal.C
@@ -25,6 +25,7 @@
 #include "matrix_free/StkSimdConnectivityMap.h"
 #include "matrix_free/StkToTpetraMap.h"
 #include "matrix_free/StkSimdGatheredElementData.h"
+#include "SetOffsetsAndCoordinates.h"
 
 #include "matrix_free/LobattoQuadratureRule.h"
 #include "matrix_free/LinearVolume.h"
@@ -52,39 +53,14 @@ make_map()
     Teuchos::make_rcp<Teuchos::MpiComm<int>>(MPI_COMM_WORLD));
 }
 
-template <int p>
-void
-set_aux_fields(elem_offset_view<p> offsets, vector_view<p> coordinates)
-{
-  constexpr auto nodes = GLL<p>::nodes;
-  Kokkos::parallel_for(
-    num_elems, KOKKOS_LAMBDA(int index) {
-      for (int k = 0; k < p + 1; ++k) {
-        const auto cz = nodes[k];
-        for (int j = 0; j < p + 1; ++j) {
-          const auto cy = nodes[j];
-          for (int i = 0; i < p + 1; ++i) {
-            const auto cx = nodes[i];
-            coordinates(index, k, j, i, 0) = cx;
-            coordinates(index, k, j, i, 1) = cy;
-            coordinates(index, k, j, i, 2) = cz;
-            offsets(index, 0, k, j, i) = 1 + index * simd_len * nodes_per_elem +
-                                         k * (order + 1) * (order + 1) +
-                                         j * (order + 1) + i;
-          }
-        }
-      }
-    });
-}
-
 } // namespace test_filter_diagonal
 class FilterDiagonal : public ::testing::Test
 {
 public:
   FilterDiagonal()
   {
     vector_view<order> coordinates{"coordinates", num_elems};
-    test_filter_diagonal::set_aux_fields<order>(offsets, coordinates);
+    set_offsets_and_coordinates<order>(offsets, coordinates, num_elems);
     volume_metric = geom::volume_metric<order>(coordinates);
   }
   static constexpr int order = 1;

diff --git a/unit_tests/matrix_free/UnitTestGradientInterior.C b/unit_tests/matrix_free/UnitTestGradientInterior.C
@@ -12,6 +12,7 @@
 #include "matrix_free/LinearAreas.h"
 #include "matrix_free/LinearVolume.h"
 #include "matrix_free/KokkosViewTypes.h"
+#include "SetOffsetsAndCoordinates.h"
 
 #include "gtest/gtest.h"
 #include "mpi.h"
@@ -34,31 +35,6 @@ make_map()
     Teuchos::make_rcp<Teuchos::MpiComm<int>>(MPI_COMM_WORLD));
 }
 
-template <int p>
-void
-set_aux_fields(elem_offset_view<p> offsets, vector_view<p> coordinates)
-{
-  constexpr auto nodes = GLL<p>::nodes;
-  Kokkos::parallel_for(
-    num_elems, KOKKOS_LAMBDA(int index) {
-      for (int k = 0; k < p + 1; ++k) {
-        const auto cz = nodes[k];
-        for (int j = 0; j < p + 1; ++j) {
-          const auto cy = nodes[j];
-          for (int i = 0; i < p + 1; ++i) {
-            const auto cx = nodes[i];
-            coordinates(index, k, j, i, 0) = cx;
-            coordinates(index, k, j, i, 1) = cy;
-            coordinates(index, k, j, i, 2) = cz;
-            offsets(index, 0, k, j, i) = 1 + index * simd_len * nodes_per_elem +
-                                         k * (order + 1) * (order + 1) +
-                                         j * (order + 1) + i;
-          }
-        }
-      }
-    });
-}
-
 } // namespace test_gradient
 class GradientResidualFixture : public ::testing::Test
 {
@@ -69,7 +45,7 @@ public:
     Kokkos::deep_copy(dqdx, 0.0);
 
     vector_view<order> coordinates{"coordinates", num_elems};
-    test_gradient::set_aux_fields<order>(offsets, coordinates);
+    set_offsets_and_coordinates<order>(offsets, coordinates, num_elems);
 
     volume_metric = geom::volume_metric<order>(coordinates);
     area_metric = geom::linear_areas<order>(coordinates);

diff --git a/unit_tests/matrix_free/UnitTestGreenGaussGradientInterior.C b/unit_tests/matrix_free/UnitTestGreenGaussGradientInterior.C
@@ -12,6 +12,7 @@
 #include "matrix_free/LinearAreas.h"
 #include "matrix_free/LinearVolume.h"
 #include "matrix_free/KokkosViewTypes.h"
+#include "SetOffsetsAndCoordinates.h"
 
 #include "gtest/gtest.h"
 #include "mpi.h"
@@ -34,31 +35,6 @@ make_map()
     Teuchos::make_rcp<Teuchos::MpiComm<int>>(MPI_COMM_WORLD));
 }
 
-template <int p>
-void
-set_aux_fields(elem_offset_view<p> offsets, vector_view<p> coordinates)
-{
-  constexpr auto nodes = GLL<p>::nodes;
-  Kokkos::parallel_for(
-    num_elems, KOKKOS_LAMBDA(int index) {
-      for (int k = 0; k < p + 1; ++k) {
-        const auto cz = nodes[k];
-        for (int j = 0; j < p + 1; ++j) {
-          const auto cy = nodes[j];
-          for (int i = 0; i < p + 1; ++i) {
-            const auto cx = nodes[i];
-            coordinates(index, k, j, i, 0) = cx;
-            coordinates(index, k, j, i, 1) = cy;
-            coordinates(index, k, j, i, 2) = cz;
-            offsets(index, 0, k, j, i) = 1 + index * simd_len * nodes_per_elem +
-                                         k * (order + 1) * (order + 1) +
-                                         j * (order + 1) + i;
-          }
-        }
-      }
-    });
-}
-
 } // namespace test_gradient
 class GradientResidualFixture : public ::testing::Test
 {
@@ -69,7 +45,7 @@ public:
     Kokkos::deep_copy(dqdx, 0.0);
 
     vector_view<order> coordinates{"coordinates", num_elems};
-    test_gradient::set_aux_fields<order>(offsets, coordinates);
+    set_offsets_and_coordinates<order>(offsets, coordinates, num_elems);
 
     volume_metric = geom::volume_metric<order>(coordinates);
     area_metric = geom::linear_areas<order>(coordinates);