pocketfft_hdronly.h 83.2 KB
Newer Older
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
This file is part of pocketfft.

Copyright (C) 2010-2019 Max-Planck-Society
Author: Martin Reinecke

All rights reserved.

Redistribution and use in source and binary forms, with or without modification,
are permitted provided that the following conditions are met:

* Redistributions of source code must retain the above copyright notice, this
  list of conditions and the following disclaimer.
* Redistributions in binary form must reproduce the above copyright notice, this
  list of conditions and the following disclaimer in the documentation and/or
  other materials provided with the distribution.
* Neither the name of the copyright holder nor the names of its contributors may
  be used to endorse or promote products derived from this software without
  specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/
32 33 34 35

#ifndef POCKETFFT_HDRONLY_H
#define POCKETFFT_HDRONLY_H

36 37 38 39
#ifndef __cplusplus
#error This file is C++ and requires a C++ compiler
#endif

Martin Reinecke's avatar
sync  
Martin Reinecke committed
40
#if !(__cplusplus >= 201103L || _MSVC_LANG+0L >= 201103L)
41 42 43
#error This file requires at least C++11 support
#endif

44 45 46 47 48 49
#include <cmath>
#include <cstring>
#include <cstdlib>
#include <stdexcept>
#include <memory>
#include <vector>
Martin Reinecke's avatar
sync  
Martin Reinecke committed
50
#include <complex>
Martin Reinecke's avatar
Martin Reinecke committed
51 52 53
#if defined(_WIN32)
#include <malloc.h>
#endif
Martin Reinecke's avatar
Martin Reinecke committed
54 55 56 57
#ifdef POCKETFFT_OPENMP
#include <omp.h>
#endif

58

Martin Reinecke's avatar
Martin Reinecke committed
59
#if defined(__GNUC__)
60 61
#define NOINLINE __attribute__((noinline))
#define restrict __restrict__
Martin Reinecke's avatar
Martin Reinecke committed
62 63 64
#elif defined(_MSC_VER)
#define NOINLINE __declspec(noinline)
#define restrict __restrict
65 66 67 68 69 70 71 72 73
#else
#define NOINLINE
#define restrict
#endif

