File mpicommreference.h

File List > code_source > templat > include > TempLat > parallel > mpi > comm > mpicommreference.h

Go to the documentation of this file

#ifndef TEMPLAT_PARALLEL_MPI_COMM_MPICOMMREFERENCE_H
#define TEMPLAT_PARALLEL_MPI_COMM_MPICOMMREFERENCE_H

/* This file is part of TempLat, available at https://cosmolattice.github.io/templat .
   Copyright 2021-2026 The TempLat authors, see AUTHORS.md.
   Released under the MIT license, see LICENSE.md. */

// File info: Main contributor(s): Wessel Valkenburg, Year: 2019

#include <map>

#include "TempLat/util/exception.h"
#include "TempLat/parallel/mpi/mpitypeconstants.h"
#include "TempLat/parallel/mpi/comm/exchange/mpiallreduce.h"

namespace TempLat
{

  MakeException(MPICommReferenceBookKeepingException);

  class MPICommReference : public MPIAllReduce
  {
  public:
    // Put public methods here. These should change very little over time.
    MPICommReference(MPI_Comm comm = MPI_COMM_WORLD)
        : MPIAllReduce(comm), mComm(comm), mSize(getSizeForComm(comm)), mRank(getRankInComm(comm))
    {
      BookKeeper(mComm, true, false);
    }

    MPICommReference(const MPICommReference &other)
        : MPIAllReduce(other.mComm), mComm(other.mComm), mSize(other.mSize), mRank(other.mRank)
    {
      BookKeeper(mComm, true, false);
    }

    MPICommReference &operator=(const MPICommReference &other)
    {
      if (this == &other) return *this;
#ifdef HAVE_MPI
      // Release the comm we currently hold before overwriting it (mirror the destructor), otherwise its
      // bookkeeping count leaks and a subgroup comm is never freed.
      ptrdiff_t referenceCount = BookKeeper(mComm, false, true);
      if (mComm != MPI_COMM_WORLD && mComm != MPI_COMM_NULL && referenceCount < 1) {
        MPI_Comm_free(&mComm);
      }
#endif
      mComm = other.mComm;
      mSize = other.mSize;
      mRank = other.mRank;
      // Keep the base reducer pointed at the new comm; otherwise Allreduce runs on the stale communicator.
      MPIAllReduce::operator=(other);
      BookKeeper(mComm, true, false);
      return *this;
    }

    ~MPICommReference()
    {
#ifdef HAVE_MPI
      ptrdiff_t referenceCount = BookKeeper(mComm, false, true);
      if (mComm != MPI_COMM_WORLD && mComm != MPI_COMM_NULL && referenceCount < 1) {
#ifdef CHECKBOUNDS
        sayMPI << "Freeing mpi comm: " << mComm << " with reference count " << referenceCount << "\n";
#endif
        MPI_Comm_free(&mComm);
      }
#endif
    }

    operator MPI_Comm() const { return mComm; }

    int size() const { return mSize; }

    int rank() const { return mRank; }
    int getRank() const { return rank(); }
    static inline int worldRank()
    {
      /* only for debug and testing purposes */
      static bool set = false;
      static int rank = -1;
      if (!set) {
        set = true;
        rank = MPICommReference().rank();
      }
      return rank;
    }

  private:
    MPI_Comm mComm;
    int mSize;
    int mRank;

    inline static ptrdiff_t BookKeeper(MPI_Comm comm, bool construct, bool destruct)
    {
      static std::map<MPI_Comm, ptrdiff_t> theBook;

      if (comm == MPI_COMM_WORLD) return 1; /* never free the world. */

      if (construct != !destruct)
        throw MPICommReferenceBookKeepingException("You called with contruct: ", construct, "destruct:", destruct);

      ptrdiff_t result = -1;

      if (construct) {

        if (theBook.count(comm) == 0) {
          theBook[comm] = 0;
        }
        result = ++theBook[comm];

      } else if (destruct) {

        if (theBook.count(comm) < 1) {
          throw MPICommReferenceBookKeepingException("Calling destruct for MPI_Comm which is not present in the Books.",
                                                     comm, theBook.count(comm));
        } else if (theBook[comm] < 1) {
          throw MPICommReferenceBookKeepingException(
              "Calling destruct for MPI_Comm which already has a reference count < 1.", comm,
              "reference count:", theBook[comm]);
        }
        result = --theBook[comm];
        if (result == 0) {
          theBook.erase(comm);
        }
      }

      return result;
    }

    inline static int getSizeForComm(MPI_Comm comm)
    {
      int result = 1;
#ifdef HAVE_MPI
      if (comm != MPI_COMM_NULL) {
        MPI_Comm_size(comm, &result);
      }
#endif
      return result;
    }

    inline static int getRankInComm(MPI_Comm comm)
    {
      int result = 0;
#ifdef HAVE_MPI
      if (comm != MPI_COMM_NULL) MPI_Comm_rank(comm, &result);
#endif
      return result;
    }
  };
} // namespace TempLat

#endif