Program Listing for File ParallelFFT.h

↰ Return to documentation for file (ADS/parallel/ParallelFFT.h)

/*--------------------------------------------------------------------------*/
/*--------------------------- File ParallelFFT.h ---------------------------*/
/*--------------------------------------------------------------------------*/
/*--------------------------------------------------------------------------*/
/*----------------------------- DEFINITIONS --------------------------------*/
/*--------------------------------------------------------------------------*/

#ifndef __ADS_ParallelFFT
 #define __ADS_ParallelFFT
                      /* self-identification: #endif constant */

/*--------------------------------------------------------------------------*/
/*------------------------------ INCLUDES ----------------------------------*/
/*--------------------------------------------------------------------------*/

#include "ADS/Contract.h"
#include "ADS/numeric/FFT.h"
#include "ADS/parallel/Executor.h"
#include "ADS/parallel/Scratch.h"

#include <bit>
#include <cmath>
#include <cstddef>
#include <numbers>
#include <vector>

/*--------------------------------------------------------------------------*/
/*--------------------------- NAMESPACE ------------------------------------*/
/*--------------------------------------------------------------------------*/

namespace ADS {

inline constexpr std::size_t kParallelFFT = std::size_t( 1 ) << 15;

template< Executor E >
void fft( E & ex , std::vector< Complex > & a , bool inverse );

/*--------------------------------------------------------------------------*/

namespace detail {

template< Executor E >
void tiledTranspose( E & ex , const Complex * S , Complex * D ,
                     std::size_t r , std::size_t c )
{
 constexpr std::size_t R = 16 , Cn = 256;
 const std::size_t rt = ( r + R - 1 ) / R , ct = ( c + Cn - 1 ) / Cn;
 parallelFor( ex , 0 , rt * ct , [ & ]( std::size_t task ) {
  const std::size_t i0 = ( task / ct ) * R , i1 = std::min( r , i0 + R ) ,
                    c0 = ( task % ct ) * Cn , c1 = std::min( c , c0 + Cn );
  for( std::size_t j0 = c0 ; j0 < c1 ; j0 += R ) {
   const std::size_t j1 = std::min( c1 , j0 + R );
   for( std::size_t i = i0 ; i < i1 ; ++i )
    for( std::size_t j = j0 ; j < j1 ; ++j )
     D[ j * r + i ] = S[ i * c + j ];
   }
  } , 1 );
 }

}  // end( namespace detail )

template< Executor E >
void fft( E & ex , std::vector< Complex > & a , bool inverse )
{
 const std::size_t n = a.size();
 ADS_PRE( ( n == 0 ) || std::has_single_bit( n ) );
 if( ( ex.concurrency() == 1 ) || ( n < kParallelFFT ) ) {
  fft( a , inverse );
  return;
  }
 const int k = std::countr_zero( n );
 const std::size_t n1 = std::size_t( 1 ) << ( k / 2 ) , n2 = n / n1;
 // the twiddles w^t = Hi[ t >> h ] Lo[ t & ( 2^h - 1 ) ]
 const int h = ( k + 1 ) / 2;
 const std::size_t lo = std::size_t( 1 ) << h , hiN = n >> h;
 const double sgn = inverse ? 1.0 : -1.0;
 std::vector< Complex > Lo( lo ) , Hi( hiN );
 parallelFor( ex , 0 , lo , [ & ]( std::size_t t ) {
  Lo[ t ] = std::polar( 1.0 , sgn * 2 * std::numbers::pi * double( t ) /
                              double( n ) ); } );
 parallelFor( ex , 0 , hiN , [ & ]( std::size_t t ) {
  Hi[ t ] = std::polar( 1.0 , sgn * 2 * std::numbers::pi *
                              double( t << h ) / double( n ) ); } );
 detail::ScratchBuffer< Complex > room( n );
 Complex * const b = room.data();
 // the columns as rows (b is n2 x n1), transformed and multiplied by
 // their twiddles; back to rows of n2 (in a), transformed; the result
 // transposed into b and copied to a
 detail::tiledTranspose( ex , a.data() , b , n1 , n2 );
 parallelForBlocks( ex , 0 , n2 , [ & ]( std::size_t r0 , std::size_t r1 ) {
  std::vector< Complex > v( n1 );
  for( std::size_t j2 = r0 ; j2 < r1 ; ++j2 ) {
   Complex * row = b + j2 * n1;
   std::copy( row , row + n1 , v.begin() );
   fft( v , inverse );
   for( std::size_t k1 = 0 ; k1 < n1 ; ++k1 ) {
    const std::size_t t = ( j2 * k1 ) & ( n - 1 );
    row[ k1 ] = v[ k1 ] * ( Hi[ t >> h ] * Lo[ t & ( lo - 1 ) ] );
    }
   }
  } , 1 );
 detail::tiledTranspose( ex , b , a.data() , n2 , n1 );
 parallelForBlocks( ex , 0 , n1 , [ & ]( std::size_t r0 , std::size_t r1 ) {
  std::vector< Complex > v( n2 );
  for( std::size_t k1 = r0 ; k1 < r1 ; ++k1 ) {
   Complex * row = a.data() + k1 * n2;
   std::copy( row , row + n2 , v.begin() );
   fft( v , inverse );
   std::copy( v.begin() , v.end() , row );
   }
  } , 1 );
 // \hat a[ k1 + n1 k2 ] is the entry k2 of the row k1
 detail::tiledTranspose( ex , a.data() , b , n1 , n2 );
 parallelForBlocks( ex , 0 , n , [ & ]( std::size_t lo2 , std::size_t hi2 ) {
  std::copy( b + lo2 , b + hi2 , a.data() + lo2 ); } , 16384 );
 }

/*--------------------------------------------------------------------------*/

}  // end( namespace ADS )

/*--------------------------------------------------------------------------*/

#endif  /* ParallelFFT.h included */

/*--------------------------------------------------------------------------*/
/*------------------------- End File ParallelFFT.h -------------------------*/
/*--------------------------------------------------------------------------*/