AMReX_FFT_R2C.H

Go to the documentation of this file.

#ifndef AMREX_FFT_R2C_H_
#define AMREX_FFT_R2C_H_
#include <AMReX_Config.H>
 
#include <AMReX_MultiFab.H>
#include <AMReX_FFT_Helper.H>
#include <algorithm>
#include <numeric>
#include <tuple>
 
namespace amrex::FFT
{
 
template <typename T> class OpenBCSolver;
 
template <typename T = Real, FFT::Direction D = FFT::Direction::both,
          FFT::DomainStrategy = FFT::DomainStrategy::slab>
          // Don't change the default. Otherwise OpenBCSolver might break.
class R2C
{
public:
    using MF = std::conditional_t<std::is_same_v<T,Real>,
                                  MultiFab, FabArray<BaseFab<T> > >;
    using cMF = FabArray<BaseFab<GpuComplex<T> > >;
 
    template <typename U> friend class OpenBCSolver;
 
    explicit R2C (Box const& domain, Info const& info = Info{});
 
    ~R2C ();
 
    R2C (R2C const&) = delete;
    R2C (R2C &&) = delete;
    R2C& operator= (R2C const&) = delete;
    R2C& operator= (R2C &&) = delete;
 
    template <typename F, Direction DIR=D,
              std::enable_if_t<DIR == Direction::both, int> = 0>
    void forwardThenBackward (MF const& inmf, MF& outmf, F const& post_forward)
    {
        BL_PROFILE("FFT::R2C::forwardbackward");
        this->forward(inmf);
        this->post_forward_doit(post_forward);
        this->backward(outmf);
    }
 
    template <Direction DIR=D, std::enable_if_t<DIR == Direction::forward ||
                                                DIR == Direction::both, int> = 0>
    void forward (MF const& inmf);
 
    template <Direction DIR=D, std::enable_if_t<DIR == Direction::forward ||
                                                DIR == Direction::both, int> = 0>
    void forward (MF const& inmf, cMF& outmf);
 
    template <Direction DIR=D, std::enable_if_t<DIR == Direction::both, int> = 0>
    void backward (MF& outmf);
 
    template <Direction DIR=D, std::enable_if_t<DIR == Direction::backward ||
                                                DIR == Direction::both, int> = 0>
    void backward (cMF const& inmf, MF& outmf);
 
    [[nodiscard]] T scalingFactor () const;
 
    template <Direction DIR=D, std::enable_if_t<DIR == Direction::forward ||
                                                DIR == Direction::both, int> = 0>
    std::pair<cMF*,IntVect> getSpectralData ();
 
    [[nodiscard]] std::pair<BoxArray,DistributionMapping> getSpectralDataLayout () const;
 
    // public for cuda
    template <typename F>
    void post_forward_doit (F const& post_forward);
 
private:
 
    static std::pair<Plan<T>,Plan<T>> make_c2c_plans (cMF& inout);
 
    void backward_doit (MF& outmf, IntVect const& ngout = IntVect(0));
 
    Plan<T> m_fft_fwd_x{};
    Plan<T> m_fft_bwd_x{};
    Plan<T> m_fft_fwd_y{};
    Plan<T> m_fft_bwd_y{};
    Plan<T> m_fft_fwd_z{};
    Plan<T> m_fft_bwd_z{};
 
    // Comm meta-data. In the forward phase, we start with (x,y,z),
    // transpose to (y,x,z) and then (z,x,y). In the backward phase, we
    // perform inverse transpose.
    std::unique_ptr<MultiBlockCommMetaData> m_cmd_x2y; // (x,y,z) -> (y,x,z)
    std::unique_ptr<MultiBlockCommMetaData> m_cmd_y2x; // (y,x,z) -> (x,y,z)
    std::unique_ptr<MultiBlockCommMetaData> m_cmd_y2z; // (y,x,z) -> (z,x,y)
    std::unique_ptr<MultiBlockCommMetaData> m_cmd_z2y; // (z,x,y) -> (y,x,z)
    std::unique_ptr<MultiBlockCommMetaData> m_cmd_x2z; // (x,y,z) -> (z,x,y)
    std::unique_ptr<MultiBlockCommMetaData> m_cmd_z2x; // (z,x,y) -> (x,y,z)
    std::unique_ptr<MultiBlockCommMetaData> m_cmd_x2z_half; // for openbc
    std::unique_ptr<MultiBlockCommMetaData> m_cmd_z2x_half; // for openbc
    Swap01 m_dtos_x2y{};
    Swap01 m_dtos_y2x{};
    Swap02 m_dtos_y2z{};
    Swap02 m_dtos_z2y{};
    RotateFwd m_dtos_x2z{};
    RotateBwd m_dtos_z2x{};
 
