content: {-# LANGUAGE GADTs               #-}
  {-# LANGUAGE OverloadedStrings   #-}
  {-# LANGUAGE PatternGuards       #-}
  {-# LANGUAGE RebindableSyntax    #-}
  {-# LANGUAGE RecordWildCards     #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TemplateHaskell     #-}
  {-# LANGUAGE TypeApplications    #-}
  {-# LANGUAGE TypeOperators       #-}
  {-# LANGUAGE ViewPatterns        #-}
  -- |
  -- Module      : Data.Array.Accelerate.LLVM.PTX.CodeGen.Scan
  -- 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.Scan (
  
    mkScan, mkScan',
  
  ) where
  
  import Data.Array.Accelerate.AST                                    ( Direction(..) )
  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
  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 Data.Array.Accelerate.LLVM.PTX.CodeGen.Generate
  import Data.Array.Accelerate.LLVM.PTX.Target
  
  import LLVM.AST.Type.Representation
  
  import qualified Foreign.CUDA.Analysis                              as CUDA
  
  import Control.Applicative
  import Control.Monad                                                ( (>=>), void )
  import Control.Monad.State                                          ( gets )
  import Data.String                                                  ( fromString )
  import Data.Coerce                                                  as Safe
  import Data.Bits                                                    as P
  import Prelude                                                      as P hiding ( last )
  
  
  
  -- 'Data.List.scanl' or 'Data.List.scanl1' style exclusive scan, but with the
  -- restriction that the combination function must be associative to enable
  -- efficient parallel implementation.
  --
  -- > scanl (+) 10 (use $ fromList (Z :. 10) [0..])
  -- >
  -- > ==> Array (Z :. 11) [10,10,11,13,16,20,25,31,38,46,55]
  --
  mkScan
      :: forall aenv sh e.
         UID
      -> Gamma            aenv
      -> ArrayR (Array (sh, Int) e)
      -> Direction
      -> IRFun2       PTX aenv (e -> e -> e)
      -> Maybe (IRExp PTX aenv e)
      -> MIRDelayed   PTX aenv (Array (sh, Int) e)
      -> CodeGen      PTX      (IROpenAcc PTX aenv (Array (sh, Int) e))
  mkScan uid aenv repr dir combine seed arr
    = foldr1 (+++) <$> sequence (codeScan ++ codeFill)
  
    where
      codeScan = case repr of
        ArrayR (ShapeRsnoc ShapeRz) tp -> [ mkScanAllP1 dir uid aenv tp   combine seed arr
                                          , mkScanAllP2 dir uid aenv tp   combine
                                          , mkScanAllP3 dir uid aenv tp   combine seed
                                          ]
        _                              -> [ mkScanDim   dir uid aenv repr combine seed arr
                                          ]
      codeFill = case seed of
        Just s  -> [ mkScanFill uid aenv repr s ]
        Nothing -> []
  
  
  -- Variant of 'scanl' where the final result is returned in a separate array.
  --
  -- > scanr' (+) 10 (use $ fromList (Z :. 10) [0..])
  -- >
  -- > ==> ( Array (Z :. 10) [10,10,11,13,16,20,25,31,38,46]
  --       , Array Z [55]
  --       )
  --
  mkScan'
      :: forall aenv sh e.
         UID
      -> Gamma          aenv
      -> ArrayR (Array (sh, Int) e)
      -> Direction
      -> IRFun2     PTX aenv (e -> e -> e)
      -> IRExp      PTX aenv e
      -> MIRDelayed PTX aenv (Array (sh, Int) e)
      -> CodeGen    PTX      (IROpenAcc PTX aenv (Array (sh, Int) e, Array sh e))
  mkScan' uid aenv repr dir combine seed arr
    | ArrayR (ShapeRsnoc ShapeRz) tp <- repr
    = foldr1 (+++) <$> sequence [ mkScan'AllP1 dir uid aenv tp combine seed arr
                                , mkScan'AllP2 dir uid aenv tp combine
                                , mkScan'AllP3 dir uid aenv tp combine
                                , mkScan'Fill      uid aenv repr seed
                                ]
    --
    | otherwise
    = (+++) <$> mkScan'Dim dir uid aenv repr combine seed arr
            <*> mkScan'Fill    uid aenv repr seed
  
  
  -- Device wide scans
  -- -----------------
  --
  -- This is a classic two-pass algorithm which proceeds in two phases and
  -- requires ~4n data movement to global memory. In future we would like to
  -- replace this with a single pass algorithm.
  --
  
  -- Parallel scan, step 1.
  --
  -- Threads scan a stripe of the input into a temporary array, incorporating the
  -- initial element and any fused functions on the way. The final reduction
  -- result of this chunk is written to a separate array.
  --
  mkScanAllP1
      :: forall aenv e.
         Direction
      -> UID
      -> Gamma          aenv                      -- ^ array environment
      -> TypeR e
      -> IRFun2     PTX aenv (e -> e -> e)        -- ^ combination function
      -> MIRExp     PTX aenv e                    -- ^ seed element, if this is an exclusive scan
      -> MIRDelayed PTX aenv (Vector e)           -- ^ input data
      -> CodeGen    PTX (IROpenAcc PTX aenv (Vector e))
  mkScanAllP1 dir uid aenv tp combine mseed marr = do
    dev <- liftCodeGen $ gets ptxDeviceProperties
    --
    let
        (arrOut, paramOut)  = mutableArray (ArrayR dim1 tp) "out"
        (arrTmp, paramTmp)  = mutableArray (ArrayR dim1 tp) "tmp"
        (arrIn,  paramIn)   = delayedArray "in" marr
        end                 = indexHead (irArrayShape arrTmp)
        paramEnv            = envParam aenv
        --
        config              = launchConfig dev (CUDA.incWarp dev) (scanSMemSize dev tp) const [|| const ||]
    --
    makeOpenAccWith config uid "scanP1" (paramTmp ++ paramOut ++ paramIn ++ paramEnv) $ do
  
      -- Size of the input array
      sz  <- indexHead <$> delayedExtent arrIn
  
      -- A thread block scans a non-empty stripe of the input, storing the final
      -- block-wide aggregate into a separate array
      --
      -- For exclusive scans, thread 0 of segment 0 must incorporate the initial
      -- element into the input and output. Threads shuffle their indices
      -- appropriately.
      --
      bid <- blockIdx
      gd  <- gridDim
      gd' <- int gd
      s0  <- int bid
  
      -- iterating over thread-block-wide segments
      -- Note that 'end' is a multiple of the gd', and the control flow is thus uniform in the loop.
      -- This is set in scanAllOp in Data.Array.Accelerate.LLVM.PTX.Execute.
      -- Hence we can run __syncthreads safely.
      imapFromStepTo s0 gd' end $ \chunk -> do
  
        -- Make sure all threads have finished previous iterations,
        -- so we can reuse (and overwrite) shared memory.
        __syncthreads
  
        bd    <- blockDim
        bd'   <- int bd
        inf   <- A.mul numType chunk bd'
  
        -- index i* is the index that this thread will read data from. Recall that
        -- the supremum index is exclusive
        tid   <- threadIdx
        tid'  <- int tid
        i0    <- case dir of
                   LeftToRight -> A.add numType inf tid'
                   RightToLeft -> do x <- A.sub numType sz inf
                                     y <- A.sub numType x tid'
                                     z <- A.sub numType y (liftInt 1)
                                     return z
  
        -- index j* is the index that we write to. Recall that for exclusive scans
        -- the output array is one larger than the input; the initial element will
        -- be written into this spot by thread 0 of the first thread block.
        j0    <- case mseed of
                   Nothing -> return i0
                   Just _  -> case dir of
                                LeftToRight -> A.add numType i0 (liftInt 1)
                                RightToLeft -> return i0
  
        -- If this thread has input, read data and participate in thread-block scan
        let valid i = case dir of
                        LeftToRight -> A.lt  singleType i sz
                        RightToLeft -> A.gte singleType i (liftInt 0)
  
        when (valid i0) $ do
          x0 <- app1 (delayedLinearIndex arrIn) i0
          x1 <- case mseed of
                  Nothing   -> return x0
                  Just seed ->
                    if (tp, A.eq singleType tid (liftInt32 0) `A.land'` A.eq singleType chunk (liftInt 0))
                      then do
                        z <- seed
                        case dir of
                          LeftToRight -> writeArray TypeInt32 arrOut (liftInt32 0) z >> app2 combine z x0
                          RightToLeft -> writeArray TypeInt   arrOut sz            z >> app2 combine x0 z
                      else
                        return x0
  
          n  <- A.sub numType sz inf
          n' <- i32 n
          x2 <- if (tp, A.gte singleType n bd')
                  then scanBlock dir dev tp combine Nothing   x1
                  else scanBlock dir dev tp combine (Just n') x1
  
          -- Write this thread's scan result to memory
          writeArray TypeInt arrOut j0 x2
  
          -- The last thread also writes its result---the aggregate for this
          -- thread block---to the temporary partial sums array. This is only
          -- necessary for full blocks in a multi-block scan; the final
          -- partially-full tile does not have a successor block.
          last <- A.sub numType bd (liftInt32 1)
          when (A.gt singleType gd (liftInt32 1) `land'` A.eq singleType tid last) $
            case dir of
              LeftToRight -> writeArray TypeInt arrTmp chunk x2
              RightToLeft -> do u <- A.sub numType end chunk
                                v <- A.sub numType u (liftInt 1)
                                writeArray TypeInt arrTmp v x2
  
      return_
  
  
  -- Parallel scan, step 2
  --
  -- A single thread block performs a scan of the per-block aggregates computed in
  -- step 1. This gives the per-block prefix which must be added to each element
  -- in step 3.
  --
  mkScanAllP2
      :: forall aenv e.
         Direction
      -> UID
      -> Gamma       aenv                         -- ^ array environment
      -> TypeR e
      -> IRFun2  PTX aenv (e -> e -> e)           -- ^ combination function
      -> CodeGen PTX      (IROpenAcc PTX aenv (Vector e))
  mkScanAllP2 dir uid aenv tp combine = do
    dev <- liftCodeGen $ gets ptxDeviceProperties
    --
    let
        (arrTmp, paramTmp)  = mutableArray (ArrayR dim1 tp) "tmp"
        paramEnv            = envParam aenv
        start               = liftInt 0
        end                 = indexHead (irArrayShape arrTmp)
        --
        config              = launchConfig dev (CUDA.incWarp dev) (scanSMemSize dev tp) grid gridQ
        grid _ _            = 1
        gridQ               = [|| \_ _ -> 1 ||]
    --
    makeOpenAccWith config uid "scanP2" (paramTmp ++ paramEnv) $ do
  
      -- The first and last threads of the block need to communicate the
      -- block-wide aggregate as a carry-in value across iterations.
      --
      -- TODO: We could optimise this a bit if we can get access to the shared
      -- memory area used by 'scanBlockSMem', and from there directly read the
      -- value computed by the last thread.
      carry <- staticSharedMem tp 1
  
      bd    <- blockDim
      bd'   <- int bd
  
      imapFromStepTo start bd' end $ \offset -> do
  
        -- Index of the partial sums array that this thread will process.
        tid   <- threadIdx
        tid'  <- int tid
        i0    <- case dir of
                   LeftToRight -> A.add numType offset tid'
                   RightToLeft -> do x <- A.sub numType end offset
                                     y <- A.sub numType x tid'
                                     z <- A.sub numType y (liftInt 1)
                                     return z
  
        let valid i = case dir of
                        LeftToRight -> A.lt  singleType i end
                        RightToLeft -> A.gte singleType i start
  
        -- wait for the carry-in value to be updated
        __syncthreads
  
        when (valid i0) $ do
          x0 <- readArray TypeInt arrTmp i0
          x1 <- if (tp, A.gt singleType offset (liftInt 0) `land'` A.eq singleType tid (liftInt32 0))
                  then do
                    c <- readArray TypeInt32 carry (liftInt32 0)
                    case dir of
                      LeftToRight -> app2 combine c x0
                      RightToLeft -> app2 combine x0 c
                  else do
                    return x0
  
          n  <- A.sub numType end offset
          n' <- i32 n
          x2 <- if (tp, A.gte singleType n bd')
                  then scanBlock dir dev tp combine Nothing   x1
                  else scanBlock dir dev tp combine (Just n') x1
  
          -- Update the temporary array with this thread's result
          writeArray TypeInt arrTmp i0 x2
  
          -- The last thread writes the carry-out value. If the last thread is not
          -- active, then this must be the last stripe anyway.
          last <- A.sub numType bd (liftInt32 1)
          when (A.eq singleType tid last) $
            writeArray TypeInt32 carry (liftInt32 0) x2
  
      return_
  
  
  -- Parallel scan, step 3.
  --
  -- Threads combine every element of the partial block results with the carry-in
  -- value computed in step 2.
  --
  mkScanAllP3
      :: forall aenv e.
         Direction
      -> UID
      -> Gamma       aenv                         -- ^ array environment
      -> TypeR e
      -> IRFun2  PTX aenv (e -> e -> e)           -- ^ combination function
      -> MIRExp  PTX aenv e                       -- ^ seed element, if this is an exclusive scan
      -> CodeGen PTX      (IROpenAcc PTX aenv (Vector e))
  mkScanAllP3 dir uid aenv tp combine mseed = do
    dev <- liftCodeGen $ gets ptxDeviceProperties
    --
    let
        (arrOut, paramOut)  = mutableArray (ArrayR dim1 tp) "out"
        (arrTmp, paramTmp)  = mutableArray (ArrayR dim1 tp) "tmp"
        paramEnv            = envParam aenv
        --
        stride              = local     (TupRsingle scalarTypeInt) "ix.stride"
        paramStride         = parameter (TupRsingle scalarTypeInt) "ix.stride"
        --
        config              = launchConfig dev (CUDA.incWarp dev) (const 0) const [|| const ||]
    --
    makeOpenAccWith config uid "scanP3" (paramTmp ++ paramOut ++ paramStride ++ paramEnv) $ do
  
      sz  <- return $ indexHead (irArrayShape arrOut)
      tid <- int =<< threadIdx
  
      -- Threads that will never contribute can just exit immediately. The size of
      -- each chunk is set by the block dimension of the step 1 kernel, which may
      -- be different from the block size of this kernel.
      when (A.lt singleType tid stride) $ do
  
        -- Iterate over the segments computed in phase 1. Note that we have one
        -- fewer chunk to process because the first has no carry-in.
        bid <- int =<< blockIdx
        gd  <- int =<< gridDim
        end <- A.sub numType (indexHead (irArrayShape arrTmp)) (liftInt 1)
  
        imapFromStepTo bid gd end $ \chunk -> do
  
          -- Determine the start and end indicies of this chunk to which we will
          -- carry-in the value. Returned for left-to-right traversal.
          (inf,sup) <- case dir of
                         LeftToRight -> do
                           a <- A.add numType chunk (liftInt 1)
                           b <- A.mul numType stride a
                           case mseed of
                             Just{}  -> do
                               c <- A.add numType    b (liftInt 1)
                               d <- A.add numType    c stride
                               e <- A.min singleType d sz
                               return (c,e)
                             Nothing -> do
                               c <- A.add numType    b stride
                               d <- A.min singleType c sz
                               return (b,d)
                         RightToLeft -> do
                           a <- A.sub numType end chunk
                           b <- A.mul numType stride a
                           c <- A.sub numType sz b
                           case mseed of
                             Just{}  -> do
                               d <- A.sub numType    c (liftInt 1)
                               e <- A.sub numType    d stride
                               f <- A.max singleType e (liftInt 0)
                               return (f,d)
                             Nothing -> do
                               d <- A.sub numType    c stride
                               e <- A.max singleType d (liftInt 0)
                               return (e,c)
  
          -- Read the carry-in value
          carry     <- case dir of
                         LeftToRight -> readArray TypeInt arrTmp chunk
                         RightToLeft -> do
                           a <- A.add numType chunk (liftInt 1)
                           b <- readArray TypeInt arrTmp a
                           return b
  
          -- Apply the carry-in value to each element in the chunk
          bd        <- int =<< blockDim
          i0        <- A.add numType inf tid
          imapFromStepTo i0 bd sup $ \i -> do
            v <- readArray TypeInt arrOut i
            u <- case dir of
                   LeftToRight -> app2 combine carry v
                   RightToLeft -> app2 combine v carry
            writeArray TypeInt arrOut i u
  
      return_
  
  
  -- Parallel scan', step 1.
  --
  -- Similar to mkScanAllP1. Threads scan a stripe of the input into a temporary
  -- array, incorporating the initial element and any fused functions on the way.
  -- The final reduction result of this chunk is written to a separate array.
  --
  mkScan'AllP1
      :: forall aenv e.
         Direction
      -> UID
      -> Gamma          aenv
      -> TypeR e
      -> IRFun2     PTX aenv (e -> e -> e)
      -> IRExp      PTX aenv e
      -> MIRDelayed PTX aenv (Vector e)
      -> CodeGen    PTX      (IROpenAcc PTX aenv (Vector e, Scalar e))
  mkScan'AllP1 dir uid aenv tp combine seed marr = do
    dev <- liftCodeGen $ gets ptxDeviceProperties
    --
    let
        (arrOut, paramOut)  = mutableArray (ArrayR dim1 tp) "out"
        (arrTmp, paramTmp)  = mutableArray (ArrayR dim1 tp) "tmp"
        (arrIn,  paramIn)   = delayedArray "in" marr
        end                 = indexHead (irArrayShape arrTmp)
        paramEnv            = envParam aenv
        --
        config              = launchConfig dev (CUDA.incWarp dev) (scanSMemSize dev tp) const [|| const ||]
    --
    makeOpenAccWith config uid "scanP1" (paramTmp ++ paramOut ++ paramIn ++ paramEnv) $ do
  
      -- Size of the input array
      sz  <- indexHead <$> delayedExtent arrIn
  
      -- A thread block scans a non-empty stripe of the input, storing the partial
      -- result and the final block-wide aggregate
      bid <- int =<< blockIdx
      gd  <- int =<< gridDim
  
      -- iterate over thread-block wide segments
      -- Note that 'end' is a multiple of the gd', and the control flow is thus uniform in the loop.
      -- This is set in scan'AllOp in Data.Array.Accelerate.LLVM.PTX.Execute.
      -- Hence we can run __syncthreads safely.
      imapFromStepTo bid gd end $ \seg -> do
  
        -- Make sure all threads have finished previous iterations,
        -- so we can reuse (and overwrite) shared memory.
        __syncthreads
  
        bd  <- int =<< blockDim
        inf <- A.mul numType seg bd
  
        -- i* is the index that this thread will read data from
        tid <- int =<< threadIdx
        i0  <- case dir of
                 LeftToRight -> A.add numType inf tid
                 RightToLeft -> do x <- A.sub numType sz inf
                                   y <- A.sub numType x tid
                                   z <- A.sub numType y (liftInt 1)
                                   return z
  
        -- j* is the index this thread will write to. This is just shifted by one
        -- to make room for the initial element
        j0  <- case dir of
                 LeftToRight -> A.add numType i0 (liftInt 1)
                 RightToLeft -> A.sub numType i0 (liftInt 1)
  
        -- If this thread has input it participates in the scan
        let valid i = case dir of
                        LeftToRight -> A.lt  singleType i sz
                        RightToLeft -> A.gte singleType i (liftInt 0)
  
        when (valid i0) $ do
          x0 <- app1 (delayedLinearIndex arrIn) i0
  
          -- Thread 0 of the first segment must also evaluate and store the
          -- initial element
          ti <- threadIdx
          x1 <- if (tp, A.eq singleType ti (liftInt32 0) `A.land'` A.eq singleType seg (liftInt 0))
                  then do
                    z <- seed
                    writeArray TypeInt arrOut i0 z
                    case dir of
                      LeftToRight -> app2 combine z x0
                      RightToLeft -> app2 combine x0 z
                  else
                    return x0
  
          -- Block-wide scan
          n  <- A.sub numType sz inf
          n' <- i32 n
          x2 <- if (tp, A.gte singleType n bd)
                  then scanBlock dir dev tp combine Nothing   x1
                  else scanBlock dir dev tp combine (Just n') x1
  
          -- Write this thread's scan result to memory. Recall that we had to make
          -- space for the initial element, so the very last thread does not store
          -- its result here.
          case dir of
            LeftToRight -> when (A.lt  singleType j0 sz)          $ writeArray TypeInt arrOut j0 x2
            RightToLeft -> when (A.gte singleType j0 (liftInt 0)) $ writeArray TypeInt arrOut j0 x2
  
          -- Last active thread writes its result to the partial sums array. These
          -- will be used to compute the carry-in value in step 2.
          m  <- do x <- A.min singleType n bd
                   y <- A.sub numType x (liftInt 1)
                   return y
          when (A.eq singleType tid m) $
            case dir of
              LeftToRight -> writeArray TypeInt arrTmp seg x2
              RightToLeft -> do x <- A.sub numType end seg
                                y <- A.sub numType x (liftInt 1)
                                writeArray TypeInt arrTmp y x2
  
      return_
  
  
  -- Parallel scan', step 2
  --
  -- A single thread block performs an inclusive scan of the partial sums array to
  -- compute the per-block carry-in values, as well as the final reduction result.
  --
  mkScan'AllP2
      :: forall aenv e.
         Direction
      -> UID
      -> Gamma aenv
      -> TypeR e
      -> IRFun2 PTX aenv (e -> e -> e)
      -> CodeGen PTX (IROpenAcc PTX aenv (Vector e, Scalar e))
  mkScan'AllP2 dir uid aenv tp combine = do
    dev <- liftCodeGen $ gets ptxDeviceProperties
    --
    let
        (arrTmp, paramTmp)  = mutableArray (ArrayR dim1 tp) "tmp"
        (arrSum, paramSum)  = mutableArray (ArrayR dim0 tp) "sum"
        paramEnv            = envParam aenv
        start               = liftInt 0
        end                 = indexHead (irArrayShape arrTmp)
        --
        config              = launchConfig dev (CUDA.incWarp dev) (scanSMemSize dev tp) grid gridQ
        grid _ _            = 1
        gridQ               = [|| \_ _ -> 1 ||]
    --
    makeOpenAccWith config uid "scanP2" (paramTmp ++ paramSum ++ paramEnv) $ do
  
      -- The first and last threads of the block need to communicate the
      -- block-wide aggregate as a carry-in value across iterations.
      carry <- staticSharedMem tp 1
  
      -- A single thread block iterates over the per-block partial results from
      -- step 1
      tid   <- threadIdx
      tid'  <- int tid
      bd    <- int =<< blockDim
  
      imapFromStepTo start bd end $ \offset -> do
  
        i0  <- case dir of
                 LeftToRight -> A.add numType offset tid'
                 RightToLeft -> do x <- A.sub numType end offset
                                   y <- A.sub numType x tid'
                                   z <- A.sub numType y (liftInt 1)
                                   return z
  
        let valid i = case dir of
                        LeftToRight -> A.lt  singleType i end
                        RightToLeft -> A.gte singleType i start
  
        -- wait for the carry-in value to be updated
        __syncthreads
  
        x0 <- if (tp, valid i0)
                then readArray TypeInt arrTmp i0
                else
                  return $ tupUndef tp
  
        x1 <- if (tp, A.gt singleType offset (liftInt 0) `A.land'` A.eq singleType tid (liftInt32 0))
                then do
                  c <- readArray TypeInt32 carry (liftInt32 0)
                  case dir of
                    LeftToRight -> app2 combine c x0
                    RightToLeft -> app2 combine x0 c
                else
                  return x0
  
        n  <- A.sub numType end offset
        n' <- i32 n
        x2 <- if (tp, A.gte singleType n bd)
                then scanBlock dir dev tp combine Nothing   x1
                else scanBlock dir dev tp combine (Just n') x1
  
        -- Update the partial results array
        when (valid i0) $
          writeArray TypeInt arrTmp i0 x2
  
        -- The last active thread saves its result as the carry-out value.
        m  <- do x <- A.min singleType bd n
                 y <- A.sub numType x (liftInt 1)
                 z <- i32 y
                 return z
        when (A.eq singleType tid m) $
          writeArray TypeInt32 carry (liftInt32 0) x2
  
      -- First thread stores the final carry-out values at the final reduction
      -- result for the entire array
      __syncthreads
  
      when (A.eq singleType tid (liftInt32 0)) $
        writeArray TypeInt32 arrSum (liftInt32 0) =<< readArray TypeInt32 carry (liftInt32 0)
  
      return_
  
  
  -- Parallel scan', step 3.
  --
  -- Threads combine every element of the partial block results with the carry-in
  -- value computed in step 2.
  --
  mkScan'AllP3
      :: forall aenv e.
         Direction
      -> UID
      -> Gamma aenv                                   -- ^ array environment
      -> TypeR e
      -> IRFun2 PTX aenv (e -> e -> e)                -- ^ combination function
      -> CodeGen PTX (IROpenAcc PTX aenv (Vector e, Scalar e))
  mkScan'AllP3 dir uid aenv tp combine = do
    dev <- liftCodeGen $ gets ptxDeviceProperties
    --
    let
        (arrOut, paramOut)  = mutableArray (ArrayR dim1 tp) "out"
        (arrTmp, paramTmp)  = mutableArray (ArrayR dim1 tp) "tmp"
        paramEnv            = envParam aenv
        --
        stride              = local     (TupRsingle scalarTypeInt) "ix.stride"
        paramStride         = parameter (TupRsingle scalarTypeInt) "ix.stride"
        --
        config              = launchConfig dev (CUDA.incWarp dev) (const 0) const [|| const ||]
    --
    makeOpenAccWith config uid "scanP3" (paramTmp ++ paramOut ++ paramStride ++ paramEnv) $ do
  
      sz  <- return $ indexHead (irArrayShape arrOut)
      tid <- int =<< threadIdx
  
      when (A.lt singleType tid stride) $ do
  
        bid <- int =<< blockIdx
        gd  <- int =<< gridDim
        end <- A.sub numType (indexHead (irArrayShape arrTmp)) (liftInt 1)
  
        imapFromStepTo bid gd end $ \chunk -> do
  
          (inf,sup) <- case dir of
                         LeftToRight -> do
                           a <- A.add numType    chunk  (liftInt 1)
                           b <- A.mul numType    stride a
                           c <- A.add numType    b      (liftInt 1)
                           d <- A.add numType    c      stride
                           e <- A.min singleType d      sz
                           return (c,e)
                         RightToLeft -> do
                           a <- A.sub numType    end    chunk
                           b <- A.mul numType    stride a
                           c <- A.sub numType    sz     b
                           d <- A.sub numType    c      (liftInt 1)
                           e <- A.sub numType    d      stride
                           f <- A.max singleType e      (liftInt 0)
                           return (f,d)
  
          carry     <- case dir of
                         LeftToRight -> readArray TypeInt arrTmp chunk
                         RightToLeft -> do
                           a <- A.add numType chunk (liftInt 1)
                           b <- readArray TypeInt arrTmp a
                           return b
  
          -- Apply the carry-in value to each element in the chunk
          bd        <- int =<< blockDim
          i0        <- A.add numType inf tid
          imapFromStepTo i0 bd sup $ \i -> do
            v <- readArray TypeInt arrOut i
            u <- case dir of
                   LeftToRight -> app2 combine carry v
                   RightToLeft -> app2 combine v carry
            writeArray TypeInt arrOut i u
  
      return_
  
  
  -- Multidimensional scans
  -- ----------------------
  
  -- Multidimensional scan along the innermost dimension
  --
  -- A thread block individually computes along each innermost dimension. This is
  -- a single-pass operation.
  --
  --  * We can assume that the array is non-empty; exclusive scans with empty
  --    innermost dimension will be instead filled with the seed element via
  --    'mkScanFill'.
  --
  --  * Small but non-empty innermost dimension arrays (size << thread
  --    block size) will have many threads which do no work.
  --
  mkScanDim
      :: forall aenv sh e.
         Direction
      -> UID
      -> Gamma          aenv                          -- ^ array environment
      -> ArrayR (Array (sh, Int) e)
      -> IRFun2     PTX aenv (e -> e -> e)            -- ^ combination function
      -> MIRExp     PTX aenv e                        -- ^ seed element, if this is an exclusive scan
      -> MIRDelayed PTX aenv (Array (sh, Int) e)      -- ^ input data
      -> CodeGen    PTX (IROpenAcc PTX aenv (Array (sh, Int) e))
  mkScanDim dir uid aenv repr@(ArrayR (ShapeRsnoc shr) tp) combine mseed marr = do
    dev <- liftCodeGen $ gets ptxDeviceProperties
    --
    let
        (arrOut, paramOut)  = mutableArray repr "out"
        (arrIn,  paramIn)   = delayedArray "in" marr
        paramEnv            = envParam aenv
        --
        config              = launchConfig dev (CUDA.incWarp dev) (scanSMemSize dev tp) const [|| const ||]
    --
    makeOpenAccWith config uid "scan" (paramOut ++ paramIn ++ paramEnv) $ do
  
      -- The first and last threads of the block need to communicate the
      -- block-wide aggregate as a carry-in value across iterations.
      --
      -- TODO: we could optimise this a bit if we can get access to the shared
      -- memory area used by 'scanBlockSMem', and from there directly read the
      -- value computed by the last thread.
      carry <- staticSharedMem tp 1
  
      -- Size of the input array
      sz  <- indexHead <$> delayedExtent arrIn
  
      -- Thread blocks iterate over the outer dimensions. Threads in a block
      -- cooperatively scan along one dimension, but thread blocks do not
      -- communicate with each other.
      --
      bid <- int =<< blockIdx
      gd  <- int =<< gridDim
      end <- shapeSize shr (indexTail (irArrayShape arrOut))
  
      -- Iterate over the outer dimensions (all but the innermost dimension).
      -- Within this loop we perform a scan over a row (innermost dimension) of
      -- the input.
      --
      -- Since 'bid', 'gd' and 'end' are uniform, the control flow within this
      -- loop is also uniform. We can thus perform __syncthreads in the loop.
      imapFromStepTo bid gd end $ \seg -> do
  
        -- Make sure all threads have finished previous iterations,
        -- so we can reuse (and overwrite) shared memory.
        __syncthreads
  
        -- Index this thread reads from
        tid   <- threadIdx
        tid'  <- int tid
        i0    <- case dir of
                   LeftToRight -> do x <- A.mul numType seg sz
                                     y <- A.add numType x tid'
                                     return y
  
                   RightToLeft -> do x <- A.add numType seg (liftInt 1)
                                     y <- A.mul numType x sz
                                     z <- A.sub numType y tid'
                                     w <- A.sub numType z (liftInt 1)
                                     return w
  
        -- Index this thread writes to
        j0  <- case mseed of
                 Nothing -> return i0
                 Just{}  -> do szp1 <- return $ indexHead (irArrayShape arrOut)
                               case dir of
                                 LeftToRight -> do x <- A.mul numType seg szp1
                                                   y <- A.add numType x tid'
                                                   return y
  
                                 RightToLeft -> do x <- A.add numType seg (liftInt 1)
                                                   y <- A.mul numType x szp1
                                                   z <- A.sub numType y tid'
                                                   w <- A.sub numType z (liftInt 1)
                                                   return w
  
        -- Stride indices by block dimension
        bd  <- blockDim
        bd' <- int bd
        let next ix = case dir of
                        LeftToRight -> A.add numType ix bd'
                        RightToLeft -> A.sub numType ix bd'
  
        -- Initialise this scan segment
        --
        -- If this is an exclusive scan then the first thread just evaluates the
        -- seed element and stores this value into the carry-in slot. All threads
        -- shift their write-to index (j) by one, to make space for this element.
        --
        -- If this is an inclusive scan then do a block-wide scan. The last thread
        -- in the block writes the carry-in value.
        --
        r <-
          case mseed of
            Just seed -> do
              when (A.eq singleType tid (liftInt32 0)) $ do
                z <- seed
                writeArray TypeInt   arrOut j0 z
                writeArray TypeInt32 carry (liftInt32 0) z
              j1 <- case dir of
                     LeftToRight -> A.add numType j0 (liftInt 1)
                     RightToLeft -> A.sub numType j0 (liftInt 1)
              return $ A.trip sz i0 j1
  
            Nothing -> do
              -- We cannot call scanBlock under non-uniform control flow.
              -- Instead, we conditionally read the input, and then
              -- unconditionally call scanBlock.
              x0 <- if (tp, A.lt singleType tid' sz)
                     then app1 (delayedLinearIndex arrIn) i0
                     else return $ tupUndef tp
              n' <- i32 sz
  
              r0 <- if (tp, A.gte singleType sz bd')
                      then scanBlock dir dev tp combine Nothing   x0
                      else scanBlock dir dev tp combine (Just n') x0
  
              when (A.lt singleType tid' sz) $ do
                writeArray TypeInt arrOut j0 r0
  
                ll <- A.sub numType bd (liftInt32 1)
                when (A.eq singleType tid ll) $
                  writeArray TypeInt32 carry (liftInt32 0) r0
  
              n1 <- A.sub numType sz bd'
              i1 <- next i0
              j1 <- next j0
              return $ A.trip n1 i1 j1
  
        -- Iterate over the remaining elements in this segment
        -- The loop condition uses the first triple of the state, 'n'.
        -- This variable is uniform (the same for all threads in the thread
        -- block), since it is initialized as 'sz' or 'sz - bd', and lowered by
        -- 'bd' each iteration. Hence the control flow in this loop is uniform,
        -- and we can thus call __syncthreads within the loop.
        void $ while
          (TupRunit `TupRpair` TupRsingle scalarTypeInt `TupRpair` TupRsingle scalarTypeInt `TupRpair` TupRsingle scalarTypeInt)
          (\(A.fst3   -> n)       -> A.gt singleType n (liftInt 0))
          (\(A.untrip -> (n,i,j)) -> do
  
            -- Wait for the carry-in value from the previous iteration to be updated
            __syncthreads
  
            -- Compute and store the next element of the scan
            --
            -- NOTE: As with 'foldSeg' we require all threads to participate in
            -- every iteration of the loop otherwise they will die prematurely.
            -- Out-of-bounds threads return 'undef' at this point, which is really
            -- unfortunate ):
            --
            x <- if (tp, A.lt singleType tid' n)
                   then app1 (delayedLinearIndex arrIn) i
                   else return $ tupUndef tp
  
            -- Thread zero incorporates the carry-in element
            y <- if (tp, A.eq singleType tid (liftInt32 0))
                   then do
                     c <- readArray TypeInt32 carry (liftInt32 0)
                     case dir of
                       LeftToRight -> app2 combine c x
                       RightToLeft -> app2 combine x c
                    else
                      return x
  
            -- Perform the scan and write the result to memory
            m <- i32 n
            z <- if (tp, A.gte singleType n bd')
                   then scanBlock dir dev tp combine Nothing  y
                   else scanBlock dir dev tp combine (Just m) y
  
            when (A.lt singleType tid' n) $ do
              writeArray TypeInt arrOut j z
  
              -- The last thread of the block writes its result as the carry-out
              -- value. If this thread is not active then we are on the last
              -- iteration of the loop and it will not be needed.
              w <- A.sub numType bd (liftInt32 1)
              when (A.eq singleType tid w) $
                writeArray TypeInt32 carry (liftInt32 0) z
  
            -- Update indices for the next iteration
            n' <- A.sub numType n bd'
            i' <- next i
            j' <- next j
            return $ A.trip n' i' j')
          r
  
      return_
  
  
  -- Multidimensional scan' along the innermost dimension
  --
  -- A thread block individually computes along each innermost dimension. This is
  -- a single-pass operation.
  --
  --  * We can assume that the array is non-empty; exclusive scans with empty
  --    innermost dimension will be instead filled with the seed element via
  --    'mkScan'Fill'.
  --
  --  * Small but non-empty innermost dimension arrays (size << thread
  --    block size) will have many threads which do no work.
  --
  mkScan'Dim
      :: forall aenv sh e.
         Direction
      -> UID
      -> Gamma          aenv                          -- ^ array environment
      -> ArrayR (Array (sh, Int) e)
      -> IRFun2     PTX aenv (e -> e -> e)            -- ^ combination function
      -> IRExp      PTX aenv e                        -- ^ seed element
      -> MIRDelayed PTX aenv (Array (sh, Int) e)      -- ^ input data
      -> CodeGen    PTX      (IROpenAcc PTX aenv (Array (sh, Int) e, Array sh e))
  mkScan'Dim dir uid aenv repr@(ArrayR (ShapeRsnoc shr) tp) combine seed marr = do
    dev <- liftCodeGen $ gets ptxDeviceProperties
    --
    let
        (arrSum, paramSum)  = mutableArray (reduceRank repr) "sum"
        (arrOut, paramOut)  = mutableArray repr "out"
        (arrIn,  paramIn)   = delayedArray "in" marr
        paramEnv            = envParam aenv
        --
        config              = launchConfig dev (CUDA.incWarp dev) (scanSMemSize dev tp) const [|| const ||]
    --
    makeOpenAccWith config uid "scan" (paramOut ++ paramSum ++ paramIn ++ paramEnv) $ do
  
      -- The first and last threads of the block need to communicate the
      -- block-wide aggregate as a carry-in value across iterations.
      --
      -- TODO: we could optimise this a bit if we can get access to the shared
      -- memory area used by 'scanBlockSMem', and from there directly read the
      -- value computed by the last thread.
      carry <- staticSharedMem tp 1
  
      -- Size of the input array
      sz    <- indexHead <$> delayedExtent arrIn
  
      -- If the innermost dimension is smaller than the number of threads in the
      -- block, those threads will never contribute to the output.
      tid   <- threadIdx
      tid'  <- int tid
      when (A.lte singleType tid' sz) $ do
  
        -- Thread blocks iterate over the outer dimensions, each thread block
        -- cooperatively scanning along each outermost index.
        bid <- int =<< blockIdx
        gd  <- int =<< gridDim
        end <- shapeSize shr (irArrayShape arrSum)
  
        imapFromStepTo bid gd end $ \seg -> do
  
          -- Not necessary to wait for threads to catch up before starting this segment
          -- __syncthreads
  
          -- Linear index bounds for this segment
          inf <- A.mul numType seg sz
          sup <- A.add numType inf sz
  
          -- Index that this thread will read from. Recall that the supremum index
          -- is exclusive.
          i0  <- case dir of
                   LeftToRight -> A.add numType inf tid'
                   RightToLeft -> do x <- A.sub numType sup tid'
                                     y <- A.sub numType x (liftInt 1)
                                     return y
  
          -- The index that this thread will write to. This is just shifted along
          -- by one to make room for the initial element.
          j0  <- case dir of
                   LeftToRight -> A.add numType i0 (liftInt 1)
                   RightToLeft -> A.sub numType i0 (liftInt 1)
  
          -- Evaluate the initial element. Store it into the carry-in slot as well
          -- as to the array as the first element. This is always valid because if
          -- the input array is empty then we will be evaluating via mkScan'Fill.
          when (A.eq singleType tid (liftInt32 0)) $ do
            z <- seed
            writeArray TypeInt   arrOut i0            z
            writeArray TypeInt32 carry  (liftInt32 0) z
  
          bd  <- blockDim
          bd' <- int bd
          let next ix = case dir of
                          LeftToRight -> A.add numType ix bd'
                          RightToLeft -> A.sub numType ix bd'
  
          -- Now, threads iterate over the elements along the innermost dimension.
          -- At each iteration the first thread incorporates the carry-in value
          -- from the previous step.
          --
          -- The index tracks how many elements remain for the thread block, since
          -- indices i* and j* are local to each thread
          n0  <- A.sub numType sup inf
          void $ while
            (TupRunit `TupRpair` TupRsingle scalarTypeInt `TupRpair` TupRsingle scalarTypeInt `TupRpair` TupRsingle scalarTypeInt)
            (\(A.fst3   -> n)       -> A.gt singleType n (liftInt 0))
            (\(A.untrip -> (n,i,j)) -> do
  
              -- Wait for threads to catch up to ensure the carry-in value from
              -- the last iteration has been updated
              __syncthreads
  
              -- If all threads in the block will participate this round we can
              -- avoid (almost) all bounds checks.
              _ <- if (TupRunit, A.gte singleType n bd')
                      -- All threads participate. No bounds checks required but
                      -- the last thread needs to update the carry-in value.
                      then do
                        x <- app1 (delayedLinearIndex arrIn) i
                        y <- if (tp, A.eq singleType tid (liftInt32 0))
                                then do
                                  c <- readArray TypeInt32 carry (liftInt32 0)
                                  case dir of
                                    LeftToRight -> app2 combine c x
                                    RightToLeft -> app2 combine x c
                                else
                                  return x
                        z <- scanBlock dir dev tp combine Nothing y
  
                        -- Write results to the output array. Note that if we
                        -- align directly on the boundary of the array this is not
                        -- valid for the last thread.
                        case dir of
                          LeftToRight -> when (A.lt  singleType j sup) $ writeArray TypeInt arrOut j z
                          RightToLeft -> when (A.gte singleType j inf) $ writeArray TypeInt arrOut j z
  
                        -- Last thread of the block also saves its result as the
                        -- carry-in value
                        bd1 <- A.sub numType bd (liftInt32 1)
                        when (A.eq singleType tid bd1) $
                          writeArray TypeInt32 carry (liftInt32 0) z
  
                        return (lift TupRunit ())
  
                      -- Only threads that are in bounds can participate. This is
                      -- the last iteration of the loop. The last active thread
                      -- still needs to store its value into the carry-in slot.
                      --
                      -- Note that all threads must call the block-wide scan.
                      -- SEE: [Synchronisation problems with SM_70 and greater]
                      else do
                        x <- if (tp, A.lt singleType tid' n)
                                then do
                                  x <- app1 (delayedLinearIndex arrIn) i
                                  y <- if (tp, A.eq singleType tid (liftInt32 0))
                                          then do
                                            c <- readArray TypeInt32 carry (liftInt32 0)
                                            case dir of
                                              LeftToRight -> app2 combine c x
                                              RightToLeft -> app2 combine x c
                                          else
                                            return x
                                  return y
                                else
                                  return $ tupUndef tp
  
                        l <- i32 n
                        y <- scanBlock dir dev tp combine (Just l) x
  
                        m <- A.sub numType l (liftInt32 1)
                        when (A.lt singleType tid m) $ writeArray TypeInt   arrOut j            y
                        when (A.eq singleType tid m) $ writeArray TypeInt32 carry (liftInt32 0) y
  
                        return (lift TupRunit ())
  
              A.trip <$> A.sub numType n bd' <*> next i <*> next j)
            (A.trip n0 i0 j0)
  
          -- Wait for the carry-in value to be updated
          __syncthreads
  
          -- Store the carry-in value to the separate final results array
          when (A.eq singleType tid (liftInt32 0)) $
            writeArray TypeInt arrSum seg =<< readArray TypeInt32 carry (liftInt32 0)
  
      return_
  
  
  
  -- Parallel scan, auxiliary
  --
  -- If this is an exclusive scan of an empty array, we just fill the result with
  -- the seed element.
  --
  mkScanFill
      :: UID
      -> Gamma aenv
      -> ArrayR (Array sh e)
      -> IRExp PTX aenv e
      -> CodeGen PTX (IROpenAcc PTX aenv (Array sh e))
  mkScanFill uid aenv repr seed =
    mkGenerate uid aenv repr (IRFun1 (const seed))
  
  mkScan'Fill
      :: UID
      -> Gamma aenv
      -> ArrayR (Array (sh, Int) e)
      -> IRExp PTX aenv e
      -> CodeGen PTX (IROpenAcc PTX aenv (Array (sh, Int) e, Array sh e))
  mkScan'Fill uid aenv repr seed =
    Safe.coerce <$> mkGenerate uid aenv (reduceRank repr) (IRFun1 (const seed))
  
  scanSMemSize :: DeviceProperties -> TypeR e -> Int -> Int
  scanSMemSize dev tp n = sharedMemorySizeAdd tp warps 0
    where
      ws        = CUDA.warpSize dev
      warps     = n `P.quot` ws
  
  -- Block wide scan
  -- ---------------
  
  -- Efficient block-wide (inclusive) scan using the specified operator.
  --
  -- Each block requires (#warps * (1 + 1.5*warp size)) elements of dynamically
  -- allocated shared memory.
  --
  -- Example: https://github.com/NVlabs/cub/blob/1.5.4/cub/block/specializations/block_scan_warp_scans.cuh
  --
  -- Must be called under uniform control flow within a thread block
  -- (as this function may use __syncthreads)
  scanBlock
      :: forall aenv e.
         Direction
      -> DeviceProperties                             -- ^ properties of the target device
      -> TypeR e
      -> IRFun2 PTX aenv (e -> e -> e)                -- ^ combination function
      -> Maybe (Operands Int32)                       -- ^ number of valid elements (may be less than block size)
      -> Operands e                                   -- ^ calling thread's input element
      -> CodeGen PTX (Operands e)
  scanBlock dir dev tp combine nelem = warpScan >=> warpPrefix
    where
      int32 :: Integral a => a -> Operands Int32
      int32 = liftInt32 . P.fromIntegral
  
      -- Temporary storage required for each warp
      -- warp_smem_elems = CUDA.warpSize dev + (CUDA.warpSize dev `P.quot` 2)
      -- warp_smem_bytes = warp_smem_elems  * bytesElt tp
  
      -- Step 1: Scan in every warp
      warpScan :: Operands e -> CodeGen PTX (Operands e)
      warpScan = scanWarp dir dev tp combine
  
      -- Step 2: Collect the aggregate results of each warp to compute the prefix
      -- values for each warp and combine with the partial result to compute each
      -- thread's final value.
      warpPrefix :: Operands e -> CodeGen PTX (Operands e)
      warpPrefix input = do
        -- Allocate #warps elements of shared memory
        bd    <- blockDim
        warps <- A.quot integralType bd (int32 (CUDA.warpSize dev))
        smem  <- dynamicSharedMem tp TypeInt32 warps (liftInt32 0)
  
        -- Share warp aggregates
        wid   <- warpId
        lane  <- laneId
        when (A.eq singleType lane (int32 (CUDA.warpSize dev - 1))) $ do
          writeArray TypeInt32 smem wid input
  
        -- Wait for each warp to finish its local scan and share the aggregate
        __syncthreads
  
        -- Compute the prefix value for this warp and add to the partial result.
        -- This step is not required for the first warp, which has no carry-in.
        if (tp, A.eq singleType wid (liftInt32 0))
          then return input
          else do
            -- Every thread sequentially scans the warp aggregates to compute
            -- their prefix value. We do this sequentially, but could also have
            -- warp 0 do it cooperatively if we limit thread block sizes to
            -- (warp size ^ 2).
            steps <- case nelem of
                        Nothing -> return wid
                        Just n  -> A.min singleType wid =<< A.quot integralType n (int32 (CUDA.warpSize dev))
  
            p0     <- readArray TypeInt32 smem (liftInt32 0)
            prefix <- iterFromStepTo tp (liftInt32 1) (liftInt32 1) steps p0 $ \step x -> do
                        y <- readArray TypeInt32 smem step
                        case dir of
                          LeftToRight -> app2 combine x y
                          RightToLeft -> app2 combine y x
  
            case dir of
              LeftToRight -> app2 combine prefix input
              RightToLeft -> app2 combine input prefix
  
  
  -- Warp-wide scan
  -- --------------
  
  -- Efficient warp-wide (inclusive) scan using the specified operator.
  --
  -- Each warp requires 48 (1.5 x warp size) elements of shared memory. The
  -- routine assumes that it is allocated individually per-warp (i.e. can be
  -- indexed in the range [0, warp size)).
  --
  -- Example: https://github.com/NVlabs/cub/blob/1.5.4/cub/warp/specializations/warp_scan_smem.cuh
  --
  scanWarp
      :: forall aenv e.
         Direction
      -> DeviceProperties                             -- ^ properties of the target device
      -> TypeR e
      -> IRFun2 PTX aenv (e -> e -> e)                -- ^ combination function
      -> Operands e                                   -- ^ calling thread's input element
      -> CodeGen PTX (Operands e)
  scanWarp dir dev tp combine = scan 0
    where
      log2 :: Double -> Double
      log2 = P.logBase 2
  
      -- Number of steps required to scan warp
      steps = P.floor (log2 (P.fromIntegral (CUDA.warpSize dev)))
  
      -- Unfold the scan as a recursive code generation function
      scan :: Int -> Operands e -> CodeGen PTX (Operands e)
      scan step x
        | step >= steps = return x
        | otherwise     = do
            let offset = 1 `P.shiftL` step
  
            -- share partial result through shared memory buffer
            y    <- __shfl_up tp x (liftWord32 offset)
            lane <- laneId
  
            -- update partial result if in range
            x'   <- if (tp, A.gte singleType lane (liftInt32 . P.fromIntegral $ offset))
                      then do
                        case dir of
                          LeftToRight -> app2 combine y x
                          RightToLeft -> app2 combine x y
  
                      else
                        return x
  
            scan (step+1) x'
  
  tupUndef :: TypeR a -> Operands a
  tupUndef TupRunit       = OP_Unit
  tupUndef (TupRpair a b) = OP_Pair (tupUndef a) (tupUndef b)
  tupUndef (TupRsingle t) = ir t (undef t)
  
  -- Utilities
  -- ---------
  
  i32 :: Operands Int -> CodeGen PTX (Operands Int32)
  i32 = A.fromIntegral integralType numType
  
  int :: Operands Int32 -> CodeGen PTX (Operands Int)
  int = A.fromIntegral integralType numType
  
  