Implement memory selection in ADIOS2

Author: franzpoeschel

include/openPMD/Dataset.hpp

@@ -34,6 +34,12 @@ namespace openPMD
 using Extent = std::vector<std::uint64_t>;
 using Offset = std::vector<std::uint64_t>;
 
+struct MemorySelection
+{
+    Offset offset;
+    Extent extent;
+};
+
 class Dataset
 {
     friend class RecordComponent;

include/openPMD/IO/ADIOS/ADIOS2File.hpp

@@ -20,11 +20,13 @@
  */
 #pragma once
 
+#include "openPMD/Dataset.hpp"
 #include "openPMD/IO/ADIOS/ADIOS2Auxiliary.hpp"
 #include "openPMD/IO/AbstractIOHandler.hpp"
 #include "openPMD/IO/IOTask.hpp"
 #include "openPMD/IO/InvalidatableFile.hpp"
 #include "openPMD/config.hpp"
+#include <optional>
 
 #if openPMD_HAVE_ADIOS2
 #include <adios2.h>
@@ -106,6 +108,7 @@ struct BufferedUniquePtrPut
     std::string name;
     Offset offset;
     Extent extent;
+    std::optional<MemorySelection> memorySelection;
     UniquePtrWithLambda<void> data;
     Datatype dtype = Datatype::UNDEFINED;
 

include/openPMD/IO/ADIOS/ADIOS2IOHandler.hpp

@@ -20,6 +20,7 @@
  */
 #pragma once
 
+#include "openPMD/Dataset.hpp"
 #include "openPMD/Error.hpp"
 #include "openPMD/IO/ADIOS/ADIOS2Auxiliary.hpp"
 #include "openPMD/IO/ADIOS/ADIOS2FilePosition.hpp"
@@ -398,6 +399,7 @@ class ADIOS2IOHandlerImpl
     adios2::Variable<T> verifyDataset(
         Offset const &offset,
         Extent const &extent,
+        std::optional<MemorySelection> const &memorySelection,
         adios2::IO &IO,
         std::string const &varName)
     {
@@ -479,6 +481,16 @@ class ADIOS2IOHandlerImpl
         var.SetSelection(
             {adios2::Dims(offset.begin(), offset.end()),
              adios2::Dims(extent.begin(), extent.end())});
+        if (memorySelection.has_value())
+        {
+            var.SetMemorySelection(
+                {adios2::Dims(
+                     memorySelection->offset.begin(),
+                     memorySelection->offset.end()),
+                 adios2::Dims(
+                     memorySelection->extent.begin(),
+                     memorySelection->extent.end())});
+        }
         return var;
     }
 

include/openPMD/IO/IOTask.hpp

@@ -33,6 +33,7 @@
 #include <cstddef>
 #include <map>
 #include <memory>
+#include <optional>
 #include <string>
 #include <utility>
 #include <variant>
@@ -419,6 +420,7 @@ struct OPENPMDAPI_EXPORT Parameter<Operation::WRITE_DATASET>
 
     Extent extent = {};
     Offset offset = {};
+    std::optional<MemorySelection> memorySelection = std::nullopt;
     Datatype dtype = Datatype::UNDEFINED;
     auxiliary::WriteBuffer data;
 };

include/openPMD/LoadStoreChunk.hpp

@@ -16,12 +16,6 @@ class ConfigureStoreChunkFromBuffer;
 template <typename T>
 class DynamicMemoryView;
 
-struct MemorySelection
-{
-    Offset offset;
-    Extent extent;
-};
-
 namespace internal
 {
     struct StoreChunkConfig

src/IO/ADIOS/ADIOS2File.cpp

@@ -23,6 +23,7 @@
 #include "openPMD/Error.hpp"
 #include "openPMD/IO/ADIOS/ADIOS2IOHandler.hpp"
 #include "openPMD/auxiliary/Environment.hpp"
+#include <optional>
 
 #if openPMD_USE_VERIFY
 #define VERIFY(CONDITION, TEXT)                                                \
@@ -55,8 +56,8 @@ void DatasetReader::call(
     adios2::Engine &engine,
     std::string const &fileName)
 {
-    adios2::Variable<T> var =
-        impl->verifyDataset<T>(bp.param.offset, bp.param.extent, IO, bp.name);
+    adios2::Variable<T> var = impl->verifyDataset<T>(
+        bp.param.offset, bp.param.extent, std::nullopt, IO, bp.name);
     if (!var)
     {
         throw std::runtime_error(
@@ -85,7 +86,11 @@ void WriteDataset::call(ADIOS2File &ba, detail::BufferedPut &bp)
                 auto ptr = static_cast<T const *>(arg.get());
 
                 adios2::Variable<T> var = ba.m_impl->verifyDataset<T>(
-                    bp.param.offset, bp.param.extent, ba.m_IO, bp.name);
+                    bp.param.offset,
+                    bp.param.extent,
+                    bp.param.memorySelection,
+                    ba.m_IO,
+                    bp.name);
 
                 ba.getEngine().Put(var, ptr);
             }
@@ -97,6 +102,7 @@ void WriteDataset::call(ADIOS2File &ba, detail::BufferedPut &bp)
                 bput.name = std::move(bp.name);
                 bput.offset = std::move(bp.param.offset);
                 bput.extent = std::move(bp.param.extent);
+                bput.memorySelection = std::move(bp.param.memorySelection);
                 /*
                  * Note: Moving is required here since it's a unique_ptr.
                  * std::forward<>() would theoretically work, but it
@@ -143,7 +149,11 @@ struct RunUniquePtrPut
     {
         auto ptr = static_cast<T const *>(bufferedPut.data.get());
         adios2::Variable<T> var = ba.m_impl->verifyDataset<T>(
-            bufferedPut.offset, bufferedPut.extent, ba.m_IO, bufferedPut.name);
+            bufferedPut.offset,
+            bufferedPut.extent,
+            bufferedPut.memorySelection,
+            ba.m_IO,
+            bufferedPut.name);
         ba.getEngine().Put(var, ptr);
     }
 

src/IO/ADIOS/ADIOS2IOHandler.cpp

@@ -39,6 +39,7 @@
 #include <iostream>
 #include <iterator>
 #include <memory>
+#include <optional>
 #include <set>
 #include <stdexcept>
 #include <string>
@@ -1076,7 +1077,7 @@ namespace detail
             auto &IO = ba.m_IO;
             auto &engine = ba.getEngine();
             adios2::Variable<T> variable = impl->verifyDataset<T>(
-                params.offset, params.extent, IO, varName);
+                params.offset, params.extent, std::nullopt, IO, varName);
             adios2::Dims offset(params.offset.begin(), params.offset.end());
             adios2::Dims extent(params.extent.begin(), params.extent.end());
             variable.SetSelection({std::move(offset), std::move(extent)});

src/RecordComponent.cpp

@@ -510,6 +510,7 @@ void RecordComponent::storeChunk_impl(
     Parameter<Operation::WRITE_DATASET> dWrite;
     dWrite.offset = std::move(o);
     dWrite.extent = std::move(e);
+    dWrite.memorySelection = memorySelection;
     dWrite.dtype = dtype;
     /* std::static_pointer_cast correctly reference-counts the pointer */
     dWrite.data = std::move(buffer);