Messing around with oddly shaped grids

parthenon-hpc-lab · Aug 17, 2023 · 91dc79d · 91dc79d
1 parent 22ebf08
commit 91dc79d
Show file tree

Hide file tree

Showing 2 changed files with 68 additions and 16 deletions.
diff --git a/example/poisson_gmg/parthenon_app_inputs.cpp b/example/poisson_gmg/parthenon_app_inputs.cpp
@@ -60,7 +60,7 @@ void ProblemGenerator(Mesh *pm, ParameterInput *pin, MeshData<Real> *md) {
         if (x > 0.25 && x < 0.75) {
           val = 100.0;
         }
-        pack(b, te, poisson_package::res_err(), k, j, i) = x + y;
+        pack(b, te, poisson_package::res_err(), k, j, i) = x + 10.0 * y;
         pack(b, te, poisson_package::rhs_base(), k, j, i) = val;
       });
 }

diff --git a/example/poisson_gmg/poisson_package.hpp b/example/poisson_gmg/poisson_package.hpp
@@ -93,22 +93,24 @@ TaskStatus AddFieldsAndStore(std::shared_ptr<MeshData<Real>> &md, Real wa = 1.0,
 
 template <class var>
 TaskStatus SetToZero(std::shared_ptr<MeshData<Real>> &md) {
-  using TE = parthenon::TopologicalElement;
-  IndexRange ib = md->GetBoundsI(IndexDomain::interior, te);
-  IndexRange jb = md->GetBoundsJ(IndexDomain::interior, te);
-  IndexRange kb = md->GetBoundsK(IndexDomain::interior, te);
-
   auto desc = parthenon::MakePackDescriptor<var>(md.get());
   auto pack = desc.GetPack(md.get());
-  parthenon::par_for(
-      DEFAULT_LOOP_PATTERN, "SetPotentialToZero", DevExecSpace(), 0,
-      pack.GetNBlocks() - 1, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e,
-      KOKKOS_LAMBDA(const int b, const int k, const int j, const int i) {
-        pack(b, te, var(), k, j, i) = 0.0;
-        // If fold expressions worked correctly in CUDA, we could make this
-        // a variadic template <class... vars>, build the pack with all vars
-        // at set everything requested to zero in the same kernel.
-        // (pack(b, te, vars(), k, j, i) = 0.0, ...);
+  const size_t scratch_size_in_bytes = 0; 
+  const int scratch_level = 1;
+  parthenon::par_for_outer(
+      DEFAULT_OUTER_LOOP_PATTERN, "Print", DevExecSpace(),
+      scratch_size_in_bytes, scratch_level, 
+      0, pack.GetNBlocks() - 1, 
+      KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int b) {
+        auto cb = GetIndexShape(pack(b, te, 0));
+        const auto &coords = pack.GetCoordinates(b);
+        IndexRange ib = cb.GetBoundsI(IndexDomain::interior, te);
+        IndexRange jb = cb.GetBoundsJ(IndexDomain::interior, te);
+        IndexRange kb = cb.GetBoundsK(IndexDomain::interior, te);
+        parthenon::par_for_inner(parthenon::inner_loop_pattern_simdfor_tag, member, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e,
+                                 [&](int k, int j, int i) {
+          pack(b, te, var(), k, j, i) = 0.0; 
+        });
       });
   return TaskStatus::complete;
 }
@@ -161,6 +163,55 @@ TaskStatus PrintChosenValues(std::shared_ptr<MeshData<Real>> &md, std::string &l
     printf("var %i: %s\n", col_num, name.c_str());
     col_num++;
   }
+  printf("i=[%i, %i] j=[%i, %i] k=[%i, %i]\n", ib.s, ib.e, jb.s, jb.e, kb.s, kb.e);
+  const size_t scratch_size_in_bytes = 0; 
+  const int scratch_level = 1;
+  parthenon::par_for(
+      DEFAULT_LOOP_PATTERN, "Print", DevExecSpace(), 
+      0, pack.GetNBlocks() - 1, 0, 0, 0, 0,
+      KOKKOS_LAMBDA(const int b, int, int) {
+        auto cb = GetIndexShape(pack(b, te, 0));
+        const auto &coords = pack.GetCoordinates(b);
+        IndexRange ib = cb.GetBoundsI(IndexDomain::interior, te);
+        IndexRange jb = cb.GetBoundsJ(IndexDomain::interior, te);
+        IndexRange kb = cb.GetBoundsK(IndexDomain::interior, te);
+        printf("b=%i i=[%i, %i] j=[%i, %i] k=[%i, %i]\n", b, ib.s, ib.e, jb.s, jb.e, kb.s, kb.e); 
+        for (int k = kb.s; k <= kb.e; ++k) { 
+          for (int j = jb.s; j <= jb.e; ++j) { 
+            for (int i = ib.s; i <= ib.e; ++i) { 
+              Real x = coords.template X<1, te>(i);
+              Real y = coords.template X<2, te>(j);
+              std::array<Real, sizeof...(vars)> vals{pack(b, te, vars(), k, j, i)...};
+              printf("b = %i i = %2i j = %2i x = %e y = %e x + 10*y = %e ", b, i, j, x, y, x + 10.0*y);
+              for (int v = 0; v < sizeof...(vars); ++v) {
+                printf("%e ", vals[v]);
+              }
+              printf("\n"); 
+            }
+          }
+        }
+      });
+  /*
+  const size_t scratch_size_in_bytes = 0; 
+  const int scratch_level = 1;
+  parthenon::par_for_outer(
+      DEFAULT_OUTER_LOOP_PATTERN, "Print", DevExecSpace(),
+      scratch_size_in_bytes, scratch_level, 
+      0, pack.GetNBlocks() - 1, 
+      KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int b) {
+        auto cb = GetIndexShape(pack(b, te, 0));
+        const auto &coords = pack.GetCoordinates(b);
+        IndexRange ib = cb.GetBoundsI(IndexDomain::interior, te);
+        IndexRange jb = cb.GetBoundsJ(IndexDomain::interior, te);
+        IndexRange kb = cb.GetBoundsK(IndexDomain::interior, te);
+        printf("b=%i i=[%i, %i] j=[%i, %i] k=[%i, %i]\n", b, ib.s, ib.e, jb.s, jb.e, kb.s, kb.e); 
+        parthenon::par_for_inner(parthenon::inner_loop_pattern_simdfor_tag, member, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e,
+                                 [&](int k, int j, int i) {
+           // Work here
+        });
+      });
+      */
+  /*
   parthenon::par_for(
       DEFAULT_LOOP_PATTERN, "SetPotentialToZero", DevExecSpace(), 0,
       pack.GetNBlocks() - 1, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e,
@@ -169,12 +220,13 @@ TaskStatus PrintChosenValues(std::shared_ptr<MeshData<Real>> &md, std::string &l
         Real x = coords.template X<1, te>(i);
         Real y = coords.template X<2, te>(j);
         std::array<Real, sizeof...(vars)> vals{pack(b, te, vars(), k, j, i)...};
-        printf("b = %i i = %2i j = %2i x = %e y = %e x+y = %e ", b, i, j, x, y, x + y);
+        printf("b = %i i = %2i j = %2i x = %e y = %e x + 10*y = %e ", b, i, j, x, y, x + 10.0*y);
         for (int v = 0; v < sizeof...(vars); ++v) {
           printf("%e ", vals[v]);
         }
         printf("\n");
       });
+  */
   printf("Done with MeshData\n\n");
   return TaskStatus::complete;
 }