1 #pragma once
2
3 #include <string>
4 #include <stdexcept>
5 #include <sstream>
6 #include <c10/core/ScalarType.h>
7 #include <c10/util/ArrayRef.h>
8 #include <c10/util/Exception.h>
9 #include <ATen/native/DispatchStub.h>
10 #include <ATen/core/TensorBase.h>
11
12 namespace at::native {
13
14 // Normalization types used in _fft_with_size
15 enum class fft_norm_mode {
16 none, // No normalization
17 by_root_n, // Divide by sqrt(signal_size)
18 by_n, // Divide by signal_size
19 };
20
21 // NOTE [ Fourier Transform Conjugate Symmetry ]
22 //
23 // Real-to-complex Fourier transform satisfies the conjugate symmetry. That is,
24 // assuming X is the transformed K-dimensionsal signal, we have
25 //
26 // X[i_1, ..., i_K] = X[j_i, ..., j_K]*,
27 //
28 // where j_k = (N_k - i_k) mod N_k, N_k being the signal size at dim k,
29 // * is the conjugate operator.
30 //
31 // Therefore, in such cases, FFT libraries return only roughly half of the
32 // values to avoid redundancy:
33 //
34 // X[:, :, ..., :floor(N / 2) + 1]
35 //
36 // This is also the assumption in cuFFT and MKL. In ATen SpectralOps, such
37 // halved signal will also be returned by default (flag onesided=True).
38 // The following infer_ft_real_to_complex_onesided_size function calculates the
39 // onesided size from the twosided size.
40 //
41 // Note that this loses some information about the size of signal at last
42 // dimension. E.g., both 11 and 10 maps to 6. Hence, the following
43 // infer_ft_complex_to_real_onesided_size function takes in optional parameter
44 // to infer the twosided size from given onesided size.
45 //
46 // cuFFT doc: http://docs.nvidia.com/cuda/cufft/index.html#multi-dimensional
47 // MKL doc: https://software.intel.com/en-us/mkl-developer-reference-c-dfti-complex-storage-dfti-real-storage-dfti-conjugate-even-storage#CONJUGATE_EVEN_STORAGE
48
infer_ft_real_to_complex_onesided_size(int64_t real_size)49 inline int64_t infer_ft_real_to_complex_onesided_size(int64_t real_size) {
50 return (real_size / 2) + 1;
51 }
52
53 inline int64_t infer_ft_complex_to_real_onesided_size(int64_t complex_size,
54 int64_t expected_size=-1) {
55 int64_t base = (complex_size - 1) * 2;
56 if (expected_size < 0) {
57 return base + 1;
58 } else if (base == expected_size) {
59 return base;
60 } else if (base + 1 == expected_size) {
61 return base + 1;
62 } else {
63 std::ostringstream ss;
64 ss << "expected real signal size " << expected_size << " is incompatible "
65 << "with onesided complex frequency size " << complex_size;
66 AT_ERROR(ss.str());
67 }
68 }
69
70 using fft_fill_with_conjugate_symmetry_fn =
71 void (*)(ScalarType dtype, IntArrayRef mirror_dims, IntArrayRef half_sizes,
72 IntArrayRef in_strides, const void* in_data,
73 IntArrayRef out_strides, void* out_data);
74 DECLARE_DISPATCH(fft_fill_with_conjugate_symmetry_fn, fft_fill_with_conjugate_symmetry_stub);
75
76 // In real-to-complex transform, cuFFT and MKL only fill half of the values
77 // due to conjugate symmetry. This function fills in the other half of the full
78 // fft by using the Hermitian symmetry in the signal.
79 // self should be the shape of the full signal and dims.back() should be the
80 // one-sided dimension.
81 // See NOTE [ Fourier Transform Conjugate Symmetry ]
82 TORCH_API void _fft_fill_with_conjugate_symmetry_(const Tensor& self, IntArrayRef dims);
83
84 } // namespace at::native
85