module Data.Array.Repa.Repr.Partitioned ( P, Array (..) , Range(..) , inRange) where import Data.Array.Repa.Base import Data.Array.Repa.Shape import Data.Array.Repa.Eval import Data.Array.Repa.Repr.Delayed -- | Partitioned arrays. -- The last partition takes priority -- -- These are produced by Repa's support functions and allow arrays to be defined -- using a different element function for each partition. -- -- The basic idea is described in ``Efficient Parallel Stencil Convolution'', -- Ben Lippmeier and Gabriele Keller, Haskell 2011 -- though the underlying -- array representation has changed since this paper was published. -- data P r1 r2 data Range sh = Range !sh !sh -- indices defining the range (sh -> Bool) -- predicate to check whether were in range -- | Check whether an index is within the given range. inRange :: Range sh -> sh -> Bool inRange (Range _ _ p) ix = p ix {-# INLINE inRange #-} -- Repr ----------------------------------------------------------------------- -- | Read elements from a partitioned array. instance (Source r1 e, Source r2 e) => Source (P r1 r2) e where data Array (P r1 r2) sh e = APart !sh -- size of the whole array !(Range sh) !(Array r1 sh e) -- if in range use this array !(Array r2 sh e) -- otherwise use this array index (APart _ range arr1 arr2) ix | inRange range ix = index arr1 ix | otherwise = index arr2 ix {-# INLINE index #-} linearIndex arr@(APart sh _ _ _) ix = index arr \$ fromIndex sh ix {-# INLINE linearIndex #-} extent (APart sh _ _ _) = sh {-# INLINE extent #-} deepSeqArray (APart sh range arr1 arr2) y = sh `deepSeq` range `deepSeqRange` arr1 `deepSeqArray` arr2 `deepSeqArray` y {-# INLINE deepSeqArray #-} deepSeqRange :: Shape sh => Range sh -> b -> b deepSeqRange (Range ix sz f) y = ix `deepSeq` sz `deepSeq` f `seq` y {-# INLINE deepSeqRange #-} -- Load ----------------------------------------------------------------------- instance (LoadRange r1 sh e, Load r2 sh e) => Load (P r1 r2) sh e where loadP (APart _ (Range ix sz _) arr1 arr2) marr = do loadRangeP arr1 marr ix sz loadP arr2 marr {-# INLINE loadP #-} loadS (APart _ (Range ix sz _) arr1 arr2) marr = do loadRangeS arr1 marr ix sz loadS arr2 marr {-# INLINE loadS #-}