Cabana_Grid_FastFourierTransform.hpp

Go to the documentation of this file.

/****************************************************************************
 * Copyright (c) 2018-2023 by the Cabana authors                            *
 * All rights reserved.                                                     *
 *                                                                          *
 * This file is part of the Cabana library. Cabana is distributed under a   *
 * BSD 3-clause license. For the licensing terms see the LICENSE file in    *
 * the top-level directory.                                                 *
 *                                                                          *
 * SPDX-License-Identifier: BSD-3-Clause                                    *
 ****************************************************************************/

#ifndef CABANA_GRID_FASTFOURIERTRANSFORM_HPP
#define CABANA_GRID_FASTFOURIERTRANSFORM_HPP
 
#include <Cabana_Grid_Array.hpp>
#include <Cabana_Grid_Types.hpp>
 
#include <Kokkos_Core.hpp>
#include <Kokkos_Profiling_ScopedRegion.hpp>
 
#include <heffte_fft3d.h>
 
#include <array>
#include <memory>
#include <type_traits>
 
namespace Cabana
{
namespace Grid
{
namespace Experimental
{
//---------------------------------------------------------------------------//

struct FFTScaleFull
{
};

struct FFTScaleNone
{
};

struct FFTScaleSymmetric
{
};

struct FFTBackendFFTW
{
};

struct FFTBackendMKL
{
};
 
namespace Impl
{
struct FFTBackendDefault
{
};
} // namespace Impl

template <class ArrayEntity, class ArrayMesh, class ArrayMemorySpace,
          class ArrayScalar, class Entity, class Mesh, class MemorySpace,
          class Scalar, typename SFINAE = void>
struct is_matching_array : public std::false_type
{
    static_assert( std::is_same<ArrayEntity, Entity>::value,
                   "Array entity type must match FFT entity type." );
    static_assert( std::is_same<ArrayMesh, Mesh>::value,
                   "Array mesh type must match FFT mesh type." );
    static_assert( std::is_same<ArrayMemorySpace, MemorySpace>::value,
                   "Array memory space must match FFT memory space." );
};

template <class ArrayEntity, class ArrayMesh, class ArrayMemorySpace,
          class ArrayScalar, class Entity, class Mesh, class MemorySpace,
          class Scalar>
struct is_matching_array<
    ArrayEntity, ArrayMesh, ArrayMemorySpace, ArrayScalar, Entity, Mesh,
    MemorySpace, Scalar,
    typename std::enable_if<
        std::is_same<ArrayEntity, Entity>::value &&
        std::is_same<ArrayMesh, Mesh>::value &&
        std::is_same<ArrayMemorySpace, MemorySpace>::value>::type>
    : public std::true_type
{
};

enum class FFTCommPattern : unsigned int
{
    alltoallv = 0u,
    p2p = 1u,
    alltoall = 2u,
    p2p_plined = 3u
};
 
//---------------------------------------------------------------------------//
class FastFourierTransformParams
{
    FFTCommPattern FFTcomm = FFTCommPattern::alltoallv;
    bool pencils = true;
    bool reorder = true;
 
  public:
    void setAlltoAll( FFTCommPattern value ) { FFTcomm = value; }
    void setAlltoAll( bool value )
    {
        if ( value )
            FFTcomm = FFTCommPattern::alltoallv;
        else
            FFTcomm = FFTCommPattern::p2p;
    }

    void setPencils( bool value ) { pencils = value; }
    void setReorder( bool value ) { reorder = value; }
    FFTCommPattern getAlltoAll() const { return FFTcomm; }
    bool getPencils() const { return pencils; }
    bool getReorder() const { return reorder; }
};
 
//---------------------------------------------------------------------------//
template <class EntityType, class MeshType, class Scalar, class MemorySpace,
          class Derived>
class FastFourierTransform
{
  public:
    using entity_type = EntityType;
    using mesh_type = MeshType;
    using value_type = Scalar;

    using memory_space = MemorySpace;
    static_assert( Kokkos::is_memory_space<MemorySpace>() );

    using execution_space = typename memory_space::execution_space;

    static constexpr std::size_t num_space_dim = mesh_type::num_space_dim;

    std::array<int, num_space_dim> global_high;
    std::array<int, num_space_dim> global_low;

