DistDataCollector.h

#include <mpi.h>
#include <set>
#include <vector>
#include <limits>
 
#ifndef DIST_DATA_COLLECTOR
#define DIST_DATA_COLLECTOR
 

class DistDataCollector {
    
    private:
    
        // MPI communicator to use for communication
        MPI_Comm comm;
 
        // number of chunks
        std::size_t numChunks;
 
        // local size of the data
        std::size_t numSize;
 
        // the array that collects the data of size numChunks * numSize on rank 0
        double* collectedData;
        
        // MPI window for remote memory access
        MPI_Win win;
 
        public:
 
        // initial values
        const double BAD_VALUE = std::numeric_limits<double>::quiet_NaN();

    DistDataCollector(MPI_Comm comm, int numChunks, int numSize);

    ~DistDataCollector();

    void inline startEpoch() {
        MPI_Win_lock_all(MPI_MODE_NOCHECK, this->win);
    }

    void inline flush() {
        MPI_Win_flush_all(this->win);
    }

    void inline endEpoch() {
        MPI_Win_unlock_all(this->win);
    }   

    void put(int chunkId, const double* data);

    void inline putAsync(int chunkId, const double* data) {
        MPI_Put(data, this->numSize, MPI_DOUBLE, 0, chunkId * this->numSize, this->numSize, MPI_DOUBLE, this->win);
    }

    std::vector<double> get(int chunkId);

    void get(int chunkId, double* buffer);

    void inline getAsync(int chunkId, double* buffer) {
        MPI_Get(buffer, this->numSize, MPI_DOUBLE, 0, chunkId * this->numSize, this->numSize, MPI_DOUBLE, this->win);
    }

    double* getCollectedDataPtr() {
        return this->collectedData;
    }

    int getNumChunks() const {
        return this->numChunks;
    }

    int getNumSize() const {
        return this->numSize;
    }

    void free() {
        if (this->win != MPI_WIN_NULL) {
            MPI_Win_free(&this->win);
        }
        // No need to free the data, MPI_Win_free will free the pointer
        //MPI_Free_mem(this->collectedData);
    }
 
    // Disable copy and assignment operations
    // to prevent accidental copying of the DistDataCollector instance
    // This is important because the class manages an MPI window and data pointer
    // which should not be copied or assigned.
    DistDataCollector(const DistDataCollector&) = delete; // Disable copy constructor
    DistDataCollector& operator=(const DistDataCollector&) = delete; // Disable assignment operator
    DistDataCollector(DistDataCollector&& other) noexcept; // Move constructor
    DistDataCollector& operator=(DistDataCollector&& other) noexcept; // Move assignment operator
};
 
#endif // DIST_DATA_COLLECTOR