ADIOS2 and array index mapping

In general file I/O maps a logically rectangular subset of the local (nx, ny, nz) array onto a part of a global array. That global array may be sparse i.e. there might be holes in it. The mapping does not cover the entire local array because we insert communication guard cells around the outside. To perform this mapping we use ADIOS Put() function to get a span of a buffer, and then iterate over fields to copy data into those buffers.
boutproject · Sep 12, 2023 · 8c6a979 · 8c6a979
1 parent 04623d4
commit 8c6a979
Show file tree

Hide file tree

Showing 3 changed files with 175 additions and 18 deletions.
diff --git a/include/bout/mesh.hxx b/include/bout/mesh.hxx
@@ -503,9 +503,20 @@ public:
   int GlobalNx, GlobalNy, GlobalNz; ///< Size of the global arrays. Note: can have holes
   /// Size of the global arrays excluding boundary points.
   int GlobalNxNoBoundaries, GlobalNyNoBoundaries, GlobalNzNoBoundaries;
+
+  /// Note: These offsets only correct if Y guards are not included in the global array
+  ///       and are corrected in gridfromfile.cxx
   int OffsetX, OffsetY, OffsetZ; ///< Offset of this mesh within the global array
                                  ///< so startx on this processor is OffsetX in global
 
+  /// Map between local and global indices
+  /// (MapGlobalX, MapGlobalY, MapGlobalZ) in the global index space maps to (MapLocalX, MapLocalY, MapLocalZ) locally.
+  /// Note that boundary cells are included in the global index space, but communication
+  /// guard cells are not.
+  int MapGlobalX, MapGlobalY, MapGlobalZ; ///< Start global indices
+  int MapLocalX, MapLocalY, MapLocalZ;    ///< Start local indices
+  int MapCountX, MapCountY, MapCountZ;    ///< Size of the mapped region
+
   /// Returns the number of unique cells (i.e., ones not used for
   /// communication) on this processor for 3D fields. Boundaries
   /// are only included to a depth of 1.

diff --git a/src/mesh/impls/bout/boutmesh.cxx b/src/mesh/impls/bout/boutmesh.cxx
@@ -333,7 +333,7 @@ void BoutMesh::setDerivedGridSizes() {
   }
 
   GlobalNx = nx;
-  GlobalNy = ny + 2 * MYG;
+  GlobalNy = ny + 2 * MYG; // Note: For double null this should be be 4 * MYG if boundary cells are stored
   GlobalNz = nz;
 
   // If we've got a second pair of diverator legs, we need an extra
@@ -374,6 +374,7 @@ void BoutMesh::setDerivedGridSizes() {
   }
 
   // Set global offsets
+  // Note: These don't properly include guard/boundary cells
   OffsetX = PE_XIND * MXSUB;
   OffsetY = PE_YIND * MYSUB;
   OffsetZ = 0;
