accelerate-cuda-0.13.0.0: Data/Array/Accelerate/CUDA/CodeGen/Reduction.hs
{-# LANGUAGE CPP #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators #-}
-- |
-- Module : Data.Array.Accelerate.CUDA.CodeGen.Reduction
-- Copyright : [2008..2010] Manuel M T Chakravarty, Gabriele Keller, Sean Lee
-- [2009..2012] Manuel M T Chakravarty, Gabriele Keller, 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.Reduction (
mkFold, mkFold1, mkFoldSeg, mkFold1Seg,
) where
import Foreign.CUDA.Analysis
import Language.C.Quote.CUDA
import qualified Language.C.Syntax as C
import Data.Array.Accelerate.Type ( IsIntegral )
import Data.Array.Accelerate.Array.Sugar ( Array, Shape, Elt, Z(..), (:.)(..) )
import Data.Array.Accelerate.Analysis.Shape
import Data.Array.Accelerate.CUDA.AST
import Data.Array.Accelerate.CUDA.CodeGen.Base
import Data.Array.Accelerate.CUDA.CodeGen.Type
#include "accelerate.h"
-- Reduce an array along the innermost dimension. The function must be
-- associative to enable efficient parallel implementation.
--
-- fold :: (Shape ix, Elt a)
-- => (Exp a -> Exp a -> Exp a)
-- -> Exp a
-- -> Acc (Array (ix :. Int) a)
-- -> Acc (Array ix a)
--
-- fold1 :: (Shape ix, Elt a)
-- => (Exp a -> Exp a -> Exp a)
-- -> Acc (Array (ix :. Int) a)
-- -> Acc (Array ix a)
--
-- If this is collapsing an array to a single value, we use a multi-pass
-- algorithm that splits the input data over several thread blocks. The first
-- kernel is executed once, and then the second recursively until a single value
-- is produced.
--
mkFold :: 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 e)]
mkFold dev aenv f z a
| expDim (undefined :: Exp aenv sh) > 0 = mkFoldDim dev aenv f (Just z) a
| otherwise = mkFoldAll dev aenv f (Just z) a
mkFold1 :: 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 e) ]
mkFold1 dev aenv f a
| expDim (undefined :: Exp aenv sh) > 0 = mkFoldDim dev aenv f Nothing a
| otherwise = mkFoldAll dev aenv f Nothing a
-- Reduction of an array of arbitrary rank to a single scalar value. Each thread
-- computes multiple elements sequentially. This reduces the overall cost of the
-- algorithm while keeping the work complexity O(n) and the step complexity
-- O(log n). c.f. Brent's Theorem optimisation.
--
-- Since the reduction occurs over multiple blocks, there are two phases. The
-- first pass incorporates any fused/embedded input arrays, while the second
-- recurses over a manifest array to produce a single value.
--
mkFoldAll
:: forall aenv sh e. (Shape sh, Elt e)
=> DeviceProperties
-> Gamma aenv
-> CUFun2 aenv (e -> e -> e)
-> Maybe (CUExp aenv e)
-> CUDelayedAcc aenv (sh :. Int) e
-> [ CUTranslSkel aenv (Array sh e) ]
mkFoldAll dev aenv f z a
= let (_, rec) = getters "Rec" (undefined :: Array (sh:.Int) e)
in
[ mkFoldAll' False dev aenv f z a
, mkFoldAll' True dev aenv f z rec ]
mkFoldAll'
:: forall aenv sh e. (Shape sh, Elt e)
=> Bool
-> DeviceProperties
-> Gamma aenv
-> CUFun2 aenv (e -> e -> e)
-> Maybe (CUExp aenv e)
-> CUDelayedAcc aenv (sh :. Int) e
-> CUTranslSkel aenv (Array sh e)
mkFoldAll' recursive dev aenv fun@(CUFun2 _ _ combine) mseed (CUDelayed (CUExp sh) _ (CUFun1 _ get))
= CUTranslSkel foldAll [cunit|
$esc:("#include <accelerate_cuda_extras.h>")
$edecls:texIn
extern "C" __global__ void
$id:foldAll
(
$params:argIn,
$params:argOut,
$params:argRec
)
{
$decls:smem
$decls:declx
$decls:decly
$items:(shIn .=. sh)
const int shapeSize = $exp:(shapeSize shIn);
const int gridSize = $exp:(gridSize dev);
int ix = $exp:(threadIdx dev);
/*
* Reduce multiple elements per thread. The number is determined by the
* number of active thread blocks (via gridDim). More blocks will result in
* a larger `gridSize', and hence fewer elements per thread
*
* The loop stride of `gridSize' is used to maintain coalescing.
*
* Note that we can't simply kill threads that won't participate in the
* reduction, as exclusive reductions of empty arrays then won't be
* initialised with their seed element.
*/
if ( ix < shapeSize )
{
/*
* Initialise the local sum, then ...
*/
$items:(y .=. get ix)
/*
* ... continue striding the array, reading new values into 'x' and
* combining into the local accumulator 'y'. The non-idiomatic
* structure of the loop below is because we have already
* initialised 'y' above.
*/
for ( ix += gridSize; ix < shapeSize; ix += gridSize )
{
$items:(x .=. get ix)
$items:(y .=. combine x y)
}
}
/*
* Each thread puts its local sum into shared memory, then threads
* cooperatively reduce the shared array to a single value.
*/
$items:(sdata "threadIdx.x" .=. y)
__syncthreads();
ix = min(shapeSize - blockIdx.x * blockDim.x, blockDim.x);
$stms:(reduceBlock dev fun x y sdata (cvar "ix"))
/*
* Write the results of this block back to global memory. If we are the last
* phase of a recursive multi-block reduction, include the seed element.
*/
if ( threadIdx.x == 0 )
{
$items:(maybe inclusive_finish exclusive_finish mseed)
}
}
|]
where
foldAll = maybe "fold1All" (const "foldAll") mseed
(texIn, argIn) = environment dev aenv
(argOut, setOut) = setters "Out" (undefined :: Array sh e)
(argRec, _)
| recursive = getters "Rec" (undefined :: Array (sh:.Int) e)
| otherwise = ([], undefined)
(_, x, declx) = locals "x" (undefined :: e)
(_, y, decly) = locals "y" (undefined :: e)
(shIn, _, _) = locals "sh" (undefined :: sh :. Int)
ix = [cvar "ix"]
(smem, sdata) = shared (undefined :: e) "sdata" [cexp| blockDim.x |] Nothing
--
inclusive_finish = setOut "blockIdx.x" .=. y
exclusive_finish (CUExp seed) = C.BlockStm [cstm|
if ( shapeSize > 0 ) {
if ( gridDim.x == 1 ) {
$items:(x .=. seed)
$items:(y .=. combine x y)
}
$items:(setOut "blockIdx.x" .=. y)
}
else {
$items:(setOut "blockIdx.x" .=. seed)
}
|] : []
-- Reduction of the innermost dimension of an array of arbitrary rank. Each
-- thread block reduces along one innermost dimension index.
--
mkFoldDim
:: forall aenv sh e. (Shape sh, Elt e)
=> DeviceProperties
-> Gamma aenv
-> CUFun2 aenv (e -> e -> e)
-> Maybe (CUExp aenv e)
-> CUDelayedAcc aenv (sh :. Int) e
-> [ CUTranslSkel aenv (Array sh e) ]
mkFoldDim dev aenv fun@(CUFun2 _ _ combine) mseed (CUDelayed (CUExp sh) _ (CUFun1 _ get))
= return
$ CUTranslSkel fold [cunit|
$esc:("#include <accelerate_cuda_extras.h>")
$edecls:texIn
extern "C" __global__ void
$id:fold
(
$params:argIn,
$params:argOut
)
{
$decls:smem
$decls:declx
$decls:decly
$items:(shIn .=. sh)
const int numIntervals = size(shOut);
const int intervalSize = $exp:(indexHead shIn);
int ix;
int seg;
/*
* If the intervals of an exclusive fold are empty, use all threads to
* map the seed value to the output array and exit.
*/
$stms:(maybe [] mapseed mseed)
/*
* Threads in a block cooperatively reduce all elements in an interval.
*/
for ( seg = blockIdx.x
; seg < numIntervals
; seg += gridDim.x )
{
const int start = seg * intervalSize;
const int end = start + intervalSize;
const int n = min(end - start, blockDim.x);
/*
* Kill threads that will not participate to avoid invalid reads.
* Take advantage of the fact that the array is rectangular.
*/
if ( threadIdx.x >= n )
return;
/*
* Ensure aligned access to global memory, and that each thread
* initialises its local sum
*/
ix = start - (start & (warpSize - 1));
if ( ix == start || intervalSize > blockDim.x)
{
ix += threadIdx.x;
if ( ix >= start )
{
$items:(y .=. get ix)
}
if ( ix + blockDim.x < end )
{
$items:(x .=. get [cvar "ix + blockDim.x"])
if ( ix >= start ) {
$items:(y .=. combine x y)
}
else {
$items:(y .=. x)
}
}
/*
* Now, iterate collecting a local sum
*/
for ( ix += 2 * blockDim.x; ix < end; ix += blockDim.x )
{
$items:(x .=. get ix)
$items:(y .=. combine x y)
}
}
else
{
$items:(y .=. get [cvar "start + threadIdx.x"])
}
/*
* Each thread puts its local sum into shared memory, and
* cooperatively reduces this to a single value.
*/
$items:(sdata "threadIdx.x" .=. y)
__syncthreads();
$stms:(reduceBlock dev fun x y sdata (cvar "n"))
/*
* Finally, the first thread writes the result for this segment. For
* exclusive reductions, we also combine with the seed element here.
*/
if ( threadIdx.x == 0 ) {
$items:(maybe [] exclusive_finish mseed)
$items:(setOut "seg" .=. y)
}
}
}
|]
where
fold = maybe "fold1" (const "fold") mseed
(texIn, argIn) = environment dev aenv
(argOut, setOut) = setters "Out" (undefined :: Array sh e)
(_, x, declx) = locals "x" (undefined :: e)
(_, y, decly) = locals "y" (undefined :: e)
(shIn, _, _) = locals "sh" (undefined :: sh :. Int)
ix = [cvar "ix"]
(smem, sdata) = shared (undefined :: e) "sdata" [cexp| blockDim.x |] Nothing
--
mapseed (CUExp seed)
= [cstm| if ( intervalSize == 0 ) {
const int gridSize = $exp:(gridSize dev);
for ( ix = $exp:(threadIdx dev)
; ix < numIntervals
; ix += gridSize )
{
$items:(setOut "ix" .=. seed)
}
} |] :[]
--
exclusive_finish (CUExp seed)
= concat [ x .=. seed
, y .=. combine x y ]
-- Segmented reduction along the innermost dimension of an array. Performs one
-- individual reduction per segment of the source array. These reductions
-- proceed in parallel.
--
-- foldSeg :: (Shape ix, Elt a)
-- => (Exp a -> Exp a -> Exp a)
-- -> Exp a
-- -> Acc (Array (ix :. Int) a)
-- -> Acc Segments
-- -> Acc (Array (ix :. Int) a)
--
-- fold1Seg :: (Shape ix, Elt a)
-- => (Exp a -> Exp a -> Exp a)
-- -> Acc (Array (ix :. Int) a)
-- -> Acc Segments
-- -> Acc (Array (ix :. Int) a)
--
-- Each segment of the vector is assigned to a warp, which computes the
-- reduction of the i-th section, in parallel. Care is taken to ensure that data
-- array access is aligned to a warp boundary.
--
-- Since an entire 32-thread warp is assigned for each segment, many threads
-- will remain idle when the segments are very small. This code relies on
-- implicit synchronisation among threads in a warp.
--
-- The offset array contains the starting index for each segment in the input
-- array. The i-th warp reduces values in the input array at indices
-- [d_offset[i], d_offset[i+1]).
--
mkFoldSeg
:: (Shape sh, Elt e, Elt i, IsIntegral i)
=> DeviceProperties
-> Gamma aenv
-> CUFun2 aenv (e -> e -> e)
-> CUExp aenv e
-> CUDelayedAcc aenv (sh :. Int) e
-> CUDelayedAcc aenv (Z :. Int) i
-> [CUTranslSkel aenv (Array (sh :. Int) e)]
mkFoldSeg dev aenv f z a s = [ mkFoldSeg' dev aenv f (Just z) a s ]
mkFold1Seg
:: (Shape sh, Elt e, Elt i, IsIntegral i)
=> DeviceProperties
-> Gamma aenv
-> CUFun2 aenv (e -> e -> e)
-> CUDelayedAcc aenv (sh :. Int) e
-> CUDelayedAcc aenv (Z :. Int) i
-> [CUTranslSkel aenv (Array (sh :. Int) e)]
mkFold1Seg dev aenv f a s = [ mkFoldSeg' dev aenv f Nothing a s ]
mkFoldSeg'
:: forall aenv sh e i. (Shape sh, Elt e, Elt i, IsIntegral i)
=> DeviceProperties
-> Gamma aenv
-> CUFun2 aenv (e -> e -> e)
-> Maybe (CUExp aenv e)
-> CUDelayedAcc aenv (sh :. Int) e
-> CUDelayedAcc aenv (Z :. Int) i
-> CUTranslSkel aenv (Array (sh :. Int) e)
mkFoldSeg' dev aenv fun@(CUFun2 _ _ combine) mseed
(CUDelayed (CUExp shIn) _ (CUFun1 _ get))
(CUDelayed _ _ (CUFun1 _ offset))
= CUTranslSkel foldSeg [cunit|
$esc:("#include <accelerate_cuda_extras.h>")
$edecls:texIn
extern "C"
__global__ void
$id:foldSeg
(
$params:argIn,
$params:argOut
)
{
const int vectors_per_block = blockDim.x / warpSize;
const int num_vectors = $exp:(umul24 dev vectors_per_block gridDim);
const int thread_id = $exp:(threadIdx dev);
const int vector_id = thread_id / warpSize;
const int thread_lane = threadIdx.x & (warpSize - 1);
const int vector_lane = threadIdx.x / warpSize;
const int num_segments = indexHead(shOut);
const int total_segments = size(shOut);
int seg;
int ix;
extern volatile __shared__ int s_ptrs[][2];
$decls:smem
$decls:declx
$decls:decly
$items:(sh .=. shIn)
/*
* Threads in a warp cooperatively reduce a segment
*/
for ( seg = vector_id
; seg < total_segments
; seg += num_vectors )
{
const int s = seg % num_segments;
const int base = (seg / num_segments) * $exp:(indexHead sh);
/*
* Use two threads to fetch the indices of the start and end of this
* segment. This results in single coalesced global read.
*/
if ( thread_lane < 2 ) {
$items:([cvar "s_ptrs[vector_lane][thread_lane]"] .=. offset [cvar "s + thread_lane"])
}
const int start = base + s_ptrs[vector_lane][0];
const int end = base + s_ptrs[vector_lane][1];
const int num_elements = end - start;
/*
* Each thread reads in values of this segment, accumulating a local sum
*/
if ( num_elements > warpSize )
{
/*
* Ensure aligned access to global memory
*/
ix = start - (start & (warpSize - 1)) + thread_lane;
if ( ix >= start )
{
$items:(y .=. get ix)
}
/*
* Subsequent reads to global memory are aligned, but make sure all
* threads have initialised their local sum.
*/
if ( ix + warpSize < end )
{
$items:(x .=. get [cvar "ix + warpSize"])
if ( ix >= start ) {
$items:(y .=. combine x y)
}
else {
$items:(y .=. x)
}
}
/*
* Now, iterate along the inner-most dimension collecting a local sum
*/
for ( ix += 2 * warpSize; ix < end; ix += warpSize )
{
$items:(x .=. get ix)
$items:(y .=. combine x y)
}
}
else if ( start + thread_lane < end )
{
$items:(y .=. get [cvar "start + thread_lane"])
}
/*
* Store local sums into shared memory and reduce to a single value
*/
ix = min(num_elements, warpSize);
$items:(sdata "threadIdx.x" .=. y)
$stms:(reduceWarp dev fun x y sdata (cvar "ix") (cvar "thread_lane"))
/*
* Finally, the first thread writes the result for this segment
*/
if ( thread_lane == 0 )
{
$items:(maybe [] exclusive_finish mseed)
$items:(setOut "seg" .=. y)
}
}
}
|]
where
foldSeg = maybe "fold1Seg" (const "foldSeg") mseed
(texIn, argIn) = environment dev aenv
(argOut, setOut) = setters "Out" (undefined :: Array (sh :. Int) e)
(_, x, declx) = locals "x" (undefined :: e)
(_, y, decly) = locals "y" (undefined :: e)
(sh, _, _) = locals "sh" (undefined :: sh :. Int)
(smem, sdata) = shared (undefined :: e) "sdata" [cexp| blockDim.x |] (Just $ [cexp| &s_ptrs[vectors_per_block][2] |])
--
ix = [cvar "ix"]
vectors_per_block = cvar "vectors_per_block"
gridDim = cvar "gridDim.x"
--
exclusive_finish (CUExp seed)
= C.BlockStm [cstm| if ( num_elements > 0 ) {
$items:(x .=. seed)
$items:(y .=. combine x y)
} else {
$items:(y .=. seed)
} |] :[]
-- Reducers
-- --------
-- Reductions of values stored in shared memory.
--
-- Two local (mutable) variables are also required to do the reduction. The
-- final result is stored in the second of these.
--
reduceWarp
:: forall aenv e. Elt e
=> DeviceProperties
-> CUFun2 aenv (e -> e -> e)
-> [C.Exp] -> [C.Exp] -- temporary variables x0 and x1
-> (Name -> [C.Exp]) -- index elements from shared memory
-> C.Exp -- number of elements
-> C.Exp -- thread identifier: usually lane or thread ID
-> [C.Stm]
reduceWarp dev fun x0 x1 sdata n tid
| shflOK dev (undefined :: e) = return
$ reduceWarpShfl dev fun x0 x1 n tid
| otherwise = reduceWarpTree dev fun x0 x1 sdata n tid
reduceBlock
:: forall aenv e. Elt e
=> DeviceProperties
-> CUFun2 aenv (e -> e -> e)
-> [C.Exp] -> [C.Exp] -- temporary variables x0 and x1
-> (Name -> [C.Exp]) -- index elements from shared memory
-> C.Exp -- number of elements
-> [C.Stm]
reduceBlock dev fun x0 x1 sdata n
| shflOK dev (undefined :: e) = reduceBlockShfl dev fun x0 x1 sdata n
| otherwise = reduceBlockTree dev fun x0 x1 sdata n
-- Tree reduction
-- --------------
reduceWarpTree
:: Elt e
=> DeviceProperties
-> CUFun2 aenv (e -> e -> e)
-> [C.Exp] -> [C.Exp] -- temporary variables x0 and x1
-> (Name -> [C.Exp]) -- index elements from shared memory
-> C.Exp -- number of elements
-> C.Exp -- thread identifier: usually lane or thread ID
-> [C.Stm]
reduceWarpTree dev (CUFun2 _ _ f) x0 x1 sdata n tid
= map (reduce . pow2) [v, v-1 .. 0]
where
v = floor (logBase 2 (fromIntegral $ warpSize dev :: Double))
pow2 :: Int -> Int
pow2 x = 2 ^ x
reduce :: Int -> C.Stm
reduce 0
= [cstm| if ( $exp:tid < $exp:n ) {
$items:(x0 .=. sdata "threadIdx.x + 1")
$items:(x1 .=. f x1 x0)
} |]
reduce i
= [cstm| if ( $exp:tid + $int:i < $exp:n ) {
$items:(x0 .=. sdata ("threadIdx.x + " ++ show i))
$items:(x1 .=. f x1 x0)
$items:(sdata "threadIdx.x" .=. x1)
} |]
reduceBlockTree
:: Elt e
=> DeviceProperties
-> CUFun2 aenv (e -> e -> e)
-> [C.Exp] -> [C.Exp] -- temporary variables x0 and x1
-> (Name -> [C.Exp]) -- index elements from shared memory
-> C.Exp -- number of elements
-> [C.Stm]
reduceBlockTree dev fun@(CUFun2 _ _ f) x0 x1 sdata n
= flip (foldr1 (.)) []
$ map (reduce . pow2) [u-1, u-2 .. v]
where
u = floor (logBase 2 (fromIntegral $ maxThreadsPerBlock dev :: Double))
v = floor (logBase 2 (fromIntegral $ warpSize dev :: Double))
pow2 :: Int -> Int
pow2 x = 2 ^ x
reduce :: Int -> [C.Stm] -> [C.Stm]
reduce i rest
| i > warpSize dev
= [cstm| if ( threadIdx.x + $int:i < $exp:n ) {
$items:(x0 .=. sdata ("threadIdx.x + " ++ show i))
$items:(x1 .=. f x1 x0)
$items:(sdata "threadIdx.x" .=. x1)
} |]
: [cstm| __syncthreads(); |]
: rest
| otherwise
= [cstm| if ( threadIdx.x < $int:(warpSize dev) ) {
$stms:(reduceWarpTree dev fun x0 x1 sdata n (cvar "threadIdx.x"))
} |]
: rest
-- Butterfly reduction
-- -------------------
shflOK :: Elt e => DeviceProperties -> e -> Bool
shflOK _dev _ = False
-- shflOK dev dummy
-- = computeCapability dev >= Compute 3 0 && all (`elem` [4,8]) (eltSizeOf dummy)
-- Reduction using the __shfl_xor() operation for exchanging variables between
-- threads of a without use of shared memory. The exchange occurs simultaneously
-- for all active threads within the wrap, moving 4 bytes of data per thread.
-- 8-byte quantities are broken into two separate transfers.
--
reduceWarpShfl
:: forall aenv e. Elt e
=> DeviceProperties
-> CUFun2 aenv (e -> e -> e)
-> [C.Exp] -> [C.Exp]
-> C.Exp
-> C.Exp
-> C.Stm
reduceWarpShfl _dev (CUFun2 _ _ f) x0 x1 n tid
= [cstm| for ( int z = warpSize/2; z >= 1; z /= 2 ) {
$items:(x0 .=. shfl_xor x1)
if ( $exp:tid + z < $exp:n ) {
$items:(x1 .=. f x1 x0)
}
} |]
where
sizeof = eltSizeOf (undefined :: e)
shfl_xor = zipWith (\s x -> ccall (shfl s) [ x, cvar "z" ]) sizeof
where
shfl 4 = "shfl_xor32"
shfl 8 = "shfl_xor64"
shfl _ = INTERNAL_ERROR(error) "shfl_xor" "I only know about 32- and 64-bit types"
-- Reduce a block of values in butterfly fashion using __shfl_xor(). Each warp
-- calculates a local reduction, and the first thread of a warp writes its
-- result into shared memory. The first warp then reduces these values to the
-- final result.
--
reduceBlockShfl
:: forall aenv e. Elt e
=> DeviceProperties
-> CUFun2 aenv (e -> e -> e)
-> [C.Exp] -> [C.Exp]
-> (Name -> [C.Exp])
-> C.Exp
-> [C.Stm]
reduceBlockShfl dev fun x0 x1 sdata n
= reduceWarpShfl dev fun x0 x1 n (cvar "threadIdx.x")
: [cstm| if ( (threadIdx.x & warpSize - 1) == 0 ) {
$items:(sdata "threadIdx.x / warpSize" .=. x1)
} |]
: [cstm| __syncthreads(); |]
: [cstm| if ( threadIdx.x < warpSize ) {
$items:(x1 .=. sdata "threadIdx.x")
$exp:n = ($exp:n + warpSize - 1) / warpSize;
$stm:(reduceWarpShfl dev fun x0 x1 n (cvar "threadIdx.x"))
} |]
: []