File mpidomainsplit.h

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

Go to the documentation of this file

#ifndef TEMPLAT_PARALLEL_MPI_COMM_MPIDOMAINSPLIT_H
#define TEMPLAT_PARALLEL_MPI_COMM_MPIDOMAINSPLIT_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 "TempLat/util/factorize.h"

#ifdef HAVE_MPI
#include <mpi.h>
#endif

namespace TempLat
{

  MakeException(MPIDomainSplitException);

  class MPIDomainSplit
  {
  public:
    // Put public methods here. These should change very little over time.
    MPIDomainSplit(ptrdiff_t total, ptrdiff_t nDimensions, ptrdiff_t nDimensionsToSplit)
        : mTotal(total), mNDimensions(nDimensions), mNDimensionsToSplit(nDimensionsToSplit), mFactors(total),
          mDomainSplitting(nDimensions, 1) /* set to nDimension entries, each with value 1. */
    {
      if (nDimensionsToSplit > nDimensions)
        throw MPIDomainSplitException(
            "You provided a larger number of dimensions to split than the number of dimensions.", nDimensionsToSplit,
            " > ", nDimensions);

#ifdef HAVE_MPI
      // Create (D-1)-dimensional Cartesian processor grid
      for (int i = 0; i < nDimensionsToSplit; ++i) {
        mDomainSplitting[i] = 0; // initialize to 1
      }
      MPI_Dims_create(mTotal, nDimensionsToSplit, mDomainSplitting.data());
      return;
#endif

      /* reverse sort */
      mFactors.sort(true);

      if ((ptrdiff_t)mFactors.size() <= mNDimensionsToSplit) {
        std::copy(mFactors.begin(), mFactors.end(), mDomainSplitting.begin());
      } else {

        /* put the first entry (largest) in the first dimension.
           Then, subsequently, put the next factor in the first smallest dimension.
           That is, smallest, and if there are several smallest, take the first. */
        for (auto &&it : mFactors) {
          /* find the next smallest */
          ptrdiff_t smallest = 0;
          while (smallest < mNDimensionsToSplit - 1 && mDomainSplitting[smallest + 1] < mDomainSplitting[smallest])
            ++smallest;
          mDomainSplitting[smallest] *= it;
        }
      }

      verifySelf();
    }

    /* for copying etc. */
    auto begin() { return mDomainSplitting.begin(); }
    auto end() { return mDomainSplitting.end(); }

    operator std::vector<int>() const { return mDomainSplitting; }

    ptrdiff_t computeSize() { return end() - begin(); }

    friend std::ostream &operator<<(std::ostream &ostream, MPIDomainSplit dom)
    {
      ostream << "Domain splitting( " << dom.mTotal << ", " << dom.mNDimensions << ", " << dom.mNDimensionsToSplit
              << " ): " << dom.mDomainSplitting;
      //            ostream << "\nfor factors: " << dom.factors;
      return ostream;
    }

  private:
    /* Put all member variables and private methods here. These may change arbitrarily. */
    const ptrdiff_t mTotal, mNDimensions, mNDimensionsToSplit;
    Factorize mFactors;
    std::vector<int> mDomainSplitting;

    void verifySelf()
    {
      /* just a check that everything worked out ok */
      ptrdiff_t checkTotal = 1;
      ptrdiff_t checkNSplit = 0;
      for (auto &&it : mDomainSplitting) {
        checkTotal *= it;
        if (it != 1) ++checkNSplit;
      }
      if (checkTotal != mTotal)
        throw MPIDomainSplitException("Domain decomposition failed: ended up with ", checkTotal,
                                      "entries, instead of the imposed", mTotal);
      if (checkNSplit > mNDimensionsToSplit) {
        throw MPIDomainSplitException("Domain decomposition failed: split more dimensions than allowed.", checkNSplit,
                                      " > ", mNDimensionsToSplit);
      }
    }
  };
} // namespace TempLat

#endif