1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112
|
#include <map>
#include <tuple>
using namespace std;
using u32 = uint32_t;
using key_type = tuple<u32, u32, u32, u32>;
map<key_type, fftw_plan> r2c_fftw_plan_cache;
map<key_type, fftw_plan> c2r_fftw_plan_cache;
// Unconditionally create a new plan for use in r2cfft1d.
//
// This function is expensive because it allows FFTW to run its benchmark code.
// The caller is intended to cache its result so that it can be re-used.
//
// WARNING(mbozzi): The returned plan must not be passed to fftw_execute.
// Use the new-array execute function fftw_execute_dft_r2c instead.
[[nodiscard]] static fftw_plan make_r2c_plan(double* in, fftw_complex* out, u32 ri, u32 ci, u32 ro, u32 co)
{
// Copy input and output buffers because FFTW's measurement process would clobber both.
//
// It's probably best to give FFTW real data to work with when measuring
// If we don't bother, parameters in, out become unnecesssary
int const i_buffer_size = ri * ci;
int const o_buffer_size = ro * co;
auto* i = fftw_alloc_real(i_buffer_size); if (!i) std::abort();
auto* o = fftw_alloc_complex(o_buffer_size); if (!o) std::abort();
static_assert(std::is_trivially_copyable<fftw_complex>::value,
"can memcpy fftw_complex"); // (but not std::copy it)
std::memcpy(i, in, i_buffer_size * sizeof * i);
std::memcpy(o, out, o_buffer_size * sizeof * o);
int const ci_ = ci; // FFTW wants an int* not a u32*.
// https://www.fftw.org/fftw3_doc/Advanced-Complex-DFTs.html
// https://www.fftw.org/fftw3_doc/Advanced-Real_002ddata-DFTs.html
fftw_plan const the_plan = fftw_plan_many_dft_r2c(
1, // compute rank-one transforms.
&ci_, // the rank-one transforms are this long in their first (only) dimension
ri, // do one transform per row
i, // the input data is here
nullptr, // the input data is not a row-major subarray of a larger array
ri, // within this input array the next element is offset ri elements from prior
1, // address of next input array is offset 1 element from start of prior
o, // the output data is here
nullptr, // the output data is not a row-major subarray of a larger array
ro, // within this output array the next element is offset ro elements from prior
1, // address of next output array is offset 1 element from start of prior
FFTW_MEASURE
);
fftw_free(i);
fftw_free(o);
return the_plan;
}
// See make_r2c_plan.
//
// WARNING(mbozzi): The returned plan must not be passed to fftw_execute.
// Use the new-array execute function fftw_execute_dft_c2r instead.
[[nodiscard]] static fftw_plan make_c2r_plan(fftw_complex* in, double* out, u32 ri, u32 ci, u32 ro, u32 co)
{
int const i_buffer_size = ri * ci;
int const o_buffer_size = ro * co;
auto* i = fftw_alloc_complex(i_buffer_size); if (!i) std::abort();
auto* o = fftw_alloc_real(o_buffer_size); if (!o) std::abort();
std::memcpy(i, in, i_buffer_size * sizeof * i);
std::memcpy(o, out, o_buffer_size * sizeof * o);
int const co_ = co;
fftw_plan const the_plan = fftw_plan_many_dft_c2r(
1, &co_, ri, i, nullptr, ri, 1, o, nullptr, ro, 1, FFTW_MEASURE);
fftw_free(i);
fftw_free(o);
return the_plan;
}
[[nodiscard]] static constexpr key_type make_key(u32 ri, u32 ci, u32 ro, u32 co)
{
return make_tuple(ri, ci, ro, co);
}
// Get an existing plan from r2c_fftw_plan_cache or make one with make_r2c_plan
// if it does not exist.
[[nodiscard]] fftw_plan get_r2c_plan(double* i, fftw_complex* o, u32 ri, u32 ci, u32 ro, u32 co)
{
auto const key = make_key(ri, ci, ro, co);
auto const iter = r2c_fftw_plan_cache.find(key);
if (iter == r2c_fftw_plan_cache.end())
r2c_fftw_plan_cache[key] = make_r2c_plan(i, o, ri, ci, ro, co);
return r2c_fftw_plan_cache[key];
}
[[nodiscard]] fftw_plan get_c2r_plan(fftw_complex* i, double* o, u32 ri, u32 ci, u32 ro, u32 co)
{
auto const key = make_key(ri, ci, ro, co);
auto const iter = c2r_fftw_plan_cache.find(key);
if (iter == c2r_fftw_plan_cache.end())
c2r_fftw_plan_cache[key] = make_c2r_plan(i, o, ri, ci, ro, co);
return c2r_fftw_plan_cache[key];
}
| |