Function fft_r2c

src/backend/cpu/fft.cpp:112–152

template<typename Tc, typename Tr>
Array<Tc> fft_r2c(const Array<Tr> &in, const int rank) {
    dim4 odims    = in.dims();
    odims[0]      = odims[0] / 2 + 1;
    Array<Tc> out = createEmptyArray<Tc>(odims);

    auto func = [=](Param<Tc> out, const af::dim4 oDataDims, CParam<Tr> in,
                    const af::dim4 iDataDims) {
        af::dim4 idims = in.dims();

        auto t_dims    = computeDims(rank, idims);
        auto in_embed  = computeDims(rank, iDataDims);
        auto out_embed = computeDims(rank, oDataDims);

        const af::dim4 istrides = in.strides();
        const af::dim4 ostrides = out.strides();

        using ctype_t = typename fftw_real_transform<Tc, Tr>::ctype_t;
        using plan_t  = typename fftw_real_transform<Tc, Tr>::plan_t;
        plan_t plan;

        fftw_real_transform<Tc, Tr> transform;

        int batch = 1;
        for (int i = rank; i < 4; i++) { batch *= idims[i]; }

        plan = transform.create(
            rank, t_dims.data(), batch, const_cast<Tr *>(in.get()),
            in_embed.data(), static_cast<int>(istrides[0]),
            static_cast<int>(istrides[rank]),
            reinterpret_cast<ctype_t *>(out.get()), out_embed.data(),
            static_cast<int>(ostrides[0]), static_cast<int>(ostrides[rank]),
            FFTW_ESTIMATE);

        transform.execute(plan);
        transform.destroy(plan);
    };

    getQueue().enqueue(func, out, out.getDataDims(), in, in.getDataDims());

    return out;
}

template<typename Tr, typename Tc>
Array<Tr> fft_c2r(const Array<Tc> &in, const dim4 &odims, const int rank) {