    FastFourierTransform( const ArrayLayout<EntityType, MeshType>& layout )
    {
        checkArrayDofs( layout.dofsPerEntity() );
 
        // Get the local dimensions of the problem.
        auto entity_space =
            layout.localGrid()->indexSpace( Own(), EntityType(), Local() );
        // Get the global grid.
        const auto& global_grid = layout.localGrid()->globalGrid();
 
        for ( std::size_t d = 0; d < num_space_dim; ++d )
        {
            // Get the low corner of the global index space on this rank.
            global_low[d] =
                (int)global_grid.globalOffset( num_space_dim - d - 1 );
 
            // Get the high corner of the global index space on this rank.
            int local_num_entity =
                (int)entity_space.extent( num_space_dim - d - 1 );
            global_high[d] = global_low[d] + local_num_entity - 1;
        }
    }

    inline void checkArrayDofs( const int dof )
    {
        if ( 2 != dof )
            throw std::logic_error(
                "Cabana::Grid::Experimental::FastFourierTransform::"
                "checkArrayDofs: Only 1 complex value per entity allowed in "
                "FFT" );
    }

    template <class Array_t, class ScaleType>
    void forward(
        const Array_t& x, const ScaleType scaling,
        typename std::enable_if<
            ( is_array<Array_t>::value &&
              is_matching_array<
                  typename Array_t::entity_type, typename Array_t::mesh_type,
                  typename Array_t::memory_space, typename Array_t::value_type,
                  entity_type, mesh_type, memory_space, value_type>::value ),
            int>::type* = 0 )
    {
        Kokkos::Profiling::ScopedRegion region( "Cabana::FFT::forward" );
 
        checkArrayDofs( x.layout()->dofsPerEntity() );
        static_cast<Derived*>( this )->forwardImpl( x, scaling );
    }

    template <class Array_t, class ScaleType>
    void reverse(
        const Array_t& x, const ScaleType scaling,
        typename std::enable_if<
            ( is_array<Array_t>::value &&
              is_matching_array<
                  typename Array_t::entity_type, typename Array_t::mesh_type,
                  typename Array_t::memory_space, typename Array_t::value_type,
                  entity_type, mesh_type, memory_space, value_type>::value ),
            int>::type* = 0 )
    {
        Kokkos::Profiling::ScopedRegion region( "Cabana::FFT::reverse" );
 
        checkArrayDofs( x.layout()->dofsPerEntity() );
        static_cast<Derived*>( this )->reverseImpl( x, scaling );
    }

    template <class ExecutionSpace, class IndexSpaceType, class LViewType,
              class LGViewType, std::size_t NSD = num_space_dim>
    std::enable_if_t<3 == NSD, void>
    copyToLocal( ExecutionSpace exec_space, const IndexSpaceType own_space,
                 LViewType& l_view, const LGViewType lg_view )
    {
        Kokkos::parallel_for(
            "Cabana::Grid::FastFourierTransform::copyTo",
            createExecutionPolicy( own_space, exec_space ),
            KOKKOS_LAMBDA( const int i, const int j, const int k ) {
                auto iw = i - own_space.min( Dim::I );
                auto jw = j - own_space.min( Dim::J );
                auto kw = k - own_space.min( Dim::K );
                l_view( iw, jw, kw, 0 ) = lg_view( i, j, k, 0 );
                l_view( iw, jw, kw, 1 ) = lg_view( i, j, k, 1 );
            } );
    }

    template <class ExecutionSpace, class IndexSpaceType, class LViewType,
              class LGViewType, std::size_t NSD = num_space_dim>
    std::enable_if_t<2 == NSD, void>
    copyToLocal( ExecutionSpace space, const IndexSpaceType own_space,
                 LViewType& l_view, const LGViewType lg_view )
    {
        Kokkos::parallel_for(
            "Cabana::Grid::FastFourierTransform::copyTo",
            createExecutionPolicy( own_space, space ),
            KOKKOS_LAMBDA( const int i, const int j ) {
                auto iw = i - own_space.min( Dim::I );
                auto jw = j - own_space.min( Dim::J );
                l_view( iw, jw, 0 ) = lg_view( i, j, 0 );
                l_view( iw, jw, 1 ) = lg_view( i, j, 1 );
            } );
    }

