/Users/mourrain/Devel/mmx/algebramix/include/algebramix/fft_threads.hpp

Go to the documentation of this file.

/******************************************************************************
* MODULE     : fft_threads.hpp
* DESCRIPTION: multi-threaded FFT
* COPYRIGHT  : (C) 2007  Joris van der Hoeven
*******************************************************************************
* This software falls under the GNU general public license and comes WITHOUT
* ANY WARRANTY WHATSOEVER. See the file $TEXMACS_PATH/LICENSE for more details.
* If you don't have this file, write to the Free Software Foundation, Inc.,
* 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
******************************************************************************/

#ifndef __MMX__FFT_THREADS__HPP
#define __MMX__FFT_THREADS__HPP
#include <basix/threads.hpp>
#include <algebramix/vector_threads.hpp>
#include <algebramix/fft_blocks.hpp>

namespace mmx {

/******************************************************************************
* Multi-threaded and dichotomic FFT transformation
******************************************************************************/

template<typename C,
         typename FFTER= fft_blocks_transformer<C,
                                                fft_naive_transformer<C>,
                                                8, 6, 10>,
         nat thr= 14>
class fft_threads_transformer {
public:
  typedef implementation<vector_linear,vector_threads<vector_naive> > Vec;
  typedef typename FFTER::R R;
  typedef typename R::U U;
  typedef typename R::S S;
  static const nat min_reps= 16;

  FFTER* ffter;
  nat    depth;
  nat    len;
  U*     roots;

public:
  inline fft_threads_transformer (nat n):
    ffter (new FFTER (n)),
    depth (ffter->depth), len (ffter->len), roots (ffter->roots) {}
  inline ~fft_threads_transformer () { delete ffter; }

  static inline void
  copy (C* d, nat drs, nat dcs, C* s, nat srs, nat scs, nat r, nat c) {
    for (nat j=0; j<c; j++, d+=dcs, s+=scs) {
      C* dd= d; C* ss= s;
      for (nat i=0; i<r; i++, dd+=drs, ss+=srs)
        *dd= *ss;
    }
  }

  struct outer_fft_task_rep: public task_rep {
    FFTER* ffter;
    bool direct;
    C *buf;    
    nat thr_reps, tot_reps;
    C *c;
    nat stride, shift, steps;
  public:
    inline outer_fft_task_rep
     (FFTER* ffter2, bool direct2,
      C* buf2, nat thr_reps2, nat tot_reps2,
      C* c2, nat stride2, nat shift2, nat steps2):
        ffter (ffter2), direct (direct2),
        buf (buf2), thr_reps (thr_reps2), tot_reps (tot_reps2),
        c (c2), stride (stride2), shift (shift2), steps (steps2) {}
    inline ~outer_fft_task_rep () {}
    void execute () {
      for (nat i=0; i*min_reps < thr_reps; i++, c += stride * min_reps) {
        C* aux= buf;
        copy (buf, 1, min_reps, c, stride, stride * tot_reps,
              min_reps, 1 << steps);
        if (direct)
          for (nat j=0; j<min_reps; j++, aux++)
            ffter->dfft (aux, min_reps, shift, steps);
        else
          for (nat j=0; j<min_reps; j++, aux++)
            ffter->ifft (aux, min_reps, shift, steps);
        copy (c, stride, stride * tot_reps, buf, 1, min_reps,
              min_reps, 1 << steps);
      }
    }
  };

