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