MPIBenchmark.hpp

/* Copyright 2018-2025 Franz Poeschel, Axel Huebl
 *
 * This file is part of openPMD-api.
 *
 * openPMD-api is free software: you can redistribute it and/or modify
 * it under the terms of of either the GNU General Public License or
 * the GNU Lesser General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * openPMD-api is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License and the GNU Lesser General Public License
 * for more details.
 *
 * You should have received a copy of the GNU General Public License
 * and the GNU Lesser General Public License along with openPMD-api.
 * If not, see <http://www.gnu.org/licenses/>.
 */
 
#pragma once
 
#include "openPMD/config.hpp"
#if openPMD_HAVE_MPI
 
#include "RandomDatasetFiller.hpp"
 
#include "openPMD/DatatypeHelpers.hpp"
#include "openPMD/benchmark/mpi/BlockSlicer.hpp"
#include "openPMD/benchmark/mpi/DatasetFiller.hpp"
#include "openPMD/benchmark/mpi/MPIBenchmarkReport.hpp"
#include "openPMD/openPMD.hpp"
 
#include <mpi.h>
 
#include <chrono>
#include <exception>
#include <iostream>
#include <set>
#include <sstream>
#include <tuple>
#include <utility>
#include <vector>
 
namespace openPMD
{
template <typename DatasetFillerProvider>
class MPIBenchmark
{
 
public:
    using extentT = Extent::value_type;
    MPI_Comm communicator = MPI_COMM_WORLD;

    Extent totalExtent;
 
    std::shared_ptr<BlockSlicer> m_blockSlicer;
 
    DatasetFillerProvider m_dfp;

    MPIBenchmark(
        std::string basePath,
        Extent tExtent,
        std::shared_ptr<BlockSlicer> blockSlicer,
        DatasetFillerProvider dfp,
        MPI_Comm comm = MPI_COMM_WORLD);

    void addConfiguration(
        std::string jsonConfig,
        std::string backend,
        Datatype dt,
        Series::IterationIndex_t iterations,
        int threadSize);

    void addConfiguration(
        std::string jsonConfig,
        std::string backend,
        Datatype dt,
        Series::IterationIndex_t iterations);
 
    void resetConfigurations();

    template <typename Clock>
    MPIBenchmarkReport<typename Clock::duration>
    runBenchmark(int rootThread = 0);
 
private:
    std::string m_basePath;
    std::vector<std::tuple<
        std::string,
        std::string,
        int,
        Datatype,
        Series::IterationIndex_t>>
        m_configurations;
 
    enum Config
    {
        JSON_CONFIG = 0,
        BACKEND,
        NRANKS,
        DTYPE,
        ITERATIONS
    };
 
    std::pair<Offset, Extent> slice(int size);

    template <typename Clock>
    struct BenchmarkExecution
    {
        MPIBenchmark<DatasetFillerProvider> *m_benchmark;
 
        explicit BenchmarkExecution(
            MPIBenchmark<DatasetFillerProvider> *benchmark)
            : m_benchmark{benchmark}
        {}

        template <typename T>
        typename Clock::duration writeBenchmark(
            std::string const &jsonConfig,
            Offset &offset,
            Extent &extent,
            std::string const &extension,
            std::shared_ptr<DatasetFiller<T>> datasetFiller,
            Series::IterationIndex_t iterations);

        template <typename T>
        typename Clock::duration readBenchmark(
            Offset &offset,
            Extent &extent,
            std::string extension,
            Series::IterationIndex_t iterations);
 
        template <typename T>
        static void call(
            BenchmarkExecution<Clock> &,
            MPIBenchmarkReport<typename Clock::duration> &report,
            int rootThread = 0);
 