@@ -392,6 +393,48 @@ void BoutMesh::setDerivedGridSizes() {
 
   zstart = MZG;
   zend = MZG + MZSUB - 1;
+
+  // Mapping local to global indices
+  if (periodicX) {
+    // No boundary cells in X
+    MapGlobalX = PE_XIND * MXSUB;
+    MapLocalX = MXG;
+    MapCountX = MXSUB;
+  } else {
+    // X boundaries stored for firstX and lastX processors
+    if (firstX()) {
+      MapGlobalX = 0;
+      MapLocalX = 0;
+      MapCountX = MXG + MXSUB;
+    } else {
+      MapGlobalX = MXG + PE_XIND * MXSUB;
+      MapLocalX = MXG; // Guard cells not included
+      MapCountX = MXSUB;
+    }
+    if (lastX()) {
+      // Doesn't change the origin, but adds outer X boundary cells
+      MapCountX += MXG;
+    }
+  }
+
+  if (PE_YIND == 0) {
+    // Include Y boundary cells
+    MapGlobalY = 0;
+    MapLocalY = 0;
+    MapCountY = MYG + MYSUB;
+  } else {
+    MapGlobalY = MYG + PE_YIND * MYSUB;
+    MapLocalY = MYG;
+    MapCountY = MYSUB;
+  }
+  if (PE_YIND == NYPE - 1) {
+    // Include Y upper boundary region.
+    MapCountY += MYG;
+  }
+
+  MapGlobalZ = 0;
+  MapLocalZ = MZG; // Omit boundary cells
+  MapCountZ = MZSUB;
 }
 
 int BoutMesh::load() {

diff --git a/src/sys/options/options_adios.cxx b/src/sys/options/options_adios.cxx
@@ -223,46 +223,149 @@ struct ADIOSPutVarVisitor {
 
 template <>
 void ADIOSPutVarVisitor::operator()<bool>(const bool& value) {
+  // Scalars are only written from processor 0
+  if (BoutComm::rank() != 0) {
+    return;
+  }
+  adios2::Variable<int> var = stream.io.DefineVariable<int>(varname);
+  stream.engine.Put<int>(var, static_cast<int>(value));
+}
+
+template <>
+void ADIOSPutVarVisitor::operator()<int>(const int& value) {
+  // Scalars are only written from processor 0
+  if (BoutComm::rank() != 0) {
+    return;
+  }
   adios2::Variable<int> var = stream.io.DefineVariable<int>(varname);
-  stream.engine.Put<int>(var, (int)value);
+  stream.engine.Put<int>(var, value);
+}
+
+template <>
+void ADIOSPutVarVisitor::operator()<BoutReal>(const BoutReal& value) {
+  // Scalars are only written from processor 0
+  if (BoutComm::rank() != 0) {
+    return;
+  }
+  adios2::Variable<BoutReal> var = stream.io.DefineVariable<BoutReal>(varname);
+  stream.engine.Put<BoutReal>(var, value);
 }
 
 template <>
 void ADIOSPutVarVisitor::operator()<std::string>(const std::string& value) {
+  // Scalars are only written from processor 0
+  if (BoutComm::rank() != 0) {
+    return;
+  }
+
   adios2::Variable<std::string> var = stream.io.DefineVariable<std::string>(varname);
   std::cout << "-- Write string variable " << var.Name() << " value = " << value << std::endl;
-  stream.engine.Put<std::string>(var, (std::string)value, adios2::Mode::Sync);
+  stream.engine.Put<std::string>(var, value, adios2::Mode::Sync);
 }
 
 template <>
 void ADIOSPutVarVisitor::operator()<Field2D>(const Field2D& value) {
-  // Pointer to data. Assumed to be contiguous array
+  // Get the mesh that describes how local data relates to global arrays
   auto mesh = value.getMesh();
-  adios2::Dims shape = {(size_t)BoutComm::size(), (size_t)value.getNx(), (size_t)value.getNy()};
-  adios2::Dims start = {(size_t)BoutComm::rank(), 0, 0};
-  adios2::Dims count = {1, shape[1], shape[2]};
+
+  // The global size of this array includes boundary cells but not communication guard cells.
+  // In general this array will be sparse because it may have gaps.
+  adios2::Dims shape = {static_cast<size_t>(mesh->GlobalNx),
+                        static_cast<size_t>(mesh->GlobalNy)};
+
+  // Offset of this processor's data into the global array
+  adios2::Dims start = {static_cast<size_t>(mesh->MapGlobalX), static_cast<size_t>(mesh->MapGlobalY)};
+
+  // The size of the mapped region
+  adios2::Dims count = {static_cast<size_t>(mesh->MapCountX), static_cast<size_t>(mesh->MapCountY)};
   adios2::Variable<BoutReal> var = stream.io.DefineVariable<BoutReal>(varname, shape, start, count);
-  stream.engine.Put<BoutReal>(var, &value(0, 0));
+
+  // Get a Span of an internal engine buffer that we can fill with data
+  auto span = stream.engine.Put<BoutReal>(var);
+
+  // Iterate over the span and the Field2D array
+  // Note: The Field2D includes communication guard cells that are not written
+  auto it = span.begin();
+  for (int x = mesh->MapLocalX; x < (mesh->MapLocalX + mesh->MapCountX); ++x) {
+    for (int y = mesh->MapLocalY; y < (mesh->MapLocalY + mesh->MapCountY); ++y) {
+      *it = value(x, y);
+      ++it;
+    }
+  }
+  ASSERT1(it == span.end());
 }
 
 template <>
 void ADIOSPutVarVisitor::operator()<Field3D>(const Field3D& value) {
-  // Pointer to data. Assumed to be contiguous array
-  adios2::Dims shape = {(size_t)BoutComm::size(), (size_t)value.getNx(), (size_t)value.getNy(), (size_t)value.getNz()};
-  adios2::Dims start = {(size_t)BoutComm::rank(), 0, 0, 0};
-  adios2::Dims count = {1, shape[1], shape[2], shape[3]};
+  // Get the mesh that describes how local data relates to global arrays
+  auto mesh = value.getMesh();
+
+  // The global size of this array includes boundary cells but not communication guard cells.
+  // In general this array will be sparse because it may have gaps.
+  adios2::Dims shape = {static_cast<size_t>(mesh->GlobalNx),
+                        static_cast<size_t>(mesh->GlobalNy),
+                        static_cast<size_t>(mesh->GlobalNz)};
+
+  // Offset of this processor's data into the global array
+  adios2::Dims start = {static_cast<size_t>(mesh->MapGlobalX),
+                        static_cast<size_t>(mesh->MapGlobalY),
+                        static_cast<size_t>(mesh->MapGlobalZ)};
+
+  // The size of the mapped region
+  adios2::Dims count = {static_cast<size_t>(mesh->MapCountX),
+                        static_cast<size_t>(mesh->MapCountY),
+                        static_cast<size_t>(mesh->MapCountZ)};
   adios2::Variable<BoutReal> var = stream.io.DefineVariable<BoutReal>(varname, shape, start, count);
-  stream.engine.Put<BoutReal>(var, &value(0, 0, 0));
+
+  // Get a Span of an internal engine buffer.
+  auto span = stream.engine.Put<BoutReal>(var);
+
+  // Iterate over the span and the Field3D array
+  // Note: The Field3D includes communication guard cells that are not written
+  auto it = span.begin();
+  for (int x = mesh->MapLocalX; x < (mesh->MapLocalX + mesh->MapCountX); ++x) {
+    for (int y = mesh->MapLocalY; y < (mesh->MapLocalY + mesh->MapCountY); ++y) {
+      for (int z = mesh->MapLocalZ; z < (mesh->MapLocalZ + mesh->MapCountZ); ++z) {
+        *it = value(x, y, z);
+        ++it;
+      }
+    }
+  }
+  ASSERT1(it == span.end());
 }
 
 template <>
 void ADIOSPutVarVisitor::operator()<FieldPerp>(const FieldPerp& value) {
-  // Pointer to data. Assumed to be contiguous array
-  adios2::Dims shape = {(size_t)BoutComm::size(), (size_t)value.getNx(), (size_t)value.getNz()};
-  adios2::Dims start = {(size_t)BoutComm::rank(), 0, 0};
-  adios2::Dims count = {1, shape[1], shape[2]};
+  // Get the mesh that describes how local data relates to global arrays
+  auto mesh = value.getMesh();
+
+  // The global size of this array includes boundary cells but not communication guard cells.
+  // In general this array will be sparse because it may have gaps.
+  adios2::Dims shape = {static_cast<size_t>(mesh->GlobalNx),
+                        static_cast<size_t>(mesh->GlobalNz)};
+
+  // Offset of this processor's data into the global array
+  adios2::Dims start = {static_cast<size_t>(mesh->MapGlobalX),
+                        static_cast<size_t>(mesh->MapGlobalZ)};
+
+  // The size of the mapped region
+  adios2::Dims count = {static_cast<size_t>(mesh->MapCountX),
+                        static_cast<size_t>(mesh->MapCountZ)};
   adios2::Variable<BoutReal> var = stream.io.DefineVariable<BoutReal>(varname, shape, start, count);
-  stream.engine.Put<BoutReal>(var, &value(0, 0));
+
+  // Get a Span of an internal engine buffer.
+  auto span = stream.engine.Put<BoutReal>(var);
+
+  // Iterate over the span and the Field3D array
+  // Note: The Field3D includes communication guard cells that are not written
+  auto it = span.begin();
+  for (int x = mesh->MapLocalX; x < (mesh->MapLocalX + mesh->MapCountX); ++x) {
+    for (int z = mesh->MapLocalZ; z < (mesh->MapLocalZ + mesh->MapCountZ); ++z) {
+      *it = value(x, z);
+      ++it;
+    }
+  }
+  ASSERT1(it == span.end());
 }
 
 template <>