content: {-# LANGUAGE GADTs               #-}
  {-# LANGUAGE OverloadedStrings   #-}
  {-# LANGUAGE RebindableSyntax    #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TemplateHaskell     #-}
  {-# LANGUAGE TypeApplications    #-}
  {-# LANGUAGE TypeOperators       #-}
  {-# LANGUAGE ViewPatterns        #-}
  -- |
  -- Module      : Data.Array.Accelerate.LLVM.PTX.CodeGen.FoldSeg
  -- Copyright   : [2016..2020] The Accelerate Team
  -- License     : BSD3
  --
  -- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
  -- Stability   : experimental
  -- Portability : non-portable (GHC extensions)
  --
  
  module Data.Array.Accelerate.LLVM.PTX.CodeGen.FoldSeg
    where
  
  import Data.Array.Accelerate.Representation.Array
  import Data.Array.Accelerate.Representation.Shape
  import Data.Array.Accelerate.Representation.Type
  
  import Data.Array.Accelerate.LLVM.CodeGen.Arithmetic                as A
  import Data.Array.Accelerate.LLVM.CodeGen.Array
  import Data.Array.Accelerate.LLVM.CodeGen.Base
  import Data.Array.Accelerate.LLVM.CodeGen.Constant
  import Data.Array.Accelerate.LLVM.CodeGen.Environment
  import Data.Array.Accelerate.LLVM.CodeGen.Exp
  import Data.Array.Accelerate.LLVM.CodeGen.IR
  import Data.Array.Accelerate.LLVM.CodeGen.Loop                      as Loop
  import Data.Array.Accelerate.LLVM.CodeGen.Monad
  import Data.Array.Accelerate.LLVM.CodeGen.Sugar
  import Data.Array.Accelerate.LLVM.Compile.Cache
  
  import Data.Array.Accelerate.LLVM.PTX.Analysis.Launch
  import Data.Array.Accelerate.LLVM.PTX.CodeGen.Base
  import qualified Data.Array.Accelerate.LLVM.PTX.CodeGen.Fold        as Fold ( reduceBlock, reduceWarp, imapFromTo )
  import Data.Array.Accelerate.LLVM.PTX.Target
  
  import LLVM.AST.Type.Representation
  
  import qualified Foreign.CUDA.Analysis                              as CUDA
  
  import Control.Monad                                                ( void )
  import Control.Monad.State                                          ( gets )
  import Data.String                                                  ( fromString )
  import Prelude                                                      as P
  
  
  -- Segmented reduction along the innermost dimension of an array. Performs one
  -- reduction per segment of the source array.
  --
  mkFoldSeg
      :: forall aenv sh i e.
         UID
      -> Gamma            aenv
      -> ArrayR (Array (sh, Int) e)
      -> IntegralType i
      -> IRFun2       PTX aenv (e -> e -> e)
      -> Maybe (IRExp PTX aenv e)
      -> MIRDelayed   PTX aenv (Array (sh, Int) e)
      -> MIRDelayed   PTX aenv (Segments i)
      -> CodeGen      PTX      (IROpenAcc PTX aenv (Array (sh, Int) e))
  mkFoldSeg uid aenv repr intTp combine seed arr seg =
    (+++) <$> mkFoldSegP_block uid aenv repr intTp combine seed arr seg
          <*> mkFoldSegP_warp  uid aenv repr intTp combine seed arr seg
  
  
  -- This implementation assumes that the segments array represents the offset
  -- indices to the source array, rather than the lengths of each segment. The
  -- segment-offset approach is required for parallel implementations.
  --
  -- Each segment is computed by a single thread block, meaning we don't have to
  -- worry about inter-block synchronisation.
  --
  mkFoldSegP_block
      :: forall aenv sh i e.
         UID
      -> Gamma          aenv
      -> ArrayR (Array (sh, Int) e)
      -> IntegralType i
      -> IRFun2     PTX aenv (e -> e -> e)
      -> MIRExp     PTX aenv e
      -> MIRDelayed PTX aenv (Array (sh, Int) e)
      -> MIRDelayed PTX aenv (Segments i)
      -> CodeGen    PTX      (IROpenAcc PTX aenv (Array (sh, Int) e))
  mkFoldSegP_block uid aenv repr@(ArrayR shr tp) intTp combine mseed marr mseg = do
    dev <- liftCodeGen $ gets ptxDeviceProperties
    --
    let
        (arrOut, paramOut)  = mutableArray repr "out"
        (arrIn,  paramIn)   = delayedArray "in"  marr
        (arrSeg, paramSeg)  = delayedArray "seg" mseg
        paramEnv            = envParam aenv
        --
        config              = launchConfig dev (CUDA.decWarp dev) dsmem const [|| const ||]
        dsmem n             = sharedMemorySizeAdd tp warps 0
          where
            ws        = CUDA.warpSize dev
            warps     = n `P.quot` ws
    --
    makeOpenAccWith config uid "foldSeg_block" (paramOut ++ paramIn ++ paramSeg ++ paramEnv) $ do
  
      -- We use a dynamically scheduled work queue in order to evenly distribute
      -- the uneven workload, due to the variable length of each segment, over the
      -- available thread blocks.
      -- queue <- globalWorkQueue
  
      -- All threads in the block need to know what the start and end indices of
      -- this segment are in order to participate in the reduction. We use
      -- variables in __shared__ memory to communicate these values between
      -- threads in the block. Furthermore, by using a 2-element array, we can
      -- have the first two threads of the block read the start and end indices as
      -- a single coalesced read, since they will be sequential in the
      -- segment-offset array.
      --
      smem  <- staticSharedMem (TupRsingle scalarTypeInt) 2
  
      -- Compute the number of segments and size of the innermost dimension. These
      -- are required if we are reducing a rank-2 or higher array, to properly
      -- compute the start and end indices of the portion of the array this thread
      -- block reduces. Note that this is a segment-offset array computed by
      -- 'scanl (+) 0' of the segment length array, so its size has increased by
      -- one.
      --
      sz    <- indexHead <$> delayedExtent arrIn
      ss    <- do n <- indexHead <$> delayedExtent arrSeg
                  A.sub numType n (liftInt 1)
  
      -- Each thread block cooperatively reduces a segment.
      -- s0    <- dequeue queue (lift 1)
      -- for s0 (\s -> A.lt singleType s end) (\_ -> dequeue queue (lift 1)) $ \s -> do
  
      start <- return (liftInt 0)
      end   <- shapeSize shr (irArrayShape arrOut)
  
      Fold.imapFromTo start end $ \s -> do
  
        -- The first two threads of the block determine the indices of the
        -- segments array that we will reduce between and distribute those values
        -- to the other threads in the block.
        tid <- threadIdx
        when (A.lt singleType tid (liftInt32 2)) $ do
          i <- case shr of
                 ShapeRsnoc ShapeRz -> return s
                 _ -> A.rem integralType s ss
          j <- A.add numType i =<< int tid
          v <- app1 (delayedLinearIndex arrSeg) j
          writeArray TypeInt32 smem tid =<< A.fromIntegral intTp numType v
  
        -- Once all threads have caught up, begin work on the new segment.
        __syncthreads
  
        u <- readArray TypeInt32 smem (liftInt32 0)
        v <- readArray TypeInt32 smem (liftInt32 1)
  
        -- Determine the index range of the input array we will reduce over.
        -- Necessary for multidimensional segmented reduction.
        (inf,sup) <- A.unpair <$> case shr of
                                    ShapeRsnoc ShapeRz -> return (A.pair u v)
                                    _ -> do q <- A.quot integralType s ss
                                            a <- A.mul numType q sz
                                            A.pair <$> A.add numType u a
                                                   <*> A.add numType v a
  
        void $
          if (TupRunit, A.eq singleType inf sup)
            -- This segment is empty. If this is an exclusive reduction the
            -- first thread writes out the initial element for this segment.
            then do
              case mseed of
                Nothing -> return (lift TupRunit ())
                Just z  -> do
                  when (A.eq singleType tid (liftInt32 0)) $ writeArray TypeInt arrOut s =<< z
                  return (lift TupRunit ())
  
            -- This is a non-empty segment.
            else do
              -- Step 1: initialise local sums
              --
              -- NOTE: We require all threads to enter this branch and execute the
              -- first step, even if they do not have a valid element and must
              -- return 'undef'. If we attempt to skip this entire section for
              -- non-participating threads (i.e. 'when (i0 < sup)'), it seems that
              -- those threads die and will not participate in the computation of
              -- _any_ further segment. I'm not sure if this is a CUDA oddity
              -- (e.g. we must have all threads convergent on __syncthreads) or
              -- a bug in NVPTX / ptxas.
              --
              i0 <- A.add numType inf =<< int tid
              x0 <- if (tp, A.lt singleType i0 sup)
                      then app1 (delayedLinearIndex arrIn) i0
                      else let
                               go :: TypeR a -> Operands a
                               go TupRunit       = OP_Unit
                               go (TupRpair a b) = OP_Pair (go a) (go b)
                               go (TupRsingle t) = ir t (undef t)
                           in
                           return $ go tp
  
              bd  <- int =<< blockDim
              v0  <- A.sub numType sup inf
              v0' <- i32 v0
              r0  <- if (tp, A.gte singleType v0 bd)
                       then Fold.reduceBlock dev tp combine Nothing    x0
                       else Fold.reduceBlock dev tp combine (Just v0') x0
  
              -- Step 2: keep walking over the input
              nxt <- A.add numType inf bd
              r   <- iterFromStepTo tp nxt bd sup r0 $ \offset r -> do
  
                       -- Wait for threads to catch up before starting the next stripe
                       __syncthreads
  
                       i' <- A.add numType offset =<< int tid
                       v' <- A.sub numType sup offset
                       r' <- if (tp, A.gte singleType v' bd)
                               -- All threads in the block are in bounds, so we
                               -- can avoid bounds checks.
                               then do
                                 x <- app1 (delayedLinearIndex arrIn) i'
                                 y <- Fold.reduceBlock dev tp combine Nothing x
                                 return y
  
                               -- Not all threads are valid. Note that we still
                               -- have all threads enter the reduction procedure
                               -- to avoid thread divergence on synchronisation
                               -- points, similar to the above NOTE.
                               else do
                                 x <- if (tp, A.lt singleType i' sup)
                                        then app1 (delayedLinearIndex arrIn) i'
                                        else let
                                                 go :: TypeR a -> Operands a
                                                 go TupRunit       = OP_Unit
                                                 go (TupRpair a b) = OP_Pair (go a) (go b)
                                                 go (TupRsingle t) = ir t (undef t)
                                             in
                                             return $ go tp
  
                                 z <- i32 v'
                                 y <- Fold.reduceBlock dev tp combine (Just z) x
                                 return y
  
                       -- first thread incorporates the result from the previous
                       -- iteration
                       if (tp, A.eq singleType tid (liftInt32 0))
                         then app2 combine r r'
                         else return r'
  
              -- Step 3: Thread zero writes the aggregate reduction for this
              -- segment to memory. If this is an exclusive fold combine with the
              -- initial element as well.
              when (A.eq singleType tid (liftInt32 0)) $
               writeArray TypeInt arrOut s =<<
                 case mseed of
                   Nothing -> return r
                   Just z  -> flip (app2 combine) r =<< z  -- Note: initial element on the left
  
              return (lift TupRunit ())
  
      return_
  
  
  -- This implementation assumes that the segments array represents the offset
  -- indices to the source array, rather than the lengths of each segment. The
  -- segment-offset approach is required for parallel implementations.
  --
  -- Each segment is computed by a single warp, meaning we don't have to worry
  -- about inter- or intra-block synchronisation.
  --
  mkFoldSegP_warp
      :: forall aenv sh i e.
         UID
      -> Gamma          aenv
      -> ArrayR (Array (sh, Int) e)
      -> IntegralType i
      -> IRFun2     PTX aenv (e -> e -> e)
      -> MIRExp     PTX aenv e
      -> MIRDelayed PTX aenv (Array (sh, Int) e)
      -> MIRDelayed PTX aenv (Segments i)
      -> CodeGen    PTX      (IROpenAcc PTX aenv (Array (sh, Int) e))
  mkFoldSegP_warp uid aenv repr@(ArrayR shr tp) intTp combine mseed marr mseg = do
    dev <- liftCodeGen $ gets ptxDeviceProperties
    --
    let
        (arrOut, paramOut)  = mutableArray repr "out"
        (arrIn,  paramIn)   = delayedArray "in"  marr
        (arrSeg, paramSeg)  = delayedArray "seg" mseg
        paramEnv            = envParam aenv
        --
        config              = launchConfig dev (CUDA.decWarp dev) dsmem grid gridQ
          where dsmem _n    = 0
        --
        grid n m            = multipleOf n (m `P.quot` ws)
        gridQ               = [|| \n m -> $$multipleOfQ n (m `P.quot` ws) ||]
        --
        ws                  = CUDA.warpSize dev
  
        int32 :: Integral a => a -> Operands Int32
        int32 = liftInt32 . P.fromIntegral
    --
    makeOpenAccWith config uid "foldSeg_warp" (paramOut ++ paramIn ++ paramSeg ++ paramEnv) $ do
  
      -- Each warp works independently.
      -- Determine the ID of this warp within the thread block.
      tid   <- threadIdx
      wid   <- A.quot integralType tid (int32 ws)
  
      -- Number of warps per thread block
      bd    <- blockDim
      wpb   <- A.quot integralType bd (int32 ws)
  
      -- ID of this warp within the grid
      bid   <- blockIdx
      gwid  <- do a <- A.mul numType bid wpb
                  b <- A.add numType wid a
                  return b
  
      -- Compute the number of segments and size of the innermost dimension. These
      -- are required if we are reducing a rank-2 or higher array, to properly
      -- compute the start and end indices of the portion of the array this warp
      -- reduces. Note that this is a segment-offset array computed by 'scanl (+) 0'
      -- of the segment length array, so its size has increased by one.
      --
      sz    <- indexHead <$> delayedExtent arrIn
      ss    <- do a <- indexHead <$> delayedExtent arrSeg
                  b <- A.sub numType a (liftInt 1)
                  return b
  
      -- Each thread reduces a segment independently
      s0    <- int gwid
      gd    <- int =<< gridDim
      wpb'  <- int wpb
      step  <- A.mul numType wpb' gd
      end   <- shapeSize shr (irArrayShape arrOut)
      imapFromStepTo s0 step end $ \s -> do
  
        __syncwarp
  
        -- The first two threads of the warp determine the indices of the segments
        -- array that we will reduce between and distribute those values to the
        -- other threads in the warp
        lane <- laneId
        idx  <- if (TupRsingle scalarTypeInt, A.lt singleType lane (liftInt32 2))
                   then do
                     a <- case shr of
                            ShapeRsnoc ShapeRz -> return s
                            _                  -> A.rem integralType s ss
                     b <- A.add numType a =<< int lane
                     c <- app1 (delayedLinearIndex arrSeg) b
                     d <- A.fromIntegral intTp numType c
                     return d
                   else
                     return (ir integralType (undef scalarType))
  
        __syncwarp
  
        -- Determine the index range of the input array we will reduce over.
        -- Necessary for multidimensional segmented reduction.
        (inf,sup) <- do
          u <- __shfl_idx (TupRsingle scalarTypeInt) idx (liftWord32 0)
  
          v <- __shfl_idx (TupRsingle scalarTypeInt) idx (liftWord32 1)
          A.unpair <$> case shr of
                         ShapeRsnoc ShapeRz -> return (A.pair u v)
                         _ -> do q <- A.quot integralType s ss
                                 a <- A.mul numType q sz
                                 A.pair <$> A.add numType u a
                                        <*> A.add numType v a
  
        __syncwarp
  
        void $
          if (TupRunit, A.eq singleType inf sup)
            -- This segment is empty. If this is an exclusive reduction the first
            -- lane writes out the initial element for this segment.
            then do
              case mseed of
                Nothing -> return (lift TupRunit ())
                Just z  -> do
                  when (A.eq singleType lane (liftInt32 0)) $ writeArray TypeInt arrOut s =<< z
                  return (lift TupRunit ())
  
            -- This is a non-empty segment.
            else do
              -- Step 1: initialise local sums
              --
              -- See comment above why we initialise the loop in this way
              --
              i0  <- A.add numType inf =<< int lane
              x0  <- if (tp, A.lt singleType i0 sup)
                       then app1 (delayedLinearIndex arrIn) i0
                       else let
                                go :: TypeR a -> Operands a
                                go TupRunit       = OP_Unit
                                go (TupRpair a b) = OP_Pair (go a) (go b)
                                go (TupRsingle t) = ir t (undef t)
                            in
                            return $ go tp
  
              v0  <- A.sub numType sup inf
              v0' <- i32 v0
              r0  <- if (tp, A.gte singleType v0 (liftInt ws))
                       then reduceWarp dev tp combine Nothing    x0
                       else reduceWarp dev tp combine (Just v0') x0
  
              -- Step 2: Keep walking over the rest of the segment
              nx  <- A.add numType inf (liftInt ws)
              r   <- iterFromStepTo tp nx (liftInt ws) sup r0 $ \offset r -> do
  
                      -- __syncwarp
                      __syncthreads -- TLM: why is this necessary?
  
                      i' <- A.add numType offset =<< int lane
                      v' <- A.sub numType sup offset
                      r' <- if (tp, A.gte singleType v' (liftInt ws))
                              then do
                                -- All lanes are in bounds, so avoid bounds checks
                                x <- app1 (delayedLinearIndex arrIn) i'
                                y <- reduceWarp dev tp combine Nothing x
                                return y
  
                              else do
                                x <- if (tp, A.lt singleType i' sup)
                                       then app1 (delayedLinearIndex arrIn) i'
                                       else let
                                                go :: TypeR a -> Operands a
                                                go TupRunit       = OP_Unit
                                                go (TupRpair a b) = OP_Pair (go a) (go b)
                                                go (TupRsingle t) = ir t (undef t)
                                            in
                                            return $ go tp
  
                                z <- i32 v'
                                y <- reduceWarp dev tp combine (Just z) x
                                return y
  
                      -- The first lane incorporates the result from the previous
                      -- iteration
                      if (tp, A.eq singleType lane (liftInt32 0))
                        then app2 combine r r'
                        else return r'
  
              -- Step 3: Lane zero writes the aggregate reduction for this
              -- segment to memory. If this is an exclusive reduction, also
              -- combine with the initial element
              when (A.eq singleType lane (liftInt32 0)) $
                writeArray TypeInt arrOut s =<<
                  case mseed of
                    Nothing -> return r
                    Just z  -> flip (app2 combine) r =<< z    -- Note: initial element on the left
  
              return (lift TupRunit ())
  
      return_
  
  
  i32 :: IsIntegral i => Operands i -> CodeGen PTX (Operands Int32)
  i32 = A.fromIntegral integralType numType
  
  int :: IsIntegral i => Operands i -> CodeGen PTX (Operands Int)
  int = A.fromIntegral integralType numType
  
  reduceWarp
      :: forall aenv e.
         DeviceProperties                         -- ^ properties of the target device
      -> TypeR e
      -> IRFun2 PTX aenv (e -> e -> e)            -- ^ combination function
      -> Maybe (Operands Int32)                   -- ^ number of items that will be reduced by this warp, otherwise all lanes are valid
      -> Operands e                               -- ^ calling thread's input element
      -> CodeGen PTX (Operands e)                 -- ^ warp-wide reduction using the specified operator (lane 0 only)
  reduceWarp dev t c = Fold.reduceWarp dev t c
  
  