    MF  m_rx;
    cMF m_cx;
    cMF m_cy;
    cMF m_cz;
 
    std::unique_ptr<char,DataDeleter> m_data_1;
    std::unique_ptr<char,DataDeleter> m_data_2;
 
    Box m_real_domain;
    Box m_spectral_domain_x;
    Box m_spectral_domain_y;
    Box m_spectral_domain_z;
 
    Info m_info;
 
    bool m_slab_decomp = false;
    bool m_openbc_half = false;
};
 
template <typename T, Direction D, DomainStrategy S>
R2C<T,D,S>::R2C (Box const& domain, Info const& info)
    : m_real_domain(domain),
      m_spectral_domain_x(IntVect(0), IntVect(AMREX_D_DECL(domain.length(0)/2,
                                                           domain.length(1)-1,
                                                           domain.length(2)-1)),
                          domain.ixType()),
#if (AMREX_SPACEDIM >= 2)
      m_spectral_domain_y(IntVect(0), IntVect(AMREX_D_DECL(domain.length(1)-1,
                                                           domain.length(0)/2,
                                                           domain.length(2)-1)),
                          domain.ixType()),
#if (AMREX_SPACEDIM == 3)
      m_spectral_domain_z(IntVect(0), IntVect(AMREX_D_DECL(domain.length(2)-1,
                                                           domain.length(0)/2,
                                                           domain.length(1)-1)),
                          domain.ixType()),
#endif
#endif
      m_info(info)
{
    BL_PROFILE("FFT::R2C");
 
    static_assert(std::is_same_v<float,T> || std::is_same_v<double,T>);
    AMREX_ALWAYS_ASSERT(m_real_domain.length(0) > 1);
#if (AMREX_SPACEDIM == 3)
    AMREX_ALWAYS_ASSERT(m_real_domain.length(2) > 1 || ! m_info.batch_mode);
    AMREX_ALWAYS_ASSERT(m_real_domain.length(1) > 1 || m_real_domain.length(2) == 1);
#else
    AMREX_ALWAYS_ASSERT(! m_info.batch_mode);
#endif
 
    int myproc = ParallelContext::MyProcSub();
    int nprocs = std::min(ParallelContext::NProcsSub(), m_info.nprocs);
 
#if (AMREX_SPACEDIM == 3)
    if (S == DomainStrategy::slab && (m_real_domain.length(1) > 1)) {
        m_slab_decomp = true;
    }
#endif
 
    //
    // make data containers
    //
 
    auto bax = amrex::decompose(m_real_domain, nprocs,
                                {AMREX_D_DECL(false,!m_slab_decomp,true)}, true);
    DistributionMapping dmx = detail::make_iota_distromap(bax.size());
    m_rx.define(bax, dmx, 1, 0, MFInfo().SetAlloc(false));
 
    {
        BoxList bl = bax.boxList();
        for (auto & b : bl) {
            b.shift(-m_real_domain.smallEnd());
            b.setBig(0, m_spectral_domain_x.bigEnd(0));
        }
        BoxArray cbax(std::move(bl));
        m_cx.define(cbax, dmx, 1, 0, MFInfo().SetAlloc(false));
    }
 
#if (AMREX_SPACEDIM >= 2)
    DistributionMapping cdmy;
    if ((m_real_domain.length(1) > 1) && !m_slab_decomp) {
        auto cbay = amrex::decompose(m_spectral_domain_y, nprocs,
                                     {AMREX_D_DECL(false,true,true)}, true);
        if (cbay.size() == dmx.size()) {
            cdmy = dmx;
        } else {
            cdmy = detail::make_iota_distromap(cbay.size());
        }
        m_cy.define(cbay, cdmy, 1, 0, MFInfo().SetAlloc(false));
    }
#endif
 
#if (AMREX_SPACEDIM == 3)
    if (m_real_domain.length(1) > 1 &&
        (! m_info.batch_mode && m_real_domain.length(2) > 1))
    {
        auto cbaz = amrex::decompose(m_spectral_domain_z, nprocs,
                                     {false,true,true}, true);
        DistributionMapping cdmz;
        if (cbaz.size() == dmx.size()) {
            cdmz = dmx;
        } else if (cbaz.size() == cdmy.size()) {
            cdmz = cdmy;
        } else {
            cdmz = detail::make_iota_distromap(cbaz.size());
        }
         m_cz.define(cbaz, cdmz, 1, 0, MFInfo().SetAlloc(false));
    }
#endif
 
    if (m_slab_decomp) {
        m_data_1 = detail::make_mfs_share(m_rx, m_cz);
        m_data_2 = detail::make_mfs_share(m_cx, m_cx);
    } else {
        m_data_1 = detail::make_mfs_share(m_rx, m_cy);
        m_data_2 = detail::make_mfs_share(m_cx, m_cz);
    }
 
    //
    // make copiers
    //
 
#if (AMREX_SPACEDIM >= 2)
    if (! m_cy.empty()) {
        // comm meta-data between x and y phases
        m_cmd_x2y = std::make_unique<MultiBlockCommMetaData>
            (m_cy, m_spectral_domain_y, m_cx, IntVect(0), m_dtos_x2y);
        m_cmd_y2x = std::make_unique<MultiBlockCommMetaData>
            (m_cx, m_spectral_domain_x, m_cy, IntVect(0), m_dtos_y2x);
    }
#endif
#if (AMREX_SPACEDIM == 3)
    if (! m_cz.empty() ) {
        if (m_slab_decomp) {
            // comm meta-data between xy and z phases
            m_cmd_x2z = std::make_unique<MultiBlockCommMetaData>
                (m_cz, m_spectral_domain_z, m_cx, IntVect(0), m_dtos_x2z);
            m_cmd_z2x = std::make_unique<MultiBlockCommMetaData>
                (m_cx, m_spectral_domain_x, m_cz, IntVect(0), m_dtos_z2x);
        } else {
            // comm meta-data between y and z phases
            m_cmd_y2z = std::make_unique<MultiBlockCommMetaData>
                (m_cz, m_spectral_domain_z, m_cy, IntVect(0), m_dtos_y2z);
            m_cmd_z2y = std::make_unique<MultiBlockCommMetaData>
                (m_cy, m_spectral_domain_y, m_cz, IntVect(0), m_dtos_z2y);
        }
    }
#endif
 
    //
    // make plans
    //
 
    if (myproc < m_rx.size())
    {
        Box const& box = m_rx.box(myproc);
        auto* pr = m_rx[myproc].dataPtr();
        auto* pc = (typename Plan<T>::VendorComplex *)m_cx[myproc].dataPtr();
#ifdef AMREX_USE_SYCL
        m_fft_fwd_x.template init_r2c<Direction::forward>(box, pr, pc, m_slab_decomp);
        m_fft_bwd_x = m_fft_fwd_x;
#else
        if constexpr (D == Direction::both || D == Direction::forward) {
            m_fft_fwd_x.template init_r2c<Direction::forward>(box, pr, pc, m_slab_decomp);
        }
        if constexpr (D == Direction::both || D == Direction::backward) {
            m_fft_bwd_x.template init_r2c<Direction::backward>(box, pr, pc, m_slab_decomp);
        }
#endif
    }
 
#if (AMREX_SPACEDIM >= 2)
    if (! m_cy.empty()) {
        std::tie(m_fft_fwd_y, m_fft_bwd_y) = make_c2c_plans(m_cy);
    }
#endif
#if (AMREX_SPACEDIM == 3)
    if (! m_cz.empty()) {
        std::tie(m_fft_fwd_z, m_fft_bwd_z) = make_c2c_plans(m_cz);
    }
#endif
}
 
template <typename T, Direction D, DomainStrategy S>
R2C<T,D,S>::~R2C<T,D,S> ()
{
    if (m_fft_bwd_x.plan != m_fft_fwd_x.plan) {
        m_fft_bwd_x.destroy();
    }
    if (m_fft_bwd_y.plan != m_fft_fwd_y.plan) {
        m_fft_bwd_y.destroy();
    }
    if (m_fft_bwd_z.plan != m_fft_fwd_z.plan) {
        m_fft_bwd_z.destroy();
    }
    m_fft_fwd_x.destroy();
    m_fft_fwd_y.destroy();
    m_fft_fwd_z.destroy();
}
 
template <typename T, Direction D, DomainStrategy S>
template <Direction DIR, std::enable_if_t<DIR == Direction::forward ||
                                          DIR == Direction::both, int> >
void R2C<T,D,S>::forward (MF const& inmf)
{
    BL_PROFILE("FFT::R2C::forward(in)");
 
    if (&m_rx != &inmf) {
        m_rx.ParallelCopy(inmf, 0, 0, 1);
    }
    m_fft_fwd_x.template compute_r2c<Direction::forward>();
 
    if (                          m_cmd_x2y) {
        ParallelCopy(m_cy, m_cx, *m_cmd_x2y, 0, 0, 1, m_dtos_x2y);
    }
    m_fft_fwd_y.template compute_c2c<Direction::forward>();
 
    if (                          m_cmd_y2z) {
        ParallelCopy(m_cz, m_cy, *m_cmd_y2z, 0, 0, 1, m_dtos_y2z);
    }
#if (AMREX_SPACEDIM == 3)
    else if (                     m_cmd_x2z) {
        if (m_openbc_half) {
            Box upper_half = m_spectral_domain_z;
            // Note that z-direction's index is 0 because we z is the unit-stride direction here.
            upper_half.growLo (0,-m_spectral_domain_z.length(0)/2);
            if (! m_cmd_x2z_half) {
                Box bottom_half = m_spectral_domain_z;
                bottom_half.growHi(0,-m_spectral_domain_z.length(0)/2);
                m_cmd_x2z_half = std::make_unique<MultiBlockCommMetaData>
                    (m_cz, bottom_half, m_cx, IntVect(0), m_dtos_x2z);
            }
            NonLocalBC::ApplyDtosAndProjectionOnReciever packing
                {NonLocalBC::PackComponents{}, m_dtos_x2z};
            auto handler = ParallelCopy_nowait(m_cz, m_cx, *m_cmd_x2z_half, packing);
            m_cz.setVal(0, upper_half, 0, 1);
            ParallelCopy_finish(m_cz, std::move(handler), *m_cmd_x2z_half, packing);
        } else {
            ParallelCopy(m_cz, m_cx, *m_cmd_x2z, 0, 0, 1, m_dtos_x2z);
        }
    }
#endif
    m_fft_fwd_z.template compute_c2c<Direction::forward>();
}
 
template <typename T, Direction D, DomainStrategy S>
template <Direction DIR, std::enable_if_t<DIR == Direction::both, int> >
void R2C<T,D,S>::backward (MF& outmf)
{
    backward_doit(outmf);
}
 
template <typename T, Direction D, DomainStrategy S>
void R2C<T,D,S>::backward_doit (MF& outmf, IntVect const& ngout)
{
    BL_PROFILE("FFT::R2C::backward(out)");
 
    m_fft_bwd_z.template compute_c2c<Direction::backward>();
    if (                          m_cmd_z2y) {
        ParallelCopy(m_cy, m_cz, *m_cmd_z2y, 0, 0, 1, m_dtos_z2y);
    }
#if (AMREX_SPACEDIM == 3)
    else if (                     m_cmd_z2x) {
        if (m_openbc_half) {
            Box upper_half = m_spectral_domain_x;
            upper_half.growLo (2,-m_spectral_domain_x.length(2)/2);
            if (! m_cmd_z2x_half) {
                Box bottom_half = m_spectral_domain_x;
                bottom_half.growHi(2,-m_spectral_domain_x.length(2)/2);
                m_cmd_z2x_half = std::make_unique<MultiBlockCommMetaData>
                    (m_cx, bottom_half, m_cz, IntVect(0), m_dtos_z2x);
            }
            NonLocalBC::ApplyDtosAndProjectionOnReciever packing
                {NonLocalBC::PackComponents{}, m_dtos_z2x};
            auto handler = ParallelCopy_nowait(m_cx, m_cz, *m_cmd_z2x_half, packing);
            ParallelCopy_finish(m_cx, std::move(handler), *m_cmd_z2x_half, packing);
        } else {
            ParallelCopy(m_cx, m_cz, *m_cmd_z2x, 0, 0, 1, m_dtos_z2x);
        }
    }
#endif
 
    m_fft_bwd_y.template compute_c2c<Direction::backward>();
    if (                          m_cmd_y2x) {
        ParallelCopy(m_cx, m_cy, *m_cmd_y2x, 0, 0, 1, m_dtos_y2x);
    }
 
    m_fft_bwd_x.template compute_r2c<Direction::backward>();
    outmf.ParallelCopy(m_rx, 0, 0, 1, IntVect(0), ngout);
}
 
template <typename T, Direction D, DomainStrategy S>
std::pair<Plan<T>, Plan<T>>
R2C<T,D,S>::make_c2c_plans (cMF& inout)
{
    Plan<T> fwd;
    Plan<T> bwd;
 
    auto* fab = detail::get_fab(inout);
    if (!fab) { return {fwd, bwd};}
 
    Box const& box = fab->box();
    auto* pio = (typename Plan<T>::VendorComplex *)fab->dataPtr();
 
#ifdef AMREX_USE_SYCL
    fwd.template init_c2c<Direction::forward>(box, pio);
    bwd = fwd;
#else
    if constexpr (D == Direction::both || D == Direction::forward) {
        fwd.template init_c2c<Direction::forward>(box, pio);
    }
    if constexpr (D == Direction::both || D == Direction::backward) {
        bwd.template init_c2c<Direction::backward>(box, pio);
    }
#endif
 
    return {fwd, bwd};
}
 
template <typename T, Direction D, DomainStrategy S>
template <typename F>
void R2C<T,D,S>::post_forward_doit (F const& post_forward)
{
    if (m_info.batch_mode) {
        amrex::Abort("xxxxx todo: post_forward");
    } else {
        if (                           ! m_cz.empty()) {
            auto* spectral_fab = detail::get_fab(m_cz);
            if (spectral_fab) {
                auto const& a = spectral_fab->array(); // m_cz's ordering is z,x,y
                ParallelFor(spectral_fab->box(),
                [=] AMREX_GPU_DEVICE (int iz, int jx, int ky)
                {
                    post_forward(jx,ky,iz,a(iz,jx,ky));
                });
            }
        } else if (                    ! m_cy.empty()) {
            auto* spectral_fab = detail::get_fab(m_cy);
            if (spectral_fab) {
                auto const& a = spectral_fab->array(); // m_cy's ordering is y,x,z
                ParallelFor(spectral_fab->box(),
                [=] AMREX_GPU_DEVICE (int iy, int jx, int k)
                {
                    post_forward(jx,iy,k,a(iy,jx,k));
                });
            }
        } else {
            auto* spectral_fab = detail::get_fab(m_cx);
            if (spectral_fab) {
                auto const& a = spectral_fab->array();
                ParallelFor(spectral_fab->box(),
                [=] AMREX_GPU_DEVICE (int i, int j, int k)
                {
                    post_forward(i,j,k,a(i,j,k));
                });
            }
        }
    }
}
 
template <typename T, Direction D, DomainStrategy S>
T R2C<T,D,S>::scalingFactor () const
{
#if (AMREX_SPACEDIM == 3)
    if (m_info.batch_mode) {
        return T(1)/T(Long(m_real_domain.length(0)) *
                      Long(m_real_domain.length(1)));
    } else
#endif
    {
        return T(1)/T(m_real_domain.numPts());
    }
}
 
template <typename T, Direction D, DomainStrategy S>
template <Direction DIR, std::enable_if_t<DIR == Direction::forward ||
                                          DIR == Direction::both, int> >
std::pair<typename R2C<T,D,S>::cMF *, IntVect>
R2C<T,D,S>::getSpectralData ()
{
    if (!m_cz.empty()) {
        return std::make_pair(&m_cz, IntVect{AMREX_D_DECL(2,0,1)});
    } else if (!m_cy.empty()) {
        return std::make_pair(&m_cy, IntVect{AMREX_D_DECL(1,0,2)});
    } else {
        return std::make_pair(&m_cx, IntVect{AMREX_D_DECL(0,1,2)});
    }
}
 
template <typename T, Direction D, DomainStrategy S>
template <Direction DIR, std::enable_if_t<DIR == Direction::forward ||
                                          DIR == Direction::both, int> >
void R2C<T,D,S>::forward (MF const& inmf, cMF& outmf)
{
    BL_PROFILE("FFT::R2C::forward(inout)");
 
    forward(inmf);
    if (!m_cz.empty()) { // m_cz's order (z,x,y) -> (x,y,z)
        RotateBwd dtos{};
        MultiBlockCommMetaData cmd
            (outmf, m_spectral_domain_x, m_cz, IntVect(0), dtos);
        ParallelCopy(outmf, m_cz, cmd, 0, 0, 1, dtos);
    } else if (!m_cy.empty()) { // m_cy's order (y,x,z) -> (x,y,z)
        MultiBlockCommMetaData cmd
            (outmf, m_spectral_domain_x, m_cy, IntVect(0), m_dtos_y2x);
        ParallelCopy(outmf, m_cy, cmd, 0, 0, 1, m_dtos_y2x);
    } else {
        outmf.ParallelCopy(m_cx, 0, 0, 1);
    }
}
 
template <typename T, Direction D, DomainStrategy S>
template <Direction DIR, std::enable_if_t<DIR == Direction::backward ||
                                          DIR == Direction::both, int> >
void R2C<T,D,S>::backward (cMF const& inmf, MF& outmf)
{
    BL_PROFILE("FFT::R2C::backward(inout)");
 
    if (!m_cz.empty()) { // (x,y,z) -> m_cz's order (z,x,y)
        RotateFwd dtos{};
        MultiBlockCommMetaData cmd
            (m_cz, m_spectral_domain_z, inmf, IntVect(0), dtos);
        ParallelCopy(m_cz, inmf, cmd, 0, 0, 1, dtos);
    } else if (!m_cy.empty()) { // (x,y,z) -> m_cy's ordering (y,x,z)
        MultiBlockCommMetaData cmd
            (m_cy, m_spectral_domain_y, inmf, IntVect(0), m_dtos_x2y);
        ParallelCopy(m_cy, inmf, cmd, 0, 0, 1, m_dtos_x2y);
    } else {
        m_cx.ParallelCopy(inmf, 0, 0, 1);
    }
    backward_doit(outmf);
}
 
template <typename T, Direction D, DomainStrategy S>
std::pair<BoxArray,DistributionMapping>
R2C<T,D,S>::getSpectralDataLayout () const
{
#if (AMREX_SPACEDIM == 3)
    if (!m_cz.empty()) {
        BoxList bl = m_cz.boxArray().boxList();
        for (auto& b : bl) {
            auto lo = b.smallEnd();
            auto hi = b.bigEnd();
            std::swap(lo[0], lo[1]);
            std::swap(lo[1], lo[2]);
            std::swap(hi[0], hi[1]);
            std::swap(hi[1], hi[2]);
            b.setSmall(lo);
            b.setBig(hi);
        }
        return std::make_pair(BoxArray(std::move(bl)), m_cz.DistributionMap());
    } else
#endif
#if (AMREX_SPACEDIM >= 2)
    if (!!m_cy.empty()) {
        BoxList bl = m_cy.boxArray().boxList();
        for (auto& b : bl) {
            auto lo = b.smallEnd();
            auto hi = b.bigEnd();
            std::swap(lo[0], lo[1]);
            std::swap(hi[0], hi[1]);
            b.setSmall(lo);
            b.setBig(hi);
        }
        return std::make_pair(BoxArray(std::move(bl)), m_cy.DistributionMap());
    } else
#endif
    {
        return std::make_pair(m_cx.boxArray(), m_cx.DistributionMap());
    }
}
 
}
 
#endif