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<detail::TransposeMetaData> m_cmd_x2y; // (x,y,z) -> (y,x,z)
463 std::unique_ptr<detail::TransposeMetaData> m_cmd_y2x; // (y,x,z) -> (x,y,z)
464 std::unique_ptr<detail::TransposeMetaData> m_cmd_y2z; // (y,x,z) -> (z,x,y)
465 std::unique_ptr<detail::TransposeMetaData> m_cmd_z2y; // (z,x,y) -> (y,x,z)
466 std::unique_ptr<detail::TransposeMetaData> m_cmd_x2z; // (x,y,z) -> (z,x,y)
467 std::unique_ptr<detail::TransposeMetaData> m_cmd_z2x; // (z,x,y) -> (x,y,z)
468 std::unique_ptr<detail::TransposeMetaData> m_cmd_x2z_half; // for openbc
469 std::unique_ptr<detail::TransposeMetaData> 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
500 int m_nprocs_active = 0;
501 MPI_Comm m_comm_active = MPI_COMM_NULL;
502 bool m_owns_comm_active = false; // true if it must be freed
503
504 Info m_info;
505
506 bool m_do_alld_fft = false;
507 bool m_slab_decomp = false;
508 bool m_openbc_half = false;
509};
510
511template <typename T, Direction D, bool C>
512R2C<T,D,C>::R2C (Box const& domain, Info const& info)
513 : m_real_domain(domain),
514 m_spectral_domain_x(make_domain_x(domain)),
515#if (AMREX_SPACEDIM >= 2)
516 m_spectral_domain_y(make_domain_y(domain)),
517#if (AMREX_SPACEDIM == 3)
518 m_spectral_domain_z(make_domain_z(domain)),
519#endif
520#endif
521 m_sub_helper(domain),
522 m_info(info)
523{
524 BL_PROFILE("FFT::R2C");
525
526 static_assert(std::is_same_v<float,T> || std::is_same_v<double,T>);
527
528 AMREX_ALWAYS_ASSERT(m_real_domain.numPts() > 1);
529#if (AMREX_SPACEDIM == 2)
530 AMREX_ALWAYS_ASSERT(!m_info.twod_mode);
531#else
532 if (m_info.twod_mode) {
533 AMREX_ALWAYS_ASSERT((int(domain.length(0) > 1) +
534 int(domain.length(1) > 1) +
535 int(domain.length(2) > 1)) >= 2);
536 }
537#endif
538
539 {
540 Box subbox = m_sub_helper.make_box(m_real_domain);
541 if (subbox.size() != m_real_domain.size()) {
542 m_r2c_sub = std::make_unique<R2C<T,D,C>>(subbox, m_info);
543 return;
544 }
545 }
546
547 int myproc = ParallelContext::MyProcSub();
548 int nprocs = std::min(ParallelContext::NProcsSub(), m_info.nprocs);
549
550#if (AMREX_SPACEDIM == 3)
551 if (m_info.domain_strategy == DomainStrategy::automatic) {
552 if (m_info.twod_mode) {
553 m_info.domain_strategy = DomainStrategy::slab;
554 } else {
555 int shortside = m_real_domain.shortside();
556 if (shortside < m_info.pencil_threshold*nprocs) {
557 m_info.domain_strategy = DomainStrategy::pencil;
558 } else {
559 m_info.domain_strategy = DomainStrategy::slab;
560 }
561 }
562 }
563
564 if (!m_info.oned_mode) {
565 if (m_info.twod_mode) {
566 m_slab_decomp = true;
567 } else if (m_info.domain_strategy == DomainStrategy::slab && (m_real_domain.length(1) > 1)) {
568 m_slab_decomp = true;
569 }
570 }
571#endif
572
573 auto const ncomp = m_info.batch_size;
574
575 auto bax = amrex::decompose(m_real_domain, nprocs,
576 {AMREX_D_DECL(false,!m_slab_decomp,m_real_domain.length(2)>1)}, true);
577
578 DistributionMapping dmx = detail::make_iota_distromap(bax.size());
579 m_rx.define(bax, dmx, ncomp, 0, MFInfo().SetAlloc(false));
580
581 {
582 BoxList bl = bax.boxList();
583 for (auto & b : bl) {
584 b.shift(-m_real_domain.smallEnd());
585 b.setBig(0, m_spectral_domain_x.bigEnd(0));
586 }
587 BoxArray cbax(std::move(bl));
588 m_cx.define(cbax, dmx, ncomp, 0, MFInfo().SetAlloc(false));
589 }
590
591 m_do_alld_fft = (ParallelDescriptor::NProcs() == 1) &&
592 (! m_info.twod_mode) && (! m_info.oned_mode);
593
594 if (!m_do_alld_fft) // do a series of 1d or 2d ffts
595 {
596 //
597 // make data containers
598 //
599
600#if (AMREX_SPACEDIM >= 2)
602 if ((m_real_domain.length(1) > 1) && !m_slab_decomp && !m_info.oned_mode)
603 {
604 auto cbay = amrex::decompose(m_spectral_domain_y, nprocs,
605 {AMREX_D_DECL(false,true,true)}, true);
606 if (cbay.size() == dmx.size()) {
607 cdmy = dmx;
608 } else {
609 cdmy = detail::make_iota_distromap(cbay.size());
610 }
611 m_cy.define(cbay, cdmy, ncomp, 0, MFInfo().SetAlloc(false));
612 }
613#endif
614
615#if (AMREX_SPACEDIM == 3)
616 if (!m_info.oned_mode && !m_info.twod_mode &&
617 m_real_domain.length(1) > 1 &&
618 m_real_domain.length(2) > 1)
619 {
620 auto cbaz = amrex::decompose(m_spectral_domain_z, nprocs,
621 {false,true,true}, true);
623 if (cbaz.size() == dmx.size()) {
624 cdmz = dmx;
625 } else if (cbaz.size() == cdmy.size()) {
626 cdmz = cdmy;
627 } else {
628 cdmz = detail::make_iota_distromap(cbaz.size());
629 }
630 m_cz.define(cbaz, cdmz, ncomp, 0, MFInfo().SetAlloc(false));
631 }
632#endif
633
634 if constexpr (C) {
635 if (m_slab_decomp) {
636 m_data_1 = detail::make_mfs_share(m_rx, m_cx);
637 m_data_2 = detail::make_mfs_share(m_cz, m_cz);
638 } else {
639 m_data_1 = detail::make_mfs_share(m_rx, m_cz);
640 m_data_2 = detail::make_mfs_share(m_cy, m_cy);
641 // make m_cx an alias to m_rx
642 if (myproc < m_cx.size()) {
643 Box const& box = m_cx.fabbox(myproc);
644 using FAB = typename cMF::FABType::value_type;
645 m_cx.setFab(myproc, FAB(box, ncomp, m_rx[myproc].dataPtr()));
646 }
647 }
648 } else {
649 if (m_slab_decomp) {
650 m_data_1 = detail::make_mfs_share(m_rx, m_cz);
651 m_data_2 = detail::make_mfs_share(m_cx, m_cx);
652 } else {
653 m_data_1 = detail::make_mfs_share(m_rx, m_cy);
654 m_data_2 = detail::make_mfs_share(m_cx, m_cz);
655 }
656 }
657
658 //
659 // make copiers
660 //
661
662 m_nprocs_active = detail::count_active_procs({&m_rx, &m_cx, &m_cy, &m_cz});
663
664#if (AMREX_SPACEDIM >= 2)
665 if (! m_cy.empty()) {
666 // comm meta-data between x and y phases
667 m_cmd_x2y = detail::make_transpose
668 (m_nprocs_active, m_cy, m_spectral_domain_y, m_cx, IntVect(0), m_dtos_x2y);
669 m_cmd_y2x = detail::make_transpose
670 (m_nprocs_active, m_cx, m_spectral_domain_x, m_cy, IntVect(0), m_dtos_y2x);
671 }
672#endif
673#if (AMREX_SPACEDIM == 3)
674 if (! m_cz.empty() ) {
675 if (m_slab_decomp) {
676 // comm meta-data between xy and z phases
677 m_cmd_x2z = detail::make_transpose
678 (m_nprocs_active, m_cz, m_spectral_domain_z, m_cx, IntVect(0), m_dtos_x2z);
679 m_cmd_z2x = detail::make_transpose
680 (m_nprocs_active, m_cx, m_spectral_domain_x, m_cz, IntVect(0), m_dtos_z2x);
681 } else {
682 // comm meta-data between y and z phases
683 m_cmd_y2z = detail::make_transpose
684 (m_nprocs_active, m_cz, m_spectral_domain_z, m_cy, IntVect(0), m_dtos_y2z);
685 m_cmd_z2y = detail::make_transpose
686 (m_nprocs_active, m_cy, m_spectral_domain_y, m_cz, IntVect(0), m_dtos_z2y);
687 }
688 }
689#endif
690
691 m_info.comm_strategy = detail::resolve_comm_strategy
692 (m_info.comm_strategy, std::max({detail::max_fab_bytes(m_cx),
693 detail::max_fab_bytes(m_cy),
694 detail::max_fab_bytes(m_cz)}),
695 detail::max_block_bytes({&m_cmd_x2y, &m_cmd_y2x, &m_cmd_y2z,
696 &m_cmd_z2y, &m_cmd_x2z, &m_cmd_z2x}),
697 detail::transposes_dense({&m_cmd_x2y, &m_cmd_y2x, &m_cmd_y2z,
698 &m_cmd_z2y, &m_cmd_x2z, &m_cmd_z2x}));
699 if (m_info.comm_strategy != CommStrategy::pointtopoint) {
700 m_comm_active = detail::make_active_comm(m_nprocs_active, m_owns_comm_active);
701 }
702
703 //
704 // make plans
705 //
706
707 if (myproc < m_rx.size())
708 {
709 if constexpr (C) {
710 int ndims = m_slab_decomp ? 2 : 1;
711 std::tie(m_fft_fwd_x, m_fft_bwd_x) = make_c2c_plans(m_cx, ndims);
712 } else {
713 Box const& box = m_rx.box(myproc);
714 auto* pr = m_rx[myproc].dataPtr();
715 auto* pc = (typename Plan<T>::VendorComplex *)m_cx[myproc].dataPtr();
716#ifdef AMREX_USE_SYCL
717 m_fft_fwd_x.template init_r2c<Direction::forward>(box, pr, pc, m_slab_decomp, ncomp);
718 m_fft_bwd_x = m_fft_fwd_x;
719#else
720 if constexpr (D == Direction::both || D == Direction::forward) {
721 m_fft_fwd_x.template init_r2c<Direction::forward>(box, pr, pc, m_slab_decomp, ncomp);
722 }
723 if constexpr (D == Direction::both || D == Direction::backward) {
724 m_fft_bwd_x.template init_r2c<Direction::backward>(box, pr, pc, m_slab_decomp, ncomp);
725 }
726#endif
727 }
728 }
729
730#if (AMREX_SPACEDIM >= 2)
731 if (! m_cy.empty()) {
732 std::tie(m_fft_fwd_y, m_fft_bwd_y) = make_c2c_plans(m_cy,1);
733 }
734#endif
735#if (AMREX_SPACEDIM == 3)
736 if (! m_cz.empty()) {
737 std::tie(m_fft_fwd_z, m_fft_bwd_z) = make_c2c_plans(m_cz,1);
738 }
739#endif
740 }
741 else // do fft in all dimensions at the same time
742 {
743 if constexpr (C) {
744 m_data_1 = detail::make_mfs_share(m_rx, m_cx);
745 std::tie(m_fft_fwd_x, m_fft_bwd_x) = make_c2c_plans(m_cx,AMREX_SPACEDIM);
746 } else {
747 m_data_1 = detail::make_mfs_share(m_rx, m_rx);
748 m_data_2 = detail::make_mfs_share(m_cx, m_cx);
749
750 auto const& len = m_real_domain.length();
751 auto* pr = (void*)m_rx[0].dataPtr();
752 auto* pc = (void*)m_cx[0].dataPtr();
753#ifdef AMREX_USE_SYCL
754 m_fft_fwd_x.template init_r2c<Direction::forward>(len, pr, pc, false, ncomp);
755 m_fft_bwd_x = m_fft_fwd_x;
756#else
757 if constexpr (D == Direction::both || D == Direction::forward) {
758 m_fft_fwd_x.template init_r2c<Direction::forward>(len, pr, pc, false, ncomp);
759 }
760 if constexpr (D == Direction::both || D == Direction::backward) {
761 m_fft_bwd_x.template init_r2c<Direction::backward>(len, pr, pc, false, ncomp);
762 }
763#endif
764 }
765 }
766}
767
768template <typename T, Direction D, bool C>
769R2C<T,D,C>::R2C (std::array<int,AMREX_SPACEDIM> const& domain_size, Info const& info)
770 : R2C<T,D,C>(Box(IntVect(0),IntVect(domain_size)-1), info)
771{}
772
773template <typename T, Direction D, bool C>
775{
776 if (m_fft_bwd_x.plan != m_fft_fwd_x.plan) {
777 m_fft_bwd_x.destroy();
778 }
779 if (m_fft_bwd_y.plan != m_fft_fwd_y.plan) {
780 m_fft_bwd_y.destroy();
781 }
782 if (m_fft_bwd_z.plan != m_fft_fwd_z.plan) {
783 m_fft_bwd_z.destroy();
784 }
785 m_fft_fwd_x.destroy();
786 m_fft_fwd_y.destroy();
787 m_fft_fwd_z.destroy();
788 if (m_fft_bwd_x_half.plan != m_fft_fwd_x_half.plan) {
789 m_fft_bwd_x_half.destroy();
790 }
791 m_fft_fwd_x_half.destroy();
792 detail::free_active_comm(m_comm_active, m_owns_comm_active);
793}
794
795template <typename T, Direction D, bool C>
796std::pair<BoxArray,DistributionMapping>
797R2C<T,D,C>::make_layout_from_local_domain (std::array<int,AMREX_SPACEDIM> const& local_start,
798 std::array<int,AMREX_SPACEDIM> const& local_size)
799{
800 IntVect lo(local_start);
801 IntVect len(local_size);
802 Box bx(lo, lo+len-1);
803#ifdef AMREX_USE_MPI
805 MPI_Allgather(&bx, 1, ParallelDescriptor::Mpi_typemap<Box>::type(),
806 allboxes.data(), 1, ParallelDescriptor::Mpi_typemap<Box>::type(),
808 Vector<int> pmap;
809 pmap.reserve(allboxes.size());
810 for (int i = 0; i < allboxes.size(); ++i) {
811 if (allboxes[i].ok()) {
812 pmap.push_back(ParallelContext::local_to_global_rank(i));
813 }
814 }
815 std::erase_if(allboxes, [=] (Box const& b) { return b.isEmpty(); });
816 BoxList bl(std::move(allboxes));
817 return std::make_pair(BoxArray(std::move(bl)), DistributionMapping(std::move(pmap)));
818#else
819 return std::make_pair(BoxArray(bx), DistributionMapping(Vector<int>({0})));
820#endif
821}
822
823template <typename T, Direction D, bool C>
824void R2C<T,D,C>::setLocalDomain (std::array<int,AMREX_SPACEDIM> const& local_start,
825 std::array<int,AMREX_SPACEDIM> const& local_size)
826{
827 auto const& [ba, dm] = make_layout_from_local_domain(local_start, local_size);
828 m_raw_mf = MF(ba, dm, m_info.batch_size, 0, MFInfo().SetAlloc(false));
829}
830
831template <typename T, Direction D, bool C>
832std::pair<std::array<int,AMREX_SPACEDIM>,std::array<int,AMREX_SPACEDIM>>
834{
835 auto const ncomp = m_info.batch_size;
836
837#if (AMREX_SPACEDIM > 1)
838 if (m_r2c_sub) {
839 m_raw_mf = MF(m_sub_helper.inverse_boxarray(m_r2c_sub->m_rx.boxArray()),
840 m_r2c_sub->m_rx.DistributionMap(), ncomp, 0,
841 MFInfo{}.SetAlloc(false));
842 } else
843#endif
844 {
845 m_raw_mf = MF(m_rx.boxArray(), m_rx.DistributionMap(), ncomp, 0,
846 MFInfo{}.SetAlloc(false));
847 }
848
849 auto const myproc = ParallelContext::MyProcSub();
850 if (myproc < m_raw_mf.size()) {
851 Box const& box = m_raw_mf.box(myproc);
852 return std::make_pair(box.smallEnd().toArray(),
853 box.length().toArray());
854 } else {
855 return std::make_pair(std::array<int,AMREX_SPACEDIM>{AMREX_D_DECL(0,0,0)},
856 std::array<int,AMREX_SPACEDIM>{AMREX_D_DECL(0,0,0)});
857 }
858}
859
860template <typename T, Direction D, bool C>
861void R2C<T,D,C>::setLocalSpectralDomain (std::array<int,AMREX_SPACEDIM> const& local_start,
862 std::array<int,AMREX_SPACEDIM> const& local_size)
863{
864 auto const& [ba, dm] = make_layout_from_local_domain(local_start, local_size);
865 m_raw_cmf = cMF(ba, dm, m_info.batch_size, 0, MFInfo().SetAlloc(false));
866}
867
868template <typename T, Direction D, bool C>
869std::pair<std::array<int,AMREX_SPACEDIM>,std::array<int,AMREX_SPACEDIM>>
871{
872 auto const ncomp = m_info.batch_size;
873 auto const& [ba, dm] = getSpectralDataLayout();
874
875 m_raw_cmf = cMF(ba, dm, ncomp, 0, MFInfo{}.SetAlloc(false));
876
877 auto const myproc = ParallelContext::MyProcSub();
878 if (myproc < m_raw_cmf.size()) {
879 Box const& box = m_raw_cmf.box(myproc);
880 return std::make_pair(box.smallEnd().toArray(), box.length().toArray());
881 } else {
882 return std::make_pair(std::array<int,AMREX_SPACEDIM>{AMREX_D_DECL(0,0,0)},
883 std::array<int,AMREX_SPACEDIM>{AMREX_D_DECL(0,0,0)});
884 }
885}
886
887template <typename T, Direction D, bool C>
888void R2C<T,D,C>::prepare_openbc ()
889{
890 if (C || m_r2c_sub) { amrex::Abort("R2C: OpenBC not supported with reduced dimensions or complex inputs"); }
891
892#if (AMREX_SPACEDIM == 3)
893 if (m_do_alld_fft) { return; }
894
895 auto const ncomp = m_info.batch_size;
896
897 if (m_slab_decomp && ! m_fft_fwd_x_half.defined) {
898 auto* fab = detail::get_fab(m_rx);
899 if (fab) {
900 Box bottom_half = m_real_domain;
901 bottom_half.growHi(2,-m_real_domain.length(2)/2);
902 Box box = fab->box() & bottom_half;
903 if (box.ok()) {
904 auto* pr = fab->dataPtr();
905 auto* pc = (typename Plan<T>::VendorComplex *)
906 detail::get_fab(m_cx)->dataPtr();
907#ifdef AMREX_USE_SYCL
908 m_fft_fwd_x_half.template init_r2c<Direction::forward>
909 (box, pr, pc, m_slab_decomp, ncomp);
910 m_fft_bwd_x_half = m_fft_fwd_x_half;
911#else
912 if constexpr (D == Direction::both || D == Direction::forward) {
913 m_fft_fwd_x_half.template init_r2c<Direction::forward>
914 (box, pr, pc, m_slab_decomp, ncomp);
915 }
916 if constexpr (D == Direction::both || D == Direction::backward) {
917 m_fft_bwd_x_half.template init_r2c<Direction::backward>
918 (box, pr, pc, m_slab_decomp, ncomp);
919 }
920#endif
921 }
922 }
923 } // else todo
924
925 if (m_cmd_x2z && ! m_cmd_x2z_half) {
926 Box bottom_half = m_spectral_domain_z;
927 // Note that z-direction's index is 0 because we z is the
928 // unit-stride direction here.
929 bottom_half.growHi(0,-m_spectral_domain_z.length(0)/2);
930 m_cmd_x2z_half = detail::make_transpose
931 (m_nprocs_active, m_cz, bottom_half, m_cx, IntVect(0), m_dtos_x2z);
932 }
933
934 if (m_cmd_z2x && ! m_cmd_z2x_half) {
935 Box bottom_half = m_spectral_domain_x;
936 bottom_half.growHi(2,-m_spectral_domain_x.length(2)/2);
937 m_cmd_z2x_half = detail::make_transpose
938 (m_nprocs_active, m_cx, bottom_half, m_cz, IntVect(0), m_dtos_z2x);
939 }
940#endif
941}
942
943template <typename T, Direction D, bool C>
944typename R2C<T,D,C>::MF&
945R2C<T,D,C>::make_m_rx_alias (IndexType const& ix_type)
946{
947 auto const key = ix_type.ixType();
948 auto it = m_rx_aliases.find(key);
949 if (it != m_rx_aliases.end()) {
950 return it->second;
951 }
952
953 BoxList bl(ix_type);
954 for (auto b : m_rx.boxArray().boxList()) {
955 b.setType(ix_type);
956 bl.push_back(b);
957 }
958
959 MF alias(BoxArray(std::move(bl)), m_rx.DistributionMap(), m_rx.nComp(),
960 m_rx.nGrowVect(), MFInfo{}.SetAlloc(false));
961
962 using FAB = typename MF::fab_type;
963 for (MFIter mfi(alias, MFItInfo().DisableDeviceSync()); mfi.isValid(); ++mfi) {
964 alias.setFab(mfi, FAB(mfi.fabbox(), m_rx.nComp(), m_rx[mfi.index()].dataPtr()));
965 }
966
967 auto [alias_it, inserted] = m_rx_aliases.emplace(key, std::move(alias));
968 amrex::ignore_unused(inserted);
969 return alias_it->second;
970}
971
972template <typename T, Direction D, bool C>
973Periodicity
974R2C<T,D,C>::make_m_rx_alias_periodicity (IndexType const& ix_type,
975 Periodicity const& period) const
976{
977 IntVect per = period.intVect();
978 for (int idim = 0; idim < AMREX_SPACEDIM; ++idim) {
979 if (ix_type.nodeCentered(idim) && per[idim] == 0) {
980 per[idim] = m_real_domain.length(idim);
981 }
982 }
983 return Periodicity(per);
984}
985
986template <typename T, Direction D, bool C>
987template <Direction DIR>
988requires ((DIR == Direction::forward) || (DIR == Direction::both))
989void R2C<T,D,C>::forward (MF const& inmf, int incomp)
990{
991 BL_PROFILE("FFT::R2C::forward(in)");
992
993 auto const ncomp = m_info.batch_size;
994
995 if (m_r2c_sub) {
996 if (m_sub_helper.ghost_safe(inmf.nGrowVect())) {
997 m_r2c_sub->forward(m_sub_helper.make_alias_mf(inmf), incomp);
998 } else {
999 MF tmp(inmf.boxArray(), inmf.DistributionMap(), ncomp, 0);
1000 tmp.LocalCopy(inmf, incomp, 0, ncomp, IntVect(0));
1001 m_r2c_sub->forward(m_sub_helper.make_alias_mf(tmp),0);
1002 }
1003 return;
1004 }
1005
1006 if (&m_rx != &inmf) {
1007 if (m_rx.boxArray().ixType() == inmf.boxArray().ixType()) {
1008 m_rx.ParallelCopy(inmf, incomp, 0, ncomp);
1009 } else {
1010 MF& rx_alias = make_m_rx_alias(inmf.boxArray().ixType());
1011 rx_alias.ParallelCopy(inmf, incomp, 0, ncomp);
1012 }
1013 }
1014
1015 if (m_do_alld_fft) {
1016 if constexpr (C) {
1017 m_fft_fwd_x.template compute_c2c<Direction::forward>();
1018 } else {
1019 m_fft_fwd_x.template compute_r2c<Direction::forward>();
1020 }
1021 return;
1022 }
1023
1024 auto& fft_x = m_openbc_half ? m_fft_fwd_x_half : m_fft_fwd_x;
1025 if constexpr (C) {
1026 fft_x.template compute_c2c<Direction::forward>();
1027 } else {
1028 fft_x.template compute_r2c<Direction::forward>();
1029 }
1030
1031 if ( m_cmd_x2y) {
1032 detail::transpose(m_cy, m_cx, *m_cmd_x2y, ncomp, m_dtos_x2y, m_info.comm_strategy, m_comm_active);
1033 }
1034 m_fft_fwd_y.template compute_c2c<Direction::forward>();
1035
1036 if ( m_cmd_y2z) {
1037 detail::transpose(m_cz, m_cy, *m_cmd_y2z, ncomp, m_dtos_y2z, m_info.comm_strategy, m_comm_active);
1038 }
1039#if (AMREX_SPACEDIM == 3)
1040 else if ( m_cmd_x2z) {
1041 if (m_openbc_half) {
1042 // Only the lower half is received; zero the upper half.
1043 Box upper_half = m_spectral_domain_z;
1044 // Note that z-direction's index is 0 because we z is the
1045 // unit-stride direction here.
1046 upper_half.growLo (0,-m_spectral_domain_z.length(0)/2);
1047 m_cz.setVal(0, upper_half, 0, ncomp);
1048 }
1049 auto const& cmd = m_openbc_half ? m_cmd_x2z_half : m_cmd_x2z;
1050 detail::transpose(m_cz, m_cx, *cmd, ncomp, m_dtos_x2z, m_info.comm_strategy, m_comm_active);
1051 }
1052#endif
1053 m_fft_fwd_z.template compute_c2c<Direction::forward>();
1054}
1055
1056template <typename T, Direction D, bool C>
1057template <typename FA, typename RT>
1058std::pair<std::unique_ptr<char,DataDeleter>,std::size_t>
1059R2C<T,D,C>::install_raw_ptr (FA& fa, RT const* p)
1060{
1061 AMREX_ALWAYS_ASSERT(!fa.empty());
1062
1063 using FAB = typename FA::FABType::value_type;
1064 using T_FAB = typename FAB::value_type;
1065 static_assert(sizeof(T_FAB) == sizeof(RT));
1066
1067 auto const ncomp = m_info.batch_size;
1068 auto const& ia = fa.IndexArray();
1069
1070 T_FAB* pp = nullptr;
1071 std::size_t sz = 0;
1072
1073 if ( ! ia.empty() ) {
1074 int K = ia[0];
1075 Box const& box = fa.fabbox(K);
1076 if ((alignof(T_FAB) == alignof(RT)) || amrex::is_aligned(p,alignof(T_FAB))) {
1077 pp = (T_FAB*)p;
1078 } else {
1079 sz = sizeof(T_FAB) * box.numPts() * ncomp;
1080 pp = (T_FAB*) The_Arena()->alloc(sz);
1081 }
1082 fa.setFab(K, FAB(box,ncomp,pp));
1083 }
1084
1085 if (sz == 0) {
1086 return std::make_pair(std::unique_ptr<char,DataDeleter>{},std::size_t(0));
1087 } else {
1088 return std::make_pair(std::unique_ptr<char,DataDeleter>
1089 {(char*)pp,DataDeleter{The_Arena()}}, sz);
1090 }
1091}
1092
1093
1094template <typename T, Direction D, bool C>
1095template <typename RT, typename CT, Direction DIR, bool CP>
1096requires ((DIR == Direction::forward) || (DIR == Direction::both))
1097 && (((sizeof(RT)*2 == sizeof(CT)) && (!CP))
1098 || ((sizeof(RT) == sizeof(CT)) && CP))
1099void R2C<T,D,C>::forward (RT const* in, CT* out)
1100{
1101 auto [rdata, rsz] = install_raw_ptr(m_raw_mf, in);
1102 auto [cdata, csz] = install_raw_ptr(m_raw_cmf, out);
1103
1104 if (rsz > 0) {
1105 Gpu::dtod_memcpy_async(rdata.get(),in,rsz);
1107 }
1108
1109 forward(m_raw_mf, m_raw_cmf);
1110
1111 if (csz) {
1112 Gpu::dtod_memcpy_async(out,cdata.get(),csz);
1114 }
1115}
1116
1117template <typename T, Direction D, bool C>
1118template <Direction DIR>
1119requires (DIR == Direction::both)
1120void R2C<T,D,C>::backward (MF& outmf, int outcomp)
1121{
1122 backward_doit(outmf, IntVect(0), Periodicity::NonPeriodic(), outcomp);
1123}
1124
1125template <typename T, Direction D, bool C>
1126void R2C<T,D,C>::backward_doit (MF& outmf, IntVect const& ngout,
1127 Periodicity const& period, int outcomp)
1128{
1129 BL_PROFILE("FFT::R2C::backward(out)");
1130
1131 auto const ncomp = m_info.batch_size;
1132
1133 if (m_r2c_sub) {
1134 if (m_sub_helper.ghost_safe(outmf.nGrowVect())) {
1135 MF submf = m_sub_helper.make_alias_mf(outmf);
1136 IntVect const& subngout = m_sub_helper.make_iv(ngout);
1137 Periodicity const& subperiod = m_sub_helper.make_periodicity(period);
1138 m_r2c_sub->backward_doit(submf, subngout, subperiod, outcomp);
1139 } else {
1140 MF tmp(outmf.boxArray(), outmf.DistributionMap(), ncomp,
1141 m_sub_helper.make_safe_ghost(outmf.nGrowVect()));
1142 this->backward_doit(tmp, ngout, period, 0);
1143 outmf.LocalCopy(tmp, 0, outcomp, ncomp, tmp.nGrowVect());
1144 }
1145 return;
1146 }
1147
1148 if (m_do_alld_fft) {
1149 if constexpr (C) {
1150 m_fft_bwd_x.template compute_c2c<Direction::backward>();
1151 } else {
1152 m_fft_bwd_x.template compute_r2c<Direction::backward>();
1153 }
1154 auto const dst_nghost = amrex::elemwiseMin(ngout,outmf.nGrowVect());
1155 if (m_rx.boxArray().ixType() == outmf.boxArray().ixType()) {
1156 outmf.ParallelCopy(m_rx, 0, outcomp, ncomp, IntVect(0),
1157 dst_nghost, period);
1158 } else {
1159 MF& rx_alias = make_m_rx_alias(outmf.boxArray().ixType());
1160 auto const copy_period = make_m_rx_alias_periodicity(outmf.boxArray().ixType(), period);
1161 outmf.ParallelCopy(rx_alias, 0, outcomp, ncomp, IntVect(0),
1162 dst_nghost, copy_period);
1163 }
1164 return;
1165 }
1166
1167 m_fft_bwd_z.template compute_c2c<Direction::backward>();
1168 if ( m_cmd_z2y) {
1169 detail::transpose(m_cy, m_cz, *m_cmd_z2y, ncomp, m_dtos_z2y, m_info.comm_strategy, m_comm_active);
1170 }
1171#if (AMREX_SPACEDIM == 3)
1172 else if ( m_cmd_z2x) {
1173 auto const& cmd = m_openbc_half ? m_cmd_z2x_half : m_cmd_z2x;
1174 detail::transpose(m_cx, m_cz, *cmd, ncomp, m_dtos_z2x, m_info.comm_strategy, m_comm_active);
1175 }
1176#endif
1177
1178 m_fft_bwd_y.template compute_c2c<Direction::backward>();
1179 if ( m_cmd_y2x) {
1180 detail::transpose(m_cx, m_cy, *m_cmd_y2x, ncomp, m_dtos_y2x, m_info.comm_strategy, m_comm_active);
1181 }
1182
1183 auto& fft_x = m_openbc_half ? m_fft_bwd_x_half : m_fft_bwd_x;
1184 if constexpr (C) {
1185 fft_x.template compute_c2c<Direction::backward>();
1186 } else {
1187 fft_x.template compute_r2c<Direction::backward>();
1188 }
1189 auto const dst_nghost = amrex::elemwiseMin(ngout,outmf.nGrowVect());
1190 if (m_rx.boxArray().ixType() == outmf.boxArray().ixType()) {
1191 outmf.ParallelCopy(m_rx, 0, outcomp, ncomp, IntVect(0),
1192 dst_nghost, period);
1193 } else {
1194 MF& rx_alias = make_m_rx_alias(outmf.boxArray().ixType());
1195 auto const copy_period = make_m_rx_alias_periodicity(outmf.boxArray().ixType(), period);
1196 outmf.ParallelCopy(rx_alias, 0, outcomp, ncomp, IntVect(0),
1197 dst_nghost, copy_period);
1198 }
1199}
1200
1201template <typename T, Direction D, bool C>
1202template <typename CT, typename RT, Direction DIR, bool CP>
1203requires ((DIR == Direction::backward) || (DIR == Direction::both))
1204 && (((sizeof(RT)*2 == sizeof(CT)) && (!CP))
1205 || ((sizeof(RT) == sizeof(CT)) && CP))
1206void R2C<T,D,C>::backward (CT const* in, RT* out)
1207{
1208 auto [rdata, rsz] = install_raw_ptr(m_raw_mf, out);
1209 auto [cdata, csz] = install_raw_ptr(m_raw_cmf, in);
1210
1211 if (csz) {
1212 Gpu::dtod_memcpy_async(cdata.get(),in,csz);
1214 }
1215
1216 backward(m_raw_cmf, m_raw_mf);
1217
1218 if (rsz > 0) {
1219 Gpu::dtod_memcpy_async(out,rdata.get(),rsz);
1221 }
1222}
1223
1224template <typename T, Direction D, bool C>
1225std::pair<Plan<T>, Plan<T>>
1226R2C<T,D,C>::make_c2c_plans (cMF& inout, int ndims) const
1227{
1228 Plan<T> fwd;
1229 Plan<T> bwd;
1230
1231 auto* fab = detail::get_fab(inout);
1232 if (!fab) { return {fwd, bwd};}
1233
1234 Box const& box = fab->box();
1235 auto* pio = (typename Plan<T>::VendorComplex *)fab->dataPtr();
1236
1237 auto const ncomp = m_info.batch_size;
1238
1239#ifdef AMREX_USE_SYCL
1240 fwd.template init_c2c<Direction::forward>(box, pio, ncomp, ndims);
1241 bwd = fwd;
1242#else
1243 if constexpr (D == Direction::both || D == Direction::forward) {
1244 fwd.template init_c2c<Direction::forward>(box, pio, ncomp, ndims);
1245 }
1246 if constexpr (D == Direction::both || D == Direction::backward) {
1247 bwd.template init_c2c<Direction::backward>(box, pio, ncomp, ndims);
1248 }
1249#endif
1250
1251 return {fwd, bwd};
1252}
1253
1254namespace fft_detail {
1255 // The trailing int/long parameter orders these two overloads. A call
1256 // passing a literal 0 prefers the GpuComplex<T>& version, which is
1257 // called once per component of the batch, and falls back to the CellData
1258 // version only when the former is not viable. Without it, a functor
1259 // taking a generic parameter by value or by const reference would make
1260 // the two overloads ambiguous.
1261 template <typename F, typename T>
1263 auto call_post_forward (F const& f, int i, int j, int k, Array4<T> const& a,
1264 int ii, int jj, int kk, int)
1265 noexcept -> decltype(void(f(0,0,0,a(0,0,0))))
1266 {
1267 for (int n = 0; n < a.nComp(); ++n) {
1268 f(i,j,k,a(ii,jj,kk,n));
1269 }
1270 }
1271
1272 template <typename F, typename T>
1274 auto call_post_forward (F const& f, int i, int j, int k, Array4<T> const& a,
1275 int ii, int jj, int kk, long)
1276 noexcept -> decltype(void(f(0,0,0,a.cellData(0,0,0))))
1277 {
1278 f(i,j,k,a.cellData(ii,jj,kk));
1279 }
1280}
1281
1282template <typename T, Direction D, bool C>
1283template <typename F>
1284void R2C<T,D,C>::post_forward_doit_0 (F const& post_forward)
1285{
1286 if (m_info.twod_mode) {
1287 amrex::Abort("xxxxx todo: post_forward");
1288#if (AMREX_SPACEDIM > 1)
1289 } else if (m_r2c_sub) {
1290 // We need to pass the originally ordered indices to post_forward.
1291#if (AMREX_SPACEDIM == 2)
1292 // The original domain is (1,ny). The sub domain is (ny,1).
1293 m_r2c_sub->post_forward_doit_1
1294 ([=] AMREX_GPU_DEVICE (int i, int, int, auto&& sp)
1295 -> decltype(void(post_forward(0,0,0,static_cast<decltype(sp)>(sp))))
1296 {
1297 post_forward(0, i, 0, static_cast<decltype(sp)>(sp));
1298 });
1299#else
1300 if (m_real_domain.length(0) == 1 && m_real_domain.length(1) == 1) {
1301 // Original domain: (1, 1, nz). Sub domain: (nz, 1, 1)
1302 m_r2c_sub->post_forward_doit_1
1303 ([=] AMREX_GPU_DEVICE (int i, int, int, auto&& sp)
1304 -> decltype(void(post_forward(0,0,0,static_cast<decltype(sp)>(sp))))
1305 {
1306 post_forward(0, 0, i, static_cast<decltype(sp)>(sp));
1307 });
1308 } else if (m_real_domain.length(0) == 1 && m_real_domain.length(2) == 1) {
1309 // Original domain: (1, ny, 1). Sub domain: (ny, 1, 1)
1310 m_r2c_sub->post_forward_doit_1
1311 ([=] AMREX_GPU_DEVICE (int i, int, int, auto&& sp)
1312 -> decltype(void(post_forward(0,0,0,static_cast<decltype(sp)>(sp))))
1313 {
1314 post_forward(0, i, 0, static_cast<decltype(sp)>(sp));
1315 });
1316 } else if (m_real_domain.length(0) == 1) {
1317 // Original domain: (1, ny, nz). Sub domain: (ny, nz, 1)
1318 m_r2c_sub->post_forward_doit_1
1319 ([=] AMREX_GPU_DEVICE (int i, int j, int, auto&& sp)
1320 -> decltype(void(post_forward(0,0,0,static_cast<decltype(sp)>(sp))))
1321 {
1322 post_forward(0, i, j, static_cast<decltype(sp)>(sp));
1323 });
1324 } else if (m_real_domain.length(1) == 1) {
1325 // Original domain: (nx, 1, nz). Sub domain: (nx, nz, 1)
1326 m_r2c_sub->post_forward_doit_1
1327 ([=] AMREX_GPU_DEVICE (int i, int j, int, auto&& sp)
1328 -> decltype(void(post_forward(0,0,0,static_cast<decltype(sp)>(sp))))
1329 {
1330 post_forward(i, 0, j, static_cast<decltype(sp)>(sp));
1331 });
1332 } else {
1333 amrex::Abort("R2c::post_forward_doit_0: how did this happen?");
1334 }
1335#endif
1336#endif
1337 } else {
1338 this->post_forward_doit_1(post_forward);
1339 }
1340}
1341
1342template <typename T, Direction D, bool C>
1343template <typename F>
1344void R2C<T,D,C>::post_forward_doit_1 (F const& post_forward)
1345{
1346 if (m_info.twod_mode) {
1347 amrex::Abort("xxxxx todo: post_forward");
1348 } else if (m_r2c_sub) {
1349 amrex::Abort("R2C::post_forward_doit_1: How did this happen?");
1350 } else {
1351 if ( ! m_cz.empty()) {
1352 auto* spectral_fab = detail::get_fab(m_cz);
1353 if (spectral_fab) {
1354 auto const& a = spectral_fab->array(); // m_cz's ordering is z,x,y
1355 ParallelForOMP(spectral_fab->box(),
1356 [=] AMREX_GPU_DEVICE (int iz, int jx, int ky)
1357 {
1358 fft_detail::call_post_forward(post_forward,
1359 jx,ky,iz,a,iz,jx,ky,0);
1360 });
1361 }
1362 } else if ( ! m_cy.empty()) {
1363 auto* spectral_fab = detail::get_fab(m_cy);
1364 if (spectral_fab) {
1365 auto const& a = spectral_fab->array(); // m_cy's ordering is y,x,z
1366 ParallelForOMP(spectral_fab->box(),
1367 [=] AMREX_GPU_DEVICE (int iy, int jx, int k)
1368 {
1369 fft_detail::call_post_forward(post_forward,
1370 jx,iy,k,a,iy,jx,k,0);
1371 });
1372 }
1373 } else {
1374 auto* spectral_fab = detail::get_fab(m_cx);
1375 if (spectral_fab) {
1376 auto const& a = spectral_fab->array();
1377 ParallelForOMP(spectral_fab->box(),
1378 [=] AMREX_GPU_DEVICE (int i, int j, int k)
1379 {
1380 fft_detail::call_post_forward(post_forward,
1381 i,j,k,a,i,j,k,0);
1382 });
1383 }
1384 }
1385 }
1386}
1387
1388template <typename T, Direction D, bool C>
1390{
1391#if (AMREX_SPACEDIM > 1)
1392 // The sub-object knows which axes are actually transformed.
1393 if (m_r2c_sub) { return m_r2c_sub->scalingFactor(); }
1394#endif
1395
1396#if (AMREX_SPACEDIM == 3)
1397 // oned_mode wins over twod_mode, like it does when the plans are made.
1398 if (m_info.oned_mode) {
1399 return T(1)/T(Long(m_real_domain.length(0)));
1400 } else if (m_info.twod_mode) {
1401 return T(1)/T(Long(m_real_domain.length(0)) *
1402 Long(m_real_domain.length(1)));
1403 } else
1404#elif (AMREX_SPACEDIM == 2)
1405 if (m_info.oned_mode) {
1406 return T(1)/T(Long(m_real_domain.length(0)));
1407 } else
1408#endif
1409 {
1410 return T(1)/T(m_real_domain.numPts());
1411 }
1412}
1413
1414template <typename T, Direction D, bool C>
1415template <Direction DIR>
1416requires ((DIR == Direction::forward) || (DIR == Direction::both))
1417std::pair<typename R2C<T,D,C>::cMF *, IntVect>
1419{
1420#if (AMREX_SPACEDIM > 1)
1421 if (m_r2c_sub) {
1422 auto [cmf, order] = m_r2c_sub->getSpectralData();
1423 return std::make_pair(cmf, m_sub_helper.inverse_order(order));
1424 } else
1425#endif
1426 if (!m_cz.empty()) {
1427 return std::make_pair(const_cast<cMF*>(&m_cz), IntVect{AMREX_D_DECL(2,0,1)});
1428 } else if (!m_cy.empty()) {
1429 return std::make_pair(const_cast<cMF*>(&m_cy), IntVect{AMREX_D_DECL(1,0,2)});
1430 } else {
1431 return std::make_pair(const_cast<cMF*>(&m_cx), IntVect{AMREX_D_DECL(0,1,2)});
1432 }
1433}
1434
1435template <typename T, Direction D, bool C>
1436template <Direction DIR>
1437requires ((DIR == Direction::forward) || (DIR == Direction::both))
1438void R2C<T,D,C>::forward (MF const& inmf, cMF& outmf, int incomp, int outcomp)
1439{
1440 BL_PROFILE("FFT::R2C::forward(inout)");
1441
1442 auto const ncomp = m_info.batch_size;
1443
1444 if (m_r2c_sub)
1445 {
1446 bool inmf_safe = m_sub_helper.ghost_safe(inmf.nGrowVect());
1447 MF inmf_sub, inmf_tmp;
1448 int incomp_sub;
1449 if (inmf_safe) {
1450 inmf_sub = m_sub_helper.make_alias_mf(inmf);
1451 incomp_sub = incomp;
1452 } else {
1453 inmf_tmp.define(inmf.boxArray(), inmf.DistributionMap(), ncomp, 0);
1454 inmf_tmp.LocalCopy(inmf, incomp, 0, ncomp, IntVect(0));
1455 inmf_sub = m_sub_helper.make_alias_mf(inmf_tmp);
1456 incomp_sub = 0;
1457 }
1458
1459 bool outmf_safe = m_sub_helper.ghost_safe(outmf.nGrowVect());
1460 cMF outmf_sub, outmf_tmp;
1461 int outcomp_sub;
1462 if (outmf_safe) {
1463 outmf_sub = m_sub_helper.make_alias_mf(outmf);
1464 outcomp_sub = outcomp;
1465 } else {
1466 outmf_tmp.define(outmf.boxArray(), outmf.DistributionMap(), ncomp, 0);
1467 outmf_sub = m_sub_helper.make_alias_mf(outmf_tmp);
1468 outcomp_sub = 0;
1469 }
1470
1471 m_r2c_sub->forward(inmf_sub, outmf_sub, incomp_sub, outcomp_sub);
1472
1473 if (!outmf_safe) {
1474 outmf.LocalCopy(outmf_tmp, 0, outcomp, ncomp, IntVect(0));
1475 }
1476 }
1477 else
1478 {
1479 forward(inmf, incomp);
1480 if (!m_cz.empty()) { // m_cz's order (z,x,y) -> (x,y,z)
1481 RotateBwd dtos{};
1483 (outmf, m_spectral_domain_x, m_cz, IntVect(0), dtos);
1484 ParallelCopy(outmf, m_cz, cmd, 0, outcomp, ncomp, dtos);
1485 } else if (!m_cy.empty()) { // m_cy's order (y,x,z) -> (x,y,z)
1487 (outmf, m_spectral_domain_x, m_cy, IntVect(0), m_dtos_y2x);
1488 ParallelCopy(outmf, m_cy, cmd, 0, outcomp, ncomp, m_dtos_y2x);
1489 } else {
1490 outmf.ParallelCopy(m_cx, 0, outcomp, ncomp);
1491 }
1492 }
1493}
1494
1495template <typename T, Direction D, bool C>
1496template <Direction DIR>
1497requires ((DIR == Direction::backward) || (DIR == Direction::both))
1498void R2C<T,D,C>::backward (cMF const& inmf, MF& outmf, int incomp, int outcomp)
1499{
1500 backward_doit(inmf, outmf, IntVect(0), Periodicity::NonPeriodic(), incomp, outcomp);
1501}
1502
1503template <typename T, Direction D, bool C>
1504void R2C<T,D,C>::backward_doit (cMF const& inmf, MF& outmf, IntVect const& ngout,
1505 Periodicity const& period, int incomp, int outcomp)
1506{
1507 BL_PROFILE("FFT::R2C::backward(inout)");
1508
1509 auto const ncomp = m_info.batch_size;
1510
1511 if (m_r2c_sub)
1512 {
1513 bool inmf_safe = m_sub_helper.ghost_safe(inmf.nGrowVect());
1514 cMF inmf_sub, inmf_tmp;
1515 int incomp_sub;
1516 if (inmf_safe) {
1517 inmf_sub = m_sub_helper.make_alias_mf(inmf);
1518 incomp_sub = incomp;
1519 } else {
1520 inmf_tmp.define(inmf.boxArray(), inmf.DistributionMap(), ncomp, 0);
1521 inmf_tmp.LocalCopy(inmf, incomp, 0, ncomp, IntVect(0));
1522 inmf_sub = m_sub_helper.make_alias_mf(inmf_tmp);
1523 incomp_sub = 0;
1524 }
1525
1526 bool outmf_safe = m_sub_helper.ghost_safe(outmf.nGrowVect());
1527 MF outmf_sub, outmf_tmp;
1528 int outcomp_sub;
1529 if (outmf_safe) {
1530 outmf_sub = m_sub_helper.make_alias_mf(outmf);
1531 outcomp_sub = outcomp;
1532 } else {
1533 IntVect const& ngtmp = m_sub_helper.make_safe_ghost(outmf.nGrowVect());
1534 outmf_tmp.define(outmf.boxArray(), outmf.DistributionMap(), ncomp, ngtmp);
1535 outmf_sub = m_sub_helper.make_alias_mf(outmf_tmp);
1536 outcomp_sub = 0;
1537 }
1538
1539 IntVect const& subngout = m_sub_helper.make_iv(ngout);
1540 Periodicity const& subperiod = m_sub_helper.make_periodicity(period);
1541 m_r2c_sub->backward_doit(inmf_sub, outmf_sub, subngout, subperiod, incomp_sub, outcomp_sub);
1542
1543 if (!outmf_safe) {
1544 outmf.LocalCopy(outmf_tmp, 0, outcomp, ncomp, outmf_tmp.nGrowVect());
1545 }
1546 }
1547 else
1548 {
1549 if (!m_cz.empty()) { // (x,y,z) -> m_cz's order (z,x,y)
1550 RotateFwd dtos{};
1551 MultiBlockCommMetaData cmd
1552 (m_cz, m_spectral_domain_z, inmf, IntVect(0), dtos);
1553 ParallelCopy(m_cz, inmf, cmd, incomp, 0, ncomp, dtos);
1554 } else if (!m_cy.empty()) { // (x,y,z) -> m_cy's ordering (y,x,z)
1555 MultiBlockCommMetaData cmd
1556 (m_cy, m_spectral_domain_y, inmf, IntVect(0), m_dtos_x2y);
1557 ParallelCopy(m_cy, inmf, cmd, incomp, 0, ncomp, m_dtos_x2y);
1558 } else {
1559 m_cx.ParallelCopy(inmf, incomp, 0, ncomp);
1560 }
1561 backward_doit(outmf, ngout, period, outcomp);
1562 }
1563}
1564
1565template <typename T, Direction D, bool C>
1566std::pair<BoxArray,DistributionMapping>
1568{
1569#if (AMREX_SPACEDIM > 1)
1570 if (m_r2c_sub) {
1571 auto const& [ba, dm] = m_r2c_sub->getSpectralDataLayout();
1572 return std::make_pair(m_sub_helper.inverse_boxarray(ba), dm);
1573 }
1574#endif
1575
1576#if (AMREX_SPACEDIM == 3)
1577 if (!m_cz.empty()) {
1578 BoxList bl = m_cz.boxArray().boxList();
1579 for (auto& b : bl) {
1580 auto lo = b.smallEnd();
1581 auto hi = b.bigEnd();
1582 std::swap(lo[0], lo[1]);
1583 std::swap(lo[1], lo[2]);
1584 std::swap(hi[0], hi[1]);
1585 std::swap(hi[1], hi[2]);
1586 b.setSmall(lo);
1587 b.setBig(hi);
1588 }
1589 return std::make_pair(BoxArray(std::move(bl)), m_cz.DistributionMap());
1590 } else
1591#endif
1592#if (AMREX_SPACEDIM >= 2)
1593 if (!m_cy.empty()) {
1594 BoxList bl = m_cy.boxArray().boxList();
1595 for (auto& b : bl) {
1596 auto lo = b.smallEnd();
1597 auto hi = b.bigEnd();
1598 std::swap(lo[0], lo[1]);
1599 std::swap(hi[0], hi[1]);
1600 b.setSmall(lo);
1601 b.setBig(hi);
1602 }
1603 return std::make_pair(BoxArray(std::move(bl)), m_cy.DistributionMap());
1604 } else
1605#endif
1606 {
1607 return std::make_pair(m_cx.boxArray(), m_cx.DistributionMap());
1608 }
1609}
1610
1612template <typename T = Real, FFT::Direction D = FFT::Direction::both>
1614
1615}
1616
1617#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:174
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:989
void forward(RT const *in, CT *out)
Forward transform.
Definition AMReX_FFT_R2C.H:1099
friend class Poisson
Definition AMReX_FFT_R2C.H:56
~R2C()
Definition AMReX_FFT_R2C.H:774
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:824
R2C(Box const &domain, Info const &info=Info{})
Constructor.
Definition AMReX_FFT_R2C.H:512
void backward(MF &outmf, int outcomp=0)
Backward transform.
Definition AMReX_FFT_R2C.H:1120
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:1344
void backward(CT const *in, RT *out)
Backward transform.
Definition AMReX_FFT_R2C.H:1206
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:1389
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:1284
std::pair< std::array< int, 3 >, std::array< int, 3 > > getLocalDomain() const
Get local domain.
Definition AMReX_FFT_R2C.H:833
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:1567
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:769
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:870
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:861
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:2498
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:2274
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
@ pointtopoint
Definition AMReX_FFT_Helper.H:62
R2C< T, D, true > C2C
FFT between complex data.
Definition AMReX_FFT_R2C.H:1613
@ backward
Definition AMReX_FFT_Helper.H:57
@ forward
Definition AMReX_FFT_Helper.H:57
@ both
Definition AMReX_FFT_Helper.H:57
@ automatic
Definition AMReX_FFT_Helper.H:59
@ pencil
Definition AMReX_FFT_Helper.H:59
@ slab
Definition AMReX_FFT_Helper.H:59
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
int MPI_Comm
Definition AMReX_ccse-mpi.H:51
static constexpr int MPI_COMM_NULL
Definition AMReX_ccse-mpi.H:59
__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
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
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:88
Definition AMReX_FFT_Helper.H:377
std::conditional_t< std::is_same_v< float, T >, cuComplex, cuDoubleComplex > VendorComplex
Definition AMReX_FFT_Helper.H:380
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 is the index mapping based on the DTOS MultiBlockDestToSrc.
Definition AMReX_NonLocalBC.H:223