accelerate-llvm-ptx-1.4.0.0: src/Data/Array/Accelerate/LLVM/PTX/CodeGen/FoldSeg.hs
{-# 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.Reader ( asks )
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 $ asks 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 $ asks 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