accelerate-llvm-1.3.0.0: src/Data/Array/Accelerate/LLVM/Array/Data.hs
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# OPTIONS_HADDOCK hide #-}
-- |
-- Module : Data.Array.Accelerate.LLVM.Array.Data
-- Copyright : [2014..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.Array.Data (
Remote(..),
newRemote, newRemoteAsync,
useRemote,
copyToRemote,
copyToHost,
copyToPeer,
indexRemote,
runIndexArray, runIndexArrayAsync,
runArray, runArrayAsync,
runArrays, runArraysAsync,
) where
import Data.Array.Accelerate.Array.Data
import Data.Array.Accelerate.Representation.Array
import Data.Array.Accelerate.Representation.Shape
import Data.Array.Accelerate.Representation.Type
import Data.Array.Accelerate.Type
import Data.Array.Accelerate.LLVM.Execute.Async
import Control.Monad ( liftM, liftM2 )
import Prelude
class Async arch => Remote arch where
-- | Allocate a new uninitialised array on the remote device.
--
allocateRemote :: ArrayR (Array sh e) -> sh -> Par arch (Array sh e)
-- | Use the given immutable array on the remote device. Since the source
-- array is immutable, the garbage collector can evict and re-upload the data
-- as necessary without copy-back.
--
{-# INLINE useRemoteR #-}
useRemoteR
:: SingleType e
-> Int -- ^ number of elements to copy
-> ArrayData e -- ^ array payload
-> Par arch (FutureR arch (ArrayData e))
useRemoteR _ _ = newFull
-- | Upload an array from the host to the remote device.
--
{-# INLINE copyToRemoteR #-}
copyToRemoteR
:: SingleType e
-> Int -- ^ number of elements to copy
-> ArrayData e -- ^ array payload
-> Par arch (FutureR arch (ArrayData e))
copyToRemoteR _ _ = newFull
-- | Copy an array from the remote device back to the host.
--
{-# INLINE copyToHostR #-}
copyToHostR
:: SingleType e
-> Int -- ^ number of elements to copy
-> ArrayData e -- ^ array payload
-> Par arch (FutureR arch (ArrayData e))
copyToHostR _ _ = newFull
-- | Copy a section of an array between two remote instances of the same type.
-- This may be more efficient than copying to the host and then to the second
-- remote instance (e.g. DMA between two CUDA devices).
--
{-# INLINE copyToPeerR #-}
copyToPeerR
:: arch -- ^ remote device to copy to
-> SingleType e
-> Int -- ^ number of elements to copy
-> ArrayData e -- ^ array payload
-> Par arch (FutureR arch (ArrayData e))
copyToPeerR _ _ _ = newFull
-- | Upload an immutable array from the host to the remote device,
-- asynchronously. Since the source array is immutable, the garbage collector
-- can evict and re-upload the data as necessary without copy-back. This may
-- upload each array payload in a separate execution stream, thereby making us
-- of multiple memcpy engines.
--
{-# INLINE useRemoteAsync #-}
useRemoteAsync :: ArraysR arrs -> arrs -> Par arch (FutureArraysR arch arrs)
useRemoteAsync repr arrs =
runArraysAsync repr arrs $ \(ArrayR shr tp) arr ->
let n = size shr (shape arr)
in runArrayAsync tp arr $ \m tp' ad ->
useRemoteR tp' (n*m) ad
-- | Upload an existing array to the remote device, asynchronously.
--
{-# INLINE copyToRemoteAsync #-}
copyToRemoteAsync :: ArraysR arrs -> arrs -> Par arch (FutureArraysR arch arrs)
copyToRemoteAsync reprs arrs =
runArraysAsync reprs arrs $ \(ArrayR shr tp) arr ->
let n = size shr (shape arr)
in runArrayAsync tp arr $ \m tp' ad ->
copyToRemoteR tp' (n*m) ad
-- | Copy an array from the remote device to the host, asynchronously
--
{-# INLINE copyToHostAsync #-}
copyToHostAsync :: ArraysR arrs -> arrs -> Par arch (FutureArraysR arch arrs)
copyToHostAsync reprs arrs =
runArraysAsync reprs arrs $ \(ArrayR shr tp) arr ->
let n = size shr (shape arr)
in runArrayAsync tp arr $ \m tp' ad ->
copyToHostR tp' (n*m) ad
-- | Copy arrays between two remote instances. This may be more efficient than
-- copying to the host and then to the second remote instance (e.g. by DMA
-- between the two remote devices).
--
{-# INLINE copyToPeerAsync #-}
copyToPeerAsync :: arch -> ArraysR arrs -> arrs -> Par arch (FutureArraysR arch arrs)
copyToPeerAsync peer reprs arrs =
runArraysAsync reprs arrs $ \(ArrayR shr tp) arr ->
let n = size shr (shape arr)
in runArrayAsync tp arr $ \m tp' ad ->
copyToPeerR peer tp' (n*m) ad
-- | Read a single element from the array at the given row-major index
--
{-# INLINE indexRemoteAsync #-}
indexRemoteAsync
:: TypeR e
-> Array sh e
-> Int
-> Par arch (FutureR arch e)
indexRemoteAsync tp (Array _ ad) i = newFull (indexArrayData tp ad i)
-- | Create a new array from its representation on the host, and upload it to
-- the remote device.
--
{-# INLINE newRemote #-}
newRemote
:: Remote arch
=> ArrayR (Array sh e)
-> sh
-> (sh -> e)
-> Par arch (Array sh e)
newRemote repr sh f =
get =<< newRemoteAsync repr sh f
-- | Create a new array from its representation on the host, and upload it as
-- a new remote array, asynchronously.
--
{-# INLINE newRemoteAsync #-}
newRemoteAsync
:: Remote arch
=> ArrayR (Array sh e)
-> sh
-> (sh -> e)
-> Par arch (FutureR arch (Array sh e))
newRemoteAsync repr sh f =
useRemoteAsync (TupRsingle repr) $! fromFunction repr sh f
-- | Upload an immutable array from the host to the remote device. This is
-- a synchronous operation in that it will not return until the transfer
-- completes, but the individual array payloads will be uploaded concurrently if
-- possible.
--
{-# INLINE useRemote #-}
useRemote :: Remote arch => ArraysR a -> a -> Par arch a
useRemote repr arrs =
getArrays repr =<< useRemoteAsync repr arrs
-- | Uploading existing arrays from the host to the remote device. This is
-- synchronous with respect to the calling thread, but the individual array
-- payloads may themselves be transferred concurrently.
--
{-# INLINE copyToRemote #-}
copyToRemote :: Remote arch => ArraysR a -> a -> Par arch a
copyToRemote repr arrs =
getArrays repr =<< copyToRemoteAsync repr arrs
-- | Copy an array from the remote device to the host. This is synchronous with
-- respect to the calling thread, but the individual array payloads may
-- themselves be transferred concurrently.
--
{-# INLINE copyToHost #-}
copyToHost :: Remote arch => ArraysR a -> a -> Par arch a
copyToHost repr arrs =
blockArrays repr =<< copyToHostAsync repr arrs
-- | Copy arrays between two remote instances of the same type. This may be more
-- efficient than copying to the host and then to the second remote instance
-- (e.g. DMA between CUDA devices).
--
{-# INLINE copyToPeer #-}
copyToPeer :: Remote arch => arch -> ArraysR a -> a -> Par arch a
copyToPeer peer repr arrs =
getArrays repr =<< copyToPeerAsync peer repr arrs
-- | Read a single element from the remote array at the given row-major index.
-- This is synchronous with respect to both the host and remote device.
--
{-# INLINE indexRemote #-}
indexRemote :: Remote arch => TypeR e -> Array sh e -> Int -> Par arch e
indexRemote tp arr i =
block =<< indexRemoteAsync tp arr i
-- Helpers for traversing the Arrays data structure
-- ------------------------------------------------
-- | Read a single element from an array at the given row-major index.
--
{-# INLINE runIndexArray #-}
runIndexArray
:: forall m sh e. Monad m
=> (forall s. ArrayData s ~ ScalarArrayData s => Int -> SingleType s -> ArrayData s -> Int -> m (ArrayData s))
-> TypeR e
-> Array sh e
-> Int
-> m e
runIndexArray worker tp (Array _ adata) i = flip (indexArrayData tp) 0 <$> indexR tp adata
where
indexR :: TypeR s -> ArrayData s -> m (ArrayData s)
indexR TupRunit !_ = return ()
indexR (TupRpair !t1 !t2) (!ad1, !ad2) = liftM2 (,) (indexR t1 ad1) (indexR t2 ad2)
indexR (TupRsingle t) !ad
| ScalarArrayDict w s <- scalarArrayDict t
, SingleArrayDict <- singleArrayDict s
= worker w s ad i
{-# INLINE runIndexArrayAsync #-}
runIndexArrayAsync
:: forall arch sh e. Async arch
=> (forall s. ArrayData s ~ ScalarArrayData s => Int -> SingleType s -> ArrayData s -> Int -> Par arch (FutureR arch (ArrayData s)))
-> TypeR e
-> Array sh e
-> Int
-> Par arch (FutureR arch e)
runIndexArrayAsync worker tp (Array _ adata) i = (flip (indexArrayData tp) 0) `liftF` indexR tp adata
where
indexR :: TypeR s -> ArrayData s -> Par arch (FutureR arch (ArrayData s))
indexR TupRunit !_ = newFull ()
indexR (TupRpair !t1 !t2) (!ad1, !ad2) = liftF2' (,) (indexR t1 ad1) (indexR t2 ad2)
indexR (TupRsingle t) !ad
| ScalarArrayDict w s <- scalarArrayDict t
, SingleArrayDict <- singleArrayDict s
= worker w s ad i
-- It is expected these transfers will be very small, so don't bother
-- creating new execution streams for them
liftF2' :: (a -> b -> c) -> Par arch (FutureR arch a) -> Par arch (FutureR arch b) -> Par arch (FutureR arch c)
liftF2' f x y = do
r <- new
x' <- x
y' <- y
put r =<< liftM2 f (get x') (get y')
return r
-- | Generalised function to traverse the Arrays structure
--
{-# INLINE runArrays #-}
runArrays
:: forall m arrs. Monad m
=> ArraysR arrs
-> arrs
-> (forall sh e. ArrayR (Array sh e) -> Array sh e -> m (Array sh e))
-> m arrs
runArrays reprs arrs worker = runR reprs arrs
where
runR :: ArraysR a -> a -> m a
runR TupRunit () = return ()
runR (TupRsingle repr@ArrayR{}) arr = worker repr arr
runR (TupRpair aeR2 aeR1) (arrs2, arrs1) = liftM2 (,) (runR aeR2 arrs2) (runR aeR1 arrs1)
{-# INLINE runArraysAsync #-}
runArraysAsync
:: forall arch arrs. Async arch
=> ArraysR arrs
-> arrs
-> (forall sh e. ArrayR (Array sh e) -> Array sh e -> Par arch (FutureR arch (Array sh e)))
-> Par arch (FutureArraysR arch arrs)
runArraysAsync reprs arrs worker = runR reprs arrs
where
runR :: ArraysR a -> a -> Par arch (FutureArraysR arch a)
runR TupRunit () = return ()
runR (TupRsingle repr@ArrayR{}) arr = worker repr arr
runR (TupRpair aeR2 aeR1) (arrs2, arrs1) = (,) <$> runR aeR2 arrs2 <*> runR aeR1 arrs1
-- | Generalised function to traverse the ArrayData structure with one
-- additional argument
--
{-# INLINE runArray #-}
runArray
:: forall m sh e. Monad m
=> TypeR e
-> Array sh e
-> (forall s. ArrayData s ~ ScalarArrayData s => Int -> SingleType s -> ScalarArrayData s -> m (ScalarArrayData s))
-> m (Array sh e)
runArray tp (Array sh adata) worker = Array sh `liftM` runR tp adata
where
runR :: TypeR s -> ArrayData s -> m (ArrayData s)
runR (TupRunit) !_ = return ()
runR (TupRpair !t2 !t1) (!ad2, !ad1) = liftM2 (,) (runR t2 ad2) (runR t1 ad1)
runR (TupRsingle !t) !ad
| ScalarArrayDict w s <- scalarArrayDict t
, SingleArrayDict <- singleArrayDict s
= worker w s ad
{-# INLINE runArrayAsync #-}
runArrayAsync
:: forall arch sh e. Async arch
=> TypeR e
-> Array sh e
-> (forall s. ArrayData s ~ ScalarArrayData s => Int -> SingleType s -> ScalarArrayData s -> Par arch (FutureR arch (ScalarArrayData s)))
-> Par arch (FutureR arch (Array sh e))
runArrayAsync tp (Array sh adata) worker = Array sh `liftF` runR tp adata
where
runR :: forall s. TypeR s -> ArrayData s -> Par arch (FutureR arch (ArrayData s))
runR (TupRunit) !_ = newFull ()
runR (TupRpair !t2 !t1) (!ad2, !ad1) = liftF2 (,) (runR t2 ad2) (runR t1 ad1)
runR (TupRsingle !t) !ad
| ScalarArrayDict w s <- scalarArrayDict t
, SingleArrayDict <- singleArrayDict s
= worker w s ad
{-# INLINE liftF #-}
liftF :: Async arch
=> (a -> b)
-> Par arch (FutureR arch a)
-> Par arch (FutureR arch b)
liftF f x = do
r <- new
x' <- x
fork $ put r . f =<< get x'
return r
{-# INLINE liftF2 #-}
liftF2 :: Async arch
=> (a -> b -> c)
-> Par arch (FutureR arch a)
-> Par arch (FutureR arch b)
-> Par arch (FutureR arch c)
liftF2 f x y = do
r <- new
x' <- spawn x
y' <- spawn y
fork $ put r =<< liftM2 f (get x') (get y')
return r