  struct inner_fft_task_rep: public task_rep {
    FFTER* ffter;
    bool direct;
    nat start, inc, total;
    C *c;
    nat stride, shift, steps;
  public:
    inline inner_fft_task_rep
     (FFTER* ffter2, bool direct2, nat start2, nat inc2, nat total2,
      C* c2, nat stride2, nat shift2, nat steps2):
        ffter (ffter2), direct (direct2),
        start (start2), inc (inc2), total (total2),
        c (c2), stride (stride2), shift (shift2), steps (steps2) {}
    inline ~inner_fft_task_rep () {}
    void execute () {
      if (direct)
        for (nat i=start; i<total; i+=inc)
          ffter->dfft (c + (i<<steps) * stride, stride,
                       shift + ((i << steps) >> 1), steps);
      else
        for (nat i=start; i<total; i+=inc)
          ffter->ifft (c + (i<<steps) * stride, stride,
                       shift + ((i << steps) >> 1), steps);     
    }
  };

  void
  fft (bool direct, C* c, nat stride, nat shift, nat steps) {
    if (steps <= thr) {
      if (direct) ffter->dfft (c, stride, shift, steps);
      else ffter->ifft (c, stride, shift, steps);
    }
    else {
      nat half1= min(log_2 (threads_number), steps);
      nat half2= steps - half1;

      for (nat stage=0; stage<2; stage++) {
        if ((stage == 0) ^ (!direct)) {
          //C*  cc1  = c;
          //for (nat i=0; i<((nat) 1<<half2); i++, cc1 += stride)
          //  ffter->fft (direct, cc1, stride << half2,
          //              shift >> half2, half1);

          nat nt = threads_number;
          nat bsz= (min_reps << half1);
          C* buf= mmx_new<C> (nt * bsz);
          task tasks[nt];
          for (nat i=0; i<nt; i++) {
            nat tot= 1 << half2;
            nat sta= min_reps * ((  i   * tot) / (min_reps * nt));
            nat end= min_reps * (((i+1) * tot) / (min_reps * nt));
            tasks[i]= new outer_fft_task_rep
              (ffter, direct, buf + i*bsz, end-sta, tot,
               c + sta*stride, stride, shift >> half2, half1);
          }
          threads_execute (tasks, nt);
          mmx_delete<C> (buf, nt * bsz);

        }
        else {
          //C* cc2= c;
          //for (nat i=0; i<((nat) 1<<half1); i++, cc2 += (stride << half2))
          //  ffter->fft (direct, cc2, stride,
          //              shift + ((i << half2) >> 1), half2);
          nat nt= threads_number;
          task tasks[nt];
          for (nat i=0; i<nt; i++)
            tasks[i]= new inner_fft_task_rep
              (ffter, direct, i, nt, 1 << half1, c, stride, shift, half2);
          threads_execute (tasks, nt);
        }
      }
    }
  }

  inline void
  dfft (C* c, nat stride, nat shift, nat steps) {
    if (steps <= thr) ffter->dfft (c, stride, shift, steps);
    else fft (true, c, stride, shift, steps);
  }

  inline void
  ifft (C* c, nat stride, nat shift, nat steps) {
    if (steps <= thr) ffter->ifft (c, stride, shift, steps);
    else fft (false, c, stride, shift, steps);
  }

  inline void
  dfft (C* c, nat stride, nat shift, nat steps, nat step1, nat step2) {
    if (step1 == 0 && step2 == steps && steps > thr)
      fft (true, c, stride, shift, steps);
    else ffter->dfft (c, stride, shift, steps, step1, step2);
  }

  inline void
  ifft (C* c, nat stride, nat shift, nat steps, nat step1, nat step2) {
    if (step1 == 0 && step2 == steps && steps > thr)
      fft (false, c, stride, shift, steps);
    else ffter->ifft (c, stride, shift, steps, step1, step2);
  }

  inline void
  direct_transform (C* c) {
    dfft (c, 1, 0, depth); }

  inline void
  inverse_transform (C* c, bool divide=true) {
    ifft (c, 1, 0, depth);
    if (divide) {
      S x= binpow (S (2), depth);
      x= invert (x);
      Vec::template vec_unary_scalar<typename R::fft_mul_sc_op> (c, x, len);
    }
  }
};

} // namespace mmx
#endif //__MMX__FFT_THREADS__HPP