Add Blocks distribution strategy

Author: franzpoeschel

include/openPMD/ChunkInfo.hpp

@@ -235,6 +235,22 @@ namespace chunk_assignment
         virtual std::unique_ptr<Strategy> clone() const override;
     };
 
+    struct Blocks : Strategy
+    {
+    private:
+        unsigned int mpi_size, mpi_rank;
+
+    public:
+        Blocks(unsigned int mpi_rank, unsigned int mpi_size);
+
+        Assignment assign(
+            PartialAssignment,
+            RankMeta const &in,
+            RankMeta const &out) override;
+
+        [[nodiscard]] std::unique_ptr<Strategy> clone() const override;
+    };
+
     /**
      * @brief Strategy that assigns chunks to be read by processes within
      *        the same host that produced the chunk.

include/openPMD/benchmark/mpi/OneDimensionalBlockSlicer.hpp

@@ -33,6 +33,9 @@ class OneDimensionalBlockSlicer : public BlockSlicer
 
     explicit OneDimensionalBlockSlicer(Extent::value_type dim = 0);
 
+    static std::pair<size_t, size_t>
+    n_th_block_inside(size_t length, size_t rank, size_t size);
+
     std::pair<Offset, Extent>
     sliceBlock(Extent &totalExtent, int size, int rank) override;
 

src/ChunkInfo.cpp

@@ -22,13 +22,15 @@
 #include "openPMD/ChunkInfo_internal.hpp"
 
 #include "openPMD/auxiliary/Mpi.hpp"
+#include "openPMD/benchmark/mpi/OneDimensionalBlockSlicer.hpp"
 
 #include <algorithm> // std::sort
 #include <deque>
 #include <iostream>
 #include <iterator>
 #include <list>
 #include <map>
+#include <memory>
 #include <optional>
 #include <set>
 #include <utility>
@@ -335,6 +337,30 @@ namespace chunk_assignment
         return std::unique_ptr<Strategy>(new RoundRobinOfSourceRanks);
     }
 
+    Blocks::Blocks(unsigned int mpi_rank_in, unsigned int mpi_size_in)
+        : mpi_size(mpi_size_in), mpi_rank(mpi_rank_in)
+    {}
+
+    Assignment
+    Blocks::assign(PartialAssignment pa, RankMeta const &, RankMeta const &)
+    {
+        auto [notAssigned, res] = std::move(pa);
+        auto [myChunksFrom, myChunksTo] =
+            OneDimensionalBlockSlicer::n_th_block_inside(
+                notAssigned.size(), mpi_rank, mpi_size);
+        std::transform(
+            notAssigned.begin() + myChunksFrom,
+            notAssigned.begin() + (myChunksFrom + myChunksTo),
+            std::back_inserter(res[mpi_rank]),
+            [](WrittenChunkInfo &chunk) { return std::move(chunk); });
+        return res;
+    }
+
+    std::unique_ptr<Strategy> Blocks::clone() const
+    {
+        return std::unique_ptr<Strategy>(new Blocks(*this));
+    }
+
     ByHostname::ByHostname(std::unique_ptr<Strategy> withinNode)
         : m_withinNode(std::move(withinNode))
     {}

src/benchmark/mpi/OneDimensionalBlockSlicer.cpp

@@ -29,29 +29,23 @@ OneDimensionalBlockSlicer::OneDimensionalBlockSlicer(Extent::value_type dim)
     : m_dim{dim}
 {}
 
-std::pair<Offset, Extent>
-OneDimensionalBlockSlicer::sliceBlock(Extent &totalExtent, int size, int rank)
+std::pair<size_t, size_t> OneDimensionalBlockSlicer::n_th_block_inside(
+    size_t length, size_t rank, size_t size)
 {
-    Offset offs(totalExtent.size(), 0);
-
     if (rank >= size)
     {
-        Extent extent(totalExtent.size(), 0);
-        return std::make_pair(std::move(offs), std::move(extent));
+        return {length, 0};
     }
 
-    auto dim = this->m_dim;
-
     // for more equal balancing, we want the start index
     // at the upper gaussian bracket of (N/n*rank)
     // where N the size of the dataset in dimension dim
     // and n the MPI size
     // for avoiding integer overflow, this is the same as:
     // (N div n)*rank + round((N%n)/n*rank)
-    auto f = [&totalExtent, size, dim](int threadRank) {
-        auto N = totalExtent[dim];
-        auto res = (N / size) * threadRank;
-        auto padDivident = (N % size) * threadRank;
+    auto f = [length, size](size_t rank_lambda) {
+        auto res = (length / size) * rank_lambda;
+        auto padDivident = (length % size) * rank_lambda;
         auto pad = padDivident / size;
         if (pad * size < padDivident)
         {
@@ -60,17 +54,31 @@ OneDimensionalBlockSlicer::sliceBlock(Extent &totalExtent, int size, int rank)
         return res + pad;
     };
 
-    offs[dim] = f(rank);
+    size_t offset = f(rank);
+    size_t extent = [&]() {
+        if (rank >= size - 1)
+        {
+            return length - offset;
+        }
+        else
+        {
+            return f(rank + 1) - offset;
+        }
+    }();
+    return {offset, extent};
+}
+
+std::pair<Offset, Extent>
+OneDimensionalBlockSlicer::sliceBlock(Extent &totalExtent, int size, int rank)
+{
+    Offset localOffset(totalExtent.size(), 0);
     Extent localExtent{totalExtent};
-    if (rank >= size - 1)
-    {
-        localExtent[dim] -= offs[dim];
-    }
-    else
-    {
-        localExtent[dim] = f(rank + 1) - offs[dim];
-    }
-    return std::make_pair(std::move(offs), std::move(localExtent));
+
+    auto [offset_dim, extent_dim] =
+        n_th_block_inside(totalExtent.at(this->m_dim), rank, size);
+    localOffset[m_dim] = offset_dim;
+    localExtent[m_dim] = extent_dim;
+    return std::make_pair(std::move(localOffset), std::move(localExtent));
 }
 
 std::unique_ptr<BlockSlicer> OneDimensionalBlockSlicer::clone() const

src/binding/python/ChunkInfo.cpp

@@ -300,6 +300,11 @@ void init_Chunk(py::module &m)
     py::class_<RoundRobin, Strategy>(m, "RoundRobin").def(py::init<>());
     py::class_<RoundRobinOfSourceRanks, Strategy>(m, "RoundRobinOfSourceRanks")
         .def(py::init<>());
+    py::class_<Blocks, Strategy>(m, "Blocks")
+        .def(
+            py::init<unsigned int, unsigned int>(),
+            py::arg("mpi_rank"),
+            py::arg("mpi_size"));
 
     py::class_<ByHostname, PartialStrategy>(m, "ByHostname")
         .def(

test/ParallelIOTest.cpp

@@ -2455,6 +2455,11 @@ void adios2_chunk_distribution()
             chunkTable, rankMetaIn, readingRanksHostnames);
         printAssignment(
             "CUBOID SLICE", cuboidSliceAssignment, readingRanksHostnames);
+
+        Blocks blocksStrategy(mpi_rank, mpi_size);
+        auto blocksAssignment = blocksStrategy.assign(
+            chunkTable, rankMetaIn, readingRanksHostnames);
+        printAssignment("BLOCKS", blocksAssignment, readingRanksHostnames);
     }
 }