d447d142c849aa443ad4f6eb87a05ee0395da949
[packages/dph.git] / dph-prim-par / Data / Array / Parallel / Unlifted / Parallel / Combinators.hs
1 {-# LANGUAGE CPP #-}
2 #include "fusion-phases.h"
3
4 -- | Parallel combinators for unlifted arrays.
5 module Data.Array.Parallel.Unlifted.Parallel.Combinators
6 ( mapUP
7 , filterUP
8 , packUP
9 , combineUP, combine2UP
10 , zipWithUP
11 , foldUP, foldlUP, fold1UP, foldl1UP
12 , scanUP)
13 where
14 import Data.Array.Parallel.Base
15 import Data.Array.Parallel.Unlifted.Distributed
16 import Data.Array.Parallel.Unlifted.Parallel.UPSel
17 import Data.Array.Parallel.Unlifted.Sequential.Vector as Seq
18
19 here :: String -> String
20 here s = "Data.Array.Parallel.Unlifted.Parallel.Combinators." Prelude.++ s
21
22
23 -- | Apply a worker to all elements of an array.
24 mapUP :: (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
25 mapUP f xs
26 = splitJoinD theGang (mapD theGang (Seq.map f)) xs
27 {-# INLINE_UP mapUP #-}
28
29
30 -- | Keep elements that match the given predicate.
31 filterUP :: Unbox a => (a -> Bool) -> Vector a -> Vector a
32 filterUP f
33 = joinD theGang unbalanced
34 . mapD theGang (Seq.filter f)
35 . splitD theGang unbalanced
36 {-# INLINE_UP filterUP #-}
37
38
39 -- | Take elements of an array where a flag value is true, and pack them into
40 -- the result.
41 --
42 -- * The souce and flag arrays must have the same length, but this is not checked.
43 --
44 packUP :: Unbox e => Vector e -> Vector Bool -> Vector e
45 packUP xs flags
46 = Seq.fsts . filterUP snd $ Seq.zip xs flags
47 {-# INLINE_UP packUP #-}
48
49
50 -- | Combine two vectors based on a selector.
51 -- If the selector is true then take the element from the first vector,
52 -- otherwise take it from the second.
53 --
54 -- * The data vectors must have enough elements to satisfy the flag vector,
55 -- but this is not checked.
56 --
57 combineUP :: Unbox a => Vector Bool -> Vector a -> Vector a -> Vector a
58 combineUP flags xs ys
59 = checkEq (here "combineUP")
60 ("tags length /= sum of args length")
61 (Seq.length flags) (Seq.length xs + Seq.length ys)
62 $ combine2UP tags (mkUPSelRep2 tags) xs ys
63 where tags = Seq.map (fromBool . not) flags
64 {-# INLINE combineUP #-}
65
66
67 -- | Combine two vectors based on a selector.
68 --
69 -- * The data vectors must have enough elements to satisfy the selector,
70 -- but this is not checked.
71 --
72 combine2UP :: Unbox a => Vector Tag -> UPSelRep2 -> Vector a -> Vector a -> Vector a
73 combine2UP tags rep !xs !ys
74 = checkEq (here "combine2UP")
75 ("tags length /= sum of args length")
76 (Seq.length tags) (Seq.length xs + Seq.length ys)
77 $ joinD theGang balanced
78 $ zipWithD theGang go rep
79 $ splitD theGang balanced tags
80 where go ((i,j), (m,n)) ts
81 = Seq.combine2ByTag ts
82 (Seq.slice (here "combine2UP") xs i m)
83 (Seq.slice (here "combine2UP") ys j n)
84 {-# INLINE_UP combine2UP #-}
85
86
87 -- | Apply a worker function to correponding elements of two arrays.
88 zipWithUP :: (Unbox a, Unbox b, Unbox c)
89 => (a -> b -> c) -> Vector a -> Vector b -> Vector c
90 zipWithUP f xs ys
91 = splitJoinD theGang
92 (mapD theGang (Seq.map (uncurry f)))
93 (Seq.zip xs ys)
94 {-# INLINE_UP zipWithUP #-}
95
96
97 -- | Undirected fold.
98 -- Note that this function has more constraints on its parameters than the
99 -- standard fold function from the Haskell Prelude.
100 --
101 -- * The worker function must be associative.
102 --
103 -- * The provided starting element must be neutral with respect to the worker.
104 -- For example 0 is neutral wrt (+) and 1 is neutral wrt (*).
105 --
106 -- We need these constraints so that we can partition the fold across
107 -- several threads. Each thread folds a chunk of the input vector,
108 -- then we fold together all the results in the main thread.
109 --
110 foldUP :: (Unbox a, DT a) => (a -> a -> a) -> a -> Vector a -> a
111 foldUP f !z xs
112 = foldD theGang f
113 (mapD theGang (Seq.fold f z)
114 (splitD theGang unbalanced xs))
115 {-# INLINE_UP foldUP #-}
116
117
118 -- | Left fold over an array.
119 --
120 -- * If the vector is empty then this returns the provided neural element.
121 --
122 -- * The worker function must be associative.
123 --
124 -- * The provided starting element must be neutral with respect to the worker,
125 -- see `foldUP` for discussion.
126 --
127 foldlUP :: (DT a, Unbox a) => (a -> a -> a) -> a -> Vector a -> a
128 foldlUP f z arr
129 | Seq.null arr = z
130 | otherwise = foldl1UP f arr
131 {-# INLINE_UP foldlUP #-}
132
133
134 -- | Alias for `foldl1UP`
135 fold1UP :: (DT a, Unbox a) => (a -> a -> a) -> Vector a -> a
136 fold1UP = foldl1UP
137 {-# INLINE_UP fold1UP #-}
138
139
140 -- | Left fold over an array, using the first element of the vector as the
141 -- neural element.
142 --
143 -- * If the vector contains no elements then you'll get a bounds-check error.
144 --
145 -- * The worker function must be associative.
146 --
147 -- * The provided starting element must be neutral with respect to the worker,
148 -- see `foldUP` for discussion.
149 --
150 foldl1UP :: (DT a, Unbox a) => (a -> a -> a) -> Vector a -> a
151 foldl1UP f arr
152 = (maybe z (f z)
153 . foldD theGang combine'
154 . mapD theGang (Seq.foldl1Maybe f)
155 . splitD theGang unbalanced) arr
156 where
157 z = Seq.index (here "fold1UP") arr 0
158 combine' (Just x) (Just y) = Just (f x y)
159 combine' (Just x) Nothing = Just x
160 combine' Nothing (Just y) = Just y
161 combine' Nothing Nothing = Nothing
162 {-# INLINE_UP foldl1UP #-}
163
164
165 -- | Prefix scan. Similar to fold, but produce an array of the intermediate states.
166 --
167 -- * The worker function must be associative.
168 --
169 -- * The provided starting element must be neutral with respect to the worker,
170 -- see `foldUP` for discussion.
171 --
172 scanUP :: (DT a, Unbox a) => (a -> a -> a) -> a -> Vector a -> Vector a
173 scanUP f z
174 = splitJoinD theGang go
175 where go xs = let (ds,zs) = unzipD $ mapD theGang (Seq.scanRes f z) xs
176 zs' = fst (scanD theGang f z zs)
177 in zipWithD theGang (Seq.map . f) zs' ds
178 {-# INLINE_UP scanUP #-}
179