sharedRegionImpl.h

// Copyright 2009-2021 NTESS. Under the terms
// of Contract DE-NA0003525 with NTESS, the U.S.
// Government retains certain rights in this software.
//
// Copyright (c) 2009-2021, NTESS
// All rights reserved.
//
// This file is part of the SST software package. For license
// information, see the LICENSE file in the top level directory of the
// distribution.

#ifndef SST_CORE_CORE_SHAREDREGIONIMPL_H
#define SST_CORE_CORE_SHAREDREGIONIMPL_H

#include "sst/core/sst_types.h"
#include "sst/core/warnmacros.h"

#include <string>
#include <vector>
#include <map>
#include <mutex>

#include "sst/core/sharedRegion.h"
#include "sst/core/output.h"
#include "sst/core/serialization/serializable.h"

namespace SST {

class SharedRegionImpl;

class ChangeSet : public SST::Core::Serialization::serializable {
public:
    ChangeSet() { }
    size_t offset;
    size_t length;
    /*const*/ uint8_t *data;

    ChangeSet(size_t offset, size_t length, /*const*/ uint8_t *data = nullptr) : offset(offset), length(length), data(data) { }

    void serialize_order(SST::Core::Serialization::serializer &ser) override {
        ser & offset;
        // ser & length;
        // if ( ser.mode() == SST::Core::Serialization::serializer::UNPACK ) {
        //     data = new uint8_t[length];
        // }
        // for ( int i = 0; i < length; i++ ) {
        //     ser & data[i];
        // }
        ser & SST::Core::Serialization::array(data,length);
        // ser.binary(data,length);

    }

    ImplementSerializable(SST::ChangeSet)


};

class RegionInfo {
public:
    class RegionMergeInfo : public SST::Core::Serialization::serializable {

    protected:
        int rank;
        std::string key;
        size_t length;
        uint8_t initByte;

        virtual void check_size(RegionInfo *ri) {
            // Check to see if sizes match.

            // If the merge info has zero length, then it doesn't know
            // the size, so just return
            if ( length == 0 ) return;

            ri->setSize(length, initByte);
        }

    public:
        RegionMergeInfo() {}
        RegionMergeInfo(int rank, const std::string& key, size_t length, uint8_t initByte) :
            rank(rank), key(key), length(length), initByte(initByte) { }
        virtual ~RegionMergeInfo() { }

        virtual bool merge(RegionInfo *ri) {
            check_size(ri);
            return true;

        }
        const std::string& getKey() const { return key; }
        size_t getLength() const { return length; }

        void serialize_order(SST::Core::Serialization::serializer &ser) override {
            ser & rank;
            ser & key;
            ser & length;
            ser & initByte;
        }

        ImplementSerializable(SST::RegionInfo::RegionMergeInfo)
    };


    class BulkMergeInfo : public RegionMergeInfo {
    protected:
        void *data;

    public:
        BulkMergeInfo() : RegionMergeInfo() {}
        BulkMergeInfo(int rank, const std::string& key, void *data, size_t length, uint8_t initByte) :
            RegionMergeInfo(rank, key, length, initByte),
            data(data)
        { }

        bool merge(RegionInfo *ri) override {
            check_size(ri);
            bool ret = ri->getMerger()->merge((uint8_t*)ri->getMemory(), (const uint8_t*)data, length);
            free(data);
            return ret;
        }

        void serialize_order(SST::Core::Serialization::serializer &ser) override {
            RegionInfo::RegionMergeInfo::serialize_order(ser);

            // if ( ser.mode() == SST::Core::Serialization::serializer::UNPACK ) {
            //     data = malloc(length);
            // }
            // for ( int i = 0; i < length; i++ ) {
            //     ser & ((uint8_t*)data)[i];
            // }
            ser & SST::Core::Serialization::array(data,length);
            // ser.binary(data,length);
        }

        ImplementSerializable(SST::RegionInfo::BulkMergeInfo)
    };

