Line data Source code
1 : //
2 : // Class FFT
3 : // The FFT class performs complex-to-complex,
4 : // real-to-complex on IPPL Fields.
5 : // FFT is templated on the type of transform to be performed,
6 : // the dimensionality of the Field to transform, and the
7 : // floating-point precision type of the Field (float or double).
8 : // Currently, we use heffte for taking the transforms and the class FFT
9 : // serves as an interface between IPPL and heffte. In making this interface,
10 : // we have referred Cabana library.
11 : // https://github.com/ECP-copa/Cabana.
12 : //
13 : // Copyright (c) 2021, Sriramkrishnan Muralikrishnan,
14 : // Paul Scherrer Institut, Villigen PSI, Switzerland
15 : // All rights reserved
16 : //
17 : // This file is part of IPPL.
18 : //
19 : /**
20 : Implementations for FFT constructor/destructor and transforms
21 : */
22 :
23 : #include "Utility/IpplTimings.h"
24 :
25 : #include "Field/BareField.h"
26 :
27 : #include "FieldLayout/FieldLayout.h"
28 :
29 : namespace ippl {
30 :
31 : template <typename Field, template <typename...> class FFT, typename Backend, typename T>
32 4 : FFTBase<Field, FFT, Backend, T>::FFTBase(const Layout_t& layout, const ParameterList& params) {
33 : std::array<long long, 3> low;
34 : std::array<long long, 3> high;
35 :
36 4 : const NDIndex<Dim> lDom = layout.getLocalNDIndex();
37 4 : domainToBounds(lDom, low, high);
38 :
39 4 : heffte::box3d<long long> inbox = {low, high};
40 4 : heffte::box3d<long long> outbox = {low, high};
41 :
42 4 : setup(inbox, outbox, params);
43 4 : }
44 :
45 : template <typename Field, template <typename...> class FFT, typename Backend, typename T>
46 12 : void FFTBase<Field, FFT, Backend, T>::domainToBounds(const NDIndex<Dim>& domain,
47 : std::array<long long, 3>& low,
48 : std::array<long long, 3>& high) {
49 12 : low.fill(0);
50 12 : high.fill(0);
51 :
52 : /**
53 : * Static cast to detail::long long (uint64_t) is necessary, as heffte::box3d requires it
54 : * like that.
55 : */
56 42 : for (size_t d = 0; d < Dim; ++d) {
57 30 : low[d] = static_cast<long long>(domain[d].first());
58 30 : high[d] = static_cast<long long>(domain[d].length() + domain[d].first() - 1);
59 : }
60 12 : }
61 :
62 : /**
63 : setup performs the initialization necessary.
64 : */
65 : template <typename Field, template <typename...> class FFT, typename Backend, typename T>
66 8 : void FFTBase<Field, FFT, Backend, T>::setup(const heffte::box3d<long long>& inbox,
67 : const heffte::box3d<long long>& outbox,
68 : const ParameterList& params) {
69 8 : heffte::plan_options heffteOptions = heffte::default_options<heffteBackend>();
70 :
71 16 : if (!params.get<bool>("use_heffte_defaults")) {
72 0 : heffteOptions.use_pencils = params.get<bool>("use_pencils");
73 0 : heffteOptions.use_reorder = params.get<bool>("use_reorder");
74 : #ifdef Heffte_ENABLE_GPU
75 : heffteOptions.use_gpu_aware = params.get<bool>("use_gpu_aware");
76 : #endif
77 :
78 0 : switch (params.get<int>("comm")) {
79 0 : case a2a:
80 0 : heffteOptions.algorithm = heffte::reshape_algorithm::alltoall;
81 0 : break;
82 0 : case a2av:
83 0 : heffteOptions.algorithm = heffte::reshape_algorithm::alltoallv;
84 0 : break;
85 0 : case p2p:
86 0 : heffteOptions.algorithm = heffte::reshape_algorithm::p2p;
87 0 : break;
88 0 : case p2p_pl:
89 0 : heffteOptions.algorithm = heffte::reshape_algorithm::p2p_plined;
90 0 : break;
91 0 : default:
92 0 : throw IpplException("FFT::setup", "Unrecognized heffte communication type");
93 : }
94 : }
95 :
96 : if constexpr (std::is_same_v<FFT<heffteBackend>, heffte::fft3d<heffteBackend>>) {
97 4 : heffte_m = std::make_shared<FFT<heffteBackend, long long>>(
98 4 : inbox, outbox, Comm->getCommunicator(), heffteOptions);
99 : } else {
100 4 : heffte_m = std::make_shared<FFT<heffteBackend, long long>>(
101 12 : inbox, outbox, params.get<int>("r2c_direction"), Comm->getCommunicator(),
102 : heffteOptions);
103 : }
104 :
105 : // heffte::gpu::device_set(Comm->rank() % heffte::gpu::device_count());
106 8 : if (workspace_m.size() < heffte_m->size_workspace()) {
107 16 : workspace_m = workspace_t(heffte_m->size_workspace());
108 : }
109 8 : }
110 :
111 : template <typename ComplexField>
112 0 : void FFT<CCTransform, ComplexField>::warmup(ComplexField& f) {
113 0 : this->transform(FORWARD, f);
114 0 : this->transform(BACKWARD, f);
115 0 : }
116 :
117 : template <typename ComplexField>
118 8 : void FFT<CCTransform, ComplexField>::transform(TransformDirection direction, ComplexField& f) {
119 : static_assert(Dim == 2 || Dim == 3, "heFFTe only supports 2D and 3D");
120 :
121 8 : auto fview = f.getView();
122 8 : const int nghost = f.getNghost();
123 :
124 : /**
125 : *This copy to a temporary Kokkos view is needed because of following
126 : *reasons:
127 : *1) heffte wants the input and output fields without ghost layers
128 : *2) heffte accepts data in layout left (by default) even though this
129 : *can be changed during heffte box creation
130 : */
131 8 : auto& tempField = this->tempField;
132 8 : if (tempField.size() != f.getOwned().size()) {
133 8 : tempField = detail::shrinkView("tempField", fview, nghost);
134 : }
135 :
136 : using index_array_type = typename RangePolicy<Dim>::index_array_type;
137 8 : ippl::parallel_for(
138 8 : "copy from Kokkos FFT", getRangePolicy(fview, nghost),
139 36880 : KOKKOS_LAMBDA(const index_array_type& args) {
140 18432 : apply(tempField, args - nghost).real(apply(fview, args).real());
141 18432 : apply(tempField, args - nghost).imag(apply(fview, args).imag());
142 : });
143 :
144 8 : if (direction == FORWARD) {
145 4 : this->heffte_m->forward(tempField.data(), tempField.data(), this->workspace_m.data(),
146 : heffte::scale::full);
147 4 : } else if (direction == BACKWARD) {
148 4 : this->heffte_m->backward(tempField.data(), tempField.data(), this->workspace_m.data(),
149 : heffte::scale::none);
150 : } else {
151 0 : throw std::logic_error("Only 1:forward and -1:backward are allowed as directions");
152 : }
153 :
154 8 : ippl::parallel_for(
155 8 : "copy to Kokkos FFT", getRangePolicy(fview, nghost),
156 36880 : KOKKOS_LAMBDA(const index_array_type& args) {
157 18432 : apply(fview, args).real() = apply(tempField, args - nghost).real();
158 18432 : apply(fview, args).imag() = apply(tempField, args - nghost).imag();
159 : });
160 8 : }
161 :
162 : //========================================================================
163 : // FFT RCTransform Constructors
164 : //========================================================================
165 :
166 : /**
167 : *Create a new FFT object of type RCTransform, with given input and output
168 : *layouts and heffte parameters.
169 : */
170 :
171 : template <typename RealField>
172 4 : FFT<RCTransform, RealField>::FFT(const Layout_t& layoutInput, const Layout_t& layoutOutput,
173 4 : const ParameterList& params) {
174 : /**
175 : * Heffte requires to pass a 3D array even for 2D and
176 : * 1D FFTs we just have to make the length in other
177 : * dimensions to be 1.
178 : */
179 : std::array<long long, 3> lowInput;
180 : std::array<long long, 3> highInput;
181 : std::array<long long, 3> lowOutput;
182 : std::array<long long, 3> highOutput;
183 :
184 4 : const NDIndex<Dim>& lDomInput = layoutInput.getLocalNDIndex();
185 4 : const NDIndex<Dim>& lDomOutput = layoutOutput.getLocalNDIndex();
186 :
187 4 : this->domainToBounds(lDomInput, lowInput, highInput);
188 4 : this->domainToBounds(lDomOutput, lowOutput, highOutput);
189 :
190 4 : heffte::box3d<long long> inbox = {lowInput, highInput};
191 4 : heffte::box3d<long long> outbox = {lowOutput, highOutput};
192 :
193 4 : this->setup(inbox, outbox, params);
194 4 : }
195 :
196 : template <typename RealField>
197 0 : void FFT<RCTransform, RealField>::warmup(RealField& f, ComplexField& g) {
198 0 : this->transform(FORWARD, f, g);
199 0 : this->transform(BACKWARD, f, g);
200 0 : }
201 :
202 : template <typename RealField>
203 8 : void FFT<RCTransform, RealField>::transform(TransformDirection direction, RealField& f,
204 : ComplexField& g) {
205 : static_assert(Dim == 2 || Dim == 3, "heFFTe only supports 2D and 3D");
206 :
207 8 : auto fview = f.getView();
208 8 : auto gview = g.getView();
209 8 : const int nghostf = f.getNghost();
210 8 : const int nghostg = g.getNghost();
211 :
212 : /**
213 : *This copy to a temporary Kokkos view is needed because of following
214 : *reasons:
215 : *1) heffte wants the input and output fields without ghost layers
216 : *2) heffte accepts data in layout left (by default) eventhough this
217 : *can be changed during heffte box creation
218 : */
219 8 : auto& tempFieldf = this->tempField;
220 8 : auto& tempFieldg = this->tempFieldComplex;
221 8 : if (tempFieldf.size() != f.getOwned().size()) {
222 8 : tempFieldf = detail::shrinkView("tempFieldf", fview, nghostf);
223 : }
224 8 : if (tempFieldg.size() != g.getOwned().size()) {
225 8 : tempFieldg = detail::shrinkView("tempFieldg", gview, nghostg);
226 : }
227 :
228 : using index_array_type = typename RangePolicy<Dim>::index_array_type;
229 8 : ippl::parallel_for(
230 8 : "copy from Kokkos f field in FFT", getRangePolicy(fview, nghostf),
231 36880 : KOKKOS_LAMBDA(const index_array_type& args) {
232 18432 : apply(tempFieldf, args - nghostf) = apply(fview, args);
233 : });
234 8 : ippl::parallel_for(
235 8 : "copy from Kokkos g field in FFT", getRangePolicy(gview, nghostg),
236 19600 : KOKKOS_LAMBDA(const index_array_type& args) {
237 9792 : apply(tempFieldg, args - nghostg).real(apply(gview, args).real());
238 9792 : apply(tempFieldg, args - nghostg).imag(apply(gview, args).imag());
239 : });
240 :
241 8 : if (direction == FORWARD) {
242 4 : this->heffte_m->forward(tempFieldf.data(), tempFieldg.data(), this->workspace_m.data(),
243 : heffte::scale::full);
244 4 : } else if (direction == BACKWARD) {
245 4 : this->heffte_m->backward(tempFieldg.data(), tempFieldf.data(), this->workspace_m.data(),
246 : heffte::scale::none);
247 : } else {
248 0 : throw std::logic_error("Only 1:forward and -1:backward are allowed as directions");
249 : }
250 :
251 8 : ippl::parallel_for(
252 8 : "copy to Kokkos f field FFT", getRangePolicy(fview, nghostf),
253 36880 : KOKKOS_LAMBDA(const index_array_type& args) {
254 18432 : apply(fview, args) = apply(tempFieldf, args - nghostf);
255 : });
256 :
257 8 : ippl::parallel_for(
258 8 : "copy to Kokkos g field FFT", getRangePolicy(gview, nghostg),
259 19600 : KOKKOS_LAMBDA(const index_array_type& args) {
260 9792 : apply(gview, args).real() = apply(tempFieldg, args - nghostg).real();
261 9792 : apply(gview, args).imag() = apply(tempFieldg, args - nghostg).imag();
262 : });
263 8 : }
264 :
265 : template <typename Field>
266 : void FFT<SineTransform, Field>::warmup(Field& f) {
267 : this->transform(FORWARD, f);
268 : this->transform(BACKWARD, f);
269 : }
270 :
271 : template <typename Field>
272 0 : void FFT<SineTransform, Field>::transform(TransformDirection direction, Field& f) {
273 : static_assert(Dim == 2 || Dim == 3, "heFFTe only supports 2D and 3D");
274 : #ifdef Heffte_ENABLE_FFTW
275 : if (direction == FORWARD) {
276 : f = f / 8.0;
277 : }
278 : #endif
279 :
280 0 : auto fview = f.getView();
281 0 : const int nghost = f.getNghost();
282 :
283 : /**
284 : *This copy to a temporary Kokkos view is needed because of following
285 : *reasons:
286 : *1) heffte wants the input and output fields without ghost layers
287 : *2) heffte accepts data in layout left (by default) eventhough this
288 : *can be changed during heffte box creation
289 : */
290 0 : auto& tempField = this->tempField;
291 0 : if (tempField.size() != f.getOwned().size()) {
292 0 : tempField = detail::shrinkView("tempField", fview, nghost);
293 : }
294 :
295 : using index_array_type = typename RangePolicy<Dim>::index_array_type;
296 0 : ippl::parallel_for(
297 0 : "copy from Kokkos FFT", getRangePolicy(fview, nghost),
298 0 : KOKKOS_LAMBDA(const index_array_type& args) {
299 0 : apply(tempField, args - nghost) = apply(fview, args);
300 : });
301 :
302 0 : if (direction == FORWARD) {
303 0 : this->heffte_m->forward(tempField.data(), tempField.data(), this->workspace_m.data(),
304 : heffte::scale::full);
305 0 : } else if (direction == BACKWARD) {
306 0 : this->heffte_m->backward(tempField.data(), tempField.data(), this->workspace_m.data(),
307 : heffte::scale::none);
308 : } else {
309 0 : throw std::logic_error("Only 1:forward and -1:backward are allowed as directions");
310 : }
311 :
312 0 : ippl::parallel_for(
313 0 : "copy to Kokkos FFT", getRangePolicy(fview, nghost),
314 0 : KOKKOS_LAMBDA(const index_array_type& args) {
315 0 : apply(fview, args) = apply(tempField, args - nghost);
316 : });
317 : #ifdef Heffte_ENABLE_FFTW
318 : if (direction == BACKWARD) {
319 : f = f * 8.0;
320 : }
321 : #endif
322 0 : }
323 :
324 : template <typename Field>
325 : void FFT<CosTransform, Field>::warmup(Field& f) {
326 : this->transform(FORWARD, f);
327 : this->transform(BACKWARD, f);
328 : }
329 :
330 : template <typename Field>
331 0 : void FFT<CosTransform, Field>::transform(TransformDirection direction, Field& f) {
332 : static_assert(Dim == 2 || Dim == 3, "heFFTe only supports 2D and 3D");
333 : #ifdef Heffte_ENABLE_FFTW
334 : if (direction == FORWARD) {
335 : f = f / 8.0;
336 : }
337 : #endif
338 :
339 0 : auto fview = f.getView();
340 0 : const int nghost = f.getNghost();
341 :
342 : /**
343 : *This copy to a temporary Kokkos view is needed because of following
344 : *reasons:
345 : *1) heffte wants the input and output fields without ghost layers
346 : *2) heffte accepts data in layout left (by default) eventhough this
347 : *can be changed during heffte box creation
348 : */
349 0 : auto& tempField = this->tempField;
350 0 : if (tempField.size() != f.getOwned().size()) {
351 0 : tempField = detail::shrinkView("tempField", fview, nghost);
352 : }
353 :
354 : using index_array_type = typename RangePolicy<Dim>::index_array_type;
355 0 : ippl::parallel_for(
356 0 : "copy from Kokkos FFT", getRangePolicy(fview, nghost),
357 0 : KOKKOS_LAMBDA(const index_array_type& args) {
358 0 : apply(tempField, args - nghost) = apply(fview, args);
359 : });
360 :
361 0 : if (direction == FORWARD) {
362 0 : this->heffte_m->forward(tempField.data(), tempField.data(), this->workspace_m.data(),
363 : heffte::scale::full);
364 0 : } else if (direction == BACKWARD) {
365 0 : this->heffte_m->backward(tempField.data(), tempField.data(), this->workspace_m.data(),
366 : heffte::scale::none);
367 : } else {
368 0 : throw std::logic_error("Only 1:forward and -1:backward are allowed as directions");
369 : }
370 :
371 0 : ippl::parallel_for(
372 0 : "copy to Kokkos FFT", getRangePolicy(fview, nghost),
373 0 : KOKKOS_LAMBDA(const index_array_type& args) {
374 0 : apply(fview, args) = apply(tempField, args - nghost);
375 : });
376 : #ifdef Heffte_ENABLE_FFTW
377 : if (direction == BACKWARD) {
378 : f = f * 8.0;
379 : }
380 : #endif
381 0 : }
382 :
383 : template <typename Field>
384 0 : void FFT<Cos1Transform, Field>::warmup(Field& f) {
385 0 : this->transform(FORWARD, f);
386 0 : this->transform(BACKWARD, f);
387 0 : }
388 :
389 : template <typename Field>
390 0 : void FFT<Cos1Transform, Field>::transform(TransformDirection direction, Field& f) {
391 : static_assert(Dim == 2 || Dim == 3, "heFFTe only supports 2D and 3D");
392 :
393 : /**
394 : * This rescaling is needed to match the normalization constant
395 : * between fftw and the other gpu interfaces. fftw rescales with an extra factor of 8.
396 : */
397 : #ifdef Heffte_ENABLE_FFTW
398 : if (direction == FORWARD) {
399 : f = f / 8.0;
400 : }
401 : #endif
402 :
403 0 : auto fview = f.getView();
404 0 : const int nghost = f.getNghost();
405 :
406 : /**
407 : *This copy to a temporary Kokkos view is needed because of following
408 : *reasons:
409 : *1) heffte wants the input and output fields without ghost layers
410 : *2) heffte accepts data in layout left (by default) eventhough this
411 : *can be changed during heffte box creation
412 : */
413 0 : auto& tempField = this->tempField;
414 0 : if (tempField.size() != f.getOwned().size()) {
415 0 : tempField = detail::shrinkView("tempField", fview, nghost);
416 : }
417 :
418 : using index_array_type = typename RangePolicy<Dim>::index_array_type;
419 0 : ippl::parallel_for(
420 0 : "copy from Kokkos FFT", getRangePolicy(fview, nghost),
421 0 : KOKKOS_LAMBDA(const index_array_type& args) {
422 0 : apply(tempField, args - nghost) = apply(fview, args);
423 : });
424 :
425 0 : if (direction == FORWARD) {
426 0 : this->heffte_m->forward(tempField.data(), tempField.data(), this->workspace_m.data(),
427 : heffte::scale::full);
428 0 : } else if (direction == BACKWARD) {
429 0 : this->heffte_m->backward(tempField.data(), tempField.data(), this->workspace_m.data(),
430 : heffte::scale::none);
431 : } else {
432 0 : throw std::logic_error("Only 1:forward and -1:backward are allowed as directions");
433 : }
434 :
435 0 : ippl::parallel_for(
436 0 : "copy to Kokkos FFT", getRangePolicy(fview, nghost),
437 0 : KOKKOS_LAMBDA(const index_array_type& args) {
438 0 : apply(fview, args) = apply(tempField, args - nghost);
439 : });
440 :
441 : /**
442 : * This rescaling is needed to match the normalization constant
443 : * between fftw and the other gpu interfaces. fftw rescales with an extra factor of 8.
444 : */
445 : #ifdef Heffte_ENABLE_FFTW
446 : if (direction == BACKWARD) {
447 : f = f * 8.0;
448 : }
449 : #endif
450 0 : }
451 :
452 : } // namespace ippl
453 :
454 : // vi: set et ts=4 sw=4 sts=4:
455 : // Local Variables:
456 : // mode:c
457 : // c-basic-offset: 4
458 : // indent-tabs-mode: nil
459 : // require-final-newline: nil
460 : // End:
|
