/* Copyright (c) 2022, NVIDIA CORPORATION. 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 NVIDIA CORPORATION 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 ``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 OWNER 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.
 */

/*
* 1D DWT for Haar wavelet and signals with a length which is a power of 2.
* The code reduces bank conflicts and non-coalesced reads / writes as
* appropriate but does not fully remove them because the computational
* overhead to achieve this would outweighs the benefit (see inline comments
* for more details).
* Large signals are subdivided into sub-signals with 512 elements and the
* wavelet transform for these is computed with one block over 10 decomposition
* levels. The resulting signal consisting of the approximation coefficients at
* level X is then processed in a subsequent step on the device. This requires
* interblock synchronization which is only possible on host side.
* Detail coefficients which have been computed are not further referenced
* during the decomposition so that they can be stored directly in their final
* position in global memory. The transform and its storing scheme preserve
* locality in the coefficients so that these writes are coalesced.
* Approximation coefficients are stored in shared memory because they are
* needed to compute the subsequent decomposition step. The top most
* approximation coefficient for a sub-signal processed by one block is stored
* in a special global memory location to simplify the processing after the
* interblock synchronization.
* Most books on wavelets explain the Haar wavelet decomposition. A good freely
* available resource is the Wavelet primer by Stollnitz et al.
* http://grail.cs.washington.edu/projects/wavelets/article/wavelet1.pdf
* http://grail.cs.washington.edu/projects/wavelets/article/wavelet2.pdf
* The basic of all Wavelet transforms is to decompose a signal into
* approximation (a) and detail (d) coefficients where the detail tends to be
* small or zero which allows / simplifies compression. The following "graphs"
* demonstrate the transform for a signal
* of length eight. The index always describes the decomposition level where
* a coefficient arises. The input signal is interpreted as approximation signal
* at level 0. The coefficients computed on the device are stored in the same
* scheme as in the example. This data structure is particularly well suited for
* compression and also preserves the hierarchical structure of the
decomposition.

-------------------------------------------------
| a_0 | a_0 | a_0 | a_0 | a_0 | a_0 | a_0 | a_0 |
-------------------------------------------------

-------------------------------------------------
| a_1 | a_1 | a_1 | a_1 | d_1 | d_1 | d_1 | d_1 |
-------------------------------------------------

-------------------------------------------------
| a_2 | a_2 | d_2 | d_2 | d_1 | d_1 | d_1 | d_1 |
-------------------------------------------------

-------------------------------------------------
| a_3 | d_3 | d_2 | d_2 | d_1 | d_1 | d_1 | d_1 |
-------------------------------------------------

* Device Code.
*/

#ifndef _DWTHAAR1D_KERNEL_H_
#define _DWTHAAR1D_KERNEL_H_

#include <cooperative_groups.h>

namespace cg = cooperative_groups;

////////////////////////////////////////////////////////////////////////////////
//! @param id  input data
//! @param od  output data
//! @param value
////////////////////////////////////////////////////////////////////////////////
__global__ void initValue(float *od, float value) {
  // Handle to thread block group
  cg::thread_block cta = cg::this_thread_block();
  // position of write into global memory
  unsigned int index = (blockIdx.x * blockDim.x) + threadIdx.x;

  od[index] = value;

  // sync after each decomposition step
  cg::sync(cta);
}