    template <class ExecutionSpace, class IndexSpaceType, class LViewType,
              class LGViewType, std::size_t NSD = num_space_dim>
    std::enable_if_t<3 == NSD, void>
    copyFromLocal( ExecutionSpace space, const IndexSpaceType own_space,
                   const LViewType l_view, LGViewType& lg_view )
    {
        Kokkos::parallel_for(
            "Cabana::Grid::FastFourierTransform::copyFrom",
            createExecutionPolicy( own_space, space ),
            KOKKOS_LAMBDA( const int i, const int j, const int k ) {
                auto iw = i - own_space.min( Dim::I );
                auto jw = j - own_space.min( Dim::J );
                auto kw = k - own_space.min( Dim::K );
                lg_view( i, j, k, 0 ) = l_view( iw, jw, kw, 0 );
                lg_view( i, j, k, 1 ) = l_view( iw, jw, kw, 1 );
            } );
    }

    template <class ExecutionSpace, class IndexSpaceType, class LViewType,
              class LGViewType, std::size_t NSD = num_space_dim>
    std::enable_if_t<2 == NSD, void>
    copyFromLocal( ExecutionSpace space, const IndexSpaceType own_space,
                   const LViewType l_view, LGViewType& lg_view )
    {
        Kokkos::parallel_for(
            "Cabana::Grid::FastFourierTransform::copyFrom",
            createExecutionPolicy( own_space, space ),
            KOKKOS_LAMBDA( const int i, const int j ) {
                auto iw = i - own_space.min( Dim::I );
                auto jw = j - own_space.min( Dim::J );
                lg_view( i, j, 0 ) = l_view( iw, jw, 0 );
                lg_view( i, j, 1 ) = l_view( iw, jw, 1 );
            } );
    }
};
 
//---------------------------------------------------------------------------//
// heFFTe
//---------------------------------------------------------------------------//
namespace Impl
{
template <class ExecutionSpace, class BackendType>
struct HeffteBackendTraits
{
};
#ifdef Heffte_ENABLE_MKL
template <class ExecutionSpace>
struct HeffteBackendTraits<ExecutionSpace, FFTBackendMKL>
{
    using backend_type = heffte::backend::mkl;
};
#endif
#ifdef Heffte_ENABLE_FFTW
template <class ExecutionSpace>
struct HeffteBackendTraits<ExecutionSpace, FFTBackendFFTW>
{
    using backend_type = heffte::backend::fftw;
};
#endif
#ifdef Heffte_ENABLE_FFTW
template <class ExecutionSpace>
struct HeffteBackendTraits<ExecutionSpace, Impl::FFTBackendDefault>
{
    using backend_type = heffte::backend::fftw;
};
#else
#ifdef Heffte_ENABLE_MKL
template <class ExecutionSpace>
struct HeffteBackendTraits<ExecutionSpace, Impl::FFTBackendDefault>
{
    using backend_type = heffte::backend::mkl;
};
#else
throw std::runtime_error( "Cabana::Grid::FastFourierTransform: Must enable at "
                          "least one heFFTe host backend." );
#endif
#endif
#ifdef Heffte_ENABLE_CUDA
#ifdef KOKKOS_ENABLE_CUDA
template <>
struct HeffteBackendTraits<Kokkos::Cuda, Impl::FFTBackendDefault>
{
    using backend_type = heffte::backend::cufft;
};
#endif
#endif
#ifdef Heffte_ENABLE_ROCM
#ifdef KOKKOS_ENABLE_HIP
template <>
struct HeffteBackendTraits<Kokkos::Experimental::HIP, Impl::FFTBackendDefault>
{
    using backend_type = heffte::backend::rocfft;
};
#endif
#endif
#ifdef Heffte_ENABLE_ONEAPI
#ifdef KOKKOS_ENABLE_SYCL
template <>
struct HeffteBackendTraits<Kokkos::Experimental::SYCL, Impl::FFTBackendDefault>
{
    using backend_type = heffte::backend::onemkl;
};
#endif
#endif
 
template <class ScaleType>
struct HeffteScalingTraits
{
};
template <>
struct HeffteScalingTraits<FFTScaleNone>
{
    static const auto scaling_type = heffte::scale::none;
};
template <>
struct HeffteScalingTraits<FFTScaleFull>
{
    static const auto scaling_type = heffte::scale::full;
};
template <>
struct HeffteScalingTraits<FFTScaleSymmetric>
{
    static const auto scaling_type = heffte::scale::symmetric;
};
 
#ifdef KOKKOS_ENABLE_SYCL
// Overload for SYCL.
template <class ExecSpace, class HeffteBackendType>
auto createHeffteFft3d(
    ExecSpace exec_space, HeffteBackendType, heffte::box3d<> inbox,
    heffte::box3d<> outbox, MPI_Comm comm, heffte::plan_options params,
    typename std::enable_if<
        std::is_same<HeffteBackendType, heffte::backend::onemkl>::value,
        int>::type* = 0 )
{
    // Set FFT options from given parameters
    // heFFTe correctly handles 2D or 3D FFTs within "fft3d"
    sycl::queue& q = exec_space.sycl_queue();
    return std::make_shared<heffte::fft3d<HeffteBackendType>>( q, inbox, outbox,
                                                               comm, params );
}
#endif
 
template <class ExecSpace, class HeffteBackendType>
auto createHeffteFft3d(
    ExecSpace, HeffteBackendType, heffte::box3d<> inbox, heffte::box3d<> outbox,
    MPI_Comm comm, heffte::plan_options params,
    typename std::enable_if<
        std::is_same<HeffteBackendType, heffte::backend::fftw>::value ||
            std::is_same<HeffteBackendType, heffte::backend::mkl>::value ||
            std::is_same<HeffteBackendType, heffte::backend::cufft>::value ||
            std::is_same<HeffteBackendType, heffte::backend::rocfft>::value,
        int>::type* = 0 )
{
    // Set FFT options from given parameters
    // heFFTe correctly handles 2D or 3D FFTs within "fft3d"
    return std::make_shared<heffte::fft3d<HeffteBackendType>>( inbox, outbox,
                                                               comm, params );
}

} // namespace Impl
 
//---------------------------------------------------------------------------//
template <class EntityType, class MeshType, class Scalar, class MemorySpace,
          class ExecSpace, class BackendType>
class HeffteFastFourierTransform
    : public FastFourierTransform<
          EntityType, MeshType, Scalar, MemorySpace,
          HeffteFastFourierTransform<EntityType, MeshType, Scalar, MemorySpace,
                                     ExecSpace, BackendType>>
{
  public:
    using value_type = Scalar;

    using memory_space = MemorySpace;
    static_assert( Kokkos::is_memory_space<MemorySpace>() );

    using execution_space = ExecSpace;
    using backend_type = BackendType;
    using mesh_type = MeshType;

    static constexpr std::size_t num_space_dim = mesh_type::num_space_dim;

    using heffte_backend_type =
        typename Impl::HeffteBackendTraits<execution_space,
                                           backend_type>::backend_type;

    execution_space heffte_execution_space;

    HeffteFastFourierTransform( execution_space exec_space,
                                const ArrayLayout<EntityType, MeshType>& layout,
                                const FastFourierTransformParams& params )
        : FastFourierTransform<
              EntityType, MeshType, Scalar, MemorySpace,
              HeffteFastFourierTransform<EntityType, MeshType, Scalar,
                                         MemorySpace, ExecSpace, BackendType>>(
              layout )
        , heffte_execution_space( exec_space )
    {
        // heFFTe correctly handles 2D or 3D domains within "box3d"
        heffte::box3d<> inbox = { this->global_low, this->global_high };
        heffte::box3d<> outbox = { this->global_low, this->global_high };
 
        heffte::plan_options heffte_params =
            heffte::default_options<heffte_backend_type>();
        auto FFTcomm = params.getAlltoAll();
        switch ( FFTcomm )
        {
        case Cabana::Grid::Experimental::FFTCommPattern::p2p:
            heffte_params.algorithm = heffte::reshape_algorithm::p2p;
            break;
        case Cabana::Grid::Experimental::FFTCommPattern::alltoallv:
            heffte_params.algorithm = heffte::reshape_algorithm::alltoallv;
            break;
        case Cabana::Grid::Experimental::FFTCommPattern::alltoall:
            heffte_params.algorithm = heffte::reshape_algorithm::alltoall;
            break;
        case Cabana::Grid::Experimental::FFTCommPattern::p2p_plined:
            heffte_params.algorithm = heffte::reshape_algorithm::p2p_plined;
            break;
        default:
            heffte_params.algorithm = heffte::reshape_algorithm::alltoallv;
            break;
        }
 
        heffte_params.use_pencils = params.getPencils();
        heffte_params.use_reorder = params.getReorder();
 
        // Create the heFFTe main class (separated to handle SYCL queue
        // correctly).
        _fft = Impl::createHeffteFft3d(
            heffte_execution_space, heffte_backend_type{}, inbox, outbox,
            layout.localGrid()->globalGrid().comm(), heffte_params );
        long long fftsize = std::max( _fft->size_outbox(), _fft->size_inbox() );
 
        // Check the size.
        auto entity_space =
            layout.localGrid()->indexSpace( Own(), EntityType(), Local() );
        if ( fftsize < (int)entity_space.size() )
            throw std::logic_error(
                "Cabana::Grid::Experimental::HeffteFastFourierTransform: "
                "Expected FFT allocation size smaller "
                "than local grid size" );
 
        _fft_work = Kokkos::View<Scalar*, memory_space>(
            Kokkos::ViewAllocateWithoutInitializing( "fft_work" ),
            2 * fftsize );
        _workspace = Kokkos::View<Scalar* [2], memory_space>(
            Kokkos::ViewAllocateWithoutInitializing( "workspace" ),
            _fft->size_workspace() );
    }

    template <class Array_t, class ScaleType>
    void forwardImpl( const Array_t& x, const ScaleType )
    {
        compute( x, 1, Impl::HeffteScalingTraits<ScaleType>().scaling_type );
    }

    template <class Array_t, class ScaleType>
    void reverseImpl( const Array_t& x, const ScaleType )
    {
        compute( x, -1, Impl::HeffteScalingTraits<ScaleType>().scaling_type );
    }

    template <class Array_t>
    void compute( const Array_t& x, const int flag, const heffte::scale scale )
    {
        // Create a subview of the work array to write the local data into.
        auto own_space =
            x.layout()->localGrid()->indexSpace( Own(), EntityType(), Local() );
        auto local_view_space = appendDimension( own_space, 2 );
        auto local_view = createView<Scalar, Kokkos::LayoutRight, memory_space>(
            local_view_space, _fft_work.data() );
 
        // TODO: pull this out to template function
        // Copy to the work array. The work array only contains owned data.
        auto localghost_view = x.view();
 
        this->copyToLocal( heffte_execution_space, own_space, local_view,
                           localghost_view );
 
        if ( flag == 1 )
        {
            _fft->forward(
                reinterpret_cast<std::complex<Scalar>*>( _fft_work.data() ),
                reinterpret_cast<std::complex<Scalar>*>( _fft_work.data() ),
                reinterpret_cast<std::complex<Scalar>*>( _workspace.data() ),
                scale );
        }
        else if ( flag == -1 )
        {
            _fft->backward(
                reinterpret_cast<std::complex<Scalar>*>( _fft_work.data() ),
                reinterpret_cast<std::complex<Scalar>*>( _fft_work.data() ),
                reinterpret_cast<std::complex<Scalar>*>( _workspace.data() ),
                scale );
        }
        else
        {
            throw std::logic_error(
                "Cabana::Grid::Experimental::HeffteFastFourierTransform::"
                "compute: Only 1:forward and -1:backward are allowed as "
                "compute flag" );
        }
 
        // Copy back to output array.
        this->copyFromLocal( heffte_execution_space, own_space, local_view,
                             localghost_view );
    }
 
  private:
    // heFFTe correctly handles 2D or 3D FFTs within "fft3d"
    std::shared_ptr<heffte::fft3d<heffte_backend_type>> _fft;
    Kokkos::View<Scalar*, memory_space> _fft_work;
    Kokkos::View<Scalar* [2], memory_space> _workspace;
};
 
//---------------------------------------------------------------------------//
// heFFTe creation
//---------------------------------------------------------------------------//
template <class Scalar, class MemorySpace, class BackendType, class EntityType,
          class MeshType, class ExecSpace>
auto createHeffteFastFourierTransform(
    ExecSpace exec_space, const ArrayLayout<EntityType, MeshType>& layout,
    const FastFourierTransformParams& params )
{
    return std::make_shared<HeffteFastFourierTransform<
        EntityType, MeshType, Scalar, MemorySpace, ExecSpace, BackendType>>(
        exec_space, layout, params );
}

template <class Scalar, class MemorySpace, class EntityType, class MeshType,
          class ExecSpace>
auto createHeffteFastFourierTransform(
    ExecSpace exec_space, const ArrayLayout<EntityType, MeshType>& layout,
    const FastFourierTransformParams& params )
{
    return createHeffteFastFourierTransform<Scalar, MemorySpace,
                                            Impl::FFTBackendDefault>(
        exec_space, layout, params );
}

template <class Scalar, class MemorySpace, class BackendType, class EntityType,
          class MeshType, class ExecSpace>
auto createHeffteFastFourierTransform(
    ExecSpace exec_space, const ArrayLayout<EntityType, MeshType>& layout )
{
    using heffte_backend_type =
        typename Impl::HeffteBackendTraits<ExecSpace,
                                           BackendType>::backend_type;
 
    // use default heFFTe params for this backend
    const heffte::plan_options heffte_params =
        heffte::default_options<heffte_backend_type>();
    FastFourierTransformParams params;
    params.setAlltoAll( FFTCommPattern::alltoallv );
    params.setPencils( heffte_params.use_pencils );
    params.setReorder( heffte_params.use_reorder );
 
    return std::make_shared<HeffteFastFourierTransform<
        EntityType, MeshType, Scalar, MemorySpace, ExecSpace, BackendType>>(
        exec_space, layout, params );
}

template <class Scalar, class MemorySpace, class EntityType, class MeshType,
          class ExecSpace>
auto createHeffteFastFourierTransform(
    ExecSpace exec_space, const ArrayLayout<EntityType, MeshType>& layout )
{
    return createHeffteFastFourierTransform<
        Scalar, MemorySpace, Impl::FFTBackendDefault, EntityType, MeshType>(
        exec_space, layout );
}

template <class Scalar, class MemorySpace, class BackendType, class EntityType,
          class MeshType>
auto createHeffteFastFourierTransform(
    const ArrayLayout<EntityType, MeshType>& layout,
    const FastFourierTransformParams& params )
{
    using exec_space = typename MemorySpace::execution_space;
    return createHeffteFastFourierTransform<Scalar, MemorySpace, BackendType,
                                            EntityType, MeshType>(
        exec_space{}, layout, params );
}

template <class Scalar, class MemorySpace, class EntityType, class MeshType>
auto createHeffteFastFourierTransform(
    const ArrayLayout<EntityType, MeshType>& layout,
    const FastFourierTransformParams& params )
{
    using exec_space = typename MemorySpace::execution_space;
    return createHeffteFastFourierTransform<
        Scalar, MemorySpace, Impl::FFTBackendDefault, EntityType, MeshType>(
        exec_space{}, layout, params );
}

template <class Scalar, class MemorySpace, class BackendType, class EntityType,
          class MeshType>
auto createHeffteFastFourierTransform(
    const ArrayLayout<EntityType, MeshType>& layout )
{
    using exec_space = typename MemorySpace::execution_space;
    return createHeffteFastFourierTransform<Scalar, MemorySpace, BackendType,
                                            EntityType, MeshType>( exec_space{},
                                                                   layout );
}

template <class Scalar, class MemorySpace, class EntityType, class MeshType>
auto createHeffteFastFourierTransform(
    const ArrayLayout<EntityType, MeshType>& layout )
{
    using exec_space = typename MemorySpace::execution_space;
    return createHeffteFastFourierTransform<
        Scalar, MemorySpace, Impl::FFTBackendDefault, EntityType, MeshType>(
        exec_space{}, layout );
}
 
//---------------------------------------------------------------------------//
 
} // end namespace Experimental
} // namespace Grid
} // namespace Cabana
 
#endif // end CABANA_GRID_FASTFOURIERTRANSFORM_HPP