accelerate-cuda-0.17.0.0: Data/Array/Accelerate/CUDA/Execute.hs
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE IncoherentInstances #-}
{-# LANGUAGE NoForeignFunctionInterface #-}
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE UndecidableInstances #-}
-- |
-- Module : Data.Array.Accelerate.CUDA.Execute
-- 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.Execute (
-- * Execute a computation under a CUDA environment
executeAcc, executeAfun1,
-- -- * Executing a sequence computation and streaming its output.
-- StreamSeq(..), streamSeq,
) where
-- friends
import Data.Array.Accelerate.CUDA.AST
import Data.Array.Accelerate.CUDA.Analysis.Shape
import Data.Array.Accelerate.CUDA.Array.Data
import Data.Array.Accelerate.CUDA.Array.Sugar
import Data.Array.Accelerate.CUDA.CodeGen.Base ( Name, namesOfArray, groupOfInt )
import Data.Array.Accelerate.CUDA.Execute.Event ( Event )
import Data.Array.Accelerate.CUDA.Execute.Stream ( Stream )
import Data.Array.Accelerate.CUDA.Foreign.Import ( canExecuteAcc )
import Data.Array.Accelerate.CUDA.State
import qualified Data.Array.Accelerate.CUDA.Array.Prim as Prim
import qualified Data.Array.Accelerate.CUDA.Debug as D
import qualified Data.Array.Accelerate.CUDA.Execute.Event as Event
import qualified Data.Array.Accelerate.CUDA.Execute.Stream as Stream
import Data.Array.Accelerate.Error
import Data.Array.Accelerate.Interpreter ( evalPrim, evalPrimConst, evalPrj )
import Data.Array.Accelerate.Array.Data ( ArrayElt, ArrayData )
import Data.Array.Accelerate.Array.Representation ( SliceIndex(..) )
import Data.Array.Accelerate.FullList ( FullList(..), List(..) )
import Data.Array.Accelerate.Lifetime ( withLifetime )
import qualified Data.Array.Accelerate.Array.Representation as R
-- standard library
import Control.Applicative hiding ( Const )
import Control.Monad ( join, when, liftM )
import Control.Monad.Reader ( asks )
import Control.Monad.State ( gets )
import Control.Monad.Trans ( MonadIO, liftIO )
import Control.Monad.Trans.Cont ( ContT(..) )
import System.IO.Unsafe ( unsafeInterleaveIO )
import Data.Int
import Data.Word
import Prelude hiding ( exp, sum, iterate )
import Foreign.CUDA.Analysis.Device ( computeCapability, Compute(..) )
import qualified Foreign.CUDA.Driver as CUDA
import qualified Data.HashMap.Strict as Map
-- Asynchronous kernel execution
-- -----------------------------
-- Arrays with an associated CUDA Event that will be signalled once the
-- computation has completed.
--
data Async a = Async {-# UNPACK #-} !Event !a
-- Valuation for an environment of asynchronous array computations
--
data Aval env where
Aempty :: Aval ()
Apush :: Aval env -> Async t -> Aval (env, t)
-- -- A suspended sequence computation.
-- newtype StreamSeq a = StreamSeq (CIO (Maybe (a, StreamSeq a)))
-- Projection of a value from a valuation using a de Bruijn index.
--
aprj :: Idx env t -> Aval env -> Async t
aprj ZeroIdx (Apush _ x) = x
aprj (SuccIdx idx) (Apush val _) = aprj idx val
aprj _ _ = $internalError "aprj" "inconsistent valuation"
-- All work submitted to the given stream will occur after the asynchronous
-- event for the given array has been fulfilled. Synchronisation is performed
-- efficiently on the device. This function returns immediately.
--
after :: MonadIO m => Stream -> Async a -> m a
after stream (Async event arr) = liftIO $ Event.after event stream >> return arr
-- Block the calling thread until the event for the given array computation
-- is recorded.
--
wait :: MonadIO m => Async a -> m a
wait (Async e x) = liftIO $ Event.block e >> return x
-- Execute the given computation in a unique execution stream.
--
streaming :: (Stream -> CIO a) -> (Async a -> CIO b) -> CIO b
streaming first second = do
context <- asks activeContext
reservoir <- gets streamReservoir
table <- gets eventTable
Stream.streaming context reservoir table first (\e a -> second (Async e a))
-- Array expression evaluation
-- ---------------------------
-- Computations are evaluated by traversing the AST bottom-up, and for each node
-- distinguishing between three cases:
--
-- 1. If it is a Use node, return a reference to the device memory holding the
-- array data
--
-- 2. If it is a non-skeleton node, such as a let-binding or shape conversion,
-- this is executed directly by updating the environment or similar
--
-- 3. If it is a skeleton node, the associated binary object is retrieved,
-- memory allocated for the result, and the kernel(s) that implement the
-- skeleton are invoked
--
executeAcc :: Arrays a => ExecAcc a -> CIO a
executeAcc !acc = streaming (executeOpenAcc acc Aempty) wait
executeAfun1 :: (Arrays a, Arrays b) => ExecAfun (a -> b) -> a -> CIO b
executeAfun1 !afun !arrs = do
streaming (useArrays (arrays arrs) (fromArr arrs))
(\(Async event ()) -> executeOpenAfun1 afun Aempty (Async event arrs))
where
useArrays :: ArraysR arrs -> arrs -> Stream -> CIO ()
useArrays ArraysRunit () _ = return ()
useArrays (ArraysRpair r1 r0) (a1, a0) st = useArrays r1 a1 st >> useArrays r0 a0 st
useArrays ArraysRarray arr st = useArrayAsync arr (Just st)
executeOpenAfun1 :: PreOpenAfun ExecOpenAcc aenv (a -> b) -> Aval aenv -> Async a -> CIO b
executeOpenAfun1 (Alam (Abody f)) aenv x = streaming (executeOpenAcc f (aenv `Apush` x)) wait
executeOpenAfun1 _ _ _ = error "the sword comes out after you swallow it, right?"
-- Evaluate an open array computation
--
executeOpenAcc
:: forall aenv arrs.
ExecOpenAcc aenv arrs
-> Aval aenv
-> Stream
-> CIO arrs
executeOpenAcc EmbedAcc{} _ _
= $internalError "execute" "unexpected delayed array"
-- executeOpenAcc (ExecSeq l) !aenv !stream
-- = executeSequence l aenv stream
executeOpenAcc (ExecAcc (FL () kernel more) !gamma !pacc) !aenv !stream
= case pacc of
-- Array introduction
Use arr -> return (toArr arr)
Unit x -> fromFunction Z . const =<< travE x
-- Environment manipulation
Avar ix -> after stream (aprj ix aenv)
Alet bnd body -> streaming (executeOpenAcc bnd aenv) (\x -> executeOpenAcc body (aenv `Apush` x) stream)
Apply f a -> streaming (executeOpenAcc a aenv) (executeOpenAfun1 f aenv)
Atuple tup -> toAtuple <$> travT tup
Aprj ix tup -> evalPrj ix . fromAtuple <$> travA tup
Acond p t e -> travE p >>= \x -> if x then travA t else travA e
Awhile p f a -> awhile p f =<< travA a
-- Foreign
Aforeign ff afun a -> aforeign ff afun =<< travA a
-- Producers
Map _ a -> executeOp =<< extent a
Generate sh _ -> executeOp =<< travE sh
Transform sh _ _ _ -> executeOp =<< travE sh
Backpermute sh _ _ -> executeOp =<< travE sh
Reshape sh a -> reshapeOp <$> travE sh <*> travA a
-- Consumers
Fold _ _ a -> foldOp =<< extent a
Fold1 _ a -> fold1Op =<< extent a
FoldSeg _ _ a s -> join $ foldSegOp <$> extent a <*> extent s
Fold1Seg _ a s -> join $ foldSegOp <$> extent a <*> extent s
Scanl1 _ a -> scan1Op =<< extent a
Scanr1 _ a -> scan1Op =<< extent a
Scanl' _ _ a -> scan'Op =<< extent a
Scanr' _ _ a -> scan'Op =<< extent a
Scanl _ _ a -> scanOp True =<< extent a
Scanr _ _ a -> scanOp False =<< extent a
Permute _ d _ a -> join $ permuteOp <$> extent a <*> travA d
Stencil _ _ a -> stencilOp =<< travA a
Stencil2 _ _ a1 _ a2 -> join $ stencil2Op <$> travA a1 <*> travA a2
-- AST nodes that should be inaccessible at this point
Replicate{} -> fusionError
Slice{} -> fusionError
ZipWith{} -> fusionError
-- Collect{} -> streamingError
where
fusionError = $internalError "executeOpenAcc" "unexpected fusible matter"
-- streamingError = $internalError "executeOpenAcc" "unexpected sequence computation"
-- term traversals
travA :: ExecOpenAcc aenv a -> CIO a
travA !acc = executeOpenAcc acc aenv stream
travE :: ExecExp aenv t -> CIO t
travE !exp = executeExp exp aenv stream
travT :: Atuple (ExecOpenAcc aenv) t -> CIO t
travT NilAtup = return ()
travT (SnocAtup !t !a) = (,) <$> travT t <*> travA a
awhile :: PreOpenAfun ExecOpenAcc aenv (a -> Scalar Bool) -> PreOpenAfun ExecOpenAcc aenv (a -> a) -> a -> CIO a
awhile p f a = do
tbl <- gets eventTable
ctx <- asks activeContext
nop <- liftIO $ Event.create ctx tbl -- record event never call, so this is a functional no-op
r <- executeOpenAfun1 p aenv (Async nop a)
ok <- indexArray r 0 -- TLM TODO: memory manager should remember what is already on the host
if ok then awhile p f =<< executeOpenAfun1 f aenv (Async nop a)
else return a
aforeign :: (Arrays as, Arrays bs, Foreign asm) => asm (as -> bs) -> PreAfun ExecOpenAcc (as -> bs) -> as -> CIO bs
aforeign ff next a =
case canExecuteAcc ff of
Just asm -> asm stream a
Nothing -> executeAfun1 next a
-- get the extent of an embedded array
extent :: Shape sh => ExecOpenAcc aenv (Array sh e) -> CIO sh
extent ExecAcc{} = $internalError "executeOpenAcc" "expected delayed array"
-- extent ExecSeq{} = $internalError "executeOpenAcc" "expected delayed array"
extent (EmbedAcc sh) = travE sh
-- Skeleton implementation
-- -----------------------
-- Execute a skeleton that has no special requirements: thread decomposition
-- is based on the given shape.
--
executeOp :: (Shape sh, Elt e) => sh -> CIO (Array sh e)
executeOp !sh = do
out <- allocateArray sh
execute kernel gamma aenv (size sh) out stream
return out
-- Change the shape of an array without altering its contents. This does not
-- execute any kernel programs.
--
reshapeOp :: Shape sh => sh -> Array sh' e -> Array sh e
reshapeOp sh (Array sh' adata)
= $boundsCheck "reshape" "shape mismatch" (size sh == R.size sh')
$ Array (fromElt sh) adata
-- Executing fold operations depend on whether we are recursively collapsing
-- to a single value using multiple thread blocks, or a multidimensional
-- single-pass reduction where there is one block per inner dimension.
--
fold1Op :: (Shape sh, Elt e) => (sh :. Int) -> CIO (Array sh e)
fold1Op !sh@(_ :. sz)
= $boundsCheck "fold1" "empty array" (sz > 0)
$ foldCore sh
foldOp :: (Shape sh, Elt e) => (sh :. Int) -> CIO (Array sh e)
foldOp !(!sh :. sz)
= foldCore ((listToShape . map (max 1) . shapeToList $ sh) :. sz)
foldCore :: (Shape sh, Elt e) => (sh :. Int) -> CIO (Array sh e)
foldCore !(!sh :. sz)
| rank sh > 0 = executeOp sh
| otherwise
= let !numElements = size sh * sz
(_,!numBlocks,_) = configure kernel numElements
in do
out <- allocateArray (sh :. numBlocks)
execute kernel gamma aenv numElements out stream
foldRec out
-- Recursive step(s) of a multi-block reduction
--
foldRec :: (Shape sh, Elt e) => Array (sh:.Int) e -> CIO (Array sh e)
foldRec arr@(Array _ !adata)
| Cons _ rec _ <- more
= let sh :. sz = shape arr
!numElements = size sh * sz
(_,!numBlocks,_) = configure rec numElements
in if sz <= 1
then return $ Array (fromElt sh) adata
else do
out <- allocateArray (sh :. numBlocks)
execute rec gamma aenv numElements (out, arr) stream
foldRec out
| otherwise
= $internalError "foldRec" "missing phase-2 kernel module"
-- Segmented reduction. Subtract one from the size of the segments vector as
-- this is the result of an exclusive scan to calculate segment offsets.
--
foldSegOp :: (Shape sh, Elt e) => (sh :. Int) -> (Z :. Int) -> CIO (Array (sh :. Int) e)
foldSegOp (!sh :. _) !(Z :. sz) = executeOp (sh :. sz - 1)
-- Scans, all variations on a theme.
--
scanOp :: forall sh e. (Shape sh, Elt e) => Bool -> (sh :. Int) -> CIO (Array (sh:.Int) e)
scanOp !left !(_ :. numElements)
| Just Refl <- matchShapeType (undefined::sh) (undefined::Z)
= do
arr@(Array _ adata) <- allocateArray (Z :. numElements + 1)
withDevicePtrs adata (Just stream) $ \out -> do
let (!body, !sum)
| left = (out, advancePtrsOfArrayData adata numElements out)
| otherwise = (advancePtrsOfArrayData adata 1 out, out)
--
scanCore numElements arr body sum
return arr
scan1Op :: forall sh e. (Shape sh, Elt e) => (sh :. Int) -> CIO (Array (sh:.Int) e)
scan1Op !(_ :. numElements)
| Just Refl <- matchShapeType (undefined::sh) (undefined::Z)
= do
arr@(Array _ adata) <- allocateArray (Z :. numElements + 1) :: CIO (Vector e)
withDevicePtrs adata (Just stream) $ \body -> do
let sum {- to fix type -} = advancePtrsOfArrayData adata numElements body
--
scanCore numElements arr body sum
return (Array ((),numElements) adata)
scan'Op :: forall sh e. (Shape sh, Elt e) => (sh :. Int) -> CIO (Array (sh:.Int) e, Array sh e)
scan'Op !(_ :. numElements)
| Just Refl <- matchShapeType (undefined::sh) (undefined::Z)
= do
vec@(Array _ ad_vec) <- allocateArray (Z :. numElements) :: CIO (Vector e)
sum@(Array _ ad_sum) <- allocateArray Z :: CIO (Scalar e)
withDevicePtrs ad_vec (Just stream) $ \d_vec ->
withDevicePtrs ad_sum (Just stream) $ \d_sum -> do
--
scanCore numElements vec d_vec d_sum
return (vec, sum)
scanCore
:: forall e. Elt e
=> Int
-> Vector e -- to fix Elt vs. EltRepr
-> Prim.DevicePtrs (EltRepr e)
-> Prim.DevicePtrs (EltRepr e)
-> CIO ()
scanCore !numElements (Array _ !adata) !body !sum
| Cons _ !upsweep1 (Cons _ !upsweep2 _) <- more
= let (_,!numIntervals,_) = configure kernel numElements
!d_body = marshalDevicePtrs adata body
!d_sum = marshalDevicePtrs adata sum
in do
blk <- allocateArray (Z :. numIntervals) :: CIO (Vector e)
-- Phase 1: Split the array over multiple thread blocks and calculate
-- the final scan result from each interval.
--
when (numIntervals > 1) $ do
execute upsweep1 gamma aenv numElements blk stream
execute upsweep2 gamma aenv numIntervals (blk, blk, d_sum) stream
-- Phase 2: Re-scan the input using the carry-in value from each
-- interval sum calculated in phase 1.
--
execute kernel gamma aenv numElements (numElements, d_body, blk, d_sum) stream
| otherwise
= $internalError "scanOp" "missing multi-block kernel module(s)"
-- Forward permutation
--
permuteOp :: forall sh sh' e. (Shape sh, Shape sh', Elt e) => sh -> Array sh' e -> CIO (Array sh' e)
permuteOp !sh !dfs = do
let sh' = shape dfs
n' = size sh'
out <- allocateArray sh'
Array _ locks <- allocateArray sh' :: CIO (Array sh' Int32)
withDevicePtrs locks (Just stream) $ \d_locks -> do
liftIO $ CUDA.memsetAsync d_locks n' 0 (Just stream) -- TLM: overlap these two operations?
copyArrayAsync dfs out (Just stream)
execute kernel gamma aenv (size sh) (out, d_locks) stream
return out
-- Stencil operations. NOTE: the arguments to 'namesOfArray' must be the
-- same as those given in the function 'mkStencil[2]'.
--
stencilOp :: forall sh a b. (Shape sh, Elt a, Elt b) => Array sh a -> CIO (Array sh b)
stencilOp !arr = do
let sh = shape arr
out <- allocateArray sh
dev <- asks deviceProperties
if computeCapability dev < Compute 2 0
then marshalAccTex (namesOfArray "Stencil" (undefined :: a)) kernel arr (Just stream) $
execute kernel gamma aenv (size sh) (out, sh) stream
else execute kernel gamma aenv (size sh) (out, arr) stream
execute kernel gamma aenv (size sh) (out, arr) stream
--
return out
stencil2Op :: forall sh a b c. (Shape sh, Elt a, Elt b, Elt c)
=> Array sh a -> Array sh b -> CIO (Array sh c)
stencil2Op !arr1 !arr2
| Cons _ spec _ <- more
= let sh1 = shape arr1
sh2 = shape arr2
(sh, op)
| fromElt sh1 == fromElt sh2 = (sh1, spec)
| otherwise = (sh1 `intersect` sh2, kernel)
in do
out <- allocateArray sh
dev <- asks deviceProperties
if computeCapability dev < Compute 2 0
then marshalAccTex (namesOfArray "Stencil1" (undefined :: a)) op arr1 (Just stream) $
marshalAccTex (namesOfArray "Stencil2" (undefined :: b)) op arr2 (Just stream) $
execute op gamma aenv (size sh) (out, sh1, sh2) stream
else execute op gamma aenv (size sh) (out, arr1, arr2) stream
execute op gamma aenv (size sh) (out, arr1, arr2) stream
--
return out
| otherwise
= $internalError "stencil2Op" "missing stencil specialisation kernel"
{--
-- Execute a streaming computation
--
executeSequence
:: forall aenv arrs. Arrays arrs
=> ExecOpenSeq aenv () arrs
-> Aval aenv
-> Stream
-> CIO arrs
executeSequence topSequence aenv stream
= initializeOpenSeq topSequence aenv stream >>= loop >>= returnOut
where
loop :: ExecOpenSeq aenv () arrs -> CIO (ExecOpenSeq aenv () arrs)
loop = loop'
where
loop' :: ExecOpenSeq aenv () arrs -> CIO (ExecOpenSeq aenv () arrs)
loop' s = do
ms <- runMaybeT (stepOpenSeq aenv s stream)
case ms of
Nothing -> return s
Just s' -> loop' s'
returnOut :: forall lenv arrs' . ExecOpenSeq aenv lenv arrs' -> CIO arrs'
returnOut !l =
case l of
ExecP _ l' -> returnOut l'
ExecC c -> readConsumer c
ExecR _ ma -> case ma of
Just a -> return [a]
Nothing -> $internalError "executeSequence" "Expected already executed Sequence"
where
readConsumer :: forall a . ExecC aenv lenv a -> CIO a
readConsumer c =
case c of
ExecFoldSeq _ _ _ (Just a) -> let a' = fromList Z [a] in useArray a' >> return a'
ExecFoldSeqFlatten _ _ _ (Just a) -> return a
ExecStuple t -> toAtuple <$> rdT t
where
rdT :: forall t. Atuple (ExecC aenv lenv) t -> CIO t
rdT NilAtup = return ()
rdT (SnocAtup t' c') = (,) <$> rdT t' <*> readConsumer c'
_ -> $internalError "executeSequence" "Expected already executed consumer"
initializeSeq
:: Aval aenv
-> ExecOpenSeq aenv () arrs'
-> CIO (ExecOpenSeq aenv () arrs')
initializeSeq aenv s = streaming (initializeOpenSeq s aenv) wait
initializeOpenSeq
:: forall lenv aenv arrs'.
ExecOpenSeq aenv lenv arrs'
-> Aval aenv
-> Stream
-> CIO (ExecOpenSeq aenv lenv arrs')
initializeOpenSeq l aenv stream =
case l of
ExecP p l' -> ExecP <$> initP p <*> initializeOpenSeq l' aenv stream
ExecC c -> ExecC <$> initC c
ExecR ix a -> return (ExecR ix a)
where
initP :: forall a. ExecP aenv lenv a -> CIO (ExecP aenv lenv a)
initP (ExecToSeq slix acc k g (_::[slix])) =
do sh <- extent acc
--Lazy evaluation will stop this entire list being generated.
let slices = enumSlices slix sh :: [slix]
return $ ExecToSeq slix acc k g slices
initP (ExecUseLazy slix arr (_::[slix])) =
let sh = shape arr
-- Same as above.
slices = enumSlices slix sh :: [slix]
in return $ ExecUseLazy slix arr slices
initP s@ExecStreamIn{} = return s
initP s@ExecMap{} = return s
initP s@ExecZipWith{} = return s
initP (ExecScanSeq f a0 ix _) =
do a <- executeExp a0 aenv stream
return $ ExecScanSeq f a0 ix (Just a)
initC :: forall a. ExecC aenv lenv a -> CIO (ExecC aenv lenv a)
initC c =
case c of
ExecFoldSeq f a0 ix _ ->
do a <- executeExp a0 aenv stream
return $ ExecFoldSeq f a0 ix (Just a)
ExecFoldSeqFlatten afun acc ix _ ->
do a <- executeOpenAcc acc aenv stream
return $ ExecFoldSeqFlatten afun acc ix (Just a)
ExecStuple t -> ExecStuple <$> initCT t
initCT :: forall t. Atuple (ExecC aenv lenv) t -> CIO (Atuple (ExecC aenv lenv) t)
initCT NilAtup = return NilAtup
initCT (SnocAtup t c) = SnocAtup <$> initCT t <*> initC c
-- get the extent of an embedded array
extent :: Shape sh => ExecOpenAcc aenv (Array sh e) -> CIO sh
extent ExecAcc{} = $internalError "executeOpenAcc" "expected delayed array"
extent ExecSeq{} = $internalError "executeOpenAcc" "expected delayed array"
extent (EmbedAcc sh) = executeExp sh aenv stream
-- Turn a sequence computation into a suspended computation that can be forced
-- an element at a time.
--
streamSeq :: ExecSeq [a] -> StreamSeq a
streamSeq (ExecS binds sequ) = StreamSeq $ do
aenv <- executeExtend binds Aempty
iseq <- initializeSeq aenv sequ
let go s = do
ms <- stepSeq aenv s
case ms of
Nothing -> return Nothing
Just (s', a) -> return (Just (a, StreamSeq (go s')))
go iseq
stepSeq
:: forall a aenv.
Aval aenv
-> ExecOpenSeq aenv () [a]
-> CIO (Maybe (ExecOpenSeq aenv () [a], a))
stepSeq aenv s = streaming step wait
where
step :: Stream -> CIO (Maybe (ExecOpenSeq aenv () [a], a))
step stream = do
ms <- runMaybeT (stepOpenSeq aenv s stream)
return $ (,) <$> ms <*> (coll <$> ms)
coll :: ExecOpenSeq aenv lenv [a] -> a
coll (ExecP _ s') = coll s'
coll (ExecC _) = $internalError "stepSeq" "Unreachable"
coll (ExecR _ ma) = case ma of
Nothing -> $internalError "stepSeq" "Trying to collect the value of an unexecuted sequence"
Just a -> a
stepOpenSeq
:: forall aenv arrs'.
Aval aenv
-> ExecOpenSeq aenv () arrs'
-> Stream
-> MaybeT CIO (ExecOpenSeq aenv () arrs')
stepOpenSeq aenv !l stream = go l Empty
where
go :: forall lenv. ExecOpenSeq aenv lenv arrs' -> Val lenv -> MaybeT CIO (ExecOpenSeq aenv lenv arrs')
go s lenv =
case s of
ExecP p s' -> do
(p', a) <- produce p
s'' <- go s' (lenv `Push` a)
return $ ExecP p' s''
ExecC c -> ExecC <$> lift (consume c)
ExecR ix _ -> return $ ExecR ix (Just (prj ix lenv))
where
produce :: forall a . ExecP aenv lenv a -> MaybeT CIO (ExecP aenv lenv a, a)
produce (ExecToSeq slix acc kernel gamma sls) =
do
sl : sls' <- return sls
lift $ do
sh <- extent acc
out <- allocateArray (sliceShape slix sh)
m <- marshalSlice slix sl
execute kernel gamma aenv (size sh) (m, out) stream
return (ExecToSeq slix acc kernel gamma sls', out)
produce (ExecUseLazy slix arr sls) =
do let sh = shape arr
lift $ do
sl : sls' <- return sls
out <- allocateArray (sliceShape slix sh)
useArraySlice slix sl arr out
return (ExecUseLazy slix arr sls', out)
produce (ExecStreamIn xs) =
let use :: ArraysR arrs -> arrs -> CIO ()
use ArraysRunit () = return ()
use ArraysRarray arr = useArrayAsync arr Nothing
use (ArraysRpair r1 r2) (a1, a2) = use r1 a1 >> use r2 a2
in case xs of
[] -> MaybeT (return Nothing)
x:xs' -> lift (use (arrays x) (fromArr x)) >> return (ExecStreamIn xs', x)
produce (ExecMap afun x) = (ExecMap afun x ,) <$> lift (travAfun1 afun (prj x lenv))
produce (ExecZipWith afun x y) = (ExecZipWith afun x y,) <$> lift (travAfun2 afun (prj x lenv) (prj y lenv))
produce (ExecScanSeq afun a0 x ma) =
do
a <- MaybeT (return ma)
a' <- lift $ travFun2 afun a (prj x lenv ! Z)
return (ExecScanSeq afun a0 x (Just a'), fromList Z [a'])
consume :: forall a . ExecC aenv lenv a -> CIO (ExecC aenv lenv a)
consume c =
case c of
ExecFoldSeq afun a0 x (Just a) ->
do b <- indexArray (prj x lenv) 0
a' <- travFun2 afun a b
return $ ExecFoldSeq afun a0 x (Just a')
ExecFoldSeqFlatten afun a0 x (Just a) ->
do useArray shapes
a' <- travAfun3 afun a shapes elems
return $ ExecFoldSeqFlatten afun a0 x (Just a')
where
Array sh adata = prj x lenv
elems = Array ((), R.size sh) adata
shapes = fromList (Z:.1) [toElt sh]
ExecStuple t -> ExecStuple <$> consumeT t
_ -> $internalError "executeSequence" "Expected already executed consumer"
consumeT :: forall t. Atuple (ExecC aenv lenv) t -> CIO (Atuple (ExecC aenv lenv) t)
consumeT NilAtup = return NilAtup
consumeT (SnocAtup t c) = SnocAtup <$> consumeT t <*> consume c
travAfun1 :: forall a b. PreOpenAfun ExecOpenAcc aenv (a -> b) -> a -> CIO b
travAfun1 (Alam (Abody afun)) a =
do nop <- join $ liftIO <$> (Event.create <$> asks activeContext <*> gets eventTable)
executeOpenAcc afun (aenv `Apush` (Async nop a)) stream
travAfun1 _ _ = error "travAfun1"
travAfun2 :: forall a b c. PreOpenAfun ExecOpenAcc aenv (a -> b -> c) -> a -> b -> CIO c
travAfun2 (Alam (Alam (Abody afun))) a b =
do nop <- join $ liftIO <$> (Event.create <$> asks activeContext <*> gets eventTable)
executeOpenAcc afun (aenv `Apush` (Async nop a) `Apush` (Async nop b)) stream
travAfun2 _ _ _ = error "travAfun2"
travAfun3 :: forall a b c d. PreOpenAfun ExecOpenAcc aenv (a -> b -> c -> d) -> a -> b -> c -> CIO d
travAfun3 (Alam (Alam (Alam (Abody afun)))) a b c =
do nop <- join $ liftIO <$> (Event.create <$> asks activeContext <*> gets eventTable)
executeOpenAcc afun (aenv `Apush` (Async nop a) `Apush` (Async nop b) `Apush` (Async nop c)) stream
travAfun3 _ _ _ _ = error "travAfun3"
travFun2 :: forall a b c. PreFun ExecOpenAcc aenv (a -> b -> c) -> a -> b -> CIO c
travFun2 (Lam (Lam (Body c))) a b = executeOpenExp c (Empty `Push` a `Push` b) aenv stream
travFun2 _ _ _ = error "travFun2"
-- get the extent of an embedded array
extent :: Shape sh => ExecOpenAcc aenv (Array sh e) -> CIO sh
extent ExecAcc{} = $internalError "executeOpenAcc" "expected delayed array"
extent ExecSeq{} = $internalError "executeOpenAcc" "expected delayed array"
extent (EmbedAcc sh) = executeExp sh aenv stream
-- Evaluating bindings
-- -------------------
executeExtend :: Extend ExecOpenAcc aenv aenv' -> Aval aenv -> CIO (Aval aenv')
executeExtend BaseEnv aenv = return aenv
executeExtend (PushEnv e a) aenv = do
aenv' <- executeExtend e aenv
streaming (executeOpenAcc a aenv') $ \a' -> return $ Apush aenv' a'
--}
-- Scalar expression evaluation
-- ----------------------------
executeExp :: ExecExp aenv t -> Aval aenv -> Stream -> CIO t
executeExp !exp !aenv !stream = executeOpenExp exp Empty aenv stream
executeOpenExp
:: forall env aenv exp.
ExecOpenExp env aenv exp
-> Val env
-> Aval aenv
-> Stream
-> CIO exp
executeOpenExp !rootExp !env !aenv !stream = travE rootExp
where
travE :: ExecOpenExp env aenv t -> CIO t
travE exp = case exp of
Var ix -> return (prj ix env)
Let bnd body -> travE bnd >>= \x -> executeOpenExp body (env `Push` x) aenv stream
Const c -> return (toElt c)
PrimConst c -> return (evalPrimConst c)
PrimApp f x -> evalPrim f <$> travE x
Tuple t -> toTuple <$> travT t
Prj ix e -> evalPrj ix . fromTuple <$> travE e
Cond p t e -> travE p >>= \x -> if x then travE t else travE e
While p f x -> while p f =<< travE x
IndexAny -> return Any
IndexNil -> return Z
IndexCons sh sz -> (:.) <$> travE sh <*> travE sz
IndexHead sh -> (\(_ :. ix) -> ix) <$> travE sh
IndexTail sh -> (\(ix :. _) -> ix) <$> travE sh
IndexSlice ix slix sh -> indexSlice ix <$> travE slix <*> travE sh
IndexFull ix slix sl -> indexFull ix <$> travE slix <*> travE sl
ToIndex sh ix -> toIndex <$> travE sh <*> travE ix
FromIndex sh ix -> fromIndex <$> travE sh <*> travE ix
Intersect sh1 sh2 -> intersect <$> travE sh1 <*> travE sh2
Union sh1 sh2 -> union <$> travE sh1 <*> travE sh2
ShapeSize sh -> size <$> travE sh
Shape acc -> shape <$> travA acc
Index acc ix -> join $ index <$> travA acc <*> travE ix
LinearIndex acc ix -> join $ indexArray <$> travA acc <*> travE ix
Foreign _ f x -> foreign f x
-- Helpers
-- -------
travT :: Tuple (ExecOpenExp env aenv) t -> CIO t
travT tup = case tup of
NilTup -> return ()
SnocTup !t !e -> (,) <$> travT t <*> travE e
travA :: ExecOpenAcc aenv a -> CIO a
travA !acc = executeOpenAcc acc aenv stream
foreign :: ExecFun () (a -> b) -> ExecOpenExp env aenv a -> CIO b
foreign (Lam (Body f)) x = travE x >>= \e -> executeOpenExp f (Empty `Push` e) Aempty stream
foreign _ _ = error "I bless the rains down in Africa"
travF1 :: ExecOpenFun env aenv (a -> b) -> a -> CIO b
travF1 (Lam (Body f)) x = executeOpenExp f (env `Push` x) aenv stream
travF1 _ _ = error "Gonna take some time to do the things we never have"
while :: ExecOpenFun env aenv (a -> Bool) -> ExecOpenFun env aenv (a -> a) -> a -> CIO a
while !p !f !x = do
ok <- travF1 p x
if ok then while p f =<< travF1 f x
else return x
indexSlice :: (Elt slix, Elt sh, Elt sl)
=> SliceIndex (EltRepr slix) (EltRepr sl) co (EltRepr sh)
-> slix
-> sh
-> sl
indexSlice !ix !slix !sh = toElt $! restrict ix (fromElt slix) (fromElt sh)
where
restrict :: SliceIndex slix sl co sh -> slix -> sh -> sl
restrict SliceNil () () = ()
restrict (SliceAll sliceIdx) (slx, ()) (sl, sz) = (restrict sliceIdx slx sl, sz)
restrict (SliceFixed sliceIdx) (slx, _) (sl, _) = restrict sliceIdx slx sl
indexFull :: (Elt slix, Elt sh, Elt sl)
=> SliceIndex (EltRepr slix) (EltRepr sl) co (EltRepr sh)
-> slix
-> sl
-> sh
indexFull !ix !slix !sl = toElt $! extend ix (fromElt slix) (fromElt sl)
where
extend :: SliceIndex slix sl co sh -> slix -> sl -> sh
extend SliceNil () () = ()
extend (SliceAll sliceIdx) (slx, ()) (sh, sz) = (extend sliceIdx slx sh, sz)
extend (SliceFixed sliceIdx) (slx, sz) sh = (extend sliceIdx slx sh, sz)
index :: (Shape sh, Elt e) => Array sh e -> sh -> CIO e
index !arr !ix = indexArray arr (toIndex (shape arr) ix)
-- Marshalling data
-- ----------------
{--
marshalSlice'
:: SliceIndex slix sl co dim
-> slix
-> CIO [CUDA.FunParam]
marshalSlice' SliceNil () = return []
marshalSlice' (SliceAll sl) (sh, ()) = marshalSlice' sl sh
marshalSlice' (SliceFixed sl) (sh, n) =
do x <- runContT (marshal n Nothing) return
xs <- marshalSlice' sl sh
return (xs ++ x)
marshalSlice
:: Elt slix => SliceIndex (EltRepr slix) sl co dim
-> slix
-> CIO [CUDA.FunParam]
marshalSlice slix = marshalSlice' slix . fromElt
--}
-- Data which can be marshalled as function arguments to a kernel invocation.
--
class Marshalable a where
marshal :: a -> Maybe Stream -> ContT b CIO [CUDA.FunParam]
instance Marshalable () where
marshal () _ = return []
instance Marshalable CUDA.FunParam where
marshal !x _ = return [x]
instance ArrayElt e => Marshalable (ArrayData e) where
marshal !ad ms = ContT $ marshalArrayData ad ms
instance Shape sh => Marshalable sh where
marshal !sh ms = marshal (reverse (shapeToList sh)) ms
instance Marshalable a => Marshalable [a] where
marshal xs ms = concatMapM (flip marshal ms) xs
instance (Marshalable sh, Elt e) => Marshalable (Array sh e) where
marshal !(Array sh ad) ms = (++) <$> marshal (toElt sh :: sh) ms <*> marshal ad ms
instance (Marshalable a, Marshalable b) => Marshalable (a, b) where
marshal (!a, !b) ms = (++) <$> marshal a ms <*> marshal b ms
instance (Marshalable a, Marshalable b, Marshalable c) => Marshalable (a, b, c) where
marshal (!a, !b, !c) ms
= concat <$> sequence [marshal a ms, marshal b ms, marshal c ms]
instance (Marshalable a, Marshalable b, Marshalable c, Marshalable d)
=> Marshalable (a, b, c, d) where
marshal (!a, !b, !c, !d) ms
= concat <$> sequence [marshal a ms, marshal b ms, marshal c ms, marshal d ms]
#define primMarshalable(ty) \
instance Marshalable (ty) where { \
marshal !x _ = return [CUDA.VArg x] }
primMarshalable(Int)
primMarshalable(Int8)
primMarshalable(Int16)
primMarshalable(Int32)
primMarshalable(Int64)
primMarshalable(Word)
primMarshalable(Word8)
primMarshalable(Word16)
primMarshalable(Word32)
primMarshalable(Word64)
primMarshalable(Float)
primMarshalable(Double)
primMarshalable(CUDA.DevicePtr a)
-- Note [Array references in scalar code]
-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
--
-- All CUDA devices have between 6-8KB of read-only texture memory per
-- multiprocessor. Since all arrays in Accelerate are immutable, we can always
-- access input arrays through the texture cache to reduce global memory demand
-- when accesses do not follow the regular patterns required for coalescing.
--
-- This is great for older 1.x series devices, but newer devices have a
-- dedicated L2 cache (device dependent, 256KB-1.5MB), as well as a configurable
-- L1 cache combined with shared memory (16-48KB).
--
-- For older 1.x series devices, we pass free array variables as texture
-- references, but for new devices we pass them as standard array arguments so
-- as to use the larger available caches.
--
marshalAccEnvTex :: AccKernel a -> Aval aenv -> Gamma aenv -> Stream -> ContT b CIO [CUDA.FunParam]
marshalAccEnvTex !kernel !aenv (Gamma !gamma) !stream
= flip concatMapM (Map.toList gamma)
$ \(Idx_ !(idx :: Idx aenv (Array sh e)), i) ->
do arr <- after stream (aprj idx aenv)
ContT $ \f -> marshalAccTex (namesOfArray (groupOfInt i) (undefined :: e)) kernel arr (Just stream) (f ())
marshal (shape arr) (Just stream)
marshalAccTex :: (Name,[Name]) -> AccKernel a -> Array sh e -> Maybe Stream -> CIO b -> CIO b
marshalAccTex (_, !arrIn) (AccKernel _ _ !lmdl _ _ _ _) (Array !sh !adata) ms run
= do
texs <- liftIO $ withLifetime lmdl $ \mdl -> (sequence' $ map (CUDA.getTex mdl) (reverse arrIn))
marshalTextureData adata (R.size sh) texs ms (const run)
marshalAccEnvArg :: Aval aenv -> Gamma aenv -> Stream -> ContT b CIO [CUDA.FunParam]
marshalAccEnvArg !aenv (Gamma !gamma) !stream
= concatMapM (\(Idx_ !idx) -> flip marshal (Just stream) =<< after stream (aprj idx aenv)) (Map.keys gamma)
-- A lazier version of 'Control.Monad.sequence'
--
sequence' :: [IO a] -> IO [a]
sequence' = foldr k (return [])
where k m ms = do { x <- m; xs <- unsafeInterleaveIO ms; return (x:xs) }
-- Generalise concatMap for teh monadz
--
concatMapM :: Monad m => (a -> m [b]) -> [a] -> m [b]
concatMapM f xs = concat `liftM` mapM f xs
-- Kernel execution
-- ----------------
-- What launch parameters should we use to execute the kernel with a number of
-- array elements?
--
configure :: AccKernel a -> Int -> (Int, Int, Int)
configure (AccKernel _ _ _ _ !cta !smem !grid) !n = (cta, grid n, smem)
-- Marshal the kernel arguments. For older 1.x devices this binds free arrays to
-- texture references, and for newer devices adds the parameters to the front of
-- the argument list
--
arguments
:: Marshalable args
=> AccKernel a
-> Aval aenv
-> Gamma aenv
-> args
-> Stream
-> ContT b CIO [CUDA.FunParam]
arguments !kernel !aenv !gamma !a !stream = do
dev <- asks deviceProperties
let marshaller | computeCapability dev < Compute 2 0 = marshalAccEnvTex kernel
| otherwise = marshalAccEnvArg
--
(++) <$> marshaller aenv gamma stream <*> marshal a (Just stream)
-- Link the binary object implementing the computation, configure the kernel
-- launch parameters, and initiate the computation. This also handles lifting
-- and binding of array references from scalar expressions.
--
execute
:: Marshalable args
=> AccKernel a -- The binary module implementing this kernel
-> Gamma aenv -- variables of arrays embedded in scalar expressions
-> Aval aenv -- the environment
-> Int -- a "size" parameter, typically number of elements in the output
-> args -- arguments to marshal to the kernel function
-> Stream -- Compute stream to execute in
-> CIO ()
execute !kernel !gamma !aenv !n !a !stream = flip runContT return $ do
args <- arguments kernel aenv gamma a stream
liftIO $ launch kernel (configure kernel n) args stream
-- Execute a device function, with the given thread configuration and function
-- parameters. The tuple contains (threads per block, grid size, shared memory)
--
launch :: AccKernel a -> (Int,Int,Int) -> [CUDA.FunParam] -> Stream -> IO ()
launch (AccKernel entry !fn _ _ _ _ _) !(cta, grid, smem) !args !stream
= D.timed D.dump_exec msg (Just stream)
$ CUDA.launchKernel fn (grid,1,1) (cta,1,1) smem (Just stream) args
where
msg gpuTime cpuTime
= "exec: " ++ entry ++ "<<< " ++ shows grid ", " ++ shows cta ", " ++ shows smem " >>> "
++ D.elapsed gpuTime cpuTime