Block-Structured AMR Software Framework
Loading...
Searching...
No Matches
AMReX_FFT_R2C.H
Go to the documentation of this file.
1#ifndef AMREX_FFT_R2C_H_
2#define AMREX_FFT_R2C_H_
3#include <AMReX_Config.H>
4
5#include <AMReX_MultiFab.H>
6#include <AMReX_FFT_Helper.H>
7#include <algorithm>
8#include <map>
9#include <numeric>
10#include <tuple>
11
12namespace amrex::FFT
13{
14
20
21template <typename T> class OpenBCSolver;
22template <typename T> class Poisson;
23template <typename T> class PoissonHybrid;
24
46template <typename T = Real, FFT::Direction D = FFT::Direction::both, bool C = false>
47class R2C
48{
49public:
51 using MF = std::conditional_t
52 <C, cMF, std::conditional_t<std::is_same_v<T,Real>,
54
55 template <typename U> friend class OpenBCSolver;
56 template <typename U> friend class Poisson;
57 template <typename U> friend class PoissonHybrid;
58 template <typename U> friend class CrossProj;
59
66 explicit R2C (Box const& domain, Info const& info = Info{});
67
77 explicit R2C (std::array<int,AMREX_SPACEDIM> const& domain_size,
78 Info const& info = Info{});
79
80 ~R2C ();
81
82 R2C (R2C const&) = delete;
83 R2C (R2C &&) = delete;
84 R2C& operator= (R2C const&) = delete;
85 R2C& operator= (R2C &&) = delete;
86
110 void setLocalDomain (std::array<int,AMREX_SPACEDIM> const& local_start,
111 std::array<int,AMREX_SPACEDIM> const& local_size);
112
126 std::pair<std::array<int,AMREX_SPACEDIM>,std::array<int,AMREX_SPACEDIM>>
128
157 void setLocalSpectralDomain (std::array<int,AMREX_SPACEDIM> const& local_start,
158 std::array<int,AMREX_SPACEDIM> const& local_size);
159
178 std::pair<std::array<int,AMREX_SPACEDIM>,std::array<int,AMREX_SPACEDIM>>
180
210 template <typename F, Direction DIR=D>
211 requires (DIR == Direction::both)
212 void forwardThenBackward (MF const& inmf, MF& outmf, F const& post_forward,
213 int incomp = 0, int outcomp = 0)
214 {
216 !m_info.twod_mode,
217 "FFT::R2C::forwardThenBackward(post_forward) currently supports only !twod_mode");
218 BL_PROFILE("FFT::R2C::forwardbackward");
219 this->forward(inmf, incomp);
220 this->post_forward_doit_0(post_forward);
221 this->backward(outmf, outcomp);
222 }
223
234 template <Direction DIR=D>
235 requires ((DIR == Direction::forward) || (DIR == Direction::both))
236 void forward (MF const& inmf, int incomp = 0);
237
249 template <Direction DIR=D>
250 requires ((DIR == Direction::forward) || (DIR == Direction::both))
251 void forward (MF const& inmf, cMF& outmf, int incomp = 0, int outcomp = 0);
252
269 template <typename RT, typename CT, Direction DIR=D, bool CP=C>
270 requires ((DIR == Direction::forward) || (DIR == Direction::both))
271 && (((sizeof(RT)*2 == sizeof(CT)) && (!CP))
272 || ((sizeof(RT) == sizeof(CT)) && CP))
273 void forward (RT const* in, CT* out);
274
284 template <Direction DIR=D>
285 requires (DIR == Direction::both)
286 void backward (MF& outmf, int outcomp = 0);
287
299 template <Direction DIR=D>
300 requires ((DIR == Direction::backward) || (DIR == Direction::both))
301 void backward (cMF const& inmf, MF& outmf, int incomp = 0, int outcomp = 0);
302
319 template <typename CT, typename RT, Direction DIR=D, bool CP=C>
320 requires ((DIR == Direction::backward) || (DIR == Direction::both))
321 && (((sizeof(RT)*2 == sizeof(CT)) && (!CP))
322 || ((sizeof(RT) == sizeof(CT)) && CP))
323 void backward (CT const* in, RT* out);
324
332 [[nodiscard]] T scalingFactor () const;
333
345 template <Direction DIR=D>
346 requires ((DIR == Direction::forward) || (DIR == Direction::both))
347 std::pair<cMF*,IntVect> getSpectralData () const;
348
349
359 [[nodiscard]] std::pair<BoxArray,DistributionMapping> getSpectralDataLayout () const;
360
366 template <typename F>
367 void post_forward_doit_0 (F const& post_forward);
368
369 template <typename F>
375 void post_forward_doit_1 (F const& post_forward);
376
377private:
378
379 void prepare_openbc ();
380
381 void backward_doit (MF& outmf, IntVect const& ngout = IntVect(0),
382 Periodicity const& period = Periodicity::NonPeriodic(),
383 int outcomp = 0);
384
385 void backward_doit (cMF const& inmf, MF& outmf,
386 IntVect const& ngout = IntVect(0),
387 Periodicity const& period = Periodicity::NonPeriodic(),
388 int incomp = 0, int outcomp = 0);
389
390 std::pair<Plan<T>,Plan<T>> make_c2c_plans (cMF& inout, int ndims) const;
391
392 [[nodiscard]] MF& make_m_rx_alias (IndexType const& ix_type);
393
394 [[nodiscard]] Periodicity make_m_rx_alias_periodicity (IndexType const& ix_type,
395 Periodicity const& period) const;
396
397 static Box make_domain_x (Box const& domain)
398 {
399 if constexpr (C) {
400 return Box(IntVect(0), IntVect(AMREX_D_DECL(domain.length(0)-1,
401 domain.length(1)-1,
402 domain.length(2)-1)),
403 domain.ixType());
404 } else {
405 return Box(IntVect(0), IntVect(AMREX_D_DECL(domain.length(0)/2,
406 domain.length(1)-1,
407 domain.length(2)-1)),
408 domain.ixType());
409 }
410 }
411
412 static Box make_domain_y (Box const& domain)
413 {
414 if constexpr (C) {
415 return Box(IntVect(0), IntVect(AMREX_D_DECL(domain.length(1)-1,
416 domain.length(0)-1,
417 domain.length(2)-1)),
418 domain.ixType());
419 } else {
420 return Box(IntVect(0), IntVect(AMREX_D_DECL(domain.length(1)-1,
421 domain.length(0)/2,
422 domain.length(2)-1)),
423 domain.ixType());
424 }
425 }
426
427 static Box make_domain_z (Box const& domain)
428 {
429 if constexpr (C) {
430 return Box(IntVect(0), IntVect(AMREX_D_DECL(domain.length(2)-1,
431 domain.length(0)-1,
432 domain.length(1)-1)),
433 domain.ixType());
434 } else {
435 return Box(IntVect(0), IntVect(AMREX_D_DECL(domain.length(2)-1,
436 domain.length(0)/2,
437 domain.length(1)-1)),
438 domain.ixType());
439 }
440 }
441
442 static std::pair<BoxArray,DistributionMapping>
443 make_layout_from_local_domain (std::array<int,AMREX_SPACEDIM> const& local_start,
444 std::array<int,AMREX_SPACEDIM> const& local_size);
445
446 template <typename FA, typename RT>
447 std::pair<std::unique_ptr<char,DataDeleter>,std::size_t>
448 install_raw_ptr (FA& fa, RT const* p);
449
450 Plan<T> m_fft_fwd_x{};
451 Plan<T> m_fft_bwd_x{};
452 Plan<T> m_fft_fwd_y{};
453 Plan<T> m_fft_bwd_y{};
454 Plan<T> m_fft_fwd_z{};
455 Plan<T> m_fft_bwd_z{};
456 Plan<T> m_fft_fwd_x_half{};
457 Plan<T> m_fft_bwd_x_half{};
458
459 // Comm meta-data. In the forward phase, we start with (x,y,z),
460 // transpose to (y,x,z) and then (z,x,y). In the backward phase, we
461 // perform inverse transpose.
462 std::unique_ptr<MultiBlockCommMetaData> m_cmd_x2y; // (x,y,z) -> (y,x,z)
463 std::unique_ptr<MultiBlockCommMetaData> m_cmd_y2x; // (y,x,z) -> (x,y,z)
464 std::unique_ptr<MultiBlockCommMetaData> m_cmd_y2z; // (y,x,z) -> (z,x,y)
465 std::unique_ptr<MultiBlockCommMetaData> m_cmd_z2y; // (z,x,y) -> (y,x,z)
466 std::unique_ptr<MultiBlockCommMetaData> m_cmd_x2z; // (x,y,z) -> (z,x,y)
467 std::unique_ptr<MultiBlockCommMetaData> m_cmd_z2x; // (z,x,y) -> (x,y,z)
468 std::unique_ptr<MultiBlockCommMetaData> m_cmd_x2z_half; // for openbc
469 std::unique_ptr<MultiBlockCommMetaData> m_cmd_z2x_half; // for openbc
470 Swap01 m_dtos_x2y{};
471 Swap01 m_dtos_y2x{};
472 Swap02 m_dtos_y2z{};
473 Swap02 m_dtos_z2y{};
474 RotateFwd m_dtos_x2z{};
475 RotateBwd m_dtos_z2x{};
476
477 MF m_rx;
478 cMF m_cx;
479 cMF m_cy;
480 cMF m_cz;
481
482 mutable MF m_raw_mf;
483 mutable cMF m_raw_cmf;
484
485 std::map<IntVect,MF> m_rx_aliases;
486
487 std::unique_ptr<char,DataDeleter> m_data_1;
488 std::unique_ptr<char,DataDeleter> m_data_2;
489
490 Box m_real_domain;
491 Box m_spectral_domain_x;
492 Box m_spectral_domain_y;
493 Box m_spectral_domain_z;
494
495 std::unique_ptr<R2C<T,D,C>> m_r2c_sub;
496 detail::SubHelper m_sub_helper;
497
498 Info m_info;
499
500 bool m_do_alld_fft = false;
501 bool m_slab_decomp = false;
502 bool m_openbc_half = false;
503};
504
505template <typename T, Direction D, bool C>
506R2C<T,D,C>::R2C (Box const& domain, Info const& info)
507 : m_real_domain(domain),
508 m_spectral_domain_x(make_domain_x(domain)),
509#if (AMREX_SPACEDIM >= 2)
510 m_spectral_domain_y(make_domain_y(domain)),
511#if (AMREX_SPACEDIM == 3)
512 m_spectral_domain_z(make_domain_z(domain)),
513#endif
514#endif
515 m_sub_helper(domain),
516 m_info(info)
517{
518 BL_PROFILE("FFT::R2C");
519
520 static_assert(std::is_same_v<float,T> || std::is_same_v<double,T>);
521
522 AMREX_ALWAYS_ASSERT(m_real_domain.numPts() > 1);
523#if (AMREX_SPACEDIM == 2)
524 AMREX_ALWAYS_ASSERT(!m_info.twod_mode);
525#else
526 if (m_info.twod_mode) {
527 AMREX_ALWAYS_ASSERT((int(domain.length(0) > 1) +
528 int(domain.length(1) > 1) +
529 int(domain.length(2) > 1)) >= 2);
530 }
531#endif
532
533 {
534 Box subbox = m_sub_helper.make_box(m_real_domain);
535 if (subbox.size() != m_real_domain.size()) {
536 m_r2c_sub = std::make_unique<R2C<T,D,C>>(subbox, m_info);
537 return;
538 }
539 }
540
541 int myproc = ParallelContext::MyProcSub();
542 int nprocs = std::min(ParallelContext::NProcsSub(), m_info.nprocs);
543
544#if (AMREX_SPACEDIM == 3)
545 if (m_info.domain_strategy == DomainStrategy::automatic) {
546 if (m_info.twod_mode) {
547 m_info.domain_strategy = DomainStrategy::slab;
548 } else {
549 int shortside = m_real_domain.shortside();
550 if (shortside < m_info.pencil_threshold*nprocs) {
551 m_info.domain_strategy = DomainStrategy::pencil;
552 } else {
553 m_info.domain_strategy = DomainStrategy::slab;
554 }
555 }
556 }
557
558 if (!m_info.oned_mode) {
559 if (m_info.twod_mode) {
560 m_slab_decomp = true;
561 } else if (m_info.domain_strategy == DomainStrategy::slab && (m_real_domain.length(1) > 1)) {
562 m_slab_decomp = true;
563 }
564 }
565#endif
566
567 auto const ncomp = m_info.batch_size;
568
569 auto bax = amrex::decompose(m_real_domain, nprocs,
570 {AMREX_D_DECL(false,!m_slab_decomp,m_real_domain.length(2)>1)}, true);
571
572 DistributionMapping dmx = detail::make_iota_distromap(bax.size());
573 m_rx.define(bax, dmx, ncomp, 0, MFInfo().SetAlloc(false));
574
575 {
576 BoxList bl = bax.boxList();
577 for (auto & b : bl) {
578 b.shift(-m_real_domain.smallEnd());
579 b.setBig(0, m_spectral_domain_x.bigEnd(0));
580 }
581 BoxArray cbax(std::move(bl));
582 m_cx.define(cbax, dmx, ncomp, 0, MFInfo().SetAlloc(false));
583 }
584
585 m_do_alld_fft = (ParallelDescriptor::NProcs() == 1) &&
586 (! m_info.twod_mode) && (! m_info.oned_mode);
587
588 if (!m_do_alld_fft) // do a series of 1d or 2d ffts
589 {
590 //
591 // make data containers
592 //
593
594#if (AMREX_SPACEDIM >= 2)
596 if ((m_real_domain.length(1) > 1) && !m_slab_decomp && !m_info.oned_mode)
597 {
598 auto cbay = amrex::decompose(m_spectral_domain_y, nprocs,
599 {AMREX_D_DECL(false,true,true)}, true);
600 if (cbay.size() == dmx.size()) {
601 cdmy = dmx;
602 } else {
603 cdmy = detail::make_iota_distromap(cbay.size());
604 }
605 m_cy.define(cbay, cdmy, ncomp, 0, MFInfo().SetAlloc(false));
606 }
607#endif
608
609#if (AMREX_SPACEDIM == 3)
610 if (!m_info.oned_mode && !m_info.twod_mode &&
611 m_real_domain.length(1) > 1 &&
612 m_real_domain.length(2) > 1)
613 {
614 auto cbaz = amrex::decompose(m_spectral_domain_z, nprocs,
615 {false,true,true}, true);
617 if (cbaz.size() == dmx.size()) {
618 cdmz = dmx;
619 } else if (cbaz.size() == cdmy.size()) {
620 cdmz = cdmy;
621 } else {
622 cdmz = detail::make_iota_distromap(cbaz.size());
623 }
624 m_cz.define(cbaz, cdmz, ncomp, 0, MFInfo().SetAlloc(false));
625 }
626#endif
627
628 if constexpr (C) {
629 if (m_slab_decomp) {
630 m_data_1 = detail::make_mfs_share(m_rx, m_cx);
631 m_data_2 = detail::make_mfs_share(m_cz, m_cz);
632 } else {
633 m_data_1 = detail::make_mfs_share(m_rx, m_cz);
634 m_data_2 = detail::make_mfs_share(m_cy, m_cy);
635 // make m_cx an alias to m_rx
636 if (myproc < m_cx.size()) {
637 Box const& box = m_cx.fabbox(myproc);
638 using FAB = typename cMF::FABType::value_type;
639 m_cx.setFab(myproc, FAB(box, ncomp, m_rx[myproc].dataPtr()));
640 }
641 }
642 } else {
643 if (m_slab_decomp) {
644 m_data_1 = detail::make_mfs_share(m_rx, m_cz);
645 m_data_2 = detail::make_mfs_share(m_cx, m_cx);
646 } else {
647 m_data_1 = detail::make_mfs_share(m_rx, m_cy);
648 m_data_2 = detail::make_mfs_share(m_cx, m_cz);
649 }
650 }
651
652 //
653 // make copiers
654 //
655
656#if (AMREX_SPACEDIM >= 2)
657 if (! m_cy.empty()) {
658 // comm meta-data between x and y phases
659 m_cmd_x2y = std::make_unique<MultiBlockCommMetaData>
660 (m_cy, m_spectral_domain_y, m_cx, IntVect(0), m_dtos_x2y);
661 m_cmd_y2x = std::make_unique<MultiBlockCommMetaData>
662 (m_cx, m_spectral_domain_x, m_cy, IntVect(0), m_dtos_y2x);
663 }
664#endif
665#if (AMREX_SPACEDIM == 3)
666 if (! m_cz.empty() ) {
667 if (m_slab_decomp) {
668 // comm meta-data between xy and z phases
669 m_cmd_x2z = std::make_unique<MultiBlockCommMetaData>
670 (m_cz, m_spectral_domain_z, m_cx, IntVect(0), m_dtos_x2z);
671 m_cmd_z2x = std::make_unique<MultiBlockCommMetaData>
672 (m_cx, m_spectral_domain_x, m_cz, IntVect(0), m_dtos_z2x);
673 } else {
674 // comm meta-data between y and z phases
675 m_cmd_y2z = std::make_unique<MultiBlockCommMetaData>
676 (m_cz, m_spectral_domain_z, m_cy, IntVect(0), m_dtos_y2z);
677 m_cmd_z2y = std::make_unique<MultiBlockCommMetaData>
678 (m_cy, m_spectral_domain_y, m_cz, IntVect(0), m_dtos_z2y);
679 }
680 }
681#endif
682
683 //
684 // make plans
685 //
686
687 if (myproc < m_rx.size())
688 {
689 if constexpr (C) {
690 int ndims = m_slab_decomp ? 2 : 1;
691 std::tie(m_fft_fwd_x, m_fft_bwd_x) = make_c2c_plans(m_cx, ndims);
692 } else {
693 Box const& box = m_rx.box(myproc);
694 auto* pr = m_rx[myproc].dataPtr();
695 auto* pc = (typename Plan<T>::VendorComplex *)m_cx[myproc].dataPtr();
696#ifdef AMREX_USE_SYCL
697 m_fft_fwd_x.template init_r2c<Direction::forward>(box, pr, pc, m_slab_decomp, ncomp);
698 m_fft_bwd_x = m_fft_fwd_x;
699#else
700 if constexpr (D == Direction::both || D == Direction::forward) {
701 m_fft_fwd_x.template init_r2c<Direction::forward>(box, pr, pc, m_slab_decomp, ncomp);
702 }
703 if constexpr (D == Direction::both || D == Direction::backward) {
704 m_fft_bwd_x.template init_r2c<Direction::backward>(box, pr, pc, m_slab_decomp, ncomp);
705 }
706#endif
707 }
708 }
709
710#if (AMREX_SPACEDIM >= 2)
711 if (! m_cy.empty()) {
712 std::tie(m_fft_fwd_y, m_fft_bwd_y) = make_c2c_plans(m_cy,1);
713 }
714#endif
715#if (AMREX_SPACEDIM == 3)
716 if (! m_cz.empty()) {
717 std::tie(m_fft_fwd_z, m_fft_bwd_z) = make_c2c_plans(m_cz,1);
718 }
719#endif
720 }
721 else // do fft in all dimensions at the same time
722 {
723 if constexpr (C) {
724 m_data_1 = detail::make_mfs_share(m_rx, m_cx);
725 std::tie(m_fft_fwd_x, m_fft_bwd_x) = make_c2c_plans(m_cx,AMREX_SPACEDIM);
726 } else {
727 m_data_1 = detail::make_mfs_share(m_rx, m_rx);
728 m_data_2 = detail::make_mfs_share(m_cx, m_cx);
729
730 auto const& len = m_real_domain.length();
731 auto* pr = (void*)m_rx[0].dataPtr();
732 auto* pc = (void*)m_cx[0].dataPtr();
733#ifdef AMREX_USE_SYCL
734 m_fft_fwd_x.template init_r2c<Direction::forward>(len, pr, pc, false, ncomp);
735 m_fft_bwd_x = m_fft_fwd_x;
736#else
737 if constexpr (D == Direction::both || D == Direction::forward) {
738 m_fft_fwd_x.template init_r2c<Direction::forward>(len, pr, pc, false, ncomp);
739 }
740 if constexpr (D == Direction::both || D == Direction::backward) {
741 m_fft_bwd_x.template init_r2c<Direction::backward>(len, pr, pc, false, ncomp);
742 }
743#endif
744 }
745 }
746}
747
748template <typename T, Direction D, bool C>
749R2C<T,D,C>::R2C (std::array<int,AMREX_SPACEDIM> const& domain_size, Info const& info)
750 : R2C<T,D,C>(Box(IntVect(0),IntVect(domain_size)-1), info)
751{}
752
753template <typename T, Direction D, bool C>
755{
756 if (m_fft_bwd_x.plan != m_fft_fwd_x.plan) {
757 m_fft_bwd_x.destroy();
758 }
759 if (m_fft_bwd_y.plan != m_fft_fwd_y.plan) {
760 m_fft_bwd_y.destroy();
761 }
762 if (m_fft_bwd_z.plan != m_fft_fwd_z.plan) {
763 m_fft_bwd_z.destroy();
764 }
765 m_fft_fwd_x.destroy();
766 m_fft_fwd_y.destroy();
767 m_fft_fwd_z.destroy();
768 if (m_fft_bwd_x_half.plan != m_fft_fwd_x_half.plan) {
769 m_fft_bwd_x_half.destroy();
770 }
771 m_fft_fwd_x_half.destroy();
772}
773
774template <typename T, Direction D, bool C>
775std::pair<BoxArray,DistributionMapping>
776R2C<T,D,C>::make_layout_from_local_domain (std::array<int,AMREX_SPACEDIM> const& local_start,
777 std::array<int,AMREX_SPACEDIM> const& local_size)
778{
779 IntVect lo(local_start);
780 IntVect len(local_size);
781 Box bx(lo, lo+len-1);
782#ifdef AMREX_USE_MPI
784 MPI_Allgather(&bx, 1, ParallelDescriptor::Mpi_typemap<Box>::type(),
785 allboxes.data(), 1, ParallelDescriptor::Mpi_typemap<Box>::type(),
787 Vector<int> pmap;
788 pmap.reserve(allboxes.size());
789 for (int i = 0; i < allboxes.size(); ++i) {
790 if (allboxes[i].ok()) {
791 pmap.push_back(ParallelContext::local_to_global_rank(i));
792 }
793 }
794 std::erase_if(allboxes, [=] (Box const& b) { return b.isEmpty(); });
795 BoxList bl(std::move(allboxes));
796 return std::make_pair(BoxArray(std::move(bl)), DistributionMapping(std::move(pmap)));
797#else
798 return std::make_pair(BoxArray(bx), DistributionMapping(Vector<int>({0})));
799#endif
800}
801
802template <typename T, Direction D, bool C>
803void R2C<T,D,C>::setLocalDomain (std::array<int,AMREX_SPACEDIM> const& local_start,
804 std::array<int,AMREX_SPACEDIM> const& local_size)
805{
806 auto const& [ba, dm] = make_layout_from_local_domain(local_start, local_size);
807 m_raw_mf = MF(ba, dm, m_info.batch_size, 0, MFInfo().SetAlloc(false));
808}
809
810template <typename T, Direction D, bool C>
811std::pair<std::array<int,AMREX_SPACEDIM>,std::array<int,AMREX_SPACEDIM>>
813{
814 auto const ncomp = m_info.batch_size;
815
816#if (AMREX_SPACEDIM > 1)
817 if (m_r2c_sub) {
818 m_raw_mf = MF(m_sub_helper.inverse_boxarray(m_r2c_sub->m_rx.boxArray()),
819 m_r2c_sub->m_rx.DistributionMap(), ncomp, 0,
820 MFInfo{}.SetAlloc(false));
821 } else
822#endif
823 {
824 m_raw_mf = MF(m_rx.boxArray(), m_rx.DistributionMap(), ncomp, 0,
825 MFInfo{}.SetAlloc(false));
826 }
827
828 auto const myproc = ParallelContext::MyProcSub();
829 if (myproc < m_raw_mf.size()) {
830 Box const& box = m_raw_mf.box(myproc);
831 return std::make_pair(box.smallEnd().toArray(),
832 box.length().toArray());
833 } else {
834 return std::make_pair(std::array<int,AMREX_SPACEDIM>{AMREX_D_DECL(0,0,0)},
835 std::array<int,AMREX_SPACEDIM>{AMREX_D_DECL(0,0,0)});
836 }
837}
838
839template <typename T, Direction D, bool C>
840void R2C<T,D,C>::setLocalSpectralDomain (std::array<int,AMREX_SPACEDIM> const& local_start,
841 std::array<int,AMREX_SPACEDIM> const& local_size)
842{
843 auto const& [ba, dm] = make_layout_from_local_domain(local_start, local_size);
844 m_raw_cmf = cMF(ba, dm, m_info.batch_size, 0, MFInfo().SetAlloc(false));
845}
846
847template <typename T, Direction D, bool C>
848std::pair<std::array<int,AMREX_SPACEDIM>,std::array<int,AMREX_SPACEDIM>>
850{
851 auto const ncomp = m_info.batch_size;
852 auto const& [ba, dm] = getSpectralDataLayout();
853
854 m_raw_cmf = cMF(ba, dm, ncomp, 0, MFInfo{}.SetAlloc(false));
855
856 auto const myproc = ParallelContext::MyProcSub();
857 if (myproc < m_raw_cmf.size()) {
858 Box const& box = m_raw_cmf.box(myproc);
859 return std::make_pair(box.smallEnd().toArray(), box.length().toArray());
860 } else {
861 return std::make_pair(std::array<int,AMREX_SPACEDIM>{AMREX_D_DECL(0,0,0)},
862 std::array<int,AMREX_SPACEDIM>{AMREX_D_DECL(0,0,0)});
863 }
864}
865
866template <typename T, Direction D, bool C>
867void R2C<T,D,C>::prepare_openbc ()
868{
869 if (C || m_r2c_sub) { amrex::Abort("R2C: OpenBC not supported with reduced dimensions or complex inputs"); }
870
871#if (AMREX_SPACEDIM == 3)
872 if (m_do_alld_fft) { return; }
873
874 auto const ncomp = m_info.batch_size;
875
876 if (m_slab_decomp && ! m_fft_fwd_x_half.defined) {
877 auto* fab = detail::get_fab(m_rx);
878 if (fab) {
879 Box bottom_half = m_real_domain;
880 bottom_half.growHi(2,-m_real_domain.length(2)/2);
881 Box box = fab->box() & bottom_half;
882 if (box.ok()) {
883 auto* pr = fab->dataPtr();
884 auto* pc = (typename Plan<T>::VendorComplex *)
885 detail::get_fab(m_cx)->dataPtr();
886#ifdef AMREX_USE_SYCL
887 m_fft_fwd_x_half.template init_r2c<Direction::forward>
888 (box, pr, pc, m_slab_decomp, ncomp);
889 m_fft_bwd_x_half = m_fft_fwd_x_half;
890#else
891 if constexpr (D == Direction::both || D == Direction::forward) {
892 m_fft_fwd_x_half.template init_r2c<Direction::forward>
893 (box, pr, pc, m_slab_decomp, ncomp);
894 }
895 if constexpr (D == Direction::both || D == Direction::backward) {
896 m_fft_bwd_x_half.template init_r2c<Direction::backward>
897 (box, pr, pc, m_slab_decomp, ncomp);
898 }
899#endif
900 }
901 }
902 } // else todo
903
904 if (m_cmd_x2z && ! m_cmd_x2z_half) {
905 Box bottom_half = m_spectral_domain_z;
906 // Note that z-direction's index is 0 because we z is the
907 // unit-stride direction here.
908 bottom_half.growHi(0,-m_spectral_domain_z.length(0)/2);
909 m_cmd_x2z_half = std::make_unique<MultiBlockCommMetaData>
910 (m_cz, bottom_half, m_cx, IntVect(0), m_dtos_x2z);
911 }
912
913 if (m_cmd_z2x && ! m_cmd_z2x_half) {
914 Box bottom_half = m_spectral_domain_x;
915 bottom_half.growHi(2,-m_spectral_domain_x.length(2)/2);
916 m_cmd_z2x_half = std::make_unique<MultiBlockCommMetaData>
917 (m_cx, bottom_half, m_cz, IntVect(0), m_dtos_z2x);
918 }
919#endif
920}
921
922template <typename T, Direction D, bool C>
923typename R2C<T,D,C>::MF&
924R2C<T,D,C>::make_m_rx_alias (IndexType const& ix_type)
925{
926 auto const key = ix_type.ixType();
927 auto it = m_rx_aliases.find(key);
928 if (it != m_rx_aliases.end()) {
929 return it->second;
930 }
931
932 BoxList bl(ix_type);
933 for (auto b : m_rx.boxArray().boxList()) {
934 b.setType(ix_type);
935 bl.push_back(b);
936 }
937
938 MF alias(BoxArray(std::move(bl)), m_rx.DistributionMap(), m_rx.nComp(),
939 m_rx.nGrowVect(), MFInfo{}.SetAlloc(false));
940
941 using FAB = typename MF::fab_type;
942 for (MFIter mfi(alias, MFItInfo().DisableDeviceSync()); mfi.isValid(); ++mfi) {
943 alias.setFab(mfi, FAB(mfi.fabbox(), m_rx.nComp(), m_rx[mfi.index()].dataPtr()));
944 }
945
946 auto [alias_it, inserted] = m_rx_aliases.emplace(key, std::move(alias));
947 amrex::ignore_unused(inserted);
948 return alias_it->second;
949}
950
951template <typename T, Direction D, bool C>
952Periodicity
953R2C<T,D,C>::make_m_rx_alias_periodicity (IndexType const& ix_type,
954 Periodicity const& period) const
955{
956 IntVect per = period.intVect();
957 for (int idim = 0; idim < AMREX_SPACEDIM; ++idim) {
958 if (ix_type.nodeCentered(idim) && per[idim] == 0) {
959 per[idim] = m_real_domain.length(idim);
960 }
961 }
962 return Periodicity(per);
963}
964
965template <typename T, Direction D, bool C>
966template <Direction DIR>
967requires ((DIR == Direction::forward) || (DIR == Direction::both))
968void R2C<T,D,C>::forward (MF const& inmf, int incomp)
969{
970 BL_PROFILE("FFT::R2C::forward(in)");
971
972 auto const ncomp = m_info.batch_size;
973
974 if (m_r2c_sub) {
975 if (m_sub_helper.ghost_safe(inmf.nGrowVect())) {
976 m_r2c_sub->forward(m_sub_helper.make_alias_mf(inmf), incomp);
977 } else {
978 MF tmp(inmf.boxArray(), inmf.DistributionMap(), ncomp, 0);
979 tmp.LocalCopy(inmf, incomp, 0, ncomp, IntVect(0));
980 m_r2c_sub->forward(m_sub_helper.make_alias_mf(tmp),0);
981 }
982 return;
983 }
984
985 if (&m_rx != &inmf) {
986 if (m_rx.boxArray().ixType() == inmf.boxArray().ixType()) {
987 m_rx.ParallelCopy(inmf, incomp, 0, ncomp);
988 } else {
989 MF& rx_alias = make_m_rx_alias(inmf.boxArray().ixType());
990 rx_alias.ParallelCopy(inmf, incomp, 0, ncomp);
991 }
992 }
993
994 if (m_do_alld_fft) {
995 if constexpr (C) {
996 m_fft_fwd_x.template compute_c2c<Direction::forward>();
997 } else {
998 m_fft_fwd_x.template compute_r2c<Direction::forward>();
999 }
1000 return;
1001 }
1002
1003 auto& fft_x = m_openbc_half ? m_fft_fwd_x_half : m_fft_fwd_x;
1004 if constexpr (C) {
1005 fft_x.template compute_c2c<Direction::forward>();
1006 } else {
1007 fft_x.template compute_r2c<Direction::forward>();
1008 }
1009
1010 if ( m_cmd_x2y) {
1011 ParallelCopy(m_cy, m_cx, *m_cmd_x2y, 0, 0, ncomp, m_dtos_x2y);
1012 }
1013 m_fft_fwd_y.template compute_c2c<Direction::forward>();
1014
1015 if ( m_cmd_y2z) {
1016 ParallelCopy(m_cz, m_cy, *m_cmd_y2z, 0, 0, ncomp, m_dtos_y2z);
1017 }
1018#if (AMREX_SPACEDIM == 3)
1019 else if ( m_cmd_x2z) {
1020 if (m_openbc_half) {
1021 NonLocalBC::PackComponents components{};
1022 components.n_components = ncomp;
1024 {components, m_dtos_x2z};
1025 auto handler = ParallelCopy_nowait(m_cz, m_cx, *m_cmd_x2z_half, packing);
1026
1027 Box upper_half = m_spectral_domain_z;
1028 // Note that z-direction's index is 0 because we z is the
1029 // unit-stride direction here.
1030 upper_half.growLo (0,-m_spectral_domain_z.length(0)/2);
1031 m_cz.setVal(0, upper_half, 0, ncomp);
1032
1033 ParallelCopy_finish(m_cz, std::move(handler), *m_cmd_x2z_half, packing);
1034 } else {
1035 ParallelCopy(m_cz, m_cx, *m_cmd_x2z, 0, 0, ncomp, m_dtos_x2z);
1036 }
1037 }
1038#endif
1039 m_fft_fwd_z.template compute_c2c<Direction::forward>();
1040}
1041
1042template <typename T, Direction D, bool C>
1043template <typename FA, typename RT>
1044std::pair<std::unique_ptr<char,DataDeleter>,std::size_t>
1045R2C<T,D,C>::install_raw_ptr (FA& fa, RT const* p)
1046{
1047 AMREX_ALWAYS_ASSERT(!fa.empty());
1048
1049 using FAB = typename FA::FABType::value_type;
1050 using T_FAB = typename FAB::value_type;
1051 static_assert(sizeof(T_FAB) == sizeof(RT));
1052
1053 auto const ncomp = m_info.batch_size;
1054 auto const& ia = fa.IndexArray();
1055
1056 T_FAB* pp = nullptr;
1057 std::size_t sz = 0;
1058
1059 if ( ! ia.empty() ) {
1060 int K = ia[0];
1061 Box const& box = fa.fabbox(K);
1062 if ((alignof(T_FAB) == alignof(RT)) || amrex::is_aligned(p,alignof(T_FAB))) {
1063 pp = (T_FAB*)p;
1064 } else {
1065 sz = sizeof(T_FAB) * box.numPts() * ncomp;
1066 pp = (T_FAB*) The_Arena()->alloc(sz);
1067 }
1068 fa.setFab(K, FAB(box,ncomp,pp));
1069 }
1070
1071 if (sz == 0) {
1072 return std::make_pair(std::unique_ptr<char,DataDeleter>{},std::size_t(0));
1073 } else {
1074 return std::make_pair(std::unique_ptr<char,DataDeleter>
1075 {(char*)pp,DataDeleter{The_Arena()}}, sz);
1076 }
1077}
1078
1079
1080template <typename T, Direction D, bool C>
1081template <typename RT, typename CT, Direction DIR, bool CP>
1082requires ((DIR == Direction::forward) || (DIR == Direction::both))
1083 && (((sizeof(RT)*2 == sizeof(CT)) && (!CP))
1084 || ((sizeof(RT) == sizeof(CT)) && CP))
1085void R2C<T,D,C>::forward (RT const* in, CT* out)
1086{
1087 auto [rdata, rsz] = install_raw_ptr(m_raw_mf, in);
1088 auto [cdata, csz] = install_raw_ptr(m_raw_cmf, out);
1089
1090 if (rsz > 0) {
1091 Gpu::dtod_memcpy_async(rdata.get(),in,rsz);
1093 }
1094
1095 forward(m_raw_mf, m_raw_cmf);
1096
1097 if (csz) {
1098 Gpu::dtod_memcpy_async(out,cdata.get(),csz);
1100 }
1101}
1102
1103template <typename T, Direction D, bool C>
1104template <Direction DIR>
1105requires (DIR == Direction::both)
1106void R2C<T,D,C>::backward (MF& outmf, int outcomp)
1107{
1108 backward_doit(outmf, IntVect(0), Periodicity::NonPeriodic(), outcomp);
1109}
1110
1111template <typename T, Direction D, bool C>
1112void R2C<T,D,C>::backward_doit (MF& outmf, IntVect const& ngout,
1113 Periodicity const& period, int outcomp)
1114{
1115 BL_PROFILE("FFT::R2C::backward(out)");
1116
1117 auto const ncomp = m_info.batch_size;
1118
1119 if (m_r2c_sub) {
1120 if (m_sub_helper.ghost_safe(outmf.nGrowVect())) {
1121 MF submf = m_sub_helper.make_alias_mf(outmf);
1122 IntVect const& subngout = m_sub_helper.make_iv(ngout);
1123 Periodicity const& subperiod = m_sub_helper.make_periodicity(period);
1124 m_r2c_sub->backward_doit(submf, subngout, subperiod, outcomp);
1125 } else {
1126 MF tmp(outmf.boxArray(), outmf.DistributionMap(), ncomp,
1127 m_sub_helper.make_safe_ghost(outmf.nGrowVect()));
1128 this->backward_doit(tmp, ngout, period, 0);
1129 outmf.LocalCopy(tmp, 0, outcomp, ncomp, tmp.nGrowVect());
1130 }
1131 return;
1132 }
1133
1134 if (m_do_alld_fft) {
1135 if constexpr (C) {
1136 m_fft_bwd_x.template compute_c2c<Direction::backward>();
1137 } else {
1138 m_fft_bwd_x.template compute_r2c<Direction::backward>();
1139 }
1140 auto const dst_nghost = amrex::elemwiseMin(ngout,outmf.nGrowVect());
1141 if (m_rx.boxArray().ixType() == outmf.boxArray().ixType()) {
1142 outmf.ParallelCopy(m_rx, 0, outcomp, ncomp, IntVect(0),
1143 dst_nghost, period);
1144 } else {
1145 MF& rx_alias = make_m_rx_alias(outmf.boxArray().ixType());
1146 auto const copy_period = make_m_rx_alias_periodicity(outmf.boxArray().ixType(), period);
1147 outmf.ParallelCopy(rx_alias, 0, outcomp, ncomp, IntVect(0),
1148 dst_nghost, copy_period);
1149 }
1150 return;
1151 }
1152
1153 m_fft_bwd_z.template compute_c2c<Direction::backward>();
1154 if ( m_cmd_z2y) {
1155 ParallelCopy(m_cy, m_cz, *m_cmd_z2y, 0, 0, ncomp, m_dtos_z2y);
1156 }
1157#if (AMREX_SPACEDIM == 3)
1158 else if ( m_cmd_z2x) {
1159 auto const& cmd = m_openbc_half ? m_cmd_z2x_half : m_cmd_z2x;
1160 ParallelCopy(m_cx, m_cz, *cmd, 0, 0, ncomp, m_dtos_z2x);
1161 }
1162#endif
1163
1164 m_fft_bwd_y.template compute_c2c<Direction::backward>();
1165 if ( m_cmd_y2x) {
1166 ParallelCopy(m_cx, m_cy, *m_cmd_y2x, 0, 0, ncomp, m_dtos_y2x);
1167 }
1168
1169 auto& fft_x = m_openbc_half ? m_fft_bwd_x_half : m_fft_bwd_x;
1170 if constexpr (C) {
1171 fft_x.template compute_c2c<Direction::backward>();
1172 } else {
1173 fft_x.template compute_r2c<Direction::backward>();
1174 }
1175 auto const dst_nghost = amrex::elemwiseMin(ngout,outmf.nGrowVect());
1176 if (m_rx.boxArray().ixType() == outmf.boxArray().ixType()) {
1177 outmf.ParallelCopy(m_rx, 0, outcomp, ncomp, IntVect(0),
1178 dst_nghost, period);
1179 } else {
1180 MF& rx_alias = make_m_rx_alias(outmf.boxArray().ixType());
1181 auto const copy_period = make_m_rx_alias_periodicity(outmf.boxArray().ixType(), period);
1182 outmf.ParallelCopy(rx_alias, 0, outcomp, ncomp, IntVect(0),
1183 dst_nghost, copy_period);
1184 }
1185}
1186
1187template <typename T, Direction D, bool C>
1188template <typename CT, typename RT, Direction DIR, bool CP>
1189requires ((DIR == Direction::backward) || (DIR == Direction::both))
1190 && (((sizeof(RT)*2 == sizeof(CT)) && (!CP))
1191 || ((sizeof(RT) == sizeof(CT)) && CP))
1192void R2C<T,D,C>::backward (CT const* in, RT* out)
1193{
1194 auto [rdata, rsz] = install_raw_ptr(m_raw_mf, out);
1195 auto [cdata, csz] = install_raw_ptr(m_raw_cmf, in);
1196
1197 if (csz) {
1198 Gpu::dtod_memcpy_async(cdata.get(),in,csz);
1200 }
1201
1202 backward(m_raw_cmf, m_raw_mf);
1203
1204 if (rsz > 0) {
1205 Gpu::dtod_memcpy_async(out,rdata.get(),rsz);
1207 }
1208}
1209
1210template <typename T, Direction D, bool C>
1211std::pair<Plan<T>, Plan<T>>
1212R2C<T,D,C>::make_c2c_plans (cMF& inout, int ndims) const
1213{
1214 Plan<T> fwd;
1215 Plan<T> bwd;
1216
1217 auto* fab = detail::get_fab(inout);
1218 if (!fab) { return {fwd, bwd};}
1219
1220 Box const& box = fab->box();
1221 auto* pio = (typename Plan<T>::VendorComplex *)fab->dataPtr();
1222
1223 auto const ncomp = m_info.batch_size;
1224
1225#ifdef AMREX_USE_SYCL
1226 fwd.template init_c2c<Direction::forward>(box, pio, ncomp, ndims);
1227 bwd = fwd;
1228#else
1229 if constexpr (D == Direction::both || D == Direction::forward) {
1230 fwd.template init_c2c<Direction::forward>(box, pio, ncomp, ndims);
1231 }
1232 if constexpr (D == Direction::both || D == Direction::backward) {
1233 bwd.template init_c2c<Direction::backward>(box, pio, ncomp, ndims);
1234 }
1235#endif
1236
1237 return {fwd, bwd};
1238}
1239
1240namespace fft_detail {
1241 // The trailing int/long parameter orders these two overloads. A call
1242 // passing a literal 0 prefers the GpuComplex<T>& version, which is
1243 // called once per component of the batch, and falls back to the CellData
1244 // version only when the former is not viable. Without it, a functor
1245 // taking a generic parameter by value or by const reference would make
1246 // the two overloads ambiguous.
1247 template <typename F, typename T>
1249 auto call_post_forward (F const& f, int i, int j, int k, Array4<T> const& a,
1250 int ii, int jj, int kk, int)
1251 noexcept -> decltype(void(f(0,0,0,a(0,0,0))))
1252 {
1253 for (int n = 0; n < a.nComp(); ++n) {
1254 f(i,j,k,a(ii,jj,kk,n));
1255 }
1256 }
1257
1258 template <typename F, typename T>
1260 auto call_post_forward (F const& f, int i, int j, int k, Array4<T> const& a,
1261 int ii, int jj, int kk, long)
1262 noexcept -> decltype(void(f(0,0,0,a.cellData(0,0,0))))
1263 {
1264 f(i,j,k,a.cellData(ii,jj,kk));
1265 }
1266}
1267
1268template <typename T, Direction D, bool C>
1269template <typename F>
1270void R2C<T,D,C>::post_forward_doit_0 (F const& post_forward)
1271{
1272 if (m_info.twod_mode) {
1273 amrex::Abort("xxxxx todo: post_forward");
1274#if (AMREX_SPACEDIM > 1)
1275 } else if (m_r2c_sub) {
1276 // We need to pass the originally ordered indices to post_forward.
1277#if (AMREX_SPACEDIM == 2)
1278 // The original domain is (1,ny). The sub domain is (ny,1).
1279 m_r2c_sub->post_forward_doit_1
1280 ([=] AMREX_GPU_DEVICE (int i, int, int, auto&& sp)
1281 -> decltype(void(post_forward(0,0,0,static_cast<decltype(sp)>(sp))))
1282 {
1283 post_forward(0, i, 0, static_cast<decltype(sp)>(sp));
1284 });
1285#else
1286 if (m_real_domain.length(0) == 1 && m_real_domain.length(1) == 1) {
1287 // Original domain: (1, 1, nz). Sub domain: (nz, 1, 1)
1288 m_r2c_sub->post_forward_doit_1
1289 ([=] AMREX_GPU_DEVICE (int i, int, int, auto&& sp)
1290 -> decltype(void(post_forward(0,0,0,static_cast<decltype(sp)>(sp))))
1291 {
1292 post_forward(0, 0, i, static_cast<decltype(sp)>(sp));
1293 });
1294 } else if (m_real_domain.length(0) == 1 && m_real_domain.length(2) == 1) {
1295 // Original domain: (1, ny, 1). Sub domain: (ny, 1, 1)
1296 m_r2c_sub->post_forward_doit_1
1297 ([=] AMREX_GPU_DEVICE (int i, int, int, auto&& sp)
1298 -> decltype(void(post_forward(0,0,0,static_cast<decltype(sp)>(sp))))
1299 {
1300 post_forward(0, i, 0, static_cast<decltype(sp)>(sp));
1301 });
1302 } else if (m_real_domain.length(0) == 1) {
1303 // Original domain: (1, ny, nz). Sub domain: (ny, nz, 1)
1304 m_r2c_sub->post_forward_doit_1
1305 ([=] AMREX_GPU_DEVICE (int i, int j, int, auto&& sp)
1306 -> decltype(void(post_forward(0,0,0,static_cast<decltype(sp)>(sp))))
1307 {
1308 post_forward(0, i, j, static_cast<decltype(sp)>(sp));
1309 });
1310 } else if (m_real_domain.length(1) == 1) {
1311 // Original domain: (nx, 1, nz). Sub domain: (nx, nz, 1)
1312 m_r2c_sub->post_forward_doit_1
1313 ([=] AMREX_GPU_DEVICE (int i, int j, int, auto&& sp)
1314 -> decltype(void(post_forward(0,0,0,static_cast<decltype(sp)>(sp))))
1315 {
1316 post_forward(i, 0, j, static_cast<decltype(sp)>(sp));
1317 });
1318 } else {
1319 amrex::Abort("R2c::post_forward_doit_0: how did this happen?");
1320 }
1321#endif
1322#endif
1323 } else {
1324 this->post_forward_doit_1(post_forward);
1325 }
1326}
1327
1328template <typename T, Direction D, bool C>
1329template <typename F>
1330void R2C<T,D,C>::post_forward_doit_1 (F const& post_forward)
1331{
1332 if (m_info.twod_mode) {
1333 amrex::Abort("xxxxx todo: post_forward");
1334 } else if (m_r2c_sub) {
1335 amrex::Abort("R2C::post_forward_doit_1: How did this happen?");
1336 } else {
1337 if ( ! m_cz.empty()) {
1338 auto* spectral_fab = detail::get_fab(m_cz);
1339 if (spectral_fab) {
1340 auto const& a = spectral_fab->array(); // m_cz's ordering is z,x,y
1341 ParallelForOMP(spectral_fab->box(),
1342 [=] AMREX_GPU_DEVICE (int iz, int jx, int ky)
1343 {
1344 fft_detail::call_post_forward(post_forward,
1345 jx,ky,iz,a,iz,jx,ky,0);
1346 });
1347 }
1348 } else if ( ! m_cy.empty()) {
1349 auto* spectral_fab = detail::get_fab(m_cy);
1350 if (spectral_fab) {
1351 auto const& a = spectral_fab->array(); // m_cy's ordering is y,x,z
1352 ParallelForOMP(spectral_fab->box(),
1353 [=] AMREX_GPU_DEVICE (int iy, int jx, int k)
1354 {
1355 fft_detail::call_post_forward(post_forward,
1356 jx,iy,k,a,iy,jx,k,0);
1357 });
1358 }
1359 } else {
1360 auto* spectral_fab = detail::get_fab(m_cx);
1361 if (spectral_fab) {
1362 auto const& a = spectral_fab->array();
1363 ParallelForOMP(spectral_fab->box(),
1364 [=] AMREX_GPU_DEVICE (int i, int j, int k)
1365 {
1366 fft_detail::call_post_forward(post_forward,
1367 i,j,k,a,i,j,k,0);
1368 });
1369 }
1370 }
1371 }
1372}
1373
1374template <typename T, Direction D, bool C>
1376{
1377#if (AMREX_SPACEDIM > 1)
1378 // The sub-object knows which axes are actually transformed.
1379 if (m_r2c_sub) { return m_r2c_sub->scalingFactor(); }
1380#endif
1381
1382#if (AMREX_SPACEDIM == 3)
1383 // oned_mode wins over twod_mode, like it does when the plans are made.
1384 if (m_info.oned_mode) {
1385 return T(1)/T(Long(m_real_domain.length(0)));
1386 } else if (m_info.twod_mode) {
1387 return T(1)/T(Long(m_real_domain.length(0)) *
1388 Long(m_real_domain.length(1)));
1389 } else
1390#elif (AMREX_SPACEDIM == 2)
1391 if (m_info.oned_mode) {
1392 return T(1)/T(Long(m_real_domain.length(0)));
1393 } else
1394#endif
1395 {
1396 return T(1)/T(m_real_domain.numPts());
1397 }
1398}
1399
1400template <typename T, Direction D, bool C>
1401template <Direction DIR>
1402requires ((DIR == Direction::forward) || (DIR == Direction::both))
1403std::pair<typename R2C<T,D,C>::cMF *, IntVect>
1405{
1406#if (AMREX_SPACEDIM > 1)
1407 if (m_r2c_sub) {
1408 auto [cmf, order] = m_r2c_sub->getSpectralData();
1409 return std::make_pair(cmf, m_sub_helper.inverse_order(order));
1410 } else
1411#endif
1412 if (!m_cz.empty()) {
1413 return std::make_pair(const_cast<cMF*>(&m_cz), IntVect{AMREX_D_DECL(2,0,1)});
1414 } else if (!m_cy.empty()) {
1415 return std::make_pair(const_cast<cMF*>(&m_cy), IntVect{AMREX_D_DECL(1,0,2)});
1416 } else {
1417 return std::make_pair(const_cast<cMF*>(&m_cx), IntVect{AMREX_D_DECL(0,1,2)});
1418 }
1419}
1420
1421template <typename T, Direction D, bool C>
1422template <Direction DIR>
1423requires ((DIR == Direction::forward) || (DIR == Direction::both))
1424void R2C<T,D,C>::forward (MF const& inmf, cMF& outmf, int incomp, int outcomp)
1425{
1426 BL_PROFILE("FFT::R2C::forward(inout)");
1427
1428 auto const ncomp = m_info.batch_size;
1429
1430 if (m_r2c_sub)
1431 {
1432 bool inmf_safe = m_sub_helper.ghost_safe(inmf.nGrowVect());
1433 MF inmf_sub, inmf_tmp;
1434 int incomp_sub;
1435 if (inmf_safe) {
1436 inmf_sub = m_sub_helper.make_alias_mf(inmf);
1437 incomp_sub = incomp;
1438 } else {
1439 inmf_tmp.define(inmf.boxArray(), inmf.DistributionMap(), ncomp, 0);
1440 inmf_tmp.LocalCopy(inmf, incomp, 0, ncomp, IntVect(0));
1441 inmf_sub = m_sub_helper.make_alias_mf(inmf_tmp);
1442 incomp_sub = 0;
1443 }
1444
1445 bool outmf_safe = m_sub_helper.ghost_safe(outmf.nGrowVect());
1446 cMF outmf_sub, outmf_tmp;
1447 int outcomp_sub;
1448 if (outmf_safe) {
1449 outmf_sub = m_sub_helper.make_alias_mf(outmf);
1450 outcomp_sub = outcomp;
1451 } else {
1452 outmf_tmp.define(outmf.boxArray(), outmf.DistributionMap(), ncomp, 0);
1453 outmf_sub = m_sub_helper.make_alias_mf(outmf_tmp);
1454 outcomp_sub = 0;
1455 }
1456
1457 m_r2c_sub->forward(inmf_sub, outmf_sub, incomp_sub, outcomp_sub);
1458
1459 if (!outmf_safe) {
1460 outmf.LocalCopy(outmf_tmp, 0, outcomp, ncomp, IntVect(0));
1461 }
1462 }
1463 else
1464 {
1465 forward(inmf, incomp);
1466 if (!m_cz.empty()) { // m_cz's order (z,x,y) -> (x,y,z)
1467 RotateBwd dtos{};
1469 (outmf, m_spectral_domain_x, m_cz, IntVect(0), dtos);
1470 ParallelCopy(outmf, m_cz, cmd, 0, outcomp, ncomp, dtos);
1471 } else if (!m_cy.empty()) { // m_cy's order (y,x,z) -> (x,y,z)
1473 (outmf, m_spectral_domain_x, m_cy, IntVect(0), m_dtos_y2x);
1474 ParallelCopy(outmf, m_cy, cmd, 0, outcomp, ncomp, m_dtos_y2x);
1475 } else {
1476 outmf.ParallelCopy(m_cx, 0, outcomp, ncomp);
1477 }
1478 }
1479}
1480
1481template <typename T, Direction D, bool C>
1482template <Direction DIR>
1483requires ((DIR == Direction::backward) || (DIR == Direction::both))
1484void R2C<T,D,C>::backward (cMF const& inmf, MF& outmf, int incomp, int outcomp)
1485{
1486 backward_doit(inmf, outmf, IntVect(0), Periodicity::NonPeriodic(), incomp, outcomp);
1487}
1488
1489template <typename T, Direction D, bool C>
1490void R2C<T,D,C>::backward_doit (cMF const& inmf, MF& outmf, IntVect const& ngout,
1491 Periodicity const& period, int incomp, int outcomp)
1492{
1493 BL_PROFILE("FFT::R2C::backward(inout)");
1494
1495 auto const ncomp = m_info.batch_size;
1496
1497 if (m_r2c_sub)
1498 {
1499 bool inmf_safe = m_sub_helper.ghost_safe(inmf.nGrowVect());
1500 cMF inmf_sub, inmf_tmp;
1501 int incomp_sub;
1502 if (inmf_safe) {
1503 inmf_sub = m_sub_helper.make_alias_mf(inmf);
1504 incomp_sub = incomp;
1505 } else {
1506 inmf_tmp.define(inmf.boxArray(), inmf.DistributionMap(), ncomp, 0);
1507 inmf_tmp.LocalCopy(inmf, incomp, 0, ncomp, IntVect(0));
1508 inmf_sub = m_sub_helper.make_alias_mf(inmf_tmp);
1509 incomp_sub = 0;
1510 }
1511
1512 bool outmf_safe = m_sub_helper.ghost_safe(outmf.nGrowVect());
1513 MF outmf_sub, outmf_tmp;
1514 int outcomp_sub;
1515 if (outmf_safe) {
1516 outmf_sub = m_sub_helper.make_alias_mf(outmf);
1517 outcomp_sub = outcomp;
1518 } else {
1519 IntVect const& ngtmp = m_sub_helper.make_safe_ghost(outmf.nGrowVect());
1520 outmf_tmp.define(outmf.boxArray(), outmf.DistributionMap(), ncomp, ngtmp);
1521 outmf_sub = m_sub_helper.make_alias_mf(outmf_tmp);
1522 outcomp_sub = 0;
1523 }
1524
1525 IntVect const& subngout = m_sub_helper.make_iv(ngout);
1526 Periodicity const& subperiod = m_sub_helper.make_periodicity(period);
1527 m_r2c_sub->backward_doit(inmf_sub, outmf_sub, subngout, subperiod, incomp_sub, outcomp_sub);
1528
1529 if (!outmf_safe) {
1530 outmf.LocalCopy(outmf_tmp, 0, outcomp, ncomp, outmf_tmp.nGrowVect());
1531 }
1532 }
1533 else
1534 {
1535 if (!m_cz.empty()) { // (x,y,z) -> m_cz's order (z,x,y)
1536 RotateFwd dtos{};
1537 MultiBlockCommMetaData cmd
1538 (m_cz, m_spectral_domain_z, inmf, IntVect(0), dtos);
1539 ParallelCopy(m_cz, inmf, cmd, incomp, 0, ncomp, dtos);
1540 } else if (!m_cy.empty()) { // (x,y,z) -> m_cy's ordering (y,x,z)
1541 MultiBlockCommMetaData cmd
1542 (m_cy, m_spectral_domain_y, inmf, IntVect(0), m_dtos_x2y);
1543 ParallelCopy(m_cy, inmf, cmd, incomp, 0, ncomp, m_dtos_x2y);
1544 } else {
1545 m_cx.ParallelCopy(inmf, incomp, 0, ncomp);
1546 }
1547 backward_doit(outmf, ngout, period, outcomp);
1548 }
1549}
1550
1551template <typename T, Direction D, bool C>
1552std::pair<BoxArray,DistributionMapping>
1554{
1555#if (AMREX_SPACEDIM > 1)
1556 if (m_r2c_sub) {
1557 auto const& [ba, dm] = m_r2c_sub->getSpectralDataLayout();
1558 return std::make_pair(m_sub_helper.inverse_boxarray(ba), dm);
1559 }
1560#endif
1561
1562#if (AMREX_SPACEDIM == 3)
1563 if (!m_cz.empty()) {
1564 BoxList bl = m_cz.boxArray().boxList();
1565 for (auto& b : bl) {
1566 auto lo = b.smallEnd();
1567 auto hi = b.bigEnd();
1568 std::swap(lo[0], lo[1]);
1569 std::swap(lo[1], lo[2]);
1570 std::swap(hi[0], hi[1]);
1571 std::swap(hi[1], hi[2]);
1572 b.setSmall(lo);
1573 b.setBig(hi);
1574 }
1575 return std::make_pair(BoxArray(std::move(bl)), m_cz.DistributionMap());
1576 } else
1577#endif
1578#if (AMREX_SPACEDIM >= 2)
1579 if (!m_cy.empty()) {
1580 BoxList bl = m_cy.boxArray().boxList();
1581 for (auto& b : bl) {
1582 auto lo = b.smallEnd();
1583 auto hi = b.bigEnd();
1584 std::swap(lo[0], lo[1]);
1585 std::swap(hi[0], hi[1]);
1586 b.setSmall(lo);
1587 b.setBig(hi);
1588 }
1589 return std::make_pair(BoxArray(std::move(bl)), m_cy.DistributionMap());
1590 } else
1591#endif
1592 {
1593 return std::make_pair(m_cx.boxArray(), m_cx.DistributionMap());
1594 }
1595}
1596
1598template <typename T = Real, FFT::Direction D = FFT::Direction::both>
1600
1601}
1602
1603#endif
#define BL_PROFILE(a)
Definition AMReX_BLProfiler.H:593
#define AMREX_ALWAYS_ASSERT_WITH_MESSAGE(EX, MSG)
Definition AMReX_BLassert.H:49
#define AMREX_ALWAYS_ASSERT(EX)
Definition AMReX_BLassert.H:50
#define AMREX_FORCE_INLINE
Definition AMReX_Extension.H:124
#define AMREX_GPU_DEVICE
Definition AMReX_GpuQualifiers.H:18
while parser_exe_t p
Definition AMReX_Parser_Exe_Body.H:33
#define AMREX_D_DECL(a, b, c)
Definition AMReX_SPACE.H:171
virtual void * alloc(std::size_t sz)=0
Allocate sz bytes from this arena.
Reference-counted collection of Boxes.
Definition AMReX_BoxArray.H:681
void define(const Box &bx)
Initialize the BoxArray from a single box.
Definition AMReX_BoxArray.cpp:352
A list of Boxes sharing a common IndexType.
Definition AMReX_BoxList.H:109
__host__ __device__ Long numPts() const noexcept
Return the number of points contained in the BoxND.
Definition AMReX_Box.H:385
__host__ __device__ IntVectND< dim > length() const noexcept
Return the length of the BoxND.
Definition AMReX_Box.H:167
__host__ __device__ IntVectND< dim > size() const noexcept
Return the length of the BoxND.
Definition AMReX_Box.H:160
__host__ __device__ BoxND & growLo(int idir, int n_cell=1) noexcept
Grow on the low end by n_cell cells in direction idir (negative shrinks).
Definition AMReX_Box.H:690
__host__ __device__ IndexTypeND< dim > ixType() const noexcept
Return the indexing type.
Definition AMReX_Box.H:148
__host__ __device__ BoxND & growHi(int idir, int n_cell=1) noexcept
Grow on the high end by n_cell cells in direction idir (negative shrinks).
Definition AMReX_Box.H:702
__host__ __device__ const IntVectND< dim > & smallEnd() const &noexcept
Return the inclusive lower bound of the box.
Definition AMReX_Box.H:124
Calculates the distribution of FABs to MPI processes.
Definition AMReX_DistributionMapping.H:51
Long size() const noexcept
Length of the underlying processor map.
Definition AMReX_DistributionMapping.H:169
void define(const BoxArray &boxes, int nprocs=ParallelDescriptor::NProcs())
Build a mapping from a BoxArray using the current strategy.
Definition AMReX_DistributionMapping.cpp:347
Convolution-based solver for open boundary conditions using Green's functions.
Definition AMReX_FFT_OpenBCSolver.H:26
3D Poisson solver for periodic, Dirichlet & Neumann boundaries in the first two dimensions,...
Definition AMReX_FFT_Poisson.H:164
Poisson solver for periodic, Dirichlet & Neumann boundaries using FFT.
Definition AMReX_FFT_Poisson.H:32
Parallel Discrete Fourier Transform.
Definition AMReX_FFT_R2C.H:48
R2C & operator=(R2C const &)=delete
void forward(MF const &inmf, int incomp=0)
Forward transform.
Definition AMReX_FFT_R2C.H:968
void forward(RT const *in, CT *out)
Forward transform.
Definition AMReX_FFT_R2C.H:1085
friend class Poisson
Definition AMReX_FFT_R2C.H:56
~R2C()
Definition AMReX_FFT_R2C.H:754
std::conditional_t< C, cMF, std::conditional_t< std::is_same_v< T, Real >, MultiFab, FabArray< BaseFab< T > > > > MF
Definition AMReX_FFT_R2C.H:51
void setLocalDomain(std::array< int, 3 > const &local_start, std::array< int, 3 > const &local_size)
Set local domain.
Definition AMReX_FFT_R2C.H:803
R2C(Box const &domain, Info const &info=Info{})
Constructor.
Definition AMReX_FFT_R2C.H:506
void backward(MF &outmf, int outcomp=0)
Backward transform.
Definition AMReX_FFT_R2C.H:1106
R2C(R2C &&)=delete
void post_forward_doit_1(F const &post_forward)
CUDA-visible helper that redistributes and applies post_forward for the batched layout.
Definition AMReX_FFT_R2C.H:1330
void backward(CT const *in, RT *out)
Backward transform.
Definition AMReX_FFT_R2C.H:1192
T scalingFactor() const
Scaling factor. If the data goes through forward and then backward, the result multiplied by the scal...
Definition AMReX_FFT_R2C.H:1375
std::pair< cMF *, IntVect > getSpectralData() const
void post_forward_doit_0(F const &post_forward)
CUDA-visible hook that walks internal spectral data and applies post_forward.
Definition AMReX_FFT_R2C.H:1270
std::pair< std::array< int, 3 >, std::array< int, 3 > > getLocalDomain() const
Get local domain.
Definition AMReX_FFT_R2C.H:812
friend class PoissonHybrid
Definition AMReX_FFT_R2C.H:57
R2C(R2C const &)=delete
friend class CrossProj
Definition AMReX_FFT_R2C.H:58
std::pair< BoxArray, DistributionMapping > getSpectralDataLayout() const
Get BoxArray and DistributionMapping for spectral data.
Definition AMReX_FFT_R2C.H:1553
void forwardThenBackward(MF const &inmf, MF &outmf, F const &post_forward, int incomp=0, int outcomp=0)
Definition AMReX_FFT_R2C.H:212
R2C(std::array< int, 3 > const &domain_size, Info const &info=Info{})
Constructor.
Definition AMReX_FFT_R2C.H:749
FabArray< BaseFab< GpuComplex< T > > > cMF
Definition AMReX_FFT_R2C.H:50
std::pair< std::array< int, 3 >, std::array< int, 3 > > getLocalSpectralDomain() const
Get local spectral domain.
Definition AMReX_FFT_R2C.H:849
void setLocalSpectralDomain(std::array< int, 3 > const &local_start, std::array< int, 3 > const &local_size)
Set local spectral domain.
Definition AMReX_FFT_R2C.H:840
friend class OpenBCSolver
Definition AMReX_FFT_R2C.H:55
IntVect nGrowVect() const noexcept
Definition AMReX_FabArrayBase.H:85
const DistributionMapping & DistributionMap() const noexcept
Return constant reference to associated DistributionMapping.
Definition AMReX_FabArrayBase.H:135
const BoxArray & boxArray() const noexcept
Return a constant reference to the BoxArray that defines the valid region associated with this FabArr...
Definition AMReX_FabArrayBase.H:100
An Array of FortranArrayBox(FAB)-like Objects.
Definition AMReX_FabArray.H:356
void ParallelCopy(const FabArray< FAB > &src, const Periodicity &period=Periodicity::NonPeriodic(), CpOp op=FabArrayBase::COPY)
Definition AMReX_FabArray.H:971
typename std::conditional_t< IsBaseFab< FAB >::value, BaseFab< GpuComplex< T > >, FABType >::value_type value_type
Definition AMReX_FabArray.H:367
void define(const BoxArray &bxs, const DistributionMapping &dm, int nvar, int ngrow, const MFInfo &info=MFInfo(), const FabFactory< FAB > &factory=DefaultFabFactory< FAB >())
Define this FabArray identically to that performed by the constructor having an analogous function si...
Definition AMReX_FabArray.H:2489
void LocalCopy(FabArray< SFAB > const &src, int scomp, int dcomp, int ncomp, IntVect const &nghost)
Perform local copy of FabArray data.
Definition AMReX_FabArray.H:2265
A collection (stored as an array) of FArrayBox objects.
Definition AMReX_MultiFab.H:40
This provides length of period for periodic domains. 0 means it is not periodic in that direction....
Definition AMReX_Periodicity.H:17
static const Periodicity & NonPeriodic() noexcept
Definition AMReX_Periodicity.cpp:52
This class is a thin wrapper around std::vector. Unlike vector, Vector::operator[] provides bound che...
Definition AMReX_Vector.H:29
amrex_long Long
Definition AMReX_INT.H:30
void ParallelForOMP(T n, L const &f) noexcept
Performance-portable kernel launch function with optional OpenMP threading.
Definition AMReX_GpuLaunch.H:328
ArrayND< T, 4, true > Array4
Convenient alias for a 4D ArrayND with three spatial dimensions plus a component dimension.
Definition AMReX_Array4.H:1340
bool is_aligned(const void *p, std::size_t alignment) noexcept
Definition AMReX_Arena.H:46
Arena * The_Arena()
Definition AMReX_Arena.cpp:839
int NProcs() noexcept
Definition AMReX_ParallelDescriptor.H:255
__host__ __device__ constexpr T elemwiseMin(T const &a, T const &b) noexcept
Definition AMReX_Algorithm.H:75
Definition AMReX_FFT_CrossProj.H:8
R2C< T, D, true > C2C
FFT between complex data.
Definition AMReX_FFT_R2C.H:1599
@ backward
Definition AMReX_FFT_Helper.H:55
@ forward
Definition AMReX_FFT_Helper.H:55
@ both
Definition AMReX_FFT_Helper.H:55
@ automatic
Definition AMReX_FFT_Helper.H:57
@ pencil
Definition AMReX_FFT_Helper.H:57
@ slab
Definition AMReX_FFT_Helper.H:57
void dtod_memcpy_async(void *p_d_dst, const void *p_d_src, const std::size_t sz) noexcept
Definition AMReX_GpuDevice.H:468
void streamSynchronize() noexcept
Definition AMReX_GpuDevice.H:329
MPI_Comm CommunicatorSub() noexcept
sub-communicator for current frame
Definition AMReX_ParallelContext.H:70
int MyProcSub() noexcept
my sub-rank in current frame
Definition AMReX_ParallelContext.H:76
int local_to_global_rank(int rank) noexcept
translate between local rank and global rank
Definition AMReX_ParallelContext.H:98
int NProcsSub() noexcept
number of ranks in current frame
Definition AMReX_ParallelContext.H:74
__host__ __device__ void ignore_unused(const Ts &...)
No-op helper that marks variables as intentionally unused.
Definition AMReX.H:273
@ C
Definition AMReX_Order.H:7
@ F
Definition AMReX_Order.H:7
void ParallelCopy_finish(FabArray< FAB > &dest, CommHandler handler, const FabArrayBase::CommMetaData &cmd, const DataPacking &data_packing)
Definition AMReX_NonLocalBC.H:821
BoxND< 3 > Box
Box is an alias for amrex::BoxND instantiated with AMREX_SPACEDIM.
Definition AMReX_BaseFwd.H:35
void ParallelCopy(MF &dst, MF const &src, int scomp, int dcomp, int ncomp, IntVect const &ng_src=IntVect(0), IntVect const &ng_dst=IntVect(0), Periodicity const &period=Periodicity::NonPeriodic())
dst = src w/ MPI communication
Definition AMReX_FabArrayUtility.H:2256
BoxArray decompose(Box const &domain, int nboxes, Array< bool, 3 > const &decomp, bool no_overlap)
Decompose domain box into BoxArray.
Definition AMReX_BoxArray.cpp:1961
IndexTypeND< 3 > IndexType
IndexType is an alias for amrex::IndexTypeND instantiated with AMREX_SPACEDIM.
Definition AMReX_BaseFwd.H:41
CommHandler ParallelCopy_nowait(NoLocalCopy, FabArray< FAB > &dest, const FabArray< FAB > &src, const FabArrayBase::CommMetaData &cmd, const DataPacking &data_packing)
Definition AMReX_NonLocalBC.H:729
IntVectND< 3 > IntVect
IntVect is an alias for amrex::IntVectND instantiated with AMREX_SPACEDIM.
Definition AMReX_BaseFwd.H:38
void Abort(const std::string &msg)
Print a fatal-error message to stderr and abort execution.
Definition AMReX.cpp:244
Definition AMReX_FFT_Helper.H:83
Definition AMReX_FFT_Helper.H:360
std::conditional_t< std::is_same_v< float, T >, cuComplex, cuDoubleComplex > VendorComplex
Definition AMReX_FFT_Helper.H:363
FabArray memory allocation information.
Definition AMReX_FabArray.H:73
MFInfo & SetAlloc(bool a) noexcept
Control whether FAB storage is allocated when the FabArray is defined.
Definition AMReX_FabArray.H:81
This class specializes behaviour on local copies and unpacking receive buffers.
Definition AMReX_NonLocalBC.H:642
This is the index mapping based on the DTOS MultiBlockDestToSrc.
Definition AMReX_NonLocalBC.H:223
Contains information about which components take part of the data transaction.
Definition AMReX_NonLocalBC.H:550
int n_components
Definition AMReX_NonLocalBC.H:553