Skip to content

File memorytoolbox.h

File List > code_source > templat > include > TempLat > lattice > memory > memorytoolbox.h

Go to the documentation of this file

#ifndef TEMPLAT_LATTICE_MEMORY_MEMORYTOOLBOX_H
#define TEMPLAT_LATTICE_MEMORY_MEMORYTOOLBOX_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/parallel/mpi/cartesian/mpicartesiangroup.h"
#include "TempLat/fft/fftlibraryselector.h"
#include "TempLat/fft/fftmpidomainsplit.h"
#include "TempLat/fft/fftnormalization.h"

#include "TempLat/lattice/memory/triplestatelayouts.h"

#include "TempLat/lattice/ghostcells/ghostbuster.h"
#include "TempLat/lattice/ghostcells/ghostupdater.h"

#include "TempLat/lattice/memory/verbositylevels.h"
#include "TempLat/util/makeuniformarray.h"

#include "TempLat/parallel/device_memory.h"

#include <cstddef>
#include <sstream>
#include <stdexcept>
// #include "TempLat/lattice/algebra/shift.h"

namespace TempLat
{
  MakeException(InconsistentDomainSplitting);
  MakeException(InconsistentDimensions);

  template <size_t _NDim> class MemoryToolBox
  {
#ifdef HAVE_MPI
    static_assert(_NDim > 1, "MPI builds do not support 1D lattices. FFTW's MPI interface returns a NULL plan "
                             "for rank=1 R2C transforms (which would segfault on execute), and ParaFaFT requires "
                             "NDim >= 2. Use NDim >= 2, or build TempLat without MPI for 1D problems.");
#endif

  public:
    // Put public methods here. These should change very little over time.

    static constexpr size_t NDim = _NDim;

    MemoryToolBox(device::Idx nGridPoints, device::Idx ghostDepth, bool forbidTransposition = false)
        : MemoryToolBox(FFTMPIDomainSplit<NDim>::makeMPIGroup(makeUniformArray<device::Idx, NDim>(nGridPoints)),
                        nGridPoints, ghostDepth, forbidTransposition)
    {
    }
    static inline device::memory::host_ptr<MemoryToolBox<NDim>>
    makeShared(device::Idx nGridPoints, device::Idx ghostDepth, bool forbidTransposition = false)
    {
      return device::memory::host_ptr<MemoryToolBox<NDim>>(nGridPoints, ghostDepth, forbidTransposition);
    }

    MemoryToolBox(device::IdxArray<NDim> nGrid, device::Idx ghostDepth, bool forbidTransposition = false)
        : MemoryToolBox(FFTMPIDomainSplit<NDim>::makeMPIGroup(nGrid), nGrid, ghostDepth, forbidTransposition)
    {
    }
    static inline device::memory::host_ptr<MemoryToolBox<NDim>>
    makeShared(device::IdxArray<NDim> nGrid, device::Idx ghostDepth, bool forbidTransposition = false)
    {
      return device::memory::host_ptr<MemoryToolBox<NDim>>(nGrid, ghostDepth, forbidTransposition);
    }

    MemoryToolBox(device::Idx nGridPoints, device::Idx ghostDepth, device::Idx nThreads, bool forbidTransposition)
        : MemoryToolBox(FFTMPIDomainSplit<NDim>::makeMPIGroup(makeUniformArray<device::Idx, NDim>(nGridPoints)),
                        nGridPoints, ghostDepth, nThreads, forbidTransposition)
    {
    }
    static inline device::memory::host_ptr<MemoryToolBox<NDim>>
    makeShared(device::Idx nGridPoints, device::Idx ghostDepth, device::Idx nThreads, bool forbidTransposition = false)
    {
      return device::memory::host_ptr<MemoryToolBox<NDim>>(nGridPoints, ghostDepth, nThreads, forbidTransposition);
    }

    MemoryToolBox(device::IdxArray<NDim> nGrid, device::Idx ghostDepth, device::Idx nThreads, bool forbidTransposition)
        : MemoryToolBox(FFTMPIDomainSplit<NDim>::makeMPIGroup(nGrid), nGrid, ghostDepth, nThreads, forbidTransposition)
    {
    }
    static inline device::memory::host_ptr<MemoryToolBox<NDim>> makeShared(device::IdxArray<NDim> nGrid,
                                                                           device::Idx ghostDepth, device::Idx nThreads,
                                                                           bool forbidTransposition = false)
    {
      return device::memory::host_ptr<MemoryToolBox<NDim>>(nGrid, ghostDepth, nThreads, forbidTransposition);
    }

    MemoryToolBox(MPICommReference comm, device::Idx nGridPoints, device::Idx ghostDepth,
                  bool forbidTransposition = false)
        : MemoryToolBox(FFTMPIDomainSplit<NDim>::makeMPIGroup(comm, makeUniformArray<device::Idx, NDim>(nGridPoints)),
                        nGridPoints, ghostDepth, forbidTransposition)
    {
    }
    static inline device::memory::host_ptr<MemoryToolBox<NDim>>
    makeShared(MPICommReference comm, device::Idx nGridPoints, device::Idx ghostDepth, bool forbidTransposition = false)
    {
      return device::memory::host_ptr<MemoryToolBox<NDim>>(comm, nGridPoints, ghostDepth, forbidTransposition);
    }

    MemoryToolBox(MPICommReference comm, device::IdxArray<NDim> nGrid, device::Idx ghostDepth,
                  bool forbidTransposition = false)
        : MemoryToolBox(FFTMPIDomainSplit<NDim>::makeMPIGroup(comm, nGrid), nGrid, ghostDepth, forbidTransposition)
    {
    }
    static inline device::memory::host_ptr<MemoryToolBox<NDim>> makeShared(MPICommReference comm,
                                                                           device::IdxArray<NDim> nGrid,
                                                                           device::Idx ghostDepth,
                                                                           bool forbidTransposition = false)
    {
      return device::memory::host_ptr<MemoryToolBox<NDim>>(comm, nGrid, ghostDepth, forbidTransposition);
    }

    MemoryToolBox(MPICartesianGroup group, device::Idx nGridPoints, device::Idx ghostDepth,
                  bool forbidTransposition = false)
        : MemoryToolBox(group, nGridPoints, ghostDepth, 1, forbidTransposition)
    {
    }
    static inline device::memory::host_ptr<MemoryToolBox<NDim>> makeShared(MPICartesianGroup group,
                                                                           device::Idx nGridPoints,
                                                                           device::Idx ghostDepth,
                                                                           bool forbidTransposition = false)
    {
      return device::memory::host_ptr<MemoryToolBox<NDim>>(group, nGridPoints, ghostDepth, 1, forbidTransposition);
    }

    MemoryToolBox(MPICartesianGroup group, device::IdxArray<NDim> nGrid, device::Idx ghostDepth,
                  bool forbidTransposition = false)
        : MemoryToolBox(group, nGrid, ghostDepth, 1, forbidTransposition)
    {
    }
    static inline device::memory::host_ptr<MemoryToolBox<NDim>> makeShared(MPICartesianGroup group,
                                                                           device::IdxArray<NDim> nGrid,
                                                                           device::Idx ghostDepth,
                                                                           bool forbidTransposition = false)
    {
      return device::memory::host_ptr<MemoryToolBox<NDim>>(group, nGrid, ghostDepth, 1, forbidTransposition);
    }

    MemoryToolBox(MPICartesianGroup group, device::Idx nGridPoints, device::Idx ghostDepth, device::Idx nThreads,
                  bool forbidTransposition = false)
        : MemoryToolBox(group, makeUniformArray<device::Idx, NDim>(nGridPoints), ghostDepth, nThreads,
                        forbidTransposition)
    {
    }

    MemoryToolBox(MPICartesianGroup group, device::IdxArray<NDim> nGrid, device::Idx ghostDepth, device::Idx nThreads,
                  bool forbidTransposition = false)
        : mGroup(group),
          // mNGridPoints(nGridPoints),
          mNGridPointsVec(nGrid), mGhostDepth(ghostDepth), mFFTLibrary(group, mNGridPointsVec, forbidTransposition),
          mLayouts(mFFTLibrary.getLayout(), mGhostDepth), mFFTNormalization(mLayouts),
          mGhostBuster_toFFTConfig(mLayouts.getConfigSpaceLayout(), mLayouts.getFFTConfigSpaceLayout()),
          mGhostBuster_toConfig(mLayouts.getFFTConfigSpaceLayout(), mLayouts.getConfigSpaceLayout()),
          mGhostUpdater(group, mLayouts.getConfigSpaceLayout())
    {
      checkParallelConsistency();
#ifndef CHECKBOUNDS
      unsetVerbose();
#endif
    }

    void checkParallelConsistency()
    {
      auto dec = mGroup.getDecomposition();
      bool check = true;
      std::stringstream tmp1, tmp2;

      for (size_t i = 0; i < dec.size(); ++i) {
        check = check && (mNGridPointsVec[i] % dec[i] == 0);
        tmp1 << mNGridPointsVec[i] << " ";
        tmp2 << dec[i] << " ";
      }
      if (!check) {

        throw(InconsistentDomainSplitting("Abort: The domain decomposition (" + tmp2.str() +
                                          ") is not consistent with the lattice (" + tmp1.str() +
                                          "). Every lattice dimension should be dividable by the corresponding number "
                                          "of processes specified in the decomposition. Two solutions: Either adjust "
                                          "the number of gridpoints to match the decomposition or vice versa."));
      } else {
        // sayMPI << "Using the following domain decomposition:(" + tmp2.str() +"), with nThreads = " <<
        // mCSpaceItThreaded.size();
        if (amIRoot()) sayMPI << "Using the following domain decomposition:(" + tmp2.str() + ").";
      }
    }

    void setVerbose() { verbosity.setAllOn(); }

    void unsetVerbose() { verbosity.setAllOff(); }

    void resetVerbose() { verbosity.setToDefault(); }

    int getMPIRank() { return mGroup.getBaseComm().rank(); }

    int getNProcesses() const
    {
      int res = 1;
      auto dec = mGroup.getDecomposition();
      for (size_t i = 0; i < dec.size(); ++i) {
        res *= dec[i];
      }
      return res;
    }

    auto getDecomposition() const { return mGroup.getDecomposition(); }

    bool amIRoot() { return mGroup.getBaseComm().rank() == 0; }

    template <typename T> bool initializeFFT()
    {
      if constexpr (std::is_same_v<T, float>)
        mFFTLibrary.getPlans_float();
      else if constexpr (std::is_same_v<T, double>)
        mFFTLibrary.getPlans_double();
      return true;
    }

    device::Idx getNGhosts() const { return mGhostDepth; }

  public:
    MPICartesianGroup mGroup;
    device::IdxArray<NDim> mNGridPointsVec; // Lattice size in each dimension
    device::Idx mGhostDepth;
    FFTLibrarySelector<NDim> mFFTLibrary;
    TripleStateLayouts<NDim> mLayouts;
    FFTNormalization<NDim> mFFTNormalization;
    GhostBuster<NDim> mGhostBuster_toFFTConfig;
    GhostBuster<NDim> mGhostBuster_toConfig;
    GhostUpdater<NDim> mGhostUpdater;

    VerbosityLevels verbosity;
  };

} // namespace TempLat

#include "TempLat/lattice/memory/memorymanager.h"

#endif