Use subviews of a single view for fluxes

CHANGELOG.md

@@ -15,6 +15,7 @@
 - [[PR 487]](https://github.com/lanl/parthenon/pull/487) Add default tiling matching `i` index range to MDRange loop pattern.
 
 ### Infrastructure (changes irrelevant to downstream codes)
+- [[PR 502]](https://github.com/lanl/parthenon/pull/502) Use subviews of a single view for fluxes
 - [[PR 505]](https://github.com/lanl/parthenon/pull/505) Can also use buffer-pack-in-one function also in `Mesh::Initialize` (and thus during load balancing/mesh refinement). Breaks sparse variables with FillGhost. Enable with `PARTHENON_ENABLE_INIT_PACKING=ON` (default OFF).
 - [[PR 493]](https://github.com/lanl/parthenon/pull/493) Use subviews of a single view for comm buffers
 - [[PR 500]](https://github.com/lanl/parthenon/pull/500) Update docker file and CI environment (for Cuda 11.3 and latest `nsys`)

src/bvals/bvals_interfaces.hpp

@@ -3,7 +3,7 @@
 // Copyright(C) 2014 James M. Stone <jmstone@princeton.edu> and other code contributors
 // Licensed under the 3-clause BSD License, see LICENSE file for details
 //========================================================================================
-// (C) (or copyright) 2020. Triad National Security, LLC. All rights reserved.
+// (C) (or copyright) 2020-2021. Triad National Security, LLC. All rights reserved.
 //
 // This program was produced under U.S. Government contract 89233218CNA000001 for Los
 // Alamos National Laboratory (LANL), which is operated by Triad National Security, LLC

src/bvals/bvals_var.cpp

@@ -3,7 +3,7 @@
 // Copyright(C) 2014 James M. Stone <jmstone@princeton.edu> and other code contributors
 // Licensed under the 3-clause BSD License, see LICENSE file for details
 //========================================================================================
-// (C) (or copyright) 2020. Triad National Security, LLC. All rights reserved.
+// (C) (or copyright) 2020-2021. Triad National Security, LLC. All rights reserved.
 //
 // This program was produced under U.S. Government contract 89233218CNA000001 for Los
 // Alamos National Laboratory (LANL), which is operated by Triad National Security, LLC

src/interface/variable.cpp

@@ -1,5 +1,5 @@
 //========================================================================================
-// (C) (or copyright) 2020. Triad National Security, LLC. All rights reserved.
+// (C) (or copyright) 2020-2021. Triad National Security, LLC. All rights reserved.
 //
 // This program was produced under U.S. Government contract 89233218CNA000001 for Los
 // Alamos National Laboratory (LANL), which is operated by Triad National Security, LLC
@@ -14,6 +14,7 @@
 #include "interface/variable.hpp"
 
 #include <iostream>
+#include <utility>
 
 #include "bvals/cc/bvals_cc.hpp"
 #include "mesh/mesh.hpp"
@@ -69,13 +70,13 @@ CellVariable<T>::AllocateCopy(const bool allocComms, std::weak_ptr<MeshBlock> wp
       // Note that vbvar->var_cc will be set when stage is selected
       cv->vbvar = vbvar;
 
-      // fluxes, etc are always a copy
+      // fluxes, coarse buffers, etc., are always a copy
+      // Rely on reference counting and shallow copy of kokkos views
+      cv->flux_data_ = flux_data_; // reference counted
       for (int i = 1; i <= 3; i++) {
-        cv->flux[i] = flux[i];
+        cv->flux[i] = flux[i]; // these are subviews
       }
-
-      // These members are pointers,      // point at same memory as src
-      cv->coarse_s = coarse_s;
+      cv->coarse_s = coarse_s; // reference counted
     }
   }
   return cv;
@@ -85,28 +86,42 @@ CellVariable<T>::AllocateCopy(const bool allocComms, std::weak_ptr<MeshBlock> wp
 /// Initialize a 6D variable
 template <typename T>
 void CellVariable<T>::allocateComms(std::weak_ptr<MeshBlock> wpmb) {
-  // set up fluxes
   std::string base_name = label();
+
+  // TODO(JMM): Note that this approach assumes LayoutRight. Otherwise
+  // the stride will mess up the types.
+
+  // Compute size of unified flux_data object and create it. A unified
+  // flux_data_ object reduces the number of memory allocations per
+  // variable per meshblock from 5 to 3.
+  int n_outer = 0;
   if (IsSet(Metadata::Independent)) {
-    flux[X1DIR] = ParArrayND<T>(base_name + ".fluxX1", GetDim(6), GetDim(5), GetDim(4),
-                                GetDim(3), GetDim(2), GetDim(1));
-    if (GetDim(2) > 1)
-      flux[X2DIR] = ParArrayND<T>(base_name + ".fluxX2", GetDim(6), GetDim(5), GetDim(4),
-                                  GetDim(3), GetDim(2), GetDim(1));
-    if (GetDim(3) > 1)
-      flux[X3DIR] = ParArrayND<T>(base_name + ".fluxX3", GetDim(6), GetDim(5), GetDim(4),
-                                  GetDim(3), GetDim(2), GetDim(1));
+    n_outer += (GetDim(1) > 1) + (GetDim(2) > 1) + (GetDim(3) > 1);
   }
+  flux_data_ = ParArray7D<T>(base_name + ".flux_data", n_outer, GetDim(6), GetDim(5),
+                             GetDim(4), GetDim(3), GetDim(2), GetDim(1));
 
-  if (wpmb.expired()) return;
+  // set up fluxes
+  int offset = 0;
+  if (IsSet(Metadata::Independent)) {
+    for (int d = X1DIR; d <= X3DIR; ++d) {
+      if (GetDim(d) > 1) {
+        flux[d] = ParArrayND<T>(
+            Kokkos::subview(flux_data_, offset++, Kokkos::ALL(), Kokkos::ALL(),
+                            Kokkos::ALL(), Kokkos::ALL(), Kokkos::ALL(), Kokkos::ALL()));
+      }
+    }
+  }
 
+  if (wpmb.expired()) return;
   std::shared_ptr<MeshBlock> pmb = wpmb.lock();
 
-  if (pmb->pmy_mesh->multilevel)
+  if (pmb->pmy_mesh != nullptr && pmb->pmy_mesh->multilevel) {
     coarse_s = ParArrayND<T>(base_name + ".coarse", GetDim(6), GetDim(5), GetDim(4),
                              pmb->c_cellbounds.ncellsk(IndexDomain::entire),
                              pmb->c_cellbounds.ncellsj(IndexDomain::entire),
                              pmb->c_cellbounds.ncellsi(IndexDomain::entire));
+  }
 
   // Create the boundary object
   vbvar = std::make_shared<CellCenteredBoundaryVariable>(pmb, data, coarse_s, flux);

src/interface/variable.hpp

@@ -106,6 +106,7 @@ class CellVariable {
  private:
   Metadata m_;
   std::string label_;
+  ParArray7D<T> flux_data_; // unified par array for the fluxes
   int sparse_id_;
 };
 

src/kokkos_abstraction.hpp

@@ -67,6 +67,8 @@ template <typename T>
 using ParArray5D = Kokkos::View<T *****, LayoutWrapper, DevMemSpace>;
 template <typename T>
 using ParArray6D = Kokkos::View<T ******, LayoutWrapper, DevMemSpace>;
+template <typename T>
+using ParArray7D = Kokkos::View<T *******, LayoutWrapper, DevMemSpace>;
 
 using team_policy = Kokkos::TeamPolicy<>;
 using team_mbr_t = Kokkos::TeamPolicy<>::member_type;