Block-Structured AMR Software Framework
AMReX_FFT_OpenBCSolver.H
Go to the documentation of this file.
1 #ifndef AMREX_FFT_OPENBC_SOLVER_H_
2 #define AMREX_FFT_OPENBC_SOLVER_H_
3 
4 #include <AMReX_FFT_R2C.H>
5 
6 namespace amrex::FFT
7 {
8 
9 template <typename T = Real>
11 {
12 public:
13  using MF = typename R2C<T>::MF;
14  using cMF = typename R2C<T>::cMF;
15 
16  explicit OpenBCSolver (Box const& domain, Info const& info = Info{});
17 
18  template <class F>
19  void setGreensFunction (F const& greens_function);
20 
21  void solve (MF& phi, MF const& rho);
22 
23  [[nodiscard]] Box const& Domain () const { return m_domain; }
24 
25 private:
26  static Box make_grown_domain (Box const& domain, Info const& info);
27 
32  std::unique_ptr<R2C<T>> m_r2c_green;
33 };
34 
35 template <typename T>
36 Box OpenBCSolver<T>::make_grown_domain (Box const& domain, Info const& info)
37 {
38  IntVect len = domain.length();
39 #if (AMREX_SPACEDIM == 3)
40  if (info.batch_mode) { len[2] = 0; }
41 #else
43 #endif
44  return Box(domain.smallEnd(), domain.bigEnd()+len, domain.ixType());
45 }
46 
47 template <typename T>
48 OpenBCSolver<T>::OpenBCSolver (Box const& domain, Info const& info)
49  : m_domain(domain),
50  m_info(info),
51  m_r2c(OpenBCSolver<T>::make_grown_domain(domain,info), info)
52 {
53 #if (AMREX_SPACEDIM == 3)
54  if (m_info.batch_mode) {
55  auto gdom = make_grown_domain(domain,m_info);
56  gdom.enclosedCells(2);
57  gdom.setSmall(2, 0);
58  int nprocs = std::min({ParallelContext::NProcsSub(),
59  m_info.nprocs,
60  m_domain.length(2)});
61  gdom.setBig(2, nprocs-1);
62  m_r2c_green = std::make_unique<R2C<T>>(gdom,info);
63  auto [sd, ord] = m_r2c_green->getSpectralData();
64  m_G_fft = cMF(*sd, amrex::make_alias, 0, 1);
65  } else
66 #endif
67  {
69  auto [sd, ord] = m_r2c.getSpectralData();
71  m_G_fft.define(sd->boxArray(), sd->DistributionMap(), 1, 0);
72  }
73 }
74 
75 template <typename T>
76 template <class F>
77 void OpenBCSolver<T>::setGreensFunction (F const& greens_function)
78 {
79  BL_PROFILE("OpenBCSolver::setGreensFunction");
80 
81  auto* infab = m_info.batch_mode ? detail::get_fab(m_r2c_green->m_rx)
82  : detail::get_fab(m_r2c.m_rx);
83  auto const& lo = m_domain.smallEnd();
84  auto const& lo3 = lo.dim3();
85  auto const& len = m_domain.length3d();
86  if (infab) {
87  auto const& a = infab->array();
88  auto box = infab->box();
89  GpuArray<int,3> nimages{1,1,1};
90  int ndims = m_info.batch_mode ? AMREX_SPACEDIM : AMREX_SPACEDIM-1;
91  for (int idim = 0; idim < ndims; ++idim) {
92  if (box.smallEnd(idim) == lo[idim] && box.length(idim) == 2*len[idim]) {
93  box.growHi(idim, -len[idim]+1); // +1 to include the middle plane
94  nimages[idim] = 2;
95  }
96  }
97  AMREX_ASSERT(nimages[0] == 2);
98  box.shift(-lo);
99  amrex::ParallelFor(box, [=] AMREX_GPU_DEVICE (int i, int j, int k)
100  {
101  T G;
102  if (i == len[0] || j == len[1] || k == len[2]) {
103  G = 0;
104  } else {
105  auto ii = i;
106  auto jj = (j > len[1]) ? 2*len[1]-j : j;
107  auto kk = (k > len[2]) ? 2*len[2]-k : k;
108  G = greens_function(ii+lo3.x,jj+lo3.y,kk+lo3.z);
109  }
110  for (int koff = 0; koff < nimages[2]; ++koff) {
111  int k2 = (koff == 0) ? k : 2*len[2]-k;
112  if ((k2 == 2*len[2]) || (koff == 1 && k == len[2])) {
113  continue;
114  }
115  for (int joff = 0; joff < nimages[1]; ++joff) {
116  int j2 = (joff == 0) ? j : 2*len[1]-j;
117  if ((j2 == 2*len[1]) || (joff == 1 && j == len[1])) {
118  continue;
119  }
120  for (int ioff = 0; ioff < nimages[0]; ++ioff) {
121  int i2 = (ioff == 0) ? i : 2*len[0]-i;
122  if ((i2 == 2*len[0]) || (ioff == 1 && i == len[0])) {
123  continue;
124  }
125  a(i2+lo3.x,j2+lo3.y,k2+lo3.z) = G;
126  }
127  }
128  }
129  });
130  }
131 
132  if (m_info.batch_mode) {
133  m_r2c_green->forward(m_r2c_green->m_rx);
134  } else {
135  m_r2c.forward(m_r2c.m_rx);
136  }
137 
138  if (!m_info.batch_mode) {
139  auto [sd, ord] = m_r2c.getSpectralData();
141  auto const* srcfab = detail::get_fab(*sd);
142  if (srcfab) {
143  auto* dstfab = detail::get_fab(m_G_fft);
144  if (dstfab) {
145 #if defined(AMREX_USE_GPU)
147 #else
149 #endif
150  (dstfab->dataPtr(), srcfab->dataPtr(), dstfab->nBytes());
151  } else {
152  amrex::Abort("FFT::OpenBCSolver: how did this happen");
153  }
154  }
155  }
156 
157  m_r2c.prepare_openbc();
158 }
159 
160 template <typename T>
161 void OpenBCSolver<T>::solve (MF& phi, MF const& rho)
162 {
163  BL_PROFILE("OpenBCSolver::solve");
164 
165  auto& inmf = m_r2c.m_rx;
166  inmf.setVal(T(0));
167  inmf.ParallelCopy(rho, 0, 0, 1);
168 
169  m_r2c.m_openbc_half = true;
170  m_r2c.forward(inmf);
171  m_r2c.m_openbc_half = false;
172 
173  auto scaling_factor = m_r2c.scalingFactor();
174 
175  auto const* gfab = detail::get_fab(m_G_fft);
176  if (gfab) {
177  auto [sd, ord] = m_r2c.getSpectralData();
179  auto* rhofab = detail::get_fab(*sd);
180  if (rhofab) {
181  auto* pdst = rhofab->dataPtr();
182  auto const* psrc = gfab->dataPtr();
183  Box const& rhobox = rhofab->box();
184 #if (AMREX_SPACEDIM == 3)
185  Long leng = gfab->box().numPts();
186  if (m_info.batch_mode) {
187  AMREX_ASSERT(gfab->box().length(2) == 1 &&
188  leng == (rhobox.length(0) * rhobox.length(1)));
189  } else {
190  AMREX_ASSERT(leng == rhobox.numPts());
191  }
192 #endif
193  amrex::ParallelFor(rhobox.numPts(), [=] AMREX_GPU_DEVICE (Long i)
194  {
195 #if (AMREX_SPACEDIM == 3)
196  Long isrc = i % leng;
197 #else
198  Long isrc = i;
199 #endif
200  pdst[i] *= psrc[isrc] * scaling_factor;
201  });
202  } else {
203  amrex::Abort("FFT::OpenBCSolver::solve: how did this happen?");
204  }
205  }
206 
207  m_r2c.m_openbc_half = true;
208  m_r2c.backward_doit(phi, phi.nGrowVect());
209  m_r2c.m_openbc_half = false;
210 }
211 
212 }
213 
214 #endif
#define BL_PROFILE(a)
Definition: AMReX_BLProfiler.H:551
#define AMREX_ASSERT(EX)
Definition: AMReX_BLassert.H:38
#define AMREX_GPU_DEVICE
Definition: AMReX_GpuQualifiers.H:18
Real * pdst
Definition: AMReX_HypreMLABecLap.cpp:1090
AMREX_GPU_HOST_DEVICE const IntVectND< dim > & smallEnd() const &noexcept
Get the smallend of the BoxND.
Definition: AMReX_Box.H:105
AMREX_GPU_HOST_DEVICE IndexTypeND< dim > ixType() const noexcept
Returns the indexing type.
Definition: AMReX_Box.H:127
AMREX_GPU_HOST_DEVICE const IntVectND< dim > & bigEnd() const &noexcept
Get the bigend.
Definition: AMReX_Box.H:116
AMREX_GPU_HOST_DEVICE IntVectND< dim > length() const noexcept
Return the length of the BoxND.
Definition: AMReX_Box.H:146
AMREX_GPU_HOST_DEVICE Long numPts() const noexcept
Returns the number of points contained in the BoxND.
Definition: AMReX_Box.H:346
Definition: AMReX_FFT_OpenBCSolver.H:11
Info m_info
Definition: AMReX_FFT_OpenBCSolver.H:29
cMF m_G_fft
Definition: AMReX_FFT_OpenBCSolver.H:31
typename R2C< T >::MF MF
Definition: AMReX_FFT_OpenBCSolver.H:13
void solve(MF &phi, MF const &rho)
Definition: AMReX_FFT_OpenBCSolver.H:161
R2C< T > m_r2c
Definition: AMReX_FFT_OpenBCSolver.H:30
void setGreensFunction(F const &greens_function)
Definition: AMReX_FFT_OpenBCSolver.H:77
Box m_domain
Definition: AMReX_FFT_OpenBCSolver.H:28
typename R2C< T >::cMF cMF
Definition: AMReX_FFT_OpenBCSolver.H:14
std::unique_ptr< R2C< T > > m_r2c_green
Definition: AMReX_FFT_OpenBCSolver.H:32
static Box make_grown_domain(Box const &domain, Info const &info)
Definition: AMReX_FFT_OpenBCSolver.H:36
OpenBCSolver(Box const &domain, Info const &info=Info{})
Definition: AMReX_FFT_OpenBCSolver.H:48
Box const & Domain() const
Definition: AMReX_FFT_OpenBCSolver.H:23
Parallel Discrete Fourier Transform.
Definition: AMReX_FFT_R2C.H:34
std::conditional_t< std::is_same_v< T, Real >, MultiFab, FabArray< BaseFab< T > > > MF
Definition: AMReX_FFT_R2C.H:37
An Array of FortranArrayBox(FAB)-like Objects.
Definition: AMReX_FabArray.H:344
FA::FABType::value_type * get_fab(FA &fa)
Definition: AMReX_FFT_Helper.H:1303
Definition: AMReX_FFT.cpp:7
void dtod_memcpy_async(void *p_d_dst, const void *p_d_src, const std::size_t sz) noexcept
Definition: AMReX_GpuDevice.H:279
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void * memcpy(void *dest, const void *src, std::size_t count)
Definition: AMReX_GpuUtility.H:214
int NProcsSub() noexcept
number of ranks in current frame
Definition: AMReX_ParallelContext.H:74
@ min
Definition: AMReX_ParallelReduce.H:18
std::enable_if_t< std::is_integral_v< T > > ParallelFor(TypeList< CTOs... > ctos, std::array< int, sizeof...(CTOs)> const &runtime_options, T N, F &&f)
Definition: AMReX_CTOParallelForImpl.H:200
@ make_alias
Definition: AMReX_MakeType.H:7
BoxND< AMREX_SPACEDIM > Box
Definition: AMReX_BaseFwd.H:27
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void ignore_unused(const Ts &...)
This shuts up the compiler about unused variables.
Definition: AMReX.H:111
void Abort(const std::string &msg)
Print out message to cerr and exit via abort().
Definition: AMReX.cpp:225
Definition: AMReX_FFT_Helper.H:56
bool batch_mode
Definition: AMReX_FFT_Helper.H:60
int nprocs
Max number of processes to use.
Definition: AMReX_FFT_Helper.H:63