content: {-# LANGUAGE CPP                 #-}
  {-# LANGUAGE GADTs               #-}
  {-# LANGUAGE ImpredicativeTypes  #-}
  {-# LANGUAGE PatternGuards       #-}
  {-# LANGUAGE QuasiQuotes         #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TypeOperators       #-}
  -- |
  -- Module      : Data.Array.Accelerate.CUDA.CodeGen.PrefixSum
  -- Copyright   : [2008..2014] Manuel M T Chakravarty, Gabriele Keller
  --               [2009..2014] Trevor L. McDonell
  -- License     : BSD3
  --
  -- Maintainer  : Trevor L. McDonell <tmcdonell@cse.unsw.edu.au>
  -- Stability   : experimental
  -- Portability : non-portable (GHC extensions)
  --
  
  module  Data.Array.Accelerate.CUDA.CodeGen.PrefixSum (
  
    -- skeletons
    mkScanl, mkScanl1, mkScanl',
    mkScanr, mkScanr1, mkScanr',
  
  ) where
  
  import Data.Maybe
  import Foreign.CUDA.Analysis
  import Language.C.Quote.CUDA
  import qualified Language.C.Syntax                      as C
  
  import Data.Array.Accelerate.Array.Sugar
  import Data.Array.Accelerate.Analysis.Match
  
  import Data.Array.Accelerate.CUDA.AST
  import Data.Array.Accelerate.CUDA.Analysis.Shape
  import Data.Array.Accelerate.CUDA.CodeGen.Base
  
  errorMsg :: String
  errorMsg
    = error
    $ unlines [ "accelerate-cuda does not support rank-polymorphic scans. Please switch to accelerate-llvm-ptx instead."
              , ""
              , "***   https://hackage.haskell.org/package/accelerate-llvm-ptx   ***"
              , "***   https://github.com/AccelerateHS/accelerate-llvm           ***"
              ]
  
  -- Wrappers
  -- --------
  
  mkScanl, mkScanr
      :: forall aenv sh e. (Shape sh, Elt e)
      => DeviceProperties
      -> Gamma aenv
      -> CUFun2 aenv (e -> e -> e)
      -> CUExp aenv e
      -> CUDelayedAcc aenv (sh:.Int) e
      -> [CUTranslSkel aenv (Array (sh:.Int) e)]
  mkScanl dev aenv f z a
    | Just Refl <- matchShapeType (undefined::sh) (undefined::Z)
    = [ mkScan    L dev aenv f (Just z) a
      , mkScanUp1 L dev aenv f a
      , mkScanUp2 L dev aenv f (Just z) ]
  
    | otherwise
    = error errorMsg
  
  mkScanr dev aenv f z a
    | Just Refl <- matchShapeType (undefined::sh) (undefined::Z)
    = [ mkScan    R dev aenv f (Just z) a
      , mkScanUp1 R dev aenv f a
      , mkScanUp2 R dev aenv f (Just z) ]
  
    | otherwise
    = error errorMsg
  
  mkScanl1, mkScanr1
      :: forall aenv sh e. (Shape sh, Elt e)
      => DeviceProperties
      -> Gamma aenv
      -> CUFun2 aenv (e -> e -> e)
      -> CUDelayedAcc aenv (sh:.Int) e
      -> [CUTranslSkel aenv (Array (sh:.Int) e)]
  mkScanl1 dev aenv f a
    | Just Refl <- matchShapeType (undefined::sh) (undefined::Z)
    = [ mkScan    L dev aenv f Nothing a
      , mkScanUp1 L dev aenv f a
      , mkScanUp2 L dev aenv f Nothing ]
  
    | otherwise
    = error errorMsg
  
  mkScanr1 dev aenv f a
    | Just Refl <- matchShapeType (undefined::sh) (undefined::Z)
    = [ mkScan    R dev aenv f Nothing a
      , mkScanUp1 R dev aenv f a
      , mkScanUp2 R dev aenv f Nothing ]
  
    | otherwise
    = error errorMsg
  
  mkScanl', mkScanr'
      :: forall aenv sh e. (Shape sh, Elt e)
      => DeviceProperties
      -> Gamma aenv
      -> CUFun2 aenv (e -> e -> e)
      -> CUExp aenv e
      -> CUDelayedAcc aenv (sh:.Int) e
      -> [CUTranslSkel aenv (Array (sh:.Int) e, Array sh e)]
  mkScanl' dev aenv f z
    | Just Refl <- matchShapeType (undefined::sh) (undefined::Z)
    = map cast . mkScanl dev aenv f z
    | otherwise
    = error errorMsg
  
  mkScanr' dev aenv f z
    | Just Refl <- matchShapeType (undefined::sh) (undefined::Z)
    = map cast . mkScanr dev aenv f z
    | otherwise
    = error errorMsg
  
  cast :: CUTranslSkel aenv a -> CUTranslSkel aenv b
  cast (CUTranslSkel entry code) = CUTranslSkel entry code
  
  
  -- Core implementation
  -- -------------------
  
  data Direction = L | R
    deriving Eq
  
  instance Show Direction where
    show L = "l"
    show R = "r"
  
  
  -- [OVERVIEW]
  --
  -- Data.List-style exclusive scan, with the additional restriction that the
  -- first argument needs to be an /associative/ function to enable efficient
  -- parallel implementation. The initial value may be arbitrary.
  --
  -- scanl :: Elt a
  --       => (Exp a -> Exp a -> Exp a)
  --       -> Exp a
  --       -> Acc (Vector a)
  --       -> Acc (Vector a)
  --
  -- > scanl (+) 10 (use xs)
  -- >   where
  -- >     xs = fromList (Z:.10) (cycle [1])
  -- >
  -- > ==> Array (Z:.11) [10,11,12,13,14,15,16,17,18,19,20]
  --
  -- Data.List-style inclusive scan without an initial value
  --
  -- scanl1 :: Elt a
  --        => (Exp a -> Exp a -> Exp a)
  --        -> Acc (Vector a)
  --        -> Acc (Vector a)
  --
  -- > scanl1 (+) (use xs)
  -- >   where
  -- >     xs = fromList (Z:.10) (cycle [1])
  -- >
  -- > ==> Array (Z:.10) [1,2,3,4,5,6,7,8,9,10]
  --
  -- Variant of 'scanl' where the final result is returned separately.
  --
  -- scanl' :: Elt a
  --        => (Exp a -> Exp a -> Exp a)
  --        -> Exp a
  --        -> Acc (Vector a)
  --        -> (Acc (Vector a), Acc (Scalar a))
  --
  -- Denotationally, we have:
  --
  -- > scanl' f z xs = (init res, last res)
  -- >   where
  -- >     res = scanl f z xs
  --
  --
  -- [IMPLEMENTATION]
  --
  -- This code handles all the above cases, in both left and right-handed
  -- variants. This is the _downsweep_ phase to a multi-block scan procedure.
  -- We require a work distribution such that there is a _single_ thread block for
  -- each interval. For multi-block scans, we have an array of interval sums that
  -- are used to determine the carry-in value from the previous interval. Note
  -- that 'argBlk' will not be accessed by a single-block scan, so may be null.
  --
  -- We require some pointer manipulation from the calling code in order to
  -- support all types of scans:
  --
  --   * scanl          : argSum should point to the last position of argOut
  --   * scanr          : argSum should be the start of argOut, argOut should be incremented by one
  --   * scanl1, scanr1 : no change (argSum is required, even though it will not be used Haskell-side)
  --   * scanl', scanr' : no change
  --
  mkScan :: forall aenv e. Elt e
         => Direction
         -> DeviceProperties
         -> Gamma aenv
         -> CUFun2 aenv (e -> e -> e)
         -> Maybe (CUExp aenv e)
         -> CUDelayedAcc aenv DIM1 e
         -> CUTranslSkel aenv (Vector e)
  mkScan dir dev aenv fun@(CUFun2 _ _ combine) mseed (CUDelayed (CUExp shIn) _ (CUFun1 _ get)) =
    CUTranslSkel scan [cunit|
  
      $esc:("#include <accelerate_cuda.h>")
      $edecls:texIn
  
      extern "C" __global__ void
      $id:scan
      (
          $params:argIn,
          $params:argOut,
          $params:argBlk,
          $params:(tail argSum)           // just the pointers, no shape information
      )
      {
          $decls:smem
          $decls:declt
          $decls:declx
          $decls:decly
          $decls:declz
          $items:(sh .=. shIn)
  
          const int shapeSize     = $exp:(csize sh);
          const int intervalSize  = (shapeSize + gridDim.x - 1) / gridDim.x;
  
          /*
           * Read in previous result partial sum. We store the carry value in
           * temporary value 'z' and read new values from the input array into
           * 'x', since 'scanBlock' will store its results into 'y' on completion.
           */
          int carryIn = 0;
  
          if ( threadIdx.x == 0 ) {
              $stm:initialise
          }
  
          const int start         = blockIdx.x * intervalSize;
          const int end           = min(start + intervalSize, shapeSize);
          const int numElements   = end - start;
                int seg;
  
          for ( seg = threadIdx.x
              ; seg < numElements
              ; seg += blockDim.x )
          {
              const int ix = $exp:firstIndex;
  
              /*
               * Generate the next set of values
               */
              $items:(x .=. get ix)
  
              /*
               * Carry in the result from the privous segment
               */
              if ( $exp:carryIn ) {
                  $items:(t .=. combine z x)
                  $items:(x .=. t)
              }
  
              /*
               * Store our input into shared memory and perform a cooperative
               * inclusive left scan.
               */
              $items:(sdata "threadIdx.x" .=. x)
              __syncthreads();
  
              $items:(scanBlock dev fun x y t sdata Nothing)
  
              /*
               * Exclusive scans write the result of the previous thread to global
               * memory. The first thread must reinstate the carry-in value which
               * is the result of the last thread from the previous interval, or
               * the carry-in/seed value for multi-block scans.
               */
              if ( $exp:(cbool (isJust mseed)) ) {
                  if ( threadIdx.x == 0 ) {
                      $items:(x .=. z)
                  } else {
                      $items:(x .=. sdata "threadIdx.x - 1")
                  }
              }
              $items:(setOut "ix" .=. x)
  
              /*
               * Carry the final result of this block through the set 'z'. If this
               * is the final interval, this is the value to write out as the
               * reduction result
               */
              if ( threadIdx.x == 0 ) {
                  const int last = min(numElements - seg, blockDim.x) - 1;
                  $items:(z .=. sdata "last")
              }
              $items:setCarry
          }
  
          /*
           * Finally, exclusive scans set the overall scan result (reduction value)
           */
          $items:setFinal
      }
    |]
    where
      scan                        = "scan" ++ show dir ++ maybe "1" (const []) mseed
      (texIn, argIn)              = environment dev aenv
      (argOut, _, setOut)         = writeArray "Out" (undefined :: Vector e)
      (argSum, _, totalSum)       = writeArray "Sum" (undefined :: Vector e)
      (argBlk, _, blkSum)         = writeArray "Blk" (undefined :: Vector e)
      (_, t, declt)               = locals "t" (undefined :: e)
      (_, x, declx)               = locals "x" (undefined :: e)
      (_, y, decly)               = locals "y" (undefined :: e)
      (_, z, declz)               = locals "z" (undefined :: e)
      (sh, _, _)                  = locals "sh" (undefined :: DIM1)
      (smem, sdata)               = shared (undefined :: e) "sdata" [cexp| blockDim.x |] Nothing
      ix                          = [cvar "ix"]
      setSum                      = totalSum "0"
  
      -- depending on whether we are inclusive/exclusive scans
      setCarry
        | isNothing mseed         = [[citem| carryIn = 1; |]]
        | otherwise               = []
  
      setFinal
        | isNothing mseed         = []
        | otherwise               = [[citem| if ( threadIdx.x == 0 && blockIdx.x == $id:lastBlock ) {
                                                 $items:(setSum .=. z)
                                             } |]]
  
      -- accessing neighbouring blocks
      firstBlock          = if dir == L then "0" else "gridDim.x - 1"
      lastBlock           = if dir == R then "0" else "gridDim.x - 1"
      prevBlock           = if dir == L then "blockIdx.x - 1" else "blockIdx.x + 1"
  
      firstIndex
        | dir == L        = [cexp| start + seg |]
        | otherwise       = [cexp| end - seg - 1 |]
  
      carryIn
        | isJust mseed    = [cexp| threadIdx.x == 0 |]
        | otherwise       = [cexp| threadIdx.x == 0 && carryIn |]
  
      -- initialise the first thread with the results of the previous block sweep
      -- or exclusive scan element
      initialise
        | Just (CUExp seed) <- mseed
        = [cstm|  if ( gridDim.x > 1 ) {
                      $items:(z .=. blkSum "blockIdx.x")
                  } else {
                      $items:(z .=. seed)
                  }
          |]
  
        | otherwise
        = [cstm|  if ( blockIdx.x != $id:firstBlock ) {
                      $items:(z .=. blkSum prevBlock)
                      carryIn = 1;
                  }
          |]
  
  
  -- This computes the _upsweep_ phase of a multi-block scan. This is much like a
  -- regular inclusive scan, except that only the final value for each interval is
  -- output, rather than the entire body of the scan. Indeed, if the combination
  -- function were commutative, this is equivalent to a parallel tree reduction.
  --
  mkScanUp1
      :: forall aenv e. Elt e
      => Direction
      -> DeviceProperties
      -> Gamma aenv
      -> CUFun2 aenv (e -> e -> e)
      -> CUDelayedAcc aenv DIM1 e
      -> CUTranslSkel aenv (Vector e)
  mkScanUp1 dir dev aenv fun@(CUFun2 _ _ combine) (CUDelayed (CUExp shIn) _ (CUFun1 _ get)) =
    CUTranslSkel scan [cunit|
  
      $esc:("#include <accelerate_cuda.h>")
      $edecls:texIn
  
      extern "C" __global__ void
      $id:scan
      (
          $params:argIn,
          $params:argOut
      )
      {
          $decls:smem
          $decls:declt
          $decls:declx
          $decls:decly
          $items:(sh .=. shIn)
  
          const int shapeSize     = $exp:(csize sh);
          const int intervalSize  = (shapeSize + gridDim.x - 1) / gridDim.x;
  
          const int start         = blockIdx.x * intervalSize;
          const int end           = min(start + intervalSize, shapeSize);
          const int numElements   = end - start;
                int carryIn       = 0;
                int seg;
  
          for ( seg = threadIdx.x
              ; seg < numElements
              ; seg += blockDim.x )
          {
              const int ix = $exp:firstIndex ;
  
              /*
               * Read in new values, combine with carry-in
               */
              $items:(x .=. get ix)
  
              if ( threadIdx.x == 0 && carryIn ) {
                  $items:(t .=. combine y x)
                  $items:(x .=. t)
              }
  
              /*
               * Store in shared memory and cooperatively scan
               */
              $items:(sdata "threadIdx.x" .=. x)
              __syncthreads();
  
              $items:(scanBlock dev fun x y t sdata Nothing)
  
              /*
               * Store the final result of the block to be carried in
               */
              if ( threadIdx.x == 0 ) {
                  const int last = min(numElements - seg, blockDim.x) - 1;
                  $items:(y .=. sdata "last")
              }
              carryIn = 1;
          }
  
          /*
           * Finally, the first thread writes the result of this interval
           */
          if ( threadIdx.x == 0 ) {
              $items:(setOut "blockIdx.x" .=. y)
          }
      }
    |]
    where
      scan                        = "scan" ++ show dir ++ "Up"
      (texIn, argIn)              = environment dev aenv
      (argOut, _, setOut)         = writeArray "Out" (undefined :: Vector e)
      (_, x, declx)               = locals "x" (undefined :: e)
      (_, y, decly)               = locals "y" (undefined :: e)
      (_, t, declt)               = locals "t" (undefined :: e)
      (sh, _, _)                  = locals "sh" (undefined :: DIM1)
      (smem, sdata)               = shared (undefined :: e) "sdata" [cexp| blockDim.x |] Nothing
      ix                          = [cvar "ix"]
  
      firstIndex
        | dir == L                = [cexp| start + seg |]
        | otherwise               = [cexp| end - seg - 1 |]
  
  
  -- Second step of the upsweep phase: scan the interval sums to produce carry-in
  -- values for each block of the final downsweep step
  --
  mkScanUp2
      :: forall aenv e. Elt e
      => Direction
      -> DeviceProperties
      -> Gamma aenv
      -> CUFun2 aenv (e -> e -> e)
      -> Maybe (CUExp aenv e)
      -> CUTranslSkel aenv (Vector e)
  mkScanUp2 dir dev aenv f z
    = let (_, _, get) = readArray "Blk" (undefined :: Vector e)
      in  mkScan dir dev aenv f z get
  
  
  -- Block scans
  -- ===========
  
  scanBlock
      :: forall aenv e. Elt e
      => DeviceProperties
      -> CUFun2 aenv (e -> e -> e)
      -> [C.Exp] -> [C.Exp] -> [C.Exp]
      -> (Name -> [C.Exp])
      -> Maybe C.Exp
      -> [C.BlockItem]
  scanBlock dev f x0 x1 x2 sdata mlim
    | shflOK dev (undefined :: e) = error "shfl-scan"
    | otherwise                   = scanBlockTree dev f x0 x1 x2 sdata mlim
  
  
  -- Use a thread block to scan values in shared memory. Each thread must have
  -- already stored its initial value into shared memory. The final result for
  -- this thread will be stored in x0 as well as the appropriate place in shared
  -- memory.
  --
  scanBlockTree
      :: forall aenv e. Elt e
      => DeviceProperties
      -> CUFun2 aenv (e -> e -> e)
      -> [C.Exp] -> [C.Exp] -> [C.Exp]    -- input variables x0 and x1, plus a temporary to store the intermediate value
      -> (Name -> [C.Exp])                -- index elements from shared memory
      -> Maybe C.Exp                      -- partially full block bounds check?
      -> [C.BlockItem]
  scanBlockTree dev (CUFun2 _ _ f) x0 x1 x2 sdata mlim = map (scan . pow2) [ 0 .. maxThreads ]
    where
      pow2 :: Int -> Int
      pow2 x      = 2 ^ x
      maxThreads  = floor (logBase 2 (fromIntegral $ maxThreadsPerBlock dev :: Double))
  
      inrange n
        | Just m <- mlim  = [cexp| threadIdx.x >= $int:n && threadIdx.x < $exp:m |]
        | otherwise       = [cexp| threadIdx.x >= $int:n |]
  
      scan n = [citem|
        if ( blockDim.x > $int:n ) {
            if ( $exp:(inrange n) ) {
                $items:(x1 .=. sdata ("threadIdx.x - " ++ show n))
                $items:(x2 .=. f x1 x0)
                $items:(x0 .=. x2)
            }
            __syncthreads();
            $items:(sdata "threadIdx.x" .=. x0)
            __syncthreads();
        }
        |]
  
  
  -- Shuffle scan
  -- ------------
  
  shflOK :: Elt e => DeviceProperties -> e -> Bool
  shflOK _dev _ = False
  -- shflOk dev dummy
  --   = computeCapability dev >= Compute 3 0 && all (`elem` [4,8]) (eltSizeOf dummy)
  
  {--
  scanWarpShfl
      :: forall aenv e. Elt e
      => DeviceProperties
      -> CUFun2 aenv (e -> e -> e)
      -> [C.Exp] -> [C.Exp]               -- temporary variables x0 and x1
      -> Maybe C.Exp                      -- partially full block bounds check
      -> C.Exp                            -- thread identified, usually lane or thread ID
      -> C.Stm
  scanWarpShfl _dev (CUFun2 f) x0 x1 mlim tid
    = [cstm|
        for ( int z = 1; z <= warpSize; z *= 2 ) {
            $items:(x0 .=. shfl_up x1)
  
            if ( $exp:inrange ) {
                $items:(x1 .=. f x1 x0)
            }
        }
      |]
    where
      inrange
        | Just m <- mlim  = [cexp| $exp:tid >= z && $exp:tid < $exp:m |]
        | otherwise       = [cexp| $exp:tid >= z |]
  
      sizeof      = eltSizeOf (undefined :: e)
      shfl_up     = zipWith (\s x -> ccall (shfl s) [ x, cvar "z" ]) sizeof
        where
          shfl 4  = "shfl_up32"
          shfl 8  = "shfl_up64"
          shfl _  = INTERNAL_ERROR(error) "shfl_up" "I only know about 32- and 64-bit types"
  --}
  
  