diff --git a/Data/Array/Repa.hs b/Data/Array/Repa.hs --- a/Data/Array/Repa.hs +++ b/Data/Array/Repa.hs @@ -1,288 +1,124 @@ -{-# LANGUAGE PatternGuards, PackageImports #-} -{-# LANGUAGE ScopedTypeVariables, RankNTypes #-} -{-# LANGUAGE TypeOperators, FlexibleContexts #-} +{-# LANGUAGE PatternGuards, PackageImports, ScopedTypeVariables, RankNTypes #-} +{-# LANGUAGE TypeOperators, FlexibleContexts, NoMonomorphismRestriction, FlexibleInstances, UndecidableInstances #-} +{-# OPTIONS -fno-warn-orphans #-} -- | See the repa-examples package for examples. -- --- More information is also at --- --- NOTE: To get decent performance you must use GHC head branch > 6.13.20100309. --- --- WARNING: Most of the functions that operate on indices don't perform bounds checks. --- Doing these checks would interfere with code optimisation and reduce performance. --- Indexing outside arrays, or failing to meet the stated obligations will --- likely cause heap corruption. --- --- +-- More information at +-- module Data.Array.Repa ( module Data.Array.Repa.Shape , module Data.Array.Repa.Index , module Data.Array.Repa.Slice - - , Array (..) - -- * Constructors - , fromUArray - , fromFunction - , unit + -- from Data.Array.Repa.Internals.Elt ----------------------- + , Elt(..) - -- * Projections - , extent - , delay - , toUArray - , index, (!:) - , toScalar + -- from Data.Array.Repa.Internals.Base ---------------------- + , Array(..) + , Region(..) + , Range(..) + , Rect(..) + , Generator(..) + , deepSeqArray, deepSeqArrays + , singleton, toScalar + , extent, delay - -- * Basic Operations - , force - , deepSeqArray - - -- * Conversion + -- + , withManifest, withManifest' + + -- * Indexing + , (!), index + , (!?), safeIndex + , unsafeIndex + + -- * Construction + , fromFunction + , fromVector , fromList + + -- from Data.Array.Repa.Interlals.Forcing ------------------- + -- * Forcing + , force, force2 + , toVector , toList - -- * Index space transformations + -- from Data.Array.Repa.Operators.IndexSpace ---------------- + -- * Index space transformations , reshape - , append, (+:+) + , append, (++) , transpose - , replicate + , extend , slice , backpermute , backpermuteDft - -- * Structure preserving operations + -- from Data.Array.Repa.Operators.Mapping ------------------- + -- * Structure preserving operations , map , zipWith + , (+^), (-^), (*^), (/^) - -- * Reductions + -- from Data.Array.Repa.Operators.Reduction ----------------- + -- * Reductions , fold, foldAll , sum, sumAll - - -- * Generic traversal + + -- from Data.Array.Repa.Operators.Traverse ------------------ + -- * Generic Traversal , traverse , traverse2 , traverse3 , traverse4 - - -- * Interleaving + , unsafeTraverse + , unsafeTraverse2 + , unsafeTraverse3 + , unsafeTraverse4 + + -- from Data.Array.Repa.Operators.Interleave ---------------- + -- * Interleaving , interleave2 , interleave3 , interleave4 + + -- from Data.Array.Repa.Operators.Select -------------------- + -- * Selection + , select) - -- * Testing - , arbitrarySmallArray - , props_DataArrayRepa) where import Data.Array.Repa.Index import Data.Array.Repa.Slice import Data.Array.Repa.Shape -import Data.Array.Repa.QuickCheck +import Data.Array.Repa.Internals.Elt +import Data.Array.Repa.Internals.Base +import Data.Array.Repa.Internals.Forcing +import Data.Array.Repa.Operators.Traverse +import Data.Array.Repa.Operators.IndexSpace +import Data.Array.Repa.Operators.Interleave +import Data.Array.Repa.Operators.Mapping +import Data.Array.Repa.Operators.Reduction +import Data.Array.Repa.Operators.Select import qualified Data.Array.Repa.Shape as S -import "dph-prim-par" Data.Array.Parallel.Unlifted (Elt) -import qualified "dph-prim-par" Data.Array.Parallel.Unlifted as U -import qualified "dph-prim-seq" Data.Array.Parallel.Unlifted.Sequential as USeq - -import Test.QuickCheck -import Prelude hiding (sum, map, zipWith, replicate) +import Prelude hiding (sum, map, zipWith, (++)) import qualified Prelude as P stage = "Data.Array.Repa" - --- | Possibly delayed arrays. -data Array sh a - = -- | An array represented as some concrete unboxed data. - Manifest sh (U.Array a) - -- | An array represented as a function that computes each element. - | Delayed sh (sh -> a) --- Constructors ---------------------------------------------------------------------------------- - --- | Create a `Manifest` array from an unboxed `U.Array`. --- The elements are in row-major order. -fromUArray - :: Shape sh - => sh - -> U.Array a - -> Array sh a - -{-# INLINE fromUArray #-} -fromUArray sh uarr - = sh `S.deepSeq` - uarr `seq` - Manifest sh uarr - - --- | Create a `Delayed` array from a function. -fromFunction - :: Shape sh - => sh - -> (sh -> a) - -> Array sh a - -{-# INLINE fromFunction #-} -fromFunction sh fnElems - = sh `S.deepSeq` Delayed sh fnElems - - --- | Wrap a scalar into a singleton array. -unit :: Elt a => a -> Array Z a -{-# INLINE unit #-} -unit = Delayed Z . const - - --- Projections ------------------------------------------------------------------------------------ - --- | Take the extent of an array. -extent :: Array sh a -> sh -{-# INLINE extent #-} -extent arr - = case arr of - Manifest sh _ -> sh - Delayed sh _ -> sh - --- | Unpack an array into delayed form. -delay :: (Shape sh, Elt a) - => Array sh a - -> (sh, sh -> a) - -{-# INLINE delay #-} -delay arr - = case arr of - Manifest sh uarr -> (sh, \i -> uarr U.!: S.toIndex sh i) - Delayed sh fn -> (sh, fn) - - --- | Convert an array to an unboxed `U.Array`, forcing it if required. --- The elements come out in row-major order. -toUArray - :: (Shape sh, Elt a) - => Array sh a - -> U.Array a -{-# INLINE toUArray #-} -toUArray arr - = case force arr of - Manifest _ uarr -> uarr - _ -> error $ stage ++ ".toList: force failed" - - --- | Get an indexed element from an array. --- --- OBLIGATION: The index must be within the array. --- --- @inRange zeroDim (shape arr) ix == True@ --- -index, (!:) - :: forall sh a - . (Shape sh, Elt a) - => Array sh a - -> sh - -> a - -{-# INLINE index #-} -index arr ix - = case arr of - Delayed _ fn -> fn ix - Manifest sh uarr -> uarr U.!: (S.toIndex sh ix) - -{-# INLINE (!:) #-} -(!:) arr ix = index arr ix - - --- | Take the scalar value from a singleton array. -toScalar :: Elt a => Array Z a -> a -{-# INLINE toScalar #-} -toScalar arr - = case arr of - Delayed _ fn -> fn Z - Manifest _ uarr -> uarr U.!: 0 - - --- Basic Operations ------------------------------------------------------------------------------- - --- | Force an array, so that it becomes `Manifest`. -force :: (Shape sh, Elt a) - => Array sh a -> Array sh a - -{-# INLINE force #-} -force arr - = case arr of - Manifest sh uarr - -> sh `S.deepSeq` uarr `seq` - Manifest sh uarr - - Delayed sh fn - -> let uarr = U.map (fn . S.fromIndex sh) - $! U.enumFromTo (0 :: Int) (S.size sh - 1) - in sh `S.deepSeq` uarr `seq` - Manifest sh uarr - - --- | Ensure an array's structure is fully evaluated. --- This evaluates the extent and outer constructor, but does not `force` the elements. -infixr 0 `deepSeqArray` -deepSeqArray - :: Shape sh - => Array sh a - -> b -> b - -{-# INLINE deepSeqArray #-} -deepSeqArray arr x - = case arr of - Delayed sh _ -> sh `S.deepSeq` x - Manifest sh uarr -> sh `S.deepSeq` uarr `seq` x - - --- Conversion ------------------------------------------------------------------------------------- --- | Convert a list to an array. --- The length of the list must be exactly the `size` of the extent given, else `error`. --- -fromList - :: (Shape sh, Elt a) - => sh - -> [a] - -> Array sh a - -{-# INLINE fromList #-} -fromList sh xx - | U.length uarr /= S.size sh - = error $ unlines - [ stage ++ ".fromList: size of array shape does not match size of list" - , " size of shape = " ++ (show $ S.size sh) ++ "\n" - , " size of list = " ++ (show $ U.length uarr) ++ "\n" ] - - | otherwise - = Manifest sh uarr - - where uarr = U.fromList xx - - --- | Convert an array to a list. -toList :: (Shape sh, Elt a) - => Array sh a - -> [a] - -{-# INLINE toList #-} -toList arr - = case force arr of - Manifest _ uarr -> U.toList uarr - _ -> error $ stage ++ ".toList: force failed" - - -- Instances -------------------------------------------------------------------------------------- -- Show instance (Shape sh, Elt a, Show a) => Show (Array sh a) where show arr = show $ toList arr + -- Eq instance (Shape sh, Elt a, Eq a) => Eq (Array sh a) where {-# INLINE (==) #-} (==) arr1 arr2 - = toScalar - $ fold (&&) True - $ (flip reshape) (Z :. (S.size $ extent arr1)) + = foldAll (&&) True + $ reshape (Z :. (S.size $ extent arr1)) $ zipWith (==) arr1 arr2 {-# INLINE (/=) #-} @@ -292,496 +128,55 @@ -- Num -- All operators apply elementwise. instance (Shape sh, Elt a, Num a) => Num (Array sh a) where + {-# INLINE (+) #-} (+) = zipWith (+) + + {-# INLINE (-) #-} (-) = zipWith (-) + + {-# INLINE (*) #-} (*) = zipWith (*) - negate = map negate - abs = map abs - signum = map signum - fromInteger n = Delayed failShape (\_ -> fromInteger n) - where failShape = error $ stage ++ ".fromInteger: Constructed array has no shape." + {-# INLINE negate #-} + negate = map negate + {-# INLINE abs #-} + abs = map abs --- Index space transformations -------------------------------------------------------------------- --- | Impose a new shape on the elements of an array. --- The new extent must be the same size as the original, else `error`. --- -reshape :: (Shape sh, Shape sh', Elt a) - => Array sh a - -> sh' - -> Array sh' a + {-# INLINE signum #-} + signum = map signum -{-# INLINE reshape #-} -reshape arr newExtent - | not $ S.size newExtent == S.size (extent arr) - = error $ stage ++ ".reshape: reshaped array will not match size of the original" - - | otherwise - = Delayed newExtent - $ ((arr !:) . (S.fromIndex (extent arr)) . (S.toIndex newExtent)) + {-# INLINE fromInteger #-} + fromInteger n = fromFunction failShape (\_ -> fromInteger n) + where failShape = error $ stage P.++ ".fromInteger: Constructed array has no shape." --- | Append two arrays. --- --- OBLIGATION: The higher dimensions of both arrays must have the same extent. --- --- @tail (listOfShape (shape arr1)) == tail (listOfShape (shape arr2))@ --- -append, (+:+) +-- | Force an array before passing it to a function. +withManifest :: (Shape sh, Elt a) - => Array (sh :. Int) a - -> Array (sh :. Int) a - -> Array (sh :. Int) a - -{-# INLINE append #-} -append arr1 arr2 - = traverse2 arr1 arr2 fnExtent fnElem - where - (_ :. n) = extent arr1 - - fnExtent (sh :. i) (_ :. j) - = sh :. (i + j) - - fnElem f1 f2 (sh :. i) - | i < n = f1 (sh :. i) - | otherwise = f2 (sh :. (i - n)) - -{-# INLINE (+:+) #-} -(+:+) arr1 arr2 = append arr1 arr2 - - --- | Transpose the lowest two dimensions of an array. --- Transposing an array twice yields the original. -transpose - :: (Shape sh, Elt a) - => Array (sh :. Int :. Int) a - -> Array (sh :. Int :. Int) a - -{-# INLINE transpose #-} -transpose arr - = traverse arr - (\(sh :. m :. n) -> (sh :. n :.m)) - (\f -> \(sh :. i :. j) -> f (sh :. j :. i)) - - --- | Replicate an array, according to a given slice specification. -replicate - :: ( Slice sl - , Shape (FullShape sl) - , Shape (SliceShape sl) - , Elt e) - => sl - -> Array (SliceShape sl) e - -> Array (FullShape sl) e - -{-# INLINE replicate #-} -replicate sl arr - = backpermute - (fullOfSlice sl (extent arr)) - (sliceOfFull sl) - arr - --- | Take a slice from an array, according to a given specification. -slice :: ( Slice sl - , Shape (FullShape sl) - , Shape (SliceShape sl) - , Elt e) - => Array (FullShape sl) e - -> sl - -> Array (SliceShape sl) e - -{-# INLINE slice #-} -slice arr sl - = backpermute - (sliceOfFull sl (extent arr)) - (fullOfSlice sl) - arr - - --- | Backwards permutation of an array's elements. --- The result array has the same extent as the original. -backpermute - :: forall sh sh' a - . (Shape sh, Shape sh', Elt a) - => sh' -- ^ Extent of result array. - -> (sh' -> sh) -- ^ Function mapping each index in the result array - -- to an index of the source array. - -> Array sh a -- ^ Source array. - -> Array sh' a + => (Array sh a -> b) -> Array sh a -> b -{-# INLINE backpermute #-} -backpermute newExtent perm arr - = traverse arr (const newExtent) (. perm) +{-# INLINE withManifest #-} +withManifest f arr + = case arr of + Array sh [Region RangeAll (GenManifest vec)] + -> vec `seq` f (Array sh [Region RangeAll (GenManifest vec)]) - --- | Default backwards permutation of an array's elements. --- If the function returns `Nothing` then the value at that index is taken --- from the default array (@arrDft@) -backpermuteDft - :: forall sh sh' a - . (Shape sh, Shape sh', Elt a) - => Array sh' a -- ^ Default values (@arrDft@) - -> (sh' -> Maybe sh) -- ^ Function mapping each index in the result array - -- to an index in the source array. - -> Array sh a -- ^ Source array. - -> Array sh' a - -{-# INLINE backpermuteDft #-} -backpermuteDft arrDft fnIndex arrSrc - = Delayed (extent arrDft) fnElem - where fnElem ix - = case fnIndex ix of - Just ix' -> arrSrc !: ix' - Nothing -> arrDft !: ix - - --- Structure Preserving Operations ---------------------------------------------------------------- --- | Apply a worker function to each element of an array, --- yielding a new array with the same extent. -map :: (Shape sh, Elt a, Elt b) - => (a -> b) - -> Array sh a - -> Array sh b - -{-# INLINE map #-} -map f arr - = Delayed (extent arr) (f . (arr !:)) - - --- | Combine two arrays, element-wise, with a binary operator. --- If the extent of the two array arguments differ, --- then the resulting array's extent is their intersection. -zipWith :: (Shape sh, Elt a, Elt b, Elt c) - => (a -> b -> c) - -> Array sh a - -> Array sh b - -> Array sh c - -{-# INLINE zipWith #-} -zipWith f arr1 arr2 - = arr1 `deepSeqArray` - arr2 `deepSeqArray` - Delayed (S.intersectDim (extent arr1) (extent arr2)) - (\ix -> f (arr1 !: ix) (arr2 !: ix)) - - --- Reductions ------------------------------------------------------------------------------------- - --- IMPORTANT: --- These reductions use the sequential version of foldU, mapU and enumFromToU. --- If we use parallel versions then we'll end up with nested parallelism --- and the gang will abort at runtime. - --- | Fold the innermost dimension of an array. --- Combine this with `transpose` to fold any other dimension. -fold :: (Shape sh, Elt a) - => (a -> a -> a) - -> a - -> Array (sh :. Int) a - -> Array sh a - -{-# INLINE fold #-} -fold f x arr - = x `seq` arr `deepSeqArray` - let sh' :. n = extent arr - elemFn i = USeq.foldU f x - $ USeq.mapU - (\ix -> arr !: (i :. ix)) - (USeq.enumFromToU 0 (n - 1)) - in Delayed sh' elemFn - - --- | Fold all the elements of an array. -foldAll :: (Shape sh, Elt a) - => (a -> a -> a) - -> a - -> Array sh a - -> a + _ -> f (force arr) -{-# INLINE foldAll #-} -foldAll f x arr - = USeq.foldU f x - $ USeq.mapU ((arr !:) . (S.fromIndex (extent arr))) - $ USeq.enumFromToU - 0 - ((S.size $ extent arr) - 1) - - --- | Sum the innermost dimension of an array. -sum :: (Shape sh, Elt a, Num a) - => Array (sh :. Int) a - -> Array sh a - -{-# INLINE sum #-} -sum arr = fold (+) 0 arr - - --- | Sum all the elements of an array. -sumAll :: (Shape sh, Elt a, Num a) - => Array sh a - -> a - -{-# INLINE sumAll #-} -sumAll arr - = USeq.foldU (+) 0 - $ USeq.mapU ((arr !:) . (S.fromIndex (extent arr))) - $ USeq.enumFromToU - 0 - ((S.size $ extent arr) - 1) - - --- Generic Traversal ----------------------------------------------------------------------------- --- | Unstructured traversal. -traverse - :: forall sh sh' a b - . (Shape sh, Shape sh', Elt a) - => Array sh a -- ^ Source array. - -> (sh -> sh') -- ^ Function to produce the extent of the result. - -> ((sh -> a) -> sh' -> b) -- ^ Function to produce elements of the result. - -- It is passed a lookup function to get elements of the source. - -> Array sh' b - -{-# INLINE traverse #-} -traverse arr transExtent newElem - = arr `deepSeqArray` - Delayed - (transExtent (extent arr)) - (newElem (arr !:)) - - --- | Unstructured traversal over two arrays at once. -traverse2 - :: forall sh sh' sh'' a b c - . ( Shape sh, Shape sh', Shape sh'' - , Elt a, Elt b, Elt c) - => Array sh a -- ^ First source array. - -> Array sh' b -- ^ Second source array. - -> (sh -> sh' -> sh'') -- ^ Function to produce the extent of the result. - -> ((sh -> a) -> (sh' -> b) - -> (sh'' -> c)) -- ^ Function to produce elements of the result. - -- It is passed lookup functions to get elements of the - -- source arrays. - -> Array sh'' c - -{-# INLINE traverse2 #-} -traverse2 arrA arrB transExtent newElem - = arrA `deepSeqArray` arrB `deepSeqArray` - Delayed - (transExtent (extent arrA) (extent arrB)) - (newElem ((!:) arrA) ((!:) arrB)) - - --- | Unstructured traversal over three arrays at once. -traverse3 - :: forall sh1 sh2 sh3 sh4 - a b c d - . ( Shape sh1, Shape sh2, Shape sh3, Shape sh4 - , Elt a, Elt b, Elt c, Elt d) - => Array sh1 a - -> Array sh2 b - -> Array sh3 c - -> (sh1 -> sh2 -> sh3 -> sh4) - -> ( (sh1 -> a) -> (sh2 -> b) - -> (sh3 -> c) - -> sh4 -> d ) - -> Array sh4 d - -{-# INLINE traverse3 #-} -traverse3 arrA arrB arrC transExtent newElem - = arrA `deepSeqArray` arrB `deepSeqArray` arrC `deepSeqArray` - Delayed - (transExtent (extent arrA) (extent arrB) (extent arrC)) - (newElem (arrA !:) (arrB !:) (arrC !:)) - - --- | Unstructured traversal over four arrays at once. -traverse4 - :: forall sh1 sh2 sh3 sh4 sh5 - a b c d e - . ( Shape sh1, Shape sh2, Shape sh3, Shape sh4, Shape sh5 - , Elt a, Elt b, Elt c, Elt d, Elt e) - => Array sh1 a - -> Array sh2 b - -> Array sh3 c - -> Array sh4 d - -> (sh1 -> sh2 -> sh3 -> sh4 -> sh5 ) - -> ( (sh1 -> a) -> (sh2 -> b) - -> (sh3 -> c) -> (sh4 -> d) - -> sh5 -> e ) - -> Array sh5 e - -{-# INLINE traverse4 #-} -traverse4 arrA arrB arrC arrD transExtent newElem - = arrA `deepSeqArray` arrB `deepSeqArray` arrC `deepSeqArray` arrD `deepSeqArray` - Delayed - (transExtent (extent arrA) (extent arrB) (extent arrC) (extent arrD)) - (newElem (arrA !:) (arrB !:) (arrC !:) (arrD !:)) - - --- Interleaving ----------------------------------------------------------------------------------- --- | Interleave the elments of two arrays. --- All the input arrays must have the same extent, else `error`. --- The lowest dimenion of the result array is twice the size of the inputs. --- --- @ --- interleave2 a1 a2 b1 b2 => a1 b1 a2 b2 --- a3 a4 b3 b4 a3 b3 a4 b4 --- @ --- -interleave2 +-- | Force an array before passing it to a function. +withManifest' :: (Shape sh, Elt a) - => Array (sh :. Int) a - -> Array (sh :. Int) a - -> Array (sh :. Int) a - -{-# INLINE interleave2 #-} -interleave2 arr1 arr2 - = arr1 `deepSeqArray` arr2 `deepSeqArray` - traverse2 arr1 arr2 shapeFn elemFn - where - shapeFn dim1 dim2 - | dim1 == dim2 - , sh :. len <- dim1 - = sh :. (len * 2) - - | otherwise - = error "Data.Array.Repa.interleave2: arrays must have same extent" - - elemFn get1 get2 (sh :. ix) - = case ix `mod` 3 of - 0 -> get1 (sh :. ix `div` 2) - 1 -> get2 (sh :. ix `div` 2) - _ -> error "Data.Array.Repa.interleave2: this never happens :-P" - + => Array sh a -> (Array sh a -> b) -> b --- | Interleave the elments of three arrays. -interleave3 - :: (Shape sh, Elt a) - => Array (sh :. Int) a - -> Array (sh :. Int) a - -> Array (sh :. Int) a - -> Array (sh :. Int) a - -{-# INLINE interleave3 #-} -interleave3 arr1 arr2 arr3 - = arr1 `deepSeqArray` arr2 `deepSeqArray` arr3 `deepSeqArray` - traverse3 arr1 arr2 arr3 shapeFn elemFn - where - shapeFn dim1 dim2 dim3 - | dim1 == dim2 - , dim1 == dim3 - , sh :. len <- dim1 - = sh :. (len * 3) +{-# INLINE withManifest' #-} +withManifest' arr f + = case arr of + Array sh [Region RangeAll (GenManifest vec)] + -> vec `seq` f (Array sh [Region RangeAll (GenManifest vec)]) - | otherwise - = error "Data.Array.Repa.interleave3: arrays must have same extent" - - elemFn get1 get2 get3 (sh :. ix) - = case ix `mod` 3 of - 0 -> get1 (sh :. ix `div` 3) - 1 -> get2 (sh :. ix `div` 3) - 2 -> get3 (sh :. ix `div` 3) - _ -> error "Data.Array.Repa.interleave3: this never happens :-P" - + _ -> f (force arr) --- | Interleave the elments of four arrays. -interleave4 - :: (Shape sh, Elt a) - => Array (sh :. Int) a - -> Array (sh :. Int) a - -> Array (sh :. Int) a - -> Array (sh :. Int) a - -> Array (sh :. Int) a -{-# INLINE interleave4 #-} -interleave4 arr1 arr2 arr3 arr4 - = arr1 `deepSeqArray` arr2 `deepSeqArray` arr3 `deepSeqArray` arr4 `deepSeqArray` - traverse4 arr1 arr2 arr3 arr4 shapeFn elemFn - where - shapeFn dim1 dim2 dim3 dim4 - | dim1 == dim2 - , dim1 == dim3 - , dim1 == dim4 - , sh :. len <- dim1 - = sh :. (len * 4) - - | otherwise - = error "Data.Array.Repa.interleave4: arrays must have same extent" - - elemFn get1 get2 get3 get4 (sh :. ix) - = case ix `mod` 4 of - 0 -> get1 (sh :. ix `div` 4) - 1 -> get2 (sh :. ix `div` 4) - 2 -> get3 (sh :. ix `div` 4) - 3 -> get4 (sh :. ix `div` 4) - _ -> error "Data.Array.Repa.interleave4: this never happens :-P" - --- Arbitrary -------------------------------------------------------------------------------------- --- | Create an arbitrary small array, restricting the size of each of the dimensions to some value. -arbitrarySmallArray - :: (Shape sh, Elt a, Arbitrary sh, Arbitrary a) - => Int - -> Gen (Array (sh :. Int) a) - -arbitrarySmallArray maxDim - = do sh <- arbitrarySmallShape maxDim - xx <- arbitraryListOfLength (S.size sh) - return $ fromList sh xx - - - --- Properties ------------------------------------------------------------------------------------- - --- | QuickCheck properties for this module and its children. -props_DataArrayRepa :: [(String, Property)] -props_DataArrayRepa - = props_DataArrayRepaIndex - ++ [(stage ++ "." ++ name, test) | (name, test) - <- [ ("id_force/DIM5", property prop_id_force_DIM5) - , ("id_toScalarUnit", property prop_id_toScalarUnit) - , ("id_toListFromList/DIM3", property prop_id_toListFromList_DIM3) - , ("id_transpose/DIM4", property prop_id_transpose_DIM4) - , ("reshapeTransposeSize/DIM3", property prop_reshapeTranspose_DIM3) - , ("appendIsAppend/DIM3", property prop_appendIsAppend_DIM3) - , ("sumAllIsSum/DIM3", property prop_sumAllIsSum_DIM3) ]] - - --- The Eq instance uses fold and zipWith. -prop_id_force_DIM5 - = forAll (arbitrarySmallArray 10) $ \(arr :: Array DIM5 Int) -> - arr == force arr - -prop_id_toScalarUnit (x :: Int) - = toScalar (unit x) == x - --- Conversions ------------------------ -prop_id_toListFromList_DIM3 - = forAll (arbitrarySmallShape 10) $ \(sh :: DIM3) -> - forAll (arbitraryListOfLength (S.size sh)) $ \(xx :: [Int]) -> - toList (fromList sh xx) == xx - --- Index Space Transforms ------------- -prop_id_transpose_DIM4 - = forAll (arbitrarySmallArray 20) $ \(arr :: Array DIM3 Int) -> - transpose (transpose arr) == arr - --- A reshaped array has the same size and sum as the original -prop_reshapeTranspose_DIM3 - = forAll (arbitrarySmallArray 20) $ \(arr :: Array DIM3 Int) -> - let arr' = transpose arr - sh' = extent arr' - in (S.size $ extent arr) == S.size (extent (reshape arr sh')) - && (sumAll arr == sumAll arr') - -prop_appendIsAppend_DIM3 - = forAll (arbitrarySmallArray 20) $ \(arr1 :: Array DIM3 Int) -> - sumAll (append arr1 arr1) == (2 * sumAll arr1) - --- Reductions -------------------------- -prop_sumAllIsSum_DIM3 - = forAll (arbitrarySmallShape 100) $ \(sh :: DIM2) -> - forAll (arbitraryListOfLength (S.size sh)) $ \(xx :: [Int]) -> - sumAll (fromList sh xx) == P.sum xx diff --git a/Data/Array/Repa/Arbitrary.hs b/Data/Array/Repa/Arbitrary.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Repa/Arbitrary.hs @@ -0,0 +1,99 @@ +{-# LANGUAGE TypeOperators, FlexibleInstances #-} +{-# OPTIONS -fno-warn-orphans #-} + +-- Utils to help with testing. Not exported. +module Data.Array.Repa.Arbitrary + ( arbitraryShape + , arbitrarySmallShape + , arbitraryListOfLength + , arbitrarySmallArray) +where +import Data.Array.Repa.Internals.Elt +import Data.Array.Repa.Internals.Base +import Data.Array.Repa.Index +import Data.Array.Repa.Shape as S +import Control.Monad +import Test.QuickCheck + + +-- Arbitrary -------------------------------------------------------------------------------------- +instance Arbitrary Z where + arbitrary = return Z + +-- | Generate an arbitrary index, which may have 0's for some components. +instance (Shape sh, Arbitrary sh) => Arbitrary (sh :. Int) where + arbitrary + = do sh1 <- arbitrary + let sh1Unit = if size sh1 == 0 then unitDim else sh1 + + -- Make sure not to create an index so big that we get + -- integer overflow when converting it to the linear form. + n <- liftM abs $ arbitrary + let nMax = maxBound `div` (size sh1Unit) + let nMaxed = n `mod` nMax + + return $ sh1 :. nMaxed + +-- | Generate an aribrary shape that does not have 0's for any component. +arbitraryShape + :: (Shape sh, Arbitrary sh) + => Gen (sh :. Int) + +arbitraryShape + = do sh1 <- arbitrary + let sh1Unit = if size sh1 == 0 then unitDim else sh1 + + -- Make sure not to create an index so big that we get + -- integer overflow when converting it to the linear form. + n <- liftM abs $ arbitrary + let nMax = maxBound `div` size sh1Unit + let nMaxed = n `mod` nMax + let nClamped = if nMaxed == 0 then 1 else nMaxed + + return $ sh1Unit :. nClamped + + +-- | Generate an arbitrary shape where each dimension is more than zero, +-- but less than a specific value. +arbitrarySmallShape + :: (Shape sh, Arbitrary sh) + => Int + -> Gen (sh :. Int) + +arbitrarySmallShape maxDim + = do sh <- arbitraryShape + let dims = listOfShape sh + + let clamp x + = case x `mod` maxDim of + 0 -> 1 + n -> n + + return $ if True + then shapeOfList $ map clamp dims + else sh + + +arbitraryListOfLength + :: Arbitrary a + => Int -> Gen [a] + +arbitraryListOfLength n + | n == 0 = return [] + | otherwise + = do i <- arbitrary + rest <- arbitraryListOfLength (n - 1) + return $ i : rest + +-- | Create an arbitrary small array, restricting the size of each of the +-- dimensions to some value. +arbitrarySmallArray + :: (Shape sh, Elt a, Arbitrary sh, Arbitrary a) + => Int + -> Gen (Array (sh :. Int) a) + +arbitrarySmallArray maxDim + = do sh <- arbitrarySmallShape maxDim + xx <- arbitraryListOfLength (S.size sh) + return $ fromList sh xx + diff --git a/Data/Array/Repa/Index.hs b/Data/Array/Repa/Index.hs --- a/Data/Array/Repa/Index.hs +++ b/Data/Array/Repa/Index.hs @@ -13,16 +13,9 @@ , DIM2 , DIM3 , DIM4 - , DIM5 - - -- * Testing - , arbitraryShape - , arbitrarySmallShape - , props_DataArrayRepaIndex) + , DIM5) where import Data.Array.Repa.Shape -import Test.QuickCheck -import Control.Monad import GHC.Base (quotInt, remInt) stage = "Data.Array.Repa.Index" @@ -48,30 +41,50 @@ -- Shape ------------------------------------------------------------------------------------------ instance Shape Z where - dim _ = 0 + {-# INLINE rank #-} + rank _ = 0 + + {-# INLINE zeroDim #-} zeroDim = Z + + {-# INLINE unitDim #-} unitDim = Z + + {-# INLINE intersectDim #-} intersectDim _ _ = Z + {-# INLINE addDim #-} + addDim _ _ = Z + + {-# INLINE size #-} size _ = 1 + + {-# INLINE sizeIsValid #-} sizeIsValid _ = True + + {-# INLINE toIndex #-} toIndex _ _ = 0 + + {-# INLINE fromIndex #-} fromIndex _ _ = Z - inRange Z Z Z = True + {-# INLINE inShapeRange #-} + inShapeRange Z Z Z = True + listOfShape _ = [] shapeOfList [] = Z shapeOfList _ = error $ stage ++ ".fromList: non-empty list when converting to Z." + {-# INLINE deepSeq #-} deepSeq Z x = x instance Shape sh => Shape (sh :. Int) where - {-# INLINE dim #-} - dim (sh :. _) - = dim sh + 1 + {-# INLINE rank #-} + rank (sh :. _) + = rank sh + 1 {-# INLINE zeroDim #-} zeroDim = zeroDim :. 0 @@ -83,6 +96,10 @@ intersectDim (sh1 :. n1) (sh2 :. n2) = (intersectDim sh1 sh2 :. (min n1 n2)) + {-# INLINE addDim #-} + addDim (sh1 :. n1) (sh2 :. n2) + = addDim sh1 sh2 :. (n1 + n2) + {-# INLINE size #-} size (sh1 :. n) = size sh1 * n @@ -107,12 +124,12 @@ -- in computing the remainder for the highest dimension since -- n < d must hold. This saves one remInt per element access which -- is quite a big deal. - r | dim ds == 0 = n + r | rank ds == 0 = n | otherwise = n `remInt` d - {-# INLINE inRange #-} - inRange (zs :. z) (sh1 :. n1) (sh2 :. n2) - = (n2 >= z) && (n2 < n1) && (inRange zs sh1 sh2) + {-# INLINE inShapeRange #-} + inShapeRange (zs :. z) (sh1 :. n1) (sh2 :. n2) + = (n2 >= z) && (n2 < n1) && (inShapeRange zs sh1 sh2) listOfShape (sh :. n) @@ -129,90 +146,6 @@ --- Arbitrary -------------------------------------------------------------------------------------- -instance Arbitrary Z where - arbitrary = return Z - --- | Generate an arbitrary index, which may have 0's for some components. -instance (Shape sh, Arbitrary sh) => Arbitrary (sh :. Int) where - arbitrary - = do sh1 <- arbitrary - let sh1Unit = if size sh1 == 0 then unitDim else sh1 - - -- Make sure not to create an index so big that we get - -- integer overflow when converting it to the linear form. - n <- liftM abs $ arbitrary - let nMax = maxBound `div` (size sh1Unit) - let nMaxed = n `mod` nMax - - return $ sh1 :. nMaxed - --- | Generate an aribrary shape that does not have 0's for any component. -arbitraryShape - :: (Shape sh, Arbitrary sh) - => Gen (sh :. Int) - -arbitraryShape - = do sh1 <- arbitrary - let sh1Unit = if size sh1 == 0 then unitDim else sh1 - - -- Make sure not to create an index so big that we get - -- integer overflow when converting it to the linear form. - n <- liftM abs $ arbitrary - let nMax = maxBound `div` size sh1Unit - let nMaxed = n `mod` nMax - let nClamped = if nMaxed == 0 then 1 else nMaxed - - return $ sh1Unit :. nClamped - - --- | Generate an arbitrary shape where each dimension is more than zero, --- but less than a specific value. -arbitrarySmallShape - :: (Shape sh, Arbitrary sh) - => Int - -> Gen (sh :. Int) - -arbitrarySmallShape maxDim - = do sh <- arbitraryShape - let dims = listOfShape sh - - let clamp x - = case x `mod` maxDim of - 0 -> 1 - n -> n - - return $ if True - then shapeOfList $ map clamp dims - else sh - - -genInShape2 :: DIM2 -> Gen DIM2 -genInShape2 (Z :. yMax :. xMax) - = do y <- liftM (`mod` yMax) $ arbitrary - x <- liftM (`mod` xMax) $ arbitrary - return $ Z :. y :. x - - --- Properties ------------------------------------------------------------------------------------- --- | QuickCheck properties for this module. -props_DataArrayRepaIndex :: [(String, Property)] -props_DataArrayRepaIndex - = [(stage ++ "." ++ name, test) | (name, test) - <- [ ("toIndexFromIndex/DIM1", property prop_toIndexFromIndex_DIM1) - , ("toIndexFromIndex/DIM2", property prop_toIndexFromIndex_DIM2) ]] - -prop_toIndexFromIndex_DIM1 sh ix - = (sizeIsValid sh) - ==> (inShape sh ix) - ==> fromIndex sh (toIndex sh ix) == ix - where _types = ( sh :: DIM1 - , ix :: DIM1) - -prop_toIndexFromIndex_DIM2 - = forAll arbitraryShape $ \(sh :: DIM2) -> - forAll (genInShape2 sh) $ \(ix :: DIM2) -> - fromIndex sh (toIndex sh ix) == ix diff --git a/Data/Array/Repa/Internals/Base.hs b/Data/Array/Repa/Internals/Base.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Repa/Internals/Base.hs @@ -0,0 +1,389 @@ +{-# LANGUAGE ExplicitForAll, TypeOperators, FlexibleInstances, UndecidableInstances, BangPatterns, + ExistentialQuantification #-} +module Data.Array.Repa.Internals.Base + ( Array (..) + , Region(..) + , Range (..) + , Rect (..) + , Generator(..) + , deepSeqArray, deepSeqArrays + + , singleton, toScalar + , extent, delay + + -- * Predicates + , inRange + + -- * Indexing + , (!), index + , (!?), safeIndex + , unsafeIndex + + -- * Construction + , fromFunction + , fromVector + , fromList) +where +import Data.Array.Repa.Index +import Data.Array.Repa.Internals.Elt +import Data.Array.Repa.Shape as S +import qualified Data.Vector.Unboxed as V +import Data.Vector.Unboxed (Vector) + +stage = "Data.Array.Repa.Array" + +-- Array ---------------------------------------------------------------------- +-- | Repa arrays. +data Array sh a + = Array + { -- | The entire extent of the array. + arrayExtent :: sh + + -- | Arrays can be partitioned into several regions. + , arrayRegions :: [Region sh a] } + + +-- | Defines the values in a region of the array. +data Region sh a + = Region + { -- | The range of elements this region applies to. + regionRange :: Range sh + + -- | How to compute the array elements in this region. + , regionGenerator :: Generator sh a } + + +-- | Represents a range of elements in the array. +data Range sh + -- | Covers the entire array. + = RangeAll + + -- | The union of a possibly disjoint set of rectangles. + | RangeRects + { rangeMatch :: sh -> Bool + , rangeRects :: [Rect sh] } + + +-- | A rectangle\/cube of arbitrary dimension. +-- The indices are of the minimum and maximim elements to fill. +data Rect sh + = Rect sh sh + +-- | Generates array elements for a particular region in the array. +data Generator sh a + -- | Elements are already computed and sitting in this vector. + = GenManifest !(Vector a) + + -- | Elements can be computed using these cursor functions. + | forall cursor + . GenCursor + { -- | Make a cursor to a particular element. + genMakeCursor :: sh -> cursor + + -- | Shift the cursor by an offset, to get to another element. + , genShiftCursor :: sh -> cursor -> cursor + + -- | Load\/compute the element at the given cursor. + , genLoadElem :: cursor -> a } + + +-- DeepSeqs ------------------------------------------------------------------- + +-- | Ensure the structure for an array is fully evaluated. +-- As we are in a lazy language, applying the @force@ function to a delayed array doesn't +-- actually compute it at that point. Rather, Haskell builds a suspension representing the +-- appliction of the @force@ function to that array. Use @deepSeqArray@ to ensure the array +-- is actually computed at a particular point in the program. +infixr 0 `deepSeqArray` +deepSeqArray :: Shape sh => Array sh a -> b -> b +{-# INLINE deepSeqArray #-} +deepSeqArray (Array ex rgns) x + = ex `S.deepSeq` rgns `deepSeqRegions` x + +-- | Like `deepSeqArray` but seqs all the arrays in a list. +-- This is specialised up to lists of 4 arrays. Using more in the list will break fusion. +infixr 0 `deepSeqArrays` +deepSeqArrays :: Shape sh => [Array sh a] -> b -> b +{-# INLINE deepSeqArrays #-} +deepSeqArrays as y + = case as of + [] -> y + [a] -> a `deepSeqArray` y + [a1, a2] -> a1 `deepSeqArray` a2 `deepSeqArray` y + [a1, a2, a3] -> a1 `deepSeqArray` a2 `deepSeqArray` a3 `deepSeqArray` y + [a1, a2, a3, a4]-> a1 `deepSeqArray` a2 `deepSeqArray` a3 `deepSeqArray` a4 `deepSeqArray` y + _ -> deepSeqArrays' as y + +deepSeqArrays' as' y + = case as' of + [] -> y + x : xs -> x `deepSeqArray` xs `deepSeqArrays` y + +-- | Ensure the structure for a region is fully evaluated. +infixr 0 `deepSeqRegion` +deepSeqRegion :: Shape sh => Region sh a -> b -> b +{-# INLINE deepSeqRegion #-} +deepSeqRegion (Region range gen) x + = range `deepSeqRange` gen `deepSeqGen` x + + +-- | Ensure the structure for some regions are fully evaluated. +infixr 0 `deepSeqRegions` +deepSeqRegions :: Shape sh => [Region sh a] -> b -> b +{-# INLINE deepSeqRegions #-} +deepSeqRegions rs y + = case rs of + [] -> y + [r] -> r `deepSeqRegion` y + [r1, r2] -> r1 `deepSeqRegion` r2 `deepSeqRegion` y + rs' -> deepSeqRegions' rs' y + +deepSeqRegions' rs' y + = case rs' of + [] -> y + x : xs -> x `deepSeqRegion` xs `deepSeqRegions'` y + + +-- | Ensure a range is fully evaluated. +infixr 0 `deepSeqRange` +deepSeqRange :: Shape sh => Range sh -> b -> b +{-# INLINE deepSeqRange #-} +deepSeqRange range x + = case range of + RangeAll -> x + RangeRects f rects -> f `seq` rects `seq` x + + +-- | Ensure a Generator's structure is fully evaluated. +infixr 0 `deepSeqGen` +deepSeqGen :: Shape sh => Generator sh a -> b -> b +{-# INLINE deepSeqGen #-} +deepSeqGen gen x + = case gen of + GenManifest vec -> vec `seq` x + GenCursor{} -> x + + +-- Predicates ------------------------------------------------------------------------------------- +inRange :: Shape sh => Range sh -> sh -> Bool +{-# INLINE inRange #-} +inRange RangeAll _ = True +inRange (RangeRects fn _) ix = fn ix + + +-- Singletons ------------------------------------------------------------------------------------- +-- | Wrap a scalar into a singleton array. +singleton :: Elt a => a -> Array Z a +{-# INLINE singleton #-} +singleton = fromFunction Z . const + +-- | Take the scalar value from a singleton array. +toScalar :: Elt a => Array Z a -> a +{-# INLINE toScalar #-} +toScalar arr = arr ! Z + + +-- Projections ------------------------------------------------------------------------------------ +-- | Take the extent of an array. +extent :: Array sh a -> sh +{-# INLINE extent #-} +extent arr = arrayExtent arr + + +-- | Unpack an array into delayed form. +delay :: (Shape sh, Elt a) + => Array sh a + -> (sh, sh -> a) + +{-# INLINE delay #-} +delay arr@(Array sh _) + = (sh, (arr !)) + + +-- Indexing --------------------------------------------------------------------------------------- +-- | Get an indexed element from an array. +-- This uses the same level of bounds checking as your Data.Vector installation. +(!), index + :: forall sh a + . (Shape sh, Elt a) + => Array sh a + -> sh + -> a + +{-# INLINE (!) #-} +(!) arr ix = index arr ix + +{-# INLINE index #-} +index arr ix + = case arr of + Array _ [] + -> zero + + Array sh [Region _ gen1] + -> indexGen sh gen1 ix + + Array sh [Region r1 gen1, Region _ gen2] + | inRange r1 ix -> indexGen sh gen1 ix + | otherwise -> indexGen sh gen2 ix + + _ -> index' arr ix + + + where {-# INLINE indexGen #-} + indexGen sh gen ix' + = case gen of + GenManifest vec + -> vec V.! (S.toIndex sh ix') + + GenCursor makeCursor _ loadElem + -> loadElem $ makeCursor ix' + + index' (Array sh (Region range gen : rs)) ix' + | inRange range ix = indexGen sh gen ix' + | otherwise = index' (Array sh rs) ix' + + index' (Array _ []) _ + = zero + + + +-- | Get an indexed element from an array. +-- If the element is out of range then `Nothing`. +(!?), safeIndex + :: forall sh a + . (Shape sh, Elt a) + => Array sh a + -> sh + -> Maybe a + +{-# INLINE (!?) #-} +(!?) arr ix = safeIndex arr ix + + +{-# INLINE safeIndex #-} +safeIndex arr ix + = case arr of + Array _ [] + -> Nothing + + Array sh [Region _ gen1] + -> indexGen sh gen1 ix + + Array sh [Region r1 gen1, Region r2 gen2] + | inRange r1 ix -> indexGen sh gen1 ix + | inRange r2 ix -> indexGen sh gen2 ix + | otherwise -> Nothing + + _ -> index' arr ix + + + where {-# INLINE indexGen #-} + indexGen sh gen ix' + = case gen of + GenManifest vec + -> vec V.!? (S.toIndex sh ix') + + GenCursor makeCursor _ loadElem + -> Just (loadElem $ makeCursor ix') + + index' (Array sh (Region range gen : rs)) ix' + | inRange range ix = indexGen sh gen ix' + | otherwise = index' (Array sh rs) ix' + + index' (Array _ []) _ + = Nothing + + +-- | Get an indexed element from an array, without bounds checking. +-- This assumes that the regions in the array give full coverage. +-- An array with no regions gets zero for every element. +unsafeIndex + :: forall sh a + . (Shape sh, Elt a) + => Array sh a + -> sh + -> a + +{-# INLINE unsafeIndex #-} +unsafeIndex arr ix + = case arr of + Array _ [] + -> zero + + Array sh [Region _ gen1] + -> unsafeIndexGen sh gen1 ix + + Array sh [Region r1 gen1, Region _ gen2] + | inRange r1 ix -> unsafeIndexGen sh gen1 ix + | otherwise -> unsafeIndexGen sh gen2 ix + + _ -> unsafeIndex' arr ix + + where {-# INLINE unsafeIndexGen #-} + unsafeIndexGen sh gen ix' + = case gen of + GenManifest vec + -> vec `V.unsafeIndex` (S.toIndex sh ix') + + GenCursor makeCursor _ loadElem + -> loadElem $ makeCursor ix' + + unsafeIndex' (Array sh (Region range gen : rs)) ix' + | inRange range ix = unsafeIndexGen sh gen ix' + | otherwise = unsafeIndex' (Array sh rs) ix' + + unsafeIndex' (Array _ []) _ + = zero + + +-- Conversions ------------------------------------------------------------------------------------ +-- | Create a `Delayed` array from a function. +fromFunction + :: Shape sh + => sh + -> (sh -> a) + -> Array sh a + +{-# INLINE fromFunction #-} +fromFunction sh fnElems + = sh `S.deepSeq` + Array sh [Region + RangeAll + (GenCursor id addDim fnElems)] + +-- | Create a `Manifest` array from an unboxed `Vector`. +-- The elements are in row-major order. +fromVector + :: Shape sh + => sh + -> Vector a + -> Array sh a + +{-# INLINE fromVector #-} +fromVector sh vec + = sh `S.deepSeq` vec `seq` + Array sh [Region RangeAll (GenManifest vec)] + + +-- Conversion ------------------------------------------------------------------------------------- +-- | Convert a list to an array. +-- The length of the list must be exactly the `size` of the extent given, else `error`. +fromList + :: (Shape sh, Elt a) + => sh + -> [a] + -> Array sh a + +{-# INLINE fromList #-} +fromList sh xx + | V.length vec /= S.size sh + = error $ unlines + [ stage ++ ".fromList: size of array shape does not match size of list" + , " size of shape = " ++ (show $ S.size sh) ++ "\n" + , " size of list = " ++ (show $ V.length vec) ++ "\n" ] + + | otherwise + = Array sh [Region RangeAll (GenManifest vec)] + + where vec = V.fromList xx + diff --git a/Data/Array/Repa/Internals/Elt.hs b/Data/Array/Repa/Internals/Elt.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Repa/Internals/Elt.hs @@ -0,0 +1,163 @@ +-- | Values that can be stored in Repa Arrays. +{-# LANGUAGE MagicHash, UnboxedTuples, TypeSynonymInstances, FlexibleInstances #-} +module Data.Array.Repa.Internals.Elt + (Elt (..)) +where +import GHC.Prim +import GHC.Exts +import GHC.Types +import GHC.Word +import GHC.Int +import Data.Vector.Unboxed + + +-- Note that the touch# function is special because we can pass it boxed or unboxed +-- values. The argument type has kind ?, not just * or #. + +-- | Element types that can be stored in Repa arrays. +class (Show a, Unbox a) => Elt a where + + -- | We use this to prevent bindings from being floated inappropriatey. + -- Doing a `seq` sometimes isn't enough, because the GHC simpplifier can + -- erase these, and/or still move around the bindings. + touch :: a -> IO () + + -- | Generic zero value, helpful for debugging. + zero :: a + + -- | Generic one value, helpful for debugging. + one :: a + + +-- Bool ----------------------------------------------------------------------- +instance Elt Bool where + {-# INLINE touch #-} + touch b + = IO (\state -> case touch# b state of + state' -> (# state', () #)) + + {-# INLINE zero #-} + zero = False + + {-# INLINE one #-} + one = True + + +-- Tuple ---------------------------------------------------------------------- +instance (Elt a, Elt b) => Elt (a, b) where + {-# INLINE touch #-} + touch (a, b) + = do touch a + touch b + + {-# INLINE zero #-} + zero = (zero, zero) + + {-# INLINE one #-} + one = (one, one) + + +-- Floating ------------------------------------------------------------------- +instance Elt Float where + {-# INLINE touch #-} + touch (F# f) + = IO (\state -> case touch# f state of + state' -> (# state', () #)) + + {-# INLINE zero #-} + zero = 0 + + {-# INLINE one #-} + one = 1 + + +instance Elt Double where + {-# INLINE touch #-} + touch (D# d) + = IO (\state -> case touch# d state of + state' -> (# state', () #)) + + {-# INLINE zero #-} + zero = 0 + + {-# INLINE one #-} + one = 1 + + +-- Int ------------------------------------------------------------------------ +instance Elt Int where + {-# INLINE touch #-} + touch (I# i) + = IO (\state -> case touch# i state of + state' -> (# state', () #)) + + {-# INLINE zero #-} + zero = 0 + + {-# INLINE one #-} + one = 1 + +instance Elt Int8 where + {-# INLINE touch #-} + touch (I8# w) + = IO (\state -> case touch# w state of + state' -> (# state', () #)) + + {-# INLINE zero #-} + zero = 0 + + {-# INLINE one #-} + one = 1 + + +instance Elt Int16 where + {-# INLINE touch #-} + touch (I16# w) + = IO (\state -> case touch# w state of + state' -> (# state', () #)) + + {-# INLINE zero #-} + zero = 0 + + {-# INLINE one #-} + one = 1 + + +-- Word ----------------------------------------------------------------------- +instance Elt Word where + {-# INLINE touch #-} + touch (W# i) + = IO (\state -> case touch# i state of + state' -> (# state', () #)) + + {-# INLINE zero #-} + zero = 0 + + {-# INLINE one #-} + one = 1 + + +instance Elt Word8 where + {-# INLINE touch #-} + touch (W8# w) + = IO (\state -> case touch# w state of + state' -> (# state', () #)) + + {-# INLINE zero #-} + zero = 0 + + {-# INLINE one #-} + one = 1 + + +instance Elt Word16 where + {-# INLINE touch #-} + touch (W16# w) + = IO (\state -> case touch# w state of + state' -> (# state', () #)) + + {-# INLINE zero #-} + zero = 0 + + {-# INLINE one #-} + one = 1 diff --git a/Data/Array/Repa/Internals/EvalBlockwise.hs b/Data/Array/Repa/Internals/EvalBlockwise.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Repa/Internals/EvalBlockwise.hs @@ -0,0 +1,149 @@ +{-# LANGUAGE BangPatterns #-} +module Data.Array.Repa.Internals.EvalBlockwise + ( fillVectorBlockwiseP + , fillVectorBlock + , fillVectorBlockP) +where +import Data.Array.Repa.Internals.Elt +import Data.Array.Repa.Internals.Gang +import Data.Vector.Unboxed.Mutable as VM +import GHC.Base (remInt, quotInt) +import Prelude as P + + +-- Blockwise filling ------------------------------------------------------------------------------ +fillVectorBlockwiseP + :: Elt a + => IOVector a -- ^ vector to write elements into + -> (Int -> a) -- ^ fn to evaluate an element at the given index + -> Int -- ^ width of image. + -> IO () + +{-# INLINE [0] fillVectorBlockwiseP #-} +fillVectorBlockwiseP !vec !getElemFVBP !imageWidth + = gangIO theGang fillBlock + + where !threads = gangSize theGang + !vecLen = VM.length vec + !imageHeight = vecLen `div` imageWidth + !colChunkLen = imageWidth `quotInt` threads + !colChunkSlack = imageWidth `remInt` threads + + + {-# INLINE colIx #-} + colIx !ix + | ix < colChunkSlack = ix * (colChunkLen + 1) + | otherwise = ix * colChunkLen + colChunkSlack + + + -- just give one column to each thread + {-# INLINE fillBlock #-} + fillBlock :: Int -> IO () + fillBlock !ix + = let !x0 = colIx ix + !x1 = colIx (ix + 1) + !y0 = 0 + !y1 = imageHeight + in fillVectorBlock vec getElemFVBP imageWidth x0 y0 x1 y1 + + +-- Block filling ---------------------------------------------------------------------------------- +-- | Fill a block in a 2D image, in parallel. +-- Coordinates given are of the filled edges of the block. +-- We divide the block into columns, and give one column to each thread. +fillVectorBlockP + :: Elt a + => IOVector a -- ^ vector to write elements into + -> (Int -> a) -- ^ fn to evaluate an element at the given index. + -> Int -- ^ width of whole image + -> Int -- ^ x0 lower left corner of block to fill + -> Int -- ^ y0 (low x and y value) + -> Int -- ^ x1 upper right corner of block + -> Int -- ^ y1 (high x and y value, last index to fill) + -> IO () + +{-# INLINE [0] fillVectorBlockP #-} +fillVectorBlockP !vec !getElem !imageWidth !x0 !y0 !x1 !y1 + = gangIO theGang fillBlock + where !threads = gangSize theGang + !blockWidth = x1 - x0 + 1 + + -- All columns have at least this many pixels. + !colChunkLen = blockWidth `quotInt` threads + + -- Extra pixels that we have to divide between some of the threads. + !colChunkSlack = blockWidth `remInt` threads + + -- Get the starting pixel of a column in the image. + {-# INLINE colIx #-} + colIx !ix + | ix < colChunkSlack = x0 + ix * (colChunkLen + 1) + | otherwise = x0 + ix * colChunkLen + colChunkSlack + + -- Give one column to each thread + {-# INLINE fillBlock #-} + fillBlock :: Int -> IO () + fillBlock !ix + = let !x0' = colIx ix + !x1' = colIx (ix + 1) - 1 + !y0' = y0 + !y1' = y1 + in fillVectorBlock vec getElem imageWidth x0' y0' x1' y1' + + +-- | Fill a block in a 2D image. +-- Coordinates given are of the filled edges of the block. +fillVectorBlock + :: Elt a + => IOVector a -- ^ vector to write elements into. + -> (Int -> a) -- ^ fn to evaluate an element at the given index. + -> Int -- ^ width of whole image + -> Int -- ^ x0 lower left corner of block to fill + -> Int -- ^ y0 (low x and y value) + -> Int -- ^ x1 upper right corner of block + -> Int -- ^ y1 (high x and y value, last index to fill) + -> IO () + +{-# INLINE [0] fillVectorBlock #-} +fillVectorBlock !vec !getElemFVB !imageWidth !x0 !y0 !x1 !y1 + = do -- putStrLn $ "fillVectorBlock: " P.++ show (x0, y0, x1, y1) + fillBlock ixStart (ixStart + (x1 - x0)) + where + -- offset from end of one line to the start of the next. + !ixStart = x0 + y0 * imageWidth + !ixFinal = x1 + y1 * imageWidth + + {-# INLINE fillBlock #-} + fillBlock !ixLineStart !ixLineEnd + | ixLineStart > ixFinal = return () + | otherwise + = do fillLine4 ixLineStart + fillBlock (ixLineStart + imageWidth) (ixLineEnd + imageWidth) + + where {-# INLINE fillLine4 #-} + fillLine4 !ix + | ix + 4 > ixLineEnd = fillLine1 ix + | otherwise + = do + let d0 = getElemFVB (ix + 0) + let d1 = getElemFVB (ix + 1) + let d2 = getElemFVB (ix + 2) + let d3 = getElemFVB (ix + 3) + + touch d0 + touch d1 + touch d2 + touch d3 + + VM.unsafeWrite vec (ix + 0) d0 + VM.unsafeWrite vec (ix + 1) d1 + VM.unsafeWrite vec (ix + 2) d2 + VM.unsafeWrite vec (ix + 3) d3 + fillLine4 (ix + 4) + + {-# INLINE fillLine1 #-} + fillLine1 !ix + | ix > ixLineEnd = return () + | otherwise + = do VM.unsafeWrite vec ix (getElemFVB ix) + fillLine1 (ix + 1) diff --git a/Data/Array/Repa/Internals/EvalChunked.hs b/Data/Array/Repa/Internals/EvalChunked.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Repa/Internals/EvalChunked.hs @@ -0,0 +1,68 @@ +-- | Evaluate a vector by breaking it up into linear chunks and filling each chunk +-- in parallel. +{-# LANGUAGE BangPatterns #-} +module Data.Array.Repa.Internals.EvalChunked + ( fillChunkedS + , fillChunkedP) +where +import Data.Array.Repa.Internals.Elt +import Data.Array.Repa.Internals.Gang +import Data.Vector.Unboxed as V +import Data.Vector.Unboxed.Mutable as VM +import GHC.Base (remInt, quotInt) +import Prelude as P + + +-- | Fill a vector sequentially. +fillChunkedS + :: Elt a + => IOVector a -- ^ Vector to fill. + -> (Int -> a) -- ^ Fn to get the value at a given index. + -> IO () + +{-# INLINE [0] fillChunkedS #-} +fillChunkedS !vec !getElem + = fill 0 + where !len = VM.length vec + + fill !ix + | ix >= len = return () + | otherwise + = do VM.unsafeWrite vec ix (getElem ix) + fill (ix + 1) + + +-- | Fill a vector in parallel. +fillChunkedP + :: Unbox a + => IOVector a -- ^ Vector to fill. + -> (Int -> a) -- ^ Fn to get the value at a given index. + -> IO () + +{-# INLINE [0] fillChunkedP #-} +fillChunkedP !vec !getElem + = gangIO theGang + $ \thread -> fill (splitIx thread) (splitIx (thread + 1)) + + where + -- Decide now to split the work across the threads. + -- If the length of the vector doesn't divide evenly among the threads, + -- then the first few get an extra element. + !threads = gangSize theGang + !len = VM.length vec + !chunkLen = len `quotInt` threads + !chunkLeftover = len `remInt` threads + + {-# INLINE splitIx #-} + splitIx thread + | thread < chunkLeftover = thread * (chunkLen + 1) + | otherwise = thread * chunkLen + chunkLeftover + + -- Evaluate the elements of a single chunk. + {-# INLINE fill #-} + fill !ix !end + | ix >= end = return () + | otherwise + = do VM.unsafeWrite vec ix (getElem ix) + fill (ix + 1) end + diff --git a/Data/Array/Repa/Internals/EvalCursored.hs b/Data/Array/Repa/Internals/EvalCursored.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Repa/Internals/EvalCursored.hs @@ -0,0 +1,137 @@ + +{-# LANGUAGE BangPatterns, UnboxedTuples #-} +module Data.Array.Repa.Internals.EvalCursored + ( fillCursoredBlock2P + , fillCursoredBlock2 ) +where +import Data.Array.Repa.Index +import Data.Array.Repa.Internals.Elt +import Data.Array.Repa.Internals.Gang +import Data.Vector.Unboxed.Mutable as VM +import GHC.Base (remInt, quotInt) +import Prelude as P + + +-- Block filling ---------------------------------------------------------------------------------- +-- | Fill a block in a 2D image, in parallel. +-- Coordinates given are of the filled edges of the block. +-- We divide the block into columns, and give one column to each thread. +fillCursoredBlock2P + :: Elt a + => IOVector a -- ^ vector to write elements into + -> (DIM2 -> cursor) -- ^ make a cursor to a particular element + -> (DIM2 -> cursor -> cursor) -- ^ shift the cursor by an offset + -> (cursor -> a) -- ^ fn to evaluate an element at the given index. + -> Int -- ^ width of whole image + -> Int -- ^ x0 lower left corner of block to fill + -> Int -- ^ y0 (low x and y value) + -> Int -- ^ x1 upper right corner of block to fill + -> Int -- ^ y1 (high x and y value, index of last elem to fill) + -> IO () + +{-# INLINE [0] fillCursoredBlock2P #-} +fillCursoredBlock2P + !vec + !makeCursorFCB !shiftCursorFCB !getElemFCB + !imageWidth !x0 !y0 !x1 !y1 + = gangIO theGang fillBlock + where !threads = gangSize theGang + !blockWidth = x1 - x0 + 1 + + -- All columns have at least this many pixels. + !colChunkLen = blockWidth `quotInt` threads + + -- Extra pixels that we have to divide between some of the threads. + !colChunkSlack = blockWidth `remInt` threads + + -- Get the starting pixel of a column in the image. + {-# INLINE colIx #-} + colIx !ix + | ix < colChunkSlack = x0 + ix * (colChunkLen + 1) + | otherwise = x0 + ix * colChunkLen + colChunkSlack + + -- Give one column to each thread + {-# INLINE fillBlock #-} + fillBlock :: Int -> IO () + fillBlock !ix + = let !x0' = colIx ix + !x1' = colIx (ix + 1) - 1 + !y0' = y0 + !y1' = y1 + in fillCursoredBlock2 + vec + makeCursorFCB shiftCursorFCB getElemFCB + imageWidth x0' y0' x1' y1' + + +-- | Fill a block in a 2D image. +-- Coordinates given are of the filled edges of the block. +fillCursoredBlock2 + :: Elt a + => IOVector a -- ^ vector to write elements into. + -> (DIM2 -> cursor) -- ^ make a cursor to a particular element + -> (DIM2 -> cursor -> cursor) -- ^ shift the cursor by an offset + -> (cursor -> a) -- ^ fn to evaluate an element at the given index. + -> Int -- ^ width of whole image + -> Int -- ^ x0 lower left corner of block to fill + -> Int -- ^ y0 (low x and y value) + -> Int -- ^ x1 upper right corner of block to fill + -> Int -- ^ y1 (high x and y value, index of last elem to fill) + -> IO () + +{-# INLINE [0] fillCursoredBlock2 #-} +fillCursoredBlock2 + !vec + !makeCursor !shiftCursor !getElem + !imageWidth !x0 !y0 !x1 !y1 + + = fillBlock y0 + + where {-# INLINE fillBlock #-} + fillBlock !y + | y > y1 = return () + | otherwise + = do fillLine4 x0 + fillBlock (y + 1) + + where {-# INLINE fillLine4 #-} + fillLine4 !x + | x + 4 > x1 = fillLine1 x + | otherwise + = do -- Compute each source cursor based on the previous one so that + -- the variable live ranges in the generated code are shorter. + let srcCur0 = makeCursor (Z :. y :. x) + let srcCur1 = shiftCursor (Z :. 0 :. 1) srcCur0 + let srcCur2 = shiftCursor (Z :. 0 :. 1) srcCur1 + let srcCur3 = shiftCursor (Z :. 0 :. 1) srcCur2 + + -- Get the result value for each cursor. + let val0 = getElem srcCur0 + let val1 = getElem srcCur1 + let val2 = getElem srcCur2 + let val3 = getElem srcCur3 + + -- Ensure that we've computed each of the result values before we + -- write into the array. If the backend code generator can't tell + -- our destination array doesn't alias with the source then writing + -- to it can prevent the sharing of intermediate computations. +{- touch val0 + touch val1 + touch val2 + touch val3 +-} + -- Compute cursor into destination array. + let !dstCur0 = x + y * imageWidth + VM.unsafeWrite vec (dstCur0) val0 + VM.unsafeWrite vec (dstCur0 + 1) val1 + VM.unsafeWrite vec (dstCur0 + 2) val2 + VM.unsafeWrite vec (dstCur0 + 3) val3 + fillLine4 (x + 4) + + {-# INLINE fillLine1 #-} + fillLine1 !x + | x > x1 = return () + | otherwise + = do VM.unsafeWrite vec (x + y * imageWidth) (getElem $ makeCursor (Z :. y :. x)) + fillLine1 (x + 1) + diff --git a/Data/Array/Repa/Internals/Forcing.hs b/Data/Array/Repa/Internals/Forcing.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Repa/Internals/Forcing.hs @@ -0,0 +1,190 @@ +{-# LANGUAGE BangPatterns #-} +module Data.Array.Repa.Internals.Forcing + ( toVector + , toList + , force, force2) +where +import Data.Array.Repa.Internals.EvalChunked +import Data.Array.Repa.Internals.EvalCursored +import Data.Array.Repa.Internals.Elt +import Data.Array.Repa.Internals.Base +import Data.Array.Repa.Index +import Data.Array.Repa.Shape as S +import qualified Data.Vector.Unboxed as V +import qualified Data.Vector.Unboxed.Mutable as VM +import Data.Vector.Unboxed (Vector) +import System.IO.Unsafe + +stage = "Data.Array.Repa.Internals.Forcing" + + +-- Conversions that also force the array ---------------------------------------------------------- +-- | Convert an array to an unboxed `Data.Vector`, forcing it if required. +-- The elements come out in row-major order. +toVector + :: (Shape sh, Elt a) + => Array sh a + -> Vector a +{-# INLINE toVector #-} +toVector arr + = case force arr of + Array _ [Region _ (GenManifest vec)] -> vec + _ -> error $ stage ++ ".toVector: force failed" + + +-- | Convert an array to a list, forcing it if required. +toList :: (Shape sh, Elt a) + => Array sh a + -> [a] + +{-# INLINE toList #-} +toList arr + = V.toList $ toVector arr + + +-- Forcing ---------------------------------------------------------------------------------------- +-- | Force an array, so that it becomes `Manifest`. +-- The array is split into linear chunks and each chunk evaluated in parallel. +force :: (Shape sh, Elt a) + => Array sh a -> Array sh a + +{-# INLINE [2] force #-} +force arr + = unsafePerformIO + $ do (sh, vec) <- forceIO arr + return $ sh `seq` vec `seq` + Array sh [Region RangeAll (GenManifest vec)] + + where forceIO arr' + = case arr' of + -- Don't force an already forced array. + Array sh [Region RangeAll (GenManifest vec)] + -> return (sh, vec) + + Array sh _ + -> do mvec <- VM.unsafeNew (S.size sh) + fillChunkedP mvec (\ix -> arr' `unsafeIndex` fromIndex sh ix) + vec <- V.unsafeFreeze mvec + return (sh, vec) + + +-- | Force an array, so that it becomes `Manifest`. +-- This forcing function is specialised for DIM2 arrays, and does blockwise filling. +force2 :: Elt a => Array DIM2 a -> Array DIM2 a +{-# INLINE [2] force2 #-} +force2 arr + = unsafePerformIO + $ do (sh, vec) <- forceIO2 arr + return $ sh `seq` vec `seq` + Array sh [Region RangeAll (GenManifest vec)] + + where forceIO2 arr' + = arr' `deepSeqArray` + case arr' of + -- Don't force an already forced array. + Array sh [Region RangeAll (GenManifest vec)] + -> return (sh, vec) + + -- Create a vector to hold the new array and load in the regions. + -- NOTE We must specialise this for the common case of two regions to enable + -- fusion for them. If we just have the next case (arbitrary region list) + -- the worker won't fuse with the filling / evaluation code. + Array sh [r1] + -> do mvec <- VM.new (S.size sh) + fillRegion2P mvec sh r1 + vec <- V.unsafeFreeze mvec + return (sh, vec) + + Array sh [r1, r2] + -> do mvec <- VM.new (S.size sh) + fillRegion2P mvec sh r1 + fillRegion2P mvec sh r2 + vec <- V.unsafeFreeze mvec + return (sh, vec) + + -- Create a vector to hold the new array and load in the regions. + Array sh regions + -> do mvec <- VM.new (S.size sh) + mapM_ (fillRegion2P mvec sh) regions + vec <- V.unsafeFreeze mvec + return (sh, vec) + + +-- FillRegion2P ----------------------------------------------------------------------------------- +-- | Fill an array region into a vector. +-- This is specialised for DIM2 regions. +-- The region is evaluated in parallel in a blockwise manner, where each block is +-- evaluated independently and in a separate thread. For delayed or cursored regions +-- access their source elements from the local neighbourhood, this specialised version +-- should given better cache performance than plain `fillRegionP`. +-- +fillRegion2P + :: Elt a + => VM.IOVector a -- ^ Vector to write elements into. + -> DIM2 -- ^ Extent of entire array. + -> Region DIM2 a -- ^ Region to fill. + -> IO () + +{-# INLINE [1] fillRegion2P #-} +fillRegion2P mvec sh@(_ :. height :. width) (Region range gen) + = mvec `seq` height `seq` width `seq` + case range of + RangeAll + -> fillRect2 mvec sh gen + (Rect (Z :. 0 :. 0) + (Z :. height - 1 :. width - 1)) + + RangeRects _ [rect] + -> fillRect2 mvec sh gen rect + + -- Specialise for the common case of 4 rectangles so we get fusion. + -- The following case with mapM_ doesn't fuse because mapM_ isn't completely unrolled. + RangeRects _ [r1, r2, r3, r4] + -> do fillRect2 mvec sh gen r1 + fillRect2 mvec sh gen r2 + fillRect2 mvec sh gen r3 + fillRect2 mvec sh gen r4 + + RangeRects _ rects + -> mapM_ (fillRect2 mvec sh gen) rects + + +-- | Fill a rectangle in a vector. +fillRect2 + :: Elt a + => VM.IOVector a -- ^ Vector to write elements into. + -> DIM2 -- ^ Extent of entire array. + -> Generator DIM2 a -- ^ Generator for array elements. + -> Rect DIM2 -- ^ Rectangle to fill. + -> IO () + +{-# INLINE fillRect2 #-} +fillRect2 mvec (_ :. _ :. width) gen (Rect (Z :. y0 :. x0) (Z :. y1 :. x1)) + = mvec `seq` width `seq` y0 `seq` x0 `seq` y1 `seq` x1 `seq` + case gen of + GenManifest{} + -> error "fillRegion2P: GenManifest, copy elements." + + -- If the region we're filling is just one pixel wide then just fill it + -- in the current thread instead of starting up the whole gang. +{- GenDelayed getElem + | x0 == x1 + -> fillVectorBlock mvec + (getElem . fromIndex sh) + width x0 y0 x1 y1 + + | y0 == y1 + -> fillVectorBlock mvec + (getElem . fromIndex sh) + width x0 y0 x1 y1 + + | otherwise + -> fillVectorBlockP mvec + (getElem . fromIndex sh) + width x0 y0 x1 y1 +-} + -- Cursor based arrays. + GenCursor makeCursor shiftCursor loadElem + -> fillCursoredBlock2P mvec + makeCursor shiftCursor loadElem + width x0 y0 x1 y1 diff --git a/Data/Array/Repa/Internals/Gang.hs b/Data/Array/Repa/Internals/Gang.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Repa/Internals/Gang.hs @@ -0,0 +1,248 @@ +{-# LANGUAGE CPP #-} + +-- | Gang Primitives. +-- Based on DPH code by Roman Leshchinskiy +-- +-- Gang primitives. +-- +#define TRACE_GANG 0 + +module Data.Array.Repa.Internals.Gang + ( Gang, seqGang, forkGang, gangSize, gangIO, gangST, traceGang, traceGangST + , theGang) +where +import GHC.IO +import GHC.ST +import GHC.Conc (forkOnIO) + +import Control.Concurrent.MVar +import Control.Exception (assert) + +import Control.Monad (zipWithM, zipWithM_) +import GHC.Conc (numCapabilities) +import System.IO + +#if TRACE_GANG +import GHC.Exts (traceEvent) +import System.Time ( ClockTime(..), getClockTime ) +#endif + +-- TheGang ---------------------------------------------------------------------------------------- +-- | The gang is shared by all computations. +theGang :: Gang +{-# NOINLINE theGang #-} +theGang = unsafePerformIO $ forkGang numCapabilities + + +-- Requests --------------------------------------------------------------------------------------- +-- | The 'Req' type encapsulates work requests for individual members of a gang. +data Req + -- | Instruct the worker to run the given action then signal it's done + -- by writing to the MVar. + = ReqDo (Int -> IO ()) (MVar ()) + + -- | Tell the worker that we're shutting the gang down. The worker should + -- signal that it's received the equest down by writing to the MVar before + -- returning to its caller (forkGang) + | ReqShutdown (MVar ()) + + +-- | Create a new request for the given action. +newReq :: (Int -> IO ()) -> IO Req +newReq p + = do mv <- newEmptyMVar + return $ ReqDo p mv + + +-- | Block until a thread request has been executed. +-- NOTE: only one thread can wait for the request. +waitReq :: Req -> IO () +waitReq req + = case req of + ReqDo _ varDone -> takeMVar varDone + ReqShutdown varDone -> takeMVar varDone + + +-- Gang ------------------------------------------------------------------------------------------ +-- | A 'Gang' is a group of threads which execute arbitrary work requests. +-- To get the gang to do work, write Req-uest valuesto its MVars +data Gang + = Gang !Int -- Number of 'Gang' threads + [MVar Req] -- One 'MVar' per thread + (MVar Bool) -- Indicates whether the 'Gang' is busy + + +instance Show Gang where + showsPrec p (Gang n _ _) + = showString "<<" + . showsPrec p n + . showString " threads>>" + + +-- | A sequential gang has no threads. +seqGang :: Gang -> Gang +seqGang (Gang n _ mv) = Gang n [] mv + + +-- | The worker thread of a 'Gang'. +-- The threads blocks on the MVar waiting for a work request. +gangWorker :: Int -> MVar Req -> IO () +gangWorker threadId varReq + = do traceGang $ "Worker " ++ show threadId ++ " waiting for request." + req <- takeMVar varReq + + case req of + ReqDo action varDone + -> do traceGang $ "Worker " ++ show threadId ++ " begin" + start <- getGangTime + action threadId + end <- getGangTime + traceGang $ "Worker " ++ show threadId ++ " end (" ++ diffTime start end ++ ")" + + putMVar varDone () + gangWorker threadId varReq + + ReqShutdown varDone + -> do traceGang $ "Worker " ++ show threadId ++ " shutting down." + putMVar varDone () + + +-- | Finaliser for worker threads. +-- We want to shutdown the corresponding thread when it's MVar becomes unreachable. +-- Without this Repa programs can complain about "Blocked indefinitely on an MVar" +-- because worker threads are still blocked on the request MVars when the program ends. +-- Whether the finalizer is called or not is very racey. It happens about 1 in 10 runs +-- when for the repa-edgedetect benchmark, and less often with the others. +-- +-- We're relying on the comment in System.Mem.Weak that says +-- "If there are no other threads to run, the runtime system will check for runnable +-- finalizers before declaring the system to be deadlocked." +-- +-- If we were creating and destroying the gang cleanly we wouldn't need this, but theGang +-- is created with a top-level unsafePerformIO. Hacks beget hacks beget hacks... +-- +finaliseWorker :: MVar Req -> IO () +finaliseWorker varReq + = do varDone <- newEmptyMVar + putMVar varReq (ReqShutdown varDone) + takeMVar varDone + return () + + +-- | Fork a 'Gang' with the given number of threads (at least 1). +forkGang :: Int -> IO Gang +forkGang n + = assert (n > 0) + $ do + -- Create the vars we'll use to issue work requests. + mvs <- sequence . replicate n $ newEmptyMVar + + -- Add finalisers so we can shut the workers down cleanly if they become unreachable. + mapM_ (\var -> addMVarFinalizer var (finaliseWorker var)) mvs + + -- Create all the worker threads + zipWithM_ forkOnIO [0..] + $ zipWith gangWorker [0 .. n-1] mvs + + -- The gang is currently idle. + busy <- newMVar False + + return $ Gang n mvs busy + + +-- | The number of threads in the 'Gang'. +gangSize :: Gang -> Int +gangSize (Gang n _ _) = n + + +-- | Issue work requests for the 'Gang' and wait until they have been executed. +-- If the gang is already busy then just run the action in the +-- requesting thread. +-- +-- TODO: We might want to print a configurable warning that this is happening. +-- +gangIO :: Gang + -> (Int -> IO ()) + -> IO () + +{-# NOINLINE gangIO #-} +gangIO (Gang n mvs busy) p + = do traceGang "gangIO: issuing work requests (SEQ_IF_GANG_BUSY)" + b <- swapMVar busy True + + traceGang $ "gangIO: gang is currently " ++ (if b then "busy" else "idle") + if b + then do + hPutStr stderr + $ unlines [ "Data.Array.Repa: Performing nested parallel computation sequentially." + , " You've probably called the 'force' function while another instance was" + , " already running. This can happen if the second version was suspended due" + , " to lazy evaluation. Use 'deepSeqArray' to ensure that each array is fully" + , " evaluated before you 'force' the next one." + , "" ] + + mapM_ p [0 .. n-1] + + else do + parIO n mvs p + _ <- swapMVar busy False + return () + + +-- | Issue some requests to the worker threads and wait for them to complete. +parIO :: Int -- ^ Number of threads in the gang. + -> [MVar Req] -- ^ Request vars for worker threads. + -> (Int -> IO ()) -- ^ Action to run in all the workers, it's given the ix of + -- the particular worker thread it's running on. + -> IO () + +parIO n mvs p + = do traceGang "parIO: begin" + + start <- getGangTime + reqs <- sequence . replicate n $ newReq p + + traceGang "parIO: issuing requests" + _ <- zipWithM putMVar mvs reqs + + traceGang "parIO: waiting for requests to complete" + mapM_ waitReq reqs + end <- getGangTime + + traceGang $ "parIO: end " ++ diffTime start end + + +-- | Same as 'gangIO' but in the 'ST' monad. +gangST :: Gang -> (Int -> ST s ()) -> ST s () +gangST g p = unsafeIOToST . gangIO g $ unsafeSTToIO . p + + +-- Tracing ---------------------------------------------------------------------------------------- +#if TRACE_GANG +getGangTime :: IO Integer +getGangTime + = do TOD sec pico <- getClockTime + return (pico + sec * 1000000000000) + +diffTime :: Integer -> Integer -> String +diffTime x y = show (y-x) + +traceGang :: String -> IO () +traceGang s + = do t <- getGangTime + traceEvent $ show t ++ " @ " ++ s + +#else +getGangTime :: IO () +getGangTime = return () + +diffTime :: () -> () -> String +diffTime _ _ = "" + +traceGang :: String -> IO () +traceGang _ = return () + +#endif + +traceGangST :: String -> ST s () +traceGangST s = unsafeIOToST (traceGang s) diff --git a/Data/Array/Repa/Internals/Select.hs b/Data/Array/Repa/Internals/Select.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Repa/Internals/Select.hs @@ -0,0 +1,118 @@ +{-# LANGUAGE BangPatterns, ExplicitForAll, ScopedTypeVariables, PatternGuards #-} +module Data.Array.Repa.Internals.Select + (selectChunkedS, selectChunkedP) +where +import Data.Array.Repa.Internals.Gang +import Data.Array.Repa.Shape +import Data.Vector.Unboxed as V +import Data.Vector.Unboxed.Mutable as VM +import GHC.Base (remInt, quotInt) +import Prelude as P +import Control.Monad as P +import Data.IORef + +-- | Select indices matching a predicate +selectChunkedS + :: (Shape sh, Unbox a) + => (sh -> Bool) -- ^ See if this predicate matches. + -> (sh -> a) -- ^ .. and apply fn to the matching index + -> IOVector a -- ^ .. then write the result into the vector. + -> sh -- ^ Extent of indices to apply to predicate. + -> IO Int -- ^ Number of elements written to destination array. + +{-# INLINE selectChunkedS #-} +selectChunkedS match produce !vDst !shSize + = fill 0 0 + where lenSrc = size shSize + lenDst = VM.length vDst + + fill !nSrc !nDst + | nSrc >= lenSrc = return nDst + | nDst >= lenDst = return nDst + + | ixSrc <- fromIndex shSize nSrc + , match ixSrc + = do VM.unsafeWrite vDst nDst (produce ixSrc) + fill (nSrc + 1) (nDst + 1) + + | otherwise + = fill (nSrc + 1) nDst + + +-- | Select indices matching a predicate, in parallel. +-- The array is chunked up, with one chunk being given to each thread. +-- The number of elements in the result array depends on how many threads +-- you're running the program with. +selectChunkedP + :: forall a + . Unbox a + => (Int -> Bool) -- ^ See if this predicate matches. + -> (Int -> a) -- .. and apply fn to the matching index + -> Int -- Extent of indices to apply to predicate. + -> IO [IOVector a] -- Chunks containing array elements. + +{-# INLINE selectChunkedP #-} +selectChunkedP !match !produce !len + = do + -- Make IORefs that the threads will write their result chunks to. + -- We start with a chunk size proportial to the number of threads we have, + -- but the threads themselves can grow the chunks if they run out of space. + refs <- P.replicateM threads + $ do vec <- VM.new $ len `div` threads + newIORef vec + + -- Fire off a thread to fill each chunk. + gangIO theGang + $ \thread -> makeChunk (refs !! thread) + (splitIx thread) + (splitIx (thread + 1) - 1) + + -- Read the result chunks back from the IORefs. + -- If a thread had to grow a chunk, then these might not be the same ones + -- we created back in the first step. + P.mapM readIORef refs + + where -- See how many threads we have available. + !threads = gangSize theGang + !chunkLen = len `quotInt` threads + !chunkLeftover = len `remInt` threads + + + -- Decide where to split the source array. + {-# INLINE splitIx #-} + splitIx thread + | thread < chunkLeftover = thread * (chunkLen + 1) + | otherwise = thread * chunkLen + chunkLeftover + + + -- Fill the given chunk with elements selected from this range of indices. + makeChunk :: IORef (IOVector a) -> Int -> Int -> IO () + makeChunk !ref !ixSrc !ixSrcEnd + = do vecDst <- VM.new (len `div` threads) + vecDst' <- fillChunk ixSrc ixSrcEnd vecDst 0 (VM.length vecDst - 1) + writeIORef ref vecDst' + + + -- The main filling loop. + fillChunk :: Int -> Int -> IOVector a -> Int -> Int -> IO (IOVector a) + fillChunk !ixSrc !ixSrcEnd !vecDst !ixDst !ixDstEnd + -- If we've finished selecting elements, then slice the vector down + -- so it doesn't have any empty space at the end. + | ixSrc >= ixSrcEnd + = return $ VM.slice 0 ixDst vecDst + + -- If we've run out of space in the chunk then grow it some more. + | ixDst >= ixDstEnd + = do let ixDstEnd' = VM.length vecDst * 2 - 1 + vecDst' <- VM.grow vecDst (ixDstEnd + 1) + fillChunk (ixSrc + 1) ixSrcEnd vecDst' (ixDst + 1) ixDstEnd' + + -- We've got a maching element, so add it to the chunk. + | match ixSrc + = do VM.unsafeWrite vecDst ixDst (produce ixSrc) + fillChunk (ixSrc + 1) ixSrcEnd vecDst (ixDst + 1) ixDstEnd + + -- The element doesnt match, so keep going. + | otherwise + = fillChunk (ixSrc + 1) ixSrcEnd vecDst ixDst ixDstEnd + diff --git a/Data/Array/Repa/Operators/IndexSpace.hs b/Data/Array/Repa/Operators/IndexSpace.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Repa/Operators/IndexSpace.hs @@ -0,0 +1,164 @@ +{-# OPTIONS_HADDOCK hide #-} +{-# LANGUAGE TypeOperators, ExplicitForAll, FlexibleContexts #-} + +module Data.Array.Repa.Operators.IndexSpace + ( reshape + , append, (++) + , transpose + , extend + , slice + , backpermute + , backpermuteDft) +where +import Data.Array.Repa.Index +import Data.Array.Repa.Slice +import Data.Array.Repa.Internals.Elt +import Data.Array.Repa.Internals.Base +import Data.Array.Repa.Operators.Traverse +import Data.Array.Repa.Shape as S +import Prelude hiding ((++)) +import qualified Prelude as P + +stage = "Data.Array.Repa.Operators.IndexSpace" + +-- Index space transformations -------------------------------------------------------------------- +-- | Impose a new shape on the elements of an array. +-- The new extent must be the same size as the original, else `error`. +-- +-- TODO: This only works for arrays with a single region. +-- +reshape :: (Shape sh, Shape sh', Elt a) + => sh' + -> Array sh a + -> Array sh' a + +{-# INLINE reshape #-} +reshape sh' arr + | not $ S.size sh' == S.size (extent arr) + = error $ stage P.++ ".reshape: reshaped array will not match size of the original" + +reshape sh' (Array sh [Region RangeAll gen]) + = Array sh' [Region RangeAll gen'] + where gen' = case gen of + GenManifest vec + -> GenManifest vec + + GenCursor makeCursor _ loadElem + -> GenCursor + id + addDim + (loadElem . makeCursor . fromIndex sh . toIndex sh') + +reshape _ _ + = error $ stage P.++ ".reshape: can't reshape a partitioned array" + + +-- | Append two arrays. +-- +append, (++) + :: (Shape sh, Elt a) + => Array (sh :. Int) a + -> Array (sh :. Int) a + -> Array (sh :. Int) a + +{-# INLINE append #-} +append arr1 arr2 + = unsafeTraverse2 arr1 arr2 fnExtent fnElem + where + (_ :. n) = extent arr1 + + fnExtent (sh :. i) (_ :. j) + = sh :. (i + j) + + fnElem f1 f2 (sh :. i) + | i < n = f1 (sh :. i) + | otherwise = f2 (sh :. (i - n)) + +{-# INLINE (++) #-} +(++) arr1 arr2 = append arr1 arr2 + + +-- | Transpose the lowest two dimensions of an array. +-- Transposing an array twice yields the original. +transpose + :: (Shape sh, Elt a) + => Array (sh :. Int :. Int) a + -> Array (sh :. Int :. Int) a + +{-# INLINE transpose #-} +transpose arr + = unsafeTraverse arr + (\(sh :. m :. n) -> (sh :. n :.m)) + (\f -> \(sh :. i :. j) -> f (sh :. j :. i)) + + +-- | Extend an array, according to a given slice specification. +-- (used to be called replicate). +extend + :: ( Slice sl + , Shape (FullShape sl) + , Shape (SliceShape sl) + , Elt e) + => sl + -> Array (SliceShape sl) e + -> Array (FullShape sl) e + +{-# INLINE extend #-} +extend sl arr + = backpermute + (fullOfSlice sl (extent arr)) + (sliceOfFull sl) + arr + +-- | Take a slice from an array, according to a given specification. +slice :: ( Slice sl + , Shape (FullShape sl) + , Shape (SliceShape sl) + , Elt e) + => Array (FullShape sl) e + -> sl + -> Array (SliceShape sl) e + +{-# INLINE slice #-} +slice arr sl + = backpermute + (sliceOfFull sl (extent arr)) + (fullOfSlice sl) + arr + + +-- | Backwards permutation of an array's elements. +-- The result array has the same extent as the original. +backpermute + :: forall sh sh' a + . (Shape sh, Shape sh', Elt a) + => sh' -- ^ Extent of result array. + -> (sh' -> sh) -- ^ Function mapping each index in the result array + -- to an index of the source array. + -> Array sh a -- ^ Source array. + -> Array sh' a + +{-# INLINE backpermute #-} +backpermute newExtent perm arr + = traverse arr (const newExtent) (. perm) + + +-- | Default backwards permutation of an array's elements. +-- If the function returns `Nothing` then the value at that index is taken +-- from the default array (@arrDft@) +backpermuteDft + :: forall sh sh' a + . (Shape sh, Shape sh', Elt a) + => Array sh' a -- ^ Default values (@arrDft@) + -> (sh' -> Maybe sh) -- ^ Function mapping each index in the result array + -- to an index in the source array. + -> Array sh a -- ^ Source array. + -> Array sh' a + +{-# INLINE backpermuteDft #-} +backpermuteDft arrDft fnIndex arrSrc + = fromFunction (extent arrDft) fnElem + where fnElem ix + = case fnIndex ix of + Just ix' -> arrSrc ! ix' + Nothing -> arrDft ! ix diff --git a/Data/Array/Repa/Operators/Interleave.hs b/Data/Array/Repa/Operators/Interleave.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Repa/Operators/Interleave.hs @@ -0,0 +1,110 @@ +{-# OPTIONS_HADDOCK hide #-} +{-# LANGUAGE TypeOperators, PatternGuards #-} + +module Data.Array.Repa.Operators.Interleave + ( interleave2 + , interleave3 + , interleave4) +where +import Data.Array.Repa.Index +import Data.Array.Repa.Internals.Elt +import Data.Array.Repa.Internals.Base +import Data.Array.Repa.Operators.Traverse +import Data.Array.Repa.Shape as S + +-- | Interleave the elements of two arrays. +-- All the input arrays must have the same extent, else `error`. +-- The lowest dimenion of the result array is twice the size of the inputs. +-- +-- @ +-- interleave2 a1 a2 b1 b2 => a1 b1 a2 b2 +-- a3 a4 b3 b4 a3 b3 a4 b4 +-- @ +-- +interleave2 + :: (Shape sh, Elt a) + => Array (sh :. Int) a + -> Array (sh :. Int) a + -> Array (sh :. Int) a + +{-# INLINE interleave2 #-} +interleave2 arr1 arr2 + = arr1 `deepSeqArray` arr2 `deepSeqArray` + unsafeTraverse2 arr1 arr2 shapeFn elemFn + where + shapeFn dim1 dim2 + | dim1 == dim2 + , sh :. len <- dim1 + = sh :. (len * 2) + + | otherwise + = error "Data.Array.Repa.interleave2: arrays must have same extent" + + elemFn get1 get2 (sh :. ix) + = case ix `mod` 3 of + 0 -> get1 (sh :. ix `div` 2) + 1 -> get2 (sh :. ix `div` 2) + _ -> error "Data.Array.Repa.interleave2: this never happens :-P" + + +-- | Interleave the elements of three arrays. +interleave3 + :: (Shape sh, Elt a) + => Array (sh :. Int) a + -> Array (sh :. Int) a + -> Array (sh :. Int) a + -> Array (sh :. Int) a + +{-# INLINE interleave3 #-} +interleave3 arr1 arr2 arr3 + = arr1 `deepSeqArray` arr2 `deepSeqArray` arr3 `deepSeqArray` + unsafeTraverse3 arr1 arr2 arr3 shapeFn elemFn + where + shapeFn dim1 dim2 dim3 + | dim1 == dim2 + , dim1 == dim3 + , sh :. len <- dim1 + = sh :. (len * 3) + + | otherwise + = error "Data.Array.Repa.interleave3: arrays must have same extent" + + elemFn get1 get2 get3 (sh :. ix) + = case ix `mod` 3 of + 0 -> get1 (sh :. ix `div` 3) + 1 -> get2 (sh :. ix `div` 3) + 2 -> get3 (sh :. ix `div` 3) + _ -> error "Data.Array.Repa.interleave3: this never happens :-P" + + +-- | Interleave the elements of four arrays. +interleave4 + :: (Shape sh, Elt a) + => Array (sh :. Int) a + -> Array (sh :. Int) a + -> Array (sh :. Int) a + -> Array (sh :. Int) a + -> Array (sh :. Int) a + +{-# INLINE interleave4 #-} +interleave4 arr1 arr2 arr3 arr4 + = arr1 `deepSeqArray` arr2 `deepSeqArray` arr3 `deepSeqArray` arr4 `deepSeqArray` + unsafeTraverse4 arr1 arr2 arr3 arr4 shapeFn elemFn + where + shapeFn dim1 dim2 dim3 dim4 + | dim1 == dim2 + , dim1 == dim3 + , dim1 == dim4 + , sh :. len <- dim1 + = sh :. (len * 4) + + | otherwise + = error "Data.Array.Repa.interleave4: arrays must have same extent" + + elemFn get1 get2 get3 get4 (sh :. ix) + = case ix `mod` 4 of + 0 -> get1 (sh :. ix `div` 4) + 1 -> get2 (sh :. ix `div` 4) + 2 -> get3 (sh :. ix `div` 4) + 3 -> get4 (sh :. ix `div` 4) + _ -> error "Data.Array.Repa.interleave4: this never happens :-P" diff --git a/Data/Array/Repa/Operators/Mapping.hs b/Data/Array/Repa/Operators/Mapping.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Repa/Operators/Mapping.hs @@ -0,0 +1,108 @@ +{-# OPTIONS_HADDOCK hide #-} +{-# LANGUAGE NoMonomorphismRestriction, PatternGuards #-} + +module Data.Array.Repa.Operators.Mapping + ( map + , zipWith + , (+^) + , (-^) + , (*^) + , (/^)) +where +import Data.Array.Repa.Internals.Elt +import Data.Array.Repa.Internals.Base +import Data.Array.Repa.Shape as S +import qualified Data.Vector.Unboxed as V +import qualified Prelude as P +import Prelude (($), (.), (+), (*), (+), (/), (-)) + +-- | Apply a worker function to each element of an array, yielding a new array with the same extent. +-- +-- This is specialised for arrays of up to four regions, using more breaks fusion. +-- +map :: (Shape sh, Elt a, Elt b) + => (a -> b) + -> Array sh a + -> Array sh b + +{-# INLINE map #-} +map f (Array sh regions) + = Array sh (mapRegions regions) + + where {-# INLINE mapRegions #-} + mapRegions rs + = case rs of + [] -> [] + [r] -> [mapRegion r] + [r1, r2] -> [mapRegion r1, mapRegion r2] + [r1, r2, r3] -> [mapRegion r1, mapRegion r2, mapRegion r3] + [r1, r2, r3, r4] -> [mapRegion r1, mapRegion r2, mapRegion r3, mapRegion r4] + _ -> mapRegions' rs + + mapRegions' rs + = case rs of + [] -> [] + (r : rs') -> mapRegion r : mapRegions' rs' + + {-# INLINE mapRegion #-} + mapRegion (Region range gen) + = Region range (mapGen gen) + + {-# INLINE mapGen #-} + mapGen gen + = case gen of + GenManifest vec + -> GenCursor + P.id + addDim + (\ix -> f $ V.unsafeIndex vec $ S.toIndex sh ix) + + GenCursor makeCursor shiftCursor loadElem + -> GenCursor makeCursor shiftCursor (f . loadElem) + + +-- | Combine two arrays, element-wise, with a binary operator. +-- If the extent of the two array arguments differ, +-- then the resulting array's extent is their intersection. +-- +zipWith :: (Shape sh, Elt a, Elt b, Elt c) + => (a -> b -> c) + -> Array sh a + -> Array sh b + -> Array sh c + +{-# INLINE zipWith #-} +zipWith f arr1 arr2 + | Array sh2 [_] <- arr1 + , Array sh1 [ Region g21 (GenCursor make21 _ load21) + , Region g22 (GenCursor make22 _ load22)] <- arr2 + + = let {-# INLINE load21' #-} + load21' ix = f (arr1 `unsafeIndex` ix) (load21 $ make21 ix) + + {-# INLINE load22' #-} + load22' ix = f (arr1 `unsafeIndex` ix) (load22 $ make22 ix) + + in Array (S.intersectDim sh1 sh2) + [ Region g21 (GenCursor P.id addDim load21') + , Region g22 (GenCursor P.id addDim load22') ] + + | P.otherwise + = let {-# INLINE getElem' #-} + getElem' ix = f (arr1 `unsafeIndex` ix) (arr2 `unsafeIndex` ix) + in fromFunction + (S.intersectDim (extent arr1) (extent arr2)) + getElem' + + +{-# INLINE (+^) #-} +(+^) = zipWith (+) + +{-# INLINE (-^) #-} +(-^) = zipWith (-) + +{-# INLINE (*^) #-} +(*^) = zipWith (*) + +{-# INLINE (/^) #-} +(/^) = zipWith (/) diff --git a/Data/Array/Repa/Operators/Reduction.hs b/Data/Array/Repa/Operators/Reduction.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Repa/Operators/Reduction.hs @@ -0,0 +1,72 @@ +{-# OPTIONS_HADDOCK hide #-} +{-# LANGUAGE ExplicitForAll, TypeOperators #-} + +module Data.Array.Repa.Operators.Reduction + ( fold, foldAll + , sum, sumAll) +where +import Data.Array.Repa.Index +import Data.Array.Repa.Internals.Elt +import Data.Array.Repa.Internals.Base +import Data.Array.Repa.Shape as S +import qualified Data.Vector.Unboxed as V +import Prelude hiding (sum) + + +-- | Sequentially fold the innermost dimension of an array. +-- Combine this with `transpose` to fold any other dimension. +fold :: (Shape sh, Elt a) + => (a -> a -> a) + -> a + -> Array (sh :. Int) a + -> Array sh a + +{-# INLINE fold #-} +fold f x arr + = x `seq` arr `deepSeqArray` + let sh' :. n = extent arr + elemFn i = V.foldl' f x + $ V.map (\ix -> arr ! (i :. ix)) + (V.enumFromTo 0 (n - 1)) + in fromFunction sh' elemFn + + +-- | Sequentially fold all the elements of an array. +foldAll :: (Shape sh, Elt a) + => (a -> a -> a) + -> a + -> Array sh a + -> a + +{-# INLINE foldAll #-} +foldAll f x arr + = V.foldl' f x + $ V.map ((arr !) . (S.fromIndex (extent arr))) + $ V.enumFromTo + 0 + ((S.size $ extent arr) - 1) + + + +-- | Sum the innermost dimension of an array. +sum :: (Shape sh, Elt a, Num a) + => Array (sh :. Int) a + -> Array sh a + +{-# INLINE sum #-} +sum arr = fold (+) 0 arr + + +-- | Sum all the elements of an array. +sumAll :: (Shape sh, Elt a, Num a) + => Array sh a + -> a + +{-# INLINE sumAll #-} +sumAll arr + = V.foldl' (+) 0 + $ V.map ((arr !) . (S.fromIndex (extent arr))) + $ V.enumFromTo + 0 + ((S.size $ extent arr) - 1) + diff --git a/Data/Array/Repa/Operators/Select.hs b/Data/Array/Repa/Operators/Select.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Repa/Operators/Select.hs @@ -0,0 +1,44 @@ +{-# OPTIONS_HADDOCK hide #-} +{-# LANGUAGE BangPatterns #-} + +module Data.Array.Repa.Operators.Select + (select) +where +import Data.Array.Repa.Index +import Data.Array.Repa.Internals.Elt +import Data.Array.Repa.Internals.Base +import Data.Array.Repa.Internals.Select +import qualified Data.Vector.Unboxed as V +import System.IO.Unsafe + + +-- | Produce an array by applying a predicate to a range of integers. +-- If the predicate matches, then use the second function to generate +-- the element. +-- +-- This is a low-level function helpful for writing filtering operations on arrays. +-- Use the integer as the index into the array you're filtering. +-- +select :: Elt a + => (Int -> Bool) -- ^ If the Int matches this predicate, + -> (Int -> a) -- ^ ... then pass it to this fn to produce a value + -> Int -- ^ Range between 0 and this maximum. + -> Array DIM1 a -- ^ Array containing produced values. + +{-# INLINE select #-} +select match produce len + = unsafePerformIO + $ do (sh, vec) <- selectIO + return $ sh `seq` vec `seq` + Array sh [Region RangeAll (GenManifest vec)] + + where {-# INLINE selectIO #-} + selectIO + = do vecs <- selectChunkedP match produce len + vecs' <- mapM V.unsafeFreeze vecs + + -- TODO: avoid copy. + let result = V.concat vecs' + + return (Z :. V.length result, result) + diff --git a/Data/Array/Repa/Operators/Traverse.hs b/Data/Array/Repa/Operators/Traverse.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Repa/Operators/Traverse.hs @@ -0,0 +1,126 @@ +{-# OPTIONS_HADDOCK hide #-} +{-# LANGUAGE ExplicitForAll #-} + +module Data.Array.Repa.Operators.Traverse + ( traverse, unsafeTraverse + , traverse2, unsafeTraverse2 + , traverse3, unsafeTraverse3 + , traverse4, unsafeTraverse4) +where +import Data.Array.Repa.Internals.Elt +import Data.Array.Repa.Internals.Base +import Data.Array.Repa.Shape as S + +-- Generic Traversal ----------------------------------------------------------------------------- +-- | Unstructured traversal. +traverse + :: forall sh sh' a b + . (Shape sh, Shape sh', Elt a) + => Array sh a -- ^ Source array. + -> (sh -> sh') -- ^ Function to produce the extent of the result. + -> ((sh -> a) -> sh' -> b) -- ^ Function to produce elements of the result. + -- It is passed a lookup function to get elements of the source. + -> Array sh' b + +{-# INLINE traverse #-} +traverse arr transExtent newElem + = arr `deepSeqArray` + fromFunction (transExtent (extent arr)) (newElem (arr !)) + + +{-# INLINE unsafeTraverse #-} +unsafeTraverse arr transExtent newElem + = arr `deepSeqArray` + fromFunction (transExtent (extent arr)) (newElem (unsafeIndex arr)) + + +-- | Unstructured traversal over two arrays at once. +traverse2, unsafeTraverse2 + :: forall sh sh' sh'' a b c + . ( Shape sh, Shape sh', Shape sh'' + , Elt a, Elt b, Elt c) + => Array sh a -- ^ First source array. + -> Array sh' b -- ^ Second source array. + -> (sh -> sh' -> sh'') -- ^ Function to produce the extent of the result. + -> ((sh -> a) -> (sh' -> b) + -> (sh'' -> c)) -- ^ Function to produce elements of the result. + -- It is passed lookup functions to get elements of the + -- source arrays. + -> Array sh'' c + +{-# INLINE traverse2 #-} +traverse2 arrA arrB transExtent newElem + = arrA `deepSeqArray` arrB `deepSeqArray` + fromFunction + (transExtent (extent arrA) (extent arrB)) + (newElem (arrA !) (arrB !)) + +{-# INLINE unsafeTraverse2 #-} +unsafeTraverse2 arrA arrB transExtent newElem + = arrA `deepSeqArray` arrB `deepSeqArray` + fromFunction + (transExtent (extent arrA) (extent arrB)) + (newElem (unsafeIndex arrA) (unsafeIndex arrB)) + + +-- | Unstructured traversal over three arrays at once. +traverse3, unsafeTraverse3 + :: forall sh1 sh2 sh3 sh4 + a b c d + . ( Shape sh1, Shape sh2, Shape sh3, Shape sh4 + , Elt a, Elt b, Elt c, Elt d) + => Array sh1 a + -> Array sh2 b + -> Array sh3 c + -> (sh1 -> sh2 -> sh3 -> sh4) + -> ( (sh1 -> a) -> (sh2 -> b) + -> (sh3 -> c) + -> sh4 -> d ) + -> Array sh4 d + +{-# INLINE traverse3 #-} +traverse3 arrA arrB arrC transExtent newElem + = arrA `deepSeqArray` arrB `deepSeqArray` arrC `deepSeqArray` + fromFunction + (transExtent (extent arrA) (extent arrB) (extent arrC)) + (newElem (arrA !) (arrB !) (arrC !)) + +{-# INLINE unsafeTraverse3 #-} +unsafeTraverse3 arrA arrB arrC transExtent newElem + = arrA `deepSeqArray` arrB `deepSeqArray` arrC `deepSeqArray` + fromFunction + (transExtent (extent arrA) (extent arrB) (extent arrC)) + (newElem (unsafeIndex arrA) (unsafeIndex arrB) (unsafeIndex arrC)) + + +-- | Unstructured traversal over four arrays at once. +traverse4, unsafeTraverse4 + :: forall sh1 sh2 sh3 sh4 sh5 + a b c d e + . ( Shape sh1, Shape sh2, Shape sh3, Shape sh4, Shape sh5 + , Elt a, Elt b, Elt c, Elt d, Elt e) + => Array sh1 a + -> Array sh2 b + -> Array sh3 c + -> Array sh4 d + -> (sh1 -> sh2 -> sh3 -> sh4 -> sh5 ) + -> ( (sh1 -> a) -> (sh2 -> b) + -> (sh3 -> c) -> (sh4 -> d) + -> sh5 -> e ) + -> Array sh5 e + +{-# INLINE traverse4 #-} +traverse4 arrA arrB arrC arrD transExtent newElem + = arrA `deepSeqArray` arrB `deepSeqArray` arrC `deepSeqArray` arrD `deepSeqArray` + fromFunction + (transExtent (extent arrA) (extent arrB) (extent arrC) (extent arrD)) + (newElem (arrA !) (arrB !) (arrC !) (arrD !)) + + +{-# INLINE unsafeTraverse4 #-} +unsafeTraverse4 arrA arrB arrC arrD transExtent newElem + = arrA `deepSeqArray` arrB `deepSeqArray` arrC `deepSeqArray` arrD `deepSeqArray` + fromFunction + (transExtent (extent arrA) (extent arrB) (extent arrC) (extent arrD)) + (newElem (unsafeIndex arrA) (unsafeIndex arrB) (unsafeIndex arrC) (unsafeIndex arrD)) + diff --git a/Data/Array/Repa/Properties.hs b/Data/Array/Repa/Properties.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Repa/Properties.hs @@ -0,0 +1,98 @@ +{-# LANGUAGE ScopedTypeVariables #-} + +module Data.Array.Repa.Properties + ( props_DataArrayRepaIndex + , props_DataArrayRepa) +where +import Data.Array.Repa as R +import qualified Data.Array.Repa.Shape as S +import Data.Array.Repa.Arbitrary +import Control.Monad +import Test.QuickCheck +import Prelude as P + +stage = "Data.Array.Repa.Properties" + + +-- Data.Array.Repa.Index -------------------------------------------------------------------------- +-- | QuickCheck properties for "Data.Array.Repa.Index". +props_DataArrayRepaIndex :: [(String, Property)] +props_DataArrayRepaIndex + = [(stage P.++ "." P.++ name, test) | (name, test) + <- [ ("toIndexFromIndex/DIM1", property prop_toIndexFromIndex_DIM1) + , ("toIndexFromIndex/DIM2", property prop_toIndexFromIndex_DIM2) ]] + +prop_toIndexFromIndex_DIM1 sh ix + = (sizeIsValid sh) + ==> (inShape sh ix) + ==> fromIndex sh (toIndex sh ix) == ix + where _types = ( sh :: DIM1 + , ix :: DIM1) + +prop_toIndexFromIndex_DIM2 + = forAll arbitraryShape $ \(sh :: DIM2) -> + forAll (genInShape2 sh) $ \(ix :: DIM2) -> + fromIndex sh (toIndex sh ix) == ix + + + + +-- Data.Array.Repa -------------------------------------------------------------------------------- +-- | QuickCheck properties for "Data.Array.Repa" and its children. +props_DataArrayRepa :: [(String, Property)] +props_DataArrayRepa + = props_DataArrayRepaIndex + P.++ [(stage P.++ "." P.++ name, test) | (name, test) + <- [ ("id_force/DIM5", property prop_id_force_DIM5) + , ("id_toScalarUnit", property prop_id_toScalarUnit) + , ("id_toListFromList/DIM3", property prop_id_toListFromList_DIM3) + , ("id_transpose/DIM4", property prop_id_transpose_DIM4) + , ("reshapeTransposeSize/DIM3", property prop_reshapeTranspose_DIM3) + , ("appendIsAppend/DIM3", property prop_appendIsAppend_DIM3) + , ("sumAllIsSum/DIM3", property prop_sumAllIsSum_DIM3) ]] + + +-- The Eq instance uses fold and zipWith. +prop_id_force_DIM5 + = forAll (arbitrarySmallArray 10) $ \(arr :: Array DIM5 Int) -> + arr == force arr + +prop_id_toScalarUnit (x :: Int) + = toScalar (singleton x) == x + +-- Conversions ------------------------ +prop_id_toListFromList_DIM3 + = forAll (arbitrarySmallShape 10) $ \(sh :: DIM3) -> + forAll (arbitraryListOfLength (S.size sh)) $ \(xx :: [Int]) -> + toList (fromList sh xx) == xx + +-- Index Space Transforms ------------- +prop_id_transpose_DIM4 + = forAll (arbitrarySmallArray 20) $ \(arr :: Array DIM3 Int) -> + transpose (transpose arr) == arr + +-- A reshaped array has the same size and sum as the original +prop_reshapeTranspose_DIM3 + = forAll (arbitrarySmallArray 20) $ \(arr :: Array DIM3 Int) -> + let arr' = transpose arr + sh' = extent arr' + in (S.size $ extent arr) == S.size (extent (reshape sh' arr)) + && (sumAll arr == sumAll arr') + +prop_appendIsAppend_DIM3 + = forAll (arbitrarySmallArray 20) $ \(arr1 :: Array DIM3 Int) -> + sumAll (append arr1 arr1) == (2 * sumAll arr1) + +-- Reductions -------------------------- +prop_sumAllIsSum_DIM3 + = forAll (arbitrarySmallShape 100) $ \(sh :: DIM2) -> + forAll (arbitraryListOfLength (S.size sh)) $ \(xx :: [Int]) -> + sumAll (fromList sh xx) == P.sum xx + + +-- Utils ------------------------------------------------------------------------------------------ +genInShape2 :: DIM2 -> Gen DIM2 +genInShape2 (Z :. yMax :. xMax) + = do y <- liftM (`mod` yMax) $ arbitrary + x <- liftM (`mod` xMax) $ arbitrary + return $ Z :. y :. x diff --git a/Data/Array/Repa/QuickCheck.hs b/Data/Array/Repa/QuickCheck.hs deleted file mode 100644 --- a/Data/Array/Repa/QuickCheck.hs +++ /dev/null @@ -1,21 +0,0 @@ - --- Utils to help with testing. Not exported. -module Data.Array.Repa.QuickCheck - (arbitraryListOfLength) -where -import Test.QuickCheck - - -arbitraryListOfLength - :: Arbitrary a - => Int -> Gen [a] - -arbitraryListOfLength n - | n == 0 = return [] - | otherwise - = do i <- arbitrary - rest <- arbitraryListOfLength (n - 1) - return $ i : rest - - - diff --git a/Data/Array/Repa/Shape.hs b/Data/Array/Repa/Shape.hs --- a/Data/Array/Repa/Shape.hs +++ b/Data/Array/Repa/Shape.hs @@ -11,7 +11,7 @@ class Eq sh => Shape sh where -- | Get the number of dimensions in a shape. - dim :: sh -> Int + rank :: sh -> Int -- | The shape of an array of size zero, with a particular dimensionality. zeroDim :: sh @@ -22,6 +22,8 @@ -- | Compute the intersection of two shapes. intersectDim :: sh -> sh -> sh + -- | Add the coordinates of two shapes componentwise + addDim :: sh -> sh -> sh -- | Get the total number of elements in an array with this shape. size :: sh -> Int @@ -44,7 +46,8 @@ -> sh -- | Check whether an index is within a given shape. - inRange :: sh -- ^ Start index for range. + inShapeRange + :: sh -- ^ Start index for range. -> sh -- ^ Final index for range. -> sh -- ^ Index to check for. -> Bool @@ -67,5 +70,6 @@ -> sh -- ^ Index. -> Bool +{-# INLINE inShape #-} inShape sh ix - = inRange zeroDim sh ix + = inShapeRange zeroDim sh ix diff --git a/Data/Array/Repa/Slice.hs b/Data/Array/Repa/Slice.hs --- a/Data/Array/Repa/Slice.hs +++ b/Data/Array/Repa/Slice.hs @@ -47,31 +47,36 @@ instance Slice Z where + {-# INLINE sliceOfFull #-} sliceOfFull _ _ = Z + + {-# INLINE fullOfSlice #-} fullOfSlice _ _ = Z instance Slice (Any sh) where + {-# INLINE sliceOfFull #-} sliceOfFull _ sh = sh + + {-# INLINE fullOfSlice #-} fullOfSlice _ sh = sh instance Slice sl => Slice (sl :. Int) where + {-# INLINE sliceOfFull #-} sliceOfFull (fsl :. _) (ssl :. _) = sliceOfFull fsl ssl + {-# INLINE fullOfSlice #-} fullOfSlice (fsl :. n) ssl = fullOfSlice fsl ssl :. n instance Slice sl => Slice (sl :. All) where + {-# INLINE sliceOfFull #-} sliceOfFull (fsl :. All) (ssl :. s) = sliceOfFull fsl ssl :. s + {-# INLINE fullOfSlice #-} fullOfSlice (fsl :. All) (ssl :. s) = fullOfSlice fsl ssl :. s - - - - - diff --git a/Data/Array/Repa/Specialised/Dim2.hs b/Data/Array/Repa/Specialised/Dim2.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Repa/Specialised/Dim2.hs @@ -0,0 +1,109 @@ +{-# LANGUAGE BangPatterns #-} + +-- | Functions specialised for arrays of dimension 2. +module Data.Array.Repa.Specialised.Dim2 + ( isInside2 + , isOutside2 + , clampToBorder2 + , makeBordered2) +where +import Data.Array.Repa + + +-- | Check if an index lies inside the given extent. +-- As opposed to `inRange` from "Data.Array.Repa.Index", +-- this is a short-circuited test that checks that lowest dimension first. +isInside2 + :: DIM2 -- ^ Extent of array. + -> DIM2 -- ^ Index to check. + -> Bool + +{-# INLINE isInside2 #-} +isInside2 ex = not . isOutside2 ex + + +-- | Check if an index lies outside the given extent. +-- As opposed to `inRange` from "Data.Array.Repa.Index", +-- this is a short-circuited test that checks the lowest dimension first. +isOutside2 + :: DIM2 -- ^ Extent of array. + -> DIM2 -- ^ Index to check. + -> Bool + +{-# INLINE isOutside2 #-} +isOutside2 (_ :. yLen :. xLen) (_ :. yy :. xx) + | xx < 0 = True + | xx >= xLen = True + | yy < 0 = True + | yy >= yLen = True + | otherwise = False + + +-- | Given the extent of an array, clamp the components of an index so they +-- lie within the given array. Outlying indices are clamped to the index +-- of the nearest border element. +clampToBorder2 + :: DIM2 -- ^ Extent of array. + -> DIM2 -- ^ Index to clamp. + -> DIM2 + +{-# INLINE clampToBorder2 #-} +clampToBorder2 (_ :. yLen :. xLen) (sh :. j :. i) + = clampX j i + where {-# INLINE clampX #-} + clampX !y !x + | x < 0 = clampY y 0 + | x >= xLen = clampY y (xLen - 1) + | otherwise = clampY y x + + {-# INLINE clampY #-} + clampY !y !x + | y < 0 = sh :. 0 :. x + | y >= yLen = sh :. (yLen - 1) :. x + | otherwise = sh :. y :. x + + +-- | Make a 2D partitioned array given two generators, one to produce elements in the +-- border region, and one to produce values in the internal region. +-- The border must be the same width on all sides. +makeBordered2 + :: Elt a + => DIM2 -- ^ Extent of array. + -> Int -- ^ Width of border. + -> Generator DIM2 a -- ^ Generator for border elements. + -> Generator DIM2 a -- ^ Generator for internal elements. + -> Array DIM2 a + +{-# INLINE makeBordered2 #-} +makeBordered2 sh@(_ :. aHeight :. aWidth) borderWidth genInternal genBorder + = let + -- minimum and maximum indicies of values in the inner part of the image. + !xMin = borderWidth + !yMin = borderWidth + !xMax = aWidth - borderWidth - 1 + !yMax = aHeight - borderWidth - 1 + + -- | Range of values where some of the data needed by the stencil is outside the image. + rectsBorder + = [ Rect (Z :. 0 :. 0) (Z :. yMin -1 :. aWidth - 1) -- bot + , Rect (Z :. yMax + 1 :. 0) (Z :. aHeight - 1 :. aWidth - 1) -- top + , Rect (Z :. yMin :. 0) (Z :. yMax :. xMin - 1) -- left + , Rect (Z :. yMin :. xMax + 1) (Z :. yMax :. aWidth - 1) ] -- right + + {-# INLINE inBorder #-} + inBorder = not . inInternal + + -- Range of values where we don't need to worry about the border + rectsInternal + = [ Rect (Z :. yMin :. xMin) (Z :. yMax :. xMax ) ] + + {-# INLINE inInternal #-} + inInternal (Z :. y :. x) + = x >= xMin && x <= xMax + && y >= yMin && y <= yMax + + in Array sh + [ Region (RangeRects inBorder rectsBorder) genInternal + , Region (RangeRects inInternal rectsInternal) genBorder ] + + diff --git a/Data/Array/Repa/Stencil.hs b/Data/Array/Repa/Stencil.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Repa/Stencil.hs @@ -0,0 +1,271 @@ +{-# LANGUAGE MagicHash, PatternGuards, BangPatterns, TemplateHaskell, QuasiQuotes, + ParallelListComp, TypeOperators, ExplicitForAll, ScopedTypeVariables #-} +{-# OPTIONS -Wnot #-} + +-- | Efficient computation of stencil based convolutions. +-- +-- This is specialised for stencils up to 7x7. +-- Due to limitations in the GHC optimiser, using larger stencils doesn't work, and will yield `error` +-- at runtime. We can probably increase the limit if required -- just ask. +-- +-- The focus of the stencil is in the center of the 7x7 tile, which has coordinates (0, 0). +-- All coefficients in the stencil must fit in the tile, so they can be given X,Y coordinates up to +-- +/- 3 positions. The stencil can be any shape, and need not be symmetric -- provided it fits in the 7x7 tile. +-- +module Data.Array.Repa.Stencil + ( Stencil (..) + , Boundary (..) + + -- * Stencil creation. + , makeStencil, makeStencil2 + + -- * Stencil operators. + , mapStencil2, forStencil2 + , mapStencilFrom2, forStencilFrom2 + + -- From Data.Array.Repa.Stencil.Template + , stencil2) +where +import Data.Array.Repa as R +import Data.Array.Repa.Internals.Base as R +import Data.Array.Repa.Stencil.Base +import Data.Array.Repa.Stencil.Template +import Data.Array.Repa.Specialised.Dim2 +import qualified Data.Array.Repa.Shape as S +import qualified Data.Vector.Unboxed as V +import Data.List as List +import GHC.Exts +import GHC.Base +import Debug.Trace + +-- | A index into the flat array. +-- Should be abstract outside the stencil modules. +data Cursor + = Cursor Int + + +-- Wrappers --------------------------------------------------------------------------------------- +-- | Like `mapStencil2` but with the parameters flipped. +forStencil2 + :: Elt a + => Boundary a + -> Array DIM2 a + -> Stencil DIM2 a + -> Array DIM2 a + +{-# INLINE forStencil2 #-} +forStencil2 boundary arr stencil + = mapStencil2 boundary stencil arr + + +-- | Like `mapStencilFrom2` but with the parameters flipped. +forStencilFrom2 + :: (Elt a, Elt b) + => Boundary a + -> Array DIM2 b + -> (b -> a) + -> Stencil DIM2 a + -> Array DIM2 a + +{-# INLINE forStencilFrom2 #-} +forStencilFrom2 boundary arr from stencil + = mapStencilFrom2 boundary stencil arr from + + +-- | Apply a stencil to every element of a 2D array. +-- The array must be manifest else `error`. +mapStencil2 + :: Elt a + => Boundary a + -> Stencil DIM2 a + -> Array DIM2 a + -> Array DIM2 a + +{-# INLINE mapStencil2 #-} +mapStencil2 boundary stencil arr + = mapStencilFrom2 boundary stencil arr id + + +--------------------------------------------------------------------------------------------------- +-- | Apply a stencil to every element of a 2D array. +-- The array must be manifest else `error`. +mapStencilFrom2 + :: (Elt a, Elt b) + => Boundary a -- ^ How to handle the boundary of the array. + -> Stencil DIM2 a -- ^ Stencil to apply. + -> Array DIM2 b -- ^ Array to apply stencil to. + -> (b -> a) -- ^ Apply this function to values read from the array before + -- transforming them with the stencil. + -> Array DIM2 a + +{-# INLINE mapStencilFrom2 #-} +mapStencilFrom2 boundary stencil@(StencilStatic sExtent zero load) arr preConvert + = let (_ :. aHeight :. aWidth) = extent arr + (_ :. sHeight :. sWidth) = sExtent + + sHeight2 = sHeight `div` 2 + sWidth2 = sWidth `div` 2 + + -- minimum and maximum indicies of values in the inner part of the image. + !xMin = sWidth2 + !yMin = sHeight2 + !xMax = aWidth - sWidth2 - 1 + !yMax = aHeight - sHeight2 - 1 + + -- Rectangles ----------------------- + -- range of values where we don't need to worry about the border + rectsInternal + = [ Rect (Z :. yMin :. xMin) (Z :. yMax :. xMax ) ] + + {-# INLINE inInternal #-} + inInternal (Z :. y :. x) + = x >= xMin && x <= xMax + && y >= yMin && y <= yMax + + + -- range of values where some of the data needed by the stencil is outside the image. + rectsBorder + = [ Rect (Z :. 0 :. 0) (Z :. yMin -1 :. aWidth - 1) -- bot + , Rect (Z :. yMax + 1 :. 0) (Z :. aHeight - 1 :. aWidth - 1) -- top + , Rect (Z :. yMin :. 0) (Z :. yMax :. xMin - 1) -- left + , Rect (Z :. yMin :. xMax + 1) (Z :. yMax :. aWidth - 1) ] -- right + + {-# INLINE inBorder #-} + inBorder = not . inInternal + + + -- Cursor functions ---------------- + {-# INLINE makeCursor' #-} + makeCursor' (Z :. y :. x) + = Cursor (x + y * aWidth) + + {-# INLINE shiftCursor' #-} + shiftCursor' ix (Cursor off) + = Cursor + $ case ix of + Z :. y :. x -> off + y * aWidth + x + + {-# INLINE getInner' #-} + getInner' cur + = unsafeAppStencilCursor2 shiftCursor' stencil + arr preConvert cur + + {-# INLINE getBorder' #-} + getBorder' cur + = case boundary of + BoundConst c -> c + BoundClamp -> unsafeAppStencilCursor2_clamp addDim stencil + arr preConvert cur + + in Array (extent arr) + [ Region (RangeRects inBorder rectsBorder) + (GenCursor id addDim getBorder') + + , Region (RangeRects inInternal rectsInternal) + (GenCursor makeCursor' shiftCursor' getInner') ] + + +unsafeAppStencilCursor2 + :: (Elt a, Elt b) + => (DIM2 -> Cursor -> Cursor) + -> Stencil DIM2 a + -> Array DIM2 b + -> (b -> a) + -> Cursor + -> a + +{-# INLINE [1] unsafeAppStencilCursor2 #-} +unsafeAppStencilCursor2 shift + stencil@(StencilStatic sExtent zero load) + arr@(Array aExtent [Region RangeAll (GenManifest vec)]) preConvert + cur@(Cursor off) + + | _ :. sHeight :. sWidth <- sExtent + , _ :. aHeight :. aWidth <- aExtent + , sHeight <= 7, sWidth <= 7 + = let + -- Get data from the manifest array. + {-# INLINE [0] getData #-} + getData (Cursor cur) = preConvert $ vec `V.unsafeIndex` cur + + -- Build a function to pass data from the array to our stencil. + {-# INLINE oload #-} + oload oy ox + = let !cur' = shift (Z :. oy :. ox) cur + in load (Z :. oy :. ox) (getData cur') + + in template7x7 oload zero + + +-- | Like above, but clamp out of bounds array values to the closest real value. +unsafeAppStencilCursor2_clamp + :: forall a b. (Elt a, Elt b) + => (DIM2 -> DIM2 -> DIM2) + -> Stencil DIM2 a + -> Array DIM2 b + -> (b -> a) + -> DIM2 + -> a + +{-# INLINE [1] unsafeAppStencilCursor2_clamp #-} +unsafeAppStencilCursor2_clamp shift + stencil@(StencilStatic sExtent zero load) + arr@(Array aExtent [Region RangeAll (GenManifest vec)]) preConvert + cur + + | _ :. sHeight :. sWidth <- sExtent + , _ :. aHeight :. aWidth <- aExtent + , sHeight <= 7, sWidth <= 7 + = let + -- Get data from the manifest array. + {-# INLINE [0] getData #-} + getData :: DIM2 -> a + getData (Z :. y :. x) + = wrapLoadX x y + + -- TODO: Inlining this into above makes SpecConstr choke + wrapLoadX :: Int -> Int -> a + wrapLoadX !x !y + | x < 0 = wrapLoadY 0 y + | x >= aWidth = wrapLoadY (aWidth - 1) y + | otherwise = wrapLoadY x y + + {-# INLINE wrapLoadY #-} + wrapLoadY :: Int -> Int -> a + wrapLoadY !x !y + | y < 0 = loadXY x 0 + | y >= aHeight = loadXY x (aHeight - 1) + | otherwise = loadXY x y + + {-# INLINE loadXY #-} + loadXY :: Int -> Int -> a + loadXY !x !y + = preConvert $ vec `V.unsafeIndex` (x + y * aWidth) + + -- Build a function to pass data from the array to our stencil. + {-# INLINE oload #-} + oload oy ox + = let !cur' = shift (Z :. oy :. ox) cur + in load (Z :. oy :. ox) (getData cur') + + in template7x7 oload zero + + + +-- | Data template for stencils up to 7x7. +template7x7 + :: (Int -> Int -> a -> a) + -> a -> a + +{-# INLINE [1] template7x7 #-} +template7x7 f zero + = f (-3) (-3) $ f (-3) (-2) $ f (-3) (-1) $ f (-3) 0 $ f (-3) 1 $ f (-3) 2 $ f (-3) 3 + $ f (-2) (-3) $ f (-2) (-2) $ f (-2) (-1) $ f (-2) 0 $ f (-2) 1 $ f (-2) 2 $ f (-2) 3 + $ f (-1) (-3) $ f (-1) (-2) $ f (-1) (-1) $ f (-1) 0 $ f (-1) 1 $ f (-1) 2 $ f (-1) 3 + $ f 0 (-3) $ f 0 (-2) $ f 0 (-1) $ f 0 0 $ f 0 1 $ f 0 2 $ f 0 3 + $ f 1 (-3) $ f 1 (-2) $ f 1 (-1) $ f 1 0 $ f 1 1 $ f 1 2 $ f 1 3 + $ f 2 (-3) $ f 2 (-2) $ f 2 (-1) $ f 2 0 $ f 2 1 $ f 2 2 $ f 2 3 + $ f 3 (-3) $ f 3 (-2) $ f 3 (-1) $ f 3 0 $ f 3 1 $ f 3 2 $ f 3 3 + $ zero + + diff --git a/Data/Array/Repa/Stencil/Base.hs b/Data/Array/Repa/Stencil/Base.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Repa/Stencil/Base.hs @@ -0,0 +1,59 @@ + +-- | Basic definitions for stencil handling. +module Data.Array.Repa.Stencil.Base + ( Boundary (..) + , Stencil (..) + , makeStencil, makeStencil2) +where +import Data.Array.Repa.Internals.Elt +import Data.Array.Repa.Index + +-- | How to handle the case when the stencil lies partly outside the array. +data Boundary a + -- | Treat points outside as having a constant value. + = BoundConst a + + -- | Clamp points outside to the same value as the edge pixel. + | BoundClamp + deriving (Show) + + +-- | Represents a convolution stencil that we can apply to array. Only statically known stencils +-- are supported right now. +data Stencil sh a + + -- | Static stencils are used when the coefficients are fixed, + -- and known at compile time. + = StencilStatic + { stencilExtent :: !sh + , stencilZero :: !a + , stencilAcc :: !(sh -> a -> a -> a) } + + +-- | Make a stencil from a function yielding coefficients at each index. +makeStencil + :: (Elt a, Num a) + => sh -- ^ Extent of stencil. + -> (sh -> Maybe a) -- ^ Get the coefficient at this index. + -> Stencil sh a + +{-# INLINE makeStencil #-} +makeStencil ex getCoeff + = StencilStatic ex 0 + $ \ix val acc + -> case getCoeff ix of + Nothing -> acc + Just coeff -> acc + val * coeff + + +-- | Wrapper for `makeStencil` that requires a DIM2 stencil. +makeStencil2 + :: (Elt a, Num a) + => Int -> Int -- ^ extent of stencil + -> (DIM2 -> Maybe a) -- ^ Get the coefficient at this index. + -> Stencil DIM2 a + +{-# INLINE makeStencil2 #-} +makeStencil2 height width getCoeff + = makeStencil (Z :. height :. width) getCoeff + diff --git a/Data/Array/Repa/Stencil/Template.hs b/Data/Array/Repa/Stencil/Template.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Repa/Stencil/Template.hs @@ -0,0 +1,99 @@ +{-# LANGUAGE TemplateHaskell, QuasiQuotes, ParallelListComp #-} + +-- | Template +module Data.Array.Repa.Stencil.Template + (stencil2) +where +import Data.Array.Repa.Index +import Language.Haskell.TH +import Language.Haskell.TH.Quote +import qualified Data.List as List + +-- | QuasiQuoter for producing a static stencil defintion. +-- +-- A definition like +-- +-- @ +-- [stencil2| 0 1 0 +-- 1 0 1 +-- 0 1 0 |] +-- @ +-- +-- Is converted to: +-- +-- @ +-- makeStencil2 (Z:.3:.3) +-- (\\ix -> case ix of +-- Z :. -1 :. 0 -> Just 1 +-- Z :. 0 :. -1 -> Just 1 +-- Z :. 0 :. 1 -> Just 1 +-- Z :. 1 :. 0 -> Just 1 +-- _ -> Nothing) +-- @ +-- +stencil2 :: QuasiQuoter +stencil2 = QuasiQuoter + { quoteExp = parseStencil2 + , quotePat = undefined + , quoteType = undefined + , quoteDec = undefined } + + +-- | Parse a stencil definition. +-- TODO: make this more robust. +parseStencil2 :: String -> Q Exp +parseStencil2 str + = let + -- Determine the extent of the stencil based on the layout. + -- TODO: make this more robust. In particular, handle blank + -- lines at the start of the definition. + line1 : _ = lines str + sizeX = fromIntegral $ length $ lines str + sizeY = fromIntegral $ length $ words line1 + + -- TODO: this probably doesn't work for stencils who's extents are even. + minX = negate (sizeX `div` 2) + minY = negate (sizeY `div` 2) + maxX = sizeX `div` 2 + maxY = sizeY `div` 2 + + -- List of coefficients for the stencil. + coeffs = (List.map read $ words str) :: [Integer] + + in makeStencil2' sizeX sizeY + $ filter (\(_, _, v) -> v /= 0) + $ [ (fromIntegral y, fromIntegral x, fromIntegral v) + | y <- [minX, minX + 1 .. maxX] + , x <- [minY, minY + 1 .. maxY] + | v <- coeffs ] + + +makeStencil2' + :: Integer -> Integer + -> [(Integer, Integer, Integer)] + -> Q Exp + +makeStencil2' sizeX sizeY coeffs + = do let makeStencil' = mkName "makeStencil2" + let dot' = mkName ":." + let just' = mkName "Just" + ix' <- newName "ix" + z' <- [p| Z |] + coeffs' <- newName "coeffs" + + let fnCoeffs + = LamE [VarP ix'] + $ CaseE (VarE ix') + $ [ Match (InfixP (InfixP z' dot' (LitP (IntegerL oy))) dot' (LitP (IntegerL ox))) + (NormalB $ ConE just' `AppE` LitE (IntegerL v)) + [] | (oy, ox, v) <- coeffs ] + ++ [Match WildP + (NormalB $ ConE (mkName "Nothing")) []] + + return + $ AppE (VarE makeStencil' `AppE` (LitE (IntegerL sizeX)) `AppE` (LitE (IntegerL sizeY))) + $ LetE [ PragmaD (InlineP coeffs' (InlineSpec True False Nothing)) + , ValD (VarP coeffs') (NormalB fnCoeffs) [] ] + (VarE coeffs') + + diff --git a/repa.cabal b/repa.cabal --- a/repa.cabal +++ b/repa.cabal @@ -1,5 +1,5 @@ Name: repa -Version: 1.1.0.0 +Version: 2.0.0.1 License: BSD3 License-file: LICENSE Author: The DPH Team @@ -11,7 +11,6 @@ Homepage: http://trac.haskell.org/repa Bug-reports: http://trac.haskell.org/repa/newticket Description: - NOTE: You must use the GHC head branch > 6.13.20100309 to get decent performance. Repa provides high performance, regular, multi-dimensional, shape polymorphic parallel arrays. All numeric data is stored unboxed. Functions written with the Repa combinators are automatically parallel provided you supply +RTS -Nwhatever on the command @@ -20,24 +19,46 @@ Synopsis: High performance, regular, shape polymorphic parallel arrays. -Tested-with: GHC == 6.13.20100309, GHC == 6.12.1 +Tested-with: GHC == 7.0.1 Library Build-Depends: base == 4.*, - dph-prim-par == 0.4.*, - dph-prim-seq == 0.4.*, - QuickCheck == 2.1.* + ghc-prim == 0.2.*, + vector >= 0.7 && < 0.8, + QuickCheck >= 2.3 && < 2.5, + template-haskell >= 2.5 && < 2.6 ghc-options: - -Odph -Wall -fno-warn-missing-signatures + -Wall -fno-warn-missing-signatures + -Odph + -funbox-strict-fields + -fcpr-off Exposed-modules: Data.Array.Repa Data.Array.Repa.Index Data.Array.Repa.Shape Data.Array.Repa.Slice + Data.Array.Repa.Stencil + Data.Array.Repa.Arbitrary + Data.Array.Repa.Properties + Data.Array.Repa.Specialised.Dim2 Other-modules: - Data.Array.Repa.QuickCheck - + Data.Array.Repa.Operators.IndexSpace + Data.Array.Repa.Operators.Traverse + Data.Array.Repa.Operators.Interleave + Data.Array.Repa.Operators.Mapping + Data.Array.Repa.Operators.Reduction + Data.Array.Repa.Operators.Select + Data.Array.Repa.Internals.Elt + Data.Array.Repa.Internals.Base + Data.Array.Repa.Internals.Gang + Data.Array.Repa.Internals.EvalChunked + Data.Array.Repa.Internals.EvalBlockwise + Data.Array.Repa.Internals.EvalCursored + Data.Array.Repa.Internals.Forcing + Data.Array.Repa.Internals.Select + Data.Array.Repa.Stencil.Base + Data.Array.Repa.Stencil.Template