namespace pocketfft {

namespace detail {

Martin Reinecke's avatar
Martin Reinecke committed
74 75
using namespace std;

Martin Reinecke's avatar
Martin Reinecke committed
76 77 78 79 80 81
using shape_t = vector<size_t>;
using stride_t = vector<ptrdiff_t>;

constexpr bool FORWARD  = true,
               BACKWARD = false;

Martin Reinecke's avatar
sync  
Martin Reinecke committed
82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97
// only enable vector support for gcc>=5.0 and clang>=5.0
#ifndef POCKETFFT_NO_VECTORS
#define POCKETFFT_NO_VECTORS
#if defined(__INTEL_COMPILER)
// do nothing. This is necessary because this compiler also sets __GNUC__
#elif defined(__clang__)
#if __clang__>=5
#undef POCKETFFT_NO_VECTORS
#endif
#elif defined(__GNUC__)
#if __GNUC__>=5
#undef POCKETFFT_NO_VECTORS
#endif
#endif
#endif

Martin Reinecke's avatar
Martin Reinecke committed
98 99 100 101 102 103
#ifndef POCKETFFT_NO_VECTORS
#if (defined(__AVX512F__))
constexpr int VBYTELEN=64;
#elif (defined(__AVX__))
constexpr int VBYTELEN=32;
#elif (defined(__SSE2__))
Martin Reinecke's avatar
sync  
Martin Reinecke committed
104 105
constexpr int VBYTELEN=16;
#elif (defined(__VSX__))
Martin Reinecke's avatar
Martin Reinecke committed
106 107 108 109 110 111
constexpr int VBYTELEN=16;
#else
#define POCKETFFT_NO_VECTORS
#endif
#endif

Martin Reinecke's avatar
Martin Reinecke committed
112
template<typename T> class arr
113 114 115 116 117
  {
  private:
    T *p;
    size_t sz;

Martin Reinecke's avatar
Martin Reinecke committed
118
#if defined(POCKETFFT_NO_VECTORS)
119 120 121 122 123
    static T *ralloc(size_t num)
      {
      if (num==0) return nullptr;
      void *res = malloc(num*sizeof(T));
      if (!res) throw bad_alloc();
Martin Reinecke's avatar
Martin Reinecke committed
124 125 126 127 128 129 130 131
      return reinterpret_cast<T *>(res);
      }
    static void dealloc(T *ptr)
      { free(ptr); }
#elif __cplusplus >= 201703L
    static T *ralloc(size_t num)
      {
      if (num==0) return nullptr;
132 133
      void *res = aligned_alloc(64,num*sizeof(T));
      if (!res) throw bad_alloc();
Martin Reinecke's avatar
Martin Reinecke committed
134 135 136 137
      return reinterpret_cast<T *>(res);
      }
    static void dealloc(T *ptr)
      { free(ptr); }
Martin Reinecke's avatar
Martin Reinecke committed
138
#elif defined(_WIN32)
Martin Reinecke's avatar
Martin Reinecke committed
139 140 141
    static T *ralloc(size_t num)
      {
      if (num==0) return nullptr;
Martin Reinecke's avatar
Martin Reinecke committed
142 143
      void *res = _aligned_malloc(num*sizeof(T), 64);
      if (!res) throw bad_alloc();
Martin Reinecke's avatar
Martin Reinecke committed
144 145 146 147
      return reinterpret_cast<T *>(res);
      }
    static void dealloc(T *ptr)
      { _aligned_free(ptr); }
148
#else
Martin Reinecke's avatar
Martin Reinecke committed
149 150 151
    static T *ralloc(size_t num)
      {
      if (num==0) return nullptr;
152 153 154 155 156 157 158
      void *res(nullptr);
      if (posix_memalign(&res, 64, num*sizeof(T))!=0)
        throw bad_alloc();
      return reinterpret_cast<T *>(res);
      }
    static void dealloc(T *ptr)
      { free(ptr); }
Martin Reinecke's avatar
Martin Reinecke committed
159
#endif
160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208

  public:
    arr() : p(0), sz(0) {}
    arr(size_t n) : p(ralloc(n)), sz(n) {}
    arr(arr &&other)
      : p(other.p), sz(other.sz)
      { other.p=nullptr; other.sz=0; }
    ~arr() { dealloc(p); }

    void resize(size_t n)
      {
      if (n==sz) return;
      dealloc(p);
      p = ralloc(n);
      sz = n;
      }

    T &operator[](size_t idx) { return p[idx]; }
    const T &operator[](size_t idx) const { return p[idx]; }

    T *data() { return p; }
    const T *data() const { return p; }

    size_t size() const { return sz; }
  };

template<typename T> struct cmplx {
  T r, i;
  cmplx() {}
  cmplx(T r_, T i_) : r(r_), i(i_) {}
  void Set(T r_, T i_) { r=r_; i=i_; }
  void Set(T r_) { r=r_; i=T(0); }
  cmplx &operator+= (const cmplx &other)
    { r+=other.r; i+=other.i; return *this; }
  template<typename T2>cmplx &operator*= (T2 other)
    { r*=other; i*=other; return *this; }
  cmplx operator+ (const cmplx &other) const
    { return cmplx(r+other.r, i+other.i); }
  cmplx operator- (const cmplx &other) const
    { return cmplx(r-other.r, i-other.i); }
  template<typename T2> auto operator* (const T2 &other) const
    -> cmplx<decltype(r*other)>
    { return {r*other, i*other}; }
  template<typename T2> auto operator* (const cmplx<T2> &other) const
    -> cmplx<decltype(r+other.r)>
    { return {r*other.r-i*other.i, r*other.i + i*other.r}; }
  template<bool bwd, typename T2> auto special_mul (const cmplx<T2> &other) const
    -> cmplx<decltype(r+other.r)>
    {
Martin Reinecke's avatar
sync  
Martin Reinecke committed
209 210 211
    using Tres = cmplx<decltype(r+other.r)>;
    return bwd ? Tres(r*other.r-i*other.i, r*other.i + i*other.r)
               : Tres(r*other.r+i*other.i, i*other.r - r*other.i);
212 213 214 215 216 217 218 219 220 221
    }
};
template<typename T> void PMC(cmplx<T> &a, cmplx<T> &b,
  const cmplx<T> &c, const cmplx<T> &d)
  { a = c+d; b = c-d; }
template<typename T> cmplx<T> conj(const cmplx<T> &a)
  { return {a.r, -a.i}; }

template<typename T> void ROT90(cmplx<T> &a)
  { auto tmp_=a.r; a.r=-a.i; a.i=tmp_; }
Martin Reinecke's avatar
sync  
Martin Reinecke committed
222 223
template<bool bwd, typename T> void ROTX90(cmplx<T> &a)
  { auto tmp_= bwd ? a.r : -a.r; a.r = bwd ? -a.i : a.i; a.i=tmp_; }
224 225 226 227 228 229 230 231

//
// twiddle factor section
//

template<typename T> class sincos_2pibyn
  {
  private:
Martin Reinecke's avatar
sync  
Martin Reinecke committed
232 233 234 235 236 237 238 239
    template<typename Ta, typename Tb, bool bigger> struct TypeSelector
      {};
    template<typename Ta, typename Tb> struct TypeSelector<Ta, Tb, true>
      { using type = Ta; };
    template<typename Ta, typename Tb> struct TypeSelector<Ta, Tb, false>
      { using type = Tb; };

    using Thigh = typename TypeSelector<T, double, (sizeof(T)>sizeof(double))>::type;
240 241
    arr<T> data;

Martin Reinecke's avatar
sync  
Martin Reinecke committed
242
    void my_sincosm1pi (Thigh a_, Thigh *restrict res)
243
      {
Martin Reinecke's avatar
sync  
Martin Reinecke committed
244 245
      if (sizeof(Thigh)>sizeof(double)) // don't have the code for long double
        {
Martin Reinecke's avatar
sync  
Martin Reinecke committed
246 247 248
        constexpr Thigh pi = Thigh(3.141592653589793238462643383279502884197L);
        auto s = sin(pi*a_);
        res[1] = s;
Martin Reinecke's avatar
sync  
Martin Reinecke committed
249 250 251
        res[0] = (s*s)/(-sqrt((1-s)*(1+s))-1);
        return;
        }
Martin Reinecke's avatar
sync  
Martin Reinecke committed
252 253 254
      // adapted from https://stackoverflow.com/questions/42792939/
      // CAUTION: this function only works for arguments in the range
      //          [-0.25; 0.25]!
Martin Reinecke's avatar
sync  
Martin Reinecke committed
255
      double a = double(a_);
256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285
      double s = a * a;
      /* Approximate cos(pi*x)-1 for x in [-0.25,0.25] */
      double r =     -1.0369917389758117e-4;
      r = fma (r, s,  1.9294935641298806e-3);
      r = fma (r, s, -2.5806887942825395e-2);
      r = fma (r, s,  2.3533063028328211e-1);
      r = fma (r, s, -1.3352627688538006e+0);
      r = fma (r, s,  4.0587121264167623e+0);
      r = fma (r, s, -4.9348022005446790e+0);
      double c = r*s;
      /* Approximate sin(pi*x) for x in [-0.25,0.25] */
      r =             4.6151442520157035e-4;
      r = fma (r, s, -7.3700183130883555e-3);
      r = fma (r, s,  8.2145868949323936e-2);
      r = fma (r, s, -5.9926452893214921e-1);
      r = fma (r, s,  2.5501640398732688e+0);
      r = fma (r, s, -5.1677127800499516e+0);
      s = s * a;
      r = r * s;
      s = fma (a, 3.1415926535897931e+0, r);
      res[0] = c;
      res[1] = s;
      }

    NOINLINE void calc_first_octant(size_t den, T * restrict res)
      {
      size_t n = (den+4)>>3;
      if (n==0) return;
      res[0]=1.; res[1]=0.;
      if (n==1) return;
Martin Reinecke's avatar
Martin Reinecke committed
286
      size_t l1 = size_t(sqrt(n));
Martin Reinecke's avatar
sync  
Martin Reinecke committed
287
      arr<Thigh> tmp(2*l1);
288 289
      for (size_t i=1; i<l1; ++i)
        {
Martin Reinecke's avatar
sync  
Martin Reinecke committed
290 291 292
        my_sincosm1pi(Thigh(2*i)/Thigh(den),&tmp[2*i]);
        res[2*i  ] = T(tmp[2*i]+1);
        res[2*i+1] = T(tmp[2*i+1]);
293 294 295 296
        }
      size_t start=l1;
      while(start<n)
        {
Martin Reinecke's avatar
sync  
Martin Reinecke committed
297
        Thigh cs[2];
Martin Reinecke's avatar
sync  
Martin Reinecke committed
298 299 300
        my_sincosm1pi((Thigh(2*start))/Thigh(den),cs);
        res[2*start] = T(cs[0]+1);
        res[2*start+1] = T(cs[1]);
301 302 303 304
        size_t end = l1;
        if (start+end>n) end = n-start;
        for (size_t i=1; i<end; ++i)
          {
Martin Reinecke's avatar
sync  
Martin Reinecke committed
305
          Thigh csx[2]={tmp[2*i], tmp[2*i+1]};
Martin Reinecke's avatar
sync  
Martin Reinecke committed
306 307
          res[2*(start+i)] = T(((cs[0]*csx[0] - cs[1]*csx[1] + cs[0]) + csx[0]) + 1);
          res[2*(start+i)+1] = T((cs[0]*csx[1] + cs[1]*csx[0]) + cs[1] + csx[1]);
308 309 310 311 312 313 314 315 316 317 318 319 320
          }
        start += l1;
        }
      }

    void calc_first_quadrant(size_t n, T * restrict res)
      {
      T * restrict p = res+n;
      calc_first_octant(n<<1, p);
      size_t ndone=(n+2)>>2;
      size_t i=0, idx1=0, idx2=2*ndone-2;
      for (; i+1<ndone; i+=2, idx1+=2, idx2-=2)
        {
Martin Reinecke's avatar
Martin Reinecke committed
321 322
        res[idx1] = p[2*i  ]; res[idx1+1] = p[2*i+1];
        res[idx2] = p[2*i+3]; res[idx2+1] = p[2*i+2];
323 324
        }
      if (i!=ndone)
Martin Reinecke's avatar
Martin Reinecke committed
325
        { res[idx1] = p[2*i]; res[idx1+1] = p[2*i+1]; }
326 327 328 329
      }

    void calc_first_half(size_t n, T * restrict res)
      {
Martin Reinecke's avatar
sync  
Martin Reinecke committed
330
      int ndone=int(n+1)>>1;
331 332
      T * p = res+n-1;
      calc_first_octant(n<<2, p);
Martin Reinecke's avatar
sync  
Martin Reinecke committed
333
      int i4=0, in=int(n), i=0;
334
      for (; i4<=in-i4; ++i, i4+=4) // octant 0
Martin Reinecke's avatar
Martin Reinecke committed
335
        { res[2*i] = p[2*i4]; res[2*i+1] = p[2*i4+1]; }
336
      for (; i4-in <= 0; ++i, i4+=4) // octant 1
Martin Reinecke's avatar
sync  
Martin Reinecke committed
337
        { auto xm = in-i4; res[2*i] = p[2*xm+1]; res[2*i+1] = p[2*xm]; }
338
      for (; i4<=3*in-i4; ++i, i4+=4) // octant 2
Martin Reinecke's avatar
sync  
Martin Reinecke committed
339
        { auto xm = i4-in; res[2*i] = -p[2*xm+1]; res[2*i+1] = p[2*xm]; }
340
      for (; i<ndone; ++i, i4+=4) // octant 3
Martin Reinecke's avatar
sync  
Martin Reinecke committed
341
        { auto xm = 2*in-i4; res[2*i] = -p[2*xm]; res[2*i+1] = p[2*xm+1]; }
342 343 344 345
      }

    void fill_first_quadrant(size_t n, T * restrict res)
      {
Martin Reinecke's avatar
sync  
Martin Reinecke committed
346
      constexpr T hsqt2 = T(0.707106781186547524400844362104849L);
347 348 349 350
      size_t quart = n>>2;
      if ((n&7)==0)
        res[quart] = res[quart+1] = hsqt2;
      for (size_t i=2, j=2*quart-2; i<quart; i+=2, j-=2)
Martin Reinecke's avatar
Martin Reinecke committed
351
        { res[j] = res[i+1]; res[j+1] = res[i]; }
352 353 354 355 356 357 358
      }

    NOINLINE void fill_first_half(size_t n, T * restrict res)
      {
      size_t half = n>>1;
      if ((n&3)==0)
        for (size_t i=0; i<half; i+=2)
Martin Reinecke's avatar
Martin Reinecke committed
359
          { res[i+half] = -res[i+1]; res[i+half+1] = res[i]; }
360 361
      else
        for (size_t i=2, j=2*half-2; i<half; i+=2, j-=2)
Martin Reinecke's avatar
Martin Reinecke committed
362
          { res[j] = -res[i]; res[j+1] = res[i+1]; }
363 364 365 366 367 368 369 370 371
      }

    void fill_second_half(size_t n, T * restrict res)
      {
      if ((n&1)==0)
        for (size_t i=0; i<n; ++i)
          res[i+n] = -res[i];
      else
        for (size_t i=2, j=2*n-2; i<n; i+=2, j-=2)
Martin Reinecke's avatar
Martin Reinecke committed
372
          { res[j] = res[i]; res[j+1] = -res[i+1]; }
373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405
      }

    NOINLINE void sincos_2pibyn_half(size_t n, T * restrict res)
      {
      if ((n&3)==0)
        {
        calc_first_octant(n, res);
        fill_first_quadrant(n, res);
        fill_first_half(n, res);
        }
      else if ((n&1)==0)
        {
        calc_first_quadrant(n, res);
        fill_first_half(n, res);
        }
      else
        calc_first_half(n, res);
      }

  public:
    NOINLINE sincos_2pibyn(size_t n, bool half)
      : data(2*n)
      {
      sincos_2pibyn_half(n, data.data());
      if (!half) fill_second_half(n, data.data());
      }

    T operator[](size_t idx) const { return data[idx]; }
    const T *rdata() const { return data; }
    const cmplx<T> *cdata() const
      { return reinterpret_cast<const cmplx<T> *>(data.data()); }
  };

Martin Reinecke's avatar
Martin Reinecke committed
406
struct util // hack to avoid duplicate symbols
407
  {
Martin Reinecke's avatar
Martin Reinecke committed
408
  static NOINLINE size_t largest_prime_factor (size_t n)
409
    {
Martin Reinecke's avatar
Martin Reinecke committed
410 411 412
    size_t res=1;
    while ((n&1)==0)
      { res=2; n>>=1; }
413

Martin Reinecke's avatar
sync  
Martin Reinecke committed
414
    size_t limit=size_t(sqrt(double(n)+0.01));
Martin Reinecke's avatar
Martin Reinecke committed
415 416 417 418 419
    for (size_t x=3; x<=limit; x+=2)
    while ((n/x)*x==n)
      {
      res=x;
      n/=x;
Martin Reinecke's avatar
sync  
Martin Reinecke committed
420
      limit=size_t(sqrt(double(n)+0.01));
Martin Reinecke's avatar
Martin Reinecke committed
421 422
      }
    if (n>1) res=n;
423

Martin Reinecke's avatar
Martin Reinecke committed
424 425 426 427
    return res;
    }

  static NOINLINE double cost_guess (size_t n)
428
    {
Martin Reinecke's avatar
sync  
Martin Reinecke committed
429
    constexpr double lfp=1.1; // penalty for non-hardcoded larger factors
Martin Reinecke's avatar
Martin Reinecke committed
430 431 432 433 434
    size_t ni=n;
    double result=0.;
    while ((n&1)==0)
      { result+=2; n>>=1; }

Martin Reinecke's avatar
sync  
Martin Reinecke committed
435
    size_t limit=size_t(sqrt(double(n)+0.01));
Martin Reinecke's avatar
Martin Reinecke committed
436 437 438
    for (size_t x=3; x<=limit; x+=2)
    while ((n/x)*x==n)
      {
Martin Reinecke's avatar
sync  
Martin Reinecke committed
439
      result+= (x<=5) ? double(x) : lfp*double(x); // penalize larger prime factors
Martin Reinecke's avatar
Martin Reinecke committed
440
      n/=x;
Martin Reinecke's avatar
sync  
Martin Reinecke committed
441
      limit=size_t(sqrt(double(n)+0.01));
Martin Reinecke's avatar
Martin Reinecke committed
442
      }
Martin Reinecke's avatar
sync  
Martin Reinecke committed
443
    if (n>1) result+=(n<=5) ? double(n) : lfp*double(n);
Martin Reinecke's avatar
Martin Reinecke committed
444

Martin Reinecke's avatar
sync  
Martin Reinecke committed
445
    return result*double(ni);
446 447
    }

Martin Reinecke's avatar
Martin Reinecke committed
448 449 450 451 452 453 454 455 456 457 458 459 460 461
  /* returns the smallest composite of 2, 3, 5, 7 and 11 which is >= n */
  static NOINLINE size_t good_size(size_t n)
    {
    if (n<=12) return n;

    size_t bestfac=2*n;
    for (size_t f2=1; f2<bestfac; f2*=2)
      for (size_t f23=f2; f23<bestfac; f23*=3)
        for (size_t f235=f23; f235<bestfac; f235*=5)
          for (size_t f2357=f235; f2357<bestfac; f2357*=7)
            for (size_t f235711=f2357; f235711<bestfac; f235711*=11)
              if (f235711>=n) bestfac=f235711;
    return bestfac;
    }
462

Martin Reinecke's avatar
Martin Reinecke committed
463 464 465 466 467 468 469
  static size_t prod(const shape_t &shape)
    {
    size_t res=1;
    for (auto sz: shape)
      res*=sz;
    return res;
    }
Martin Reinecke's avatar
Martin Reinecke committed
470 471 472 473 474 475 476 477

  static NOINLINE void sanity_check(const shape_t &shape,
    const stride_t &stride_in, const stride_t &stride_out, bool inplace)
    {
    auto ndim = shape.size();
    if (ndim<1) throw runtime_error("ndim must be >= 1");
    if ((stride_in.size()!=ndim) || (stride_out.size()!=ndim))
      throw runtime_error("stride dimension mismatch");
Martin Reinecke's avatar
sync  
Martin Reinecke committed
478 479
    if (inplace && (stride_in!=stride_out))
      throw runtime_error("stride mismatch");
Martin Reinecke's avatar
Martin Reinecke committed
480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502
    }

  static NOINLINE void sanity_check(const shape_t &shape,
    const stride_t &stride_in, const stride_t &stride_out, bool inplace,
    const shape_t &axes)
    {
    sanity_check(shape, stride_in, stride_out, inplace);
    auto ndim = shape.size();
    shape_t tmp(ndim,0);
    for (auto ax : axes)
      {
      if (ax>=ndim) throw runtime_error("bad axis number");
      if (++tmp[ax]>1) throw runtime_error("axis specified repeatedly");
      }
    }

  static NOINLINE void sanity_check(const shape_t &shape,
    const stride_t &stride_in, const stride_t &stride_out, bool inplace,
    size_t axis)
    {
    sanity_check(shape, stride_in, stride_out, inplace);
    if (axis>=shape.size()) throw runtime_error("bad axis number");
    }
Martin Reinecke's avatar
Martin Reinecke committed
503 504

#ifdef POCKETFFT_OPENMP
Martin Reinecke's avatar
sync  
Martin Reinecke committed
505 506
    static size_t nthreads() { return size_t(omp_get_num_threads()); }
    static size_t thread_num() { return size_t(omp_get_thread_num()); }
Martin Reinecke's avatar
Martin Reinecke committed
507 508 509 510 511
    static size_t thread_count (size_t nthreads, const shape_t &shape,
      size_t axis)
      {
      if (nthreads==1) return 1;
      if (prod(shape)/shape[axis] < 20) return 1;
Martin Reinecke's avatar
sync  
Martin Reinecke committed
512
      return (nthreads==0) ? size_t(omp_get_max_threads()) : nthreads;
Martin Reinecke's avatar
Martin Reinecke committed
513
      }
Martin Reinecke's avatar
Martin Reinecke committed
514
#else
Martin Reinecke's avatar
Martin Reinecke committed
515 516
    static size_t nthreads() { return 1; }
    static size_t thread_num() { return 0; }
Martin Reinecke's avatar
Martin Reinecke committed
517
#endif
Martin Reinecke's avatar
Martin Reinecke committed
518
  };
519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536

#define CH(a,b,c) ch[(a)+ido*((b)+l1*(c))]
#define CC(a,b,c) cc[(a)+ido*((b)+cdim*(c))]
#define WA(x,i) wa[(i)-1+(x)*(ido-1)]

//
// complex FFTPACK transforms
//

template<typename T0> class cfftp
  {
  private:
    struct fctdata
      {
      size_t fct;
      cmplx<T0> *tw, *tws;
      };

Martin Reinecke's avatar
Martin Reinecke committed
537
    size_t length;
538
    arr<cmplx<T0>> mem;
Martin Reinecke's avatar
Martin Reinecke committed
539
    vector<fctdata> fact;
540 541

    void add_factor(size_t factor)
Martin Reinecke's avatar
Martin Reinecke committed
542
      { fact.push_back({factor, nullptr, nullptr}); }
543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591

template<bool bwd, typename T> void pass2 (size_t ido, size_t l1,
  const T * restrict cc, T * restrict ch, const cmplx<T0> * restrict wa)
  {
  constexpr size_t cdim=2;

  if (ido==1)
    for (size_t k=0; k<l1; ++k)
      {
      CH(0,k,0) = CC(0,0,k)+CC(0,1,k);
      CH(0,k,1) = CC(0,0,k)-CC(0,1,k);
      }
  else
    for (size_t k=0; k<l1; ++k)
      {
      CH(0,k,0) = CC(0,0,k)+CC(0,1,k);
      CH(0,k,1) = CC(0,0,k)-CC(0,1,k);
      for (size_t i=1; i<ido; ++i)
        {
        CH(i,k,0) = CC(i,0,k)+CC(i,1,k);
        CH(i,k,1) = (CC(i,0,k)-CC(i,1,k)).template special_mul<bwd>(WA(0,i));
        }
      }
  }

#define PREP3(idx) \
        T t0 = CC(idx,0,k), t1, t2; \
        PMC (t1,t2,CC(idx,1,k),CC(idx,2,k)); \
        CH(idx,k,0)=t0+t1;
#define PARTSTEP3a(u1,u2,twr,twi) \
        { \
        T ca,cb; \
        ca=t0+t1*twr; \
        cb=t2*twi; ROT90(cb); \
        PMC(CH(0,k,u1),CH(0,k,u2),ca,cb) ;\
        }
#define PARTSTEP3b(u1,u2,twr,twi) \
        { \
        T ca,cb,da,db; \
        ca=t0+t1*twr; \
        cb=t2*twi; ROT90(cb); \
        PMC(da,db,ca,cb); \
        CH(i,k,u1) = da.template special_mul<bwd>(WA(u1-1,i)); \
        CH(i,k,u2) = db.template special_mul<bwd>(WA(u2-1,i)); \
        }
template<bool bwd, typename T> void pass3 (size_t ido, size_t l1,
  const T * restrict cc, T * restrict ch, const cmplx<T0> * restrict wa)
  {
  constexpr size_t cdim=3;
Martin Reinecke's avatar
sync  
Martin Reinecke committed
592 593
  constexpr T0 tw1r=-0.5,
               tw1i= (bwd ? 1: -1) * T0(0.8660254037844386467637231707529362L);
594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630

  if (ido==1)
    for (size_t k=0; k<l1; ++k)
      {
      PREP3(0)
      PARTSTEP3a(1,2,tw1r,tw1i)
      }
  else
    for (size_t k=0; k<l1; ++k)
      {
      {
      PREP3(0)
      PARTSTEP3a(1,2,tw1r,tw1i)
      }
      for (size_t i=1; i<ido; ++i)
        {
        PREP3(i)
        PARTSTEP3b(1,2,tw1r,tw1i)
        }
      }
  }

#undef PARTSTEP3b
#undef PARTSTEP3a
#undef PREP3

template<bool bwd, typename T> void pass4 (size_t ido, size_t l1,
  const T * restrict cc, T * restrict ch, const cmplx<T0> * restrict wa)
  {
  constexpr size_t cdim=4;

  if (ido==1)
    for (size_t k=0; k<l1; ++k)
      {
      T t1, t2, t3, t4;
      PMC(t2,t1,CC(0,0,k),CC(0,2,k));
      PMC(t3,t4,CC(0,1,k),CC(0,3,k));
Martin Reinecke's avatar
sync  
Martin Reinecke committed
631
      ROTX90<bwd>(t4);
632 633 634 635 636 637 638 639 640 641
      PMC(CH(0,k,0),CH(0,k,2),t2,t3);
      PMC(CH(0,k,1),CH(0,k,3),t1,t4);
      }
  else
    for (size_t k=0; k<l1; ++k)
      {
      {
      T t1, t2, t3, t4;
      PMC(t2,t1,CC(0,0,k),CC(0,2,k));
      PMC(t3,t4,CC(0,1,k),CC(0,3,k));
Martin Reinecke's avatar
sync  
Martin Reinecke committed
642
      ROTX90<bwd>(t4);
643 644 645 646 647 648 649 650 651
      PMC(CH(0,k,0),CH(0,k,2),t2,t3);
      PMC(CH(0,k,1),CH(0,k,3),t1,t4);
      }
      for (size_t i=1; i<ido; ++i)
        {
        T c2, c3, c4, t1, t2, t3, t4;
        T cc0=CC(i,0,k), cc1=CC(i,1,k),cc2=CC(i,2,k),cc3=CC(i,3,k);
        PMC(t2,t1,cc0,cc2);
        PMC(t3,t4,cc1,cc3);
Martin Reinecke's avatar
sync  
Martin Reinecke committed
652
        ROTX90<bwd>(t4);
653 654
        PMC(CH(i,k,0),c3,t2,t3);
        PMC(c2,c4,t1,t4);
Martin Reinecke's avatar
sync  
Martin Reinecke committed
655 656 657
        CH(i,k,1) = c2.template special_mul<bwd>(WA(0,i));
        CH(i,k,2) = c3.template special_mul<bwd>(WA(1,i));
        CH(i,k,3) = c4.template special_mul<bwd>(WA(2,i));
658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693
        }
      }
  }

#define PREP5(idx) \
        T t0 = CC(idx,0,k), t1, t2, t3, t4; \
        PMC (t1,t4,CC(idx,1,k),CC(idx,4,k)); \
        PMC (t2,t3,CC(idx,2,k),CC(idx,3,k)); \
        CH(idx,k,0).r=t0.r+t1.r+t2.r; \
        CH(idx,k,0).i=t0.i+t1.i+t2.i;

#define PARTSTEP5a(u1,u2,twar,twbr,twai,twbi) \
        { \
        T ca,cb; \
        ca.r=t0.r+twar*t1.r+twbr*t2.r; \
        ca.i=t0.i+twar*t1.i+twbr*t2.i; \
        cb.i=twai*t4.r twbi*t3.r; \
        cb.r=-(twai*t4.i twbi*t3.i); \
        PMC(CH(0,k,u1),CH(0,k,u2),ca,cb); \
        }

#define PARTSTEP5b(u1,u2,twar,twbr,twai,twbi) \
        { \
        T ca,cb,da,db; \
        ca.r=t0.r+twar*t1.r+twbr*t2.r; \
        ca.i=t0.i+twar*t1.i+twbr*t2.i; \
        cb.i=twai*t4.r twbi*t3.r; \
        cb.r=-(twai*t4.i twbi*t3.i); \
        PMC(da,db,ca,cb); \
        CH(i,k,u1) = da.template special_mul<bwd>(WA(u1-1,i)); \
        CH(i,k,u2) = db.template special_mul<bwd>(WA(u2-1,i)); \
        }
template<bool bwd, typename T> void pass5 (size_t ido, size_t l1,
  const T * restrict cc, T * restrict ch, const cmplx<T0> * restrict wa)
  {
  constexpr size_t cdim=5;
Martin Reinecke's avatar
sync  
Martin Reinecke committed
694 695 696 697
  constexpr T0 tw1r= T0(0.3090169943749474241022934171828191L),
               tw1i= (bwd ? 1: -1) * T0(0.9510565162951535721164393333793821L),
               tw2r= T0(-0.8090169943749474241022934171828191L),
               tw2i= (bwd ? 1: -1) * T0(0.5877852522924731291687059546390728L);
698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757

  if (ido==1)
    for (size_t k=0; k<l1; ++k)
      {
      PREP5(0)
      PARTSTEP5a(1,4,tw1r,tw2r,+tw1i,+tw2i)
      PARTSTEP5a(2,3,tw2r,tw1r,+tw2i,-tw1i)
      }
  else
    for (size_t k=0; k<l1; ++k)
      {
      {
      PREP5(0)
      PARTSTEP5a(1,4,tw1r,tw2r,+tw1i,+tw2i)
      PARTSTEP5a(2,3,tw2r,tw1r,+tw2i,-tw1i)
      }
      for (size_t i=1; i<ido; ++i)
        {
        PREP5(i)
        PARTSTEP5b(1,4,tw1r,tw2r,+tw1i,+tw2i)
        PARTSTEP5b(2,3,tw2r,tw1r,+tw2i,-tw1i)
        }
      }
  }

#undef PARTSTEP5b
#undef PARTSTEP5a
#undef PREP5

#define PREP7(idx) \
        T t1 = CC(idx,0,k), t2, t3, t4, t5, t6, t7; \
        PMC (t2,t7,CC(idx,1,k),CC(idx,6,k)); \
        PMC (t3,t6,CC(idx,2,k),CC(idx,5,k)); \
        PMC (t4,t5,CC(idx,3,k),CC(idx,4,k)); \
        CH(idx,k,0).r=t1.r+t2.r+t3.r+t4.r; \
        CH(idx,k,0).i=t1.i+t2.i+t3.i+t4.i;

#define PARTSTEP7a0(u1,u2,x1,x2,x3,y1,y2,y3,out1,out2) \
        { \
        T ca,cb; \
        ca.r=t1.r+x1*t2.r+x2*t3.r+x3*t4.r; \
        ca.i=t1.i+x1*t2.i+x2*t3.i+x3*t4.i; \
        cb.i=y1*t7.r y2*t6.r y3*t5.r; \
        cb.r=-(y1*t7.i y2*t6.i y3*t5.i); \
        PMC(out1,out2,ca,cb); \
        }
#define PARTSTEP7a(u1,u2,x1,x2,x3,y1,y2,y3) \
        PARTSTEP7a0(u1,u2,x1,x2,x3,y1,y2,y3,CH(0,k,u1),CH(0,k,u2))
#define PARTSTEP7(u1,u2,x1,x2,x3,y1,y2,y3) \
        { \
        T da,db; \
        PARTSTEP7a0(u1,u2,x1,x2,x3,y1,y2,y3,da,db) \
        CH(i,k,u1) = da.template special_mul<bwd>(WA(u1-1,i)); \
        CH(i,k,u2) = db.template special_mul<bwd>(WA(u2-1,i)); \
        }

template<bool bwd, typename T> void pass7(size_t ido, size_t l1,
  const T * restrict cc, T * restrict ch, const cmplx<T0> * restrict wa)
  {
  constexpr size_t cdim=7;
Martin Reinecke's avatar
sync  
Martin Reinecke committed
758 759 760 761 762 763
  constexpr T0 tw1r= T0(0.6234898018587335305250048840042398L),
               tw1i= (bwd ? 1 : -1) * T0(0.7818314824680298087084445266740578L),
               tw2r= T0(-0.2225209339563144042889025644967948L),
               tw2i= (bwd ? 1 : -1) * T0(0.9749279121818236070181316829939312L),
               tw3r= T0(-0.9009688679024191262361023195074451L),
               tw3i= (bwd ? 1 : -1) * T0(0.433883739117558120475768332848359L);
764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796

  if (ido==1)
    for (size_t k=0; k<l1; ++k)
      {
      PREP7(0)
      PARTSTEP7a(1,6,tw1r,tw2r,tw3r,+tw1i,+tw2i,+tw3i)
      PARTSTEP7a(2,5,tw2r,tw3r,tw1r,+tw2i,-tw3i,-tw1i)
      PARTSTEP7a(3,4,tw3r,tw1r,tw2r,+tw3i,-tw1i,+tw2i)
      }
  else
    for (size_t k=0; k<l1; ++k)
      {
      {
      PREP7(0)
      PARTSTEP7a(1,6,tw1r,tw2r,tw3r,+tw1i,+tw2i,+tw3i)
      PARTSTEP7a(2,5,tw2r,tw3r,tw1r,+tw2i,-tw3i,-tw1i)
      PARTSTEP7a(3,4,tw3r,tw1r,tw2r,+tw3i,-tw1i,+tw2i)
      }
      for (size_t i=1; i<ido; ++i)
        {
        PREP7(i)
        PARTSTEP7(1,6,tw1r,tw2r,tw3r,+tw1i,+tw2i,+tw3i)
        PARTSTEP7(2,5,tw2r,tw3r,tw1r,+tw2i,-tw3i,-tw1i)
        PARTSTEP7(3,4,tw3r,tw1r,tw2r,+tw3i,-tw1i,+tw2i)
        }
      }
  }

#undef PARTSTEP7
#undef PARTSTEP7a0
#undef PARTSTEP7a
#undef PREP7

Martin Reinecke's avatar
sync  
Martin Reinecke committed
797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897
template <bool bwd, typename T> void ROTX45(T &a)
  {
  constexpr T0 hsqt2=T0(0.707106781186547524400844362104849L);
  if (bwd)
    { auto tmp_=a.r; a.r=hsqt2*(a.r-a.i); a.i=hsqt2*(a.i+tmp_); }
  else
    { auto tmp_=a.r; a.r=hsqt2*(a.r+a.i); a.i=hsqt2*(a.i-tmp_); }
  }
template <bool bwd, typename T> void ROTX135(T &a)
  {
  constexpr T0 hsqt2=T0(0.707106781186547524400844362104849L);
  if (bwd)
    { auto tmp_=a.r; a.r=hsqt2*(-a.r-a.i); a.i=hsqt2*(tmp_-a.i); }
  else
    { auto tmp_=a.r; a.r=hsqt2*(a.i-a.r); a.i=hsqt2*(-tmp_-a.i); }
  }
template<typename T> inline void PMINPLACE(T &a, T &b)
  { T t = a; a.r+=b.r; a.i+=b.i; b.r=t.r-b.r; b.i=t.i-b.i; }

template<bool bwd, typename T> void pass8 (size_t ido, size_t l1,
  const T * restrict cc, T * restrict ch, const cmplx<T0> * restrict wa)
  {
  constexpr size_t cdim=8;

  if (ido==1)
    for (size_t k=0; k<l1; ++k)
      {
      T a0, a1, a2, a3, a4, a5, a6, a7;
      PMC(a0,a4,CC(0,0,k),CC(0,4,k));
      PMC(a1,a5,CC(0,1,k),CC(0,5,k));
      PMC(a2,a6,CC(0,2,k),CC(0,6,k));
      PMC(a3,a7,CC(0,3,k),CC(0,7,k));
      ROTX90<bwd>(a6);
      ROTX90<bwd>(a7);
      PMINPLACE(a0,a2);
      PMINPLACE(a1,a3);
      PMINPLACE(a4,a6);
      PMINPLACE(a5,a7);
      ROTX45<bwd>(a5);
      ROTX90<bwd>(a3);
      ROTX135<bwd>(a7);
      PMC(CH(0,k,0),CH(0,k,4),a0,a1);
      PMC(CH(0,k,1),CH(0,k,5),a4,a5);
      PMC(CH(0,k,2),CH(0,k,6),a2,a3);
      PMC(CH(0,k,3),CH(0,k,7),a6,a7);
      }
  else
    for (size_t k=0; k<l1; ++k)
      {
      T a0, a1, a2, a3, a4, a5, a6, a7;
      PMC(a0,a4,CC(0,0,k),CC(0,4,k));
      PMC(a1,a5,CC(0,1,k),CC(0,5,k));
      PMC(a2,a6,CC(0,2,k),CC(0,6,k));
      PMC(a3,a7,CC(0,3,k),CC(0,7,k));
      ROTX90<bwd>(a6);
      ROTX90<bwd>(a7);
      PMINPLACE(a0,a2);
      PMINPLACE(a1,a3);
      PMINPLACE(a4,a6);
      PMINPLACE(a5,a7);
      ROTX45<bwd>(a5);
      ROTX90<bwd>(a3);
      ROTX135<bwd>(a7);
      PMC(CH(0,k,0),CH(0,k,4),a0,a1);
      PMC(CH(0,k,1),CH(0,k,5),a4,a5);
      PMC(CH(0,k,2),CH(0,k,6),a2,a3);
      PMC(CH(0,k,3),CH(0,k,7),a6,a7);

      for (size_t i=1; i<ido; ++i)
        {
        T a0, a1, a2, a3, a4, a5, a6, a7;
        PMC(a0,a4,CC(i,0,k),CC(i,4,k));
        PMC(a1,a5,CC(i,1,k),CC(i,5,k));
        PMC(a2,a6,CC(i,2,k),CC(i,6,k));
        PMC(a3,a7,CC(i,3,k),CC(i,7,k));
        ROTX90<bwd>(a6);
        ROTX90<bwd>(a7);
        PMINPLACE(a0,a2);
        PMINPLACE(a1,a3);
        PMINPLACE(a4,a6);
        PMINPLACE(a5,a7);
        ROTX45<bwd>(a5);
        ROTX90<bwd>(a3);
        ROTX135<bwd>(a7);
        PMINPLACE(a0,a1);
        PMINPLACE(a2,a3);
        PMINPLACE(a4,a5);
        PMINPLACE(a6,a7);
        CH(i,k,0) = a0;
        CH(i,k,1) = a4.template special_mul<bwd>(WA(0,i));
        CH(i,k,2) = a2.template special_mul<bwd>(WA(1,i));
        CH(i,k,3) = a6.template special_mul<bwd>(WA(2,i));
        CH(i,k,4) = a1.template special_mul<bwd>(WA(3,i));
        CH(i,k,5) = a5.template special_mul<bwd>(WA(4,i));
        CH(i,k,6) = a3.template special_mul<bwd>(WA(5,i));
        CH(i,k,7) = a7.template special_mul<bwd>(WA(6,i));
        }
      }
   }


898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921 922 923 924 925 926 927 928
#define PREP11(idx) \
        T t1 = CC(idx,0,k), t2, t3, t4, t5, t6, t7, t8, t9, t10, t11; \
        PMC (t2,t11,CC(idx,1,k),CC(idx,10,k)); \
        PMC (t3,t10,CC(idx,2,k),CC(idx, 9,k)); \
        PMC (t4,t9 ,CC(idx,3,k),CC(idx, 8,k)); \
        PMC (t5,t8 ,CC(idx,4,k),CC(idx, 7,k)); \
        PMC (t6,t7 ,CC(idx,5,k),CC(idx, 6,k)); \
        CH(idx,k,0).r=t1.r+t2.r+t3.r+t4.r+t5.r+t6.r; \
        CH(idx,k,0).i=t1.i+t2.i+t3.i+t4.i+t5.i+t6.i;

#define PARTSTEP11a0(u1,u2,x1,x2,x3,x4,x5,y1,y2,y3,y4,y5,out1,out2) \
        { \
        T ca = t1 + t2*x1 + t3*x2 + t4*x3 + t5*x4 +t6*x5, \
          cb; \
        cb.i=y1*t11.r y2*t10.r y3*t9.r y4*t8.r y5*t7.r; \
        cb.r=-(y1*t11.i y2*t10.i y3*t9.i y4*t8.i y5*t7.i ); \
        PMC(out1,out2,ca,cb); \
        }
#define PARTSTEP11a(u1,u2,x1,x2,x3,x4,x5,y1,y2,y3,y4,y5) \
        PARTSTEP11a0(u1,u2,x1,x2,x3,x4,x5,y1,y2,y3,y4,y5,CH(0,k,u1),CH(0,k,u2))
#define PARTSTEP11(u1,u2,x1,x2,x3,x4,x5,y1,y2,y3,y4,y5) \
        { \
        T da,db; \
        PARTSTEP11a0(u1,u2,x1,x2,x3,x4,x5,y1,y2,y3,y4,y5,da,db) \
        CH(i,k,u1) = da.template special_mul<bwd>(WA(u1-1,i)); \
        CH(i,k,u2) = db.template special_mul<bwd>(WA(u2-1,i)); \
        }

template<bool bwd, typename T> void pass11 (size_t ido, size_t l1,
  const T * restrict cc, T * restrict ch, const cmplx<T0> * restrict wa)
  {
Martin Reinecke's avatar
sync  
Martin Reinecke committed
929 930 931 932 933 934 935 936 937 938 939
  constexpr size_t cdim=11;
  constexpr T0 tw1r= T0(0.8412535328311811688618116489193677L),
               tw1i= (bwd ? 1 : -1) * T0(0.5406408174555975821076359543186917L),
               tw2r= T0(0.4154150130018864255292741492296232L),
               tw2i= (bwd ? 1 : -1) * T0(0.9096319953545183714117153830790285L),
               tw3r= T0(-0.1423148382732851404437926686163697L),
               tw3i= (bwd ? 1 : -1) * T0(0.9898214418809327323760920377767188L),
               tw4r= T0(-0.6548607339452850640569250724662936L),
               tw4i= (bwd ? 1 : -1) * T0(0.7557495743542582837740358439723444L),
               tw5r= T0(-0.9594929736144973898903680570663277L),
               tw5i= (bwd ? 1 : -1) * T0(0.2817325568414296977114179153466169L);
940 941 942 943 944 945 946 947 948 949 950 951 952 953 954 955 956 957 958 959 960 961 962 963 964 965 966 967 968 969 970 971 972 973 974 975 976 977 978 979 980 981 982 983 984 985 986 987 988 989 990 991 992 993 994 995 996 997 998 999 1000 1001 1002 1003 1004 1005 1006 1007 1008 1009 1010 1011 1012 1013 1014 1015 1016 1017 1018 1019 1020 1021 1022 1023 1024 1025 1026 1027 1028 1029 1030 1031 1032 1033 1034 1035 1036 1037 1038 1039 1040 1041 1042 1043 1044 1045 1046 1047 1048 1049 1050 1051 1052 1053 1054 1055 1056 1057 1058 1059 1060 1061 1062 1063 1064 1065 1066 1067 1068 1069 1070 1071 1072 1073 1074 1075 1076 1077 1078 1079 1080 1081 1082 1083

  if (ido==1)
    for (size_t k=0; k<l1; ++k)
      {
      PREP11(0)
      PARTSTEP11a(1,10,tw1r,tw2r,tw3r,tw4r,tw5r,+tw1i,+tw2i,+tw3i,+tw4i,+tw5i)
      PARTSTEP11a(2, 9,tw2r,tw4r,tw5r,tw3r,tw1r,+tw2i,+tw4i,-tw5i,-tw3i,-tw1i)
      PARTSTEP11a(3, 8,tw3r,tw5r,tw2r,tw1r,tw4r,+tw3i,-tw5i,-tw2i,+tw1i,+tw4i)
      PARTSTEP11a(4, 7,tw4r,tw3r,tw1r,tw5r,tw2r,+tw4i,-tw3i,+tw1i,+tw5i,-tw2i)
      PARTSTEP11a(5, 6,tw5r,tw1r,tw4r,tw2r,tw3r,+tw5i,-tw1i,+tw4i,-tw2i,+tw3i)
      }
  else
    for (size_t k=0; k<l1; ++k)
      {
      {
      PREP11(0)
      PARTSTEP11a(1,10,tw1r,tw2r,tw3r,tw4r,tw5r,+tw1i,+tw2i,+tw3i,+tw4i,+tw5i)
      PARTSTEP11a(2, 9,tw2r,tw4r,tw5r,tw3r,tw1r,+tw2i,+tw4i,-tw5i,-tw3i,-tw1i)
      PARTSTEP11a(3, 8,tw3r,tw5r,tw2r,tw1r,tw4r,+tw3i,-tw5i,-tw2i,+tw1i,+tw4i)
      PARTSTEP11a(4, 7,tw4r,tw3r,tw1r,tw5r,tw2r,+tw4i,-tw3i,+tw1i,+tw5i,-tw2i)
      PARTSTEP11a(5, 6,tw5r,tw1r,tw4r,tw2r,tw3r,+tw5i,-tw1i,+tw4i,-tw2i,+tw3i)
      }
      for (size_t i=1; i<ido; ++i)
        {
        PREP11(i)
        PARTSTEP11(1,10,tw1r,tw2r,tw3r,tw4r,tw5r,+tw1i,+tw2i,+tw3i,+tw4i,+tw5i)
        PARTSTEP11(2, 9,tw2r,tw4r,tw5r,tw3r,tw1r,+tw2i,+tw4i,-tw5i,-tw3i,-tw1i)
        PARTSTEP11(3, 8,tw3r,tw5r,tw2r,tw1r,tw4r,+tw3i,-tw5i,-tw2i,+tw1i,+tw4i)
        PARTSTEP11(4, 7,tw4r,tw3r,tw1r,tw5r,tw2r,+tw4i,-tw3i,+tw1i,+tw5i,-tw2i)
        PARTSTEP11(5, 6,tw5r,tw1r,tw4r,tw2r,tw3r,+tw5i,-tw1i,+tw4i,-tw2i,+tw3i)
        }
      }
  }

#undef PARTSTEP11
#undef PARTSTEP11a0
#undef PARTSTEP11a
#undef PREP11

#define CX(a,b,c) cc[(a)+ido*((b)+l1*(c))]
#define CX2(a,b) cc[(a)+idl1*(b)]
#define CH2(a,b) ch[(a)+idl1*(b)]

template<bool bwd, typename T> void passg (size_t ido, size_t ip,
  size_t l1, T * restrict cc, T * restrict ch, const cmplx<T0> * restrict wa,
  const cmplx<T0> * restrict csarr)
  {
  const size_t cdim=ip;
  size_t ipph = (ip+1)/2;
  size_t idl1 = ido*l1;

  arr<cmplx<T0>> wal(ip);
  wal[0] = cmplx<T0>(1., 0.);
  for (size_t i=1; i<ip; ++i)
    wal[i]=cmplx<T0>(csarr[i].r,bwd ? csarr[i].i : -csarr[i].i);

  for (size_t k=0; k<l1; ++k)
    for (size_t i=0; i<ido; ++i)
      CH(i,k,0) = CC(i,0,k);
  for (size_t j=1, jc=ip-1; j<ipph; ++j, --jc)
    for (size_t k=0; k<l1; ++k)
      for (size_t i=0; i<ido; ++i)
        PMC(CH(i,k,j),CH(i,k,jc),CC(i,j,k),CC(i,jc,k));
  for (size_t k=0; k<l1; ++k)
    for (size_t i=0; i<ido; ++i)
      {
      T tmp = CH(i,k,0);
      for (size_t j=1; j<ipph; ++j)
        tmp+=CH(i,k,j);
      CX(i,k,0) = tmp;
      }
  for (size_t l=1, lc=ip-1; l<ipph; ++l, --lc)
    {
    // j=0
    for (size_t ik=0; ik<idl1; ++ik)
      {
      CX2(ik,l).r = CH2(ik,0).r+wal[l].r*CH2(ik,1).r+wal[2*l].r*CH2(ik,2).r;
      CX2(ik,l).i = CH2(ik,0).i+wal[l].r*CH2(ik,1).i+wal[2*l].r*CH2(ik,2).i;
      CX2(ik,lc).r=-wal[l].i*CH2(ik,ip-1).i-wal[2*l].i*CH2(ik,ip-2).i;
      CX2(ik,lc).i=wal[l].i*CH2(ik,ip-1).r+wal[2*l].i*CH2(ik,ip-2).r;
      }

    size_t iwal=2*l;
    size_t j=3, jc=ip-3;
    for (; j<ipph-1; j+=2, jc-=2)
      {
      iwal+=l; if (iwal>ip) iwal-=ip;
      cmplx<T0> xwal=wal[iwal];
      iwal+=l; if (iwal>ip) iwal-=ip;
      cmplx<T0> xwal2=wal[iwal];
      for (size_t ik=0; ik<idl1; ++ik)
        {
        CX2(ik,l).r += CH2(ik,j).r*xwal.r+CH2(ik,j+1).r*xwal2.r;
        CX2(ik,l).i += CH2(ik,j).i*xwal.r+CH2(ik,j+1).i*xwal2.r;
        CX2(ik,lc).r -= CH2(ik,jc).i*xwal.i+CH2(ik,jc-1).i*xwal2.i;
        CX2(ik,lc).i += CH2(ik,jc).r*xwal.i+CH2(ik,jc-1).r*xwal2.i;
        }
      }
    for (; j<ipph; ++j, --jc)
      {
      iwal+=l; if (iwal>ip) iwal-=ip;
      cmplx<T0> xwal=wal[iwal];
      for (size_t ik=0; ik<idl1; ++ik)
        {
        CX2(ik,l).r += CH2(ik,j).r*xwal.r;
        CX2(ik,l).i += CH2(ik,j).i*xwal.r;
        CX2(ik,lc).r -= CH2(ik,jc).i*xwal.i;
        CX2(ik,lc).i += CH2(ik,jc).r*xwal.i;
        }
      }
    }

  // shuffling and twiddling
  if (ido==1)
    for (size_t j=1, jc=ip-1; j<ipph; ++j, --jc)
      for (size_t ik=0; ik<idl1; ++ik)
        {
        T t1=CX2(ik,j), t2=CX2(ik,jc);
        PMC(CX2(ik,j),CX2(ik,jc),t1,t2);
        }
  else
    {
    for (size_t j=1, jc=ip-1; j<ipph; ++j,--jc)
      for (size_t k=0; k<l1; ++k)
        {
        T t1=CX(0,k,j), t2=CX(0,k,jc);
        PMC(CX(0,k,j),CX(0,k,jc),t1,t2);
        for (size_t i=1; i<ido; ++i)
          {
          T x1, x2;
          PMC(x1,x2,CX(i,k,j),CX(i,k,jc));
          size_t idij=(j-1)*(ido-1)+i-1;
          CX(i,k,j) = x1.template special_mul<bwd>(wa[idij]);
          idij=(jc-1)*(ido-1)+i-1;
          CX(i,k,jc) = x2.template special_mul<bwd>(wa[idij]);
          }
        }
    }
  }

#undef CH2
#undef CX2
#undef CX

Martin Reinecke's avatar
Martin Reinecke committed
1084
template<bool bwd, typename T> void pass_all(T c[], T0 fct)
1085
  {
Martin Reinecke's avatar
Martin Reinecke committed
1086
  if (length==1) { c[0]*=fct; return; }
1087 1088 1089 1090
  size_t l1=1;
  arr<T> ch(length);
  T *p1=c, *p2=ch.data();

Martin Reinecke's avatar
Martin Reinecke committed
1091
  for(size_t k1=0; k1<fact.size(); k1++)
1092
    {
Martin Reinecke's avatar
Martin Reinecke committed
1093
    size_t ip=fact[k1].fct;
1094 1095 1096
    size_t l2=ip*l1;
    size_t ido = length/l2;
    if     (ip==4)
Martin Reinecke's avatar
Martin Reinecke committed
1097
      pass4<bwd> (ido, l1, p1, p2, fact[k1].tw);
Martin Reinecke's avatar
sync  
Martin Reinecke committed
1098 1099
    else if(ip==8)
      pass8<bwd>(ido, l1, p1, p2, fact[k1].tw);
1100
    else if(ip==2)
Martin Reinecke's avatar
Martin Reinecke committed
1101
      pass2<bwd>(ido, l1, p1, p2, fact[k1].tw);
1102
    else if(ip==3)
Martin Reinecke's avatar
Martin Reinecke committed
1103
      pass3<bwd> (ido, l1, p1, p2, fact[k1].tw);
1104
    else if(ip==5)
Martin Reinecke's avatar
Martin Reinecke committed
1105
      pass5<bwd> (ido, l1, p1, p2, fact[k1].tw);
1106
    else if(ip==7)
Martin Reinecke's avatar
Martin Reinecke committed
1107
      pass7<bwd> (ido, l1, p1, p2, fact[k1].tw);
1108
    else if(ip==11)
Martin Reinecke's avatar
Martin Reinecke committed
1109
      pass11<bwd> (ido, l1, p1, p2, fact[k1].tw);
1110 1111
    else
      {
Martin Reinecke's avatar
Martin Reinecke committed
1112
      passg<bwd>(ido, ip, l1, p1, p2, fact[k1].tw, fact[k1].tws);
1113 1114 1115 1116 1117 1118 1119
      swap(p1,p2);
      }
    swap(p1,p2);
    l1=l2;
    }
  if (p1!=c)
    {
Martin Reinecke's avatar
Martin Reinecke committed
1120
    if (fct!=1.)
1121
      for (size_t i=0; i<length; ++i)
Martin Reinecke's avatar
Martin Reinecke committed
1122
        c[i] = ch[i]*fct;
1123 1124 1125 1126
    else
      memcpy (c,p1,length*sizeof(T));
    }
  else
Martin Reinecke's avatar
Martin Reinecke committed
1127
    if (fct!=1.)
1128
      for (size_t i=0; i<length; ++i)
Martin Reinecke's avatar
Martin Reinecke committed
1129
        c[i] *= fct;
1130 1131 1132 1133 1134 1135 1136 1137 1138 1139 1140 1141 1142 1143 1144 1145 1146
  }

#undef WA
#undef CC
#undef CH

  public:
    template<typename T> void forward(T c[], T0 fct)
      { pass_all<false>(c, fct); }

    template<typename T> void backward(T c[], T0 fct)
      { pass_all<true>(c, fct); }

  private:
    NOINLINE void factorize()
      {
      size_t len=length;
Martin Reinecke's avatar
sync  
Martin Reinecke committed
1147 1148
      while ((len&7)==0)
        { add_factor(8); len>>=3; }
1149 1150 1151 1152 1153 1154 1155
      while ((len&3)==0)
        { add_factor(4); len>>=2; }
      if ((len&1)==0)
        {
        len>>=1;
        // factor 2 should be at the front of the factor list
        add_factor(2);
Martin Reinecke's avatar
Martin Reinecke committed
1156
        swap(fact[0].fct, fact.back().fct);
1157
        }
Martin Reinecke's avatar
Martin Reinecke committed
1158
      size_t maxl = size_t(sqrt(double(len)))+1;
1159 1160 1161 1162 1163 1164 1165 1166
      for (size_t divisor=3; (len>1)&&(divisor<maxl); divisor+=2)
        if ((len%divisor)==0)
          {
          while ((len%divisor)==0)
            {
            add_factor(divisor);
            len/=divisor;
            }
Martin Reinecke's avatar
Martin Reinecke committed
1167
          maxl=size_t(sqrt(double(len)))+1;
1168 1169 1170 1171 1172 1173 1174
          }
      if (len>1) add_factor(len);
      }

    size_t twsize() const
      {
      size_t twsize=0, l1=1;
Martin Reinecke's avatar
Martin Reinecke committed
1175
      for (size_t k=0; k<fact.size(); ++k)
1176
        {
Martin Reinecke's avatar
Martin Reinecke committed
1177
        size_t ip=fact[k].fct, ido= length/(l1*ip);
1178 1179 1180 1181 1182 1183 1184 1185 1186 1187 1188 1189 1190 1191
        twsize+=(ip-1)*(ido-1);
        if (ip>11)
          twsize+=ip;
        l1*=ip;
        }
      return twsize;
      }

    void comp_twiddle()
      {
      sincos_2pibyn<T0> twid(length, false);
      auto twiddle = twid.cdata();
      size_t l1=1;
      size_t memofs=0;
Martin Reinecke's avatar
Martin Reinecke committed
1192
      for (size_t k=0; k<fact.size(); ++k)
1193
        {
Martin Reinecke's avatar
Martin Reinecke committed
1194 1195
        size_t ip=fact[k].fct, ido=length/(l1*ip);
        fact[k].tw=mem.data()+memofs;
1196 1197 1198
        memofs+=(ip-1)*(ido-1);
        for (size_t j=1; j<ip; ++j)
          for (size_t i=1; i<ido; ++i)
Martin Reinecke's avatar
Martin Reinecke committed
1199
            fact[k].tw[(j-1)*(ido-1)+i-1] = twiddle[j*l1*i];
1200 1201
        if (ip>11)
          {
Martin Reinecke's avatar
Martin Reinecke committed
1202
          fact[k].tws=mem.data()+memofs;
1203 1204
          memofs+=ip;
          for (size_t j=0; j<ip; ++j)
Martin Reinecke's avatar
Martin Reinecke committed
1205
            fact[k].tws[j] = twiddle[j*l1*ido];
1206 1207 1208 1209 1210 1211 1212
          }
        l1*=ip;
        }
      }

  public:
    NOINLINE cfftp(size_t length_)
Martin Reinecke's avatar
Martin Reinecke committed
1213
      : length(length_)
1214 1215 1216 1217 1218 1219 1220 1221 1222 1223 1224 1225 1226 1227 1228 1229 1230 1231 1232 1233 1234 1235
      {
      if (length==0) throw runtime_error("zero length FFT requested");
      if (length==1) return;
      factorize();
      mem.resize(twsize());
      comp_twiddle();
      }
  };

//
// real-valued FFTPACK transforms
//

template<typename T0> class rfftp
  {
  private:
    struct fctdata
      {
      size_t fct;
      T0 *tw, *tws;
      };

Martin Reinecke's avatar
Martin Reinecke committed
1236
    size_t length;
1237
    arr<T0> mem;
Martin Reinecke's avatar
Martin Reinecke committed
1238
    vector<fctdata> fact;
1239 1240

    void add_factor(size_t factor)
Martin Reinecke's avatar
Martin Reinecke committed
1241
      { fact.push_back({factor, nullptr, nullptr}); }
1242 1243

#define WA(x,i) wa[(i)+(x)*(ido-1)]
Martin Reinecke's avatar
Martin Reinecke committed
1244 1245 1246
template<typename T> inline void PM(T &a, T &b, T c, T d)
  { a=c+d; b=c-d; }

1247
/* (a+ib) = conj(c+id) * (e+if) */
Martin Reinecke's avatar
Martin Reinecke committed
1248 1249 1250
template<typename T1, typename T2, typename T3> inline void MULPM
  (T1 &a, T1 &b, T2 c, T2 d, T3 e, T3 f)
  {  a=c*e+d*f; b=c*f-d*e; }
1251 1252 1253 1254 1255 1256 1257 1258 1259 1260

#define CC(a,b,c) cc[(a)+ido*((b)+l1*(c))]
#define CH(a,b,c) ch[(a)+ido*((b)+cdim*(c))]

template<typename T> void radf2 (size_t ido, size_t l1,
  const T * restrict cc, T * restrict ch, const T0 * restrict wa)
  {
  constexpr size_t cdim=2;

  for (size_t k=0; k<l1; k++)
Martin Reinecke's avatar
Martin Reinecke committed
1261
    PM (CH(0,0,k),CH(ido-1,1,k),CC(0,k,0),CC(0,k,1));
1262 1263 1264 1265 1266 1267 1268 1269 1270 1271 1272 1273
  if ((ido&1)==0)
    for (size_t k=0; k<l1; k++)
      {
      CH(    0,1,k) = -CC(ido-1,k,1);
      CH(ido-1,0,k) =  CC(ido-1,k,0);
      }
  if (ido<=2) return;
  for (size_t k=0; k<l1; k++)
    for (size_t i=2; i<ido; i+=2)
      {
      size_t ic=ido-i;
      T tr2, ti2;
Martin Reinecke's avatar
Martin Reinecke committed
1274 1275 1276
      MULPM (tr2,ti2,WA(0,i-2),WA(0,i-1),CC(i-1,k,1),CC(i,k,1));
      PM (CH(i-1,0,k),CH(ic-1,1,k),CC(i-1,k,0),tr2);
      PM (CH(i  ,0,k),CH(ic  ,1,k),ti2,CC(i  ,k,0));
1277 1278 1279 1280 1281 1282 1283
      }
  }

template<typename T> void radf3(size_t ido, size_t l1,
  const T * restrict cc, T * restrict ch, const T0 * restrict wa)
  {
  constexpr size_t cdim=3;
Martin Reinecke's avatar
sync  
Martin Reinecke committed
1284
  constexpr T0 taur=-0.5, taui=T0(0.8660254037844386467637231707529362L);
1285 1286 1287 1288 1289 1290 1291 1292 1293 1294 1295 1296 1297 1298

  for (size_t k=0; k<l1; k++)
    {
    T cr2=CC(0,k,1)+CC(0,k,2);
    CH(0,0,k) = CC(0,k,0)+cr2;
    CH(0,2,k) = taui*(CC(0,k,2)-CC(0,k,1));
    CH(ido-1,1,k) = CC(0,k,0)+taur*cr2;
    }
  if (ido==1) return;
  for (size_t k=0; k<l1; k++)
    for (size_t i=2; i<ido; i+=2)
      {
      size_t ic=ido-i;
      T di2, di3, dr2, dr3;
Martin Reinecke's avatar
Martin Reinecke committed
1299 1300
      MULPM (dr2,di2,WA(0,i-2),WA(0,i-1),CC(i-1,k,1),CC(i,k,1)); // d2=conj(WA0)*CC1
      MULPM (dr3,di3,WA(1,i-2),WA(1,i-1),CC(i-1,k,2),CC(i,k,2)); // d3=conj(WA1)*CC2
1301 1302 1303 1304 1305 1306 1307 1308
      T cr2=dr2+dr3; // c add
      T ci2=di2+di3;
      CH(i-1,0,k) = CC(i-1,k,0)+cr2; // c add
      CH(i  ,0,k) = CC(i  ,k,0)+ci2;
      T tr2 = CC(i-1,k,0)+taur*cr2; // c add
      T ti2 = CC(i  ,k,0)+taur*ci2;
      T tr3 = taui*(di2-di3);  // t3 = taui*i*(d3-d2)?
      T ti3 = taui*(dr3-dr2);
Martin Reinecke's avatar
Martin Reinecke committed
1309 1310
      PM(CH(i-1,2,k),CH(ic-1,1,k),tr2,tr3); // PM(i) = t2+t3
      PM(CH(i  ,2,k),CH(ic  ,1,k),ti3,ti2); // PM(ic) = conj(t2-t3)
1311 1312 1313 1314 1315 1316 1317
      }
  }

template<typename T> void radf4(size_t ido, size_t l1,
  const T * restrict cc, T * restrict ch, const T0 * restrict wa)
  {
  constexpr size_t cdim=4;
Martin Reinecke's avatar
sync  
Martin Reinecke committed
1318
  constexpr T0 hsqt2=T0(0.707106781186547524400844362104849L);
1319 1320 1321 1322

  for (size_t k=0; k<l1; k++)
    {
    T tr1,tr2;
Martin Reinecke's avatar
Martin Reinecke committed
1323 1324 1325
    PM (tr1,CH(0,2,k),CC(0,k,3),CC(0,k,1));
    PM (tr2,CH(ido-1,1,k),CC(0,k,0),CC(0,k,2));
    PM (CH(0,0,k),CH(ido-1,3,k),tr2,tr1);
1326 1327 1328 1329 1330 1331
    }
  if ((ido&1)==0)
    for (size_t k=0; k<l1; k++)
      {
      T ti1=-hsqt2*(CC(ido-1,k,1)+CC(ido-1,k,3));
      T tr1= hsqt2*(CC(ido-1,k,1)-CC(ido-1,k,3));
Martin Reinecke's avatar
Martin Reinecke committed
1332 1333
      PM (CH(ido-1,0,k),CH(ido-1,2,k),CC(ido-1,k,0),tr1);
      PM (CH(    0,3,k),CH(    0,1,k),ti1,CC(ido-1,k,2));
1334 1335 1336 1337 1338 1339 1340
      }
  if (ido<=2) return;
  for (size_t k=0; k<l1; k++)
    for (size_t i=2; i<ido; i+=2)
      {
      size_t ic=ido-i;
      T ci2, ci3, ci4, cr2, cr3, cr4, ti1, ti2, ti3, ti4, tr1, tr2, tr3, tr4;