////////////////////////////////////////////////////////////////////////////////
//! Compute partial wavelet decomposition on the GPU using Haar basis
//! For each thread block the full decomposition is computed but these results
//! have to be combined
//! Use one thread to perform the full decomposition
//! @param id  input data
//! @param od  output data
//! @param approx_final  place to store the final approximation coefficient for
//!                      the subsignal
//! @param dlevels  number of decomposition levels for this transform
//! @param slength_step_half   half signal length for current decomposition
//!                            level (offset for storing detail coefficients in
//!                            global memory
//! @param bdim  block dimension
////////////////////////////////////////////////////////////////////////////////
__global__ void dwtHaar1D(float *id, float *od, float *approx_final,
                          const unsigned int dlevels,
                          const unsigned int slength_step_half,
                          const int bdim) {
  // Handle to thread block group
  cg::thread_block cta = cg::this_thread_block();

  // shared memory for part of the signal
  extern __shared__ float shared[];

  // thread runtime environment, 1D parametrization
  const int gdim = gridDim.x;
  // const int bdim = blockDim.x;
  const int bid = blockIdx.x;
  const int tid = threadIdx.x;

  // global thread id (w.r.t. to total data set)
  const int tid_global = (bid * bdim) + tid;
  unsigned int idata = (bid * (2 * bdim)) + tid;

  // read data from global memory
  shared[tid] = id[idata];
  shared[tid + bdim] = id[idata + bdim];
  cg::sync(cta);

  // this operation has a two way bank conflicts for all threads, this are two
  // additional cycles for each warp -- all alternatives to avoid this bank
  // conflict are more expensive than the one cycle introduced by serialization
  float data0 = shared[2 * tid];
  float data1 = shared[(2 * tid) + 1];
  cg::sync(cta);

  // detail coefficient, not further referenced so directly store in
  // global memory
  od[tid_global + slength_step_half] = (data0 - data1) * INV_SQRT_2;

  // offset to avoid bank conflicts
  // see the scan example for a more detailed description
  unsigned int atid = tid + (tid >> LOG_NUM_BANKS);

  // approximation coefficient
  // store in shared memory for further decomposition steps in this global step
  shared[atid] = (data0 + data1) * INV_SQRT_2;

  // all threads have to write approximation coefficient to shared memory before
  // next steps can take place
  cg::sync(cta);

  // early out if possible
  // the compiler removes this part from the source because dlevels is
  // a constant shader input
  // note: syncthreads in bodies of branches can lead to dead-locks unless
  // the condition evaluates the same way for ALL threads of a block, as in
  // this case
  if (dlevels > 1) {
    // offset to second element in shared element which has to be used for the
    // decomposition, effectively 2^(i - 1)
    unsigned int offset_neighbor = 1;
    // number of active threads per decomposition level
    // identical to the offset for the detail coefficients
    unsigned int num_threads = bdim >> 1;

    // index for the first element of the pair to process
    // the representation is still compact (and therefore still tid * 2)
    // because the first step operated on registers and only the result has been
    // written to shared memory
    unsigned int idata0 = tid * 2;

    // offset levels to make the loop more efficient
    for (unsigned int i = 1; i < dlevels; ++i) {
      // Non-coalesced writes occur if the number of active threads becomes
      // less than 16 for a block because the start address for the first
      // block is not always aligned with 64 byte which is necessary for
      // coalesced access. However, the problem only occurs at high levels
      // with only a small number of active threads so that the total number of
      // non-coalesced access is rather small and does not justify the
      // computations which are necessary to avoid these uncoalesced writes
      // (this has been tested and verified)
      if (tid < num_threads) {
        // update stride, with each decomposition level the stride grows by a
        // factor of 2
        unsigned int idata1 = idata0 + offset_neighbor;

        // position of write into global memory
        unsigned int g_wpos = (num_threads * gdim) + (bid * num_threads) + tid;

        // compute wavelet decomposition step

        // offset to avoid bank conflicts
        unsigned int c_idata0 = idata0 + (idata0 >> LOG_NUM_BANKS);
        unsigned int c_idata1 = idata1 + (idata1 >> LOG_NUM_BANKS);

        // detail coefficient, not further modified so directly store
        // in global memory
        od[g_wpos] = (shared[c_idata0] - shared[c_idata1]) * INV_SQRT_2;

        // approximation coefficient
        // note that the representation in shared memory becomes rather sparse
        // (with a lot of holes inbetween) but the storing scheme in global
        // memory guarantees that the common representation (approx, detail_0,
        // detail_1, ...)
        // is achieved
        shared[c_idata0] = (shared[c_idata0] + shared[c_idata1]) * INV_SQRT_2;

        // update storage offset for details
        num_threads = num_threads >> 1;  // div 2
        offset_neighbor <<= 1;           // mul 2
        idata0 = idata0 << 1;            // mul 2
      }

      // sync after each decomposition step
      cg::sync(cta);
    }

    // write the top most level element for the next decomposition steps
    // which are performed after an interlock synchronization on host side
    if (0 == tid) {
      approx_final[bid] = shared[0];
    }

  }  // end early out if possible
}

#endif  // #ifndef _DWTHAAR1D_KERNEL_H_