        static constexpr char const *errorMsg = "BenchmarkExecution";
    };
};
 
// Implementation
 
template <typename DatasetFillerProvider>
template <typename Clock>
MPIBenchmarkReport<typename Clock::duration>
MPIBenchmark<DatasetFillerProvider>::runBenchmark(int rootThread)
{
    MPIBenchmarkReport<typename Clock::duration> res{this->communicator};
    BenchmarkExecution<Clock> exec{this};
 
    std::set<Datatype> datatypes;
    for (auto const &conf : m_configurations)
    {
        datatypes.insert(std::get<DTYPE>(conf));
    }
    for (Datatype dt : datatypes)
    {
        switchDatasetType<BenchmarkExecution<Clock>>(dt, exec, res, rootThread);
    }
 
    return res;
}
 
template <typename DatasetFillerProvider>
MPIBenchmark<DatasetFillerProvider>::MPIBenchmark(
    std::string basePath,
    Extent tExtent,
    std::shared_ptr<BlockSlicer> blockSlicer,
    DatasetFillerProvider dfp,
    MPI_Comm comm)
    : communicator{comm}
    , totalExtent{std::move(tExtent)}
    , m_blockSlicer{std::move(blockSlicer)}
    , m_dfp{dfp}
    , m_basePath{std::move(basePath)}
{
    if (m_blockSlicer == nullptr)
        throw std::runtime_error("Argument blockSlicer cannot be a nullptr!");
}
 
template <typename DatasetFillerProvider>
std::pair<Offset, Extent> MPIBenchmark<DatasetFillerProvider>::slice(int size)
{
    int actualSize;
    MPI_Comm_size(this->communicator, &actualSize);
    int rank;
    MPI_Comm_rank(this->communicator, &rank);
    size = std::min(size, actualSize);
    return m_blockSlicer->sliceBlock(totalExtent, size, rank);
}
 
template <typename DatasetFillerProvider>
void MPIBenchmark<DatasetFillerProvider>::addConfiguration(
    std::string jsonConfig,
    std::string backend,
    Datatype dt,
    Series::IterationIndex_t iterations,
    int threadSize)
{
    this->m_configurations.emplace_back(
        std::move(jsonConfig), backend, threadSize, dt, iterations);
}
 
template <typename DatasetFillerProvider>
void MPIBenchmark<DatasetFillerProvider>::addConfiguration(
    std::string jsonConfig,
    std::string backend,
    Datatype dt,
    Series::IterationIndex_t iterations)
{
    int size;
    MPI_Comm_size(communicator, &size);
    addConfiguration(std::move(jsonConfig), backend, dt, iterations, size);
}
 
template <typename DatasetFillerProvider>
void MPIBenchmark<DatasetFillerProvider>::resetConfigurations()
{
    this->m_configurations.clear();
}
 
template <typename DatasetFillerProvider>
template <typename Clock>
template <typename T>
typename Clock::duration
MPIBenchmark<DatasetFillerProvider>::BenchmarkExecution<Clock>::writeBenchmark(
    std::string const &jsonConfig,
    Offset &offset,
    Extent &extent,
    std::string const &extension,
    std::shared_ptr<DatasetFiller<T>> datasetFiller,
    Series::IterationIndex_t iterations)
{
    MPI_Barrier(m_benchmark->communicator);
    auto start = Clock::now();
 
    // open file for writing
    Series series = Series(
        m_benchmark->m_basePath + "." + extension,
        Access::CREATE,
        m_benchmark->communicator,
        jsonConfig);
 
    typename Clock::duration ignore{};
 
    for (Series::IterationIndex_t i = 0; i < iterations; i++)
    {
        auto ignore_start = Clock::now();
        auto writeData = datasetFiller->produceData();
        auto ignore_end = Clock::now();
        ignore += ignore_end - ignore_start;
 
        MeshRecordComponent id = series.writeIterations()[i]
                                     .meshes["id"][MeshRecordComponent::SCALAR];
 
        Datatype datatype = determineDatatype(writeData);
        Dataset dataset = Dataset(datatype, m_benchmark->totalExtent);
 
        id.resetDataset(dataset);
 
        series.flush();
 
        id.storeChunk<T>(writeData, offset, extent);
        series.flush();
    }
    series.close();
    MPI_Barrier(m_benchmark->communicator);
    auto end = Clock::now();
 
    return (end - start) - ignore;
}
 
template <typename DatasetFillerProvider>
template <typename Clock>
template <typename T>
typename Clock::duration
MPIBenchmark<DatasetFillerProvider>::BenchmarkExecution<Clock>::readBenchmark(
    Offset &offset,
    Extent &extent,
    std::string extension,
    Series::IterationIndex_t iterations)
{
    MPI_Barrier(m_benchmark->communicator);
    // let every thread measure time
    auto start = Clock::now();
 
    Series series = Series(
        m_benchmark->m_basePath + "." + extension,
        Access::READ_ONLY,
        m_benchmark->communicator);
 
    for (Series::IterationIndex_t i = 0; i < iterations; i++)
    {
        MeshRecordComponent id =
            series.snapshots()[i].meshes["id"][MeshRecordComponent::SCALAR];
 
        auto chunk_data = id.loadChunk<T>(offset, extent);
        series.flush();
    }
 
    MPI_Barrier(m_benchmark->communicator);
    auto end = Clock::now();
    return end - start;
}
 
template <typename DatasetFillerProvider>
template <typename Clock>
template <typename T>
void MPIBenchmark<DatasetFillerProvider>::BenchmarkExecution<Clock>::call(
    BenchmarkExecution<Clock> &exec,
    MPIBenchmarkReport<typename Clock::duration> &report,
    int rootThread)
{
    Datatype dt = determineDatatype<T>();
    auto dsf = std::dynamic_pointer_cast<DatasetFiller<T>>(
        exec.m_benchmark->m_dfp.template operator()<T>());
    for (auto const &config : exec.m_benchmark->m_configurations)
    {
        std::string jsonConfig;
        std::string backend;
        int size;
        Datatype dt2;
        Series::IterationIndex_t iterations;
        std::tie(jsonConfig, backend, size, dt2, iterations) = config;
 
        if (dt != dt2)
        {
            continue;
        }
 
        auto localCuboid = exec.m_benchmark->slice(size);
 
        extentT blockSize = 1;
        for (auto ext : localCuboid.second)
        {
            blockSize *= ext;
        }
        dsf->setNumberOfItems(blockSize);
 
        auto writeTime = exec.writeBenchmark<T>(
            jsonConfig,
            localCuboid.first,
            localCuboid.second,
            backend,
            dsf,
            iterations);
        auto readTime = exec.readBenchmark<T>(
            localCuboid.first, localCuboid.second, backend, iterations);
        report.addReport(
            rootThread,
            jsonConfig,
            backend,
            size,
            dt2,
            iterations,
            std::make_pair(writeTime, readTime));
    }
}
} // namespace openPMD
 
#endif