16#ifndef CABANA_GRID_FASTFOURIERTRANSFORM_HPP
17#define CABANA_GRID_FASTFOURIERTRANSFORM_HPP
22#include <Kokkos_Core.hpp>
23#include <Kokkos_Profiling_ScopedRegion.hpp>
25#include <heffte_fft3d.h>
64struct FFTBackendDefault
71template <
class ArrayEntity,
class ArrayMesh,
class ArrayMemorySpace,
72 class ArrayScalar,
class Entity,
class Mesh,
class MemorySpace,
73 class Scalar,
typename SFINAE =
void>
76 static_assert( std::is_same<ArrayEntity, Entity>::value,
77 "Array entity type must match FFT entity type." );
78 static_assert( std::is_same<ArrayMesh, Mesh>::value,
79 "Array mesh type must match FFT mesh type." );
80 static_assert( std::is_same<ArrayMemorySpace, MemorySpace>::value,
81 "Array memory space must match FFT memory space." );
85template <
class ArrayEntity,
class ArrayMesh,
class ArrayMemorySpace,
86 class ArrayScalar,
class Entity,
class Mesh,
class MemorySpace,
89 ArrayEntity, ArrayMesh, ArrayMemorySpace, ArrayScalar, Entity, Mesh,
91 typename std::enable_if<
92 std::is_same<ArrayEntity, Entity>::value &&
93 std::is_same<ArrayMesh, Mesh>::value &&
94 std::is_same<ArrayMemorySpace, MemorySpace>::value>::type>
95 :
public std::true_type
133 FFTcomm = FFTCommPattern::alltoallv;
135 FFTcomm = FFTCommPattern::p2p;
170template <
class EntityType,
class MeshType,
class Scalar,
class MemorySpace,
188 Cabana::Impl::deprecated( Kokkos::is_device<MemorySpace>() ) );
195 using device_type [[deprecated]] =
typename memory_space::device_type;
217 const auto& global_grid = layout.
localGrid()->globalGrid();
226 int local_num_entity =
239 throw std::logic_error(
240 "Cabana::Grid::Experimental::FastFourierTransform::"
241 "checkArrayDofs: Only 1 complex value per entity allowed in "
250 template <
class Array_t,
class ScaleType>
252 const Array_t& x,
const ScaleType scaling,
253 typename std::enable_if<
256 typename Array_t::entity_type,
typename Array_t::mesh_type,
257 typename Array_t::memory_space,
typename Array_t::value_type,
261 Kokkos::Profiling::ScopedRegion region(
"Cabana::FFT::forward" );
264 static_cast<Derived*
>( this )->forwardImpl( x, scaling );
272 template <
class Array_t,
class ScaleType>
274 const Array_t& x,
const ScaleType scaling,
275 typename std::enable_if<
278 typename Array_t::entity_type,
typename Array_t::mesh_type,
279 typename Array_t::memory_space,
typename Array_t::value_type,
283 Kokkos::Profiling::ScopedRegion region(
"Cabana::FFT::reverse" );
286 static_cast<Derived*
>( this )->reverseImpl( x, scaling );
292 template <
class ExecutionSpace,
class IndexSpaceType,
class LViewType,
294 std::enable_if_t<3 == NSD, void>
295 copyToLocal( ExecutionSpace exec_space,
const IndexSpaceType own_space,
296 LViewType& l_view,
const LGViewType lg_view )
298 Kokkos::parallel_for(
299 "Cabana::Grid::FastFourierTransform::copyTo",
301 KOKKOS_LAMBDA(
const int i,
const int j,
const int k ) {
302 auto iw = i - own_space.min( Dim::I );
303 auto jw = j - own_space.min( Dim::J );
304 auto kw = k - own_space.min( Dim::K );
305 l_view( iw, jw, kw, 0 ) = lg_view( i, j, k, 0 );
306 l_view( iw, jw, kw, 1 ) = lg_view( i, j, k, 1 );
313 template <
class ExecutionSpace,
class IndexSpaceType,
class LViewType,
315 std::enable_if_t<2 == NSD, void>
316 copyToLocal( ExecutionSpace space,
const IndexSpaceType own_space,
317 LViewType& l_view,
const LGViewType lg_view )
319 Kokkos::parallel_for(
320 "Cabana::Grid::FastFourierTransform::copyTo",
322 KOKKOS_LAMBDA(
const int i,
const int j ) {
323 auto iw = i - own_space.min( Dim::I );
324 auto jw = j - own_space.min( Dim::J );
325 l_view( iw, jw, 0 ) = lg_view( i, j, 0 );
326 l_view( iw, jw, 1 ) = lg_view( i, j, 1 );
333 template <
class ExecutionSpace,
class IndexSpaceType,
class LViewType,
335 std::enable_if_t<3 == NSD, void>
337 const LViewType l_view, LGViewType& lg_view )
339 Kokkos::parallel_for(
340 "Cabana::Grid::FastFourierTransform::copyFrom",
342 KOKKOS_LAMBDA(
const int i,
const int j,
const int k ) {
343 auto iw = i - own_space.min( Dim::I );
344 auto jw = j - own_space.min( Dim::J );
345 auto kw = k - own_space.min( Dim::K );
346 lg_view( i, j, k, 0 ) = l_view( iw, jw, kw, 0 );
347 lg_view( i, j, k, 1 ) = l_view( iw, jw, kw, 1 );
354 template <
class ExecutionSpace,
class IndexSpaceType,
class LViewType,
356 std::enable_if_t<2 == NSD, void>
358 const LViewType l_view, LGViewType& lg_view )
360 Kokkos::parallel_for(
361 "Cabana::Grid::FastFourierTransform::copyFrom",
363 KOKKOS_LAMBDA(
const int i,
const int j ) {
364 auto iw = i - own_space.min( Dim::I );
365 auto jw = j - own_space.min( Dim::J );
366 lg_view( i, j, 0 ) = l_view( iw, jw, 0 );
367 lg_view( i, j, 1 ) = l_view( iw, jw, 1 );
378template <
class ExecutionSpace,
class BackendType>
379struct HeffteBackendTraits
382#ifdef Heffte_ENABLE_MKL
383template <
class ExecutionSpace>
384struct HeffteBackendTraits<ExecutionSpace, FFTBackendMKL>
386 using backend_type = heffte::backend::mkl;
389#ifdef Heffte_ENABLE_FFTW
390template <
class ExecutionSpace>
391struct HeffteBackendTraits<ExecutionSpace, FFTBackendFFTW>
393 using backend_type = heffte::backend::fftw;
396#ifdef Heffte_ENABLE_FFTW
397template <
class ExecutionSpace>
398struct HeffteBackendTraits<ExecutionSpace, Impl::FFTBackendDefault>
400 using backend_type = heffte::backend::fftw;
403#ifdef Heffte_ENABLE_MKL
404template <
class ExecutionSpace>
405struct HeffteBackendTraits<ExecutionSpace, Impl::FFTBackendDefault>
407 using backend_type = heffte::backend::mkl;
410throw std::runtime_error(
"Cabana::Grid::FastFourierTransform: Must enable at "
411 "least one heFFTe host backend." );
414#ifdef Heffte_ENABLE_CUDA
415#ifdef KOKKOS_ENABLE_CUDA
417struct HeffteBackendTraits<Kokkos::Cuda, Impl::FFTBackendDefault>
419 using backend_type = heffte::backend::cufft;
423#ifdef Heffte_ENABLE_ROCM
424#ifdef KOKKOS_ENABLE_HIP
426struct HeffteBackendTraits<Kokkos::Experimental::HIP, Impl::FFTBackendDefault>
428 using backend_type = heffte::backend::rocfft;
432#ifdef Heffte_ENABLE_ONEAPI
433#ifdef KOKKOS_ENABLE_SYCL
435struct HeffteBackendTraits<Kokkos::Experimental::SYCL, Impl::FFTBackendDefault>
437 using backend_type = heffte::backend::onemkl;
442template <
class ScaleType>
443struct HeffteScalingTraits
447struct HeffteScalingTraits<FFTScaleNone>
449 static const auto scaling_type = heffte::scale::none;
452struct HeffteScalingTraits<FFTScaleFull>
454 static const auto scaling_type = heffte::scale::full;
457struct HeffteScalingTraits<FFTScaleSymmetric>
459 static const auto scaling_type = heffte::scale::symmetric;
462#ifdef KOKKOS_ENABLE_SYCL
464template <
class ExecSpace,
class HeffteBackendType>
465auto createHeffteFft3d(
466 ExecSpace exec_space, HeffteBackendType, heffte::box3d<> inbox,
467 heffte::box3d<> outbox, MPI_Comm comm, heffte::plan_options params,
468 typename std::enable_if<
469 std::is_same<HeffteBackendType, heffte::backend::onemkl>::value,
474 sycl::queue& q = exec_space.sycl_queue();
475 return std::make_shared<heffte::fft3d<HeffteBackendType>>( q, inbox, outbox,
480template <
class ExecSpace,
class HeffteBackendType>
481auto createHeffteFft3d(
482 ExecSpace, HeffteBackendType, heffte::box3d<> inbox, heffte::box3d<> outbox,
483 MPI_Comm comm, heffte::plan_options params,
484 typename std::enable_if<
485 std::is_same<HeffteBackendType, heffte::backend::fftw>::value ||
486 std::is_same<HeffteBackendType, heffte::backend::mkl>::value ||
487 std::is_same<HeffteBackendType, heffte::backend::cufft>::value ||
488 std::is_same<HeffteBackendType, heffte::backend::rocfft>::value,
493 return std::make_shared<heffte::fft3d<HeffteBackendType>>( inbox, outbox,
504template <
class EntityType,
class MeshType,
class Scalar,
class MemorySpace,
505 class ExecSpace,
class BackendType>
508 EntityType, MeshType, Scalar, MemorySpace,
509 HeffteFastFourierTransform<EntityType, MeshType, Scalar, MemorySpace,
510 ExecSpace, BackendType>>
522 Cabana::Impl::deprecated( Kokkos::is_device<MemorySpace>() ) );
529 using device_type [[deprecated]] =
typename memory_space::device_type;
556 EntityType, MeshType, Scalar, MemorySpace,
558 MemorySpace, ExecSpace, BackendType>>(
563 heffte::box3d<> inbox = { this->global_low, this->global_high };
564 heffte::box3d<> outbox = { this->global_low, this->global_high };
566 heffte::plan_options heffte_params =
567 heffte::default_options<heffte_backend_type>();
571 case Cabana::Grid::Experimental::FFTCommPattern::p2p:
572 heffte_params.algorithm = heffte::reshape_algorithm::p2p;
574 case Cabana::Grid::Experimental::FFTCommPattern::alltoallv:
575 heffte_params.algorithm = heffte::reshape_algorithm::alltoallv;
577 case Cabana::Grid::Experimental::FFTCommPattern::alltoall:
578 heffte_params.algorithm = heffte::reshape_algorithm::alltoall;
580 case Cabana::Grid::Experimental::FFTCommPattern::p2p_plined:
581 heffte_params.algorithm = heffte::reshape_algorithm::p2p_plined;
584 heffte_params.algorithm = heffte::reshape_algorithm::alltoallv;
588 heffte_params.use_pencils = params.
getPencils();
589 heffte_params.use_reorder = params.
getReorder();
593 _fft = Impl::createHeffteFft3d(
595 layout.
localGrid()->globalGrid().comm(), heffte_params );
596 long long fftsize = std::max( _fft->size_outbox(), _fft->size_inbox() );
601 if ( fftsize < (
int)entity_space.size() )
602 throw std::logic_error(
603 "Cabana::Grid::Experimental::HeffteFastFourierTransform: "
604 "Expected FFT allocation size smaller "
605 "than local grid size" );
607 _fft_work = Kokkos::View<Scalar*, memory_space>(
608 Kokkos::ViewAllocateWithoutInitializing(
"fft_work" ),
610 _workspace = Kokkos::View<Scalar* [2], memory_space>(
611 Kokkos::ViewAllocateWithoutInitializing(
"workspace" ),
612 _fft->size_workspace() );
620 template <
class Array_t,
class ScaleType>
623 compute( x, 1, Impl::HeffteScalingTraits<ScaleType>().scaling_type );
631 template <
class Array_t,
class ScaleType>
634 compute( x, -1, Impl::HeffteScalingTraits<ScaleType>().scaling_type );
643 template <
class Array_t>
644 void compute(
const Array_t& x,
const int flag,
const heffte::scale scale )
648 x.layout()->localGrid()->indexSpace(
Own(), EntityType(),
Local() );
651 local_view_space, _fft_work.data() );
655 auto localghost_view = x.view();
657 this->
copyToLocal( heffte_execution_space, own_space, local_view,
663 reinterpret_cast<std::complex<Scalar>*
>( _fft_work.data() ),
664 reinterpret_cast<std::complex<Scalar>*
>( _fft_work.data() ),
665 reinterpret_cast<std::complex<Scalar>*
>( _workspace.data() ),
668 else if ( flag == -1 )
671 reinterpret_cast<std::complex<Scalar>*
>( _fft_work.data() ),
672 reinterpret_cast<std::complex<Scalar>*
>( _fft_work.data() ),
673 reinterpret_cast<std::complex<Scalar>*
>( _workspace.data() ),
678 throw std::logic_error(
679 "Cabana::Grid::Experimental::HeffteFastFourierTransform::"
680 "compute: Only 1:forward and -1:backward are allowed as "
685 this->
copyFromLocal( heffte_execution_space, own_space, local_view,
691 std::shared_ptr<heffte::fft3d<heffte_backend_type>> _fft;
692 Kokkos::View<Scalar*, memory_space> _fft_work;
693 Kokkos::View<Scalar* [2], memory_space> _workspace;
703template <
class Scalar,
class MemorySpace,
class BackendType,
class EntityType,
704 class MeshType,
class ExecSpace>
710 EntityType, MeshType, Scalar, MemorySpace, ExecSpace, BackendType>>(
711 exec_space, layout, params );
718template <
class Scalar,
class MemorySpace,
class EntityType,
class MeshType,
725 Impl::FFTBackendDefault>(
726 exec_space, layout, params );
733template <
class Scalar,
class MemorySpace,
class BackendType,
class EntityType,
734 class MeshType,
class ExecSpace>
738 using heffte_backend_type =
739 typename Impl::HeffteBackendTraits<ExecSpace,
740 BackendType>::backend_type;
743 const heffte::plan_options heffte_params =
744 heffte::default_options<heffte_backend_type>();
747 params.
setPencils( heffte_params.use_pencils );
748 params.
setReorder( heffte_params.use_reorder );
751 EntityType, MeshType, Scalar, MemorySpace, ExecSpace, BackendType>>(
752 exec_space, layout, params );
759template <
class Scalar,
class MemorySpace,
class EntityType,
class MeshType,
765 Scalar, MemorySpace, Impl::FFTBackendDefault, EntityType, MeshType>(
766 exec_space, layout );
772template <
class Scalar,
class MemorySpace,
class BackendType,
class EntityType,
778 using exec_space =
typename MemorySpace::execution_space;
780 EntityType, MeshType>(
781 exec_space{}, layout, params );
787template <
class Scalar,
class MemorySpace,
class EntityType,
class MeshType>
792 using exec_space =
typename MemorySpace::execution_space;
794 Scalar, MemorySpace, Impl::FFTBackendDefault, EntityType, MeshType>(
795 exec_space{}, layout, params );
801template <
class Scalar,
class MemorySpace,
class BackendType,
class EntityType,
806 using exec_space =
typename MemorySpace::execution_space;
808 EntityType, MeshType>( exec_space{},
815template <
class Scalar,
class MemorySpace,
class EntityType,
class MeshType>
819 using exec_space =
typename MemorySpace::execution_space;
821 Scalar, MemorySpace, Impl::FFTBackendDefault, EntityType, MeshType>(
822 exec_space{}, layout );
Kokkos::RangePolicy< ExecutionSpace > createExecutionPolicy(const IndexSpace< 1 > &index_space, const ExecutionSpace &)
Create a multi-dimensional execution policy over an index space.
Definition Cabana_Grid_IndexSpace.hpp:175
Kokkos::View< Scalar *, Params... > createView(const std::string &label, const IndexSpace< 1 > &index_space)
Given an index space create a view over the extent of that index space.
Definition Cabana_Grid_IndexSpace.hpp:235
IndexSpace< N+1 > appendDimension(const IndexSpace< N > &index_space, const long size)
Given an N-dimensional index space append an additional dimension with the given size.
Definition Cabana_Grid_IndexSpace.hpp:442
Entity layout for array data on the local mesh.
Definition Cabana_Grid_Array.hpp:46
int dofsPerEntity() const
Get the number of degrees-of-freedom on each grid entity.
Definition Cabana_Grid_Array.hpp:77
const std::shared_ptr< LocalGrid< MeshType > > localGrid() const
Get the local grid over which this layout is defined.
Definition Cabana_Grid_Array.hpp:71
Core: particle data structures and algorithms.
Definition Cabana_AoSoA.hpp:36
Tag specifying FFTW backend for FFT (host default).
Definition Cabana_Grid_FastFourierTransform.hpp:54
Tag specifying MKL backend for FFT.
Definition Cabana_Grid_FastFourierTransform.hpp:58
Tag for full scaling of FFT.
Definition Cabana_Grid_FastFourierTransform.hpp:41
Tag for no scaling of FFT.
Definition Cabana_Grid_FastFourierTransform.hpp:45
Tag for symmetric scaling of FFT.
Definition Cabana_Grid_FastFourierTransform.hpp:49
Matching Array static type checker.
Definition Cabana_Grid_FastFourierTransform.hpp:75
Local index tag.
Definition Cabana_Grid_Types.hpp:208
Owned decomposition tag.
Definition Cabana_Grid_Types.hpp:190
Definition Cabana_Grid_Array.hpp:324