accelerate-cuda-0.17.0.0: Data/Array/Accelerate/CUDA/CodeGen/IndexSpace.hs
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ImpredicativeTypes #-}
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE ViewPatterns #-}
-- |
-- Module : Data.Array.Accelerate.CUDA.CodeGen.IndexSpace
-- Copyright : [2008..2014] Manuel M T Chakravarty, Gabriele Keller
-- [2009..2014] Trevor L. McDonell
-- License : BSD3
--
-- Maintainer : Trevor L. McDonell <tmcdonell@cse.unsw.edu.au>
-- Stability : experimental
-- Portability : non-portable (GHC extensions)
--
module Data.Array.Accelerate.CUDA.CodeGen.IndexSpace (
-- Array construction
mkGenerate,
-- Permutations
mkTransform, mkPermute,
) where
import Language.C.Quote.CUDA
import Foreign.CUDA.Analysis.Device
import qualified Language.C.Syntax as C
import Data.Array.Accelerate.Array.Sugar ( Array, Shape, Elt, ignore, shapeToList )
import Data.Array.Accelerate.Error ( internalError )
import Data.Array.Accelerate.CUDA.AST ( Gamma )
import Data.Array.Accelerate.CUDA.CodeGen.Base
-- Construct a new array by applying a function to each index. Each thread
-- processes multiple elements, striding the array by the grid size.
--
-- generate :: (Shape ix, Elt e)
-- => Exp ix -- dimension of the result
-- -> (Exp ix -> Exp a) -- function to apply at each index
-- -> Acc (Array ix a)
--
mkGenerate
:: forall aenv sh e. (Shape sh, Elt e)
=> DeviceProperties
-> Gamma aenv
-> CUFun1 aenv (sh -> e)
-> [CUTranslSkel aenv (Array sh e)]
mkGenerate dev aenv (CUFun1 dce f)
= return
$ CUTranslSkel "generate" [cunit|
$esc:("#include <accelerate_cuda.h>")
$edecls:texIn
extern "C" __global__ void
generate
(
$params:argIn,
$params:argOut
)
{
const int shapeSize = $exp:(csize shOut);
const int gridSize = $exp:(gridSize dev);
int ix;
for ( ix = $exp:(threadIdx dev)
; ix < shapeSize
; ix += gridSize )
{
$items:(dce sh .=. cfromIndex shOut "ix" "tmp")
$items:(setOut "ix" .=. f sh)
}
}
|]
where
(sh, _, _) = locals "sh" (undefined :: sh)
(texIn, argIn) = environment dev aenv
(argOut, shOut, setOut) = writeArray "Out" (undefined :: Array sh e)
-- A combination map/backpermute, where the index and value transformations have
-- been separated.
--
-- transform :: (Elt a, Elt b, Shape sh, Shape sh')
-- => PreExp acc aenv sh' -- dimension of the result
-- -> PreFun acc aenv (sh' -> sh) -- index permutation function
-- -> PreFun acc aenv (a -> b) -- function to apply at each element
-- -> acc aenv (Array sh a) -- source array
-- -> PreOpenAcc acc aenv (Array sh' b)
--
mkTransform
:: forall aenv sh sh' a b. (Shape sh, Shape sh', Elt a, Elt b)
=> DeviceProperties
-> Gamma aenv
-> CUFun1 aenv (sh' -> sh)
-> CUFun1 aenv (a -> b)
-> CUDelayedAcc aenv sh a
-> [CUTranslSkel aenv (Array sh' b)]
mkTransform dev aenv perm fun arr
| CUFun1 dce_p p <- perm
, CUFun1 dce_f f <- fun
, CUDelayed _ (CUFun1 dce_g get) _ <- arr
= return
$ CUTranslSkel "transform" [cunit|
$esc:("#include <accelerate_cuda.h>")
$edecls:texIn
extern "C" __global__ void
transform
(
$params:argIn,
$params:argOut
)
{
const int shapeSize = $exp:(csize shOut);
const int gridSize = $exp:(gridSize dev);
int ix;
for ( ix = $exp:(threadIdx dev)
; ix < shapeSize
; ix += gridSize )
{
$items:(dce_p sh' .=. cfromIndex shOut "ix" "tmp")
$items:(dce_g sh .=. p sh')
$items:(dce_f x0 .=. get sh)
$items:(setOut "ix" .=. f x0)
}
}
|]
where
(texIn, argIn) = environment dev aenv
(argOut, shOut, setOut) = writeArray "Out" (undefined :: Array sh' b)
(x0, _, _) = locals "x" (undefined :: a)
(sh, _, _) = locals "sh" (undefined :: sh)
(sh', _, _) = locals "sh_" (undefined :: sh')
-- Forward permutation specified by an index mapping that determines for each
-- element in the source array where it should go in the target. The resultant
-- array is initialised with the given defaults and any further values that are
-- permuted into the result array are added to the current value using the given
-- combination function.
--
-- The combination function must be associative. Extents that are mapped to the
-- magic value 'ignore' by the permutation function are dropped.
--
-- permute :: (Shape ix, Shape ix', Elt a)
-- => (Exp a -> Exp a -> Exp a) -- combination function
-- -> Acc (Array ix' a) -- array of default values
-- -> (Exp ix -> Exp ix') -- permutation
-- -> Acc (Array ix a) -- permuted array
-- -> Acc (Array ix' a)
--
mkPermute
:: forall aenv sh sh' e. (Shape sh, Shape sh', Elt e)
=> DeviceProperties
-> Gamma aenv
-> CUFun2 aenv (e -> e -> e)
-> CUFun1 aenv (sh -> sh')
-> CUDelayedAcc aenv sh e
-> [CUTranslSkel aenv (Array sh' e)]
mkPermute dev aenv (CUFun2 dce_x dce_y combine) (CUFun1 dce_p prj) arr
| CUDelayed (CUExp shIn) _ (CUFun1 _ get) <- arr
= return
$ CUTranslSkel "permute" [cunit|
$esc:("#include <accelerate_cuda.h>")
$edecls:texIn
extern "C" __global__ void
permute
(
$params:argIn,
$params:argOut,
typename Int32 * __restrict__ lock
)
{
/*
* The input shape might be a complex expression. Evaluate it first to reuse the result.
*/
$items:(sh .=. shIn)
const int shapeSize = $exp:(csize sh);
const int gridSize = $exp:(gridSize dev);
int ix;
for ( ix = $exp:(threadIdx dev)
; ix < shapeSize
; ix += gridSize )
{
$items:(dce_p src .=. cfromIndex sh "ix" "srcTmp")
$items:(dst .=. prj src)
if ( ! $exp:(cignore dst) )
{
$items:(jx .=. ctoIndex shOut dst)
$items:(dce_x x .=. get ix)
$items:(atomically jx
[ dce_y y .=. setOut jx
, setOut jx .=. combine x y
]
)
}
}
}
|]
where
(texIn, argIn) = environment dev aenv
(argOut, shOut, setOut) = writeArray "Out" (undefined :: Array sh' e)
(x, _, _) = locals "x" (undefined :: e)
(y, _, _) = locals "y" (undefined :: e)
(sh, _, _) = locals "shIn" (undefined :: sh)
(src, _, _) = locals "sh" (undefined :: sh)
(dst, _, _) = locals "sh_" (undefined :: sh')
([jx], _, _) = locals "jx" (undefined :: Int)
ix = [cvar "ix"]
sm = computeCapability dev
-- If the destination index resolves to the magic index "ignore", the result
-- is dropped from the output array.
--
cignore :: Rvalue x => [x] -> C.Exp
cignore [] = $internalError "permute" "singleton arrays not supported"
cignore xs = foldl1 (\a b -> [cexp| $exp:a && $exp:b |])
$ zipWith (\a b -> [cexp| $exp:(rvalue a) == $int:b |]) xs
$ shapeToList (ignore :: sh')
-- If we can determine that the old values are not used in the combination
-- function (e.g. filter) then the lock and unlock fragments can be replaced
-- with a NOP.
--
-- If locking is required but the hardware does not support it (compute 1.0)
-- then we issue a runtime error immediately instead of silently failing.
--
mustLock = or . fst . unzip $ dce_y y
-- The atomic section is acquired using a spin lock. This requires a
-- temporary array to represent the lock state for each element of the
-- output. We use 1 to represent the locked state, and 0 to represent
-- unlocked elements.
--
-- do {
-- old = atomicExch(&lock[i], 1); // atomic exchange
-- } while (old == 1);
--
-- /* critical section */
--
-- atomicExch(&lock[i], 0);
--
-- The initial loop repeatedly attempts to take the lock by writing a 1 into
-- the slot. Once the 'old' state of the lock returns 0 (unlocked), we have
-- just acquired the lock, and the atomic section can be computed. Finally,
-- atomically write a 0 back into the slot to unlock the element.
--
-- However, there is a complication with CUDA devices because all threads in
-- the warp must execute in lockstep (with predicated execution). Once a
-- thread acquires a lock, then it will be disabled and stop participating
-- in the first loop, waiting until all other threads in the warp acquire
-- their locks. If two threads in a warp are attempting to acquire the same
-- lock, once the lock is acquired by the first thread, it sits idle while
-- the second thread spins attempting to grab a lock that will never be
-- released, because the first thread can not progress. DEADLOCK.
--
-- So, we need to invert the algorithm so that threads can always make
-- progress, until each thread in the warp has committed their result.
--
-- done = 0;
-- do {
-- if (atomicExch(&lock[i], 1) == 0) {
--
-- /* critical section */
--
-- done = 1;
-- atomicExch(&lock[i], 0);
-- }
-- } while (done == 0)
--
atomically :: (C.Type, Name) -> [[C.BlockItem]] -> [C.BlockItem]
atomically (_,i) (concat -> body)
| not mustLock = body
| sm < Compute 1 1 = $internalError "permute" "Requires at least compute compatibility 1.1"
| otherwise =
[ [citem| typename Int32 done = 0; |]
, [citem| do {
typename Int32 *addr = &lock[ $exp:(cvar i) ];
if ( atomicExch( addr, 1 ) == 0 ) {
$items:body
done = 1;
atomicExch( addr, 0 );
}
__threadfence();
} while (done == 0);
|]
]