    class ChangeSetMergeInfo : public RegionMergeInfo {
    protected:
        // std::vector<SharedRegionMerger::ChangeSet> changeSets;
        std::vector<ChangeSet> changeSets;
    public:
        ChangeSetMergeInfo() : RegionMergeInfo() {}
        ChangeSetMergeInfo(int rank, const std::string& key, size_t length, uint8_t initByte,
                           std::vector<ChangeSet> & changeSets) :
            RegionMergeInfo(rank, key, length, initByte),
            changeSets(changeSets)
        { }
        bool merge(RegionInfo *ri) override {
            check_size(ri);
            return ri->getMerger()->merge((uint8_t*)ri->getMemory(), ri->getSize(), changeSets);
        }

        void serialize_order(SST::Core::Serialization::serializer &ser) override {
            RegionInfo::RegionMergeInfo::serialize_order(ser);
            ser & changeSets;
        }

        ImplementSerializable(SST::RegionInfo::ChangeSetMergeInfo)
    };


private:
    std::string myKey;
    size_t realSize;
    size_t apparentSize;
    void *memory;

    size_t shareCount;
    size_t publishCount;

    std::vector<SharedRegionImpl*> sharers;

    SharedRegionMerger *merger; // If null, no multi-rank merging
    // std::vector<SharedRegionMerger::ChangeSet> changesets;
    std::vector<ChangeSet> changesets;

    bool didBulk;
    bool initialized;
    bool ready;
    uint8_t initByte;

public:
    RegionInfo() :
        realSize(0), apparentSize(0), memory(nullptr),
        shareCount(0), publishCount(0), merger(nullptr),
        didBulk(false), initialized(false), ready(false), initByte(0)
    { }
    ~RegionInfo();
    bool initialize(const std::string& key, size_t size, uint8_t initByte_in, SharedRegionMerger *mergeObj);
    bool setSize(size_t size, uint8_t initByte_in);
    bool isInitialized() const { return initialized; }
    bool isReady() const { return ready; }

    SharedRegionImpl* addSharer(SharedRegionManager *manager);
    void removeSharer(SharedRegionImpl *sri);


    void modifyRegion(size_t offset, size_t length, const void *data);
    void publish();

    void updateState(bool finalize);

    const std::string& getKey() const { return myKey; }
    void* getMemory() { didBulk = true; return memory; }
    const void* getConstPtr() const { return memory; }
    size_t getSize() const { return apparentSize; }
    size_t getNumSharers() const { return shareCount; }

    bool shouldMerge() const { return (nullptr != merger); }
    SharedRegionMerger* getMerger() { return merger; }
    /** Returns the size of the data to be transferred */
    RegionMergeInfo* getMergeInfo();

    void setProtected(bool readOnly);
};


class SharedRegionImpl : public SharedRegion {
    bool published;
    RegionInfo *region;

protected:
    size_t getSize() { return region->getSize(); }

public:
    SharedRegionImpl(SharedRegionManager *manager, size_t id,
            RegionInfo *region) : SharedRegion(manager, id),
        published(false), region(region)
    { }

    bool isPublished() const { return published; }
    void setPublished() { published = true; }
    RegionInfo* getRegion() const { return region; }
    void notifySetSize() {
        if ( deferred_pointer != nullptr ) {
            deferred_pointer->setPointer(getPtr<const void*>());
            delete deferred_pointer;
            deferred_pointer = nullptr;
        }
    }
};



class SharedRegionManagerImpl : public SharedRegionManager {

    struct CommInfo_t {
        RegionInfo *region;
        std::vector<int> tgtRanks;
        std::vector<char> sendBuffer;
    };

    std::map<std::string, RegionInfo> regions;
    std::mutex mtx;

protected:
    void modifyRegion(SharedRegion *sr, size_t offset, size_t length, const void *data) override;
    void* getMemory(SharedRegion *sr) override;
    const void* getConstPtr(const SharedRegion *sr) const override;
    size_t getSize(const SharedRegion *sr) const override;

public:
    SharedRegionManagerImpl();
    ~SharedRegionManagerImpl();

    virtual SharedRegion* getLocalSharedRegion(const std::string& key, size_t size, uint8_t initByte = 0) override;
    virtual SharedRegion* getGlobalSharedRegion(const std::string& key, size_t size, SharedRegionMerger *merger = nullptr, uint8_t initByte = 0) override;

    virtual void publishRegion(SharedRegion*) override;
    virtual bool isRegionReady(const SharedRegion*) override;
    virtual void shutdownSharedRegion(SharedRegion*) override;

    void updateState(bool finalize) override;
};


}
#endif