Add BlocksOfSourceRanks strategy

franzpoeschel · Nov 15, 2024 · e509ed6 · e509ed6
1 parent adbdfaa
commit e509ed6
Show file tree

Hide file tree

Showing 4 changed files with 70 additions and 0 deletions.
diff --git a/include/openPMD/ChunkInfo.hpp b/include/openPMD/ChunkInfo.hpp
@@ -251,6 +251,22 @@ namespace chunk_assignment
         [[nodiscard]] std::unique_ptr<Strategy> clone() const override;
     };
 
+    struct BlocksOfSourceRanks : Strategy
+    {
+    private:
+        unsigned int mpi_size, mpi_rank;
+
+    public:
+        BlocksOfSourceRanks(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.

diff --git a/src/ChunkInfo.cpp b/src/ChunkInfo.cpp
@@ -361,6 +361,47 @@ namespace chunk_assignment
         return std::unique_ptr<Strategy>(new Blocks(*this));
     }
 
+    BlocksOfSourceRanks::BlocksOfSourceRanks(
+        unsigned int mpi_rank_in, unsigned int mpi_size_in)
+        : mpi_size(mpi_size_in), mpi_rank(mpi_rank_in)
+    {}
+
+    Assignment BlocksOfSourceRanks::assign(
+        PartialAssignment pa, RankMeta const &, RankMeta const &)
+    {
+        auto [notAssigned, res] = std::move(pa);
+        std::map<unsigned int, std::deque<WrittenChunkInfo>>
+            sortSourceChunksBySourceRank;
+        for (auto &chunk : notAssigned)
+        {
+            auto sourceID = chunk.sourceID;
+            sortSourceChunksBySourceRank[sourceID].push_back(std::move(chunk));
+        }
+        notAssigned.clear();
+        auto [myChunksFrom, myChunksTo] =
+            OneDimensionalBlockSlicer::n_th_block_inside(
+                sortSourceChunksBySourceRank.size(), mpi_rank, mpi_size);
+        auto it = sortSourceChunksBySourceRank.begin();
+        for (size_t i = 0; i < myChunksFrom; ++i)
+        {
+            ++it;
+        }
+        for (size_t i = 0; i < myChunksTo; ++i, ++it)
+        {
+            std::transform(
+                it->second.begin(),
+                it->second.end(),
+                std::back_inserter(res[mpi_rank]),
+                [](WrittenChunkInfo &chunk) { return std::move(chunk); });
+        }
+        return res;
+    }
+
+    std::unique_ptr<Strategy> BlocksOfSourceRanks::clone() const
+    {
+        return std::unique_ptr<Strategy>(new BlocksOfSourceRanks(*this));
+    }
+
     ByHostname::ByHostname(std::unique_ptr<Strategy> withinNode)
         : m_withinNode(std::move(withinNode))
     {}

diff --git a/src/binding/python/ChunkInfo.cpp b/src/binding/python/ChunkInfo.cpp
@@ -305,6 +305,11 @@ void init_Chunk(py::module &m)
             py::init<unsigned int, unsigned int>(),
             py::arg("mpi_rank"),
             py::arg("mpi_size"));
+    py::class_<BlocksOfSourceRanks, Strategy>(m, "BlocksOfSourceRanks")
+        .def(
+            py::init<unsigned int, unsigned int>(),
+            py::arg("mpi_rank"),
+            py::arg("mpi_size"));
 
     py::class_<ByHostname, PartialStrategy>(m, "ByHostname")
         .def(

diff --git a/test/ParallelIOTest.cpp b/test/ParallelIOTest.cpp
@@ -2460,6 +2460,14 @@ void adios2_chunk_distribution()
         auto blocksAssignment = blocksStrategy.assign(
             chunkTable, rankMetaIn, readingRanksHostnames);
         printAssignment("BLOCKS", blocksAssignment, readingRanksHostnames);
+
+        BlocksOfSourceRanks blocksOfSourceRanksStrategy(mpi_rank, mpi_size);
+        auto blocksOfSourceRanksAssignment = blocksOfSourceRanksStrategy.assign(
+            chunkTable, rankMetaIn, readingRanksHostnames);
+        printAssignment(
+            "BLOCKS OF SOURCE RANKS",
+            blocksOfSourceRanksAssignment,
+            readingRanksHostnames);
     }
 }