Branch data 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 : : //
14 : :
15 : : #ifndef IPPL_FFT_FFT_H
16 : : #define IPPL_FFT_FFT_H
17 : :
18 : : #include <Kokkos_Complex.hpp>
19 : : #include <array>
20 : : #include <heffte_fft3d.h>
21 : : #include <heffte_fft3d_r2c.h>
22 : : #include <memory>
23 : : #include <type_traits>
24 : :
25 : : #include "Utility/IpplException.h"
26 : : #include "Utility/ParameterList.h"
27 : :
28 : : #include "Field/Field.h"
29 : :
30 : : #include "FieldLayout/FieldLayout.h"
31 : : #include "Index/NDIndex.h"
32 : :
33 : : namespace heffte {
34 : : template <>
35 : : struct is_ccomplex<Kokkos::complex<float>> : std::true_type {};
36 : :
37 : : template <>
38 : : struct is_zcomplex<Kokkos::complex<double>> : std::true_type {};
39 : : } // namespace heffte
40 : :
41 : : namespace ippl {
42 : :
43 : : /**
44 : : Tag classes for Fourier transforms
45 : : */
46 : : class CCTransform {};
47 : : class RCTransform {};
48 : : class SineTransform {};
49 : : class CosTransform {};
50 : : /**
51 : : Tag classes for Cosine of type 1 transforms
52 : : */
53 : : class Cos1Transform {};
54 : :
55 : : enum FFTComm {
56 : : a2av = 0,
57 : : a2a = 1,
58 : : p2p = 2,
59 : : p2p_pl = 3
60 : : };
61 : :
62 : : enum TransformDirection {
63 : : FORWARD,
64 : : BACKWARD
65 : : };
66 : :
67 : : namespace detail {
68 : : /*!
69 : : * Wrapper type for heFFTe backends, templated
70 : : * on the Kokkos memory space
71 : : */
72 : : template <typename>
73 : : struct HeffteBackendType;
74 : :
75 : : #if defined(Heffte_ENABLE_FFTW)
76 : : template <>
77 : : struct HeffteBackendType<Kokkos::HostSpace> {
78 : : using backend = heffte::backend::fftw;
79 : : using backendSine = heffte::backend::fftw_sin;
80 : : using backendCos = heffte::backend::fftw_cos;
81 : : using backendCos1 = heffte::backend::fftw_cos1;
82 : : };
83 : : #elif defined(Heffte_ENABLE_MKL)
84 : : template <>
85 : : struct HeffteBackendType<Kokkos::HostSpace> {
86 : : using backend = heffte::backend::mkl;
87 : : using backendSine = heffte::backend::mkl_sin;
88 : : using backendCos = heffte::backend::mkl_cos;
89 : : };
90 : : #endif
91 : :
92 : : #ifdef Heffte_ENABLE_CUDA
93 : : #ifdef KOKKOS_ENABLE_CUDA
94 : : template <>
95 : : struct HeffteBackendType<Kokkos::CudaSpace> {
96 : : using backend = heffte::backend::cufft;
97 : : using backendSine = heffte::backend::cufft_sin;
98 : : using backendCos = heffte::backend::cufft_cos;
99 : : using backendCos1 = heffte::backend::cufft_cos1;
100 : : };
101 : : #else
102 : : #error cuFFT backend is enabled for heFFTe but CUDA is not enabled for Kokkos!
103 : : #endif
104 : : #endif
105 : :
106 : : #ifdef KOKKOS_ENABLE_HIP
107 : : #ifdef Heffte_ENABLE_ROCM
108 : : template <>
109 : : struct HeffteBackendType<Kokkos::HIPSpace> {
110 : : using backend = heffte::backend::rocfft;
111 : : using backendSine = heffte::backend::rocfft_sin;
112 : : using backendCos = heffte::backend::rocfft_cos;
113 : : using backendCos1 = heffte::backend::rocfft_cos1;
114 : : };
115 : : #else
116 : : template <>
117 : : struct HeffteBackendType<Kokkos::HIPSpace> {
118 : : using backend = heffte::backend::stock;
119 : : using backendSine = heffte::backend::stock_sin;
120 : : using backendCos = heffte::backend::stock_cos;
121 : : using backendCos1 = heffte::backend::stock_cos1;
122 : : };
123 : : #endif
124 : : #endif
125 : :
126 : : #if !defined(Heffte_ENABLE_MKL) && !defined(Heffte_ENABLE_FFTW)
127 : : /**
128 : : * Use heFFTe's inbuilt 1D fft computation on CPUs if no
129 : : * vendor specific or optimized backend is found
130 : : */
131 : : template <>
132 : : struct HeffteBackendType<Kokkos::HostSpace> {
133 : : using backend = heffte::backend::stock;
134 : : using backendSine = heffte::backend::stock_sin;
135 : : using backendCos = heffte::backend::stock_cos;
136 : : using backendCos1 = heffte::backend::stock_cos1;
137 : : };
138 : : #endif
139 : :
140 : : } // namespace detail
141 : :
142 : : template <typename Field, template <typename...> class FFT, typename Backend,
143 : : typename BufferType = typename Field::value_type>
144 : : class FFTBase {
145 : : constexpr static unsigned Dim = Field::dim;
146 : :
147 : : public:
148 : : using heffteBackend = Backend;
149 : : using workspace_t = typename FFT<heffteBackend>::template buffer_container<BufferType>;
150 : : using Layout_t = FieldLayout<Dim>;
151 : :
152 : : FFTBase(const Layout_t& layout, const ParameterList& params);
153 : 8 : ~FFTBase() = default;
154 : :
155 : : protected:
156 [ + - ]: 4 : FFTBase() = default;
157 : :
158 : : void domainToBounds(const NDIndex<Dim>& domain, std::array<long long, 3>& low,
159 : : std::array<long long, 3>& high);
160 : : void setup(const heffte::box3d<long long>& inbox, const heffte::box3d<long long>& outbox,
161 : : const ParameterList& params);
162 : :
163 : : std::shared_ptr<FFT<heffteBackend, long long>> heffte_m;
164 : : workspace_t workspace_m;
165 : :
166 : : template <typename FieldType>
167 : : using temp_view_type =
168 : : typename Kokkos::View<typename FieldType::view_type::data_type, Kokkos::LayoutLeft,
169 : : typename FieldType::memory_space>::uniform_type;
170 : : temp_view_type<Field> tempField;
171 : : };
172 : :
173 : : #define IN_PLACE_FFT_BASE_CLASS(Field, Backend) \
174 : : FFTBase<Field, heffte::fft3d, \
175 : : typename detail::HeffteBackendType<typename Field::memory_space>::Backend>
176 : : #define EXT_FFT_BASE_CLASS(Field, Backend, Type) \
177 : : FFTBase<Field, heffte::fft3d_r2c, \
178 : : typename detail::HeffteBackendType<typename Field::memory_space>::Backend, \
179 : : typename Type>
180 : :
181 : : /**
182 : : Non-specialized FFT class. We specialize based on Transform tag class
183 : : */
184 : : template <class Transform, typename Field>
185 : : class FFT {};
186 : :
187 : : /**
188 : : complex-to-complex FFT class
189 : : */
190 : : template <typename ComplexField>
191 : : class FFT<CCTransform, ComplexField> : public IN_PLACE_FFT_BASE_CLASS(ComplexField, backend) {
192 : : constexpr static unsigned Dim = ComplexField::dim;
193 : : using Base = IN_PLACE_FFT_BASE_CLASS(ComplexField, backend);
194 : :
195 : : public:
196 : : using Complex_t = typename ComplexField::value_type;
197 : :
198 : : using Base::Base;
199 : : using typename Base::heffteBackend, typename Base::workspace_t, typename Base::Layout_t;
200 : :
201 : : /*!
202 : : * Warmup the FFT object by forward & backward FFT on an empty field
203 : : * @param f Field whose transformation to compute (and overwrite)
204 : : */
205 : : void warmup(ComplexField& f);
206 : :
207 : : /*!
208 : : * Perform in-place FFT
209 : : * @param direction Forward or backward transformation
210 : : * @param f Field whose transformation to compute (and overwrite)
211 : : */
212 : : void transform(TransformDirection direction, ComplexField& f);
213 : : };
214 : :
215 : : /**
216 : : real-to-complex FFT class
217 : : */
218 : : template <typename RealField>
219 : : class FFT<RCTransform, RealField>
220 : : : public EXT_FFT_BASE_CLASS(RealField, backend,
221 : : Kokkos::complex<typename RealField::value_type>) {
222 : : constexpr static unsigned Dim = RealField::dim;
223 : : using Real_t = typename RealField::value_type;
224 : : using Base = EXT_FFT_BASE_CLASS(RealField, backend,
225 : : Kokkos::complex<typename RealField::value_type>);
226 : :
227 : : public:
228 : : using Complex_t = Kokkos::complex<Real_t>;
229 : : using ComplexField = typename Field<Complex_t, Dim, typename RealField::Mesh_t,
230 : : typename RealField::Centering_t,
231 : : typename RealField::execution_space>::uniform_type;
232 : :
233 : : using typename Base::heffteBackend, typename Base::workspace_t, typename Base::Layout_t;
234 : :
235 : : /** Create a new FFT object with the layout for the input and output Fields
236 : : * and parameters for heffte.
237 : : */
238 : : FFT(const Layout_t& layoutInput, const Layout_t& layoutOutput, const ParameterList& params);
239 : :
240 : : /*!
241 : : * Warmup the FFT object by forward & backward FFT on an empty field
242 : : * @param f Field whose transformation to compute
243 : : * @param g Field in which to store the transformation
244 : : */
245 : : void warmup(RealField& f, ComplexField& g);
246 : :
247 : : /*!
248 : : * Perform FFT
249 : : * @param direction Forward or backward transformation
250 : : * @param f Field whose transformation to compute
251 : : * @param g Field in which to store the transformation
252 : : */
253 : : void transform(TransformDirection direction, RealField& f, ComplexField& g);
254 : :
255 : : private:
256 : : typename Base::template temp_view_type<ComplexField> tempFieldComplex;
257 : : };
258 : :
259 : : /**
260 : : Sine transform class
261 : : */
262 : : template <typename Field>
263 : : class FFT<SineTransform, Field> : public IN_PLACE_FFT_BASE_CLASS(Field, backendSine) {
264 : : constexpr static unsigned Dim = Field::dim;
265 : : using Base = IN_PLACE_FFT_BASE_CLASS(Field, backendSine);
266 : :
267 : : public:
268 : : using Base::Base;
269 : : using typename Base::heffteBackend, typename Base::workspace_t, typename Base::Layout_t;
270 : :
271 : : /*!
272 : : * Warmup the FFT object by forward & backward FFT on an empty field
273 : : * @param f Field whose transformation to compute (and overwrite)
274 : : */
275 : : void warmup(Field& f);
276 : :
277 : : /*!
278 : : * Perform in-place FFT
279 : : * @param direction Forward or backward transformation
280 : : * @param f Field whose transformation to compute (and overwrite)
281 : : */
282 : : void transform(TransformDirection direction, Field& f);
283 : : };
284 : : /**
285 : : Cosine transform class
286 : : */
287 : : template <typename Field>
288 : : class FFT<CosTransform, Field> : public IN_PLACE_FFT_BASE_CLASS(Field, backendCos) {
289 : : constexpr static unsigned Dim = Field::dim;
290 : : using Base = IN_PLACE_FFT_BASE_CLASS(Field, backendCos);
291 : :
292 : : public:
293 : : using Base::Base;
294 : : using typename Base::heffteBackend, typename Base::workspace_t, typename Base::Layout_t;
295 : :
296 : : /*!
297 : : * Warmup the FFT object by forward & backward FFT on an empty field
298 : : * @param f Field whose transformation to compute (and overwrite)
299 : : */
300 : : void warmup(Field& f);
301 : :
302 : : /*!
303 : : * Perform in-place FFT
304 : : * @param direction Forward or backward transformation
305 : : * @param f Field whose transformation to compute (and overwrite)
306 : : */
307 : : void transform(TransformDirection direction, Field& f);
308 : : };
309 : : /**
310 : : Cosine type 1 transform class
311 : : */
312 : : template <typename Field>
313 : : class FFT<Cos1Transform, Field> : public IN_PLACE_FFT_BASE_CLASS(Field, backendCos1) {
314 : : constexpr static unsigned Dim = Field::dim;
315 : : using Base = IN_PLACE_FFT_BASE_CLASS(Field, backendCos1);
316 : :
317 : : public:
318 : : using Base::Base;
319 : : using typename Base::heffteBackend, typename Base::workspace_t, typename Base::Layout_t;
320 : :
321 : : /*!
322 : : * Warmup the FFT object by forward & backward FFT on an empty field
323 : : * @param f Field whose transformation to compute (and overwrite)
324 : : */
325 : : void warmup(Field& f);
326 : :
327 : : /*!
328 : : * Perform in-place FFT
329 : : * @param direction Forward or backward transformation
330 : : * @param f Field whose transformation to compute (and overwrite)
331 : : */
332 : : void transform(TransformDirection direction, Field& f);
333 : : };
334 : : } // namespace ippl
335 : :
336 : : #include "FFT/FFT.hpp"
337 : :
338 : : #endif // IPPL_FFT_FFT_H
339 : :
340 : : // vi: set et ts=4 sw=4 sts=4:
341 : : // Local Variables:
342 : : // mode:c
343 : : // c-basic-offset: 4
344 : : // indent-tabs-mode: nil
345 : : // require-final-newline: nil
346 : : // End:
|
