/***************************************************************************
                          Reduction.hpp  -  description
                             -------------------
    begin                : Mar 24, 2013
    copyright            : (C) 2013 by Tomas Oberhuber et al.
    email                : tomas.oberhuber@fjfi.cvut.cz
 ***************************************************************************/

/* See Copyright Notice in tnl/Copyright */

// Implemented by: Tomas Oberhuber, Jakub Klinkovsky

#pragma once

#include <memory>  // std::unique_ptr

//#define CUDA_REDUCTION_PROFILING

#include <TNL/Algorithms/Reduction.h>
#include <TNL/Algorithms/CudaReductionKernel.h>
#include <TNL/Algorithms/MultiDeviceMemoryOperations.h>

#ifdef CUDA_REDUCTION_PROFILING
#include <iostream>
#include <TNL/Timer.h>
#endif

namespace TNL {
namespace Algorithms {

/****
 * Arrays smaller than the following constant
 * are reduced on CPU. The constant must not be larger
 * than maximal CUDA grid size.
 */
static constexpr int Reduction_minGpuDataSize = 256;//65536; //16384;//1024;//256;

template< typename Index,
          typename Result,
          typename ReductionOperation,
          typename DataFetcher >
constexpr Result
Reduction< Devices::Sequential >::
reduce( const Index begin,
        const Index end,
        const ReductionOperation& reduction,
        DataFetcher& dataFetcher,
        const Result& zero )
{
   constexpr int block_size = 128;
   const Index size = end - begin;
   const Index blocks = size / block_size;

   if( blocks > 1 ) {
      // initialize array for unrolled results
      Result r[ 4 ] = { zero, zero, zero, zero };

      // main reduction (explicitly unrolled loop)
      for( Index b = 0; b < blocks; b++ ) {
         const Index offset = begin + b * block_size;
         for( int i = 0; i < block_size; i += 4 ) {
            r[ 0 ] = reduction( r[ 0 ], dataFetcher( offset + i ) );
            r[ 1 ] = reduction( r[ 1 ], dataFetcher( offset + i + 1 ) );
            r[ 2 ] = reduction( r[ 2 ], dataFetcher( offset + i + 2 ) );
            r[ 3 ] = reduction( r[ 3 ], dataFetcher( offset + i + 3 ) );
         }
      }

      // reduction of the last, incomplete block (not unrolled)
      for( Index i = begin + blocks * block_size; i < end; i++ )
         r[ 0 ] = reduction( r[ 0 ], dataFetcher( i ) );

      // reduction of unrolled results
      r[ 0 ] = reduction( r[ 0 ], r[ 2 ] );
      r[ 1 ] = reduction( r[ 1 ], r[ 3 ] );
      r[ 0 ] = reduction( r[ 0 ], r[ 1 ] );
      return r[ 0 ];
   }
   else {
      Result result = zero;
      for( Index i = begin; i < end; i++ )
         result = reduction( result, dataFetcher( i ) );
      return result;
   }
}

template< typename Index,
          typename Result,
          typename ReductionOperation,
          typename DataFetcher >
constexpr std::pair< Result, Index >
Reduction< Devices::Sequential >::
reduceWithArgument( const Index begin,
                    const Index end,
                    const ReductionOperation& reduction,
                    DataFetcher& dataFetcher,
                    const Result& zero )
{
   constexpr int block_size = 128;
   const Index size = end - begin;
   const Index blocks = size / block_size;

   if( blocks > 1 ) {
      // initialize array for unrolled results
      Index arg[ 4 ] = { 0, 0, 0, 0 };
      Result r[ 4 ] = { zero, zero, zero, zero };
      bool initialized( false );

      // main reduction (explicitly unrolled loop)
      for( Index b = 0; b < blocks; b++ ) {
         const Index offset = begin + b * block_size;
         for( int i = 0; i < block_size; i += 4 ) {
            if( ! initialized )
            {
               arg[ 0 ] = offset + i;
               arg[ 1 ] = offset + i + 1;
               arg[ 2 ] = offset + i + 2;
               arg[ 3 ] = offset + i + 3;
               r[ 0 ] = dataFetcher( offset + i );
               r[ 1 ] = dataFetcher( offset + i + 1 );
               r[ 2 ] = dataFetcher( offset + i + 2 );
               r[ 3 ] = dataFetcher( offset + i + 3 );
               initialized = true;
               continue;
            }
            reduction( r[ 0 ], dataFetcher( offset + i ),     arg[ 0 ], offset + i );
            reduction( r[ 1 ], dataFetcher( offset + i + 1 ), arg[ 1 ], offset + i + 1 );
            reduction( r[ 2 ], dataFetcher( offset + i + 2 ), arg[ 2 ], offset + i + 2 );
            reduction( r[ 3 ], dataFetcher( offset + i + 3 ), arg[ 3 ], offset + i + 3 );
         }
      }

      // reduction of the last, incomplete block (not unrolled)
      for( Index i = begin + blocks * block_size; i < size; i++ )
         reduction( r[ 0 ], dataFetcher( i ), arg[ 0 ], i );

      // reduction of unrolled results
      reduction( r[ 0 ], r[ 2 ], arg[ 0 ], arg[ 2 ] );
      reduction( r[ 1 ], r[ 3 ], arg[ 1 ], arg[ 3 ] );
      reduction( r[ 0 ], r[ 1 ], arg[ 0 ], arg[ 1 ] );
      return std::make_pair( r[ 0 ], arg[ 0 ] );
   }
   else {
      std::pair< Result, Index > result( dataFetcher( begin ), begin );
      for( Index i = begin + 1; i < end; i++ )
         reduction( result.first, dataFetcher( i ), result.second, i );
      return result;
   }
}

template< typename Index,
          typename Result,
          typename ReductionOperation,
          typename DataFetcher >
Result
Reduction< Devices::Host >::
reduce( const Index begin,
        const Index end,
        const ReductionOperation& reduction,
        DataFetcher& dataFetcher,
        const Result& zero )
{
#ifdef HAVE_OPENMP
   constexpr int block_size = 128;
   const Index size = end - begin;
   const Index blocks = size / block_size;

   if( Devices::Host::isOMPEnabled() && blocks >= 2 ) {
      // global result variable
      Result result = zero;
      const int threads = TNL::min( blocks, Devices::Host::getMaxThreadsCount() );
#pragma omp parallel num_threads(threads)
      {
         // initialize array for thread-local results
         Result r[ 4 ] = { zero, zero, zero, zero  };

         #pragma omp for nowait
         for( Index b = 0; b < blocks; b++ ) {
            const Index offset = begin + b * block_size;
            for( int i = 0; i < block_size; i += 4 ) {
               r[ 0 ] = reduction( r[ 0 ], dataFetcher( offset + i ) );
               r[ 1 ] = reduction( r[ 1 ], dataFetcher( offset + i + 1 ) );
               r[ 2 ] = reduction( r[ 2 ], dataFetcher( offset + i + 2 ) );
               r[ 3 ] = reduction( r[ 3 ], dataFetcher( offset + i + 3 ) );
            }
         }

         // the first thread that reaches here processes the last, incomplete block
         #pragma omp single nowait
         {
            for( Index i = begin + blocks * block_size; i < end; i++ )
               r[ 0 ] = reduction( r[ 0 ], dataFetcher( i ) );
         }

         // local reduction of unrolled results
         r[ 0 ] = reduction( r[ 0 ], r[ 2 ] );
         r[ 1 ] = reduction( r[ 1 ], r[ 3 ] );
         r[ 0 ] = reduction( r[ 0 ], r[ 1 ] );

         // inter-thread reduction of local results
         #pragma omp critical
         {
            result = reduction( result, r[ 0 ] );
         }
      }
      return result;
   }
   else
#endif
      return Reduction< Devices::Sequential >::reduce( begin, end, reduction, dataFetcher, zero );
}

template< typename Index,
          typename Result,
          typename ReductionOperation,
          typename DataFetcher >
std::pair< Result, Index >
Reduction< Devices::Host >::
reduceWithArgument( const Index begin,
                    const Index end,
                    const ReductionOperation& reduction,
                    DataFetcher& dataFetcher,
                    const Result& zero )
{
#ifdef HAVE_OPENMP
   constexpr int block_size = 128;
   const Index size = end - begin;
   const Index blocks = size / block_size;

   if( Devices::Host::isOMPEnabled() && blocks >= 2 ) {
      // global result variable
      std::pair< Result, Index > result( zero, -1 );
      const int threads = TNL::min( blocks, Devices::Host::getMaxThreadsCount() );
#pragma omp parallel num_threads(threads)
      {
         // initialize array for thread-local results
         Index arg[ 4 ] = { 0, 0, 0, 0 };
         Result r[ 4 ] = { zero, zero, zero, zero  };
         bool initialized( false );

         #pragma omp for nowait
         for( Index b = 0; b < blocks; b++ ) {
            const Index offset = begin + b * block_size;
            for( int i = 0; i < block_size; i += 4 ) {
               if( ! initialized ) {
                  arg[ 0 ] = offset + i;
                  arg[ 1 ] = offset + i + 1;
                  arg[ 2 ] = offset + i + 2;
                  arg[ 3 ] = offset + i + 3;
                  r[ 0 ] = dataFetcher( offset + i );
                  r[ 1 ] = dataFetcher( offset + i + 1 );
                  r[ 2 ] = dataFetcher( offset + i + 2 );
                  r[ 3 ] = dataFetcher( offset + i + 3 );
                  initialized = true;
                  continue;
               }
               reduction( r[ 0 ], dataFetcher( offset + i ),     arg[ 0 ], offset + i );
               reduction( r[ 1 ], dataFetcher( offset + i + 1 ), arg[ 1 ], offset + i + 1 );
               reduction( r[ 2 ], dataFetcher( offset + i + 2 ), arg[ 2 ], offset + i + 2 );
               reduction( r[ 3 ], dataFetcher( offset + i + 3 ), arg[ 3 ], offset + i + 3 );
            }
         }

         // the first thread that reaches here processes the last, incomplete block
         #pragma omp single nowait
         {
            for( Index i = begin + blocks * block_size; i < end; i++ )
               reduction( r[ 0 ], dataFetcher( i ), arg[ 0 ], i );
         }

         // local reduction of unrolled results
         reduction( r[ 0 ], r[ 2 ], arg[ 0 ], arg[ 2 ] );
         reduction( r[ 1 ], r[ 3 ], arg[ 1 ], arg[ 3 ] );
         reduction( r[ 0 ], r[ 1 ], arg[ 0 ], arg[ 1 ] );

         // inter-thread reduction of local results
         #pragma omp critical
         {
            if( result.second == -1 )
               result.second = arg[ 0 ];
            reduction( result.first, r[ 0 ], result.second, arg[ 0 ] );
         }
      }
      return result;
   }
   else
#endif
      return Reduction< Devices::Sequential >::reduceWithArgument( begin, end, reduction, dataFetcher, zero );
}

template< typename Index,
          typename Result,
          typename ReductionOperation,
          typename DataFetcher >
Result
Reduction< Devices::Cuda >::
reduce( const Index begin,
        const Index end,
        const ReductionOperation& reduction,
        DataFetcher& dataFetcher,
        const Result& zero )
{
   // Only fundamental and pointer types can be safely reduced on host. Complex
   // objects stored on the device might contain pointers into the device memory,
   // in which case reduction on host might fail.
   constexpr bool can_reduce_later_on_host = std::is_fundamental< Result >::value || std::is_pointer< Result >::value;

   #ifdef CUDA_REDUCTION_PROFILING
      Timer timer;
      timer.reset();
      timer.start();
   #endif

   CudaReductionKernelLauncher< Index, Result > reductionLauncher( begin, end );

   // start the reduction on the GPU
   Result* deviceAux1( 0 );
   const int reducedSize = reductionLauncher.start(
      reduction,
      dataFetcher,
      zero,
      deviceAux1 );

   #ifdef CUDA_REDUCTION_PROFILING
      timer.stop();
      std::cout << "   Reduction on GPU to size " << reducedSize << " took " << timer.getRealTime() << " sec. " << std::endl;
      timer.reset();
      timer.start();
   #endif

   if( can_reduce_later_on_host ) {
      // transfer the reduced data from device to host
      std::unique_ptr< Result[] > resultArray{
         // Workaround for nvcc 10.1.168 - it would modifie the simple expression
         // `new Result[reducedSize]` in the source code to `new (Result[reducedSize])`
         // which is not correct - see e.g. https://stackoverflow.com/a/39671946
         // Thus, the host compiler would spit out hundreds of warnings...
         // Funnily enough, nvcc's behaviour depends on the context rather than the
         // expression, because exactly the same simple expression in different places
         // does not produce warnings.
         #ifdef __NVCC__
         new Result[ static_cast<const int&>(reducedSize) ]
         #else
         new Result[ reducedSize ]
         #endif
      };
      MultiDeviceMemoryOperations< void, Devices::Cuda >::copy( resultArray.get(), deviceAux1, reducedSize );

      #ifdef CUDA_REDUCTION_PROFILING
         timer.stop();
         std::cout << "   Transferring data to CPU took " << timer.getRealTime() << " sec. " << std::endl;
         timer.reset();
         timer.start();
      #endif

      // finish the reduction on the host
      auto fetch = [&] ( Index i ) { return resultArray[ i ]; };
      const Result result = Reduction< Devices::Sequential >::reduce( 0, reducedSize, reduction, fetch, zero );

      #ifdef CUDA_REDUCTION_PROFILING
         timer.stop();
         std::cout << "   Reduction of small data set on CPU took " << timer.getRealTime() << " sec. " << std::endl;
      #endif
      return result;
   }
   else {
      // data can't be safely reduced on host, so continue with the reduction on the GPU
      auto result = reductionLauncher.finish( reduction, zero );

      #ifdef CUDA_REDUCTION_PROFILING
         timer.stop();
         std::cout << "   Reduction of small data set on GPU took " << timer.getRealTime() << " sec. " << std::endl;
         timer.reset();
         timer.start();
      #endif

      return result;
   }
}

template< typename Index,
          typename Result,
          typename ReductionOperation,
          typename DataFetcher >
std::pair< Result, Index >
Reduction< Devices::Cuda >::
reduceWithArgument( const Index begin,
                    const Index end,
                    const ReductionOperation& reduction,
                    DataFetcher& dataFetcher,
                    const Result& zero )
{
   // Only fundamental and pointer types can be safely reduced on host. Complex
   // objects stored on the device might contain pointers into the device memory,
   // in which case reduction on host might fail.
   constexpr bool can_reduce_later_on_host = std::is_fundamental< Result >::value || std::is_pointer< Result >::value;

   #ifdef CUDA_REDUCTION_PROFILING
      Timer timer;
      timer.reset();
      timer.start();
   #endif

   CudaReductionKernelLauncher< Index, Result > reductionLauncher( begin, end );

   // start the reduction on the GPU
   Result* deviceAux1( nullptr );
   Index* deviceIndexes( nullptr );
   const int reducedSize = reductionLauncher.startWithArgument(
      reduction,
      dataFetcher,
      zero,
      deviceAux1,
      deviceIndexes );

   #ifdef CUDA_REDUCTION_PROFILING
      timer.stop();
      std::cout << "   Reduction on GPU to size " << reducedSize << " took " << timer.getRealTime() << " sec. " << std::endl;
      timer.reset();
      timer.start();
   #endif

   if( can_reduce_later_on_host ) {
      // transfer the reduced data from device to host
      std::unique_ptr< Result[] > resultArray{
         // Workaround for nvcc 10.1.168 - it would modifie the simple expression
         // `new Result[reducedSize]` in the source code to `new (Result[reducedSize])`
         // which is not correct - see e.g. https://stackoverflow.com/a/39671946
         // Thus, the host compiler would spit out hundreds of warnings...
         // Funnily enough, nvcc's behaviour depends on the context rather than the
         // expression, because exactly the same simple expression in different places
         // does not produce warnings.
         #ifdef __NVCC__
         new Result[ static_cast<const int&>(reducedSize) ]
         #else
         new Result[ reducedSize ]
         #endif
      };
      std::unique_ptr< Index[] > indexArray{
         // Workaround for nvcc 10.1.168 - it would modifie the simple expression
         // `new Index[reducedSize]` in the source code to `new (Index[reducedSize])`
         // which is not correct - see e.g. https://stackoverflow.com/a/39671946
         // Thus, the host compiler would spit out hundreds of warnings...
         // Funnily enough, nvcc's behaviour depends on the context rather than the
         // expression, because exactly the same simple expression in different places
         // does not produce warnings.
         #ifdef __NVCC__
         new Index[ static_cast<const int&>(reducedSize) ]
         #else
         new Index[ reducedSize ]
         #endif
      };
      MultiDeviceMemoryOperations< void, Devices::Cuda >::copy( resultArray.get(), deviceAux1, reducedSize );
      MultiDeviceMemoryOperations< void, Devices::Cuda >::copy( indexArray.get(), deviceIndexes, reducedSize );

      #ifdef CUDA_REDUCTION_PROFILING
         timer.stop();
         std::cout << "   Transferring data to CPU took " << timer.getRealTime() << " sec. " << std::endl;
         timer.reset();
         timer.start();
      #endif

      // finish the reduction on the host
//      auto fetch = [&] ( Index i ) { return resultArray[ i ]; };
//      const Result result = Reduction< Devices::Sequential >::reduceWithArgument( reducedSize, argument, reduction, fetch, zero );
      for( Index i = 1; i < reducedSize; i++ )
         reduction( resultArray[ 0 ], resultArray[ i ], indexArray[ 0 ], indexArray[ i ]  );

      #ifdef CUDA_REDUCTION_PROFILING
         timer.stop();
         std::cout << "   Reduction of small data set on CPU took " << timer.getRealTime() << " sec. " << std::endl;
      #endif
      return std::make_pair( resultArray[ 0 ], indexArray[ 0 ] );
   }
   else {
      // data can't be safely reduced on host, so continue with the reduction on the GPU
      auto result = reductionLauncher.finishWithArgument( reduction, zero );

      #ifdef CUDA_REDUCTION_PROFILING
         timer.stop();
         std::cout << "   Reduction of small data set on GPU took " << timer.getRealTime() << " sec. " << std::endl;
         timer.reset();
         timer.start();
      #endif

      return result;
   }
}

} // namespace Algorithms
} // namespace TNL
